Instrument Pydantic AI embedders and embedding models

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)}')

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