diff --git a/logfire/_internal/integrations/pydantic_ai.py b/logfire/_internal/integrations/pydantic_ai.py index d077d9ae6..8b751574d 100644 --- a/logfire/_internal/integrations/pydantic_ai.py +++ b/logfire/_internal/integrations/pydantic_ai.py @@ -1,7 +1,8 @@ from __future__ import annotations +from typing import overload import inspect -from typing import Any, Literal +from typing import Any, Literal, TYPE_CHECKING from pydantic_ai import Agent from pydantic_ai.agent import InstrumentationSettings @@ -10,16 +11,59 @@ from logfire import Logfire +if TYPE_CHECKING: + from pydantic_ai.embeddings import Embedder, EmbeddingModel, InstrumentedEmbeddingModel +else: + try: + from pydantic_ai.embeddings import Embedder, EmbeddingModel, InstrumentedEmbeddingModel + except ImportError: + Embedder = None + EmbeddingModel = None + InstrumentedEmbeddingModel = None + +@overload +def instrument_pydantic_ai( + logfire_instance: Logfire, + obj: Agent | Embedder | None, + include_binary_content: bool | None, + include_content: bool | None, + version: Literal[1, 2, 3] | None, + event_mode: Literal['attributes', 'logs'] | None, + **kwargs: Any, +) -> None: ... + +@overload +def instrument_pydantic_ai( + logfire_instance: Logfire, + obj: Model, + include_binary_content: bool | None, + include_content: bool | None, + version: Literal[1, 2, 3] | None, + event_mode: Literal['attributes', 'logs'] | None, + **kwargs: Any, +) -> InstrumentedModel: ... + +@overload def instrument_pydantic_ai( logfire_instance: Logfire, - obj: Agent | Model | None, + obj: EmbeddingModel, include_binary_content: bool | None, include_content: bool | None, version: Literal[1, 2, 3] | None, event_mode: Literal['attributes', 'logs'] | None, **kwargs: Any, -) -> None | InstrumentedModel: +) -> InstrumentedEmbeddingModel: ... + +def instrument_pydantic_ai( + logfire_instance: Logfire, + obj: Agent | Model | Embedder | EmbeddingModel | None, + include_binary_content: bool | None, + include_content: bool | None, + version: Literal[1, 2, 3] | None, + event_mode: Literal['attributes', 'logs'] | None, + **kwargs: Any, +) -> None | InstrumentedModel | InstrumentedEmbeddingModel: # Correctly handling all past and future versions is tricky. # Since we provide these rather than the user, only include them if # InstrumentationSettings has such parameters to prevent errors/warnings. @@ -59,7 +103,13 @@ def instrument_pydantic_ai( obj.instrument = settings elif isinstance(obj, Model): return InstrumentedModel(obj, settings) + elif Embedder and isinstance(obj, Embedder): + obj.instrument = settings + elif EmbeddingModel and InstrumentedEmbeddingModel and isinstance(obj, EmbeddingModel): + return InstrumentedEmbeddingModel(obj, settings) elif obj is None: Agent.instrument_all(settings) + if Embedder: + Embedder.instrument_all(settings) else: raise TypeError(f'Cannot instrument object of type {type(obj)}') diff --git a/logfire/_internal/main.py b/logfire/_internal/main.py index 0cfc1fc1e..56df3d136 100644 --- a/logfire/_internal/main.py +++ b/logfire/_internal/main.py @@ -1032,13 +1032,13 @@ def instrument_pydantic_ai( Args: obj: What to instrument. - By default, all agents are instrumented. - You can also pass a specific model or agent. + By default, all agents and embedders are instrumented. + You can also pass a specific agent, model, embedder, or embedding model. If you pass a model, a new instrumented model will be returned. include_binary_content: Whether to include base64 encoded binary content (e.g. images) in the telemetry. On by default. Requires Pydantic AI 0.2.5 or newer. - include_content: Whether to include prompts, completions, and tool call arguments and responses - in the telemetry. On by default. Requires Pydantic AI 0.3.4 or newer. + include_content: Whether to include prompts, completions, tool call arguments and responses, + and embedded documents and embeddings in the telemetry. On by default. Requires Pydantic AI 0.3.4 or newer. version: Version of the data format. This is unrelated to the Pydantic AI package version. Requires Pydantic AI 0.7.5 or newer. Version 1 is based on the legacy event-based OpenTelemetry GenAI spec