[opentelemetry-instrumentation-genai-agno] Add instrumentation for agno embedding methods

instrumentation/opentelemetry-instrumentation-genai-agno/.changelog/738.added

@@ -0,0 +1 @@
+Add instrumentation for standalone Agno Embedder operations.

instrumentation/opentelemetry-instrumentation-genai-agno/README.rst

@@ -35,6 +35,7 @@ The instrumentation automatically traces:
 * ``Workflow.run``, ``Workflow.arun`` (including background executions), ``Workflow.continue_run``, and ``Workflow.acontinue_run``
 * ``FunctionCall.execute`` and ``FunctionCall.aexecute``
 * ``Knowledge.search`` and ``Knowledge.asearch``
+* ``Embedder.get_embedding``, ``Embedder.get_embedding_and_usage``, ``Embedder.async_get_embedding``, and ``Embedder.async_get_embedding_and_usage``
 
 Configuration
 -------------

instrumentation/opentelemetry-instrumentation-genai-agno/src/opentelemetry/instrumentation/genai/agno/patch.py

@@ -24,6 +24,7 @@
 
 if TYPE_CHECKING:
     from agno.knowledge.document.base import Document
+    from agno.knowledge.embedder.base import Embedder
     from agno.knowledge.knowledge import Knowledge
     from agno.models.response import ToolExecution
     from agno.tools.function import FunctionCall, FunctionExecutionResult
@@ -49,6 +50,8 @@ class _HasToolExecution(Protocol):
     format_content,
     format_retrieval_document,
     prepare_tool_definitions,
+    resolve_embedder_provider,
+    safe_int,
     set_invocation_user_id,
 )
 from opentelemetry.instrumentation.utils import unwrap
@@ -100,6 +103,43 @@ class _HasToolExecution(Protocol):
 _ACTIVE_FOREGROUND_WORKFLOWS: contextvars.ContextVar[frozenset[int]] = (
     contextvars.ContextVar("_ACTIVE_FOREGROUND_WORKFLOWS", default=frozenset())
 )
+_SUPPRESS_EMBEDDING: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "_SUPPRESS_EMBEDDING", default=False
+)
+_patched_embedder_classes: set[type[Any]] = set()
+_KNOWN_EMBEDDERS: tuple[tuple[str, str], ...] = (
+    ("agno.knowledge.embedder.base", "Embedder"),
+    ("agno.knowledge.embedder.openai", "OpenAIEmbedder"),
+    ("agno.knowledge.embedder.azure_openai", "AzureOpenAIEmbedder"),
+    ("agno.knowledge.embedder.aws_bedrock", "AwsBedrockEmbedder"),
+    ("agno.knowledge.embedder.google", "GoogleEmbedder"),
+    ("agno.knowledge.embedder.google", "GeminiEmbedder"),
+    ("agno.knowledge.embedder.ollama", "OllamaEmbedder"),
+    ("agno.knowledge.embedder.mistral", "MistralEmbedder"),
+    ("agno.knowledge.embedder.cohere", "CohereEmbedder"),
+    ("agno.knowledge.embedder.fireworks", "FireworksEmbedder"),
+    ("agno.knowledge.embedder.jina", "JinaEmbedder"),
+    ("agno.knowledge.embedder.together", "TogetherEmbedder"),
+    ("agno.knowledge.embedder.voyageai", "VoyageAIEmbedder"),
+    ("agno.knowledge.embedder.fastembed", "FastEmbedEmbedder"),
+    (
+        "agno.knowledge.embedder.sentence_transformer",
+        "SentenceTransformerEmbedder",
+    ),
+    ("agno.knowledge.embedder.huggingface", "HuggingfaceCustomEmbedder"),
+    ("agno.knowledge.embedder.langdb", "LangDBEmbedder"),
+    ("agno.knowledge.embedder.nebius", "NebiusEmbedder"),
+    ("agno.knowledge.embedder.openai_like", "OpenAILikeEmbedder"),
+    ("agno.knowledge.embedder.vllm", "VLLMEmbedder"),
+)
+_EMBEDDER_METHODS: frozenset[str] = frozenset(
+    {
+        "get_embedding",
+        "get_embedding_and_usage",
+        "async_get_embedding",
+        "async_get_embedding_and_usage",
+    }
+)
 
 # wrapt has no unregister API for post-import hooks; monotonic generations
 # invalidate deferred hooks registered during prior instrumentation cycles.
@@ -262,6 +302,39 @@ def patch_agent(handler: TelemetryHandler) -> None:
         current_generation,
     )
 
