[SDK] feat: add native Pydantic AI integration#7536
Conversation
github-actions
Bot
added
documentation
![baz-reviewer[bot]](https://avatars.githubusercontent.com/in/933528?s=60&v=4){kind=small}
| assert result.output is not None | ||
| assert len(fake_backend.trace_trees) == 1 | ||
| trace_tree = fake_backend.trace_trees[0] | ||
|
|
||
| assert trace_tree.name == "run my_agent" | ||
| assert trace_tree.project_name == "pydantic-ai-test" | ||
| assert trace_tree.input == {"prompt": "hello"} | ||
| assert trace_tree.output is not None | ||
|
|
||
| root_span = trace_tree.spans[0] | ||
| assert root_span.name == "run my_agent" | ||
| assert root_span.type == "general" |
There was a problem hiding this comment.
Misses span nesting regressions
According to the PR description, trace_tree.output is not None, root_span.name == "run my_agent", and llm_spans = [s for s in _flatten(trace_tree.spans) if s.type == "llm"] only prove a trace exists and that at least one llm span was emitted, and _flatten() drops the parent/child shape so flat or mis-parented spans would still pass. Can we assert the expected tree layout and concrete output value here, and apply the same pattern to the tool, structured-output, and nested-trace tests below?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/pydantic_ai/test_pydantic_ai.py around lines 35-60
in `test_pydantic_ai__standalone_run__creates_trace_with_nested_spans`, replace the
current `_flatten(trace_tree.spans)`-based `llm_spans` assertions with checks that
verify the actual parent/child span hierarchy from `trace_tree.spans` (e.g., assert the
root span named “run my_agent” directly contains the expected `llm` span(s), and
fail if spans are only present in a flat list). Also capture the concrete run output by
asserting `trace_tree.output` matches the actual `result.output` (not just `is not
None`), using the same output shape conventions already used in the structured-output
test. Apply the same “don’t flatten; assert tree structure + concrete output”
pattern to the tool test (lines 62-79), the structured-output test (lines 81-95), and
the nested-trace test (lines 97-117) so they all validate correct nesting and precise
traced values.
| def test_track_pydantic_ai__returns_instrumented_agent(): | ||
| agent = Agent(TestModel(), name="return_agent") | ||
| returned = track_pydantic_ai(agent) | ||
| assert returned is agent | ||
| assert isinstance(agent.instrument, InstrumentationSettings) |
There was a problem hiding this comment.
Default instrumentation path untested
According to the PR description, track_pydantic_ai(agent=None) is the primary entry point, but this test only covers track_pydantic_ai(agent), so regressions in Agent.instrument_all(settings) would still pass CI — can we add a library-integration test for track_pydantic_ai() with no agent that creates a fresh Agent(TestModel()) and asserts the trace/span tree?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/tests/library_integration/pydantic_ai/test_pydantic_ai.py around lines
131-135, extend the coverage for `test_track_pydantic_ai__returns_instrumented_agent` by
adding a new test for the `track_pydantic_ai(agent=None)` branch. Create a test that
calls `track_pydantic_ai()` with no arguments, then instantiates a fresh
`Agent(TestModel())`, runs it (e.g., `run_sync("hello")`), flushes with
`opik.flush_tracker()`, and asserts the fake backend has exactly one trace tree with at
least the expected root span name/type (and optionally nested llm/tool spans as in the
other tests). Refactor or reuse existing helpers like `_instrument`,
`_provider_with_opik_processor`, and `_flatten` to keep assertions consistent and ensure
the test is isolated from other tests’ global instrumentation side effects.
| from typing import Any, Dict, List, Optional | ||
|
|
||
| from opentelemetry.sdk.trace import TracerProvider | ||
| from pydantic_ai import Agent | ||
| from pydantic_ai.models.instrumented import InstrumentationSettings |
There was a problem hiding this comment.
Agent and InstrumentationSettings are imported from pydantic_ai at module load time, so import opik.integrations.pydantic_ai raises ModuleNotFoundError on installs without the optional dependency even when track_pydantic_ai() is never called — should we move those imports inside track_pydantic_ai(), use TYPE_CHECKING/string annotations for type hints, and raise a clear actionable error only when the function is invoked without the dependency?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/pydantic_ai/opik_tracker.py around lines 1-5, remove
the module-level imports of `Agent` and `InstrumentationSettings` from `pydantic_ai`.
Move those imports inside `track_pydantic_ai()` (lazy import), and catch
`ModuleNotFoundError` to raise a clear, actionable error (e.g. "pydantic-ai is required
for this integration — install it with `pip install pydantic-ai`") only when the
function is called. Update any type hints that reference `Agent` or
`InstrumentationSettings` to use `TYPE_CHECKING` guards or string annotations so they
don't require `pydantic_ai` at runtime. Also check
sdks/python/src/opik/integrations/pydantic_ai/__init__.py to ensure `track_pydantic_ai`
is re-exported without eagerly importing `pydantic_ai`, so that `import
opik.integrations.pydantic_ai` succeeds even when the optional dependency is absent.
Ensure `track_pydantic_ai()` continues to work identically once the dependency is
available.
| try: | ||
| self._handle_start(span) | ||
| except Exception: # per OTel contract on_start must not raise | ||
| LOGGER.debug("Failed to start Opik span for Pydantic AI", exc_info=True) |
There was a problem hiding this comment.
Tracing failures are hidden
on_start() and on_end() catch every exception and log at LOGGER.debug(...), so span-processing failures stay invisible in normal deployments while tracing data is dropped for that run — should we raise these to warning or error?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/pydantic_ai/opik_span_processor.py around lines
120-123 in `on_start()` and lines 318-323 in `on_end()`, the code catches all exceptions
and logs them only with `LOGGER.debug(...)`, which makes span-processing failures
effectively invisible and can drop spans/trace data silently. Refactor by changing these
log levels to `warning` or `error` (and ensure `exc_info=True` is preserved), so
operators can see when the integration fails to emit spans. Keep the “must not
raise” behavior for OTEL callbacks, but make the failure observable via the higher
severity logs.
| client = opik.get_global_client() | ||
| if client.config.log_start_trace_span: | ||
| client.__internal_api__span__(**result.span_data.as_start_parameters) | ||
|
|
||
| if result.trace_data is not None: | ||
| context_storage.set_trace_data(result.trace_data) | ||
| if client.config.log_start_trace_span: | ||
| client.__internal_api__trace__(**result.trace_data.as_start_parameters) |
There was a problem hiding this comment.
Leaked span context on start failure
context_storage.add_span_data(result.span_data) runs before the synchronous client.__internal_api__span__ / __internal_api__trace__ calls, and if either export raises on_start() swallows the error so _handle_start() never rolls back the pushed state, leaving stale trace context behind for later spans — can we wrap the start-path bookkeeping in a try/finally or otherwise remove it on failure?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/pydantic_ai/opik_span_processor.py around lines
149-168 in the _handle_start() method, the code calls
context_storage.add_span_data(result.span_data) and optionally
context_storage.set_trace_data(result.trace_data) before the synchronous
client.__internal_api__span__ / __internal_api__trace__ exports. If either export
raises, on_start() swallows the exception and the span is never added to self._span_map,
so on_end() can’t pop the shared context—leaving stale trace/span state for later
spans. Refactor _handle_start() to ensure atomic start-path bookkeeping: wrap the client
export calls in try/except (or try/finally) and, on failure, immediately remove any
already-pushed span/trace state via context_storage.pop_span_data(ensure_id=...) /
context_storage.pop_trace_data(ensure_id=...) and do not write to self._span_map unless
the start exports succeeded.
| if metadata.get("opik.provider"): | ||
| result["provider"] = metadata["opik.provider"] |
There was a problem hiding this comment.
Mismatched override disables pricing
metadata["opik.provider"] is copied into result["provider"] and later fed to validate_and_parse_usage(..., provider=provider), so a mismatched override like "openai" on an Anthropic/Bedrock-shaped payload can take the wrong parser path and leave the span with misleading provider data and no correct cost attribution — should we keep the override as display metadata only, or validate/normalize it before storing it as the cost-tracking provider?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
`sdks/python/src/opik/integrations/pydantic_ai/opik_span_processor.py` around lines
235-252, in `_extract_agent_run_data()`, the code copies `metadata[
| if metadata.get("opik.metadata"): | ||
| result["opik_metadata"] = metadata["opik.metadata"] |
There was a problem hiding this comment.
Malformed metadata drops the run
metadata["opik.metadata"] is copied into result["opik_metadata"], so span_data.update(metadata=...)/trace_data.update(metadata=...) pass non-dict values into data_helpers.merge_metadata(), where dict_utils.deepmerge() calls new_data.items() and _handle_end() aborts before client.__internal_api__span__ or client.__internal_api__trace__ runs — should we validate opik.metadata is a dict first?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/pydantic_ai/opik_span_processor.py around lines
235-252 in the `_extract_agent_run_data` logic, the code assigns
`result["opik_metadata"] = metadata["opik.metadata"]` without verifying that
`metadata["opik.metadata"]` is a mapping/dict. Update this branch to check
`isinstance(metadata["opik.metadata"], dict)` (or Mapping) before assigning; if it’s
not a dict, either ignore it or coerce to an empty dict. Then ensure `_handle_end` can
always proceed so `span_data.init_end_time().update(... metadata=...)` and
`trace_data.init_end_time().update(... metadata=...)` never receive a non-mapping value
that would make `merge_metadata()`/`deepmerge()` raise.
| if msg.get("role") == "user": | ||
| text = _text_from_parts(msg.get("parts", [])) | ||
| if text: | ||
| result["input"] = {"prompt": text} | ||
| break |
There was a problem hiding this comment.
Stale prompt captured from history
_extract_agent_run_data() picks the first role == "user" from pydantic_ai.all_messages, so history-carrying runs can store an earlier prompt in trace.input/span.input while the recorded output belongs to the latest turn — should we use the current-run user message, or the last user message in all_messages, instead?
Want Baz to fix this for you? Activate Fixer
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
sdks/python/src/opik/integrations/pydantic_ai/opik_span_processor.py around lines
235-284, in the `_extract_agent_run_data` logic that derives `result['input']` from
`pydantic_ai.all_messages`, you currently iterate messages from the start and `break` on
the first `role == 'user'`. Because `all_messages` can include history from prior runs,
this incorrectly captures an earlier prompt as the completed run’s
`trace.input`/`span.input` (while output is taken from the latest model/assistant
messages). Refactor this section to select the current-run user message
instead—preferably by iterating the messages in reverse and taking the first user
message encountered (last user in the list), or if `new_messages` is available in the
span attrs, use that list for the user-input extraction—so `result['input']` matches
the run that just finished.
| prompt = opik.Prompt(name="assistant", prompt="Be concise, reply with one sentence.") | ||
|
|
||
| result = agent.run_sync( | ||
| 'Where does "hello world" come from?', | ||
| metadata={ | ||
| "opik.metadata": {"experiment": "baseline"}, | ||
| "opik.prompts": [prompt.__internal_api__to_info_dict__()], | ||
| }, |
There was a problem hiding this comment.
prompt.__internal_api__to_info_dict__() is a private hook, so the user-facing opik.prompts guide depends on an unstable API that will drift when the name changes — should we add a public prompt-info helper and update this page and the duplicate call in apps/opik-documentation/documentation/fern/docs-v2/integrations/pydantic-ai.mdx, as apps/opik-documentation/AGENTS.md suggests?
Want Baz to fix this for you? Activate Fixer You can also update your AI coding guidelines based on this comment by apply pr to [branch name]
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
apps/opik-documentation/documentation/fern/docs/tracing/integrations/pydantic-ai.mdx
around lines 148-155, the "Enriching traces with metadata and prompts" example uses
prompt.__internal_api__to_info_dict__(), which is a private/internal API. Add a public
prompt-info helper in the Opik SDK (e.g., a method on the Prompt type or a small utility
function) that returns the public "prompt info" representation needed for the
opik.prompts metadata, then update this doc to call that new public helper instead of
the internal method. Alternatively, switch to the documented prompt serialization path
or link to the autogenerated prompt-library docs. Also apply the same fix to the
duplicate call in
apps/opik-documentation/documentation/fern/docs-v2/integrations/pydantic-ai.mdx to keep
both guides consistent, and remove any other occurrences of
__internal_api__to_info_dict__ across the documentation to comply with AGENTS.md.
| metadata={ | ||
| "opik.metadata": {"experiment": "baseline"}, | ||
| "opik.prompts": [prompt.__internal_api__to_info_dict__()], | ||
| }, |
There was a problem hiding this comment.
Internal prompt API leaked in docs
The docs example tells users to build opik.prompts with prompt.__internal_api__to_info_dict__(), so it teaches an internal SDK hook instead of a public prompt-to-dict API — should we expose a helper on Prompt or track_pydantic_ai first and point the docs at that, as apps/opik-documentation/AGENTS.md suggests?
Want Baz to fix this for you? Activate Fixer You can also update your AI coding guidelines based on this comment by apply pr to [branch name]
Other fix methods
Prompt for AI Agents
Before applying, verify this suggestion against the current code. In
apps/opik-documentation/documentation/fern/docs-v2/integrations/pydantic-ai.mdx around
lines 151-154 (within “Enriching traces with metadata and prompts”), the example
builds `metadata["opik.prompts"]` using the internal
`prompt.__internal_api__to_info_dict__()` serialization hook. Replace this with a
public, user-facing helper by first adding an SDK method/function (e.g.,
`Prompt.to_info_dict()` or a top-level `opik.prompts.to_info_dict(prompt)`) that returns
the same dict format required by the ingestion logic. Then update the docs example to
call the new public helper and confirm the generated dict shape matches what
`track_pydantic_ai`/OTEL ingestion expects. Also search the documentation for any other
`__internal_api__*` usages related to prompt serialization and update them to the public
helper.
alexkuzmik
left a comment
alexkuzmik
left a comment
There was a problem hiding this comment.
Thanks @aduverger! The integration PR looks very promising, please see some comments I left.
| @@ -0,0 +1,41 @@ | |||
| from opik.integrations.pydantic_ai.opik_span_processor import _token_usage, _try_json | |||
There was a problem hiding this comment.
- Even though it's close to a classic unit test, please put it in the
tests/library_integration/pydantic_aifor consistency with the other integrations and cleaner dependencies managements. - Avoid accessing private functions/attributes in tests. Either test the existing public API, or, if you need to cover a specific piece of logic, extract it to a separate module. If it deserves to be tested directly, it deserves to be in a more specialized module.
| self._span_map: Dict[int, Tuple[SpanData, Optional[TraceData]]] = {} | ||
|
|
||
| @staticmethod | ||
| def _classify_span(name: str) -> SpanType: |
There was a problem hiding this comment.
Let's re-arrange the logic, the class handles too much. We could at least split it into span_parsers.py (or smth like that) and opik_span_processor.py. No need to put all this logic into this class; it reduces cohesion and increases the information context.
| return None | ||
|
|
||
|
|
||
| def _token_usage(attrs: Dict[str, Any]) -> Optional[Dict[str, int]]: |
There was a problem hiding this comment.
please take a look at opik.llm_usage, let's convert tokens to this format instead of the anthropic one.
|
|
||
|
|
||
| def track_pydantic_ai( | ||
| agent: Optional[Any] = None, |
There was a problem hiding this comment.
Please use strict types whenever possible.
2 participants
Context
We've been using Pydantic AI with Opik on production for about 6 months now. The logfire workaround was not good enough as we needed better integration and customisation of Opik traces.
The custom Pydantic AI
OpikSpanProcessorwas initially developed in our in-house agent framework. I moved it into Opik with Claude as part of this PR.What I would like to discuss in particular is token tracking: I had to implement a few rules to retrieve cached tokens (and deduce them from input tokens). Do you have a more generalisable way to do it for any LLM provider ?
Details
Adds a native
opik.integrations.pydantic_aimodule so Pydantic AI agents are traced with a singletrack_pydantic_ai()call — no logfire and no manual OpenTelemetry env vars. Everyagent.runbecomes an Opik trace with correctly named, nestedllm/toolspans, token usage and cost. This replaces the documented logfire workaround, which required an extra dependency, hand-written OTEL env-var glue, and left span naming/enrichment to the backend.OpikSpanProcessormapsgen_ai.*spans:invoke_agent→general(renamedrun <agent>),chat <model>→llm,execute_tool→tool. It creates a trace when none is active, or nests under the surrounding@opik.tracktrace.final_result), tool args/result, and error info. Reservedopik.*keys (opik.metadata,opik.provider,opik.prompts,opik.thread_id) are read from the nativeagent.run(..., metadata=...).track_pydantic_ai(agent=None, *, project_name, metadata, tags)instruments all agents or a single one. Following convention,pydantic-aiis imported lazily and pinned only in the integration's testrequirements.txt(notsetup.py).Change checklist
Issues
AI-WATERMARK
AI-WATERMARK: yes
Testing
pytest tests/unit/integrations/pydantic_ai tests/library_integration/pydantic_ai— 10 passing (6 library-integration via Pydantic AI'sTestModel, hermetic/no API keys; 4 unit for cache-token math and JSON parsing).pre-commit run --files <changed py files>— ruff, ruff-format, mypy all clean.@opik.trackyields a single trace; tool spans; structured output;thread_idgrouping via metadata.Documentation
docs-v2/integrations/pydantic-ai.mdxanddocs/tracing/integrations/pydantic-ai.mdx(nativeopik.configure()→track_pydantic_ai()flow; cost tracking; thread grouping; metadata/prompt enrichment;@tracknesting).img/tracing/pydantic-ai/pydantic-ai_trace.png) showingrun <agent>→chat <model>→execute_toolspans from a real run.