diff --git a/.claude/.codebase-info/.map-state.json b/.claude/.codebase-info/.map-state.json index 0d5046f0..461aa09a 100644 --- a/.claude/.codebase-info/.map-state.json +++ b/.claude/.codebase-info/.map-state.json @@ -2,8 +2,8 @@ "schemaVersion": 1, "tool": "codebase-mapper", "version": "2.15.10", - "mappedAt": "2026-10-20", - "gitCommit": "609bacb464177649bb73b5d4eea56e631c3c91f0", + "mappedAt": "2026-10-21", + "gitCommit": "d1c1bc8a910d3a80a98d08d3a35e49d748611cea", "documents": [ "architecture.md", "coding-style.md", @@ -17,16 +17,16 @@ "tech-landscape.md" ], "hashes": { - "INDEX.md": "fa005d57802311a643c11e92ea8788e62f8d7eb5eb0001952d26e0e91deb9a49", - "architecture.md": "59d1265c89d3b42671b7c47165408625018537c016a72df421113cae3f80bb70", + "INDEX.md": "8fec73a52b4a1a7b5c9f2ab7de1e02f273aa49aeb73dc5e29a4d5e7049b2c3ec", + "architecture.md": "0ff15b04928495b48cdcf40845b5bf9e3bb548c602b734321c4b10f630d83150", "coding-style.md": "b63e8dff06d16a2d882f02624c62f490291317aec8b3032f6d79b3ee035e8567", - "communication.md": "4c49f01a99b25a2f56b2b3ab77464f05c7d70d057aa8e998deffe845afc3e7e6", - "dependencies.md": "22b8d208df66c973659a34cea10011a7ac315ddaa65b243c8aaa3adc4eed1a9e", - "directory-structure.md": "4efb931cdefecf2c733f0ffb905ad519aff3a81e245a978c0b62b2c5cf0ad49a", - "entry-points.md": "e03c248c63823276c21dbbe894b98488643ee401c2db136ad2467a00ba57bc9a", - "modules.md": "cd4f81f8ce74dfc4a7bab15ba91dc569245bc342ddf494845d658b51a0a1dc51", + "communication.md": "0a61b6090f5e67341ad605517111270bdd04caede037e70ec6565de8dcdf5f28", + "dependencies.md": "abadb58027b4092afaaa0efcaec92893eac7bb598318f7f54a542fd77d607dc1", + "directory-structure.md": "c8d654506843622fd394e0a998fddba769e4ec534449ccd6755d312192184bfd", + "entry-points.md": "65ebc321a35f1ae36f587adce8b7b04b71cc4006ec7d518f4c092836c2984705", + "modules.md": "e1185cdddfbcc7a05a69b80661a6c720aa2a1fc10bc445e319983fffb875524e", "onboarding.md": "3d28c671dc095a2aadb297bd7ab200c677779a99c7534e03e55a3d64c935f07e", "patterns.md": "09d04b9b8d079ac4ce9e1f5ddafabb6f3faf53b5ef48dca1c2d62d0875e149b0", - "tech-landscape.md": "919d254fae6fa5212a18072ee0d26a1165ce2aeef861faace6c34b8804c025b1" + "tech-landscape.md": "8a4459f7309dd13fa9e2425507fec35d21b8b3f691bc6a39b6626fffb8d1744d" } } diff --git a/.claude/.codebase-info/INDEX.md b/.claude/.codebase-info/INDEX.md index d18ba749..70a6c336 100644 --- a/.claude/.codebase-info/INDEX.md +++ b/.claude/.codebase-info/INDEX.md @@ -1,12 +1,12 @@ # Codebase Map — Atomic Agents -*Last Updated: 2026-10-20* +*Last Updated: 2026-10-21* Atomic Agents is a lightweight, modular Python framework for building agentic AI applications as composable, schema-driven building blocks (built on Instructor + Pydantic). This repository is a `uv`-workspace **monorepo**: the core framework, a TUI tool installer, a tool library, and examples. -**Stack:** Python ≥3.12 · Instructor · Pydantic v2 · LiteLLM · MCP · Textual · uv + Hatchling · Pyright · `VideoURL` multimodal content · optional OpenInference tracing (Phoenix) +**Stack:** Python ≥3.12 · Instructor · Pydantic v2 · LiteLLM · MCP · Textual · uv + Hatchling · Pyright · `VideoURL` multimodal content · optional OpenTelemetry tracing (GenAI attributes; Phoenix/OpenInference optional) **Shape:** Monorepo — `atomic-agents/` (core lib) · `atomic-assembler/` (CLI) · `atomic-forge/` (tools) · `atomic-examples/` (examples) **Package:** `atomic-agents` v2.14.0 · core import package is `atomic_agents` diff --git a/.claude/.codebase-info/architecture.md b/.claude/.codebase-info/architecture.md index 4db0b973..97a8b425 100644 --- a/.claude/.codebase-info/architecture.md +++ b/.claude/.codebase-info/architecture.md @@ -1,6 +1,6 @@ # Architecture -*Last Updated: 2026-10-15* +*Last Updated: 2026-10-21* ## Summary @@ -50,7 +50,7 @@ Tool install (atomic-assembler TUI): | Context | `atomic-agents/atomic_agents/context/` | `SystemPromptGenerator`, `BaseDynamicContextProvider`, `BaseChatHistory` (pluggable memory contract) + `ChatHistory` with Instructor media and `VideoURL` | | Connectors | `atomic-agents/atomic_agents/connectors/mcp/` | Model Context Protocol tools / resources / prompts | | Utils | `atomic-agents/atomic_agents/utils/` | Token counting (LiteLLM), tool-message formatting | -| Observability (opt-in) | `atomic-agents/atomic_agents/observability/` | OpenInference AGENT/TOOL spans for Phoenix / OpenTelemetry; `observability` extra + `openinference_instrumentor` entry point | +| Observability (opt-in) | `atomic-agents/atomic_agents/observability/` | OpenTelemetry AGENT / TOOL / LLM-attempt spans with GenAI attributes for any backend (OpenInference attributes optional via the `phoenix` extra); `observability` extra = `opentelemetry-api`; `openinference_instrumentor` + `opentelemetry_instrumentor` entry points | | Assembler (CLI) | `atomic-assembler/atomic_assembler/` | Textual TUI to fetch/install forge tools | | Forge (tools) | `atomic-forge/tools/` | 13 standalone tools, each `BaseTool`-based | diff --git a/.claude/.codebase-info/communication.md b/.claude/.codebase-info/communication.md index 7325f5e5..b4fbce37 100644 --- a/.claude/.codebase-info/communication.md +++ b/.claude/.codebase-info/communication.md @@ -1,6 +1,6 @@ # Communication & Integrations -*Last Updated: 2026-10-19* +*Last Updated: 2026-10-21* The framework exposes no HTTP API of its own; "communication" here means how it talks to LLM providers and external tools. @@ -37,15 +37,18 @@ providers and external tools. - Built on Instructor's hook system: `register_hook(event, handler)`, `unregister_hook`, `clear_hooks`. Events include `parse:error`, `completion:kwargs`, `completion:response`, `token:counted`. See the `hooks-example` example and `docs/guides/hooks.md`. -- OpenInference tracing (`atomic_agents.observability`, optional `observability` extra): an AGENT span per - `AtomicAgent` run and a TOOL span per `BaseTool` run, auto-activated by `phoenix.otel.register(auto_instrument=True)` - through `openinference_instrumentor`, by the standard `opentelemetry-instrument` launcher through - `opentelemetry_instrumentor`, or by `AtomicAgentsInstrumentor().instrument(tracer_provider=...)` - for an explicitly configured backend. Both discovery paths share the first installation's provider/config. - Missing optional tracing dependencies are DEBUG-skipped when the launcher passes `skip_dep_check=True`; - manual setup retains its import error and install hint. Choose one configuration path or share the - provider explicitly. LLM spans come from one provider-level OpenInference instrumentor and nest under - the agent span. See `docs/guides/observability.md` and the `phoenix-tracing` example. +- OpenTelemetry tracing (`atomic_agents.observability`, `observability` extra = `opentelemetry-api` only): an AGENT + span per `AtomicAgent` run, a TOOL span per `BaseTool` run and one CLIENT `chat` span per Instructor attempt + (driven by Instructor hooks registered once per client), all with OpenTelemetry GenAI attributes; OpenInference + attributes are added only when its package is importable (`phoenix` extra). Activated by + `phoenix.otel.register(auto_instrument=True)` through `openinference_instrumentor`, by the `opentelemetry-instrument` + launcher through `opentelemetry_instrumentor`, or by `AtomicAgentsInstrumentor().instrument(tracer_provider=..., + config=TracingConfig(...))`. A later install with different explicit values replaces the provider/config in place + with a WARNING; `None` keeps the installed value. Missing `opentelemetry-api` is DEBUG-skipped when the launcher + passes `skip_dep_check=True`; manual setup raises `ModuleNotFoundError` with the install hint. Application hooks on + the same client keep running; `AtomicAgent.clear_hooks()` removes only its own. Set `llm_spans=False` (or + `ATOMIC_AGENTS_TRACE_LLM_SPANS=false`) when a provider instrumentor also emits LLM spans. See + `docs/guides/observability.md` and the `phoenix-tracing` example. - Forge tools that run their async body through a `ThreadPoolExecutor` submit via `contextvars.copy_context().run` so the caller's OpenTelemetry context (and the tracing re-entrancy guard) survives in the pool thread. diff --git a/.claude/.codebase-info/dependencies.md b/.claude/.codebase-info/dependencies.md index ec1ab3c4..a2e76b38 100644 --- a/.claude/.codebase-info/dependencies.md +++ b/.claude/.codebase-info/dependencies.md @@ -1,6 +1,6 @@ # Dependencies -*Last Updated: 2026-10-19* +*Last Updated: 2026-10-21* Declared in `pyproject.toml` (Hatchling build, `uv` workspace; versions locked in `uv.lock`). @@ -12,15 +12,15 @@ Declared in `pyproject.toml` (Hatchling build, `uv` workspace; versions locked i | litellm (≥1.50,<2) | Provider-agnostic token counting & model metadata | | rich (≥13.7.1,<14) | Console formatting | -## Optional extra — `observability` -| Package | Purpose | -|---------|---------| -| openinference-instrumentation (≥0.1.38,<0.2) | `OITracer` / `TraceConfig` used by `atomic_agents.observability` | -| openinference-semantic-conventions (≥0.1.20,<0.2) | Span attribute names (`agent.name` needs ≥0.1.20) | -| opentelemetry-api (≥1.30,<2) | Tracer / context API | +## Optional extras — `observability` and `phoenix` +| Extra | Package | Purpose | +|-------|---------|---------| +| `observability` | opentelemetry-api (≥1.30,<2) | Tracer / context API; the only runtime dependency of `atomic_agents.observability` | +| `phoenix` | opentelemetry-api, openinference-instrumentation (≥0.1.38,<0.2), openinference-semantic-conventions (≥0.1.20,<0.2) | Adds the OpenInference attribute family and `using_session` / `using_user` / `using_metadata`; needed for Arize Phoenix views | -Also listed in the dev group (plus `opentelemetry-sdk`) so CI runs the observability tests. `arize-phoenix-otel` -already brings the two OpenInference packages, so the Phoenix auto path works without the extra. +Users bring their own SDK and exporter (`opentelemetry-sdk`, `opentelemetry-exporter-otlp-proto-http`, ...). The dev +group carries the OpenInference packages plus `opentelemetry-sdk` so the tests run; `arize-phoenix-otel` already +brings the OpenInference packages, so the Phoenix auto path works without the `phoenix` extra. ## Runtime — CLI / tooling | Package | Purpose | diff --git a/.claude/.codebase-info/directory-structure.md b/.claude/.codebase-info/directory-structure.md index b74bb241..1602c3cb 100644 --- a/.claude/.codebase-info/directory-structure.md +++ b/.claude/.codebase-info/directory-structure.md @@ -1,6 +1,6 @@ # Directory Structure -*Last Updated: 2026-10-15* +*Last Updated: 2026-10-21* ## Root Layout @@ -13,7 +13,7 @@ atomic-agents/ # repo root (uv workspace) │ ├── context/ # SystemPromptGenerator, BaseChatHistory/ChatHistory, context providers │ ├── connectors/mcp/ # Model Context Protocol integration │ ├── utils/ # token counter, tool-message formatting -│ └── observability/ # opt-in OpenInference instrumentor (Phoenix / OpenTelemetry) +│ └── observability/ # opt-in OpenTelemetry instrumentor (GenAI attributes; OpenInference optional) │ └── tests/ # pytest suite (agents/, base/, context/, connectors/, observability/, utils/) ├── atomic-assembler/ # Textual TUI (`atomic` command) to install forge tools │ └── atomic_assembler/ # main.py, app.py, screens/, widgets/, utils.py, constants.py diff --git a/.claude/.codebase-info/entry-points.md b/.claude/.codebase-info/entry-points.md index a1037d48..bb5feb31 100644 --- a/.claude/.codebase-info/entry-points.md +++ b/.claude/.codebase-info/entry-points.md @@ -1,6 +1,6 @@ # Entry Points -*Last Updated: 2026-10-19* +*Last Updated: 2026-10-21* ## 1. Library API (the primary entry point) @@ -49,7 +49,7 @@ MCP server tools/resources/prompts become agent components via `atomic_agents.co | `fastapi-memory` | multi-user / multi-session memory behind a FastAPI service | | `persistent-memory` | custom `BaseChatHistory` backend: cross-session recall via stdlib `sqlite3` (no extra deps) | | `hooks-example` | monitoring, error handling, retry via hooks | -| `phoenix-tracing` | Arize Phoenix traces: planner agent -> calculator tool -> answer agent under one CHAIN span, via `register(auto_instrument=True)` | +| `phoenix-tracing` | Arize Phoenix traces: planner agent -> calculator tool -> answer agent under one CHAIN span, via `register(auto_instrument=True)`; depends on `atomic-agents[phoenix]` and sets `ATOMIC_AGENTS_TRACE_LLM_SPANS=false` because the OpenAI instrumentor owns LLM spans there | | `basic-multimodal` / `nested-multimodal` | images / PDFs in (nested) schemas | | `basic-pdf-analysis` | PDF analysis with a multimodal model | | `youtube-summarizer` / `youtube-to-recipe` | transcript → summary / structured recipe | @@ -60,9 +60,10 @@ MCP server tools/resources/prompts become agent components via `atomic_agents.co `atomic_agents.observability:AtomicAgentsInstrumentor` is registered as `atomic_agents` in both `openinference_instrumentor` (Phoenix) and `opentelemetry_instrumentor` (the standard `opentelemetry-instrument` launcher). Both use the same installation state. The launcher needs -its own SDK/exporter setup and the optional OpenInference dependencies; it does not install them -through `opentelemetry-bootstrap`. Set `OTEL_PYTHON_DISABLED_INSTRUMENTATIONS=atomic_agents` to -disable launcher discovery. See `docs/guides/observability.md`. +its own SDK/exporter setup; only `opentelemetry-api` is required for the instrumentor, and +`opentelemetry-bootstrap` does not install Atomic Agents. Set `OTEL_PYTHON_DISABLED_INSTRUMENTATIONS=atomic_agents` +to disable launcher discovery; `ATOMIC_AGENTS_TRACE_*` environment variables configure hiding, size cap and +LLM attempt spans on both automatic paths. See `docs/guides/observability.md`. ## Representative Flow diff --git a/.claude/.codebase-info/modules.md b/.claude/.codebase-info/modules.md index ab7154f6..f70bbe5a 100644 --- a/.claude/.codebase-info/modules.md +++ b/.claude/.codebase-info/modules.md @@ -1,6 +1,6 @@ # Key Modules -*Last Updated: 2026-10-15* +*Last Updated: 2026-10-21* ## Core framework — `atomic-agents/atomic_agents/` @@ -46,19 +46,37 @@ ### observability (opt-in) - **Location:** `atomic-agents/atomic_agents/observability/` -- **Purpose:** OpenInference tracing for Arize Phoenix and any OpenTelemetry backend. `AtomicAgentsInstrumentor` - (`__init__.py`, stdlib-only imports) lazily loads `openinference_tracing.py`, which patches the four - `AtomicAgent.run*` methods into AGENT spans and every `BaseTool` subclass's `run`/`run_async`/`arun` into TOOL - spans (existing classes via a `__subclasses__` walk, future ones via a wrapped `BaseTool.__init_subclass__`). -- **Contract:** `instrument(tracer_provider=, config=TraceConfig)` / `uninstrument()` / - `is_instrumented_by_opentelemetry`; idempotent and fully reversible. Needs the `observability` extra - (`openinference-instrumentation>=0.1.38`); older versions are refused with a warning so Phoenix's - `register(auto_instrument=True)` (which loads the `openinference_instrumentor` entry point) never fails. -- **Span facts:** `agent.name` = `Class[InputSchema, OutputSchema]`; namespaced `atomic_agents.*` attributes - (model, mode, schema names, history size, context provider names, SHA-256 of the static prompt sections); - streams record the final/last partial and `atomic_agents.stream.closed_early`; a `ContextVar` re-entrancy guard - keeps `run -> arun` / `to_thread` delegation to one TOOL span. Nothing in `atomic_agents/__init__.py` imports it. -- **Tests:** `atomic-agents/tests/observability/` (in-memory span exporter). +- **Purpose:** plain OpenTelemetry tracing for any backend (OTLP collector, Grafana Tempo, Sentry, Datadog, + Honeycomb, PostHog, Arize Phoenix, Langfuse). Only `opentelemetry-api` is required; the OpenTelemetry GenAI + semantic-convention attributes are primary and the OpenInference family is an optional layer, auto-on when the + `openinference-instrumentation` package is importable (`phoenix` extra). +- **Modules:** `__init__.py` (loader, `AtomicAgentsInstrumentor`, stdlib-only imports), `tracing.py` (thin facade), + `tracing_config.py` (`TracingConfig`, `PayloadVisibility`, env parsing), `span_payloads.py` (serialize, hide before + serialize, truncate with `.original_length`, bounded/redacted error text), `span_attributes.py` (key constants, + Instructor-to-semconv provider map, None filtering, best-effort attribute collection, OpenInference context seam, + `using_conversation`), `agent_spans.py` (four `AtomicAgent.run*` wrappers, stream cursors, run state), + `tool_spans.py` (`BaseTool` wrappers via `__subclasses__` walk + wrapped `__init_subclass__`, re-entrancy guard), + `installation.py` (install / uninstall / replace-in-place), `llm_attempt_spans.py` (Instructor hook handlers). +- **Contract:** `instrument(tracer_provider=, config=TracingConfig | PayloadVisibility | None, **options)` / + `uninstrument()` / `is_instrumented_by_opentelemetry`. Same arguments: no-op; different provider or config: + replaced in place with a WARNING, `None` keeps the installed value; in-flight spans finish on their original + provider. Hooks are registered once per Instructor client under a lock and removed on uninstall; + `suppress_instrumentation` is honored; hooks never raise into user code. +- **Spans:** `invoke_agent {AgentConfig.name or class}` (`gen_ai.operation.name=invoke_agent`, `gen_ai.agent.name`, + `gen_ai.request.model`, `gen_ai.provider.name`, `gen_ai.conversation.id`, `atomic_agents.agent.class`, `.input`, + `.output`, model/mode/schema/history/provider/prompt-hash keys, usage totals and `atomic_agents.llm.attempts` at end); + `execute_tool {tool.tool_name}` (`gen_ai.tool.*`, `atomic_agents.tool.class`); one `chat {model}` CLIENT span per + Instructor attempt from `completion:kwargs`/`completion:response`/`completion:error` with per-attempt usage + (cached/reasoning tokens for both chat-completion and Responses shapes), `parse:error` as a `validation_error` + event; `error.type` on failures; cancellation = `atomic_agents.cancelled`, never ERROR. Spans always attach to the + current context. `llm_spans=False` keeps totals and drops chat spans. +- **Config:** `TracingConfig(hide_inputs, hide_outputs, max_attribute_length=4096, openinference_attributes=None, + llm_spans=True)`; env `ATOMIC_AGENTS_TRACE_HIDE_INPUTS/_OUTPUTS` (fallback `OPENINFERENCE_HIDE_*`), + `ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH` (tolerant, 4096 fallback), `ATOMIC_AGENTS_TRACE_LLM_SPANS`. Under + `hide_outputs` every exception becomes `: attempts (__REDACTED__)`. Nothing in `atomic_agents/__init__.py` + imports the package. +- **Tests:** `atomic-agents/tests/observability/` (in-memory exporter; `instructor_transport.py` builds real Instructor + clients over `httpx.MockTransport`; `test_instructor_modes.py` covers TOOLS/JSON/JSON_SCHEMA/RESPONSES_TOOLS). ## Subprojects diff --git a/.claude/.codebase-info/tech-landscape.md b/.claude/.codebase-info/tech-landscape.md index fe1776d6..9dcf0ee9 100644 --- a/.claude/.codebase-info/tech-landscape.md +++ b/.claude/.codebase-info/tech-landscape.md @@ -1,6 +1,6 @@ # Technology Landscape -*Last Updated: 2026-10-20* +*Last Updated: 2026-10-21* ## Source-of-Truth Files @@ -40,10 +40,12 @@ - `code-quality.yml` — on push/PR to `main`/`v2.0`, three jobs: `quality` (`uv sync`, `black --check` and `flake8` across all four project directories, a focused Pyright 1.1.414 constructor check, then `pytest --cov=atomic_agents` for the core); `fresh-install` (build the wheel, install - it with the observability extra into a clean venv without the lockfile, import every module under - `atomic_agents`, which is what a new user's `pip install` sees); `lowest-versions` (resolve runtime - dependencies with `uv pip compile --resolution lowest-direct`, install current pytest tooling plus - `opentelemetry-instrumentation==0.51b0` compatible with the OpenTelemetry API floor, then run the + it with the `observability` extra plus `opentelemetry-sdk` into a clean venv without the lockfile, prove + `openinference` is not importable, run a real `instructor.from_openai` client over `httpx.MockTransport` + through a tool and an agent and assert the GenAI span names, usage keys and the absence of every OpenInference + key, then import every module under `atomic_agents`); `lowest-versions` (resolve the runtime dependencies plus + the `observability` and `phoenix` extras with `uv pip compile --resolution lowest-direct`, install current pytest + tooling plus `opentelemetry-instrumentation==0.51b0` compatible with the OpenTelemetry API floor, then run the suite). Observability tests exercise real installed discovery metadata, the auto-loader and enabled/disabled launcher subprocesses with local exporters; no live model/backend calls are needed. - `docs.yml` — on push to `main`: build Sphinx docs, run diff --git a/.github/workflows/code-quality.yml b/.github/workflows/code-quality.yml index 9730e160..ba02e24e 100644 --- a/.github/workflows/code-quality.yml +++ b/.github/workflows/code-quality.yml @@ -84,6 +84,115 @@ jobs: run: | uv venv fresh-install-venv uv pip install --python fresh-install-venv "atomic-agents[observability] @ $(ls dist/*.whl)" + uv pip install --python fresh-install-venv opentelemetry-sdk + + - name: Verify OpenInference is absent + run: | + if fresh-install-venv/bin/python -c "import openinference"; then + echo "OpenInference must not be installed in the api-only environment" + exit 1 + fi + + - name: Trace a tool and agent with only OpenTelemetry + run: | + fresh-install-venv/bin/python - <<'PY' + import httpx + import instructor + import openai + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + + from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BaseTool + from atomic_agents.observability import AtomicAgentsInstrumentor, TracingConfig, using_conversation + + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=TracingConfig(hide_inputs=True)) + + class SmokeInput(BaseIOSchema): + """Text for the fresh-install tracing check.""" + + text: str + + class SmokeOutput(BaseIOSchema): + """Result of the fresh-install tracing check.""" + + text: str + + class SmokeTool(BaseTool[SmokeInput, SmokeOutput]): + def run(self, params: SmokeInput) -> SmokeOutput: + return SmokeOutput(text=params.text) + + def respond(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/v1/chat/completions" + return httpx.Response(200, json={ + "id": "chatcmpl_smoke", "object": "chat.completion", "created": 1, + "model": "response-model", + "choices": [{"index": 0, "finish_reason": "tool_calls", "message": { + "role": "assistant", "content": None, "tool_calls": [{ + "id": "call_smoke", "type": "function", "function": { + "name": "SmokeOutput", "arguments": '{"text": "offline response"}' + } + }] + }}], + "usage": { + "prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5, + "prompt_tokens_details": {"cached_tokens": 1}, + "completion_tokens_details": {"reasoning_tokens": 2}, + }, + }) + + client = instructor.from_openai(openai.OpenAI( + api_key="test", http_client=httpx.Client(transport=httpx.MockTransport(respond)) + )) + + class SmokeAgent(AtomicAgent[SmokeInput, SmokeOutput]): + pass + + assert SmokeTool().run(SmokeInput(text="tool input")).text == "tool input" + agent = SmokeAgent(config=AgentConfig(client=client, model="offline-model")) + with using_conversation("smoke-conversation"): + assert agent.run(SmokeInput(text="private agent input")).text == "offline response" + + spans = exporter.get_finished_spans() + for span in spans: + print(f"{span.name}: {sorted(span.attributes)}") + spans_by_operation = {span.attributes["gen_ai.operation.name"]: span for span in spans} + assert len(spans) == 3 + assert len(spans_by_operation) == 3 + chat_span = spans_by_operation["chat"] + tool_span = spans_by_operation["execute_tool"] + agent_span = spans_by_operation["invoke_agent"] + assert tool_span.name == f"execute_tool {SmokeTool().tool_name}" + assert agent_span.name == "invoke_agent SmokeAgent" + assert tool_span.attributes["gen_ai.tool.name"] == SmokeTool().tool_name + assert tool_span.attributes["atomic_agents.tool.class"] == "SmokeTool" + assert agent_span.attributes["atomic_agents.input"] == "__REDACTED__" + assert agent_span.attributes["gen_ai.agent.name"] == "SmokeAgent" + assert agent_span.attributes["atomic_agents.llm.attempts"] == 1 + assert chat_span.name == "chat offline-model" + assert chat_span.parent.span_id == agent_span.context.span_id + for span in (agent_span, chat_span): + assert span.attributes["gen_ai.conversation.id"] == "smoke-conversation" + assert span.attributes["gen_ai.request.model"] == "offline-model" + assert span.attributes["gen_ai.response.model"] == "response-model" + assert span.attributes["gen_ai.provider.name"] == "openai" + assert span.attributes["gen_ai.usage.input_tokens"] == 3 + assert span.attributes["gen_ai.usage.output_tokens"] == 2 + assert span.attributes["gen_ai.usage.cache_read.input_tokens"] == 1 + assert span.attributes["gen_ai.usage.reasoning.output_tokens"] == 2 + forbidden_keys = { + "input.value", "output.value", "input.mime_type", "output.mime_type", + "agent.name", "tool.name", "tool.description", "tool.parameters", "session.id", + } + for span in spans: + assert forbidden_keys.isdisjoint(span.attributes) + assert not any(key.startswith("openinference.") for key in span.attributes) + client.client.close() + print("api-only tracing ok: tool, agent and chat spans, real Instructor usage, redacted input, no OpenInference keys") + PY - name: Import every subpackage the way a new user would run: | @@ -126,7 +235,7 @@ jobs: - name: Resolve runtime dependencies at their declared floors run: | echo 'opentelemetry-instrumentation==0.51b0' > telemetry-test-requirements.txt - uv pip compile pyproject.toml telemetry-test-requirements.txt --resolution lowest-direct --extra observability -o lowest-requirements.txt + uv pip compile pyproject.toml telemetry-test-requirements.txt --resolution lowest-direct --extra observability --extra phoenix -o lowest-requirements.txt - name: Install the floors plus current test tooling run: | diff --git a/README.md b/README.md index 1db7d87e..408bc5e0 100644 --- a/README.md +++ b/README.md @@ -287,7 +287,7 @@ For full, runnable examples, please refer to the following files in the `atomic- In addition to the quickstart examples, we have more complex examples demonstrating the power of Atomic Agents: - [Hooks System](/atomic-examples/hooks-example/README.md): Comprehensive demonstration of the AtomicAgent hook system for monitoring, error handling, and performance metrics with intelligent retry mechanisms. -- [Phoenix Tracing](/atomic-examples/phoenix-tracing/README.md): Traces agents and tools in Arize Phoenix through the optional OpenInference instrumentor (`pip install "atomic-agents[observability]"`). +- [Phoenix Tracing](/atomic-examples/phoenix-tracing/README.md): Traces agents and tools in Arize Phoenix with the OpenTelemetry instrumentor (`pip install "atomic-agents[phoenix]"`). - [Basic Multimodal](/atomic-examples/basic-multimodal/README.md): Demonstrates how to analyze images with text, focusing on extracting structured information from nutrition labels using GPT-4 Vision capabilities. - [Deep Research](/atomic-examples/deep-research/README.md): An advanced example showing how to perform deep research tasks. - [Orchestration Agent](/atomic-examples/orchestration-agent/README.md): Shows how to create an Orchestrator Agent that intelligently decides between using different tools (search or calculator) based on user input. @@ -298,6 +298,10 @@ In addition to the quickstart examples, we have more complex examples demonstrat For a complete list of examples, see the [examples directory](/atomic-examples/). +## Observability + +Every agent and tool run is an OpenTelemetry span. Works with any OpenTelemetry backend (Grafana, Sentry, Phoenix, Langfuse, ...), including Datadog, Honeycomb, Jaeger or a plain OTLP collector. Install it with `pip install "atomic-agents[observability]"` and see the [Observability guide](docs/guides/observability.md). + ## AI-Assisted Development Building with an AI coding assistant? Atomic Agents ships first-class support so your assistant knows the framework's current API and conventions instead of guessing. diff --git a/atomic-agents/atomic_agents/agents/atomic_agent.py b/atomic-agents/atomic_agents/agents/atomic_agent.py index bf314055..76d2423b 100644 --- a/atomic-agents/atomic_agents/agents/atomic_agent.py +++ b/atomic-agents/atomic_agents/agents/atomic_agent.py @@ -68,6 +68,10 @@ class BasicChatOutputSchema(BaseIOSchema): class AgentConfig(BaseModel): client: instructor.Instructor = Field(..., description="Client for interacting with the language model.") + name: Optional[str] = Field( + default=None, + description="Stable name for this agent in logs and traces, for example the role it plays. Defaults to the class name.", + ) model: str = Field(default="gpt-5-mini", description="The model to use for generating responses.") history: Optional[BaseChatHistory] = Field(default=None, description="History component for storing chat history.") system_prompt_generator: Optional[BaseSystemPromptGenerator] = Field( @@ -190,6 +194,7 @@ def __init__(self, config: AgentConfig): Args: config (AgentConfig): Configuration for the chat agent. """ + self.name = config.name self.client = config.client self.model = config.model self.history = config.history if config.history is not None else ChatHistory() @@ -794,16 +799,11 @@ def clear_hooks(self, event: Optional[str] = None) -> None: Args: event (Optional[str]): The event name to clear, or None to clear all """ - if event: - if event in self._hook_handlers: - # Clear from instructor client first - if hasattr(self.client, "clear"): - self.client.clear(event) - self._hook_handlers[event].clear() - else: - # Clear all hooks - if hasattr(self.client, "clear"): - self.client.clear() + events = [event] if event else list(self._hook_handlers) + for registered_event in events: + for handler in list(self._hook_handlers.get(registered_event, [])): + self.unregister_hook(registered_event, handler) + if not event: self._hook_handlers.clear() def _dispatch_hook(self, event: str, *args, **kwargs) -> None: diff --git a/atomic-agents/atomic_agents/observability/__init__.py b/atomic-agents/atomic_agents/observability/__init__.py index 19bb0e22..254230a0 100644 --- a/atomic-agents/atomic_agents/observability/__init__.py +++ b/atomic-agents/atomic_agents/observability/__init__.py @@ -1,34 +1,34 @@ -""" -Opt-in OpenInference tracing for Atomic Agents. +"""Opt-in tracing for any OpenTelemetry backend using GenAI attributes. -Emits an AGENT span per ``AtomicAgent`` run and a TOOL span per ``BaseTool`` run, rendered by -Arize Phoenix and any other OpenTelemetry backend. Needs the ``observability`` extra: -``pip install "atomic-agents[observability]"``. Phoenix discovers :class:`AtomicAgentsInstrumentor` -through the ``openinference_instrumentor`` entry point, so ``phoenix.otel.register(auto_instrument=True)`` -activates it without further code. See ``docs/guides/observability.md``. +Install the observability extra for tracing. The optional phoenix extra adds +Arize Phoenix integration and OpenInference attributes. """ import importlib.util import logging from importlib.metadata import PackageNotFoundError, version -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING if TYPE_CHECKING: - from openinference.instrumentation import TraceConfig + from atomic_agents.observability.tracing import PayloadVisibility, TracingConfig, using_conversation from opentelemetry.trace import TracerProvider logger = logging.getLogger(__name__) _INSTALL_HINT = ( - "Atomic Agents observability needs the OpenInference packages. " - 'Install them with: pip install "atomic-agents[observability]"' + "Atomic Agents tracing needs opentelemetry-api. " + 'Install it with: pip install "atomic-agents[observability]" ' + '(add "atomic-agents[phoenix]" for Arize Phoenix / OpenInference attributes)' ) _OPENINFERENCE_DISTRIBUTION = "openinference-instrumentation" _MINIMUM_OPENINFERENCE_VERSION = (0, 1, 38) -def _openinference_is_installed() -> bool: - return importlib.util.find_spec("openinference") is not None +def _opentelemetry_is_installed() -> bool: + try: + return importlib.util.find_spec("opentelemetry.trace") is not None + except ModuleNotFoundError: + return False def _openinference_is_too_old() -> bool: @@ -41,21 +41,26 @@ def _openinference_is_too_old() -> bool: def _load_tracing_module(): - if _openinference_is_too_old(): - logger.warning( - "Atomic Agents tracing is disabled: %s %s or newer is required.", - _OPENINFERENCE_DISTRIBUTION, - ".".join(str(part) for part in _MINIMUM_OPENINFERENCE_VERSION), - ) - return None - from atomic_agents.observability import openinference_tracing + from atomic_agents.observability import tracing + + return tracing + + +def __getattr__(name: str) -> object: + if name in ("TracingConfig", "using_conversation"): + return getattr(_load_tracing_module(), name) + raise AttributeError(name) + - return openinference_tracing +def _warn_unknown_options(options: dict[str, object]) -> None: + for key in options: + if key != "skip_dep_check": + logger.warning("Unknown Atomic Agents instrumentation option %s ignored", key) class AtomicAgentsInstrumentor: """ - Installs OpenInference AGENT and TOOL spans around ``AtomicAgent`` and ``BaseTool`` calls. + Installs AGENT, TOOL and LLM-attempt spans around ``AtomicAgent`` and ``BaseTool`` calls. The public surface mirrors OpenTelemetry's ``BaseInstrumentor`` (``instrument`` / ``uninstrument``) without depending on it, because Arize Phoenix only guarantees ``opentelemetry-api`` and @@ -65,31 +70,30 @@ class AtomicAgentsInstrumentor: def instrument( self, *, - tracer_provider: Optional["TracerProvider"] = None, - config: Optional["TraceConfig"] = None, + tracer_provider: "TracerProvider | None" = None, + config: "TracingConfig | PayloadVisibility | None" = None, **options: object, ) -> None: """ Patch the framework classes. Safe to call more than once; a second call with different - arguments logs a warning and keeps the first installation. + arguments replaces the tracer provider/config and logs a warning. Args: tracer_provider: An OpenTelemetry ``TracerProvider``. Defaults to the global provider. - config: An ``openinference.instrumentation.TraceConfig`` (for example ``hide_inputs=True``). + config: Payload visibility and truncation settings. OpenInference TraceConfig is also accepted. **options: Launcher options accepted for OpenTelemetry auto-instrumentation compatibility. Raises: ModuleNotFoundError: When the ``observability`` extra is not installed, except during OpenTelemetry auto-discovery, which skips unavailable optional instrumentation. """ - if not _openinference_is_installed(): + _warn_unknown_options(options) + if not _opentelemetry_is_installed(): if options.get("skip_dep_check") is True: logger.debug(_INSTALL_HINT) return raise ModuleNotFoundError(_INSTALL_HINT) - tracing = _load_tracing_module() - if tracing is not None: - tracing.install(tracer_provider=tracer_provider, config=config) + _load_tracing_module().install(tracer_provider=tracer_provider, config=config) def uninstrument(self) -> None: """Restore the original framework methods. No-op when nothing is installed.""" @@ -99,9 +103,9 @@ def uninstrument(self) -> None: @property def is_instrumented_by_opentelemetry(self) -> bool: - if not _openinference_is_installed() or _openinference_is_too_old(): + if not _opentelemetry_is_installed(): return False return _load_tracing_module().is_installed() -__all__ = ["AtomicAgentsInstrumentor"] +__all__ = ["AtomicAgentsInstrumentor", "TracingConfig", "using_conversation"] diff --git a/atomic-agents/atomic_agents/observability/agent_spans.py b/atomic-agents/atomic_agents/observability/agent_spans.py new file mode 100644 index 00000000..d49000f4 --- /dev/null +++ b/atomic-agents/atomic_agents/observability/agent_spans.py @@ -0,0 +1,239 @@ +"""Agent call and stream spans.""" + +from __future__ import annotations + +import asyncio +import functools +from contextlib import contextmanager +from typing import AsyncGenerator, Awaitable, Callable, Coroutine, Generator, Generic, Iterator, TypeVar + +from opentelemetry import context as context_api +from opentelemetry import trace as trace_api +from opentelemetry.trace import Span, Status, StatusCode + +from atomic_agents.agents.atomic_agent import AtomicAgent + +from . import installation, llm_attempt_spans +from .span_attributes import ATTR_STREAM_CLOSED_EARLY, _agent_display_name, _agent_start_attributes, _suppressed +from .span_payloads import _record_failure, _set_input_safely, _set_output_safely + +OutputSchema = TypeVar("OutputSchema") +CallResult = TypeVar("CallResult") + +AGENT_SYNC_METHOD = "run" +AGENT_ASYNC_METHOD = "run_async" +AGENT_SYNC_STREAM_METHOD = "run_stream" +AGENT_ASYNC_STREAM_METHOD = "run_async_stream" + + +def _user_input_from_call(args: tuple[object, ...], kwargs: dict[str, object]) -> object: + return args[0] if args else kwargs.get("user_input") + + +def _finish_agent_span(span: Span, output: object) -> None: + _set_output_safely(span, output) + span.set_status(Status(StatusCode.OK)) + + +def _mark_closed_early(span: Span, output: object) -> None: + span.set_attribute(ATTR_STREAM_CLOSED_EARLY, True) + _set_output_safely(span, output) + + +@contextmanager +def _agent_span(agent: AtomicAgent, args: tuple[object, ...], kwargs: dict[str, object]) -> Iterator[Span]: + installed = installation._current_installation() + with installed.tracer.start_as_current_span( + f"invoke_agent {_agent_display_name(agent)}", + attributes=_agent_start_attributes(agent), + record_exception=False, + set_status_on_exception=False, + ) as span: + _set_input_safely(span, _user_input_from_call(args, kwargs)) + run = llm_attempt_spans.new_run(span, agent) + try: + with llm_attempt_spans.active_run(run): + yield span + except (asyncio.CancelledError, GeneratorExit): + span.set_attribute("atomic_agents.cancelled", True) + raise + except BaseException as error: + _record_failure(span, error) + raise + finally: + llm_attempt_spans.finish_run(run) + + +def _trace_agent_call(wrapped: Callable[..., OutputSchema]) -> Callable[..., OutputSchema]: + @functools.wraps(wrapped) + def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> OutputSchema: + if _suppressed(): + return wrapped(agent, *args, **kwargs) + with _agent_span(agent, args, kwargs) as span: + response = wrapped(agent, *args, **kwargs) + _finish_agent_span(span, response) + return response + + return traced + + +def _trace_agent_async_call( + wrapped: Callable[..., Awaitable[OutputSchema]], +) -> Callable[..., Coroutine[object, object, OutputSchema]]: + @functools.wraps(wrapped) + async def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> OutputSchema: + if _suppressed(): + return await wrapped(agent, *args, **kwargs) + with _agent_span(agent, args, kwargs) as span: + response = await wrapped(agent, *args, **kwargs) + _finish_agent_span(span, response) + return response + + return traced + + +def _start_agent_span(agent: AtomicAgent, args: tuple[object, ...], kwargs: dict[str, object]) -> Span: + if _suppressed(): + return trace_api.INVALID_SPAN + span = installation._current_installation().tracer.start_span( + f"invoke_agent {_agent_display_name(agent)}", + attributes=_agent_start_attributes(agent), + ) + _set_input_safely(span, _user_input_from_call(args, kwargs)) + return span + + +def _call_within_span(run: llm_attempt_spans._AgentRun, function: Callable[..., CallResult], *args: object) -> CallResult: + span = run.span + token = context_api.attach(trace_api.set_span_in_context(span)) + try: + with llm_attempt_spans.active_run(run): + return function(*args) + finally: + context_api.detach(token) + + +async def _await_within_span(run: llm_attempt_spans._AgentRun, awaitable: Awaitable[CallResult]) -> CallResult: + span = run.span + token = context_api.attach(trace_api.set_span_in_context(span)) + try: + with llm_attempt_spans.active_run(run): + return await awaitable + finally: + context_api.detach(token) + + +class _SyncStreamCursor(Generic[OutputSchema]): + def __init__(self, run: llm_attempt_spans._AgentRun, inner: Generator[OutputSchema, None, OutputSchema | None]): + self.run = run + self.span = run.span + self.inner = inner + self.last: OutputSchema + self.output: OutputSchema | None = None + self.final: OutputSchema | None = None + + def advance(self) -> bool: + try: + self.last = _call_within_span(self.run, next, self.inner) + self.output = self.last + except StopIteration as stop: + self.final = stop.value + return False + return True + + def close_early(self) -> None: + _call_within_span(self.run, self.inner.close) + _mark_closed_early(self.span, self.output) + self.span.set_status(Status(StatusCode.OK)) + + +class _AsyncStreamCursor(Generic[OutputSchema]): + """Async twin of :class:`_SyncStreamCursor`; async generators carry no return value, so the last partial is the output.""" + + def __init__(self, run: llm_attempt_spans._AgentRun, inner: AsyncGenerator[OutputSchema, None]): + self.run = run + self.span = run.span + self.inner = inner + self.last: OutputSchema + self.output: OutputSchema | None = None + + async def advance(self) -> bool: + try: + self.last = await _await_within_span(self.run, self.inner.__anext__()) + self.output = self.last + except StopAsyncIteration: + return False + return True + + async def close_early(self) -> None: + await _await_within_span(self.run, self.inner.aclose()) + _mark_closed_early(self.span, self.output) + self.span.set_status(Status(StatusCode.OK)) + + def mark_cancelled(self) -> None: + _mark_closed_early(self.span, self.output) + + +def _trace_agent_stream( + wrapped: Callable[..., Generator[OutputSchema, None, OutputSchema | None]], +) -> Callable[..., Generator[OutputSchema, None, OutputSchema | None]]: + @functools.wraps(wrapped) + def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> Generator[OutputSchema, None, OutputSchema | None]: + span = _start_agent_span(agent, args, kwargs) + run = llm_attempt_spans.new_run(span, agent) + cursor = _SyncStreamCursor(run, wrapped(agent, *args, **kwargs)) + try: + while cursor.advance(): + yield cursor.last + _finish_agent_span(span, cursor.final) + return cursor.final + except GeneratorExit: + cursor.close_early() + raise + except asyncio.CancelledError: + _mark_closed_early(span, cursor.output) + raise + except BaseException as error: + _record_failure(span, error) + raise + finally: + llm_attempt_spans.finish_run(run) + span.end() + + return traced + + +def _trace_agent_async_stream( + wrapped: Callable[..., AsyncGenerator[OutputSchema, None]], +) -> Callable[..., AsyncGenerator[OutputSchema, None]]: + @functools.wraps(wrapped) + async def traced(agent: AtomicAgent, *args: object, **kwargs: object) -> AsyncGenerator[OutputSchema, None]: + span = _start_agent_span(agent, args, kwargs) + run = llm_attempt_spans.new_run(span, agent) + cursor = _AsyncStreamCursor(run, wrapped(agent, *args, **kwargs)) + try: + while await cursor.advance(): + yield cursor.last + _finish_agent_span(span, cursor.output) + except GeneratorExit: + await cursor.close_early() + raise + except asyncio.CancelledError: + cursor.mark_cancelled() + raise + except BaseException as error: + _record_failure(span, error) + raise + finally: + llm_attempt_spans.finish_run(run) + span.end() + + return traced + + +_AGENT_WRAPPERS = { + AGENT_SYNC_METHOD: _trace_agent_call, + AGENT_ASYNC_METHOD: _trace_agent_async_call, + AGENT_SYNC_STREAM_METHOD: _trace_agent_stream, + AGENT_ASYNC_STREAM_METHOD: _trace_agent_async_stream, +} diff --git a/atomic-agents/atomic_agents/observability/installation.py b/atomic-agents/atomic_agents/observability/installation.py new file mode 100644 index 00000000..88052359 --- /dev/null +++ b/atomic-agents/atomic_agents/observability/installation.py @@ -0,0 +1,130 @@ +"""Tracing installation and method restoration.""" + +import logging +from dataclasses import replace +from typing import TYPE_CHECKING, Callable, Optional, Type + +from opentelemetry import trace as trace_api +from opentelemetry.trace import Tracer +from opentelemetry.util.types import AttributeValue + +from atomic_agents import __version__ +from atomic_agents.agents.atomic_agent import AtomicAgent +from atomic_agents.base.base_tool import BaseTool + +from . import agent_spans, llm_attempt_spans, span_attributes, tool_spans +from .tracing_config import PayloadVisibility, TracingConfig + +if TYPE_CHECKING: + from opentelemetry.trace import TracerProvider + + +logger = logging.getLogger("atomic_agents.observability.tracing") + + +TRACED_MARKER = "__atomic_agents_traced__" +INHERITED_MARKER = "__atomic_agents_inherited__" +TRACER_NAME = "atomic_agents" + + +def _current_installation() -> "_Installation": + assert _installation is not None + return _installation + + +class _Installation: + def __init__( + self, + tracer: Tracer, + tracer_provider: Optional["TracerProvider"], + config: TracingConfig, + original_init_subclass: classmethod, + context_attributes: Callable[[], dict[str, AttributeValue]] | None, + ): + self.context_attributes = context_attributes + self.tracer = tracer + self.tracer_provider = tracer_provider + self.config = config + self.original_init_subclass = original_init_subclass + + +_installation: Optional[_Installation] = None + + +def is_installed() -> bool: + return _installation is not None + + +def install( + *, tracer_provider: Optional["TracerProvider"] = None, config: TracingConfig | PayloadVisibility | None = None +) -> None: + global _installation + if _installation is not None: + tracer_provider = _installation.tracer_provider if tracer_provider is None else tracer_provider + _replace_installation(_installation, tracer_provider, config) + return + config = TracingConfig.coerce(config) + context_attributes = span_attributes._import_context_attributes() + if config.openinference_attributes is None: + config = replace(config, openinference_attributes=context_attributes is not None) + tracer = trace_api.get_tracer(TRACER_NAME, __version__, tracer_provider) + original_init_subclass = BaseTool.__dict__["__init_subclass__"] + _installation = _Installation(tracer, tracer_provider, config, original_init_subclass, context_attributes) + _instrument_agent_class() + tool_spans._instrument_existing_tool_classes() + tool_spans._hook_future_tool_classes(original_init_subclass) + + +def uninstall() -> None: + global _installation + if _installation is None: + return + setattr(BaseTool, "__init_subclass__", _installation.original_init_subclass) + for tool_class in tool_spans._all_subclasses(BaseTool): + tool_spans._restore_tool_class(tool_class) + for method_name in agent_spans._AGENT_WRAPPERS: + _restore_method(AtomicAgent, method_name) + llm_attempt_spans.unregister_clients() + _installation = None + + +def _replace_installation( + installed: _Installation, tracer_provider: Optional["TracerProvider"], config: TracingConfig | PayloadVisibility | None +) -> None: + config = installed.config if config is None else TracingConfig.coerce(config) + if config.openinference_attributes is None: + config = replace(config, openinference_attributes=installed.config.openinference_attributes) + if tracer_provider is installed.tracer_provider and config == installed.config: + return + installed.tracer = trace_api.get_tracer(TRACER_NAME, __version__, tracer_provider) + installed.tracer_provider = tracer_provider + installed.config = config + installed.context_attributes = span_attributes._import_context_attributes() + logger.warning( + "AtomicAgentsInstrumentor tracer provider/config replaced; " + "in-flight spans finish on the provider they started with." + ) + + +def _instrument_agent_class() -> None: + for method_name, wrap in agent_spans._AGENT_WRAPPERS.items(): + original = AtomicAgent.__dict__[method_name] + if not getattr(original, TRACED_MARKER, False): + setattr(AtomicAgent, method_name, _mark_traced(wrap(original), original, False)) + + +def _mark_traced(wrapper: Callable, original: Callable, inherited: bool) -> Callable: + setattr(wrapper, TRACED_MARKER, True) + setattr(wrapper, INHERITED_MARKER, inherited) + setattr(wrapper, "__wrapped__", original) + return wrapper + + +def _restore_method(owner: Type, method_name: str) -> None: + current = owner.__dict__.get(method_name) + if current is None or not getattr(current, TRACED_MARKER, False): + return + if getattr(current, INHERITED_MARKER, False): + delattr(owner, method_name) + else: + setattr(owner, method_name, current.__wrapped__) diff --git a/atomic-agents/atomic_agents/observability/llm_attempt_spans.py b/atomic-agents/atomic_agents/observability/llm_attempt_spans.py new file mode 100644 index 00000000..05e4debc --- /dev/null +++ b/atomic-agents/atomic_agents/observability/llm_attempt_spans.py @@ -0,0 +1,295 @@ +"""Instructor attempt spans and cumulative agent usage.""" + +from __future__ import annotations + +import contextvars +import functools +import logging +import threading +import weakref +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Callable, Iterator, ParamSpec, Protocol + +from opentelemetry import context as context_api, trace as trace_api +from opentelemetry.trace import Span, SpanKind, Status, StatusCode, Tracer +from opentelemetry.util.types import AttributeValue + +from atomic_agents.agents.atomic_agent import AtomicAgent + +from . import installation +from .span_attributes import _best_effort_attributes, _optional_attributes, _provider_attributes, _without_none +from .span_payloads import _bounded_text, _failure_message, _record_failure, _write_payload +from .tracing_config import TracingConfig + +logger = logging.getLogger("atomic_agents.observability.tracing") +_PARAMETERS = ParamSpec("_PARAMETERS") +_RESPONSE_MODEL = "gen_ai.response.model" +_ATTEMPT = "atomic_agents.llm.attempt" +_VALIDATION_ERROR = "atomic_agents.llm.last_validation_error" + + +class _PromptTokenDetails(Protocol): + cached_tokens: int | None + + +class _CompletionTokenDetails(Protocol): + reasoning_tokens: int | None + + +class _InputOutputUsage(Protocol): + input_tokens: int + output_tokens: int + input_tokens_details: _PromptTokenDetails | None + output_tokens_details: _CompletionTokenDetails | None + + +class _PromptCompletionUsage(Protocol): + prompt_tokens: int + completion_tokens: int + prompt_tokens_details: _PromptTokenDetails | None + completion_tokens_details: _CompletionTokenDetails | None + + +_Usage = _InputOutputUsage | _PromptCompletionUsage + + +class _Response(Protocol): + model: str + usage: _Usage + + +class _HookClient(Protocol): + def on(self, event: str, handler: Callable[..., None]) -> object: + """Append a handler without replacing application hooks.""" + + def off(self, event: str, handler: Callable[..., None]) -> object: + """Remove only this handler, preserving application hooks.""" + + +@dataclass +class _AgentRun: + span: Span + tracer: Tracer + config: TracingConfig + provider: dict[str, AttributeValue] + usage_totals: dict[str, int] = field(default_factory=dict) + attempt_count: int = 0 + suppressed_attempt: bool = False + response_received: bool = False + open_attempt: Span | None = None + context_token: contextvars.Token | None = None + + +_active_run: contextvars.ContextVar[_AgentRun | None] = contextvars.ContextVar("atomic_agents_llm_run", default=None) +_clients: weakref.WeakSet[_HookClient] = weakref.WeakSet() +_registration_lock = threading.Lock() + + +def _never_raise(handler: Callable[_PARAMETERS, None]) -> Callable[_PARAMETERS, None]: + @functools.wraps(handler) + def guarded(*args: _PARAMETERS.args, **kwargs: _PARAMETERS.kwargs) -> None: + try: + handler(*args, **kwargs) + except Exception: + logger.warning("Instructor tracing hook failed: %s", handler.__name__, exc_info=True) + + return guarded + + +@_never_raise +def register_client(client: _HookClient) -> None: + if not hasattr(client, "on"): + return + # Shared clients must acquire all handlers once before another run can use them. + with _registration_lock: + if client not in _clients: + _register_handlers(client) + _clients.add(client) + + +def _register_handlers(client: _HookClient) -> None: + registered: list[tuple[str, Callable[..., None]]] = [] + try: + for event, handler in _HANDLERS.items(): + client.on(event, handler) + registered.append((event, handler)) + except Exception: + for event, handler in registered: + client.off(event, handler) + raise + + +@_never_raise +def unregister_clients() -> None: + with _registration_lock: + for client in list(_clients): + _unregister_client(client) + _clients.clear() + + +@_never_raise +def _unregister_client(client: _HookClient) -> None: + for event, handler in _HANDLERS.items(): + client.off(event, handler) + + +def new_run(span: Span, agent: AtomicAgent) -> _AgentRun: + installed = installation._current_installation() + register_client(getattr(agent, "client")) + provider = _best_effort_attributes( + {"gen_ai.provider.name": lambda: _provider_attributes(agent).get("gen_ai.provider.name")} + ) + return _AgentRun(span, installed.tracer, installed.config, provider) + + +@contextmanager +def active_run(run: _AgentRun) -> Iterator[None]: + token = _active_run.set(None if run.span is trace_api.INVALID_SPAN else run) + _attach_attempt(run) + try: + yield + finally: + _detach_attempt(run) + _active_run.reset(token) + + +def _attach_attempt(run: _AgentRun) -> None: + if run.open_attempt is not None: + run.context_token = context_api.attach(trace_api.set_span_in_context(run.open_attempt)) + + +def _detach_attempt(run: _AgentRun) -> None: + if run.context_token is not None: + context_api.detach(run.context_token) + run.context_token = None + + +@_never_raise +def finish_attempt(run: _AgentRun, status: Status | None = None) -> None: + if run.open_attempt is None: + return + if run.response_received: + status = Status(StatusCode.OK) + run.response_received = False + if status is not None: + run.open_attempt.set_status(status) + run.open_attempt.end() + run.open_attempt = None + _detach_attempt(run) + + +@_never_raise +def _completion_kwargs(*args: object, **kwargs: object) -> None: + run = _active_run.get() + if run is None: + return + finish_attempt(run) + # Suppressed kwargs must also swallow later response events outside that context. + run.suppressed_attempt = bool(context_api.get_value("suppress_instrumentation")) + if run.suppressed_attempt: + return + run.attempt_count += 1 + if not run.config.llm_spans: + return + model = _request_model(kwargs) + attributes: dict[str, AttributeValue | None] = { + "gen_ai.operation.name": "chat", + "gen_ai.request.model": model, + _ATTEMPT: run.attempt_count, + } + attributes.update(run.provider) + attributes.update(_optional_attributes("LLM", {})) + run.open_attempt = run.tracer.start_span( + f"chat {model}" if model else "chat", kind=SpanKind.CLIENT, attributes=_without_none(attributes) + ) + _attach_attempt(run) + + +def _request_model(kwargs: dict[str, object]) -> str | None: + model = kwargs.get("model") + return str(model) if model is not None else None + + +def _usage_attributes(usage: _Usage | None) -> dict[str, int]: + readers = { + "input_tokens": getattr(usage, "input_tokens", getattr(usage, "prompt_tokens", None)), + "output_tokens": getattr(usage, "output_tokens", getattr(usage, "completion_tokens", None)), + "cache_read.input_tokens": getattr( + getattr(usage, "input_tokens_details", getattr(usage, "prompt_tokens_details", None)), "cached_tokens", None + ), + "reasoning.output_tokens": getattr( + getattr(usage, "output_tokens_details", getattr(usage, "completion_tokens_details", None)), + "reasoning_tokens", + None, + ), + } + return {"gen_ai.usage." + key: value for key, value in readers.items() if isinstance(value, int)} + + +@_never_raise +def _completion_response(response: _Response, **kwargs: object) -> None: + run = _active_run.get() + if run is None or run.suppressed_attempt: + return + attributes = _without_none({_RESPONSE_MODEL: getattr(response, "model", None)}) + run.span.set_attributes(attributes) + usage = _usage_attributes(getattr(response, "usage", None)) + for key, value in usage.items(): + run.usage_totals[key] = run.usage_totals.get(key, 0) + value + _complete_response_attempt(run, usage, attributes) + + +def _complete_response_attempt(run: _AgentRun, usage: dict[str, int], attributes: dict[str, AttributeValue]) -> None: + if run.open_attempt is None: + return + run.open_attempt.set_attributes(attributes) + run.open_attempt.set_attributes(usage) + run.response_received = True + + +@_never_raise +def _completion_error(error: BaseException, **kwargs: object) -> None: + run = _active_run.get() + if run is None or run.suppressed_attempt or run.open_attempt is None: + return + run.response_received = False + _record_failure(run.open_attempt, error, run.config) + finish_attempt(run) + + +@_never_raise +def _parse_error(error: BaseException, **kwargs: object) -> None: + run = _active_run.get() + if run is None or run.suppressed_attempt: + return + text = "__REDACTED__" if run.config.hide_outputs else _failure_message(error, run.config) + _write_payload(run.span, _VALIDATION_ERROR, None, text, run.config.hide_outputs, run.config) + event = {_ATTEMPT: run.attempt_count, _VALIDATION_ERROR: _bounded_text(text, run.config)} + run.span.add_event("validation_error", event) + _mark_attempt_validation_error(run, error) + + +def _mark_attempt_validation_error(run: _AgentRun, error: BaseException) -> None: + if run.open_attempt is None: + return + run.response_received = False + message = _bounded_text(_failure_message(error, run.config), run.config) + run.open_attempt.add_event("validation_error", {_ATTEMPT: run.attempt_count, _VALIDATION_ERROR: message}) + run.open_attempt.set_attribute("error.type", type(error).__qualname__) + run.open_attempt.set_status(Status(StatusCode.ERROR, message)) + + +@_never_raise +def finish_run(run: _AgentRun) -> None: + finish_attempt(run) + run.span.set_attributes(run.usage_totals) + run.span.set_attribute("atomic_agents.llm.attempts", run.attempt_count) + + +_HANDLERS: dict[str, Callable[..., None]] = { + "completion:kwargs": _completion_kwargs, + "completion:response": _completion_response, + "completion:error": _completion_error, + "parse:error": _parse_error, +} diff --git a/atomic-agents/atomic_agents/observability/openinference_tracing.py b/atomic-agents/atomic_agents/observability/openinference_tracing.py deleted file mode 100644 index 7959be40..00000000 --- a/atomic-agents/atomic_agents/observability/openinference_tracing.py +++ /dev/null @@ -1,475 +0,0 @@ -""" -OpenInference span emission for ``AtomicAgent`` runs and ``BaseTool`` runs. - -Loaded lazily by :class:`atomic_agents.observability.AtomicAgentsInstrumentor`; importing this module -requires the ``observability`` extra. -""" - -import asyncio -import contextvars -import functools -import hashlib -import inspect -import json -import logging -from types import FunctionType -from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Dict, Generator, List, Optional, Type - -from openinference.instrumentation import OITracer, TraceConfig -from openinference.semconv.trace import SpanAttributes -from opentelemetry import context as context_api -from opentelemetry import trace as trace_api -from opentelemetry.trace import Status, StatusCode - -from atomic_agents import __version__ -from atomic_agents.agents.atomic_agent import AtomicAgent -from atomic_agents.base.base_tool import BaseTool - -if TYPE_CHECKING: - from opentelemetry.trace import TracerProvider - -logger = logging.getLogger(__name__) - -TRACED_MARKER = "__atomic_agents_traced__" -INHERITED_MARKER = "__atomic_agents_inherited__" -TRACER_NAME = "atomic_agents" - -AGENT_SYNC_METHOD = "run" -AGENT_ASYNC_METHOD = "run_async" -AGENT_SYNC_STREAM_METHOD = "run_stream" -AGENT_ASYNC_STREAM_METHOD = "run_async_stream" -TOOL_METHOD_NAMES = ("run", "run_async", "arun") - -ATTR_MODEL = "atomic_agents.model" -ATTR_MODE = "atomic_agents.mode" -ATTR_INPUT_SCHEMA = "atomic_agents.input_schema" -ATTR_OUTPUT_SCHEMA = "atomic_agents.output_schema" -ATTR_HISTORY_MESSAGE_COUNT = "atomic_agents.history.message_count" -ATTR_CONTEXT_PROVIDERS = "atomic_agents.context_providers" -ATTR_SYSTEM_PROMPT_SHA256 = "atomic_agents.system_prompt.sha256" -ATTR_STREAM_CLOSED_EARLY = "atomic_agents.stream.closed_early" - -_active_tool: contextvars.ContextVar[Optional[BaseTool]] = contextvars.ContextVar("atomic_agents_active_tool", default=None) - - -class _Installation: - def __init__( - self, - tracer: OITracer, - tracer_provider: Optional["TracerProvider"], - config: TraceConfig, - original_init_subclass: classmethod, - ): - self.tracer = tracer - self.tracer_provider = tracer_provider - self.config = config - self.original_init_subclass = original_init_subclass - - -_installation: Optional[_Installation] = None - - -def is_installed() -> bool: - return _installation is not None - - -def install(*, tracer_provider: Optional["TracerProvider"] = None, config: Optional[TraceConfig] = None) -> None: - global _installation - config = config or TraceConfig() - if _installation is not None: - _warn_if_reinstalled_differently(_installation, tracer_provider, config) - return - tracer = OITracer(trace_api.get_tracer(TRACER_NAME, __version__, tracer_provider), config=config) - original_init_subclass = BaseTool.__dict__["__init_subclass__"] - _installation = _Installation(tracer, tracer_provider, config, original_init_subclass) - _instrument_agent_class(tracer) - _instrument_existing_tool_classes(tracer) - _hook_future_tool_classes(tracer, original_init_subclass) - - -def uninstall() -> None: - global _installation - if _installation is None: - return - setattr(BaseTool, "__init_subclass__", _installation.original_init_subclass) - for tool_class in _all_subclasses(BaseTool): - _restore_tool_class(tool_class) - for method_name in _AGENT_WRAPPERS: - _restore_method(AtomicAgent, method_name) - _installation = None - - -def _warn_if_reinstalled_differently( - installed: _Installation, tracer_provider: Optional["TracerProvider"], config: TraceConfig -) -> None: - same_provider = tracer_provider is None or tracer_provider is installed.tracer_provider - same_config = config == installed.config - if not (same_provider and same_config): - logger.warning( - "AtomicAgentsInstrumentor is already installed; the new tracer_provider/config are ignored. " - "Call uninstrument() first to change them." - ) - - -def _instrument_agent_class(tracer: OITracer) -> None: - for method_name, wrap in _AGENT_WRAPPERS.items(): - original = AtomicAgent.__dict__[method_name] - if not getattr(original, TRACED_MARKER, False): - setattr(AtomicAgent, method_name, _mark_traced(wrap(tracer, method_name, original), original, False)) - - -def _mark_traced(wrapper: Callable, original: Callable, inherited: bool) -> Callable: - setattr(wrapper, TRACED_MARKER, True) - setattr(wrapper, INHERITED_MARKER, inherited) - setattr(wrapper, "__wrapped__", original) - return wrapper - - -def _restore_method(owner: Type, method_name: str) -> None: - current = owner.__dict__.get(method_name) - if current is None or not getattr(current, TRACED_MARKER, False): - return - if getattr(current, INHERITED_MARKER, False): - delattr(owner, method_name) - else: - setattr(owner, method_name, current.__wrapped__) - - -def _all_subclasses(base: Type) -> List[Type]: - found: List[Type] = [] - pending = list(base.__subclasses__()) - while pending: - current = pending.pop(0) - if current not in found: - found.append(current) - pending.extend(current.__subclasses__()) - return found - - -def _instrument_existing_tool_classes(tracer: OITracer) -> None: - for tool_class in _all_subclasses(BaseTool): - _instrument_tool_class(tool_class, tracer) - - -def _hook_future_tool_classes(tracer: OITracer, original_init_subclass: classmethod) -> None: - def traced_init_subclass(cls, **kwargs): - original_init_subclass.__func__(cls, **kwargs) - _instrument_tool_class(cls, tracer) - - setattr(BaseTool, "__init_subclass__", classmethod(traced_init_subclass)) - - -def _instrument_tool_class(tool_class: Type, tracer: OITracer) -> None: - for method_name in TOOL_METHOD_NAMES: - method = inspect.getattr_static(tool_class, method_name, None) - if _is_traceable_function(method): - inherited = method_name not in tool_class.__dict__ - setattr(tool_class, method_name, _mark_traced(_trace_tool_method(tracer, method), method, inherited)) - - -def _is_traceable_function(method: Any) -> bool: - if not isinstance(method, FunctionType) or inspect.isasyncgenfunction(method): - return False - if getattr(method, "__isabstractmethod__", False): - return False - return not getattr(method, TRACED_MARKER, False) - - -def _restore_tool_class(tool_class: Type) -> None: - for method_name in TOOL_METHOD_NAMES: - _restore_method(tool_class, method_name) - - -def _agent_display_name(agent: AtomicAgent) -> str: - return f"{type(agent).__name__}[{agent.input_schema.__name__}, {agent.output_schema.__name__}]" - - -def _agent_start_attributes(agent: AtomicAgent) -> Dict[str, Any]: - try: - return { - SpanAttributes.AGENT_NAME: _agent_display_name(agent), - ATTR_MODEL: agent.model, - ATTR_MODE: getattr(agent.mode, "value", str(agent.mode)), - ATTR_INPUT_SCHEMA: agent.input_schema.__name__, - ATTR_OUTPUT_SCHEMA: agent.output_schema.__name__, - ATTR_HISTORY_MESSAGE_COUNT: len(agent.history.history), - ATTR_CONTEXT_PROVIDERS: _context_provider_names(agent), - **_system_prompt_version(agent), - } - except Exception as error: - logger.warning("Could not describe agent for tracing: %s", error) - return {} - - -def _context_provider_names(agent: AtomicAgent) -> List[str]: - providers = getattr(agent.system_prompt_generator, "context_providers", None) or {} - return sorted(str(name) for name in providers) - - -def _system_prompt_version(agent: AtomicAgent) -> Dict[str, str]: - """Hashes only the static prompt sections, so the value identifies the prompt version across turns - even though context providers change the rendered prompt on every call.""" - generator = agent.system_prompt_generator - sections = {name: getattr(generator, name, None) for name in ("background", "steps", "output_instructions")} - if agent.system_role is None or all(value is None for value in sections.values()): - return {} - digest = hashlib.sha256(json.dumps(sections, sort_keys=True, default=str).encode("utf-8")).hexdigest() - return {ATTR_SYSTEM_PROMPT_SHA256: digest} - - -def _user_input_from_call(args: tuple, kwargs: dict) -> Any: - return args[0] if args else kwargs.get("user_input") - - -def _set_input_safely(span: Any, value: Any) -> None: - if value is None: - return - try: - span.set_input(value) - except Exception as error: - logger.warning("Could not record span input: %s", error) - - -def _set_output_safely(span: Any, value: Any) -> None: - if value is None: - return - try: - span.set_output(value) - except Exception as error: - logger.warning("Could not record span output: %s", error) - - -def _finish_agent_span(span: Any, output: Any) -> None: - _set_output_safely(span, output) - span.set_status(Status(StatusCode.OK)) - - -def _mark_closed_early(span: Any, output: Any) -> None: - span.set_attribute(ATTR_STREAM_CLOSED_EARLY, True) - _set_output_safely(span, output) - - -def _record_failure(span: Any, error: BaseException) -> None: - span.record_exception(error) - span.set_status(Status(StatusCode.ERROR, str(error))) - - -def _trace_agent_call(tracer: OITracer, method_name: str, wrapped: Callable) -> Callable: - @functools.wraps(wrapped) - def traced(agent, *args, **kwargs): - span_name = f"{_agent_display_name(agent)}.{method_name}" - with tracer.start_as_current_span( - span_name, openinference_span_kind="agent", attributes=_agent_start_attributes(agent) - ) as span: - _set_input_safely(span, _user_input_from_call(args, kwargs)) - response = wrapped(agent, *args, **kwargs) - _finish_agent_span(span, response) - return response - - return traced - - -def _trace_agent_async_call(tracer: OITracer, method_name: str, wrapped: Callable) -> Callable: - @functools.wraps(wrapped) - async def traced(agent, *args, **kwargs): - span_name = f"{_agent_display_name(agent)}.{method_name}" - with tracer.start_as_current_span( - span_name, openinference_span_kind="agent", attributes=_agent_start_attributes(agent) - ) as span: - _set_input_safely(span, _user_input_from_call(args, kwargs)) - response = await wrapped(agent, *args, **kwargs) - _finish_agent_span(span, response) - return response - - return traced - - -def _start_agent_span(tracer: OITracer, agent: AtomicAgent, method_name: str, args: tuple, kwargs: dict) -> Any: - span = tracer.start_span( - f"{_agent_display_name(agent)}.{method_name}", - openinference_span_kind="agent", - attributes=_agent_start_attributes(agent), - ) - _set_input_safely(span, _user_input_from_call(args, kwargs)) - return span - - -def _call_within_span(span: Any, function: Callable, *args: Any) -> Any: - token = context_api.attach(trace_api.set_span_in_context(span)) - try: - return function(*args) - finally: - context_api.detach(token) - - -async def _await_within_span(span: Any, awaitable: Any) -> Any: - token = context_api.attach(trace_api.set_span_in_context(span)) - try: - return await awaitable - finally: - context_api.detach(token) - - -class _SyncStreamCursor: - def __init__(self, span: Any, inner: Generator): - self.span = span - self.inner = inner - self.last: Any = None - self.final: Any = None - - def advance(self) -> bool: - try: - self.last = _call_within_span(self.span, next, self.inner) - except StopIteration as stop: - self.final = stop.value - return False - return True - - def close_early(self) -> None: - _call_within_span(self.span, self.inner.close) - _mark_closed_early(self.span, self.last) - self.span.set_status(Status(StatusCode.OK)) - - -class _AsyncStreamCursor: - """Async twin of :class:`_SyncStreamCursor`; async generators carry no return value, so the last partial is the output.""" - - def __init__(self, span: Any, inner: AsyncGenerator): - self.span = span - self.inner = inner - self.last: Any = None - - async def advance(self) -> bool: - try: - self.last = await _await_within_span(self.span, self.inner.__anext__()) - except StopAsyncIteration: - return False - return True - - async def close_early(self) -> None: - await _await_within_span(self.span, self.inner.aclose()) - _mark_closed_early(self.span, self.last) - self.span.set_status(Status(StatusCode.OK)) - - def mark_cancelled(self) -> None: - _mark_closed_early(self.span, self.last) - - -def _trace_agent_stream(tracer: OITracer, method_name: str, wrapped: Callable) -> Callable: - @functools.wraps(wrapped) - def traced(agent, *args, **kwargs): - span = _start_agent_span(tracer, agent, method_name, args, kwargs) - cursor = _SyncStreamCursor(span, wrapped(agent, *args, **kwargs)) - try: - while cursor.advance(): - yield cursor.last - _finish_agent_span(span, cursor.final) - return cursor.final - except GeneratorExit: - cursor.close_early() - raise - except BaseException as error: - _record_failure(span, error) - raise - finally: - span.end() - - return traced - - -def _trace_agent_async_stream(tracer: OITracer, method_name: str, wrapped: Callable) -> Callable: - @functools.wraps(wrapped) - async def traced(agent, *args, **kwargs): - span = _start_agent_span(tracer, agent, method_name, args, kwargs) - cursor = _AsyncStreamCursor(span, wrapped(agent, *args, **kwargs)) - try: - while await cursor.advance(): - yield cursor.last - _finish_agent_span(span, cursor.last) - except GeneratorExit: - await cursor.close_early() - raise - except asyncio.CancelledError: - cursor.mark_cancelled() - raise - except BaseException as error: - _record_failure(span, error) - raise - finally: - span.end() - - return traced - - -def _describe_tool(span: Any, tool: BaseTool) -> None: - try: - span.set_tool( - name=tool.tool_name, - description=tool.tool_description, - parameters=tool.input_schema.model_json_schema(), - ) - except Exception as error: - logger.warning("Could not describe tool %s for tracing: %s", type(tool).__name__, error) - - -def _tool_input_from_call(args: tuple, kwargs: dict) -> Any: - return args[0] if args else kwargs.get("params") - - -def _trace_tool_method(tracer: OITracer, wrapped: Callable) -> Callable: - if inspect.iscoroutinefunction(wrapped): - return _trace_tool_async_method(tracer, wrapped) - - @functools.wraps(wrapped) - def traced(tool, *args, **kwargs): - if _active_tool.get() is tool: - return wrapped(tool, *args, **kwargs) - with tracer.start_as_current_span(type(tool).__name__, openinference_span_kind="tool") as span: - _describe_tool(span, tool) - _set_input_safely(span, _tool_input_from_call(args, kwargs)) - result = _run_as_active_tool(tool, wrapped, args, kwargs) - _set_output_safely(span, result) - span.set_status(Status(StatusCode.OK)) - return result - - return traced - - -def _trace_tool_async_method(tracer: OITracer, wrapped: Callable) -> Callable: - @functools.wraps(wrapped) - async def traced(tool, *args, **kwargs): - if _active_tool.get() is tool: - return await wrapped(tool, *args, **kwargs) - with tracer.start_as_current_span(type(tool).__name__, openinference_span_kind="tool") as span: - _describe_tool(span, tool) - _set_input_safely(span, _tool_input_from_call(args, kwargs)) - result = await _await_as_active_tool(tool, wrapped, args, kwargs) - _set_output_safely(span, result) - span.set_status(Status(StatusCode.OK)) - return result - - return traced - - -def _run_as_active_tool(tool: BaseTool, wrapped: Callable, args: tuple, kwargs: dict) -> Any: - token = _active_tool.set(tool) - try: - return wrapped(tool, *args, **kwargs) - finally: - _active_tool.reset(token) - - -async def _await_as_active_tool(tool: BaseTool, wrapped: Callable, args: tuple, kwargs: dict) -> Any: - token = _active_tool.set(tool) - try: - return await wrapped(tool, *args, **kwargs) - finally: - _active_tool.reset(token) - - -_AGENT_WRAPPERS: Dict[str, Callable[[OITracer, str, Callable], Callable]] = { - AGENT_SYNC_METHOD: _trace_agent_call, - AGENT_ASYNC_METHOD: _trace_agent_async_call, - AGENT_SYNC_STREAM_METHOD: _trace_agent_stream, - AGENT_ASYNC_STREAM_METHOD: _trace_agent_async_stream, -} diff --git a/atomic-agents/atomic_agents/observability/span_attributes.py b/atomic-agents/atomic_agents/observability/span_attributes.py new file mode 100644 index 00000000..ccbf615b --- /dev/null +++ b/atomic-agents/atomic_agents/observability/span_attributes.py @@ -0,0 +1,190 @@ +"""Span attributes and conversation context.""" + +import contextvars +import hashlib +import json +import logging +from collections.abc import Mapping +from contextlib import contextmanager +from typing import Callable, Dict, Iterator, List + +from opentelemetry import context as context_api +from opentelemetry.util.types import AttributeValue + +from atomic_agents.agents.atomic_agent import AtomicAgent + +from . import installation + + +logger = logging.getLogger("atomic_agents.observability.tracing") + + +ATTR_MODEL = "atomic_agents.model" +ATTR_MODE = "atomic_agents.mode" +ATTR_INPUT_SCHEMA = "atomic_agents.input_schema" +ATTR_OUTPUT_SCHEMA = "atomic_agents.output_schema" +ATTR_HISTORY_MESSAGE_COUNT = "atomic_agents.history.message_count" +ATTR_CONTEXT_PROVIDERS = "atomic_agents.context_providers" +ATTR_SYSTEM_PROMPT_SHA256 = "atomic_agents.system_prompt.sha256" +ATTR_STREAM_CLOSED_EARLY = "atomic_agents.stream.closed_early" +ATTR_OPERATION = "gen_ai.operation.name" +ATTR_AGENT_NAME = "gen_ai.agent.name" +ATTR_REQUEST_MODEL = "gen_ai.request.model" +ATTR_PROVIDER_NAME = "gen_ai.provider.name" +ATTR_CONVERSATION_ID = "gen_ai.conversation.id" +ATTR_TOOL_NAME = "gen_ai.tool.name" +ATTR_TOOL_DESCRIPTION = "gen_ai.tool.description" +ATTR_TOOL_TYPE = "gen_ai.tool.type" +ATTR_TOOL_ARGUMENTS = "gen_ai.tool.call.arguments" +ATTR_TOOL_RESULT = "gen_ai.tool.call.result" +ATTR_SPAN_KIND = "openinference.span.kind" +ATTR_INPUT_VALUE = "input.value" +ATTR_INPUT_MIME_TYPE = "input.mime_type" +ATTR_OUTPUT_VALUE = "output.value" +ATTR_OUTPUT_MIME_TYPE = "output.mime_type" +ATTR_OI_AGENT_NAME = "agent.name" +ATTR_OI_TOOL_NAME = "tool.name" +ATTR_OI_TOOL_DESCRIPTION = "tool.description" +ATTR_TOOL_PARAMETERS = "tool.parameters" +ATTR_SESSION_ID = "session.id" +ATTR_INPUT = "atomic_agents.input" +ATTR_OUTPUT = "atomic_agents.output" +ATTR_AGENT_CLASS = "atomic_agents.agent.class" +ATTR_ERROR_TYPE = "error.type" +ATTR_TOOL_CLASS = "atomic_agents.tool.class" +PROVIDER_NAMES = { + "azure_openai": "azure.ai.openai", + "vertexai": "gcp.vertex_ai", + "gemini": "gcp.gemini", + "genai": "gcp.gemini", + "bedrock": "aws.bedrock", + "mistral": "mistral_ai", + "xai": "x_ai", +} + + +def _import_context_attributes() -> Callable[[], dict[str, AttributeValue]] | None: + try: + from openinference.instrumentation import get_attributes_from_context + except ImportError: + return None + from atomic_agents.observability import _openinference_is_too_old + + if _openinference_is_too_old(): + logger.warning("OpenInference context attributes disabled: 0.1.38 or newer is required.") + return None + return lambda: dict(get_attributes_from_context()) + + +def _context_attributes() -> dict[str, AttributeValue]: + reader = installation._current_installation().context_attributes + return reader() if reader is not None else {} + + +def _suppressed() -> bool: + return bool(context_api.get_value("suppress_instrumentation")) + + +_conversation_id: contextvars.ContextVar[str | None] = contextvars.ContextVar("atomic_agents_conversation_id", default=None) + + +@contextmanager +def using_conversation(conversation_id: str) -> Iterator[None]: + token = _conversation_id.set(conversation_id) + try: + yield + finally: + _conversation_id.reset(token) + + +def _agent_display_name(agent: AtomicAgent) -> str: + attributes = _best_effort_attributes({ATTR_AGENT_NAME: lambda: getattr(agent, "name", None)}) + return str(attributes.get(ATTR_AGENT_NAME) or type(agent).__name__) + + +def _agent_class_name(agent: AtomicAgent) -> str: + return f"{type(agent).__name__}[{agent.input_schema.__name__}, {agent.output_schema.__name__}]" + + +def _best_effort_attributes( + readers: dict[str, Callable[[], AttributeValue | None]], +) -> dict[str, AttributeValue]: + attributes: dict[str, AttributeValue] = {} + for key, read in readers.items(): + try: + value = read() + if value is not None: + attributes[key] = value + except Exception as error: + logger.debug("Could not collect tracing attribute %s: %s", key, error) + return attributes + + +def _without_none(attributes: Mapping[str, AttributeValue | None]) -> dict[str, AttributeValue]: + return {key: value for key, value in attributes.items() if value is not None} + + +def _agent_start_attributes(agent: AtomicAgent) -> dict[str, AttributeValue]: + attributes: dict[str, AttributeValue] = { + ATTR_OPERATION: "invoke_agent", + ATTR_AGENT_NAME: _agent_display_name(agent), + ATTR_AGENT_CLASS: _agent_class_name(agent), + } + attributes.update( + _best_effort_attributes( + { + ATTR_REQUEST_MODEL: lambda: agent.model, + ATTR_MODEL: lambda: agent.model, + ATTR_MODE: lambda: getattr(agent.mode, "value", None), + ATTR_INPUT_SCHEMA: lambda: agent.input_schema.__name__, + ATTR_OUTPUT_SCHEMA: lambda: agent.output_schema.__name__, + ATTR_HISTORY_MESSAGE_COUNT: lambda: len(agent.history.history), + ATTR_CONTEXT_PROVIDERS: lambda: _context_provider_names(agent), + ATTR_SYSTEM_PROMPT_SHA256: lambda: _system_prompt_version(agent).get(ATTR_SYSTEM_PROMPT_SHA256), + ATTR_PROVIDER_NAME: lambda: _provider_attributes(agent).get(ATTR_PROVIDER_NAME), + } + ) + ) + attributes.update(_optional_attributes("AGENT", {ATTR_OI_AGENT_NAME: _agent_display_name(agent)})) + return _without_none(attributes) + + +def _context_provider_names(agent: AtomicAgent) -> List[str]: + providers = getattr(agent.system_prompt_generator, "context_providers", None) or {} + return sorted(str(name) for name in providers) + + +def _system_prompt_version(agent: AtomicAgent) -> Dict[str, str]: + """Hashes only the static prompt sections, so the value identifies the prompt version across turns + even though context providers change the rendered prompt on every call.""" + generator = agent.system_prompt_generator + sections = {name: getattr(generator, name, None) for name in ("background", "steps", "output_instructions")} + if agent.system_role is None or all(value is None for value in sections.values()): + return {} + digest = hashlib.sha256(json.dumps(sections, sort_keys=True, default=str).encode("utf-8")).hexdigest() + return {ATTR_SYSTEM_PROMPT_SHA256: digest} + + +def _provider_attributes(agent: AtomicAgent) -> dict[str, str]: + provider = getattr(getattr(agent.client, "provider", None), "value", None) + return {ATTR_PROVIDER_NAME: PROVIDER_NAMES.get(provider, provider)} if isinstance(provider, str) else {} + + +def _optional_attributes(kind: str, attributes: dict[str, AttributeValue]) -> dict[str, AttributeValue]: + context = _context_attributes() + conversation = _conversation_attributes(context) + if not installation._current_installation().config.openinference_attributes: + return conversation + return _without_none({ATTR_SPAN_KIND: kind, **attributes, **context, **conversation}) + + +def _conversation_attributes(context: dict[str, AttributeValue]) -> dict[str, AttributeValue]: + conversation_id: AttributeValue | None = _conversation_id.get() + if conversation_id is None: + conversation_id = context.get(ATTR_SESSION_ID) + if conversation_id is None: + return {} + attributes = {ATTR_CONVERSATION_ID: conversation_id} + if installation._current_installation().config.openinference_attributes: + attributes[ATTR_SESSION_ID] = conversation_id + return attributes diff --git a/atomic-agents/atomic_agents/observability/span_payloads.py b/atomic-agents/atomic_agents/observability/span_payloads.py new file mode 100644 index 00000000..62317f29 --- /dev/null +++ b/atomic-agents/atomic_agents/observability/span_payloads.py @@ -0,0 +1,143 @@ +"""Payload serialization, redaction, and bounded failures.""" + +import json +import logging +import traceback +from collections.abc import Iterator + +from instructor.core.exceptions import InstructorError + +from opentelemetry.trace import Span, Status, StatusCode +from pydantic import BaseModel, ValidationError + +from . import installation +from .tracing_config import TracingConfig +from .span_attributes import ( + ATTR_INPUT_VALUE, + ATTR_INPUT_MIME_TYPE, + ATTR_OUTPUT_VALUE, + ATTR_OUTPUT_MIME_TYPE, + ATTR_INPUT, + ATTR_OUTPUT, + ATTR_ERROR_TYPE, +) + + +logger = logging.getLogger("atomic_agents.observability.tracing") + + +TRUNCATION_MARKER = "...[truncated]" + + +def _serialize_payload(value: object) -> str | None: + try: + return value.model_dump_json() if isinstance(value, BaseModel) else json.dumps(value, default=str) + except Exception as error: + logger.warning("Could not serialize tracing payload: %s", error) + return None + + +def _record_payload( + span: Span, key_pairs: tuple[tuple[str, str | None], ...], value: object, hide: bool, config: TracingConfig +) -> None: + if value is None: + return + serialized = "__REDACTED__" if hide else _serialize_payload(value) + if serialized is None: + return + for key, mime_key in key_pairs: + _write_payload(span, key, mime_key, serialized, hide, config) + + +def _write_payload(span: Span, key: str, mime_key: str | None, serialized: str, hide: bool, config: TracingConfig) -> None: + if hide: + span.set_attribute(key, "__REDACTED__") + return + if len(serialized) > config.max_attribute_length: + span.set_attribute(key + ".original_length", len(serialized)) + serialized = _bounded_text(serialized, config) + span.set_attribute(key, serialized) + if mime_key is not None: + span.set_attribute(mime_key, "application/json") + + +def _payload_keys(primary: str, output: bool) -> tuple[tuple[str, str | None], ...]: + keys: tuple[tuple[str, str | None], ...] = ((primary, None),) + if installation._current_installation().config.openinference_attributes: + keys += ((ATTR_OUTPUT_VALUE, ATTR_OUTPUT_MIME_TYPE),) if output else ((ATTR_INPUT_VALUE, ATTR_INPUT_MIME_TYPE),) + return keys + + +def _set_input_safely(span: Span, value: object, primary: str = ATTR_INPUT) -> None: + config = installation._current_installation().config + _record_payload(span, _payload_keys(primary, False), value, config.hide_inputs, config) + + +def _set_output_safely(span: Span, value: object, primary: str = ATTR_OUTPUT) -> None: + config = installation._current_installation().config + _record_payload(span, _payload_keys(primary, True), value, config.hide_outputs, config) + + +def _bounded_text(value: str, config: TracingConfig) -> str: + limit = config.max_attribute_length + if len(value) > limit: + marker = TRUNCATION_MARKER[:limit] + return value[: limit - len(marker)] + marker + return value + + +def _exception_chain(error: BaseException) -> Iterator[BaseException]: + seen: set[int] = set() + while id(error) not in seen: + seen.add(id(error)) + yield error + following = error.__cause__ or error.__context__ + if following is None: + return + error = following + + +def _instructor_attempt_count(error: InstructorError) -> int: + if error.failed_attempts: + return len(error.failed_attempts) + return getattr(error, "n_attempts", 0) + + +def _failure_count(error: BaseException) -> str | None: + if isinstance(error, InstructorError): + return f"{_instructor_attempt_count(error)} attempts" + if isinstance(error, ValidationError): + return f"{error.error_count()} errors" + return None + + +def _failure_message(error: BaseException, config: TracingConfig) -> str: + if not config.hide_outputs: + return str(error) + count = next((count for cause in _exception_chain(error) if (count := _failure_count(cause)) is not None), "1 errors") + return f"{type(error).__qualname__}: {count} (__REDACTED__)" + + +def _failure_stacktrace(error: BaseException, config: TracingConfig) -> str: + lines = traceback.format_exception(error, chain=not config.hide_outputs) + if config.hide_outputs: + lines = _redacted_traceback(error, config) + return _bounded_text("".join(lines), config) + + +def _redacted_traceback(error: BaseException, config: TracingConfig) -> list[str]: + frames = [ + traceback.FrameSummary(frame.filename, frame.lineno, frame.name, line="") + for frame in traceback.extract_tb(error.__traceback__) + ] + return ["Traceback (most recent call last):\n", *traceback.format_list(frames), _failure_message(error, config) + "\n"] + + +def _record_failure(span: Span, error: BaseException, config: TracingConfig | None = None) -> None: + config = config or installation._current_installation().config + message = _bounded_text(_failure_message(error, config), config) + span.set_attribute(ATTR_ERROR_TYPE, type(error).__qualname__) + span.record_exception( + error, attributes={"exception.message": message, "exception.stacktrace": _failure_stacktrace(error, config)} + ) + span.set_status(Status(StatusCode.ERROR, message)) diff --git a/atomic-agents/atomic_agents/observability/tool_spans.py b/atomic-agents/atomic_agents/observability/tool_spans.py new file mode 100644 index 00000000..7a140aae --- /dev/null +++ b/atomic-agents/atomic_agents/observability/tool_spans.py @@ -0,0 +1,208 @@ +"""Tool spans and tool class instrumentation.""" + +import asyncio +import contextvars +import functools +import inspect +import json +import logging +from contextlib import contextmanager +from types import FunctionType +from typing import Awaitable, Callable, Coroutine, Iterator, List, Optional, Type, TypeVar + +from opentelemetry.trace import Span, Status, StatusCode +from opentelemetry.util.types import AttributeValue + +from atomic_agents.base.base_tool import BaseTool + +from . import installation +from .span_attributes import _best_effort_attributes, _optional_attributes, _suppressed, _without_none +from .span_payloads import _bounded_text, _record_failure, _set_input_safely, _set_output_safely +from .span_attributes import ( + ATTR_OPERATION, + ATTR_TOOL_NAME, + ATTR_TOOL_DESCRIPTION, + ATTR_TOOL_TYPE, + ATTR_TOOL_ARGUMENTS, + ATTR_TOOL_RESULT, + ATTR_OI_TOOL_NAME, + ATTR_OI_TOOL_DESCRIPTION, + ATTR_TOOL_PARAMETERS, + ATTR_TOOL_CLASS, +) + + +logger = logging.getLogger("atomic_agents.observability.tracing") + + +OutputSchema = TypeVar("OutputSchema") + +TOOL_METHOD_NAMES = ("run", "run_async", "arun") + + +_active_tool: contextvars.ContextVar[Optional[BaseTool]] = contextvars.ContextVar("atomic_agents_active_tool", default=None) + + +def _all_subclasses(base: Type) -> List[Type]: + found: List[Type] = [] + pending = list(base.__subclasses__()) + while pending: + current = pending.pop(0) + if current not in found: + found.append(current) + pending.extend(current.__subclasses__()) + return found + + +def _instrument_existing_tool_classes() -> None: + for tool_class in _all_subclasses(BaseTool): + _instrument_tool_class(tool_class) + + +def _hook_future_tool_classes(original_init_subclass: classmethod) -> None: + def traced_init_subclass(cls, **kwargs): + original_init_subclass.__func__(cls, **kwargs) + _instrument_tool_class(cls) + + setattr(BaseTool, "__init_subclass__", classmethod(traced_init_subclass)) + + +def _instrument_tool_class(tool_class: Type) -> None: + for method_name in TOOL_METHOD_NAMES: + method = inspect.getattr_static(tool_class, method_name, None) + if _is_traceable_function(method): + inherited = method_name not in tool_class.__dict__ + setattr(tool_class, method_name, installation._mark_traced(_trace_tool_method(method), method, inherited)) + + +def _is_traceable_function(method: object) -> bool: + if not isinstance(method, FunctionType) or inspect.isasyncgenfunction(method): + return False + if getattr(method, "__isabstractmethod__", False): + return False + return not getattr(method, installation.TRACED_MARKER, False) + + +def _restore_tool_class(tool_class: Type) -> None: + for method_name in TOOL_METHOD_NAMES: + installation._restore_method(tool_class, method_name) + + +def _tool_parameter_attributes(tool: BaseTool) -> dict[str, AttributeValue]: + parameters = _best_effort_attributes({ATTR_TOOL_PARAMETERS: lambda: json.dumps(tool.input_schema.model_json_schema())}) + value = parameters.get(ATTR_TOOL_PARAMETERS) + if isinstance(value, str): + config = installation._current_installation().config + parameters[ATTR_TOOL_PARAMETERS] = _bounded_text(value, config) + if len(value) > config.max_attribute_length: + parameters[ATTR_TOOL_PARAMETERS + ".original_length"] = len(value) + return parameters + + +def _tool_display_name(tool: BaseTool) -> str: + try: + return tool.tool_name + except Exception as error: + logger.debug("Could not collect tracing tool name: %s", error) + return type(tool).__name__ + + +def _tool_start_attributes(tool: BaseTool) -> dict[str, AttributeValue]: + attributes: dict[str, AttributeValue] = { + ATTR_OPERATION: "execute_tool", + ATTR_TOOL_NAME: _tool_display_name(tool), + ATTR_TOOL_CLASS: type(tool).__name__, + ATTR_TOOL_TYPE: "function", + **_optional_attributes("TOOL", {}), + } + attributes.update(_best_effort_attributes({ATTR_TOOL_DESCRIPTION: lambda: tool.tool_description})) + if installation._current_installation().config.openinference_attributes: + attributes.update( + _best_effort_attributes( + { + ATTR_OI_TOOL_NAME: lambda: attributes[ATTR_TOOL_NAME], + ATTR_OI_TOOL_DESCRIPTION: lambda: tool.tool_description, + } + ) + ) + attributes.update(_tool_parameter_attributes(tool)) + return _without_none(attributes) + + +@contextmanager +def _tool_span(tool: BaseTool, args: tuple[object, ...], kwargs: dict[str, object]) -> Iterator[Span]: + attributes = _tool_start_attributes(tool) + with installation._current_installation().tracer.start_as_current_span( + f"execute_tool {attributes[ATTR_TOOL_NAME]}", + attributes=attributes, + record_exception=False, + set_status_on_exception=False, + ) as span: + _set_input_safely(span, _tool_input_from_call(args, kwargs), ATTR_TOOL_ARGUMENTS) + try: + yield span + except (asyncio.CancelledError, GeneratorExit): + span.set_attribute("atomic_agents.cancelled", True) + raise + except BaseException as error: + _record_failure(span, error) + raise + + +def _tool_input_from_call(args: tuple[object, ...], kwargs: dict[str, object]) -> object: + return args[0] if args else kwargs.get("params") + + +def _trace_tool_method( + wrapped: Callable[..., OutputSchema], +) -> Callable[..., OutputSchema] | Callable[..., Coroutine[object, object, object]]: + if inspect.iscoroutinefunction(wrapped): + return _trace_tool_async_method(wrapped) + + @functools.wraps(wrapped) + def traced(tool: BaseTool, *args: object, **kwargs: object) -> OutputSchema: + if _active_tool.get() is tool or _suppressed(): + return wrapped(tool, *args, **kwargs) + with _tool_span(tool, args, kwargs) as span: + result = _run_as_active_tool(tool, wrapped, args, kwargs) + _set_output_safely(span, result, ATTR_TOOL_RESULT) + span.set_status(Status(StatusCode.OK)) + return result + + return traced + + +def _trace_tool_async_method( + wrapped: Callable[..., Awaitable[OutputSchema]], +) -> Callable[..., Coroutine[object, object, OutputSchema]]: + @functools.wraps(wrapped) + async def traced(tool: BaseTool, *args: object, **kwargs: object) -> OutputSchema: + if _active_tool.get() is tool or _suppressed(): + return await wrapped(tool, *args, **kwargs) + with _tool_span(tool, args, kwargs) as span: + result = await _await_as_active_tool(tool, wrapped, args, kwargs) + _set_output_safely(span, result, ATTR_TOOL_RESULT) + span.set_status(Status(StatusCode.OK)) + return result + + return traced + + +def _run_as_active_tool( + tool: BaseTool, wrapped: Callable[..., OutputSchema], args: tuple[object, ...], kwargs: dict[str, object] +) -> OutputSchema: + token = _active_tool.set(tool) + try: + return wrapped(tool, *args, **kwargs) + finally: + _active_tool.reset(token) + + +async def _await_as_active_tool( + tool: BaseTool, wrapped: Callable[..., Awaitable[OutputSchema]], args: tuple[object, ...], kwargs: dict[str, object] +) -> OutputSchema: + token = _active_tool.set(tool) + try: + return await wrapped(tool, *args, **kwargs) + finally: + _active_tool.reset(token) diff --git a/atomic-agents/atomic_agents/observability/tracing.py b/atomic-agents/atomic_agents/observability/tracing.py new file mode 100644 index 00000000..4cb43e57 --- /dev/null +++ b/atomic-agents/atomic_agents/observability/tracing.py @@ -0,0 +1,11 @@ +"""Compatibility exports for Atomic Agents tracing.""" + +from . import installation, span_attributes, tracing_config + +TRACER_NAME = installation.TRACER_NAME +install = installation.install +is_installed = installation.is_installed +uninstall = installation.uninstall +PayloadVisibility = tracing_config.PayloadVisibility +TracingConfig = tracing_config.TracingConfig +using_conversation = span_attributes.using_conversation diff --git a/atomic-agents/atomic_agents/observability/tracing_config.py b/atomic-agents/atomic_agents/observability/tracing_config.py new file mode 100644 index 00000000..be16d881 --- /dev/null +++ b/atomic-agents/atomic_agents/observability/tracing_config.py @@ -0,0 +1,57 @@ +"""Tracing configuration and payload visibility.""" + +import logging +import os +from dataclasses import dataclass, replace +from typing import Protocol + + +logger = logging.getLogger("atomic_agents.observability.tracing") +DEFAULT_MAX_ATTRIBUTE_LENGTH = 4096 + + +class PayloadVisibility(Protocol): + hide_inputs: bool + hide_outputs: bool + + +@dataclass(frozen=True) +class TracingConfig: + hide_inputs: bool = False + hide_outputs: bool = False + max_attribute_length: int = DEFAULT_MAX_ATTRIBUTE_LENGTH + openinference_attributes: bool | None = None + llm_spans: bool = True + + @classmethod + def from_environment(cls) -> "TracingConfig": + return cls( + hide_inputs=_visibility_environment("INPUTS").lower() in ("true", "1"), + hide_outputs=_visibility_environment("OUTPUTS").lower() in ("true", "1"), + max_attribute_length=_attribute_length_environment(), + llm_spans=os.environ.get("ATOMIC_AGENTS_TRACE_LLM_SPANS", "true").lower() not in ("false", "0"), + ) + + @classmethod + def coerce(cls, config: "TracingConfig | PayloadVisibility | None") -> "TracingConfig": + if config is None: + return cls.from_environment() + if isinstance(config, cls): + return config + return replace(cls.from_environment(), hide_inputs=config.hide_inputs, hide_outputs=config.hide_outputs) + + +def _visibility_environment(direction: str) -> str: + return os.environ.get(f"ATOMIC_AGENTS_TRACE_HIDE_{direction}", os.environ.get(f"OPENINFERENCE_HIDE_{direction}", "")) + + +def _attribute_length_environment() -> int: + value = os.environ.get("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", str(DEFAULT_MAX_ATTRIBUTE_LENGTH)) + try: + length = int(value) + if length >= 1: + return length + except ValueError: + pass + logger.warning("Invalid ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH value %r; using %d", value, DEFAULT_MAX_ATTRIBUTE_LENGTH) + return DEFAULT_MAX_ATTRIBUTE_LENGTH diff --git a/atomic-agents/tests/agents/test_agent_config_name.py b/atomic-agents/tests/agents/test_agent_config_name.py new file mode 100644 index 00000000..707e339d --- /dev/null +++ b/atomic-agents/tests/agents/test_agent_config_name.py @@ -0,0 +1,22 @@ +import instructor +from openai import OpenAI + +from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema + + +def test_agent_name_defaults_to_none() -> None: + with OpenAI(api_key="synthetic-test-key") as provider_client: + client = instructor.from_openai(provider_client) + configuration = AgentConfig(client=client) + agent = AtomicAgent[BaseIOSchema, BaseIOSchema](configuration) + assert configuration.name is None + assert agent.name is None + + +def test_explicit_agent_name_is_stored() -> None: + with OpenAI(api_key="synthetic-test-key") as provider_client: + client = instructor.from_openai(provider_client) + configuration = AgentConfig(client=client, name="orchestrator_planning") + agent = AtomicAgent[BaseIOSchema, BaseIOSchema](configuration) + assert configuration.name == "orchestrator_planning" + assert agent.name == "orchestrator_planning" diff --git a/atomic-agents/tests/agents/test_agent_hooks.py b/atomic-agents/tests/agents/test_agent_hooks.py new file mode 100644 index 00000000..0afedfbf --- /dev/null +++ b/atomic-agents/tests/agents/test_agent_hooks.py @@ -0,0 +1,26 @@ +import instructor +import pytest +from unittest.mock import Mock + +from atomic_agents import AgentConfig, AtomicAgent + + +@pytest.mark.parametrize("event", [None, "completion:kwargs"]) +def test_clear_hooks_preserves_foreign_handlers(event): + client = instructor.Instructor(client=Mock(), create=Mock()) + agent = AtomicAgent(AgentConfig(client=client)) + foreign = Mock() + owned = Mock() + other_owned = Mock() + client.on("completion:kwargs", foreign) + agent.register_hook("completion:kwargs", owned) + agent.register_hook("completion:response", other_owned) + + agent.clear_hooks(event) + client.hooks.emit_completion_arguments(model="test-model") + client.hooks.emit_completion_response(Mock()) + + foreign.assert_called_once_with(model="test-model") + owned.assert_not_called() + assert other_owned.call_count == (0 if event is None else 1) + assert agent._hook_handlers.get("completion:kwargs", []) == [] diff --git a/atomic-agents/tests/observability/instructor_transport.py b/atomic-agents/tests/observability/instructor_transport.py new file mode 100644 index 00000000..c1d2a9ad --- /dev/null +++ b/atomic-agents/tests/observability/instructor_transport.py @@ -0,0 +1,104 @@ +"""Offline Instructor clients using real OpenAI request and validation paths.""" + +from dataclasses import dataclass + +import httpx +import instructor +import openai +from openai.types import CompletionUsage +from openai.types.completion_usage import CompletionTokensDetails, PromptTokensDetails + + +@dataclass(frozen=True) +class ScriptedResponse: + arguments: str = '{"answer": "42"}' + input_tokens: int = 3 + output_tokens: int = 2 + cached_tokens: int = 1 + reasoning_tokens: int = 2 + schema_name: str = "Answer" + + def chat_usage(self) -> CompletionUsage: + return CompletionUsage( + prompt_tokens=self.input_tokens, + completion_tokens=self.output_tokens, + total_tokens=self.input_tokens + self.output_tokens, + prompt_tokens_details=PromptTokensDetails(cached_tokens=self.cached_tokens), + completion_tokens_details=CompletionTokensDetails(reasoning_tokens=self.reasoning_tokens), + ) + + def responses_body(self) -> dict[str, object]: + return { + "id": "resp_test", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "response-model", + "output": [ + { + "type": "function_call", + "id": "fc_test", + "call_id": "call_test", + "name": self.schema_name, + "arguments": self.arguments, + "status": "completed", + } + ], + "usage": { + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "total_tokens": self.input_tokens + self.output_tokens, + "input_tokens_details": {"cached_tokens": self.cached_tokens}, + "output_tokens_details": {"reasoning_tokens": self.reasoning_tokens}, + }, + } + + def chat_body(self, mode: instructor.Mode) -> dict[str, object]: + message: dict[str, object] = {"role": "assistant", "content": self.arguments} + if mode == instructor.Mode.TOOLS: + message = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_test", + "type": "function", + "function": {"name": self.schema_name, "arguments": self.arguments}, + } + ], + } + return { + "id": "chatcmpl_test", + "object": "chat.completion", + "created": 1, + "model": "response-model", + "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": self.chat_usage().model_dump(), + } + + +def response_transport(mode: instructor.Mode, responses: list[ScriptedResponse]) -> httpx.MockTransport: + scripted = iter(responses) + + def respond(request: httpx.Request) -> httpx.Response: + response = next(scripted) + body = response.responses_body() if mode.value == "responses_tools" else response.chat_body(mode) + expected_path = "/v1/responses" if mode.value == "responses_tools" else "/v1/chat/completions" + if request.url.path != expected_path: + raise AssertionError(f"Unexpected transport path: {request.url.path}") + return httpx.Response(200, json=body) + + return httpx.MockTransport(respond) + + +def instructor_client(mode: instructor.Mode, responses: list[ScriptedResponse]) -> instructor.Instructor: + return instructor.from_openai( + openai.OpenAI(api_key="test", http_client=httpx.Client(transport=response_transport(mode, responses))), mode=mode + ) + + +def async_instructor_client(mode: instructor.Mode, responses: list[ScriptedResponse]) -> instructor.AsyncInstructor: + return instructor.from_openai( + openai.AsyncOpenAI(api_key="test", http_client=httpx.AsyncClient(transport=response_transport(mode, responses))), + mode=mode, + ) diff --git a/atomic-agents/tests/observability/test_instructor_modes.py b/atomic-agents/tests/observability/test_instructor_modes.py new file mode 100644 index 00000000..7c2e4dce --- /dev/null +++ b/atomic-agents/tests/observability/test_instructor_modes.py @@ -0,0 +1,141 @@ +"""Real Instructor transport, validation and tracing across Ada's request modes.""" + +import asyncio + +import instructor +import pytest +from instructor.core.exceptions import InstructorRetryException +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode + +from instructor_transport import ScriptedResponse, async_instructor_client, instructor_client + +from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema +from atomic_agents.observability import TracingConfig +from atomic_agents.observability import installation + + +class Question(BaseIOSchema): + """A question sent through the real client.""" + + question: str + + +class Answer(BaseIOSchema): + """The validated model answer.""" + + answer: str + + +@pytest.fixture(params=["TOOLS", "JSON", "JSON_SCHEMA", "RESPONSES_TOOLS"]) +def mode(request): + selected = getattr(instructor.Mode, request.param, None) + if selected is None: + pytest.skip(f"Installed Instructor has no Mode.{request.param} for from_openai") + return selected + + +@pytest.fixture(params=[False, True], ids=["run", "run_async"]) +def asynchronous(request): + return request.param + + +@pytest.fixture +def exporter(): + exported = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exported)) + installation.install(tracer_provider=provider, config=TracingConfig(openinference_attributes=False)) + yield exported + installation.uninstall() + provider.shutdown() + + +def run_script(mode, asynchronous, responses): + construct_client = async_instructor_client if asynchronous else instructor_client + client = construct_client(mode, responses) + configured = AtomicAgent[Question, Answer](AgentConfig(client=client, model="test-model")) + configured.model_api_parameters["max_retries"] = 1 + question = Question(question="What is the answer?") + if asynchronous: + return asyncio.run(configured.run_async(question)) + return configured.run(question) + + +def assert_usage(span, input_tokens, output_tokens, cached_tokens, reasoning_tokens): + assert span.attributes["gen_ai.usage.input_tokens"] == input_tokens + assert span.attributes["gen_ai.usage.output_tokens"] == output_tokens + assert span.attributes["gen_ai.usage.cache_read.input_tokens"] == cached_tokens + assert span.attributes["gen_ai.usage.reasoning.output_tokens"] == reasoning_tokens + + +def assert_attempt_spans(exporter, responses): + spans = exporter.get_finished_spans() + parent = spans[-1] + attempts = spans[:-1] + assert parent.attributes["gen_ai.operation.name"] == "invoke_agent" + assert len(attempts) == len(responses) + assert len(spans) == len(responses) + 1 + assert parent.attributes["atomic_agents.llm.attempts"] == len(responses) + assert_models(parent) + for number, (attempt, response) in enumerate(zip(attempts, responses), start=1): + assert attempt.attributes["gen_ai.operation.name"] == "chat" + assert attempt.name == "chat test-model" + assert attempt.kind == SpanKind.CLIENT + assert attempt.parent.span_id == parent.context.span_id + assert attempt.attributes["atomic_agents.llm.attempt"] == number + assert_models(attempt) + assert_usage(attempt, response.input_tokens, response.output_tokens, response.cached_tokens, response.reasoning_tokens) + return attempts, parent + + +def assert_models(span): + assert span.attributes["gen_ai.request.model"] == "test-model" + assert span.attributes["gen_ai.response.model"] == "response-model" + assert span.attributes["gen_ai.provider.name"] == "openai" + + +def test_success_model_usage_and_attempt_count(mode, asynchronous, exporter): + responses = [ScriptedResponse()] + assert run_script(mode, asynchronous, responses).answer == "42" + attempts, parent = assert_attempt_spans(exporter, responses) + assert attempts[0].status.status_code == parent.status.status_code == StatusCode.OK + assert_usage(parent, 3, 2, 1, 2) + + +def test_validation_retry_preserves_each_attempt_and_total_usage(mode, asynchronous, exporter): + responses = [ + ScriptedResponse(arguments='{"answer": {"secret": "private output"}}'), + ScriptedResponse(input_tokens=7, output_tokens=5, cached_tokens=3, reasoning_tokens=4), + ] + assert run_script(mode, asynchronous, responses).answer == "42" + attempts, parent = assert_attempt_spans(exporter, responses) + assert [attempt.status.status_code for attempt in attempts] == [StatusCode.ERROR, StatusCode.OK] + assert parent.status.status_code == StatusCode.OK + assert any(event.name == "validation_error" for event in attempts[0].events) + assert any(event.name == "validation_error" for event in parent.events) + assert_usage(parent, 10, 7, 4, 6) + + +def test_hidden_retry_exhaustion_redacts_all_span_failures(mode, asynchronous, exporter): + installed = installation._current_installation() + installation.install( + tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=True, openinference_attributes=False) + ) + responses = [ScriptedResponse(arguments='{"answer": {"secret": "private model output"}}')] * 2 + with pytest.raises(InstructorRetryException): + run_script(mode, asynchronous, responses) + attempts, parent = assert_attempt_spans(exporter, responses) + assert [attempt.status.status_code for attempt in attempts] == [StatusCode.ERROR, StatusCode.ERROR] + assert parent.status.status_code == StatusCode.ERROR + assert "__REDACTED__" in parent.attributes["atomic_agents.llm.last_validation_error"] + assert_usage(parent, 6, 4, 2, 4) + for span in (*attempts, parent): + assert "__REDACTED__" in span.status.description + assert "private model output" not in str(span.attributes) + str(span.events) + span.status.description + for event in span.events: + if event.name == "exception": + assert "__REDACTED__" in event.attributes["exception.message"] + assert "__REDACTED__" in event.attributes["exception.stacktrace"] diff --git a/atomic-agents/tests/observability/test_llm_attempt_spans.py b/atomic-agents/tests/observability/test_llm_attempt_spans.py new file mode 100644 index 00000000..8c21b822 --- /dev/null +++ b/atomic-agents/tests/observability/test_llm_attempt_spans.py @@ -0,0 +1,539 @@ +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor, TimeoutError +from types import SimpleNamespace +from unittest.mock import Mock + +import instructor +import pytest +from instructor.core.hooks import Hooks +from opentelemetry import context as context_api, trace as trace_api +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode +from pydantic import ValidationError +from instructor.core.exceptions import InstructorRetryException + +from instructor_transport import ScriptedResponse, instructor_client + +from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema +from atomic_agents.observability import TracingConfig, using_conversation +from atomic_agents.observability import agent_spans, installation, llm_attempt_spans + + +class Question(BaseIOSchema): + """A test question.""" + + question: str + + +class Answer(BaseIOSchema): + """A test answer.""" + + answer: str + + +@pytest.fixture +def exporter(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + installation.install(tracer_provider=provider, config=TracingConfig(openinference_attributes=False)) + yield exporter + installation.uninstall() + + +def usage(input_tokens=3, output_tokens=2): + return ScriptedResponse(input_tokens=input_tokens, output_tokens=output_tokens).chat_usage() + + +def finished_spans(exporter, count): + spans = exporter.get_finished_spans() + assert len(spans) == count + return spans + + +def response(): + return SimpleNamespace(model="response-model", usage=usage()) + + +def scripted_client(script): + client = Mock(spec=instructor.Instructor) + client.provider = instructor.Provider.OPENAI + client.hooks = Hooks() + client.on.side_effect = client.hooks.on + client.off.side_effect = client.hooks.off + client.chat = Mock() + client.chat.completions = Mock() + + def create(**kwargs): + script(client.hooks, kwargs) + return Answer(answer="42") + + client.chat.completions.create.side_effect = create + return client + + +def agent(client): + return AtomicAgent[Question, Answer](AgentConfig(client=client, model="test-model")) + + +def retry_script(hooks, kwargs): + hooks.emit_completion_arguments(**kwargs) + hooks.emit_completion_response(response()) + hooks.emit_parse_error(ValueError("invalid output")) + hooks.emit_completion_arguments(**kwargs) + hooks.emit_completion_response(response()) + + +def run_agent(client): + return agent(client).run(Question(question="test")) + + +def test_retry_spans_usage_validation_and_parentage(exporter): + client = instructor_client( + instructor.Mode.TOOLS, [ScriptedResponse(arguments='{"answer": {"secret": "private output"}}'), ScriptedResponse()] + ) + with using_conversation("conversation"): + run_agent(client) + first, second, parent = finished_spans(exporter, 3) + assert [first.name, second.name] == ["chat test-model", "chat test-model"] + assert [first.attributes["atomic_agents.llm.attempt"], second.attributes["atomic_agents.llm.attempt"]] == [1, 2] + assert first.parent.span_id == second.parent.span_id == parent.context.span_id + assert first.status.status_code == StatusCode.ERROR + assert second.status.status_code == StatusCode.OK + assert first.kind == second.kind == SpanKind.CLIENT + assert first.events[0].name == "validation_error" + assert first.attributes["error.type"] == "ValidationError" + assert first.attributes["gen_ai.provider.name"] == "openai" + assert first.attributes["gen_ai.conversation.id"] == "conversation" + assert first.attributes["gen_ai.usage.input_tokens"] == second.attributes["gen_ai.usage.input_tokens"] == 3 + assert first.attributes["gen_ai.usage.output_tokens"] == 2 + assert first.attributes["gen_ai.usage.cache_read.input_tokens"] == 1 + assert first.attributes["gen_ai.usage.reasoning.output_tokens"] == 2 + assert parent.attributes["gen_ai.usage.input_tokens"] == 6 + assert parent.attributes["gen_ai.usage.output_tokens"] == 4 + assert parent.attributes.get("atomic_agents.llm.attempts") == 2 + assert "private output" in parent.attributes["atomic_agents.llm.last_validation_error"] + assert parent.attributes["gen_ai.response.model"] == first.attributes["gen_ai.response.model"] == "response-model" + event = parent.events[0] + assert event.name == "validation_error" + assert event.attributes["atomic_agents.llm.attempt"] == 1 + + +def test_disabled_attempt_spans_keep_totals(exporter): + installed = installation._current_installation() + installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(llm_spans=False)) + run_agent(scripted_client(retry_script)) + (parent,) = finished_spans(exporter, 1) + assert parent.attributes.get("atomic_agents.llm.attempts") == 2 + assert parent.attributes["gen_ai.usage.input_tokens"] == 6 + + +def test_completion_error_is_bounded_and_ends_attempt(exporter): + installed = installation._current_installation() + installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(max_attribute_length=32)) + error = ValueError("failed " * 20) + expected = str(error)[:18] + "...[truncated]" + + def script(hooks, kwargs): + hooks.emit_completion_arguments(**kwargs) + hooks.emit_parse_error(error) + hooks.emit_completion_error(error, attempt_number=1, max_attempts=1, is_last_attempt=True) + + run_agent(scripted_client(script)) + child, parent = finished_spans(exporter, 2) + assert child.status.status_code == StatusCode.ERROR + assert child.attributes["error.type"] == "ValueError" + assert child.events[-1].attributes["exception.message"] == expected + assert len(child.events[-1].attributes["exception.stacktrace"]) <= 32 + assert child.status.description == expected + assert len(child.status.description) == 32 + assert parent.attributes["atomic_agents.llm.last_validation_error"] == expected + assert parent.attributes["atomic_agents.llm.last_validation_error.original_length"] == len(str(error)) + assert parent.status.status_code == StatusCode.OK + + +def test_registration_once_and_uninstall_preserves_application_hooks(exporter): + client = scripted_client(retry_script) + application_hook = Mock() + client.hooks.on("completion:response", application_hook) + for _ in range(3): + run_agent(client) + assert client.on.call_count == 4 + assert application_hook.call_count == 6 + installation.uninstall() + assert client.off.call_count == 4 + client.hooks.emit_completion_response(response()) + assert application_hook.call_count == 7 + assert not llm_attempt_spans._clients + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_concurrent_shared_client_runs_are_isolated(exporter, asynchronous): + client = scripted_client(retry_script) + + async def invoke(number): + async def complete(agent, *args, **kwargs): + client.hooks.emit_completion_arguments(model="test-model") + await asyncio.sleep(0) + client.hooks.emit_completion_response(SimpleNamespace(model="response-model", usage=usage(number))) + return Answer(answer="ok") + + if asynchronous: + return await agent_spans._trace_agent_async_call(complete)(agent(client)) + return await asyncio.to_thread(run_agent, client) + + async def gather(): + await asyncio.gather(invoke(10), invoke(20)) + + asyncio.run(gather()) + parents = [span for span in exporter.get_finished_spans() if span.name.startswith("invoke_agent")] + expected = [10, 20] if asynchronous else [6, 6] + assert sorted(span.attributes.get("gen_ai.usage.input_tokens", -1) for span in parents) == expected + assert all(span.attributes["atomic_agents.llm.attempts"] == (1 if asynchronous else 2) for span in parents) + + +@pytest.mark.parametrize("kind", ["sync", "async", "stream", "async_stream"]) +@pytest.mark.parametrize("fails", [False, True]) +def test_open_attempt_ends_at_agent_end(exporter, kind, fails): + client = scripted_client(retry_script) + + def start(): + client.hooks.emit_completion_arguments(model="test-model") + if fails: + raise ValueError("agent failed") + return Answer(answer="partial") + + def call(*args, **kwargs): + return start() + + async def async_call(*args, **kwargs): + return start() + + def stream(*args, **kwargs): + yield start() + yield Answer(answer="final") + + async def async_stream(*args, **kwargs): + yield start() + yield Answer(answer="final") + + def invoke(): + if kind == "sync": + return agent_spans._trace_agent_call(call)(agent(client)) + if kind == "async": + return asyncio.run(agent_spans._trace_agent_async_call(async_call)(agent(client))) + if kind == "stream": + iterator = agent_spans._trace_agent_stream(stream)(agent(client)) + next(iterator) + assert trace_api.get_current_span() is trace_api.INVALID_SPAN + iterator.close() + return None + return asyncio.run(consume_async_stream()) + + async def consume_async_stream(): + iterator = agent_spans._trace_agent_async_stream(async_stream)(agent(client)) + await anext(iterator) + assert trace_api.get_current_span() is trace_api.INVALID_SPAN + await iterator.aclose() + + if fails: + with pytest.raises(ValueError, match="agent failed"): + invoke() + else: + invoke() + child, parent = finished_spans(exporter, 2) + assert child.status.status_code == StatusCode.UNSET + assert child.parent.span_id == parent.context.span_id + assert child.end_time <= parent.end_time + + +def test_validation_text_redacted_and_anthropic_usage(exporter): + installed = installation._current_installation() + installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=True)) + + def script(hooks, kwargs): + hooks.emit_completion_arguments(**kwargs) + with pytest.raises(ValidationError) as invalid: + Answer.model_validate({"answer": {"secret": "private output"}}) + hooks.emit_completion_response(SimpleNamespace(model="response-model", usage=usage(9, 4))) + hooks.emit_parse_error(invalid.value) + hooks.emit_completion_error(invalid.value, attempt_number=1, max_attempts=1, is_last_attempt=True) + + run_agent(scripted_client(script)) + child, parent = finished_spans(exporter, 2) + assert parent.attributes["atomic_agents.llm.last_validation_error"] == "__REDACTED__" + assert parent.events[0].attributes["atomic_agents.llm.last_validation_error"] == "__REDACTED__" + assert parent.attributes["gen_ai.usage.input_tokens"] == 9 + assert child.status.description == "ValidationError: 1 errors (__REDACTED__)" + assert child.events[-1].attributes["exception.message"] == child.status.description + assert "ValidationError: 1 errors (__REDACTED__)" in child.events[-1].attributes["exception.stacktrace"] + assert "private output" not in str(child.events) + + +def test_hook_failure_warns_without_breaking_agent(exporter, caplog): + class BrokenResponse: + @property + def model(self): + raise RuntimeError("broken metadata") + + def script(hooks, kwargs): + hooks.emit_completion_arguments(**kwargs) + hooks.emit_completion_response(BrokenResponse()) + + assert run_agent(scripted_client(script)).answer == "42" + child, parent = finished_spans(exporter, 2) + assert child.status.status_code == StatusCode.UNSET + assert "Instructor tracing hook failed" in caplog.text + + +@pytest.mark.parametrize("asynchronous", [False, True]) +def test_stream_attempt_context_survives_cursor_advances(exporter, asynchronous): + client = scripted_client(retry_script) + + def begin(): + client.hooks.emit_completion_arguments(model="test-model") + return trace_api.get_current_span().get_span_context().span_id + + def complete(expected): + assert trace_api.get_current_span().get_span_context().span_id == expected + client.hooks.emit_completion_response(response()) + return Answer(answer="final") + + def stream(*args, **kwargs): + expected = begin() + yield Answer(answer="partial") + yield complete(expected) + + async def async_stream(*args, **kwargs): + expected = begin() + yield Answer(answer="partial") + await asyncio.sleep(0) + yield complete(expected) + + async def consume(): + return [partial async for partial in agent_spans._trace_agent_async_stream(async_stream)(agent(client))] + + if asynchronous: + assert len(asyncio.run(consume())) == 2 + else: + assert len(list(agent_spans._trace_agent_stream(stream)(agent(client)))) == 2 + child, parent = finished_spans(exporter, 2) + assert child.status.status_code == StatusCode.OK + assert parent.status.status_code == StatusCode.OK + + +def test_validation_text_truncation_and_extra_hook_keywords(exporter): + installed = installation._current_installation() + installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(max_attribute_length=4)) + + def script(hooks, kwargs): + hooks.emit_completion_arguments(**kwargs) + llm_attempt_spans._parse_error(ValueError("long error"), future_option=True) + llm_attempt_spans._completion_response(response(), future_option=True) + + run_agent(scripted_client(script)) + child, parent = finished_spans(exporter, 2) + assert parent.attributes["atomic_agents.llm.last_validation_error"] == "...[" + assert parent.attributes["atomic_agents.llm.last_validation_error.original_length"] == 10 + + +def test_threaded_first_use_registers_once_and_uninstalls(exporter, monkeypatch): + client = scripted_client(retry_script) + first_registration = threading.Barrier(2, timeout=5) + release_registration = threading.Event() + original_on = Hooks.on + + def blocked_on(hooks, event, handler): + if event == "completion:kwargs": + first_registration.wait() + assert release_registration.wait(timeout=5) + return original_on(hooks, event, handler) + + monkeypatch.setattr(Hooks, "on", blocked_on) + client.on.side_effect = client.hooks.on + with ThreadPoolExecutor(max_workers=2) as executor: + first = executor.submit(run_agent, client) + first_registration.wait() + second_started = threading.Event() + + def second_run(): + second_started.set() + return run_agent(client) + + second = executor.submit(second_run) + assert second_started.wait(timeout=5) + try: + with pytest.raises(TimeoutError): + second.result(timeout=0.1) + assert client.on.call_count == 1 + finally: + release_registration.set() + first.result(timeout=5) + second.result(timeout=5) + assert len(client.hooks._handlers) == 4 + assert all(len(handlers) == 1 for handlers in client.hooks._handlers.values()) + installation.uninstall() + assert all(not handlers for handlers in client.hooks._handlers.values()) + + +def test_partial_registration_is_rolled_back(exporter): + client = scripted_client(retry_script) + + def fail_registration(event, handler): + if event == "parse:error": + raise RuntimeError("registration failed") + client.hooks.on(event, handler) + + client.on.side_effect = fail_registration + llm_attempt_spans.register_client(client) + assert client not in llm_attempt_spans._clients + assert all(not handlers for handlers in client.hooks._handlers.values()) + client.on.side_effect = client.hooks.on + llm_attempt_spans.register_client(client) + assert all(len(handlers) == 1 for handlers in client.hooks._handlers.values()) + + +@pytest.mark.parametrize("completion", ["response", "error"]) +def test_suppressed_attempt_events_are_ignored_then_normal_attempt_works(exporter, completion): + def script(hooks, kwargs): + run = llm_attempt_spans._active_run.get() + token = context_api.attach(context_api.set_value("suppress_instrumentation", True)) + try: + hooks.emit_completion_arguments(**kwargs) + finally: + context_api.detach(token) + if completion == "response": + hooks.emit_completion_response(response()) + else: + hooks.emit_completion_error(ValueError("suppressed"), attempt_number=1, max_attempts=2, is_last_attempt=False) + hooks.emit_parse_error(ValueError("suppressed")) + assert run.attempt_count == 0 + assert run.open_attempt is None + assert not exporter.get_finished_spans() + hooks.emit_completion_arguments(**kwargs) + hooks.emit_completion_response(response()) + + run_agent(scripted_client(script)) + child, parent = finished_spans(exporter, 2) + assert child.attributes["atomic_agents.llm.attempt"] == 1 + assert parent.attributes["atomic_agents.llm.attempts"] == 1 + assert parent.attributes["gen_ai.usage.input_tokens"] == 3 + assert not parent.events + + +@pytest.mark.parametrize( + "mode", [instructor.Mode.TOOLS, instructor.Mode.JSON, instructor.Mode.JSON_SCHEMA, instructor.Mode.RESPONSES_TOOLS] +) +def test_real_transport_usage_is_accumulated_per_response(exporter, mode): + client = instructor_client( + mode, + [ + ScriptedResponse(arguments='{"answer": {"private": "model text"}}'), + ScriptedResponse(input_tokens=7, output_tokens=5, cached_tokens=3, reasoning_tokens=4), + ], + ) + assert run_agent(client).answer == "42" + first, second, parent = finished_spans(exporter, 3) + assert first.status.status_code == StatusCode.ERROR + assert first.events[0].name == "validation_error" + assert second.status.status_code == StatusCode.OK + assert first.attributes["gen_ai.usage.input_tokens"] == 3 + assert second.attributes["gen_ai.usage.input_tokens"] == 7 + assert parent.attributes["gen_ai.usage.input_tokens"] == 10 + assert parent.attributes["gen_ai.usage.output_tokens"] == 7 + assert parent.attributes["gen_ai.usage.cache_read.input_tokens"] == 4 + assert parent.attributes["gen_ai.usage.reasoning.output_tokens"] == 6 + assert parent.attributes["atomic_agents.llm.attempts"] == 2 + assert "completion:usage" not in llm_attempt_spans._HANDLERS + + +@pytest.mark.parametrize("hidden", [False, True]) +def test_real_retry_exhaustion_redacts_or_bounds_every_exception(exporter, hidden): + installed = installation._current_installation() + installation.install( + tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=hidden, max_attribute_length=256) + ) + secret = "private model text " * 100 + client = instructor_client( + instructor.Mode.TOOLS, [ScriptedResponse(arguments='{"answer": {"secret": "' + secret + '"}}')] * 3 + ) + configured = agent(client) + configured.model_api_parameters["max_retries"] = 2 + with pytest.raises(InstructorRetryException): + configured.run(Question(question="test")) + first, second, third, parent = finished_spans(exporter, 4) + assert parent.status.status_code == StatusCode.ERROR + assert parent.attributes["atomic_agents.llm.attempts"] == 3 + assert parent.attributes["gen_ai.usage.input_tokens"] == 9 + message = parent.events[-1].attributes["exception.message"] + assert len(message) <= 256 + assert len(parent.status.description) <= 256 + assert len(parent.events[-1].attributes["exception.stacktrace"]) <= 256 + if hidden: + assert message == parent.status.description == "InstructorRetryException: 3 attempts (__REDACTED__)" + assert "private model text" not in str(parent.events) + assert "private model text" not in str(first.events) + str(second.events) + str(third.events) + assert ( + first.events[0].attributes["atomic_agents.llm.last_validation_error"] == "ValidationError: 1 errors (__REDACTED__)" + ) + else: + assert message.endswith("...[truncated]") + + +def test_direct_instructor_calls_inside_run_add_to_agent_totals(exporter): + client = instructor_client(instructor.Mode.TOOLS, [ScriptedResponse(), ScriptedResponse(input_tokens=11)]) + + def call_twice(configured): + for _ in range(2): + client.chat.completions.create( + model="test-model", response_model=Answer, messages=[{"role": "user", "content": "test"}] + ) + return Answer(answer="42") + + agent_spans._trace_agent_call(call_twice)(agent(client)) + first, second, parent = finished_spans(exporter, 3) + assert parent.attributes["gen_ai.usage.input_tokens"] == 14 + assert parent.attributes["atomic_agents.llm.attempts"] == 2 + assert first.parent.span_id == second.parent.span_id == parent.context.span_id + + +def test_missing_model_and_final_completion_error_still_count(exporter): + def script(hooks, kwargs): + hooks.emit_completion_arguments() + hooks.emit_completion_response(response()) + hooks.emit_completion_arguments() + hooks.emit_completion_error(RuntimeError("transport failed")) + + run_agent(scripted_client(script)) + first, second, parent = finished_spans(exporter, 3) + assert first.name == second.name == "chat" + assert first.kind == second.kind == SpanKind.CLIENT + assert second.status.status_code == StatusCode.ERROR + assert parent.attributes["atomic_agents.llm.attempts"] == 2 + assert parent.attributes["gen_ai.usage.input_tokens"] == 3 + + +@pytest.mark.parametrize("chain", ["__cause__", "__context__"]) +def test_wrapped_validation_error_and_tool_exception_are_redacted(exporter, chain): + from atomic_agents.observability import tool_spans + from atomic_agents.base.base_tool import BaseTool + + installed = installation._current_installation() + installation.install(tracer_provider=installed.tracer_provider, config=TracingConfig(hide_outputs=True)) + with pytest.raises(ValidationError) as invalid: + Answer.model_validate({"answer": {"secret": "private model text"}}) + wrapped = RuntimeError("private wrapper text") + setattr(wrapped, chain, invalid.value) + with pytest.raises(RuntimeError), tool_spans._tool_span(Mock(spec=BaseTool, tool_name="test"), (), {}): + raise wrapped + (span,) = finished_spans(exporter, 1) + assert span.status.description == "RuntimeError: 1 errors (__REDACTED__)" + assert span.events[0].attributes["exception.message"] == span.status.description + stacktrace = span.events[0].attributes["exception.stacktrace"] + assert "Traceback (most recent call last)" in stacktrace + assert stacktrace.rstrip().endswith(span.status.description) + assert "private model text" not in stacktrace + assert "private wrapper text" not in stacktrace diff --git a/atomic-agents/tests/observability/test_openinference_tracing.py b/atomic-agents/tests/observability/test_openinference_tracing.py deleted file mode 100644 index 6fb3bc88..00000000 --- a/atomic-agents/tests/observability/test_openinference_tracing.py +++ /dev/null @@ -1,851 +0,0 @@ -import abc -import asyncio -import contextvars -import gc -import json -import logging -import types -from concurrent.futures import ThreadPoolExecutor -from importlib.metadata import entry_points -from typing import Dict, List -from unittest.mock import Mock - -import instructor -import pytest -from openinference.instrumentation import OITracer, TraceConfig, using_metadata, using_session -from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes -from opentelemetry import trace as trace_api -from opentelemetry.sdk.trace import ReadableSpan, TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from opentelemetry.trace import StatusCode -from pydantic import Field - -from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BaseTool -from atomic_agents.context import BaseDynamicContextProvider, SystemPromptGenerator -from atomic_agents.observability import AtomicAgentsInstrumentor -import atomic_agents.observability as observability -from atomic_agents.observability import openinference_tracing as tracing - - -class QuestionSchema(BaseIOSchema): - """A question for the agent.""" - - question: str = Field(...) - - -class AnswerSchema(BaseIOSchema): - """An answer from the agent.""" - - answer: str = Field(...) - - -class CalculationInputSchema(BaseIOSchema): - """Expression to evaluate.""" - - expression: str = Field(...) - - -class CalculationOutputSchema(BaseIOSchema): - """Result of the evaluation.""" - - result: str = Field(...) - - -class CalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema, precision: int = 2) -> CalculationOutputSchema: - return CalculationOutputSchema(result=f"{len(params.expression):.{precision}f}") - - -class AsyncCalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return asyncio.run(self.arun(params)) - - async def arun(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result=str(len(params.expression))) - - -class ThreadedCalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result=str(len(params.expression))) - - async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return await asyncio.to_thread(self.run, params) - - -class RunMixin: - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="mixin") - - -class MixinCalculatorTool(RunMixin, BaseTool[CalculationInputSchema, CalculationOutputSchema]): - pass - - -class StaticRunTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - @staticmethod - def run(params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="static") - - -class FailingTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - raise ValueError("boom") - - -class AnswerStore: - """Container object pydantic cannot serialize to JSON.""" - - def __init__(self): - self.values: List[int] = [] - - -class UnserializableInputSchema(BaseIOSchema): - """Input carrying an arbitrary object.""" - - model_config = {"arbitrary_types_allowed": True} - store: AnswerStore = Field(...) - - -class EchoTool(BaseTool[UnserializableInputSchema, CalculationOutputSchema]): - def run(self, params: UnserializableInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="echo") - - -@pytest.fixture -def exporter(): - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - exporter.provider = provider - instrumentor = AtomicAgentsInstrumentor() - instrumentor.instrument(tracer_provider=provider) - yield exporter - instrumentor.uninstrument() - - -@pytest.fixture -def provider_tracer(exporter): - return exporter.provider.get_tracer("test") - - -def sync_client(response=None, partials=None): - client = Mock(spec=instructor.Instructor) - client.chat = Mock() - client.chat.completions = Mock() - client.chat.completions.create = Mock(return_value=response or AnswerSchema(answer="42")) - - def create_partial(**kwargs): - for partial in partials or [AnswerSchema(answer="4"), AnswerSchema(answer="42")]: - yield partial - - client.chat.completions.create_partial = create_partial - return client - - -def async_client(response=None, partials=None): - client = Mock(spec=instructor.AsyncInstructor) - client.chat = Mock() - client.chat.completions = Mock() - - async def create(**kwargs): - return response or AnswerSchema(answer="42") - - async def create_partial(**kwargs): - for partial in partials or [AnswerSchema(answer="4"), AnswerSchema(answer="42")]: - yield partial - - client.chat.completions.create = create - client.chat.completions.create_partial = create_partial - return client - - -def make_agent(client) -> AtomicAgent[QuestionSchema, AnswerSchema]: - return AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=client, model="test-model")) - - -def spans_of(exporter) -> List[ReadableSpan]: - return list(exporter.get_finished_spans()) - - -def kind_of(span: ReadableSpan) -> str: - return span.attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] - - -def input_json(span: ReadableSpan) -> Dict: - return json.loads(span.attributes[SpanAttributes.INPUT_VALUE]) - - -def output_json(span: ReadableSpan) -> Dict: - return json.loads(span.attributes[SpanAttributes.OUTPUT_VALUE]) - - -def test_run_emits_agent_span_with_io_and_framework_attributes(exporter): - agent = make_agent(sync_client()) - - response = agent.run(QuestionSchema(question="6*7?")) - - assert response.answer == "42" - (span,) = spans_of(exporter) - assert span.name == "AtomicAgent[QuestionSchema, AnswerSchema].run" - assert kind_of(span) == OpenInferenceSpanKindValues.AGENT.value - assert span.attributes[SpanAttributes.AGENT_NAME] == "AtomicAgent[QuestionSchema, AnswerSchema]" - assert input_json(span) == {"question": "6*7?"} - assert output_json(span) == {"answer": "42"} - assert span.attributes[SpanAttributes.INPUT_MIME_TYPE] == "application/json" - assert span.attributes[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" - assert span.attributes[tracing.ATTR_MODEL] == "test-model" - assert span.attributes[tracing.ATTR_MODE] == agent.mode.value - assert span.attributes[tracing.ATTR_INPUT_SCHEMA] == "QuestionSchema" - assert span.attributes[tracing.ATTR_OUTPUT_SCHEMA] == "AnswerSchema" - assert span.attributes[tracing.ATTR_HISTORY_MESSAGE_COUNT] == 0 - assert span.attributes[tracing.ATTR_CONTEXT_PROVIDERS] == () - assert len(span.attributes[tracing.ATTR_SYSTEM_PROMPT_SHA256]) == 64 - assert span.status.status_code == StatusCode.OK - - -def test_second_turn_reports_history_size_entering_the_turn(exporter): - agent = make_agent(sync_client()) - agent.run(QuestionSchema(question="first")) - - agent.run(QuestionSchema(question="second")) - - first, second = spans_of(exporter) - assert first.attributes[tracing.ATTR_HISTORY_MESSAGE_COUNT] == 0 - assert second.attributes[tracing.ATTR_HISTORY_MESSAGE_COUNT] == 2 - - -def test_agent_without_system_role_has_no_prompt_hash(exporter): - agent = AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=sync_client(), model="m", system_role=None)) - - agent.run(QuestionSchema(question="?")) - - (span,) = spans_of(exporter) - assert tracing.ATTR_SYSTEM_PROMPT_SHA256 not in span.attributes - - -def test_abstract_intermediate_tool_base_is_not_wrapped(exporter): - class AbstractCalculator(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - @abc.abstractmethod - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - raise NotImplementedError - - assert not hasattr(AbstractCalculator.__dict__["run"], tracing.TRACED_MARKER) - - -def test_run_without_input_sets_no_input_attribute(exporter): - agent = make_agent(sync_client()) - - agent.run() - - (span,) = spans_of(exporter) - assert SpanAttributes.INPUT_VALUE not in span.attributes - assert output_json(span) == {"answer": "42"} - - -def test_run_failure_propagates_and_marks_span_error(exporter): - client = sync_client() - client.chat.completions.create = Mock(side_effect=RuntimeError("provider down")) - agent = make_agent(client) - - with pytest.raises(RuntimeError, match="provider down"): - agent.run(QuestionSchema(question="?")) - - (span,) = spans_of(exporter) - assert span.status.status_code == StatusCode.ERROR - assert [event.name for event in span.events] == ["exception"] - - -def test_subclassed_agent_is_named_after_the_subclass_and_emits_one_span(exporter): - class ResearchAgent(AtomicAgent[QuestionSchema, AnswerSchema]): - def run(self, user_input=None): - return super().run(user_input) - - agent = ResearchAgent(AgentConfig(client=sync_client(), model="test-model")) - - agent.run(QuestionSchema(question="?")) - - (span,) = spans_of(exporter) - assert span.name == "ResearchAgent[QuestionSchema, AnswerSchema].run" - - -def test_agent_span_nests_under_a_user_chain_span(exporter, provider_tracer): - agent = make_agent(sync_client()) - - with provider_tracer.start_as_current_span("turn") as turn: - agent.run(QuestionSchema(question="?")) - - agent_span, turn_span = spans_of(exporter) - assert agent_span.parent.span_id == turn.get_span_context().span_id - assert turn_span.name == "turn" - - -def test_manual_setup_from_the_guide_groups_agents_and_tools_under_one_chain_span(): - exporter = InMemorySpanExporter() - tracer_provider = TracerProvider() - tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) - trace_config = TraceConfig() - instrumentor = AtomicAgentsInstrumentor() - instrumentor.instrument(tracer_provider=tracer_provider, config=trace_config) - tracer = OITracer(tracer_provider.get_tracer(__name__), config=trace_config) - planner = make_agent(sync_client()) - answerer = make_agent(sync_client()) - calculator = CalculatorTool() - try: - with tracer.start_as_current_span("answer_question", openinference_span_kind="chain") as turn: - turn.set_input("what is 2 + 2?") - plan = planner.run(QuestionSchema(question="what is 2 + 2?")) - calculation = calculator.run(CalculationInputSchema(expression=plan.answer)) - answer = answerer.run(QuestionSchema(question=str(calculation.result))) - turn.set_output(answer.answer) - finally: - instrumentor.uninstrument() - - spans = spans_of(exporter) - chain_span = next(span for span in spans if kind_of(span) == OpenInferenceSpanKindValues.CHAIN.value) - children = [span for span in spans if span is not chain_span] - assert [kind_of(span) for span in children] == ["AGENT", "TOOL", "AGENT"] - assert all(span.parent.span_id == chain_span.get_span_context().span_id for span in children) - assert chain_span.attributes[SpanAttributes.INPUT_VALUE] == "what is 2 + 2?" - assert chain_span.attributes[SpanAttributes.OUTPUT_VALUE] == answer.answer - - -def test_llm_call_inside_run_becomes_child_of_agent_span(exporter, provider_tracer): - def create(**kwargs): - with provider_tracer.start_as_current_span("llm"): - return AnswerSchema(answer="42") - - client = sync_client() - client.chat.completions.create = create - agent = make_agent(client) - - agent.run(QuestionSchema(question="?")) - - llm_span, agent_span = spans_of(exporter) - assert llm_span.parent.span_id == agent_span.context.span_id - - -def test_run_stream_records_final_output_and_keeps_child_context_per_step(exporter, provider_tracer): - current_between_yields = [] - - def create_partial(**kwargs): - yield AnswerSchema(answer="4") - with provider_tracer.start_as_current_span("llm-chunk"): - yield AnswerSchema(answer="42") - - client = sync_client() - client.chat.completions.create_partial = create_partial - agent = make_agent(client) - - partials = [] - for partial in agent.run_stream(QuestionSchema(question="?")): - partials.append(partial.answer) - current_between_yields.append(trace_api.get_current_span()) - - assert partials == ["4", "42"] - assert all(span is trace_api.INVALID_SPAN for span in current_between_yields) - chunk_span, agent_span = spans_of(exporter) - assert agent_span.name == "AtomicAgent[QuestionSchema, AnswerSchema].run_stream" - assert chunk_span.parent.span_id == agent_span.context.span_id - assert output_json(agent_span) == {"answer": "42"} - assert agent_span.status.status_code == StatusCode.OK - - -def test_run_stream_preserves_the_generator_return_value(exporter): - agent = make_agent(sync_client()) - - def consume(): - final = yield from agent.run_stream(QuestionSchema(question="?")) - return final - - generator = consume() - with pytest.raises(StopIteration) as stop: - while True: - next(generator) - - assert stop.value.value.answer == "42" - - -def test_run_stream_with_no_partials_ends_ok_without_output(exporter): - client = sync_client() - client.chat.completions.create_partial = lambda **kwargs: iter([]) - agent = make_agent(client) - - assert list(agent.run_stream(QuestionSchema(question="?"))) == [] - - (span,) = spans_of(exporter) - assert span.status.status_code == StatusCode.OK - assert SpanAttributes.OUTPUT_VALUE not in span.attributes - - -def test_run_stream_closed_early_ends_span_and_closes_inner_generator(exporter): - inner_closed = [] - - def create_partial(**kwargs): - try: - yield AnswerSchema(answer="4") - yield AnswerSchema(answer="42") - finally: - inner_closed.append(True) - - client = sync_client() - client.chat.completions.create_partial = create_partial - agent = make_agent(client) - - for partial in agent.run_stream(QuestionSchema(question="?")): - break - - assert inner_closed == [True] - (span,) = spans_of(exporter) - assert span.attributes[tracing.ATTR_STREAM_CLOSED_EARLY] is True - assert output_json(span) == {"answer": "4"} - assert span.status.status_code == StatusCode.OK - - -def test_run_stream_failure_mid_stream_marks_error(exporter): - def create_partial(**kwargs): - yield AnswerSchema(answer="4") - raise RuntimeError("stream broke") - - client = sync_client() - client.chat.completions.create_partial = create_partial - agent = make_agent(client) - - with pytest.raises(RuntimeError, match="stream broke"): - list(agent.run_stream(QuestionSchema(question="?"))) - - (span,) = spans_of(exporter) - assert span.status.status_code == StatusCode.ERROR - - -@pytest.mark.asyncio -async def test_run_async_emits_agent_span(exporter): - agent = make_agent(async_client()) - - response = await agent.run_async(QuestionSchema(question="?")) - - assert response.answer == "42" - (span,) = spans_of(exporter) - assert span.name == "AtomicAgent[QuestionSchema, AnswerSchema].run_async" - assert output_json(span) == {"answer": "42"} - assert span.status.status_code == StatusCode.OK - - -@pytest.mark.asyncio -async def test_run_async_stream_records_last_partial_and_nests_children(exporter, provider_tracer): - async def create_partial(**kwargs): - yield AnswerSchema(answer="4") - with provider_tracer.start_as_current_span("llm-chunk"): - yield AnswerSchema(answer="42") - - client = async_client() - client.chat.completions.create_partial = create_partial - agent = make_agent(client) - - partials = [partial.answer async for partial in agent.run_async_stream(QuestionSchema(question="?"))] - - assert partials == ["4", "42"] - chunk_span, agent_span = spans_of(exporter) - assert agent_span.name == "AtomicAgent[QuestionSchema, AnswerSchema].run_async_stream" - assert chunk_span.parent.span_id == agent_span.context.span_id - assert output_json(agent_span) == {"answer": "42"} - - -@pytest.mark.asyncio -async def test_run_async_stream_closed_early_ends_span(exporter): - agent = make_agent(async_client()) - - stream = agent.run_async_stream(QuestionSchema(question="?")) - first = await stream.__anext__() - await stream.aclose() - - assert first.answer == "4" - (span,) = spans_of(exporter) - assert span.attributes[tracing.ATTR_STREAM_CLOSED_EARLY] is True - assert span.status.status_code == StatusCode.OK - - -@pytest.mark.asyncio -async def test_concurrent_async_runs_keep_separate_spans(exporter): - first = make_agent(async_client(response=AnswerSchema(answer="one"))) - second = make_agent(async_client(response=AnswerSchema(answer="two"))) - - await asyncio.gather(first.run_async(QuestionSchema(question="1")), second.run_async(QuestionSchema(question="2"))) - - outputs = {input_json(span)["question"]: output_json(span)["answer"] for span in spans_of(exporter)} - assert outputs == {"1": "one", "2": "two"} - assert all(span.parent is None for span in spans_of(exporter)) - - -def test_tool_defined_before_instrumentation_emits_tool_span(exporter): - tool = CalculatorTool() - - result = tool.run(CalculationInputSchema(expression="6*7"), precision=1) - - assert result.result == "3.0" - (span,) = spans_of(exporter) - assert span.name == "CalculatorTool" - assert kind_of(span) == OpenInferenceSpanKindValues.TOOL.value - assert span.attributes[SpanAttributes.TOOL_NAME] == "CalculationInputSchema" - assert span.attributes[SpanAttributes.TOOL_DESCRIPTION] == "Expression to evaluate." - assert json.loads(span.attributes[SpanAttributes.TOOL_PARAMETERS])["title"] == "CalculationInputSchema" - assert input_json(span) == {"expression": "6*7"} - assert output_json(span) == {"result": "3.0"} - assert span.status.status_code == StatusCode.OK - - -def test_tool_accepts_keyword_params(exporter): - result = CalculatorTool().run(params=CalculationInputSchema(expression="1+1")) - - assert result.result == "3.00" - (span,) = spans_of(exporter) - assert input_json(span) == {"expression": "1+1"} - - -def test_tool_defined_after_instrumentation_emits_tool_span(exporter): - class LateTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="late") - - LateTool().run(CalculationInputSchema(expression="x")) - - (span,) = spans_of(exporter) - assert span.name == "LateTool" - assert LateTool().input_schema is CalculationInputSchema - - -def test_subclass_calling_super_run_emits_one_span(exporter): - class LoudCalculatorTool(CalculatorTool): - def run(self, params: CalculationInputSchema, precision: int = 2) -> CalculationOutputSchema: - return super().run(params, precision) - - LoudCalculatorTool().run(CalculationInputSchema(expression="2*2")) - - (span,) = spans_of(exporter) - assert span.name == "LoudCalculatorTool" - - -def test_sync_run_delegating_to_arun_emits_one_span(exporter): - AsyncCalculatorTool().run(CalculationInputSchema(expression="2*2")) - - (span,) = spans_of(exporter) - assert span.name == "AsyncCalculatorTool" - assert output_json(span) == {"result": "3"} - - -@pytest.mark.asyncio -async def test_run_async_delegating_to_thread_emits_one_span(exporter): - await ThreadedCalculatorTool().run_async(CalculationInputSchema(expression="3*3")) - - (span,) = spans_of(exporter) - assert span.name == "ThreadedCalculatorTool" - assert output_json(span) == {"result": "3"} - - -def test_mixin_defined_run_is_traced(exporter): - MixinCalculatorTool().run(CalculationInputSchema(expression="x")) - - (span,) = spans_of(exporter) - assert span.name == "MixinCalculatorTool" - - -def test_static_run_is_left_alone(exporter): - assert StaticRunTool.run(CalculationInputSchema(expression="x")).result == "static" - assert spans_of(exporter) == [] - - -def test_tool_failure_propagates_same_exception_and_marks_error(exporter): - with pytest.raises(ValueError, match="boom") as raised: - FailingTool().run(CalculationInputSchema(expression="x")) - - (span,) = spans_of(exporter) - assert span.status.status_code == StatusCode.ERROR - assert raised.value.args == ("boom",) - - -def test_unserializable_input_still_returns_result(exporter, caplog): - with caplog.at_level(logging.WARNING): - result = EchoTool().run(UnserializableInputSchema(store=AnswerStore())) - - assert result.result == "echo" - (span,) = spans_of(exporter) - assert span.status.status_code == StatusCode.OK - - -def test_agent_run_inside_tool_nests_under_tool_span(exporter): - agent = make_agent(sync_client()) - - class AgentTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result=agent.run(QuestionSchema(question=params.expression)).answer) - - AgentTool().run(CalculationInputSchema(expression="?")) - - agent_span, tool_span = spans_of(exporter) - assert agent_span.parent.span_id == tool_span.context.span_id - - -def test_hide_inputs_and_outputs_redacts_values_and_drops_mime_types(): - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - instrumentor = AtomicAgentsInstrumentor() - instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True, hide_outputs=True)) - try: - make_agent(sync_client()).run(QuestionSchema(question="secret")) - CalculatorTool().run(CalculationInputSchema(expression="1+1")) - finally: - instrumentor.uninstrument() - - for span in spans_of(exporter): - assert span.attributes[SpanAttributes.INPUT_VALUE] == "__REDACTED__" - assert span.attributes[SpanAttributes.OUTPUT_VALUE] == "__REDACTED__" - assert SpanAttributes.INPUT_MIME_TYPE not in span.attributes - assert SpanAttributes.OUTPUT_MIME_TYPE not in span.attributes - - -def test_session_and_user_metadata_survive_on_agent_span(exporter): - agent = make_agent(sync_client()) - - with using_session("session-1"), using_metadata({"tenant": "acme"}): - agent.run(QuestionSchema(question="?")) - - (span,) = spans_of(exporter) - assert span.attributes[SpanAttributes.SESSION_ID] == "session-1" - assert json.loads(span.attributes[SpanAttributes.METADATA]) == {"tenant": "acme"} - - -def test_second_instrument_with_other_settings_warns_and_keeps_first(exporter, caplog): - other_provider = TracerProvider() - - with caplog.at_level(logging.WARNING): - AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) - - assert "already installed" in caplog.text - make_agent(sync_client()).run(QuestionSchema(question="?")) - assert len(spans_of(exporter)) == 1 - - -def test_uninstrument_restores_originals_and_stops_tracing(): - original_run = AtomicAgent.__dict__["run"] - original_tool_run = CalculatorTool.__dict__["run"] - original_init_subclass = BaseTool.__dict__["__init_subclass__"] - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - instrumentor = AtomicAgentsInstrumentor() - - instrumentor.instrument(tracer_provider=provider) - assert instrumentor.is_instrumented_by_opentelemetry - instrumentor.uninstrument() - instrumentor.uninstrument() - - assert not instrumentor.is_instrumented_by_opentelemetry - assert AtomicAgent.__dict__["run"] is original_run - assert CalculatorTool.__dict__["run"] is original_tool_run - assert BaseTool.__dict__["__init_subclass__"] is original_init_subclass - assert "run" not in MixinCalculatorTool.__dict__ - - class AfterwardsTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="plain") - - assert AfterwardsTool().input_schema is CalculationInputSchema - make_agent(sync_client()).run(QuestionSchema(question="?")) - AfterwardsTool().run(CalculationInputSchema(expression="x")) - assert spans_of(exporter) == [] - - -def test_phoenix_entry_point_loads_and_instruments_with_only_a_tracer_provider(): - (entry_point,) = [point for point in entry_points(group="openinference_instrumentor") if point.name == "atomic_agents"] - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - - instrumentor = entry_point.load()() - instrumentor.instrument(tracer_provider=provider) - try: - make_agent(sync_client()).run(QuestionSchema(question="?")) - finally: - instrumentor.uninstrument() - - assert isinstance(instrumentor, AtomicAgentsInstrumentor) - assert len(spans_of(exporter)) == 1 - - -@pytest.mark.parametrize("options", [{}, {"skip_dep_check": False}]) -def test_instrument_without_openinference_raises_install_hint(monkeypatch: pytest.MonkeyPatch, options: dict[str, object]): - monkeypatch.setattr(observability, "_openinference_is_installed", lambda: False) - - with pytest.raises(ImportError, match="atomic-agents\\[observability\\]"): - AtomicAgentsInstrumentor().instrument(**options) - - -def test_uninstrument_without_openinference_is_a_no_op(monkeypatch): - monkeypatch.setattr(observability, "_openinference_is_installed", lambda: False) - - AtomicAgentsInstrumentor().uninstrument() - - assert not AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry - - -def test_openinference_below_the_floor_disables_tracing_with_a_warning(monkeypatch, caplog): - monkeypatch.setattr(observability, "version", lambda distribution: "0.1.20") - - with caplog.at_level(logging.WARNING): - AtomicAgentsInstrumentor().instrument() - - assert "0.1.38 or newer is required" in caplog.text - assert not AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry - assert AtomicAgent.__dict__["run"].__name__ == "run" - assert not hasattr(AtomicAgent.__dict__["run"], tracing.TRACED_MARKER) - - -class CountingProvider(BaseDynamicContextProvider): - def __init__(self): - super().__init__(title="Counter") - self.calls = 0 - - def get_info(self) -> str: - self.calls += 1 - return f"call {self.calls}" - - -def test_prompt_hash_identifies_static_sections_despite_dynamic_providers(exporter): - generator = SystemPromptGenerator(background=["v1"], context_providers={"counter": CountingProvider()}) - agent = AtomicAgent[QuestionSchema, AnswerSchema]( - AgentConfig(client=sync_client(), model="m", system_prompt_generator=generator) - ) - - agent.run(QuestionSchema(question="a")) - agent.run(QuestionSchema(question="b")) - - first, second = spans_of(exporter) - assert first.attributes[tracing.ATTR_SYSTEM_PROMPT_SHA256] == second.attributes[tracing.ATTR_SYSTEM_PROMPT_SHA256] - assert first.attributes[tracing.ATTR_CONTEXT_PROVIDERS] == ("counter",) - - -def test_prompt_hash_changes_when_static_sections_change(exporter): - agent = make_agent(sync_client()) - agent.run(QuestionSchema(question="a")) - agent.system_prompt_generator.background = ["v2"] - - agent.run(QuestionSchema(question="b")) - - first, second = spans_of(exporter) - assert first.attributes[tracing.ATTR_SYSTEM_PROMPT_SHA256] != second.attributes[tracing.ATTR_SYSTEM_PROMPT_SHA256] - - -def test_run_stream_dropped_without_close_still_ends_span(exporter): - agent = make_agent(sync_client()) - stream = agent.run_stream(QuestionSchema(question="?")) - next(stream) - - del stream - gc.collect() - - (span,) = spans_of(exporter) - assert span.attributes[tracing.ATTR_STREAM_CLOSED_EARLY] is True - - -@pytest.mark.asyncio -async def test_run_async_stream_cancelled_mid_await_is_not_an_error(exporter): - release = asyncio.Event() - - async def create_partial(**kwargs): - yield AnswerSchema(answer="4") - await release.wait() - yield AnswerSchema(answer="42") - - client = async_client() - client.chat.completions.create_partial = create_partial - agent = make_agent(client) - - async def consume(): - async for _ in agent.run_async_stream(QuestionSchema(question="?")): - pass - - task = asyncio.create_task(consume()) - await asyncio.sleep(0.01) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - - (span,) = spans_of(exporter) - assert span.attributes[tracing.ATTR_STREAM_CLOSED_EARLY] is True - assert output_json(span) == {"answer": "4"} - assert span.status.status_code != StatusCode.ERROR - - -def test_forge_style_thread_pool_tool_emits_one_span_in_the_callers_trace(exporter, provider_tracer): - class ForgeStyleTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - with ThreadPoolExecutor() as executor: - return executor.submit(contextvars.copy_context().run, asyncio.run, self.run_async(params)).result() - - async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="pooled") - - with provider_tracer.start_as_current_span("root") as root: - ForgeStyleTool().run(CalculationInputSchema(expression="x")) - - tool_span, root_span = spans_of(exporter) - assert tool_span.name == "ForgeStyleTool" - assert tool_span.parent.span_id == root.get_span_context().span_id - assert tool_span.context.trace_id == root_span.context.trace_id - - -def test_mcp_style_generated_tool_class_emits_one_span(exporter): - async def arun(self, params): - return CalculationOutputSchema(result="mcp") - - def run(self, params): - return asyncio.run(self.arun(params)) - - GeneratedTool = types.new_class( - "GeneratedTool", - (BaseTool[CalculationInputSchema, CalculationOutputSchema],), - {}, - lambda namespace: namespace.update({"run": run, "arun": arun}), - ) - - GeneratedTool().run(CalculationInputSchema(expression="x")) - - (span,) = spans_of(exporter) - assert span.name == "GeneratedTool" - assert output_json(span) == {"result": "mcp"} - - -def test_async_generator_tool_method_is_left_alone(exporter): - class StreamingTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="x") - - async def run_async(self, params: CalculationInputSchema): - yield CalculationOutputSchema(result="chunk") - - assert not hasattr(StreamingTool.__dict__["run_async"], tracing.TRACED_MARKER) - assert hasattr(StreamingTool.__dict__["run"], tracing.TRACED_MARKER) - - -def test_instrument_uninstrument_instrument_cycle_traces_tools_defined_between_cycles(): - exporter = InMemorySpanExporter() - provider = TracerProvider() - provider.add_span_processor(SimpleSpanProcessor(exporter)) - instrumentor = AtomicAgentsInstrumentor() - instrumentor.instrument(tracer_provider=provider) - instrumentor.uninstrument() - - class BetweenCyclesTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): - def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: - return CalculationOutputSchema(result="between") - - original_run = BetweenCyclesTool.__dict__["run"] - instrumentor.instrument(tracer_provider=provider) - try: - BetweenCyclesTool().run(CalculationInputSchema(expression="x")) - finally: - instrumentor.uninstrument() - - assert [span.name for span in spans_of(exporter)] == ["BetweenCyclesTool"] - assert BetweenCyclesTool.__dict__["run"] is original_run diff --git a/atomic-agents/tests/observability/test_opentelemetry_discovery.py b/atomic-agents/tests/observability/test_opentelemetry_discovery.py index 6c40f555..b201c1fe 100644 --- a/atomic-agents/tests/observability/test_opentelemetry_discovery.py +++ b/atomic-agents/tests/observability/test_opentelemetry_discovery.py @@ -20,7 +20,7 @@ import atomic_agents.observability as observability from atomic_agents.observability import AtomicAgentsInstrumentor -from test_openinference_tracing import CalculationInputSchema, CalculatorTool, QuestionSchema, make_agent, sync_client +from test_tracing import CalculationInputSchema, CalculatorTool, QuestionSchema, make_agent, sync_client @pytest.fixture @@ -72,6 +72,7 @@ def test_installed_opentelemetry_entry_point_activates_real_loader( def test_unknown_launcher_options_preserve_provider_and_privacy( local_tracing: tuple[TracerProvider, InMemorySpanExporter], + caplog: pytest.LogCaptureFixture, ) -> None: provider, exporter = local_tracing AtomicAgentsInstrumentor().instrument( @@ -82,13 +83,16 @@ def test_unknown_launcher_options_preserve_provider_and_privacy( ) run_agent_and_tool() + assert "custom_launcher_option" in caplog.text + assert any(record.levelno == logging.WARNING for record in caplog.records) + assert "skip_dep_check" not in caplog.text spans = exporter.get_finished_spans() assert len(spans) == 2 assert [exported_attributes(span)[SpanAttributes.INPUT_VALUE] for span in spans] == ["__REDACTED__", "__REDACTED__"] assert [exported_attributes(span)[SpanAttributes.OUTPUT_VALUE] for span in spans] == ["__REDACTED__", "__REDACTED__"] -def test_loader_skips_missing_openinference_and_continues_without_error( +def test_loader_skips_missing_opentelemetry_and_continues_without_error( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: instrumentor = installed_entry_point("opentelemetry_instrumentor") @@ -98,20 +102,23 @@ def test_loader_skips_missing_openinference_and_continues_without_error( sentinel.name = "sentinel" points = (instrumentor, sentinel) monkeypatch.setattr(_load, "entry_points", lambda group: points if group == "opentelemetry_instrumentor" else ()) - monkeypatch.setattr(observability, "_openinference_is_installed", lambda: False) + monkeypatch.setattr(observability, "_opentelemetry_is_installed", lambda: False) monkeypatch.delenv("OTEL_PYTHON_DISABLED_INSTRUMENTATIONS", raising=False) with caplog.at_level(logging.DEBUG): _load._load_instrumentors(DefaultDistro()) sentinel.load.return_value.return_value.instrument.assert_called_once_with(skip_dep_check=True) - assert "atomic-agents[observability]" in caplog.text + assert ( + 'Atomic Agents tracing needs opentelemetry-api. Install it with: pip install "atomic-agents[observability]" ' + '(add "atomic-agents[phoenix]" for Arize Phoenix / OpenInference attributes)' + ) in caplog.text assert not [record for record in caplog.records if record.levelno >= logging.ERROR] assert not AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry @pytest.mark.parametrize("first_group", ["openinference_instrumentor", "opentelemetry_instrumentor"]) -def test_both_entry_points_keep_first_provider_and_config( +def test_both_entry_points_replace_provider_and_config( first_group: str, local_tracing: tuple[TracerProvider, InMemorySpanExporter], caplog: pytest.LogCaptureFixture ) -> None: provider, exporter = local_tracing @@ -129,12 +136,11 @@ def test_both_entry_points_keep_first_provider_and_config( with caplog.at_level(logging.WARNING): second.instrument(tracer_provider=other_provider, config=TraceConfig(hide_inputs=False), skip_dep_check=True) run_agent_and_tool() - assert "already installed" in caplog.text - assert [exported_attributes(span)[SpanAttributes.INPUT_VALUE] for span in exporter.get_finished_spans()] == [ - "__REDACTED__", - "__REDACTED__", - ] - assert other_exporter.get_finished_spans() == () + assert "tracer provider/config replaced" in caplog.text + assert exporter.get_finished_spans() == () + spans = other_exporter.get_finished_spans() + assert len(spans) == 2 + assert all(exported_attributes(span)[SpanAttributes.INPUT_VALUE] != "__REDACTED__" for span in spans) finally: second.uninstrument() other_provider.shutdown() @@ -223,7 +229,7 @@ def test_normal_import_does_not_require_optional_observability_packages() -> Non sys.modules['opentelemetry'] = None import atomic_agents import atomic_agents.observability -assert 'atomic_agents.observability.openinference_tracing' not in sys.modules +assert 'atomic_agents.observability.tracing' not in sys.modules print('optional-free import passed') """ result = subprocess.run([sys.executable, "-c", script], capture_output=True, text=True, timeout=45) diff --git a/atomic-agents/tests/observability/test_tracing.py b/atomic-agents/tests/observability/test_tracing.py new file mode 100644 index 00000000..96d0c8e3 --- /dev/null +++ b/atomic-agents/tests/observability/test_tracing.py @@ -0,0 +1,1506 @@ +import abc +import asyncio +import contextvars +import gc +import json +import logging +import types +from concurrent.futures import ThreadPoolExecutor +from importlib.metadata import entry_points +from typing import Dict, List +from unittest.mock import Mock, PropertyMock + +import instructor +import pytest +from openinference.instrumentation import OITracer, TraceConfig, using_metadata, using_session +from openinference.semconv.trace import OpenInferenceSpanKindValues, SpanAttributes +from opentelemetry import trace as trace_api +from opentelemetry.sdk.trace import ReadableSpan, TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import StatusCode +from pydantic import Field, ValidationError + +from atomic_agents import AgentConfig, AtomicAgent, BaseIOSchema, BaseTool +from atomic_agents.context import BaseDynamicContextProvider, SystemPromptGenerator +from atomic_agents.observability import AtomicAgentsInstrumentor +import atomic_agents.observability as observability +from atomic_agents.observability import agent_spans, installation, span_attributes, span_payloads, tracing as tracing + + +class QuestionSchema(BaseIOSchema): + """A question for the agent.""" + + question: str = Field(...) + + +class AnswerSchema(BaseIOSchema): + """An answer from the agent.""" + + answer: str = Field(...) + + +class CalculationInputSchema(BaseIOSchema): + """Expression to evaluate.""" + + expression: str = Field(...) + + +class CalculationOutputSchema(BaseIOSchema): + """Result of the evaluation.""" + + result: str = Field(...) + + +class CalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema, precision: int = 2) -> CalculationOutputSchema: + return CalculationOutputSchema(result=f"{len(params.expression):.{precision}f}") + + +class AsyncCalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return asyncio.run(self.arun(params)) + + async def arun(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result=str(len(params.expression))) + + +class ThreadedCalculatorTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result=str(len(params.expression))) + + async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return await asyncio.to_thread(self.run, params) + + +class RunMixin: + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="mixin") + + +class MixinCalculatorTool(RunMixin, BaseTool[CalculationInputSchema, CalculationOutputSchema]): + pass + + +class StaticRunTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + @staticmethod + def run(params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="static") + + +class FailingTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + raise ValueError("boom") + + +class AnswerStore: + """Container object pydantic cannot serialize to JSON.""" + + def __init__(self): + self.values: List[int] = [] + + +class UnserializableInputSchema(BaseIOSchema): + """Input carrying an arbitrary object.""" + + model_config = {"arbitrary_types_allowed": True} + store: AnswerStore = Field(...) + + +class EchoTool(BaseTool[UnserializableInputSchema, CalculationOutputSchema]): + def run(self, params: UnserializableInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="echo") + + +@pytest.fixture +def exporter(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + exporter.provider = provider + instrumentor = AtomicAgentsInstrumentor() + instrumentor.instrument(tracer_provider=provider) + yield exporter + instrumentor.uninstrument() + + +@pytest.fixture +def provider_tracer(exporter): + return exporter.provider.get_tracer("test") + + +def sync_client(response=None, partials=None): + client = Mock(spec=instructor.Instructor) + client.chat = Mock() + client.chat.completions = Mock() + client.chat.completions.create = Mock(return_value=response or AnswerSchema(answer="42")) + + def create_partial(**kwargs): + for partial in partials or [AnswerSchema(answer="4"), AnswerSchema(answer="42")]: + yield partial + + client.chat.completions.create_partial = create_partial + return client + + +def async_client(response=None, partials=None): + client = Mock(spec=instructor.AsyncInstructor) + client.chat = Mock() + client.chat.completions = Mock() + + async def create(**kwargs): + return response or AnswerSchema(answer="42") + + async def create_partial(**kwargs): + for partial in partials or [AnswerSchema(answer="4"), AnswerSchema(answer="42")]: + yield partial + + client.chat.completions.create = create + client.chat.completions.create_partial = create_partial + return client + + +def make_agent(client) -> AtomicAgent[QuestionSchema, AnswerSchema]: + return AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=client, model="test-model")) + + +def spans_of(exporter) -> List[ReadableSpan]: + return list(exporter.get_finished_spans()) + + +def kind_of(span: ReadableSpan) -> str: + return span.attributes[SpanAttributes.OPENINFERENCE_SPAN_KIND] + + +def input_json(span: ReadableSpan) -> Dict: + return json.loads(span.attributes[SpanAttributes.INPUT_VALUE]) + + +def output_json(span: ReadableSpan) -> Dict: + return json.loads(span.attributes[SpanAttributes.OUTPUT_VALUE]) + + +def test_run_emits_agent_span_with_io_and_framework_attributes(exporter): + agent = make_agent(sync_client()) + + response = agent.run(QuestionSchema(question="6*7?")) + + assert response.answer == "42" + (span,) = spans_of(exporter) + assert span.name == "invoke_agent AtomicAgent" + assert kind_of(span) == OpenInferenceSpanKindValues.AGENT.value + assert span.attributes[SpanAttributes.AGENT_NAME] == "AtomicAgent" + assert span.attributes[span_attributes.ATTR_AGENT_CLASS] == "AtomicAgent[QuestionSchema, AnswerSchema]" + assert input_json(span) == {"question": "6*7?"} + assert output_json(span) == {"answer": "42"} + assert span.attributes[SpanAttributes.INPUT_MIME_TYPE] == "application/json" + assert span.attributes[SpanAttributes.OUTPUT_MIME_TYPE] == "application/json" + assert span.attributes[span_attributes.ATTR_MODEL] == "test-model" + assert span.attributes[span_attributes.ATTR_MODE] == agent.mode.value + assert span.attributes[span_attributes.ATTR_INPUT_SCHEMA] == "QuestionSchema" + assert span.attributes[span_attributes.ATTR_OUTPUT_SCHEMA] == "AnswerSchema" + assert span.attributes[span_attributes.ATTR_HISTORY_MESSAGE_COUNT] == 0 + assert span.attributes[span_attributes.ATTR_CONTEXT_PROVIDERS] == () + assert len(span.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256]) == 64 + assert span.status.status_code == StatusCode.OK + + +def test_second_turn_reports_history_size_entering_the_turn(exporter): + agent = make_agent(sync_client()) + agent.run(QuestionSchema(question="first")) + + agent.run(QuestionSchema(question="second")) + + first, second = spans_of(exporter) + assert first.attributes[span_attributes.ATTR_HISTORY_MESSAGE_COUNT] == 0 + assert second.attributes[span_attributes.ATTR_HISTORY_MESSAGE_COUNT] == 2 + + +def test_agent_without_system_role_has_no_prompt_hash(exporter): + agent = AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=sync_client(), model="m", system_role=None)) + + agent.run(QuestionSchema(question="?")) + + (span,) = spans_of(exporter) + assert span_attributes.ATTR_SYSTEM_PROMPT_SHA256 not in span.attributes + + +def test_abstract_intermediate_tool_base_is_not_wrapped(exporter): + class AbstractCalculator(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + @abc.abstractmethod + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + raise NotImplementedError + + assert not hasattr(AbstractCalculator.__dict__["run"], installation.TRACED_MARKER) + + +def test_run_without_input_sets_no_input_attribute(exporter): + agent = make_agent(sync_client()) + + agent.run() + + (span,) = spans_of(exporter) + assert SpanAttributes.INPUT_VALUE not in span.attributes + assert output_json(span) == {"answer": "42"} + + +def test_run_failure_propagates_and_marks_span_error(exporter): + client = sync_client() + client.chat.completions.create = Mock(side_effect=RuntimeError("provider down")) + agent = make_agent(client) + + with pytest.raises(RuntimeError, match="provider down"): + agent.run(QuestionSchema(question="?")) + + (span,) = spans_of(exporter) + assert span.status.status_code == StatusCode.ERROR + assert [event.name for event in span.events] == ["exception"] + + +def test_subclassed_agent_is_named_after_the_subclass_and_emits_one_span(exporter): + class ResearchAgent(AtomicAgent[QuestionSchema, AnswerSchema]): + def run(self, user_input=None): + return super().run(user_input) + + agent = ResearchAgent(AgentConfig(client=sync_client(), model="test-model")) + + agent.run(QuestionSchema(question="?")) + + (span,) = spans_of(exporter) + assert span.name == "invoke_agent ResearchAgent" + + +def test_agent_span_nests_under_a_user_chain_span(exporter, provider_tracer): + agent = make_agent(sync_client()) + + with provider_tracer.start_as_current_span("turn") as turn: + agent.run(QuestionSchema(question="?")) + + agent_span, turn_span = spans_of(exporter) + assert agent_span.parent.span_id == turn.get_span_context().span_id + assert turn_span.name == "turn" + + +def test_manual_setup_from_the_guide_groups_agents_and_tools_under_one_chain_span(): + exporter = InMemorySpanExporter() + tracer_provider = TracerProvider() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + trace_config = TraceConfig() + instrumentor = AtomicAgentsInstrumentor() + instrumentor.instrument(tracer_provider=tracer_provider, config=trace_config) + tracer = OITracer(tracer_provider.get_tracer(__name__), config=trace_config) + planner = make_agent(sync_client()) + answerer = make_agent(sync_client()) + calculator = CalculatorTool() + try: + with tracer.start_as_current_span("answer_question", openinference_span_kind="chain") as turn: + turn.set_input("what is 2 + 2?") + plan = planner.run(QuestionSchema(question="what is 2 + 2?")) + calculation = calculator.run(CalculationInputSchema(expression=plan.answer)) + answer = answerer.run(QuestionSchema(question=str(calculation.result))) + turn.set_output(answer.answer) + finally: + instrumentor.uninstrument() + + spans = spans_of(exporter) + chain_span = next(span for span in spans if kind_of(span) == OpenInferenceSpanKindValues.CHAIN.value) + children = [span for span in spans if span is not chain_span] + assert [kind_of(span) for span in children] == ["AGENT", "TOOL", "AGENT"] + assert all(span.parent.span_id == chain_span.get_span_context().span_id for span in children) + assert chain_span.attributes[SpanAttributes.INPUT_VALUE] == "what is 2 + 2?" + assert chain_span.attributes[SpanAttributes.OUTPUT_VALUE] == answer.answer + + +def test_llm_call_inside_run_becomes_child_of_agent_span(exporter, provider_tracer): + def create(**kwargs): + with provider_tracer.start_as_current_span("llm"): + return AnswerSchema(answer="42") + + client = sync_client() + client.chat.completions.create = create + agent = make_agent(client) + + agent.run(QuestionSchema(question="?")) + + llm_span, agent_span = spans_of(exporter) + assert llm_span.parent.span_id == agent_span.context.span_id + + +def test_run_stream_records_final_output_and_keeps_child_context_per_step(exporter, provider_tracer): + current_between_yields = [] + + def create_partial(**kwargs): + yield AnswerSchema(answer="4") + with provider_tracer.start_as_current_span("llm-chunk"): + yield AnswerSchema(answer="42") + + client = sync_client() + client.chat.completions.create_partial = create_partial + agent = make_agent(client) + + partials = [] + for partial in agent.run_stream(QuestionSchema(question="?")): + partials.append(partial.answer) + current_between_yields.append(trace_api.get_current_span()) + + assert partials == ["4", "42"] + assert all(span is trace_api.INVALID_SPAN for span in current_between_yields) + chunk_span, agent_span = spans_of(exporter) + assert agent_span.name == "invoke_agent AtomicAgent" + assert chunk_span.parent.span_id == agent_span.context.span_id + assert output_json(agent_span) == {"answer": "42"} + assert agent_span.status.status_code == StatusCode.OK + + +def test_run_stream_preserves_the_generator_return_value(exporter): + agent = make_agent(sync_client()) + + def consume(): + final = yield from agent.run_stream(QuestionSchema(question="?")) + return final + + generator = consume() + with pytest.raises(StopIteration) as stop: + while True: + next(generator) + + assert stop.value.value.answer == "42" + + +def test_run_stream_with_no_partials_ends_ok_without_output(exporter): + client = sync_client() + client.chat.completions.create_partial = lambda **kwargs: iter([]) + agent = make_agent(client) + + assert list(agent.run_stream(QuestionSchema(question="?"))) == [] + + (span,) = spans_of(exporter) + assert span.status.status_code == StatusCode.OK + assert SpanAttributes.OUTPUT_VALUE not in span.attributes + + +def test_run_stream_closed_early_ends_span_and_closes_inner_generator(exporter): + inner_closed = [] + + def create_partial(**kwargs): + try: + yield AnswerSchema(answer="4") + yield AnswerSchema(answer="42") + finally: + inner_closed.append(True) + + client = sync_client() + client.chat.completions.create_partial = create_partial + agent = make_agent(client) + + for partial in agent.run_stream(QuestionSchema(question="?")): + break + + assert inner_closed == [True] + (span,) = spans_of(exporter) + assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True + assert output_json(span) == {"answer": "4"} + assert span.status.status_code == StatusCode.OK + + +def test_run_stream_failure_mid_stream_marks_error(exporter): + def create_partial(**kwargs): + yield AnswerSchema(answer="4") + raise RuntimeError("stream broke") + + client = sync_client() + client.chat.completions.create_partial = create_partial + agent = make_agent(client) + + with pytest.raises(RuntimeError, match="stream broke"): + list(agent.run_stream(QuestionSchema(question="?"))) + + (span,) = spans_of(exporter) + assert span.status.status_code == StatusCode.ERROR + + +@pytest.mark.asyncio +async def test_run_async_emits_agent_span(exporter): + agent = make_agent(async_client()) + + response = await agent.run_async(QuestionSchema(question="?")) + + assert response.answer == "42" + (span,) = spans_of(exporter) + assert span.name == "invoke_agent AtomicAgent" + assert output_json(span) == {"answer": "42"} + assert span.status.status_code == StatusCode.OK + + +@pytest.mark.asyncio +async def test_run_async_stream_records_last_partial_and_nests_children(exporter, provider_tracer): + async def create_partial(**kwargs): + yield AnswerSchema(answer="4") + with provider_tracer.start_as_current_span("llm-chunk"): + yield AnswerSchema(answer="42") + + client = async_client() + client.chat.completions.create_partial = create_partial + agent = make_agent(client) + + partials = [partial.answer async for partial in agent.run_async_stream(QuestionSchema(question="?"))] + + assert partials == ["4", "42"] + chunk_span, agent_span = spans_of(exporter) + assert agent_span.name == "invoke_agent AtomicAgent" + assert chunk_span.parent.span_id == agent_span.context.span_id + assert output_json(agent_span) == {"answer": "42"} + + +@pytest.mark.asyncio +async def test_run_async_stream_closed_early_ends_span(exporter): + agent = make_agent(async_client()) + + stream = agent.run_async_stream(QuestionSchema(question="?")) + first = await stream.__anext__() + await stream.aclose() + + assert first.answer == "4" + (span,) = spans_of(exporter) + assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True + assert span.status.status_code == StatusCode.OK + + +@pytest.mark.asyncio +async def test_concurrent_async_runs_keep_separate_spans(exporter): + first = make_agent(async_client(response=AnswerSchema(answer="one"))) + second = make_agent(async_client(response=AnswerSchema(answer="two"))) + + await asyncio.gather(first.run_async(QuestionSchema(question="1")), second.run_async(QuestionSchema(question="2"))) + + outputs = {input_json(span)["question"]: output_json(span)["answer"] for span in spans_of(exporter)} + assert outputs == {"1": "one", "2": "two"} + assert all(span.parent is None for span in spans_of(exporter)) + + +def test_tool_defined_before_instrumentation_emits_tool_span(exporter): + tool = CalculatorTool() + + result = tool.run(CalculationInputSchema(expression="6*7"), precision=1) + + assert result.result == "3.0" + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + assert kind_of(span) == OpenInferenceSpanKindValues.TOOL.value + assert span.attributes[SpanAttributes.TOOL_NAME] == "CalculationInputSchema" + assert span.attributes[SpanAttributes.TOOL_DESCRIPTION] == "Expression to evaluate." + assert json.loads(span.attributes[SpanAttributes.TOOL_PARAMETERS])["title"] == "CalculationInputSchema" + assert input_json(span) == {"expression": "6*7"} + assert output_json(span) == {"result": "3.0"} + assert span.status.status_code == StatusCode.OK + + +def test_tool_accepts_keyword_params(exporter): + result = CalculatorTool().run(params=CalculationInputSchema(expression="1+1")) + + assert result.result == "3.00" + (span,) = spans_of(exporter) + assert input_json(span) == {"expression": "1+1"} + + +def test_tool_defined_after_instrumentation_emits_tool_span(exporter): + class LateTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="late") + + LateTool().run(CalculationInputSchema(expression="x")) + + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + assert LateTool().input_schema is CalculationInputSchema + + +def test_subclass_calling_super_run_emits_one_span(exporter): + class LoudCalculatorTool(CalculatorTool): + def run(self, params: CalculationInputSchema, precision: int = 2) -> CalculationOutputSchema: + return super().run(params, precision) + + LoudCalculatorTool().run(CalculationInputSchema(expression="2*2")) + + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + + +def test_sync_run_delegating_to_arun_emits_one_span(exporter): + AsyncCalculatorTool().run(CalculationInputSchema(expression="2*2")) + + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + assert output_json(span) == {"result": "3"} + + +@pytest.mark.asyncio +async def test_run_async_delegating_to_thread_emits_one_span(exporter): + await ThreadedCalculatorTool().run_async(CalculationInputSchema(expression="3*3")) + + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + assert output_json(span) == {"result": "3"} + + +def test_mixin_defined_run_is_traced(exporter): + MixinCalculatorTool().run(CalculationInputSchema(expression="x")) + + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + + +def test_static_run_is_left_alone(exporter): + assert StaticRunTool.run(CalculationInputSchema(expression="x")).result == "static" + assert spans_of(exporter) == [] + + +def test_tool_failure_propagates_same_exception_and_marks_error(exporter): + with pytest.raises(ValueError, match="boom") as raised: + FailingTool().run(CalculationInputSchema(expression="x")) + + (span,) = spans_of(exporter) + assert span.status.status_code == StatusCode.ERROR + assert raised.value.args == ("boom",) + assert span.attributes["error.type"] == "ValueError" + assert [event.name for event in span.events] == ["exception"] + + +@pytest.mark.asyncio +async def test_tool_run_async_cancellation_leaves_span_unset(exporter): + started = asyncio.Event() + release = asyncio.Event() + + class WaitingTool(CalculatorTool): + async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: + started.set() + await release.wait() + return CalculationOutputSchema(result="never") + + task = asyncio.create_task(WaitingTool().run_async(CalculationInputSchema(expression="x"))) + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + (span,) = spans_of(exporter) + assert span.attributes.get("atomic_agents.cancelled") is True + assert span.status.status_code == StatusCode.UNSET + assert span.events == () + assert "error.type" not in span.attributes + + +@pytest.mark.parametrize("error_type", [GeneratorExit, asyncio.CancelledError]) +@pytest.mark.parametrize("method_name", ["run", "run_async"]) +@pytest.mark.asyncio +async def test_tool_cancellation_propagates_without_error(exporter, monkeypatch, error_type, method_name): + def cancelled_result(*args, **kwargs): + raise error_type + + monkeypatch.setattr(CalculationOutputSchema, "__init__", cancelled_result) + tool = ThreadedCalculatorTool() + with pytest.raises(error_type): + result = getattr(tool, method_name)(CalculationInputSchema(expression="x")) + if method_name == "run_async": + await result + (span,) = spans_of(exporter) + assert span.attributes.get("atomic_agents.cancelled") is True + assert span.status.status_code == StatusCode.UNSET + assert span.events == () + assert "error.type" not in span.attributes + + +def test_unserializable_input_still_returns_result(exporter, caplog): + with caplog.at_level(logging.WARNING): + result = EchoTool().run(UnserializableInputSchema(store=AnswerStore())) + + assert result.result == "echo" + (span,) = spans_of(exporter) + assert span.status.status_code == StatusCode.OK + + +def test_agent_run_inside_tool_nests_under_tool_span(exporter): + agent = make_agent(sync_client()) + + class AgentTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result=agent.run(QuestionSchema(question=params.expression)).answer) + + AgentTool().run(CalculationInputSchema(expression="?")) + + agent_span, tool_span = spans_of(exporter) + assert agent_span.parent.span_id == tool_span.context.span_id + + +def test_hide_inputs_and_outputs_redacts_values_and_drops_mime_types(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = AtomicAgentsInstrumentor() + instrumentor.instrument(tracer_provider=provider, config=TraceConfig(hide_inputs=True, hide_outputs=True)) + try: + make_agent(sync_client()).run(QuestionSchema(question="secret")) + CalculatorTool().run(CalculationInputSchema(expression="1+1")) + finally: + instrumentor.uninstrument() + + for span in spans_of(exporter): + assert span.attributes[SpanAttributes.INPUT_VALUE] == "__REDACTED__" + assert span.attributes[SpanAttributes.OUTPUT_VALUE] == "__REDACTED__" + assert SpanAttributes.INPUT_MIME_TYPE not in span.attributes + assert SpanAttributes.OUTPUT_MIME_TYPE not in span.attributes + + +def test_session_and_user_metadata_survive_on_agent_span(exporter): + agent = make_agent(sync_client()) + + with using_session("session-1"), using_metadata({"tenant": "acme"}): + agent.run(QuestionSchema(question="?")) + + (span,) = spans_of(exporter) + assert span.attributes[SpanAttributes.SESSION_ID] == "session-1" + assert json.loads(span.attributes[SpanAttributes.METADATA]) == {"tenant": "acme"} + + +def test_second_instrument_replaces_provider_and_warns(exporter, caplog): + other_provider = TracerProvider() + other_exporter = InMemorySpanExporter() + other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) + agent = make_agent(sync_client()) + agent.run(QuestionSchema(question="first")) + installed = installation._installation + wrapped_run = AtomicAgent.run + + with caplog.at_level(logging.WARNING): + AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) + + agent.run(QuestionSchema(question="second")) + assert installation._installation is installed + assert AtomicAgent.run is wrapped_run + assert len(spans_of(exporter)) == 1 + assert len(spans_of(other_exporter)) == 1 + assert "tracer provider/config replaced" in caplog.text + assert "in-flight spans finish on the provider they started with" in caplog.text + assert caplog.records[-1].levelno == logging.WARNING + other_provider.shutdown() + + +def test_uninstrument_restores_originals_and_stops_tracing(): + original_run = AtomicAgent.__dict__["run"] + original_tool_run = CalculatorTool.__dict__["run"] + original_init_subclass = BaseTool.__dict__["__init_subclass__"] + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = AtomicAgentsInstrumentor() + + instrumentor.instrument(tracer_provider=provider) + assert instrumentor.is_instrumented_by_opentelemetry + instrumentor.uninstrument() + instrumentor.uninstrument() + + assert not instrumentor.is_instrumented_by_opentelemetry + assert AtomicAgent.__dict__["run"] is original_run + assert CalculatorTool.__dict__["run"] is original_tool_run + assert BaseTool.__dict__["__init_subclass__"] is original_init_subclass + assert "run" not in MixinCalculatorTool.__dict__ + + class AfterwardsTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="plain") + + assert AfterwardsTool().input_schema is CalculationInputSchema + make_agent(sync_client()).run(QuestionSchema(question="?")) + AfterwardsTool().run(CalculationInputSchema(expression="x")) + assert spans_of(exporter) == [] + + +def test_phoenix_entry_point_loads_and_instruments_with_only_a_tracer_provider(): + (entry_point,) = [point for point in entry_points(group="openinference_instrumentor") if point.name == "atomic_agents"] + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + + instrumentor = entry_point.load()() + instrumentor.instrument(tracer_provider=provider) + try: + make_agent(sync_client()).run(QuestionSchema(question="?")) + finally: + instrumentor.uninstrument() + + assert isinstance(instrumentor, AtomicAgentsInstrumentor) + assert len(spans_of(exporter)) == 1 + + +@pytest.mark.parametrize("options", [{}, {"skip_dep_check": False}]) +def test_instrument_without_opentelemetry_raises_install_hint(monkeypatch: pytest.MonkeyPatch, options: dict[str, object]): + monkeypatch.setattr(observability, "_opentelemetry_is_installed", lambda: False) + + with pytest.raises(ImportError, match="atomic-agents\\[observability\\]"): + AtomicAgentsInstrumentor().instrument(**options) + + +def test_uninstrument_without_opentelemetry_is_a_no_op(monkeypatch): + monkeypatch.setattr(observability, "_opentelemetry_is_installed", lambda: False) + + AtomicAgentsInstrumentor().uninstrument() + + assert not AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry + + +def test_openinference_below_the_floor_only_disables_context_with_a_warning(monkeypatch, caplog): + monkeypatch.setattr(observability, "version", lambda distribution: "0.1.20") + + with caplog.at_level(logging.WARNING): + AtomicAgentsInstrumentor().instrument() + + assert "0.1.38 or newer is required" in caplog.text + assert AtomicAgentsInstrumentor().is_instrumented_by_opentelemetry + assert span_attributes._context_attributes() == {} + AtomicAgentsInstrumentor().uninstrument() + + +class CountingProvider(BaseDynamicContextProvider): + def __init__(self): + super().__init__(title="Counter") + self.calls = 0 + + def get_info(self) -> str: + self.calls += 1 + return f"call {self.calls}" + + +def test_prompt_hash_identifies_static_sections_despite_dynamic_providers(exporter): + generator = SystemPromptGenerator(background=["v1"], context_providers={"counter": CountingProvider()}) + agent = AtomicAgent[QuestionSchema, AnswerSchema]( + AgentConfig(client=sync_client(), model="m", system_prompt_generator=generator) + ) + + agent.run(QuestionSchema(question="a")) + agent.run(QuestionSchema(question="b")) + + first, second = spans_of(exporter) + assert ( + first.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] + == second.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] + ) + assert first.attributes[span_attributes.ATTR_CONTEXT_PROVIDERS] == ("counter",) + + +def test_prompt_hash_changes_when_static_sections_change(exporter): + agent = make_agent(sync_client()) + agent.run(QuestionSchema(question="a")) + agent.system_prompt_generator.background = ["v2"] + + agent.run(QuestionSchema(question="b")) + + first, second = spans_of(exporter) + assert ( + first.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] + != second.attributes[span_attributes.ATTR_SYSTEM_PROMPT_SHA256] + ) + + +def test_run_stream_dropped_without_close_still_ends_span(exporter): + agent = make_agent(sync_client()) + stream = agent.run_stream(QuestionSchema(question="?")) + next(stream) + + del stream + gc.collect() + + (span,) = spans_of(exporter) + assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True + + +@pytest.mark.asyncio +async def test_run_async_stream_cancelled_mid_await_is_not_an_error(exporter): + release = asyncio.Event() + + async def create_partial(**kwargs): + yield AnswerSchema(answer="4") + await release.wait() + yield AnswerSchema(answer="42") + + client = async_client() + client.chat.completions.create_partial = create_partial + agent = make_agent(client) + + async def consume(): + async for _ in agent.run_async_stream(QuestionSchema(question="?")): + pass + + task = asyncio.create_task(consume()) + await asyncio.sleep(0.01) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + (span,) = spans_of(exporter) + assert span.attributes[span_attributes.ATTR_STREAM_CLOSED_EARLY] is True + assert output_json(span) == {"answer": "4"} + assert span.status.status_code != StatusCode.ERROR + + +def test_forge_style_thread_pool_tool_emits_one_span_in_the_callers_trace(exporter, provider_tracer): + class ForgeStyleTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + with ThreadPoolExecutor() as executor: + return executor.submit(contextvars.copy_context().run, asyncio.run, self.run_async(params)).result() + + async def run_async(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="pooled") + + with provider_tracer.start_as_current_span("root") as root: + ForgeStyleTool().run(CalculationInputSchema(expression="x")) + + tool_span, root_span = spans_of(exporter) + assert tool_span.name == "execute_tool CalculationInputSchema" + assert tool_span.parent.span_id == root.get_span_context().span_id + assert tool_span.context.trace_id == root_span.context.trace_id + + +def test_mcp_style_generated_tool_class_emits_one_span(exporter): + async def arun(self, params): + return CalculationOutputSchema(result="mcp") + + def run(self, params): + return asyncio.run(self.arun(params)) + + GeneratedTool = types.new_class( + "GeneratedTool", + (BaseTool[CalculationInputSchema, CalculationOutputSchema],), + {}, + lambda namespace: namespace.update({"run": run, "arun": arun}), + ) + + GeneratedTool().run(CalculationInputSchema(expression="x")) + + (span,) = spans_of(exporter) + assert span.name == "execute_tool CalculationInputSchema" + assert output_json(span) == {"result": "mcp"} + + +def test_async_generator_tool_method_is_left_alone(exporter): + class StreamingTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="x") + + async def run_async(self, params: CalculationInputSchema): + yield CalculationOutputSchema(result="chunk") + + assert not hasattr(StreamingTool.__dict__["run_async"], installation.TRACED_MARKER) + assert hasattr(StreamingTool.__dict__["run"], installation.TRACED_MARKER) + + +def test_instrument_uninstrument_instrument_cycle_traces_tools_defined_between_cycles(): + exporter = InMemorySpanExporter() + provider = TracerProvider() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + instrumentor = AtomicAgentsInstrumentor() + instrumentor.instrument(tracer_provider=provider) + instrumentor.uninstrument() + + class BetweenCyclesTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="between") + + original_run = BetweenCyclesTool.__dict__["run"] + instrumentor.instrument(tracer_provider=provider) + try: + BetweenCyclesTool().run(CalculationInputSchema(expression="x")) + finally: + instrumentor.uninstrument() + + assert [span.name for span in spans_of(exporter)] == ["execute_tool CalculationInputSchema"] + assert BetweenCyclesTool.__dict__["run"] is original_run + + +@pytest.fixture +def configured_tracing(): + provider = TracerProvider() + exporter = InMemorySpanExporter() + provider.add_span_processor(SimpleSpanProcessor(exporter)) + yield provider, exporter + AtomicAgentsInstrumentor().uninstrument() + provider.shutdown() + + +@pytest.mark.parametrize("openinference_enabled", [False, True]) +def test_genai_attributes_survive_without_openinference(monkeypatch, configured_tracing, openinference_enabled): + provider, exporter = configured_tracing + monkeypatch.setattr(span_attributes, "_import_context_attributes", lambda: None) + config = tracing.TracingConfig(openinference_attributes=True if openinference_enabled else None) + AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=config) + client = sync_client() + client.provider = types.SimpleNamespace(value="openai") + make_agent(client).run(QuestionSchema(question="question")) + CalculatorTool().run(CalculationInputSchema(expression="1+1")) + + agent_span, tool_span = spans_of(exporter) + assert [agent_span.name, tool_span.name] == ["invoke_agent AtomicAgent", "execute_tool CalculationInputSchema"] + assert agent_span.attributes["gen_ai.operation.name"] == "invoke_agent" + assert agent_span.attributes["gen_ai.agent.name"] == "AtomicAgent" + assert agent_span.attributes["gen_ai.request.model"] == "test-model" + assert agent_span.attributes["gen_ai.provider.name"] == "openai" + assert agent_span.attributes["atomic_agents.agent.class"] == "AtomicAgent[QuestionSchema, AnswerSchema]" + assert json.loads(agent_span.attributes["atomic_agents.input"]) == {"question": "question"} + assert json.loads(agent_span.attributes["atomic_agents.output"]) == {"answer": "42"} + assert tool_span.attributes["gen_ai.operation.name"] == "execute_tool" + assert tool_span.attributes["gen_ai.tool.name"] == "CalculationInputSchema" + assert tool_span.attributes["atomic_agents.tool.class"] == "CalculatorTool" + assert tool_span.attributes["gen_ai.tool.description"] == "Expression to evaluate." + assert tool_span.attributes["gen_ai.tool.type"] == "function" + assert json.loads(tool_span.attributes["gen_ai.tool.call.arguments"]) == {"expression": "1+1"} + assert json.loads(tool_span.attributes["gen_ai.tool.call.result"]) == {"result": "3.00"} + assert ("openinference.span.kind" in agent_span.attributes) is openinference_enabled + assert ("openinference.span.kind" in tool_span.attributes) is openinference_enabled + if openinference_enabled: + assert agent_span.attributes["input.value"] == agent_span.attributes["atomic_agents.input"] + assert tool_span.attributes["input.value"] == tool_span.attributes["gen_ai.tool.call.arguments"] + assert tool_span.attributes["output.value"] == tool_span.attributes["gen_ai.tool.call.result"] + else: + assert not ({"input.value", "output.value", "agent.name"} & agent_span.attributes.keys()) + assert not ({"input.value", "output.value", "tool.name", "tool.parameters"} & tool_span.attributes.keys()) + + +@pytest.mark.parametrize("config_class", [tracing.TracingConfig, TraceConfig]) +def test_both_visibility_configs_redact_primary_payloads(configured_tracing, config_class): + provider, exporter = configured_tracing + AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=config_class(hide_inputs=True, hide_outputs=True)) + make_agent(sync_client()).run(QuestionSchema(question="secret")) + CalculatorTool().run(CalculationInputSchema(expression="secret")) + + agent_span, tool_span = spans_of(exporter) + assert agent_span.attributes["atomic_agents.input"] == "__REDACTED__" + assert agent_span.attributes["atomic_agents.output"] == "__REDACTED__" + assert tool_span.attributes["gen_ai.tool.call.arguments"] == "__REDACTED__" + assert tool_span.attributes["gen_ai.tool.call.result"] == "__REDACTED__" + + +@pytest.mark.parametrize("hide_outputs", [False, True]) +def test_large_tool_result_is_truncated_unless_hidden(configured_tracing, hide_outputs): + provider, exporter = configured_tracing + AtomicAgentsInstrumentor().instrument( + tracer_provider=provider, config=tracing.TracingConfig(hide_outputs=hide_outputs, openinference_attributes=True) + ) + + class LargeResultTool(BaseTool[CalculationInputSchema, CalculationOutputSchema]): + def run(self, params: CalculationInputSchema) -> CalculationOutputSchema: + return CalculationOutputSchema(result="x" * 10240) + + result = LargeResultTool().run(CalculationInputSchema(expression="large")) + (span,) = spans_of(exporter) + serialized = result.model_dump_json() + expected = "__REDACTED__" if hide_outputs else serialized[: 4096 - len("...[truncated]")] + "...[truncated]" + assert span.attributes["gen_ai.tool.call.result"] == expected + assert span.attributes["output.value"] == expected + if hide_outputs: + assert "gen_ai.tool.call.result.original_length" not in span.attributes + assert "output.value.original_length" not in span.attributes + else: + assert span.attributes["gen_ai.tool.call.result.original_length"] == len(serialized) + assert span.attributes["output.value.original_length"] == len(serialized) + + +@pytest.mark.parametrize("agent_method", ["run", "run_async", "run_stream", "run_async_stream"]) +@pytest.mark.asyncio +async def test_suppression_skips_agent_and_tool_spans(exporter, agent_method): + from opentelemetry import context + + client = async_client() if "async" in agent_method else sync_client() + agent = make_agent(client) + token = context.attach(context.set_value("suppress_instrumentation", True)) + try: + result = getattr(agent, agent_method)(QuestionSchema(question="suppressed")) + assert await final_agent_answer(agent_method, result) == "42" + assert CalculatorTool().run(CalculationInputSchema(expression="1+1")).result == "3.00" + assert (await ThreadedCalculatorTool().run_async(CalculationInputSchema(expression="1+1"))).result == "3" + finally: + context.detach(token) + assert spans_of(exporter) == [] + + +def test_failures_record_qualified_error_type(exporter): + client = sync_client() + client.chat.completions.create.side_effect = ValueError("agent failure") + with pytest.raises(ValueError, match="agent failure"): + make_agent(client).run(QuestionSchema(question="fail")) + with pytest.raises(ValueError, match="boom"): + FailingTool().run(CalculationInputSchema(expression="fail")) + for span in spans_of(exporter): + assert span.attributes["error.type"] == "ValueError" + assert span.status.status_code == StatusCode.ERROR + assert len(span.events) == 1 + + +def test_tracing_config_environment_and_coercion(monkeypatch): + monkeypatch.setenv("OPENINFERENCE_HIDE_INPUTS", "TrUe") + monkeypatch.setenv("OPENINFERENCE_HIDE_OUTPUTS", "1") + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", "128") + from atomic_agents.observability import TracingConfig + + config = TracingConfig.from_environment() + assert config == TracingConfig(hide_inputs=True, hide_outputs=True, max_attribute_length=128) + assert TracingConfig.coerce(None) == config + assert TracingConfig.coerce(config) is config + assert TracingConfig.coerce(TraceConfig(hide_inputs=True)) == TracingConfig(hide_inputs=True, max_attribute_length=128) + + +@pytest.mark.parametrize( + "value, expected", [("false", False), ("FaLsE", False), ("0", False), ("true", True), ("1", True), ("", True)] +) +def test_tracing_config_llm_spans_environment(monkeypatch, value, expected): + from atomic_agents.observability import TracingConfig + + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_LLM_SPANS", value) + assert TracingConfig.from_environment().llm_spans is expected + assert TracingConfig.coerce(None).llm_spans is expected + + +def test_payload_visibility_preserves_environment_settings(monkeypatch): + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_LLM_SPANS", "false") + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", "128") + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS", "true") + config = tracing.TracingConfig.coerce(TraceConfig(hide_inputs=True, hide_outputs=False)) + assert config.llm_spans is False + assert config.max_attribute_length == 128 + assert config.hide_inputs is True + assert config.hide_outputs is False + + +def test_tracing_config_llm_spans_defaults_to_enabled(monkeypatch): + from atomic_agents.observability import TracingConfig + + monkeypatch.delenv("ATOMIC_AGENTS_TRACE_LLM_SPANS", raising=False) + assert TracingConfig.from_environment().llm_spans is True + assert TracingConfig().llm_spans is True + + +def test_session_is_mirrored_to_genai_conversation(exporter): + with using_session("session-1"): + make_agent(sync_client()).run(QuestionSchema(question="?")) + CalculatorTool().run(CalculationInputSchema(expression="?")) + for span in spans_of(exporter): + assert span.attributes["gen_ai.conversation.id"] == "session-1" + + +async def final_agent_answer(agent_method, result): + if agent_method == "run_async": + return (await result).answer + elif agent_method == "run_stream": + return list(result)[-1].answer + elif agent_method == "run_async_stream": + return [partial async for partial in result][-1].answer + else: + return result.answer + + +@pytest.mark.parametrize( + "name, expected", [(None, "AtomicAgent"), ("", "AtomicAgent"), ("orchestrator_planning", "orchestrator_planning")] +) +def test_stable_agent_name_is_separate_from_schema_class(exporter, name, expected): + agent = AtomicAgent[QuestionSchema, AnswerSchema](AgentConfig(client=sync_client(), name=name)) + agent.run(QuestionSchema(question="?")) + (span,) = spans_of(exporter) + assert span.name == f"invoke_agent {expected}" + assert span.attributes["gen_ai.agent.name"] == expected + assert span.attributes["agent.name"] == expected + assert span.attributes["atomic_agents.agent.class"] == "AtomicAgent[QuestionSchema, AnswerSchema]" + + +@pytest.mark.parametrize("openinference_enabled", [False, True]) +def test_conversation_context_overrides_session_on_agent_and_tool(exporter, openinference_enabled): + config = tracing.TracingConfig(openinference_attributes=openinference_enabled) + AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) + with using_session("session-1"), observability.using_conversation("c-1"): + make_agent(sync_client()).run(QuestionSchema(question="?")) + CalculatorTool().run(CalculationInputSchema(expression="?")) + agent_span, tool_span = spans_of(exporter) + assert agent_span.attributes["gen_ai.conversation.id"] == "c-1" + assert tool_span.attributes["gen_ai.conversation.id"] == "c-1" + assert agent_span.attributes.get("session.id") == ("c-1" if openinference_enabled else None) + assert tool_span.attributes.get("session.id") == ("c-1" if openinference_enabled else None) + + +def test_session_is_mirrored_even_with_openinference_attributes_disabled(exporter): + config = tracing.TracingConfig(openinference_attributes=False) + AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) + with using_session("session-1"): + make_agent(sync_client()).run(QuestionSchema(question="?")) + CalculatorTool().run(CalculationInputSchema(expression="?")) + for span in spans_of(exporter): + assert span.attributes["gen_ai.conversation.id"] == "session-1" + assert "session.id" not in span.attributes + + +def test_conversation_context_restores_after_nested_exception(exporter): + with observability.using_conversation("outer"): + with pytest.raises(ValueError), observability.using_conversation("inner"): + CalculatorTool().run(CalculationInputSchema(expression="inner")) + raise ValueError("reset context") + CalculatorTool().run(CalculationInputSchema(expression="outer")) + CalculatorTool().run(CalculationInputSchema(expression="unset")) + inner, outer, unset = spans_of(exporter) + assert inner.attributes["gen_ai.conversation.id"] == "inner" + assert outer.attributes["gen_ai.conversation.id"] == "outer" + assert "gen_ai.conversation.id" not in unset.attributes + + +@pytest.mark.asyncio +async def test_conversation_context_is_task_local(exporter): + async def run_in_conversation(conversation_id): + with observability.using_conversation(conversation_id): + await asyncio.sleep(0) + await make_agent(async_client()).run_async(QuestionSchema(question="?")) + CalculatorTool().run(CalculationInputSchema(expression="?")) + + await asyncio.gather(run_in_conversation("first"), run_in_conversation("second")) + assert [span.attributes["gen_ai.conversation.id"] for span in spans_of(exporter)] == ["first", "first", "second", "second"] + + +def test_plain_parent_span_owns_agent_and_tool_trace(exporter): + with exporter.provider.get_tracer("app").start_as_current_span("turn") as turn: + make_agent(sync_client()).run(QuestionSchema(question="?")) + CalculatorTool().run(CalculationInputSchema(expression="?")) + agent_span, tool_span, parent = spans_of(exporter) + assert agent_span.context.trace_id == turn.get_span_context().trace_id + assert tool_span.context.trace_id == turn.get_span_context().trace_id + assert agent_span.parent.span_id == parent.context.span_id + assert tool_span.parent.span_id == parent.context.span_id + + +def test_same_provider_and_config_is_a_silent_no_op(exporter, caplog): + installed = installation._installation + tracer = installed.tracer + with caplog.at_level(logging.DEBUG, logger=tracing.__name__): + AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=tracing.TracingConfig()) + assert installation._installation is installed + assert installed.tracer is tracer + assert not caplog.records + + +def test_reinstrument_replaces_config_without_rewrapping(exporter, caplog): + wrapped_run = AtomicAgent.run + agent = make_agent(sync_client()) + config = tracing.TracingConfig(hide_inputs=True, openinference_attributes=False) + with caplog.at_level(logging.WARNING): + AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) + agent.run(QuestionSchema(question="secret")) + CalculatorTool().run(CalculationInputSchema(expression="secret")) + agent_span, tool_span = spans_of(exporter) + assert agent_span.attributes["atomic_agents.input"] == "__REDACTED__" + assert tool_span.attributes["gen_ai.tool.call.arguments"] == "__REDACTED__" + assert "agent.name" not in agent_span.attributes + assert AtomicAgent.run is wrapped_run + assert "tracer provider/config replaced" in caplog.text + + +def test_in_flight_span_finishes_on_original_provider(exporter): + other_provider = TracerProvider() + other_exporter = InMemorySpanExporter() + other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) + client = sync_client() + + def replace_provider(**kwargs): + AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) + CalculatorTool().run(CalculationInputSchema(expression="?")) + return AnswerSchema(answer="42") + + client.chat.completions.create.side_effect = replace_provider + make_agent(client).run(QuestionSchema(question="?")) + assert len(spans_of(exporter)) == 1 + assert len(spans_of(other_exporter)) == 1 + (agent_span,) = spans_of(exporter) + (tool_span,) = spans_of(other_exporter) + assert agent_span.name == "invoke_agent AtomicAgent" + assert tool_span.parent.span_id == agent_span.context.span_id + assert tool_span.context.trace_id == agent_span.context.trace_id + other_provider.shutdown() + + +@pytest.mark.parametrize("name", ["model", "mode"]) +def test_none_descriptions_are_omitted_without_otel_warnings(exporter, caplog, name): + agent = make_agent(sync_client()) + setattr(agent, name, None) + with agent_spans._agent_span(agent, (), {}): + pass + (span,) = spans_of(exporter) + assert f"atomic_agents.{name}" not in span.attributes + assert "gen_ai.request.model" not in span.attributes if name == "model" else True + assert "Invalid type" not in caplog.text + + +def test_broken_history_preserves_mandatory_agent_attributes(exporter, monkeypatch, caplog): + agent = make_agent(sync_client()) + monkeypatch.setattr(type(agent), "history", PropertyMock(side_effect=RuntimeError("history unavailable")), raising=False) + with caplog.at_level(logging.DEBUG), agent_spans._agent_span(agent, (), {}): + pass + (span,) = spans_of(exporter) + assert span.attributes["gen_ai.operation.name"] == "invoke_agent" + assert span.attributes["gen_ai.agent.name"] == "AtomicAgent" + assert span.attributes["atomic_agents.agent.class"] == "AtomicAgent[QuestionSchema, AnswerSchema]" + assert span.attributes["gen_ai.request.model"] == "test-model" + assert "atomic_agents.history.message_count" not in span.attributes + assert sum("history.message_count" in record.message for record in caplog.records) == 1 + + +def test_broken_tool_description_preserves_other_attributes(exporter, monkeypatch): + monkeypatch.setattr(CalculatorTool, "tool_description", PropertyMock(side_effect=RuntimeError("unavailable"))) + CalculatorTool().run(CalculationInputSchema(expression="1+1")) + (span,) = spans_of(exporter) + assert span.attributes["gen_ai.tool.name"] == "CalculationInputSchema" + assert span.attributes["atomic_agents.tool.class"] == "CalculatorTool" + assert span.attributes["tool.name"] == "CalculationInputSchema" + assert "tool.parameters" in span.attributes + assert "gen_ai.tool.description" not in span.attributes + + +def test_configured_tool_name_is_the_span_join_key(exporter): + tool = CalculatorTool() + tool.config.title = "calculate" + tool.run(CalculationInputSchema(expression="1+1")) + (span,) = spans_of(exporter) + assert span.name == "execute_tool calculate" + assert span.attributes["gen_ai.tool.name"] == span.attributes["tool.name"] == "calculate" + assert span.attributes["atomic_agents.tool.class"] == "CalculatorTool" + + +@pytest.mark.parametrize("hide", [False, True]) +def test_hidden_payload_is_never_serialized(exporter, caplog, hide): + payload = UnserializableInputSchema(store=AnswerStore()) + with exporter.provider.get_tracer("test").start_as_current_span("payload") as span: + span_payloads._record_payload(span, (("payload", "payload.mime_type"),), payload, hide, tracing.TracingConfig()) + (span,) = spans_of(exporter) + assert ("Could not serialize tracing payload" in caplog.text) is not hide + assert span.attributes.get("payload") == ("__REDACTED__" if hide else None) + assert "payload.mime_type" not in span.attributes + + +@pytest.mark.parametrize("hide", [False, True]) +def test_hidden_non_model_payload_does_not_call_json_dumps(exporter, monkeypatch, caplog, hide): + serialize = Mock(side_effect=ValueError("serialization failed")) + monkeypatch.setattr(span_payloads.json, "dumps", serialize) + with exporter.provider.get_tracer("test").start_as_current_span("payload") as span: + span_payloads._record_payload(span, (("payload", None),), object(), hide, tracing.TracingConfig()) + assert serialize.call_count == (0 if hide else 1) + assert ("Could not serialize tracing payload" in caplog.text) is not hide + + +@pytest.mark.parametrize( + "provider,expected", + [ + ("azure_openai", "azure.ai.openai"), + ("vertexai", "gcp.vertex_ai"), + ("gemini", "gcp.gemini"), + ("genai", "gcp.gemini"), + ("bedrock", "aws.bedrock"), + ("mistral", "mistral_ai"), + ("xai", "x_ai"), + ("openai", "openai"), + ("anthropic", "anthropic"), + ("cohere", "cohere"), + ("groq", "groq"), + ("deepseek", "deepseek"), + ("perplexity", "perplexity"), + ("custom", "custom"), + ], +) +def test_instructor_provider_names_follow_genai_conventions(exporter, provider, expected): + client = sync_client() + client.provider = types.SimpleNamespace(value=provider) + make_agent(client).run(QuestionSchema(question="question")) + (span,) = spans_of(exporter) + assert span.attributes["gen_ai.provider.name"] == expected + + +def test_none_context_attributes_are_filtered(exporter, monkeypatch, caplog): + monkeypatch.setattr(span_attributes, "_context_attributes", lambda: {"optional": None}) + make_agent(sync_client()).run(QuestionSchema(question="question")) + (span,) = spans_of(exporter) + assert "optional" not in span.attributes + assert "gen_ai.provider.name" not in span.attributes + assert "Invalid type" not in caplog.text + + +@pytest.mark.parametrize("prefix", ["ATOMIC_AGENTS_TRACE", "OPENINFERENCE"]) +def test_visibility_environment_supports_both_spellings(monkeypatch, prefix): + monkeypatch.delenv("ATOMIC_AGENTS_TRACE_HIDE_INPUTS", raising=False) + monkeypatch.delenv("ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS", raising=False) + monkeypatch.setenv(f"{prefix}_HIDE_INPUTS", "TRUE") + monkeypatch.setenv(f"{prefix}_HIDE_OUTPUTS", "1") + config = tracing.TracingConfig.from_environment() + assert config.hide_inputs is True + assert config.hide_outputs is True + + +def test_atomic_visibility_environment_overrides_openinference(monkeypatch): + monkeypatch.setenv("OPENINFERENCE_HIDE_INPUTS", "true") + monkeypatch.setenv("OPENINFERENCE_HIDE_OUTPUTS", "true") + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_HIDE_INPUTS", "false") + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS", "") + config = tracing.TracingConfig.from_environment() + assert config.hide_inputs is False + assert config.hide_outputs is False + + +def test_error_status_and_recorded_message_are_bounded(configured_tracing): + provider, exporter = configured_tracing + AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=tracing.TracingConfig(max_attribute_length=12)) + error = ValueError("private" * 100) + agent = make_agent(sync_client()) + with pytest.raises(ValueError) as raised, agent_spans._agent_span(agent, (), {}): + raise error + (span,) = spans_of(exporter) + expected = "...[truncate" + assert raised.value is error + assert span.status.description == expected + assert span.events[0].attributes["exception.message"] == expected + assert len(span.events[0].attributes["exception.stacktrace"]) <= 12 + + +def test_hidden_validation_errors_do_not_export_model_output(configured_tracing): + provider, exporter = configured_tracing + AtomicAgentsInstrumentor().instrument(tracer_provider=provider, config=tracing.TracingConfig(hide_outputs=True)) + try: + QuestionSchema.model_validate({"question": ["secret model output"]}) + except ValidationError as error: + with provider.get_tracer("test").start_as_current_span("failure") as span: + span_payloads._record_failure(span, error) + (span,) = spans_of(exporter) + assert span.status.description == "ValidationError: 1 errors (__REDACTED__)" + assert span.events[0].attributes["exception.message"] == span.status.description + assert "secret model output" not in str(span.events[0].attributes) + + +def test_tool_parameters_use_the_payload_size_limit(configured_tracing): + provider, exporter = configured_tracing + AtomicAgentsInstrumentor().instrument( + tracer_provider=provider, config=tracing.TracingConfig(max_attribute_length=12, openinference_attributes=True) + ) + CalculatorTool().run(CalculationInputSchema(expression="x")) + (span,) = spans_of(exporter) + parameters = json.dumps(CalculationInputSchema.model_json_schema()) + assert span.attributes["tool.parameters"] == "...[truncate" + assert span.attributes["tool.parameters.original_length"] == len(parameters) + + +@pytest.mark.parametrize("value, expected", [("", 4096), ("4k", 4096), ("0", 4096), ("-5", 4096), ("512", 512)]) +@pytest.mark.parametrize("visibility", [None, TraceConfig(hide_inputs=True)]) +def test_attribute_length_environment_is_tolerant(monkeypatch, caplog, value, expected, visibility): + monkeypatch.setenv("ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH", value) + with caplog.at_level(logging.WARNING): + config = tracing.TracingConfig.coerce(visibility) + assert config.max_attribute_length == expected + assert config.hide_inputs is (visibility is not None) + warnings = [record for record in caplog.records if record.levelno == logging.WARNING] + if expected == 4096: + assert len(warnings) == 1 + assert repr(value) in warnings[0].message + else: + assert warnings == [] + + +def test_reinstall_without_config_preserves_redaction(exporter, caplog): + config = tracing.TracingConfig(hide_inputs=True, openinference_attributes=False) + AtomicAgentsInstrumentor().instrument(tracer_provider=exporter.provider, config=config) + other_provider = TracerProvider() + other_exporter = InMemorySpanExporter() + other_provider.add_span_processor(SimpleSpanProcessor(other_exporter)) + caplog.clear() + try: + with caplog.at_level(logging.WARNING): + AtomicAgentsInstrumentor().instrument(tracer_provider=other_provider) + assert installation._current_installation().config is config + assert installation._current_installation().tracer_provider is other_provider + make_agent(sync_client()).run(QuestionSchema(question="secret")) + assert spans_of(other_exporter)[0].attributes["atomic_agents.input"] == "__REDACTED__" + assert len(caplog.records) == 1 + assert "replaced" in caplog.records[0].message + finally: + other_provider.shutdown() + + +def test_reinstall_without_provider_preserves_provider(exporter, caplog): + config = tracing.TracingConfig(hide_outputs=True) + with caplog.at_level(logging.WARNING): + AtomicAgentsInstrumentor().instrument(config=config) + assert installation._current_installation().tracer_provider is exporter.provider + caplog.clear() + AtomicAgentsInstrumentor().instrument() + assert caplog.records == [] + make_agent(sync_client()).run(QuestionSchema(question="?")) + assert spans_of(exporter)[0].attributes["atomic_agents.output"] == "__REDACTED__" + + +@pytest.mark.asyncio +async def test_run_async_cancellation_leaves_span_unset(exporter): + started = asyncio.Event() + release = asyncio.Event() + + async def create(**kwargs): + started.set() + await release.wait() + return AnswerSchema(answer="never") + + client = async_client() + client.chat.completions.create = create + task = asyncio.create_task(make_agent(client).run_async(QuestionSchema(question="?"))) + await asyncio.wait_for(started.wait(), timeout=2) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + (span,) = spans_of(exporter) + assert span.attributes.get("atomic_agents.cancelled") is True + assert span.status.status_code == StatusCode.UNSET + assert span.events == () + assert "error.type" not in span.attributes + + +@pytest.mark.parametrize("error_type", [GeneratorExit, asyncio.CancelledError]) +def test_run_cancellation_leaves_span_unset(exporter, error_type): + client = sync_client() + client.chat.completions.create.side_effect = error_type + with pytest.raises(error_type): + make_agent(client).run(QuestionSchema(question="?")) + (span,) = spans_of(exporter) + assert span.attributes.get("atomic_agents.cancelled") is True + assert span.status.status_code == StatusCode.UNSET + assert span.events == () + + +def test_run_stream_cancellation_marks_closed_early(exporter): + def create_partial(**kwargs): + yield AnswerSchema(answer="4") + raise asyncio.CancelledError + + client = sync_client() + client.chat.completions.create_partial = create_partial + stream = make_agent(client).run_stream(QuestionSchema(question="?")) + next(stream) + with pytest.raises(asyncio.CancelledError): + next(stream) + (span,) = spans_of(exporter) + assert span.attributes.get(span_attributes.ATTR_STREAM_CLOSED_EARLY) is True + assert span.status.status_code == StatusCode.UNSET + assert span.events == () + + +@pytest.mark.parametrize("broken_name", [False, True]) +def test_agent_name_falls_back_without_initialized_base(exporter, broken_name): + class UninitializedAgent(AtomicAgent[QuestionSchema, AnswerSchema]): + def __init__(self): + self.client = sync_client() + + if broken_name: + UninitializedAgent.name = property(lambda self: 1 / 0) + agent = UninitializedAgent() + try: + display_name = span_attributes._agent_display_name(agent) + except (AttributeError, ZeroDivisionError): + display_name = None + assert display_name == "UninitializedAgent" + with agent_spans._agent_span(agent, (), {}): + pass + (span,) = spans_of(exporter) + assert span.name == "invoke_agent UninitializedAgent" + assert span.attributes["gen_ai.agent.name"] == "UninitializedAgent" diff --git a/atomic-examples/phoenix-tracing/README.md b/atomic-examples/phoenix-tracing/README.md index 077801da..24bc3da5 100644 --- a/atomic-examples/phoenix-tracing/README.md +++ b/atomic-examples/phoenix-tracing/README.md @@ -6,17 +6,30 @@ produces this tree in Phoenix: ``` CHAIN answer_question -├── AGENT AtomicAgent[BasicChatInputSchema, CalculationPlanSchema].run +├── AGENT invoke_agent {agent name} +│ ├── LLM chat {model} (Atomic Agents attempt span, one per Instructor attempt) │ └── LLM ChatCompletion (from openinference-instrumentation-openai) -├── TOOL CalculatorTool -└── AGENT AtomicAgent[CalculationResultSchema, BasicChatOutputSchema].run +├── TOOL execute_tool {tool name} +└── AGENT invoke_agent {agent name} + ├── LLM chat {model} └── LLM ChatCompletion ``` +Running the OpenAI instrumentor next to the Atomic Agents attempt spans gives two LLM layers. To keep one, pass +`TracingConfig(llm_spans=False)` when you call `AtomicAgentsInstrumentor().instrument(...)` yourself. + The AGENT and TOOL spans come from `atomic_agents.observability`, which Phoenix activates automatically through `register(auto_instrument=True)`. The LLM spans come from the OpenAI instrumentor. The CHAIN span is the only manual tracing in `main.py`; without it every agent and tool call would be its own trace. +## Install + +Inside this repository `uv sync` installs everything. In your own project: + +```bash +pip install "atomic-agents[phoenix]" arize-phoenix-otel openinference-instrumentation-openai +``` + ## Run it 1. Start Phoenix locally (or point `PHOENIX_COLLECTOR_ENDPOINT` at a hosted instance): diff --git a/atomic-examples/phoenix-tracing/phoenix_tracing/main.py b/atomic-examples/phoenix-tracing/phoenix_tracing/main.py index ae1f9a3d..13d119c8 100644 --- a/atomic-examples/phoenix-tracing/phoenix_tracing/main.py +++ b/atomic-examples/phoenix-tracing/phoenix_tracing/main.py @@ -15,6 +15,7 @@ load_dotenv() +os.environ.setdefault("ATOMIC_AGENTS_TRACE_LLM_SPANS", "false") tracer_provider = register(project_name="atomic-agents-example", auto_instrument=True) tracer = tracer_provider.get_tracer(__name__) console = Console() diff --git a/atomic-examples/phoenix-tracing/pyproject.toml b/atomic-examples/phoenix-tracing/pyproject.toml index c626725f..e09a5bf6 100644 --- a/atomic-examples/phoenix-tracing/pyproject.toml +++ b/atomic-examples/phoenix-tracing/pyproject.toml @@ -15,7 +15,7 @@ authors = [ ] requires-python = ">=3.12" dependencies = [ - "atomic-agents[observability]", + "atomic-agents[phoenix]", "arize-phoenix-otel>=0.17.0,<1.0.0", "openinference-instrumentation-openai>=0.1.30,<0.2.0", "instructor>=1.17.0,<2", diff --git a/docs/guides/observability.md b/docs/guides/observability.md index 1bcff648..fe479272 100644 --- a/docs/guides/observability.md +++ b/docs/guides/observability.md @@ -1,69 +1,148 @@ -# Observability with Arize Phoenix and OpenTelemetry +# Observability with OpenTelemetry -Atomic Agents ships an optional [OpenInference](https://github.com/Arize-ai/openinference) instrumentor. Once it -is installed, every `AtomicAgent` run and every `BaseTool` run becomes a span that -[Arize Phoenix](https://github.com/Arize-ai/phoenix), Arize AX, Langfuse, or any other OpenTelemetry backend can -display. Nothing changes for applications that do not opt in. +Atomic Agents can turn every agent run, tool run and LLM attempt into an OpenTelemetry span. The spans follow the +OpenTelemetry GenAI semantic conventions, so they show up on any OpenTelemetry backend: a plain OTLP collector, +Grafana Tempo, Sentry, Jaeger, Honeycomb, Datadog, Arize Phoenix, Langfuse and so on. Only `opentelemetry-api` is +required. Nothing changes for applications that do not opt in. -## What you get +You get three kinds of spans: -| Call | Span | Attributes | -|------|------|------------| -| `agent.run`, `run_async`, `run_stream`, `run_async_stream` | `AGENT`, named `AtomicAgent[InputSchema, OutputSchema].run` (a subclass reports its own name) | `agent.name`, `input.value` / `output.value` as JSON, `atomic_agents.model`, `atomic_agents.mode`, `atomic_agents.input_schema`, `atomic_agents.output_schema`, `atomic_agents.history.message_count` (history size entering the turn), `atomic_agents.context_providers`, `atomic_agents.system_prompt.sha256` (hash of the static prompt sections: background, steps, output instructions) | -| `tool.run`, `run_async`, `arun` on any `BaseTool` subclass, including MCP tools | `TOOL`, named after the tool class | `tool.name`, `tool.description`, `tool.parameters` (the input JSON schema), `input.value` / `output.value` as JSON | +- **AGENT**: one per `AtomicAgent` run (`run`, `run_async`, `run_stream`, `run_async_stream`). +- **TOOL**: one per `BaseTool` run, including MCP tools. +- **LLM attempt**: one per call Instructor makes to the model, so retries after a validation error are visible. -Spans carry `session.id`, `user.id`, `metadata` and tags when you wrap calls in the OpenInference context -managers (`using_session`, `using_user`, `using_metadata`, `using_attributes`). Exceptions set the span status to -`ERROR` and are recorded as span events. Streaming spans record the final response when the stream completes, -or the last partial plus `atomic_agents.stream.closed_early = true` when the consumer stops early or the task -is cancelled (a cancelled stream is not reported as an error). +## Install -LLM spans are not produced by Atomic Agents. Install exactly one LLM instrumentor for the provider SDK you use -(`openinference-instrumentation-openai`, `-anthropic`, `-litellm`, ...). Installing both a provider instrumentor and -`openinference-instrumentation-instructor` gives you two layers of LLM spans. Because the agent span is the current -span while the provider call runs, the LLM span (with prompts, tokens and the model name) nests under it. Instructor -retries show up as several LLM spans under one agent span. +```bash +pip install "atomic-agents[observability]" opentelemetry-sdk opentelemetry-exporter-otlp-proto-http +``` -## Install +The `observability` extra installs `opentelemetry-api` and nothing else. You bring the SDK and an exporter for your +backend. The `phoenix` extra adds the OpenInference packages: ```bash -pip install "atomic-agents[observability]" arize-phoenix-otel openinference-instrumentation-openai +pip install "atomic-agents[phoenix]" ``` -The `observability` extra installs `openinference-instrumentation`, `openinference-semantic-conventions` and -`opentelemetry-api`. The Phoenix path below works without the extra because `arize-phoenix-otel` already depends on -those packages; the manual path needs the extra plus an OpenTelemetry SDK, an exporter, and the instrumentor for -your LLM client (the example below uses the OpenAI one): +With those installed, Atomic Agents also writes the OpenInference attribute family (turned on automatically) and the +OpenInference context managers `using_session`, `using_user` and `using_metadata` start to work. Backends that read +OpenInference attributes, such as Arize Phoenix, use them for their own views. Every other backend works from the GenAI +attributes alone. + +## Setup with any OpenTelemetry backend (manual) + +```python +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + +from atomic_agents.observability import AtomicAgentsInstrumentor, TracingConfig + +tracer_provider = TracerProvider() +tracer_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter())) + +AtomicAgentsInstrumentor().instrument(tracer_provider=tracer_provider, config=TracingConfig()) +``` + +`OTLPSpanExporter()` reads the standard `OTEL_EXPORTER_OTLP_*` environment variables, so this code stays the same for +every backend below. Only the environment changes. `tracer_provider` and `config` are both optional: without a provider, +the global one is used, and without a config you get the defaults described under "Privacy and size". + +Call `instrument()` once, at import time or in your service's startup hook, before any agent runs. + +## Backends + +The GenAI attributes are the OpenTelemetry standard, and backends with LLM views (Datadog LLM Observability, Langfuse, +Phoenix, Sentry AI monitoring, PostHog) read them. Every recipe below uses the setup above. Only the exporter +configuration differs. + +What has actually been run: the span output is tested offline against an in-memory exporter, and the Phoenix recipe is +the [phoenix-tracing example](https://github.com/eigenwise/atomic-agents/tree/main/atomic-examples/phoenix-tracing). +Everything marked "untested" follows the vendor's documented OTLP settings but has not been run against a live account. +Check the vendor's page before relying on it. + +**Plain OTLP collector** (untested against a live collector): ```bash -pip install "atomic-agents[observability]" opentelemetry-sdk opentelemetry-exporter-otlp-proto-http openinference-instrumentation-openai +export OTEL_SERVICE_NAME=my-app +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 +export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf ``` -## Setup with Phoenix (automatic) +**Grafana Tempo** (untested): send OTLP to Tempo directly or through a collector or Grafana Alloy. Self-hosted, point the +endpoint at Tempo's OTLP HTTP port (4318 by default). Grafana Cloud uses its OTLP gateway plus a basic-auth header built +from your instance ID and an access token: + +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=https://otlp-gateway-.grafana.net/otlp +export OTEL_EXPORTER_OTLP_HEADERS="Authorization=Basic%20" +``` + +**Jaeger** (untested): current Jaeger releases accept OTLP directly. + +```bash +docker run --rm -p 16686:16686 -p 4318:4318 jaegertracing/all-in-one +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 +``` + +**Sentry** (follows Sentry's Python OTLP integration docs, untested here). The integration sets up the exporter itself, +so you do not build a `TracerProvider` and do not set `traces_sample_rate`: + +```bash +pip install "atomic-agents[observability]" "sentry-sdk[opentelemetry-otlp]" +``` ```python -from phoenix.otel import register +import sentry_sdk +from sentry_sdk.integrations.otlp import OTLPIntegration -tracer_provider = register(project_name="my-app", auto_instrument=True) +from atomic_agents.observability import AtomicAgentsInstrumentor + +sentry_sdk.init(dsn="https://@o.ingest.sentry.io/", integrations=[OTLPIntegration()]) +AtomicAgentsInstrumentor().instrument() ``` -`register(auto_instrument=True)` discovers the Atomic Agents instrumentor through the `openinference_instrumentor` -entry point and activates it together with every other OpenInference instrumentor in the environment. Do this once, -at import time or in your service's startup hook, before any agent runs. Phoenix reads `PHOENIX_COLLECTOR_ENDPOINT`, -`PHOENIX_PROJECT` (alias `PHOENIX_PROJECT_NAME`) and `PHOENIX_API_KEY` from the environment. With no endpoint set, -`register()` exports over gRPC to `localhost:4317`; set `PHOENIX_COLLECTOR_ENDPOINT=http://localhost:6006` to send -HTTP to the port the Phoenix UI serves on. +Sentry's OTLP ingest drops span events. The `validation_error` event and recorded exceptions do not arrive there, while +`error.type`, the span status and `atomic_agents.llm.last_validation_error` do. Whether Sentry's AI Agents dashboard +fills from these spans is unverified. -Two things to know about the automatic path: +**Honeycomb** (untested): use `https://api.eu1.honeycomb.io` for the EU region. Honeycomb's Agent Timeline groups by +`gen_ai.conversation.id`, so wrap calls in `using_conversation` (see "Naming and conversations"). -- Phoenix imports `atomic_agents` to load the entry point. That import is the same as any `import atomic_agents` in - your application, including the Instructor streaming patch the framework applies at import. -- Existing `auto_instrument=True` users start exporting agent and tool payloads when they upgrade to a version with the - instrumentor. Use the privacy controls below if that is not wanted. +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=https://api.honeycomb.io +export OTEL_EXPORTER_OTLP_HEADERS="x-honeycomb-team=" +``` + +**Datadog** (untested): send OTLP to a Datadog Agent with OTLP ingestion enabled (HTTP on 4318 by default) and set +`OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318`. Datadog LLM Observability maps `invoke_agent` spans to agents, +`execute_tool` to tools and `chat` to LLM calls, and reads the `gen_ai.usage.*` tokens and `gen_ai.conversation.id`. +It reads prompt and response text only from `gen_ai.input.messages` and `gen_ai.output.messages`, which Atomic Agents +does not emit, so you see spans, tokens and sessions but not the agent text. PostHog behaves the same way. + +**Arize Phoenix**: see "Setup with Phoenix" below. + +**Langfuse** (untested): Langfuse takes OTLP at `/api/public/otel` with basic auth built from your public and secret +key. Its default filter keeps spans that carry `gen_ai.*` attributes, which these spans do. + +```bash +export OTEL_EXPORTER_OTLP_ENDPOINT=https://cloud.langfuse.com/api/public/otel +export OTEL_EXPORTER_OTLP_HEADERS="Authorization=Basic%20" +``` + +**PostHog** (untested): PostHog AI Observability ingests OTLP over HTTP only (protobuf or JSON, no gRPC). Use the +`TRACES` variables, because the generic endpoint variable appends `/v1/traces` and gets a 404. Use your project's client +API host in the URL, see [PostHog's OpenTelemetry page](https://posthog.com/docs/ai-observability/installation/opentelemetry). +`chat` spans become `$ai_generation` events and the others become `$ai_span`. + +```bash +export OTEL_EXPORTER_OTLP_TRACES_ENDPOINT="/i/v0/ai/otel" +export OTEL_EXPORTER_OTLP_TRACES_HEADERS="Authorization=Bearer%20" +``` ## Setup with the OpenTelemetry launcher (automatic) -Install the optional tracing adapter, launcher, SDK, and exporter explicitly: +Install the tracing adapter, launcher, SDK and exporter explicitly: ```bash pip install "atomic-agents[observability]" opentelemetry-distro opentelemetry-exporter-otlp-proto-http @@ -75,115 +154,294 @@ Run your existing application with environment-only configuration: export OTEL_SERVICE_NAME=my-app export OTEL_TRACES_EXPORTER=otlp export OTEL_EXPORTER_OTLP_PROTOCOL=http/protobuf -export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:6006 +export OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318 export OTEL_METRICS_EXPORTER=none export OTEL_LOGS_EXPORTER=none -export OPENINFERENCE_HIDE_INPUTS=true -export OPENINFERENCE_HIDE_OUTPUTS=true +export ATOMIC_AGENTS_TRACE_HIDE_INPUTS=true +export ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS=true +export ATOMIC_AGENTS_TRACE_LLM_SPANS=false opentelemetry-instrument python app.py ``` -The launcher discovers `atomic_agents` through the `opentelemetry_instrumentor` entry point and activates the -same adapter used by Phoenix and manual setup. Without the OpenInference dependencies, the loader skips the adapter. -`opentelemetry-bootstrap` does not automatically install Atomic Agents or its observability extra. +The launcher discovers `atomic_agents` through the `opentelemetry_instrumentor` entry point and activates the same +adapter as the manual setup. Without `opentelemetry-api` installed, the loader skips the adapter. +`opentelemetry-bootstrap` does not automatically install Atomic Agents or its extras. -To disable this adapter, set `OTEL_PYTHON_DISABLED_INSTRUMENTATIONS=atomic_agents` (or add `atomic_agents` to an -existing comma-separated list). Ordinary imports without a launcher or explicit setup remain uninstrumented. +To disable the adapter, set `OTEL_PYTHON_DISABLED_INSTRUMENTATIONS=atomic_agents` (or add `atomic_agents` to an +existing comma-separated list). Ordinary imports without a launcher or explicit setup stay uninstrumented. -Choose one provider configuration path: the launcher, Phoenix registration, or manual setup. Both discovery entry -points share one installation, so activating both does not duplicate Atomic Agents spans. The first provider/config -wins; a later call with different settings warns and keeps the existing installation. OpenTelemetry also keeps its -first global provider. Calling Phoenix registration after launcher startup therefore does not switch the export -destination. To combine explicit activations, configure a shared provider, or uninstrument before changing the -Atomic Agents provider/config as described below. +Choose one provider configuration path: the launcher, Phoenix registration, or manual setup. The discovery entry points +share one installation, so activating several does not duplicate spans. OpenTelemetry itself keeps its first global +provider, so calling Phoenix registration after launcher startup does not switch where spans are exported. See +"Re-instrumentation" for what a second `instrument()` call does. -The launcher adds AGENT and TOOL spans, with no parent CHAIN or provider LLM spans of its own. Install exactly one -instrumentor for your LLM provider and activate it through its supported setup. Use the grouping example below when -you want several agent/tool calls in one trace. +The launcher adds AGENT, TOOL and LLM-attempt spans and no parent span. The example above disables Atomic Agents' LLM attempt +spans for applications that also use a provider instrumentor. Leave `ATOMIC_AGENTS_TRACE_LLM_SPANS` unset to keep +attempt spans enabled. Use the grouping example below when you want several agent and tool calls in one trace. -## Setup with any OpenTelemetry backend (manual) +## Setup with Phoenix (automatic) + +`register(auto_instrument=True)` also activates the OpenAI instrumentor, so set +`ATOMIC_AGENTS_TRACE_LLM_SPANS=false` (before import) to avoid two LLM layers: ```python -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import BatchSpanProcessor -from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter -from openinference.instrumentation import OITracer, TraceConfig -from openinference.instrumentation.openai import OpenAIInstrumentor +import os -from atomic_agents.observability import AtomicAgentsInstrumentor +os.environ.setdefault("ATOMIC_AGENTS_TRACE_LLM_SPANS", "false") -tracer_provider = TracerProvider() -tracer_provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter("http://localhost:6006/v1/traces"))) -trace_config = TraceConfig() +from phoenix.otel import register -AtomicAgentsInstrumentor().instrument(tracer_provider=tracer_provider, config=trace_config) -OpenAIInstrumentor().instrument(tracer_provider=tracer_provider, config=trace_config) -tracer = OITracer(tracer_provider.get_tracer(__name__), config=trace_config) +tracer_provider = register(project_name="my-app", auto_instrument=True) ``` -The last line matters for the grouping section below: a plain SDK provider hands out a plain OpenTelemetry tracer, -and the `openinference_span_kind` argument, `set_input` and `set_output` only exist on an OpenInference `OITracer`. -Phoenix's `register()` returns a provider whose tracers already are `OITracer`s, so the automatic path skips this -step. Pass the same `TraceConfig` to every layer; the config given to `instrument()` configures only the Atomic Agents -spans. +`register(auto_instrument=True)` discovers the Atomic Agents instrumentor through the `openinference_instrumentor` entry +point and activates it together with every other OpenInference instrumentor in the environment. Install +`atomic-agents[phoenix]` (or `arize-phoenix-otel`, which brings the same packages) and do this once, at import time or in +your startup hook, before any agent runs. Phoenix reads `PHOENIX_COLLECTOR_ENDPOINT`, `PHOENIX_PROJECT` (alias +`PHOENIX_PROJECT_NAME`) and `PHOENIX_API_KEY` from the environment. With no endpoint set, `register()` exports over gRPC +to `localhost:4317`; set `PHOENIX_COLLECTOR_ENDPOINT=http://localhost:6006` to send HTTP to the port the Phoenix UI +serves on. -`instrument()` is idempotent. Calling it again with a different provider or config logs a warning and keeps the first -installation; call `uninstrument()` first to change it. `uninstrument()` restores the original methods on -`AtomicAgent` and on every tool class. +Two things to know about the automatic path: -## Grouping agent and tool calls into one trace +- Phoenix imports `atomic_agents` to load the entry point. That import is the same as any `import atomic_agents` in + your application, including the Instructor streaming patch the framework applies at import. +- Existing `auto_instrument=True` users start exporting agent and tool payloads when they upgrade to a version with the + instrumentor. Use the privacy controls below if that is not wanted. -Atomic Agents has no tool-calling loop: your code calls `agent.run(...)`, then `tool.run(...)`, then another agent. -Each of those is a root span, so without a parent they become separate traces. Open one parent span per unit of work: +Langfuse and other OpenInference-aware backends use the same OpenInference attributes, so the Phoenix-style setup, or +the manual setup with the `phoenix` extra installed, works for them too. + +## What you get + +| Call | Span name | Kind | +|------|-----------|------| +| `AtomicAgent.run`, `run_async`, `run_stream`, `run_async_stream` | `invoke_agent {agent name}` | AGENT | +| `run`, `run_async`, `arun` on any `BaseTool` subclass | `execute_tool {tool name}` | TOOL | +| Each model call Instructor makes during an agent run | `chat {request model}` | LLM attempt | + +The agent name is `AgentConfig.name`, or the class name when you do not set one. The tool name is the name the model +sees (`tool_name` on the tool). Spans attach to the current context and never start a new trace when a parent span +exists, so your own request or job span becomes their parent. An agent span is the parent of its LLM attempt spans, and +a tool span nests under whatever span is current when the tool runs. A span that fails gets +status `ERROR`, the exception recorded, and `error.type` set to the exception class name. + +Attributes, GenAI keys first. All of them are written whether or not the OpenInference packages are installed. + +| Key | On | Value | +|-----|----|-------| +| `gen_ai.operation.name` | all | `invoke_agent`, `execute_tool` or `chat` | +| `gen_ai.agent.name` | AGENT | Agent name | +| `gen_ai.tool.name` | TOOL | Tool name the model sees | +| `gen_ai.tool.description` | TOOL | Tool description | +| `gen_ai.tool.type` | TOOL | `function` | +| `gen_ai.tool.call.arguments` | TOOL | Tool input as JSON (hidden by `hide_inputs`, size-capped) | +| `gen_ai.tool.call.result` | TOOL | Tool output as JSON (hidden by `hide_outputs`, size-capped) | +| `gen_ai.request.model` | AGENT, LLM | Model the agent was configured with, or the model in the completion call | +| `gen_ai.provider.name` | AGENT, LLM | Instructor provider mapped to the OpenTelemetry value, for example `azure.ai.openai`, `gcp.vertex_ai`, `aws.bedrock`; `openai` and `anthropic` pass through | +| `gen_ai.response.model` | AGENT, LLM | Model the provider reported; on the AGENT span, the last response's model | +| `gen_ai.usage.input_tokens` | AGENT, LLM | Input tokens (AGENT: total across attempts) | +| `gen_ai.usage.output_tokens` | AGENT, LLM | Output tokens | +| `gen_ai.usage.cache_read.input_tokens` | AGENT, LLM | Cached input tokens, when the provider reports them | +| `gen_ai.usage.reasoning.output_tokens` | AGENT, LLM | Reasoning tokens, when the provider reports them | +| `gen_ai.conversation.id` | all | Set by `using_conversation` | +| `error.type` | failed spans | Exception class name | + +Atomic Agents' own keys: + +| Key | On | Value | +|-----|----|-------| +| `atomic_agents.agent.class` | AGENT | For example `AtomicAgent[InputSchema, OutputSchema]` (a subclass reports its own class) | +| `atomic_agents.tool.class` | TOOL | Tool class name | +| `atomic_agents.model` | AGENT | Same as the request model | +| `atomic_agents.mode` | AGENT | Instructor mode | +| `atomic_agents.input_schema`, `atomic_agents.output_schema` | AGENT | Schema class names | +| `atomic_agents.history.message_count` | AGENT | History size entering the turn | +| `atomic_agents.context_providers` | AGENT | Registered context provider names | +| `atomic_agents.system_prompt.sha256` | AGENT | Hash of the static prompt sections (background, steps, output instructions). Context provider output is not included, so it identifies the prompt version even when inputs are hidden | +| `atomic_agents.input`, `atomic_agents.output` | AGENT | Agent input and output as JSON (hidden and size-capped as described under "Privacy and size"). A stream records the final response, or the last partial one | +| `atomic_agents.stream.closed_early` | AGENT | `true` when a stream consumer stopped early or the task was cancelled (not reported as an error) | +| `atomic_agents.llm.attempts` | AGENT | Number of model calls in the run | +| `atomic_agents.llm.last_validation_error` | AGENT | Text of the last validation error | +| `atomic_agents.llm.attempt` | LLM | 1-based attempt number | + +OpenInference keys, optional. These are written only when the OpenInference packages are installed (the `phoenix` +extra), or when you force them with `TracingConfig(openinference_attributes=True)`; set it to `False` to drop them even +when the packages are present. + +| Key | On | Value | +|-----|----|-------| +| `openinference.span.kind` | all | `AGENT`, `TOOL` or `LLM` | +| `input.value`, `input.mime_type` | AGENT, TOOL | Same payload as `atomic_agents.input` or `gen_ai.tool.call.arguments` | +| `output.value`, `output.mime_type` | AGENT, TOOL | Same payload as `atomic_agents.output` or `gen_ai.tool.call.result` | +| `agent.name` | AGENT | Agent name | +| `tool.name`, `tool.description` | TOOL | Same as the GenAI tool keys | +| `tool.parameters` | TOOL | The input JSON schema (size-capped, not hidden) | +| `session.id` | all | The conversation id, plus anything from `using_session`, `using_user` and `using_metadata` | + +## LLM attempt spans + +Instructor drives the model call and its retry loop, so the detail of each attempt only exists at the Atomic Agents +layer. The instrumentor registers hooks on the agent's Instructor client (once per client, on the first instrumented +run) and removes them again in `uninstrument()`: + +- `completion:kwargs` starts the `chat {model}` span as a child of the agent span and counts the attempt. +- `completion:response` sets `gen_ai.response.model` and the token usage on that span and ends it with status OK. +- `completion:error` records the exception, sets `error.type` and ends the span with status `ERROR`. +- `parse:error` adds a `validation_error` event (with `atomic_agents.llm.attempt`) to the AGENT span and sets + `atomic_agents.llm.last_validation_error`. Instructor then retries, which opens the next attempt span. +- `completion:usage` writes the running totals to the AGENT span: the four `gen_ai.usage.*` keys and + `atomic_agents.llm.attempts`. + +An attempt is one call to the model. A retry is any attempt after the first, normally because the response failed +validation, so two attempt spans under one agent span means one retry. Hooks you registered on the same client yourself +keep running next to these, because Instructor keeps a list of handlers per event. If the agent span ends while an +attempt is still open, the attempt span is ended with an unset status. The hooks never raise into your code. + +OpenTelemetry's GenAI guidance says an agent framework should not emit model-call spans by default when the LLM library +already does. Atomic Agents emits them by default because the attempt and validation detail comes from Instructor, which +no provider instrumentor sees. If you also run a provider instrumentor (`openinference-instrumentation-openai` and +friends), turn the attempt spans off or you get two LLM layers: ```python -with tracer.start_as_current_span("answer_question", openinference_span_kind="chain") as turn: - turn.set_input(question) - plan = planner.run(BasicChatInputSchema(chat_message=question)) - calculation = calculator.run(CalculatorToolInputSchema(expression=plan.expression)) - answer = answerer.run(CalculationResultSchema(question=question, expression=plan.expression, result=calculation.result)) - turn.set_output(answer.chat_message) +AtomicAgentsInstrumentor().instrument(tracer_provider=tracer_provider, config=TracingConfig(llm_spans=False)) +``` + +With `llm_spans=False` the usage totals and attempt count still land on the AGENT span, and the `validation_error` event +still fires. + +## Naming and conversations + +Give each agent a stable name, so spans group the same way across deployments and class renames: + +```python +agent = AtomicAgent[InputSchema, OutputSchema]( + config=AgentConfig(client=client, model="gpt-5-mini", name="orchestrator_planning") +) ``` -`tracer` is the `OITracer` from the manual setup, or `tracer_provider.get_tracer(__name__)` on the Phoenix path. The -[phoenix-tracing example](https://github.com/eigenwise/atomic-agents/tree/main/atomic-examples/phoenix-tracing) -is exactly this flow, and the test suite runs the manual setup plus this grouping offline with an in-memory exporter. +The span becomes `invoke_agent orchestrator_planning`. Without a name, the class name is used. -## Sessions, users and metadata +Group related runs into a conversation with `using_conversation`: ```python -from openinference.instrumentation import using_session, using_user, using_metadata +from atomic_agents.observability import using_conversation -with using_session(session_id=conversation_id), using_user(user_id=user.id), using_metadata({"tenant": tenant}): +with using_conversation(conversation_id): response = agent.run(user_message) ``` -Phoenix groups spans that share a `session.id` into a conversation view. For a streaming endpoint, put the context -manager inside the generator that yields the chunks, not around the endpoint function; the generator runs after the -endpoint returned, so a context manager around the endpoint body no longer applies. +Every AGENT, TOOL and LLM attempt span started inside the block gets `gen_ai.conversation.id` (and `session.id` when +OpenInference keys are on). With the `phoenix` extra installed, `using_session`, `using_user` and `using_metadata` from +`openinference.instrumentation` keep working, and a `using_session` value is copied to `gen_ai.conversation.id` when you +have not set one with `using_conversation`. + +For a streaming endpoint, put the context manager inside the generator that yields the chunks, not around the endpoint +function. The generator runs after the endpoint returned, so a context manager around the endpoint body no longer +applies. + +## Privacy and size + +By default, agent and tool payloads are exported. `TracingConfig` controls that: + +```python +TracingConfig(hide_inputs=True, hide_outputs=True, max_attribute_length=2048) +``` + +| Field | Default | Effect | +|-------|---------|--------| +| `hide_inputs` | `False` | Input payload values become `__REDACTED__` | +| `hide_outputs` | `False` | Output payload values become `__REDACTED__` | +| `max_attribute_length` | `4096` | Longer values are cut to this many characters and end in `...[truncated]` | +| `openinference_attributes` | auto | See "What you get" | +| `llm_spans` | `True` | See "LLM attempt spans" | + +When no config is passed to `instrument()`, the settings come from the environment. `ATOMIC_AGENTS_TRACE_HIDE_INPUTS` +and `ATOMIC_AGENTS_TRACE_HIDE_OUTPUTS` are read first, then `OPENINFERENCE_HIDE_INPUTS` and `OPENINFERENCE_HIDE_OUTPUTS` +as a fallback (`true` or `1`). `ATOMIC_AGENTS_TRACE_MAX_ATTRIBUTE_LENGTH` sets the size cap. +`ATOMIC_AGENTS_TRACE_LLM_SPANS` disables attempt spans for `false` or `0` (case-insensitive); any other value or unset +keeps them enabled. If you pass your own +`TracingConfig(...)`, the environment is not read; use `TracingConfig.from_environment()` to start from it. An +OpenInference `TraceConfig` is also accepted for `config`; its `hide_inputs` and `hide_outputs` override those +environment settings, while the other settings still come from the environment. + +A truncated value gets a companion attribute with the original size, for example `atomic_agents.input.original_length`. +Hiding wins over truncation, and a hidden payload is never serialized in the first place. + +Covered by hiding and truncation: `atomic_agents.input`, `atomic_agents.output`, `input.value`, `output.value`, +`gen_ai.tool.call.arguments`, `gen_ai.tool.call.result` and `atomic_agents.llm.last_validation_error`. `tool.parameters` +is truncated but never hidden. Status descriptions and recorded exception messages are truncated too, and under +`hide_outputs` a Pydantic validation error message is replaced by `__REDACTED__` plus the error count, because the raw +model output is quoted inside it. + +Instructor's own logging is separate from Atomic Agents' tracing: in the installed Instructor 1.17.0 source, +`instructor.v2.retry` (re-exported through `instructor.core.retry`) logs per-attempt validation failures at DEBUG +and retry exhaustion's raw last error, including model output, at ERROR regardless of `hide_outputs`. +Setting its level to ERROR does not silence exhaustion. To keep model output out of logs, use CRITICAL below +or a filter/handler that drops those records. The logger name and levels may change in later Instructor versions. + +```python +import logging +logging.getLogger("instructor.v2.retry").setLevel(logging.CRITICAL) +``` + +Not covered: tool names and descriptions, the system prompt hash, and whatever you put in `using_metadata`. A provider +instrumentor and your own tracer apply their own settings, so configure them the same way if you want matching masking. +Multimodal inputs (images, audio, PDFs) are serialized into the input payload in full, so `max_attribute_length` is the +control for those. + +## Grouping agent and tool calls into one trace + +Atomic Agents has no tool-calling loop: your code calls `agent.run(...)`, then `tool.run(...)`, then another agent. +Each of those is a root span, so without a parent they become separate traces. Open one parent span per unit of work +with a plain OpenTelemetry tracer. It also answers "where do the seconds go": add child spans for retrieval or database +work next to the agent and tool calls. + +```python +tracer = tracer_provider.get_tracer(__name__) + +with tracer.start_as_current_span("answer_question"): + with tracer.start_as_current_span("retrieve_documents"): + documents = search_index(question) + plan = planner.run(BasicChatInputSchema(chat_message=question)) + calculation = calculator.run(CalculatorToolInputSchema(expression=plan.expression)) + answer = answerer.run(CalculationResultSchema(question=question, expression=plan.expression, result=calculation.result)) +``` + +The agent, tool and attempt spans become children of `answer_question`, and the slowest step stands out in the +waterfall. The [phoenix-tracing example](https://github.com/eigenwise/atomic-agents/tree/main/atomic-examples/phoenix-tracing) +is this flow. On the Phoenix path the provider's tracers are OpenInference `OITracer`s, which also take +`openinference_span_kind="chain"` and `set_input`/`set_output` if you want Phoenix to label the parent as a CHAIN. -## Privacy +## Re-instrumentation -Agent and tool payloads are exported as `input.value` / `output.value`. To redact them: +`instrument()` with the same provider and config as the active installation does nothing. With a different provider or +config it replaces them in place and logs a WARNING that says so. Spans already in flight finish on the provider they +started with; new spans use the new one. `uninstrument()` followed by `instrument()` is the clean way to change +settings, and `uninstrument()` restores the original methods on `AtomicAgent` and on every tool class. -- On the automatic Phoenix or OpenTelemetry launcher path, set `OPENINFERENCE_HIDE_INPUTS=true` and/or - `OPENINFERENCE_HIDE_OUTPUTS=true`. `TraceConfig()` reads these variables before the first installation. -- On the manual path, pass `TraceConfig(hide_inputs=True, hide_outputs=True)` to `instrument()`. +In a test suite, call `instrument(tracer_provider=provider)` with a fresh in-memory provider in each test. The replace +behavior means you do not need to uninstrument between tests, and the warning is expected there: -Redacted values are exported as `__REDACTED__`. Masking covers the `input.value` and `output.value` attributes of -Atomic Agents spans and nothing else: exception events and status descriptions (a validation error can quote input -values), tool names and descriptions, and whatever you put in `using_metadata` are exported as they are. The LLM -instrumentor and your own `OITracer` each take their own `TraceConfig`, so pass the same one to all three if you want -the same masking everywhere. `atomic_agents.system_prompt.sha256` stays available either way. It -hashes the static sections of a `SystemPromptGenerator` (background, steps, output instructions) and ignores what -context providers add per call, so it identifies the prompt version even when inputs are hidden. Multimodal inputs (images, audio, PDFs) are serialized into -`input.value` in full; cap attribute size with `OTEL_ATTRIBUTE_VALUE_LENGTH_LIMIT` if that is too much. +```python +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter + +exporter = InMemorySpanExporter() +provider = TracerProvider() +provider.add_span_processor(SimpleSpanProcessor(exporter)) +AtomicAgentsInstrumentor().instrument(tracer_provider=provider) +agent.run(user_input) +assert [span.name for span in exporter.get_finished_spans()][-1] == "invoke_agent my_agent" +``` ## Services and async -- Call `register(...)` (or `instrument(...)`) at import or in the startup hook so every worker process is instrumented. -- Use `register(batch=True)` (or a `BatchSpanProcessor`) in services so span export does not block request handling. -- Concurrent `run_async` calls, including `asyncio.gather` over several agents, get independent spans. +- Call `instrument(...)` (or `register(...)`) at import or in the startup hook so every worker process is instrumented. +- Use a `BatchSpanProcessor` (or `register(batch=True)`) in services so span export does not block request handling. +- Concurrent `run_async` calls, including `asyncio.gather` over several agents sharing one client, get independent + spans and their own attempt spans. - Streaming token counts need `model_api_parameters={"stream_options": {"include_usage": True}}` for OpenAI. ## Limitations @@ -191,12 +449,15 @@ context providers add per call, so it identifies the prompt version even when in - Methods bound before instrumentation stay untraced: `callback = agent.run` taken before `instrument()` keeps calling the original. - An `AtomicAgent` subclass that overrides `run` (or a stream method) without calling `super()` bypasses the wrapper - and emits no span; add `@tracer.agent` to such overrides yourself. -- Tool methods that are async generators are left untraced. + and emits no span. Open a span in the override yourself, with `tracer.start_as_current_span(...)` or `@tracer.agent` + on an `OITracer`. +- Tool methods that are async generators are left untraced. Open a span inside the generator body with + `tracer.start_as_current_span(...)`, or use `@tracer.tool` on an `OITracer`. - A plain `ThreadPoolExecutor` starts its worker with an empty context, so a tool that submits work to one loses the current span and the re-entrancy guard inside the pool. Submit through a copied context instead, as the forge tools do: `executor.submit(contextvars.copy_context().run, asyncio.run, self.run_async(params))`. `asyncio.to_thread` and `asyncio.run` already copy the context. - The AGENT span records the final (or last partial) response of a stream, not every chunk. -- The system prompt text is not on the AGENT span; it is the first message on the LLM span. The AGENT span carries its - SHA-256 instead. +- The system prompt text is not on any span. The AGENT span carries its SHA-256 instead. +- Datadog and PostHog read prompt and response text from `gen_ai.input.messages` and `gen_ai.output.messages`, which + are not emitted. diff --git a/pyproject.toml b/pyproject.toml index 87e34649..e720bc5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,10 +37,11 @@ dependencies = [ ] [project.optional-dependencies] -observability = [ +observability = ["opentelemetry-api>=1.30.0,<2.0.0"] +phoenix = [ + "opentelemetry-api>=1.30.0,<2.0.0", "openinference-instrumentation>=0.1.38,<0.2.0", "openinference-semantic-conventions>=0.1.20,<0.2.0", - "opentelemetry-api>=1.30.0,<2.0.0", ] [project.urls] diff --git a/uv.lock b/uv.lock index aa3aaa5a..8ced42d6 100644 --- a/uv.lock +++ b/uv.lock @@ -284,6 +284,9 @@ dependencies = [ [package.optional-dependencies] observability = [ + { name = "opentelemetry-api" }, +] +phoenix = [ { name = "openinference-instrumentation" }, { name = "openinference-semantic-conventions" }, { name = "opentelemetry-api" }, @@ -320,9 +323,10 @@ requires-dist = [ { name = "instructor", specifier = ">=1.17.0,<2" }, { name = "litellm", specifier = ">=1.50.0,<2.0.0" }, { name = "mcp", extras = ["cli"], specifier = ">=1.8.0,<2" }, - { name = "openinference-instrumentation", marker = "extra == 'observability'", specifier = ">=0.1.38,<0.2.0" }, - { name = "openinference-semantic-conventions", marker = "extra == 'observability'", specifier = ">=0.1.20,<0.2.0" }, + { name = "openinference-instrumentation", marker = "extra == 'phoenix'", specifier = ">=0.1.38,<0.2.0" }, + { name = "openinference-semantic-conventions", marker = "extra == 'phoenix'", specifier = ">=0.1.20,<0.2.0" }, { name = "opentelemetry-api", marker = "extra == 'observability'", specifier = ">=1.30.0,<2.0.0" }, + { name = "opentelemetry-api", marker = "extra == 'phoenix'", specifier = ">=1.30.0,<2.0.0" }, { name = "pydantic", specifier = ">=2.11.4,<3.0.0" }, { name = "pyfiglet", specifier = ">=1.0.2,<2.0.0" }, { name = "pyyaml", specifier = ">=6.0.2,<7.0.0" }, @@ -330,7 +334,7 @@ requires-dist = [ { name = "rich", specifier = ">=13.7.1,<14.0.0" }, { name = "textual", specifier = ">=5.3.0,<6.0.0" }, ] -provides-extras = ["observability"] +provides-extras = ["observability", "phoenix"] [package.metadata.requires-dev] dev = [ @@ -3161,7 +3165,7 @@ version = "1.0.0" source = { editable = "atomic-examples/phoenix-tracing" } dependencies = [ { name = "arize-phoenix-otel" }, - { name = "atomic-agents", extra = ["observability"] }, + { name = "atomic-agents", extra = ["phoenix"] }, { name = "instructor" }, { name = "openai" }, { name = "openinference-instrumentation-openai" }, @@ -3172,7 +3176,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "arize-phoenix-otel", specifier = ">=0.17.0,<1.0.0" }, - { name = "atomic-agents", extras = ["observability"], editable = "." }, + { name = "atomic-agents", extras = ["phoenix"], editable = "." }, { name = "instructor", specifier = ">=1.17.0,<2" }, { name = "openai", specifier = ">=2.0.0,<3.0.0" }, { name = "openinference-instrumentation-openai", specifier = ">=0.1.30,<0.2.0" },