+    embedder_wrappers = _embedder_method_wrappers(handler)
+    for mod_name, cls_name in _KNOWN_EMBEDDERS:
+        for method_name, wrapper_fn in embedder_wrappers:
+            _safe_wrap_function(
+                mod_name,
+                f"{cls_name}.{method_name}",
+                wrapper_fn,
+                current_generation,
+            )
+
+    try:
+        from agno.knowledge.embedder.base import Embedder
+
+        # Wrap all existing custom embedder classes
+        def _wrap_subclasses(base_cls: type[Any]) -> None:
+            for sub in base_cls.__subclasses__():
+                _wrap_embedder_class(sub, handler)
+                _wrap_subclasses(sub)
+
+        _wrap_subclasses(Embedder)
+
+        # Wrap new embedder classes created during runtime
+        def _traced_init_subclass(cls: type[Embedder], **kwargs: Any) -> None:
+            super(Embedder, cls).__init_subclass__(**kwargs)
+            if _is_instrumented:
+                _wrap_embedder_class(cls, handler)
+
+        setattr(
+            Embedder, "__init_subclass__", classmethod(_traced_init_subclass)
+        )
+    except Exception:
+        pass
+
 
 def unpatch_agent() -> None:
     """Remove patches from Agno class methods."""
@@ -275,6 +348,26 @@ def _safe_unwrap(target: Any, attr: str) -> None:
         except (AttributeError, ValueError):
             pass
 
+    try:
+        from agno.knowledge.embedder.base import Embedder
+
+        if "__init_subclass__" in Embedder.__dict__:
+            delattr(Embedder, "__init_subclass__")
+    except Exception:
+        pass
+
+    for cls in list(_patched_embedder_classes):
+        for attr in _EMBEDDER_METHODS:
+            _safe_unwrap(cls, attr)
+    _patched_embedder_classes.clear()
+    for mod_name, cls_name in _KNOWN_EMBEDDERS:
+        if mod_name in sys.modules:
+            mod = sys.modules[mod_name]
+            cls = getattr(mod, cls_name, None)
+            if cls is not None:
+                for attr in _EMBEDDER_METHODS:
+                    _safe_unwrap(cls, attr)
+
     if _AGNO_MODULE in sys.modules:
         try:
             import agno.agent
@@ -325,6 +418,251 @@ def _safe_unwrap(target: Any, attr: str) -> None:
             pass
 
 
