diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md new file mode 100644 index 0000000000..e81c668d6b --- /dev/null +++ b/.claude/CLAUDE.md @@ -0,0 +1,204 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## Overview + +The New Relic Python Agent is an Application Performance Monitoring (APM) agent that instruments Python applications for performance monitoring and analytics. It supports Python 3.9+ and monitors applications by wrapping functions, tracking transactions, collecting metrics, and sending telemetry data to New Relic's backend services. + +## Development Commands + +### Running Tests + +Tests are run using `tox` from the root of the repository. The test suite is organized by component/framework. Each component's test suite is a separate subdirectory under `tests/`. To run that test suite, call the tox environment that contains the folder name of the test suite. `tox` will handle calling `cd` into the correct directory (via the `changedir` setting in `tox.ini`). + +```bash +# List test suites available to run +tox -l + +# Run tests for a specific component (from repository root) +tox run -e linux-agent_features-py312-with_extensions + +# Run tests for a framework integration (run `tox -l` for exact env names) +tox run -e python-framework_django-py312-Djangolatest +``` + +To iterate on a single test file with pytest directly, `cd` into that suite's directory first: + +```bash +cd tests/agent_features +pytest test_agent_control_health_check.py -v +``` + +### Running Tests with Coverage + +Running tests with `tox` automatically produces a coverage data file for that environment combination. + +```bash +# Clear out old coverage files +rm .tox/**/.coverage.* +# Run tox to collect coverage for any number of environments +tox run-parallel -e python-framework_falcon-py313-falconlatest,python-framework_falcon-py314-falconlatest +# Combine data files (required, even for 1 input file) +coverage combine .tox/**/.coverage.* +# Generate the XML report, scoped to the file(s) you care about +coverage xml --include=newrelic/hooks/framework_falcon.py +# Read output file +cat coverage.xml +``` + +### Tox Environment Naming Convention + +Tox environments follow this pattern: +`services_required-tests_folder-python_version-library_version[optional]-extensions[optional]` + +Examples: +- `linux-agent_features-py312-with_extensions` +- `postgres-datastore_psycopg-py313-psycopg0302` +- `python-framework_flask-py311-flasklatest` + +### Linting and Formatting + +```bash +# Run ruff linter and formatter (line length is 120; see [tool.ruff] in pyproject.toml) +ruff check --fix && ruff format + +# Run pre-commit hooks manually (ruff hooks are registered at the pre-push stage) +pre-commit run --all-files --hook-stage pre-push +``` + +### Building the Agent + +```bash +# Install the agent in development mode +pip install -e . + +# Install with C extensions explicitly enabled +NEW_RELIC_EXTENSIONS=true pip install -e . + +# Install without C extensions +NEW_RELIC_EXTENSIONS=false pip install -e . +``` + +## Architecture + +### Core Components + +#### 1. **Import Hook System** (`newrelic/api/import_hook.py`) +The agent uses Python's import hook mechanism to automatically instrument third-party libraries. Import hooks are registered for specific modules and fire when those modules are first imported, allowing the agent to wrap functions before they're used. These are registered by calling `_process_module_definition(target, module, function)` where `target` is a string form of the instrumented library's module path, `module` is a string form of the module containing the instrument function under `newrelic.hooks.*`, and `function` is the name of the instrumentation hook to run on that module. + +#### 2. **Instrumentation Hooks** (`newrelic/hooks/`) +Each file in the `hooks/` directory contains instrumentation for a specific library or framework. Hooks use `wrap_function_wrapper` and similar utilities to instrument code without modifying the original source. Each instrument function requires at least 1 import hook which will apply it, made by a call to `_process_module_definition` in the file `newrelic/config.py` under the function `_process_module_builtin_defaults`. + +Pattern: +```python +def instrument_module_name(module): + wrap_function_wrapper(module, 'ClassName.method', wrapper_function) +``` + +#### 3. **API Layer** (`newrelic/api/`) +Provides public APIs for: +- Transaction management (`transaction.py`, `background_task.py`) +- Trace decorators/context managers (`function_trace.py`, `datastore_trace.py`, `external_trace.py`) +- Error tracking (`error_trace.py`) +- Custom instrumentation points + +#### 4. **Core Engine** (`newrelic/core/`) +Contains the core agent logic: +- `agent.py` - Main agent singleton and lifecycle management +- `application.py` - Application instance management +- `*_node.py` - Node types for the transaction trace tree (database, external, function, etc.) +- `config.py` - Configuration management + +#### 5. **Transaction Model** +Transactions are represented as trees of nodes: +- Each trace type (function, database, external call) creates a node +- Nodes track timing, metadata, and relationships +- The root transaction aggregates all nodes and generates metrics +- Transactions are context-local using thread-local or async context storage + +#### 6. **Wrapper Architecture** (`newrelic/common/object_wrapper.py`) +The agent extensively uses function wrapping to inject instrumentation: +- `wrap_function_wrapper()` - Wraps functions/methods +- `FunctionWrapper` - Wrapper object that preserves function metadata +- Wrappers can be nested and maintain proper call order + +#### 7. **Data Collection & Streaming** +- `data_collector.py` - HTTP-based protocol for sending telemetry +- `agent_streaming.py` - gRPC-based infinite tracing for distributed tracing +- Harvest cycle collects and sends metrics periodically (default: 60 seconds) + +### Key Design Patterns + +1. **Lazy Initialization**: The agent must be initialized before importing instrumented libraries for best results, but handles late initialization gracefully. + +2. **Manual Instrumentation API**: Decorators and context managers (`@background_task`, `@function_trace`, `with` blocks) mark transaction and trace boundaries. + +3. **Thread Safety**: Heavily uses thread-local storage and locks for managing per-thread transaction state. + +4. **Async Support**: Special handling for asyncio, gevent, and other async frameworks with context propagation. + +### Test Organization + +Tests are organized by component type: +- `agent_features/` - Core agent functionality +- `agent_unittests/` - Unit tests +- `adapter_*/` - WSGI/ASGI server adapters +- `datastore_*/` - Database client libraries +- `framework_*/` - Web frameworks +- `external_*/` - HTTP client libraries +- `messagebroker_*/` - Message queue libraries +- `mlmodel_*/` - ML/AI framework integrations +- `logger_*/` - Logging framework integrations +- `testing_support/` - Test utilities and fixtures + +### Configuration + +Configuration sources (in order of precedence): +1. Environment variables (`NEW_RELIC_*`) +2. `newrelic.ini` config file +3. Programmatic configuration via `newrelic.agent.global_settings()` + +Common environment variables: +- `NEW_RELIC_LICENSE_KEY` - License key for authentication +- `NEW_RELIC_APP_NAME` - Application name in APM +- `NEW_RELIC_CONFIG_FILE` - Path to config file +- `NEW_RELIC_DEVELOPER_MODE` - Enable developer mode for testing +- `NEW_RELIC_EXTENSIONS` - Control C extension compilation (true/false) + +## Important Conventions + +### Adding New Instrumentation + +When adding support for a new library: + +1. Create a hook file in `newrelic/hooks/` (e.g., `newrelic/hooks/framework_newlib.py`) +2. Implement `instrument_module_name()` function +3. Add test directory under `tests/` with matching name +4. Add tox environment definition in `tox.ini` +5. Register import hooks that trigger instrumentation (via `_process_module_definition` in `newrelic/config.py`) +6. Create tests that validate metrics, traces, and attributes + +### Wrapper Function Signature + +Wrappers should follow this pattern: +```python +def wrapper(wrapped, instance, args, kwargs): + # wrapped: original function + # instance: object instance (for bound methods) or object class (for class methods), or None (for functions and static methods) + # args, kwargs: original call arguments for the wrapped function + return wrapped(*args, **kwargs) +``` + +### Error Handling in Instrumentation + +Instrumentation code must never break the application: +- Wrap instrumentation in try/except blocks +- Log instrumentation errors at debug level +- Always call the original wrapped function + +### Testing Requirements + +- Tests must be runnable via tox +- Use `tests/testing_support/validators/` for validating collected metrics/traces +- Each test directory needs its own `conftest.py` with necessary fixtures diff --git a/.github/containers/firestore/docker-compose.yml b/.github/containers/firestore/docker-compose.yml new file mode 100644 index 0000000000..56c68687b4 --- /dev/null +++ b/.github/containers/firestore/docker-compose.yml @@ -0,0 +1,31 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +services: + firestore: + image: gcr.io/google.com/cloudsdktool/google-cloud-cli:437.0.1-emulators + command: + [ + "/bin/bash", + "-c", + "gcloud emulators firestore start --host-port=0.0.0.0:8080", + ] + ports: + - 8080:8080 + healthcheck: + test: ["CMD-SHELL", "curl -fsS http://localhost:8080/ || exit 1"] + interval: 5s + timeout: 3s + retries: 12 + start_period: 10s diff --git a/.github/containers/rediscluster/Dockerfile b/.github/containers/rediscluster/Dockerfile new file mode 100644 index 0000000000..5ed3d85a3f --- /dev/null +++ b/.github/containers/rediscluster/Dockerfile @@ -0,0 +1,30 @@ +# syntax=docker/dockerfile:1.4 +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +FROM redis:7.0.12 + +COPY <<"EOF" /etc/redis.conf +bind 0.0.0.0 +cluster-enabled yes +cluster-config-file nodes.conf +cluster-node-timeout 5000 +cluster-preferred-endpoint-type hostname +appendonly yes +EOF + +ENV REDIS_PORT=6379 \ + REDIS_ANNOUNCE_HOSTNAME=localhost + +CMD ["sh", "-c", "exec redis-server /etc/redis.conf --port \"$REDIS_PORT\" --cluster-announce-hostname \"$REDIS_ANNOUNCE_HOSTNAME\" --cluster-announce-port \"$REDIS_PORT\""] diff --git a/.github/containers/rediscluster/docker-compose.yml b/.github/containers/rediscluster/docker-compose.yml new file mode 100644 index 0000000000..9771fc3703 --- /dev/null +++ b/.github/containers/rediscluster/docker-compose.yml @@ -0,0 +1,103 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +x-redis-node: &redis-node + build: . + image: redis-cluster-node:local + tmpfs: + - /data + healthcheck: + test: ["CMD-SHELL", "redis-cli -p $$REDIS_PORT ping"] + interval: 2s + timeout: 3s + retries: 15 + +services: + redis1: + <<: *redis-node + environment: + REDIS_PORT: "6379" + REDIS_ANNOUNCE_HOSTNAME: ${REDIS_ANNOUNCE_HOSTNAME:-localhost} + ports: + - 6379:6379 + + redis2: + <<: *redis-node + environment: + REDIS_PORT: "6380" + REDIS_ANNOUNCE_HOSTNAME: ${REDIS_ANNOUNCE_HOSTNAME:-localhost} + ports: + - 6380:6380 + + redis3: + <<: *redis-node + environment: + REDIS_PORT: "6381" + REDIS_ANNOUNCE_HOSTNAME: ${REDIS_ANNOUNCE_HOSTNAME:-localhost} + ports: + - 6381:6381 + + redis4: + <<: *redis-node + environment: + REDIS_PORT: "6382" + REDIS_ANNOUNCE_HOSTNAME: ${REDIS_ANNOUNCE_HOSTNAME:-localhost} + ports: + - 6382:6382 + + redis5: + <<: *redis-node + environment: + REDIS_PORT: "6383" + REDIS_ANNOUNCE_HOSTNAME: ${REDIS_ANNOUNCE_HOSTNAME:-localhost} + ports: + - 6383:6383 + + redis6: + <<: *redis-node + environment: + REDIS_PORT: "6384" + REDIS_ANNOUNCE_HOSTNAME: ${REDIS_ANNOUNCE_HOSTNAME:-localhost} + ports: + - 6384:6384 + + cluster-setup: + image: redis-cluster-node:local + restart: "no" + command: + - bash + - -c + - >- + redis-cli --cluster create + redis1:6379 redis2:6380 redis3:6381 redis4:6382 redis5:6383 redis6:6384 + --cluster-replicas 1 + --cluster-yes && + exec sleep infinity + healthcheck: + test: + [ + "CMD-SHELL", + "redis-cli -h redis1 cluster info | grep -q cluster_state:ok", + ] + interval: 2s + timeout: 3s + retries: 30 + start_period: 30s + depends_on: + redis1: { condition: service_healthy } + redis2: { condition: service_healthy } + redis3: { condition: service_healthy } + redis4: { condition: service_healthy } + redis5: { condition: service_healthy } + redis6: { condition: service_healthy } diff --git a/.github/workflows/build-ci-image.yml b/.github/workflows/build-ci-image.yml index 7a0b3d8403..96c3d3ac8a 100644 --- a/.github/workflows/build-ci-image.yml +++ b/.github/workflows/build-ci-image.yml @@ -22,6 +22,8 @@ on: paths: - ".github/containers/ci/**" - ".github/workflows/build-ci-image.yml" + branches: + - main push: # Trigger rebuilds when pushing relevant files to main. paths: - ".github/containers/ci/**" diff --git a/.github/workflows/delete-ci-images.yml b/.github/workflows/delete-ci-images.yml new file mode 100644 index 0000000000..6032d5f1bf --- /dev/null +++ b/.github/workflows/delete-ci-images.yml @@ -0,0 +1,51 @@ +# Copyright 2010 New Relic, Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +--- +name: Delete CI Images + +on: + schedule: + - cron: "15 17 1 * *" # At 17:15 UTC on the 1st day of every month + workflow_dispatch: # Allow manual trigger + inputs: + dry-run: + description: "Dry run mode" + required: false + default: false + type: boolean + +permissions: + contents: read + packages: write + +concurrency: + group: ${{ github.workflow }} # Only 1 job ever allowed at a time + cancel-in-progress: true + +jobs: + ghcr-cleanup-image: + name: ghcr cleanup action + runs-on: ubuntu-latest + steps: + - uses: dataaxiom/ghcr-cleanup-action@d52806a0dc70b430571a37da1fde39733ffd640f # 1.2.2 + with: + package: newrelic-python-agent-ci + exclude-tags: latest + delete-untagged: true + delete-ghost-images: true + delete-partial-images: true + delete-orphaned-images: true + older-than: 90 days + validate: true + dry-run: ${{ inputs.dry-run }} diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 4933a28050..6999b805e7 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -751,26 +751,6 @@ jobs: options: >- --add-host=host.docker.internal:host-gateway timeout-minutes: 30 - services: - firestore: - # Image set here MUST be repeated down below in options. See comment below. - image: gcr.io/google.com/cloudsdktool/google-cloud-cli:437.0.1-emulators - ports: - - 8080:8080 - # Set health checks to wait until container has started - options: >- - --health-cmd "echo success" - --health-interval 10s - --health-timeout 5s - --health-retries 5 - --health-start-period 5s - gcr.io/google.com/cloudsdktool/google-cloud-cli:437.0.1-emulators /bin/bash -c "gcloud emulators firestore start --host-port=0.0.0.0:8080" || - # This is a very hacky solution. GitHub Actions doesn't provide APIs for setting commands on services, but allows adding arbitrary options. - # --entrypoint won't work as it only accepts an executable and not the [] syntax. - # Instead, we specify the image again the command afterwards like a call to docker create. The result is a few environment variables - # and the original command being appended to our hijacked docker create command. We can avoid any issues by adding || to prevent that - # from every being executed as bash commands. - steps: - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # 6.0.3 @@ -784,6 +764,18 @@ jobs: mkdir -p /github/home/.cache/pip chown -R "$(whoami)" /github/home/.cache/pip + - name: Start firestore + run: | + docker compose \ + -f .github/containers/firestore/docker-compose.yml \ + up -d \ + --wait \ + --wait-timeout 60 \ + && exit 0 + echo "firestore did not become healthy in time" + docker compose -f .github/containers/firestore/docker-compose.yml logs + exit 1 + - name: Get Environments id: get-envs run: | @@ -821,6 +813,11 @@ jobs: if-no-files-found: error retention-days: 1 + - name: Stop firestore + if: always() + run: | + docker compose -f .github/containers/firestore/docker-compose.yml down + grpc: env: TOTAL_GROUPS: 1 @@ -1874,60 +1871,6 @@ jobs: options: >- --add-host=host.docker.internal:host-gateway timeout-minutes: 30 - services: - redis1: - image: hmstepanek/redis-cluster-node:1.0.0 - ports: - - 6379:6379 - - 16379:16379 - options: >- - --add-host=host.docker.internal:host-gateway - - redis2: - image: hmstepanek/redis-cluster-node:1.0.0 - ports: - - 6380:6379 - - 16380:16379 - options: >- - --add-host=host.docker.internal:host-gateway - - redis3: - image: hmstepanek/redis-cluster-node:1.0.0 - ports: - - 6381:6379 - - 16381:16379 - options: >- - --add-host=host.docker.internal:host-gateway - - redis4: - image: hmstepanek/redis-cluster-node:1.0.0 - ports: - - 6382:6379 - - 16382:16379 - options: >- - --add-host=host.docker.internal:host-gateway - - redis5: - image: hmstepanek/redis-cluster-node:1.0.0 - ports: - - 6383:6379 - - 16383:16379 - options: >- - --add-host=host.docker.internal:host-gateway - - redis6: - image: hmstepanek/redis-cluster-node:1.0.0 - ports: - - 6384:6379 - - 16384:16379 - options: >- - --add-host=host.docker.internal:host-gateway - - cluster-setup: - image: hmstepanek/redis-cluster:1.0.0 - options: >- - --add-host=host.docker.internal:host-gateway - steps: - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # 6.0.3 @@ -1941,6 +1884,25 @@ jobs: mkdir -p /github/home/.cache/pip chown -R "$(whoami)" /github/home/.cache/pip + - name: Start rediscluster + env: + REDIS_ANNOUNCE_HOSTNAME: host.docker.internal + run: | + # Build only the redis1 image to prevent fighting, then start the entire cluster + docker compose \ + -f .github/containers/rediscluster/docker-compose.yml \ + build \ + redis1 && \ + docker compose \ + -f .github/containers/rediscluster/docker-compose.yml \ + up -d \ + --wait \ + --wait-timeout 120 \ + && exit 0 + echo "rediscluster did not become healthy in time" + docker compose -f .github/containers/rediscluster/docker-compose.yml logs + exit 1 + - name: Get Environments id: get-envs run: | @@ -1978,6 +1940,11 @@ jobs: if-no-files-found: error retention-days: 1 + - name: Stop rediscluster + if: always() + run: | + docker compose -f .github/containers/rediscluster/docker-compose.yml down + solr: env: TOTAL_GROUPS: 1 diff --git a/.gitignore b/.gitignore index 63a3cb8877..30e4d2c1a1 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,9 @@ .DS_Store .DS_Store/ +# Claude +.claude/ + # Linter megalinter-reports/ diff --git a/newrelic/hooks/mlmodel_anthropic.py b/newrelic/hooks/mlmodel_anthropic.py index f461c9029f..576418a2ea 100644 --- a/newrelic/hooks/mlmodel_anthropic.py +++ b/newrelic/hooks/mlmodel_anthropic.py @@ -37,6 +37,7 @@ "Please report this issue to New Relic Support." ) STREAM_PARSING_FAILURE_LOG_MESSAGE = "Exception occurred in Anthropic instrumentation: Failed to process event stream information. Please report this issue to New Relic Support." +TOKEN_COUNTING_CALLBACK_FAILURE_LOG_MESSAGE = "Exception occurred in llm_token_count_callback for Anthropic %s tokens. Please check your callback implementation and ensure it can handle the provided input. Falling back to token counts from response usage if available." # noqa: S105 _logger = logging.getLogger(__name__) @@ -425,6 +426,8 @@ def _record_completion_error(*, transaction, linking_metadata, completion_id, kw request_model=request_model, llm_metadata=llm_metadata, response_content=None, + # We do not record token counts in error cases, so set all_token_counts to True so the pipeline tokenizer does not run + all_token_counts=True, request_timestamp=request_timestamp, ) except Exception: @@ -447,6 +450,7 @@ def _record_completion_success( request_timestamp=None, time_to_first_token=None, ): + settings = transaction.settings or global_settings() span_id = linking_metadata.get("span.id") trace_id = linking_metadata.get("trace.id") try: @@ -455,10 +459,45 @@ def _record_completion_success( request_temperature = kwargs.get("temperature") request_max_tokens = kwargs.get("max_tokens") - # TODO: Complete token counting - # total_tokens = ( - # (input_tokens + output_tokens) if (input_tokens is not None and output_tokens is not None) else None - # ) + # Token counts default to those reported in the response object if available, + # but the user registered callback below may override them. + # Anthropic does not include a total in usage, so it is always recomputed from the parts below. + response_prompt_tokens = input_tokens + response_completion_tokens = output_tokens + response_total_tokens = None + + # If the user has registered a callback to compute token counts it should always be preferred. + token_count_callback = settings.ai_monitoring.llm_token_count_callback + if token_count_callback: + input_message_content = " ".join( + content + for msg in messages + if ( + content := _extract_message_content( + msg.get("content") if isinstance(msg, dict) else getattr(msg, "content", None) + ) + ) + ) + if input_message_content: + try: + response_prompt_tokens = token_count_callback(request_model, input_message_content) + except Exception: + _logger.exception(TOKEN_COUNTING_CALLBACK_FAILURE_LOG_MESSAGE, "prompt") + response_text = _extract_message_content(response_content) + if response_text: + try: + response_completion_tokens = token_count_callback(response_model, response_text) + except Exception: + _logger.exception(TOKEN_COUNTING_CALLBACK_FAILURE_LOG_MESSAGE, "completion") + + # Prefer the sum of individual counts as the total whenever both are available. + # This ensures consistency in the event that the token counting callback has reported + # different values for prompt or completion tokens. + if response_prompt_tokens and response_completion_tokens: + response_total_tokens = response_prompt_tokens + response_completion_tokens + + all_token_counts = bool(response_prompt_tokens and response_completion_tokens and response_total_tokens) + number_of_messages = len(messages) + (1 if response_content else 0) full_chat_completion_summary_dict = { @@ -474,13 +513,15 @@ def _record_completion_success( "response.model": response_model, "response.choices.finish_reason": stop_reason, "response.number_of_messages": number_of_messages, - # "response.usage.total_tokens": total_tokens, - # "response.usage.prompt_tokens": input_tokens, - # "response.usage.completion_tokens": output_tokens, "timestamp": request_timestamp, "time_to_first_token": time_to_first_token, } + if all_token_counts: + full_chat_completion_summary_dict["response.usage.prompt_tokens"] = response_prompt_tokens + full_chat_completion_summary_dict["response.usage.completion_tokens"] = response_completion_tokens + full_chat_completion_summary_dict["response.usage.total_tokens"] = response_total_tokens + llm_metadata = _get_llm_attributes(transaction) full_chat_completion_summary_dict.update(llm_metadata) transaction.record_custom_event("LlmChatCompletionSummary", full_chat_completion_summary_dict) @@ -496,6 +537,7 @@ def _record_completion_success( request_model=request_model, llm_metadata=llm_metadata, response_content=response_content, + all_token_counts=all_token_counts, request_timestamp=request_timestamp, ) except Exception: @@ -514,6 +556,7 @@ def create_chat_completion_message_event( request_model, llm_metadata, response_content, + all_token_counts, request_timestamp=None, ): try: @@ -530,11 +573,6 @@ def create_chat_completion_message_event( "id": message_id, "span_id": span_id, "trace_id": trace_id, - "token_count": ( - settings.ai_monitoring.llm_token_count_callback(request_model, message_content) - if settings.ai_monitoring.llm_token_count_callback and message_content - else None - ), "role": role, "completion_id": completion_id, "sequence": sequence, @@ -542,6 +580,8 @@ def create_chat_completion_message_event( "vendor": "anthropic", "ingest_source": "Python", } + if all_token_counts: + input_message_dict["token_count"] = 0 if settings.ai_monitoring.record_content.enabled and message_content is not None: input_message_dict["content"] = message_content if request_timestamp: @@ -551,26 +591,14 @@ def create_chat_completion_message_event( transaction.record_custom_event("LlmChatCompletionMessage", input_message_dict) # Record one event for the response - if response_content: + response_text = _extract_message_content(response_content) + if response_text: response_sequence = len(messages) - # response_content may be a plain string (streaming path) or a list of content blocks (non-streaming). - if isinstance(response_content, str): - response_text = response_content - else: - response_text = " ".join( - block.text for block in response_content if getattr(block, "type", None) == "text" - ) - response_message_id = f"{response_id}-{response_sequence}" if response_id else str(uuid.uuid4()) output_message_dict = { "id": response_message_id, "span_id": span_id, "trace_id": trace_id, - "token_count": ( - settings.ai_monitoring.llm_token_count_callback(response_model, response_text) - if settings.ai_monitoring.llm_token_count_callback and response_text - else None - ), "role": "assistant", "completion_id": completion_id, "sequence": response_sequence, @@ -579,6 +607,8 @@ def create_chat_completion_message_event( "ingest_source": "Python", "is_response": True, } + if all_token_counts: + output_message_dict["token_count"] = 0 if settings.ai_monitoring.record_content.enabled and response_text: output_message_dict["content"] = response_text diff --git a/tests/mlmodel_anthropic/test_chat_completion.py b/tests/mlmodel_anthropic/test_chat_completion.py index 435e4d3f26..1e229460c8 100644 --- a/tests/mlmodel_anthropic/test_chat_completion.py +++ b/tests/mlmodel_anthropic/test_chat_completion.py @@ -16,7 +16,7 @@ from conftest import ANTHROPIC_VERSION_METRIC from testing_support.fixtures import override_llm_token_callback_settings, reset_core_stats_engine, validate_attributes from testing_support.ml_testing_utils import ( - add_token_count_to_events, + add_token_counts_to_chat_events, disabled_ai_monitoring_record_content_settings, disabled_ai_monitoring_settings, events_sans_content, @@ -51,6 +51,9 @@ def chat_completion_events(is_streaming): "response.model": "claude-sonnet-4-5-20250929", "request.temperature": 0.7, "request.max_tokens": 100, + "response.usage.prompt_tokens": 16, + "response.usage.completion_tokens": 26, + "response.usage.total_tokens": 42, "response.choices.finish_reason": "end_turn", "vendor": "anthropic", "ingest_source": "Python", @@ -71,6 +74,7 @@ def chat_completion_events(is_streaming): "completion_id": None, "sequence": 0, "response.model": "claude-sonnet-4-5-20250929", + "token_count": 0, "vendor": "anthropic", "ingest_source": "Python", }, @@ -88,6 +92,7 @@ def chat_completion_events(is_streaming): "completion_id": None, "sequence": 1, "response.model": "claude-sonnet-4-5-20250929", + "token_count": 0, "vendor": "anthropic", "is_response": True, "ingest_source": "Python", @@ -238,7 +243,7 @@ def _test(): def test_anthropic_chat_completion_with_token_count( exercise_model, chat_completion_metrics, set_trace_info, chat_completion_events ): - @validate_custom_events(add_token_count_to_events(chat_completion_events)) + @validate_custom_events(add_token_counts_to_chat_events(chat_completion_events)) @validate_custom_event_count(count=3) @validate_transaction_metrics( name="test_anthropic_chat_completion_with_token_count", diff --git a/tests/mlmodel_anthropic/test_chat_completion_error.py b/tests/mlmodel_anthropic/test_chat_completion_error.py index 44f9bdbf98..7e2f9e64c6 100644 --- a/tests/mlmodel_anthropic/test_chat_completion_error.py +++ b/tests/mlmodel_anthropic/test_chat_completion_error.py @@ -19,7 +19,6 @@ from conftest import ANTHROPIC_VERSION_METRIC from testing_support.fixtures import dt_enabled, override_llm_token_callback_settings, reset_core_stats_engine from testing_support.ml_testing_utils import ( - add_token_count_to_events, disabled_ai_monitoring_record_content_settings, events_sans_content, events_with_context_attrs, @@ -69,6 +68,7 @@ "role": "user", "completion_id": None, "sequence": 0, + "token_count": 0, "vendor": "anthropic", "ingest_source": "Python", }, @@ -190,6 +190,7 @@ def _test(): "completion_id": None, "response.model": "does-not-exist", "sequence": 0, + "token_count": 0, "vendor": "anthropic", "ingest_source": "Python", }, @@ -230,7 +231,7 @@ def test_chat_completion_invalid_request_error_invalid_model_with_token_count( custom_metrics=[(ANTHROPIC_VERSION_METRIC, 1)], background_task=True, ) - @validate_custom_events(add_token_count_to_events(expected_events_on_invalid_model_error)) + @validate_custom_events(expected_events_on_invalid_model_error) @validate_custom_event_count(count=2) @background_task(name="test_chat_completion_invalid_request_error_invalid_model_with_token_count") def _test(): @@ -277,6 +278,7 @@ def _test(): "response.model": "claude-4-5-sonnet", "completion_id": None, "sequence": 0, + "token_count": 0, "vendor": "anthropic", "ingest_source": "Python", }, diff --git a/tests/testing_support/ml_testing_utils.py b/tests/testing_support/ml_testing_utils.py index 4ff70c7ed4..8c2c0444f0 100644 --- a/tests/testing_support/ml_testing_utils.py +++ b/tests/testing_support/ml_testing_utils.py @@ -29,6 +29,7 @@ def llm_token_count_callback(model, content): return 105 +# This will be removed once all LLM instrumentations have been converted to use new token count design def add_token_count_to_events(expected_events): events = copy.deepcopy(expected_events) for event in events: @@ -37,6 +38,32 @@ def add_token_count_to_events(expected_events): return events +def add_token_count_to_embedding_events(expected_events): + events = copy.deepcopy(expected_events) + for event in events: + if event[0]["type"] == "LlmEmbedding": + event[1]["response.usage.total_tokens"] = 105 + return events + + +def add_token_count_streaming_events(expected_events): + events = copy.deepcopy(expected_events) + for event in events: + if event[0]["type"] == "LlmChatCompletionMessage": + event[1]["token_count"] = 0 + return events + + +def add_token_counts_to_chat_events(expected_events): + events = copy.deepcopy(expected_events) + for event in events: + if event[0]["type"] == "LlmChatCompletionSummary": + event[1]["response.usage.prompt_tokens"] = 105 + event[1]["response.usage.completion_tokens"] = 105 + event[1]["response.usage.total_tokens"] = 210 + return events + + def events_sans_content(event): new_event = copy.deepcopy(event) for _event in new_event: diff --git a/tox.ini b/tox.ini index 7c05195724..28437fa92c 100644 --- a/tox.ini +++ b/tox.ini @@ -145,7 +145,7 @@ envlist = python-external_urllib3-{py312,py313,py314,py314t,pypy311}-urllib30126, python-framework_aiohttp-{py39,py310,py311,py312,py313,py314,py314t,pypy311}-aiohttp03, python-framework_ariadne-{py39,py310,py311,py312,py313,py314,py314t}-ariadnelatest, - ; python-framework_azurefunctions-{py39,py310,py311,py312}, + python-framework_azurefunctions-{py39,py310,py311,py312}, python-framework_bottle-{py39,py310,py311,py312,py313,py314,py314t,pypy311}-bottle0012, python-framework_cherrypy-{py39,py310,py311,py312,py313,py314,py314t,pypy311}-CherryPylatest, python-framework_django-{py39,py310,py311,py312,py313,py314,py314t}-Djangolatest, @@ -349,8 +349,8 @@ deps = framework_aiohttp-aiohttp03: aiohttp<4 framework_aiohttp-aiohttp030900rc0: aiohttp==3.9.0rc0 framework_ariadne-ariadnelatest: ariadne - ; framework_azurefunctions: azure-functions - ; framework_azurefunctions: requests + framework_azurefunctions: azure-functions + framework_azurefunctions: requests framework_bottle-bottle0012: bottle<0.13.0 framework_bottle: jinja2<3.1 framework_bottle: markupsafe<2.1 @@ -520,7 +520,7 @@ commands = framework_grpc: --grpc_python_out={toxinidir}/tests/framework_grpc/sample_application \ framework_grpc: /{toxinidir}/tests/framework_grpc/sample_application/sample_application.proto - ; framework_azurefunctions: {toxinidir}/.github/scripts/install_azure_functions_worker.sh + framework_azurefunctions: {toxinidir}/.github/scripts/install_azure_functions_worker.sh coverage run -m pytest -v [] @@ -590,7 +590,7 @@ changedir = external_urllib3: tests/external_urllib3 framework_aiohttp: tests/framework_aiohttp framework_ariadne: tests/framework_ariadne - ; framework_azurefunctions: tests/framework_azurefunctions + framework_azurefunctions: tests/framework_azurefunctions framework_bottle: tests/framework_bottle framework_cherrypy: tests/framework_cherrypy framework_django: tests/framework_django