Address copilot

Author: DylanRussell

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

@@ -94,6 +94,9 @@
     "_SUPPRESS_EMBEDDING", default=False
 )
 _patched_embedder_classes: set[type[Any]] = set()
+_embedder_init_subclass_saved: bool = False
+_had_orig_embedder_init_subclass: bool = False
+_orig_embedder_init_subclass: Any = None
 
 
 _KNOWN_EMBEDDERS: tuple[tuple[str, str], ...] = (
@@ -319,8 +322,30 @@ def _wrap_subclasses(base_cls: type[Any]) -> None:
 
         _wrap_subclasses(Embedder)
 
+        global \
+            _embedder_init_subclass_saved, \
+            _had_orig_embedder_init_subclass, \
+            _orig_embedder_init_subclass
+        if not _embedder_init_subclass_saved:
+            _had_orig_embedder_init_subclass = (
+                "__init_subclass__" in Embedder.__dict__
+            )
+            _orig_embedder_init_subclass = Embedder.__dict__.get(
+                "__init_subclass__"
+            )
+            _embedder_init_subclass_saved = True
+
+        orig_init_subclass = _orig_embedder_init_subclass
+
         # Wrap new embedder classes created during runtime
-        def _traced_init_subclass(cls: type[Any], **kwargs: Any) -> None:
+        def _traced_init_subclass(cls: type[Embedder], **kwargs: Any) -> None:
+            if orig_init_subclass is not None:
+                func = getattr(
+                    orig_init_subclass, "__func__", orig_init_subclass
+                )
+                func(cls, **kwargs)
+            else:
+                super(Embedder, cls).__init_subclass__(**kwargs)
             if _is_instrumented:
                 _wrap_embedder_class(cls, handler)
 
@@ -346,8 +371,28 @@ def _safe_unwrap(target: Any, attr: str) -> None:
     try:
         from agno.knowledge.embedder.base import Embedder
 
-        if "__init_subclass__" in Embedder.__dict__:
-            delattr(Embedder, "__init_subclass__")
+        global \
+            _embedder_init_subclass_saved, \
+            _had_orig_embedder_init_subclass, \
+            _orig_embedder_init_subclass
+        if _embedder_init_subclass_saved:
+            if (
+                _had_orig_embedder_init_subclass
+                and _orig_embedder_init_subclass is not None
+            ):
+                setattr(
+                    Embedder,
+                    "__init_subclass__",
+                    _orig_embedder_init_subclass,
+                )
+            elif (
+                hasattr(Embedder, "__init_subclass__")
+                and "__init_subclass__" in Embedder.__dict__
+            ):
+                delattr(Embedder, "__init_subclass__")
+            _embedder_init_subclass_saved = False
+            _orig_embedder_init_subclass = None
+            _had_orig_embedder_init_subclass = False
     except Exception:
         pass
 
@@ -449,6 +494,17 @@ def _extract_embedder_model(embedder: Any) -> str | 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"):
+        val = usage.get(key)
+        if val is not None:
+            try:
+                return int(val)
+            except (ValueError, TypeError):
+                return None
+    return None
+
+
 def _wrap_embedder_class(
     cls: type[Any],
     handler: TelemetryHandler,
@@ -559,16 +615,9 @@ def traced_method(
                     pass
 
             if isinstance(usage, dict):
-                input_tokens = (
-                    usage.get("prompt_tokens")
-                    or usage.get("input_tokens")
-                    or usage.get("total_tokens")
-                )
+                input_tokens = _extract_embedder_input_tokens(usage)
                 if input_tokens is not None:
-                    try:
-                        invocation.input_tokens = int(input_tokens)
-                    except (ValueError, TypeError):
-                        pass
+                    invocation.input_tokens = input_tokens
                 if model := usage.get("model"):
                     invocation.response_model_name = str(model)
 
@@ -667,16 +716,9 @@ async def traced_method(
                     pass
 
             if isinstance(usage, dict):
-                input_tokens = (
-                    usage.get("prompt_tokens")
-                    or usage.get("input_tokens")
-                    or usage.get("total_tokens")
-                )
+                input_tokens = _extract_embedder_input_tokens(usage)
                 if input_tokens is not None:
-                    try:
-                        invocation.input_tokens = int(input_tokens)
-                    except (ValueError, TypeError):
-                        pass
+                    invocation.input_tokens = input_tokens
                 if model := usage.get("model"):
                     invocation.response_model_name = str(model)
 

instrumentation/opentelemetry-instrumentation-genai-agno/tests/conformance/embedding.py

@@ -0,0 +1,67 @@
+# Copyright The OpenTelemetry Authors
+# SPDX-License-Identifier: Apache-2.0
+
+"""Conformance scenario: embedding generation for Agno."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from agno.knowledge.embedder.base import Embedder
+
+from opentelemetry.instrumentation.genai.agno import AgnoInstrumentor
+from opentelemetry.sdk._logs import LoggerProvider
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.test_util_genai.conformance import (
+    ExpectedViolation,
+    Scenario,
+)
+from opentelemetry.test_util_genai.instrumentor import instrument
+
+
+@dataclass
+class ConformanceEmbedder(Embedder):
+    id: str = "text-embedding-3-small"
+    provider: str = "openai"
+
+    def get_embedding_and_usage(
+        self, text: str
+    ) -> tuple[list[float], dict[str, Any]]:
+        return [0.1, 0.2, 0.3, 0.4], {
+            "input_tokens": 8,
+            "model": "text-embedding-3-small",
+        }
+
+
+class EmbeddingScenario(Scenario):
+    expected_spans = {"embeddings": 1}
+    expected_metrics = (
+        "gen_ai.client.operation.duration",
+        "gen_ai.client.token.usage",
+    )
+    expected_violations = (
+        ExpectedViolation(
+            advice_id="genai_expected_attribute_missing",
+            message_substring="server.address",
+        ),
+    )
+
+    def run(
+        self,
+        *,
+        tracer_provider: TracerProvider,
+        meter_provider: MeterProvider,
+        logger_provider: LoggerProvider,
+        vcr: Any,
+    ) -> None:
+        with instrument(
+            AgnoInstrumentor(),
+            tracer_provider=tracer_provider,
+            logger_provider=logger_provider,
+            meter_provider=meter_provider,
+            content_capture="SPAN_ONLY",
+        ):
+            embedder = ConformanceEmbedder()
+            embedder.get_embedding_and_usage("conformance test embedding input")

instrumentation/opentelemetry-instrumentation-genai-agno/tests/test_conformance.py

@@ -21,6 +21,7 @@
 )
 
 from .conformance.agent import AgentScenario
+from .conformance.embedding import EmbeddingScenario
 from .conformance.retrieval import RetrievalScenario
 from .conformance.workflow import WorkflowScenario
 from .conformance.workflow_streaming import WorkflowStreamingScenario
@@ -30,6 +31,7 @@
     "scenario",
     [
         pytest.param(AgentScenario()),
+        pytest.param(EmbeddingScenario()),
         pytest.param(RetrievalScenario()),
         pytest.param(WorkflowScenario()),
         pytest.param(WorkflowStreamingScenario()),

instrumentation/opentelemetry-instrumentation-genai-agno/tests/test_embedder.py

@@ -7,6 +7,7 @@
 
 import asyncio
 from dataclasses import dataclass
+from typing import Any
 
 import pytest
 from agno.knowledge.embedder.base import Embedder
@@ -236,6 +237,21 @@ def test_embedder_error_handling(
     assert span.attributes.get(ERROR_TYPE) == "ValueError"
 
 
+@dataclass
+class ZeroUsageEmbedder(Embedder):
+    model: str = "zero-tokens-model"
+
+    def get_embedding_and_usage(
+        self, text: str
+    ) -> tuple[list[float], dict[str, int]]:
+        return [0.1, 0.2], {"prompt_tokens": 0, "total_tokens": 12}
+
+    async def async_get_embedding_and_usage(
+        self, text: str
+    ) -> tuple[list[float], dict[str, int]]:
+        return [0.3, 0.4], {"prompt_tokens": 0, "total_tokens": 10}
+
+
 def test_embedder_uninstrument_restores_untraced_behavior(
     instrument_agno,
     span_exporter,
@@ -250,3 +266,72 @@ def test_embedder_uninstrument_restores_untraced_behavior(
 
     embedder.get_embedding("untraced")
     assert len(span_exporter.get_finished_spans()) == 0
+
+
+def test_embedder_zero_input_tokens_usage(
+    instrument_agno,
+    span_exporter,
+) -> None:
+    """Test that 0 prompt_tokens is recorded as 0 and not overridden by total_tokens."""
+    embedder = ZeroUsageEmbedder()
+    _, usage = embedder.get_embedding_and_usage("zero text")
+    assert usage.get("prompt_tokens") == 0
+
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].attributes.get(GEN_AI_USAGE_INPUT_TOKENS) == 0
+
+    span_exporter.clear()
+
+    async def _test() -> None:
+        await embedder.async_get_embedding_and_usage("zero text async")
+
+    asyncio.run(_test())
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    assert spans[0].attributes.get(GEN_AI_USAGE_INPUT_TOKENS) == 0
+
+
+def test_embedder_init_subclass_calls_original_and_unpatch_restores(
+    instrument_agno,
+    span_exporter,
+) -> None:
+    """Test that custom base __init_subclass__ logic is called and restored on uninstrument."""
+    called_with: list[tuple[type[Any], str | None]] = []
+
+    class TrackedBaseEmbedder(Embedder):
+        def __init_subclass__(
+            cls, custom_arg: str | None = None, **kwargs: Any
+        ) -> None:
+            super().__init_subclass__(**kwargs)
+            called_with.append((cls, custom_arg))
+
+    class SubTrackedEmbedder(TrackedBaseEmbedder, custom_arg="value"):
+        def get_embedding(self, text: str) -> list[float]:
+            return [0.1]
+
+    # Verify original __init_subclass__ was called with subclass and kwargs
+    assert len(called_with) == 1
+    assert called_with[0] == (SubTrackedEmbedder, "value")
+
+    # Verify newly created subclass was instrumented
+    sub = SubTrackedEmbedder()
+    sub.get_embedding("traced")
+    spans = span_exporter.get_finished_spans()
+    assert len(spans) == 1
+    span_exporter.clear()
+
+    # Uninstrument and verify restoration
+    instrument_agno.uninstrument()
+    called_with.clear()
+
+    class SubTrackedEmbedderAfter(TrackedBaseEmbedder, custom_arg="after"):
+        def get_embedding(self, text: str) -> list[float]:
+            return [0.2]
+
+    assert len(called_with) == 1
+    assert called_with[0] == (SubTrackedEmbedderAfter, "after")
+
+    sub_after = SubTrackedEmbedderAfter()
+    sub_after.get_embedding("untraced")
+    assert len(span_exporter.get_finished_spans()) == 0