+_extract_embedder_provider = resolve_embedder_provider
+
+
+def _extract_embedder_model(embedder: Embedder) -> str | None:
+    model = (
+        getattr(embedder, "id", None)
+        or getattr(embedder, "model", None)
+        or getattr(embedder, "name", None)
+    )
+    return str(model) if model is not None else None
+
+
+def _extract_embedder_input_tokens(usage: dict[str, Any]) -> int | None:
+    for key in ("prompt_tokens", "input_tokens", "total_tokens"):
+        if (tok := safe_int(usage.get(key))) is not None:
+            return tok
+    return None
+
+
+def _embedder_get_embedding(
+    handler: TelemetryHandler,
+) -> Callable[..., Any]:
+    def traced_method(
+        wrapped: Callable[..., Sequence[float]],
+        instance: Embedder,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Sequence[float]:
+        if _SUPPRESS_EMBEDDING.get():
+            return wrapped(*args, **kwargs)
+
+        token = _SUPPRESS_EMBEDDING.set(True)
+        provider = _extract_embedder_provider(instance)
+        request_model = _extract_embedder_model(instance)
+        invocation = handler.embedding(
+            provider=provider,
+            request_model=request_model,
+        )
+        if request_model:
+            invocation.response_model_name = request_model
+        set_invocation_user_id(
+            invocation, instance, args, kwargs, wrapped=wrapped
+        )
+        if encoding_format := getattr(instance, "encoding_format", None):
+            invocation.encoding_formats = [str(encoding_format)]
+
+        try:
+            result = wrapped(*args, **kwargs)
+            if result:
+                invocation.dimension_count = len(result)
+            else:
+                invocation.dimension_count = safe_int(instance.dimensions)
+            invocation.stop()
+            return result
+        except BaseException as error:
+            invocation.fail(error)
+            raise
+        finally:
+            _SUPPRESS_EMBEDDING.reset(token)
+
+    return traced_method
+
+
+def _embedder_get_embedding_and_usage(
+    handler: TelemetryHandler,
+) -> Callable[..., Any]:
+    def traced_method(
+        wrapped: Callable[..., tuple[Sequence[float], dict[str, Any] | None]],
+        instance: Embedder,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> tuple[Sequence[float], dict[str, Any] | None]:
+        if _SUPPRESS_EMBEDDING.get():
+            return wrapped(*args, **kwargs)
+
+        token = _SUPPRESS_EMBEDDING.set(True)
+        provider = _extract_embedder_provider(instance)
+        request_model = _extract_embedder_model(instance)
+        invocation = handler.embedding(
+            provider=provider,
+            request_model=request_model,
+        )
+        if request_model:
+            invocation.response_model_name = request_model
+        set_invocation_user_id(
+            invocation, instance, args, kwargs, wrapped=wrapped
+        )
+        if encoding_format := getattr(instance, "encoding_format", None):
+            invocation.encoding_formats = [str(encoding_format)]
+
+        try:
+            result = wrapped(*args, **kwargs)
+            embedding, usage = result
+            if embedding:
+                invocation.dimension_count = len(embedding)
+            else:
+                invocation.dimension_count = safe_int(instance.dimensions)
+
+            if isinstance(usage, dict):
+                invocation.input_tokens = _extract_embedder_input_tokens(usage)
+                if model := usage.get("model"):
+                    invocation.response_model_name = str(model)
+
+            invocation.stop()
+            return result
+        except BaseException as error:
+            invocation.fail(error)
+            raise
+        finally:
+            _SUPPRESS_EMBEDDING.reset(token)
+
+    return traced_method
+
+
+def _embedder_async_get_embedding(
+    handler: TelemetryHandler,
+) -> Callable[..., Any]:
+    async def traced_method(
+        wrapped: Callable[..., Awaitable[Sequence[float]]],
+        instance: Embedder,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> Sequence[float]:
+        if _SUPPRESS_EMBEDDING.get():
+            return await wrapped(*args, **kwargs)
+
+        token = _SUPPRESS_EMBEDDING.set(True)
+        provider = _extract_embedder_provider(instance)
+        request_model = _extract_embedder_model(instance)
+        invocation = handler.embedding(
+            provider=provider,
+            request_model=request_model,
+        )
+        if request_model:
+            invocation.response_model_name = request_model
+        set_invocation_user_id(
+            invocation, instance, args, kwargs, wrapped=wrapped
+        )
+        if encoding_format := getattr(instance, "encoding_format", None):
+            invocation.encoding_formats = [str(encoding_format)]
+
+        try:
+            result = await wrapped(*args, **kwargs)
+            if result:
+                invocation.dimension_count = len(result)
+            else:
+                invocation.dimension_count = safe_int(instance.dimensions)
+            invocation.stop()
+            return result
+        except BaseException as error:
+            invocation.fail(error)
+            raise
+        finally:
+            _SUPPRESS_EMBEDDING.reset(token)
+
+    return cast(Callable[..., Any], traced_method)
+
+
+def _embedder_async_get_embedding_and_usage(
+    handler: TelemetryHandler,
+) -> Callable[..., Any]:
+    async def traced_method(
+        wrapped: Callable[
+            ..., Awaitable[tuple[Sequence[float], dict[str, Any] | None]]
+        ],
+        instance: Embedder,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> tuple[Sequence[float], dict[str, Any] | None]:
+        if _SUPPRESS_EMBEDDING.get():
+            return await wrapped(*args, **kwargs)
+
+        token = _SUPPRESS_EMBEDDING.set(True)
+        provider = _extract_embedder_provider(instance)
+        request_model = _extract_embedder_model(instance)
+        invocation = handler.embedding(
+            provider=provider,
+            request_model=request_model,
+        )
+        if request_model:
+            invocation.response_model_name = request_model
+        set_invocation_user_id(
+            invocation, instance, args, kwargs, wrapped=wrapped
+        )
+        if encoding_format := getattr(instance, "encoding_format", None):
+            invocation.encoding_formats = [str(encoding_format)]
+
+        try:
+            result = await wrapped(*args, **kwargs)
+            embedding, usage = result
+            if embedding:
+                invocation.dimension_count = len(embedding)
+            else:
+                invocation.dimension_count = safe_int(instance.dimensions)
+
+            if isinstance(usage, dict):
+                invocation.input_tokens = _extract_embedder_input_tokens(usage)
+                if model := usage.get("model"):
+                    invocation.response_model_name = str(model)
+
+            invocation.stop()
+            return result
+        except BaseException as error:
+            invocation.fail(error)
+            raise
+        finally:
+            _SUPPRESS_EMBEDDING.reset(token)
+
+    return cast(Callable[..., Any], traced_method)
+
+
+def _embedder_method_wrappers(
+    handler: TelemetryHandler,
+) -> tuple[tuple[str, Callable[..., Any]], ...]:
+    return (
+        ("get_embedding", _embedder_get_embedding(handler)),
+        (
+            "get_embedding_and_usage",
+            _embedder_get_embedding_and_usage(handler),
+        ),
+        ("async_get_embedding", _embedder_async_get_embedding(handler)),
+        (
+            "async_get_embedding_and_usage",
+            _embedder_async_get_embedding_and_usage(handler),
+        ),
+    )
+
+
+def _wrap_embedder_class(
+    cls: type[Any],
+    handler: TelemetryHandler,
+) -> None:
+    for method_name, wrapper_fn in _embedder_method_wrappers(handler):
+        if hasattr(cls, method_name):
+            target = getattr(cls, method_name)
+            if not hasattr(target, "__wrapped__"):
+                try:
+                    wrap_function_wrapper(
+                        cast(Any, cls), method_name, wrapper_fn
+                    )
+                    _patched_embedder_classes.add(cls)
+                except Exception:
+                    pass
+
+
 def _extract_input_content(input_val: Any) -> str:
     if input_val is None:
         return ""