refactor: simplify embedding tracer handlers (OPEN-10480)

Code review findings addressed:
- Move per-call imports of _openai_embedding_common to module-level (was in
  hot path of every embedding call).
- Extract build_embedding_step_kwargs into _openai_embedding_common so that
  sync and async OpenAI handlers each become ~10 lines instead of ~50, and
  LiteLLM reuses the same kwargs assembly.
- Drop LiteLLM's local _parse_embedding_response and
  _get_embedding_model_parameters; both now delegate to the shared helpers
  (LiteLLM-specific timeout/api_base/api_version/cost/metadata are layered
  on top of the common kwargs).
- Type Bedrock _parse_embedding_output return as
  Tuple[Union[List[float], List[List[float]]], int, int] instead of bare
  tuple.

Net: -34 lines across the 5 touched source files. Tests unchanged, all
77 embedding tests + 448 lib tests still green.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: viniciusdsmello

src/openlayer/lib/integrations/_openai_embedding_common.py

@@ -1,19 +1,21 @@
-"""Shared parsing helpers for OpenAI sync + async embedding tracers."""
+"""Shared parsing helpers for OpenAI-shaped embedding tracers (OpenAI, AsyncOpenAI, LiteLLM)."""
 
-from typing import Any, Dict, List, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 
 def parse_embedding_response(
     response: Any,
 ) -> Tuple[Union[List[float], List[List[float]]], int, int]:
-    """Extract (embeddings, dimensions, count) from an OpenAI EmbeddingResponse.
+    """Extract (embeddings, dimensions, count) from an OpenAI-shaped EmbeddingResponse.
 
     For a single input, returns the vector directly.
     For a batch, returns a list of vectors.
     """
     try:
         data = getattr(response, "data", None)
-        if data is None:
+        if data is None and isinstance(response, dict):
+            data = response.get("data", [])
+        if not data:
             return [], 0, 0
         embeddings = [
             item["embedding"] if isinstance(item, dict) else item.embedding
@@ -35,3 +37,48 @@ def get_embedding_model_parameters(kwargs: Dict[str, Any]) -> Dict[str, Any]:
         "encoding_format": kwargs.get("encoding_format"),
         "user": kwargs.get("user"),
     }
+
+
+def build_embedding_step_kwargs(
+    response: Any,
+    call_kwargs: Dict[str, Any],
+    start_time: float,
+    end_time: float,
+    *,
+    name: str,
+    provider: str,
+    inference_id: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Build the kwargs to pass to ``tracer.add_embedding_step_to_trace``.
+
+    Common boilerplate for OpenAI-shaped responses (OpenAI sync/async, LiteLLM).
+    Callers may layer extra fields (cost, extra_metadata, model_parameters) on
+    top of the returned dict before invoking the tracer helper.
+    """
+    model_name = getattr(response, "model", call_kwargs.get("model", "unknown"))
+    embeddings, dim, count = parse_embedding_response(response)
+    usage = getattr(response, "usage", None)
+    prompt_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0
+    total_tokens = getattr(usage, "total_tokens", prompt_tokens) if usage else prompt_tokens
+
+    return {
+        "name": name,
+        "end_time": end_time,
+        "inputs": {"input": call_kwargs.get("input")},
+        "output": embeddings,
+        "latency": (end_time - start_time) * 1000,
+        "tokens": total_tokens,
+        "prompt_tokens": prompt_tokens,
+        "model": model_name,
+        "model_parameters": get_embedding_model_parameters(call_kwargs),
+        "embedding_dimensions": dim,
+        "embedding_count": count,
+        "raw_output": (
+            response.model_dump()
+            if hasattr(response, "model_dump")
+            else str(response)
+        ),
+        "provider": provider,
+        "id": inference_id,
+        "metadata": {"provider": provider},
+    }

src/openlayer/lib/integrations/async_openai_tracer.py

@@ -17,6 +17,7 @@
     import openai
 
 from ..tracing import tracer
+from ._openai_embedding_common import build_embedding_step_kwargs
 from .openai_tracer import (
     get_model_parameters,
     create_trace_args,
@@ -725,46 +726,21 @@ async def handle_embedding_async(
     **kwargs,
 ) -> Any:
     """Trace an async AsyncOpenAI client.embeddings.create() call."""
-    from ._openai_embedding_common import (
-        get_embedding_model_parameters as _get_embedding_model_parameters,
-    )
-    from ._openai_embedding_common import (
-        parse_embedding_response as _parse_embedding_response,
-    )
-
     start_time = time.time()
     response = await original_func(*args, **kwargs)
     end_time = time.time()
 
     try:
-        model_name = getattr(response, "model", kwargs.get("model", "unknown"))
-        embeddings, dim, count = _parse_embedding_response(response)
-        usage = getattr(response, "usage", None)
-        prompt_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0
-        total_tokens = (
-            getattr(usage, "total_tokens", prompt_tokens) if usage else prompt_tokens
-        )
-
         tracer.add_embedding_step_to_trace(
-            name="OpenAI Embedding",
-            end_time=end_time,
-            inputs={"input": kwargs.get("input")},
-            output=embeddings,
-            latency=(end_time - start_time) * 1000,
-            tokens=total_tokens,
-            prompt_tokens=prompt_tokens,
-            model=model_name,
-            model_parameters=_get_embedding_model_parameters(kwargs),
-            embedding_dimensions=dim,
-            embedding_count=count,
-            raw_output=(
-                response.model_dump()
-                if hasattr(response, "model_dump")
-                else str(response)
-            ),
-            provider="OpenAI",
-            id=inference_id,
-            metadata={"provider": "OpenAI"},
+            **build_embedding_step_kwargs(
+                response,
+                kwargs,
+                start_time,
+                end_time,
+                name="OpenAI Embedding",
+                provider="OpenAI",
+                inference_id=inference_id,
+            )
         )
     except Exception as e:
         logger.error(

src/openlayer/lib/integrations/bedrock_tracer.py

@@ -5,7 +5,7 @@
 import logging
 import time
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Iterator, List, Optional, Tuple, Union
 
 from botocore.response import StreamingBody
 
@@ -237,7 +237,7 @@ def _parse_embedding_input(body_data: Dict[str, Any], model_id: str) -> Dict[str
 
 def _parse_embedding_output(
     response_data: Dict[str, Any], model_id: str
-) -> tuple:
+) -> Tuple[Union[List[float], List[List[float]]], int, int]:
     """Returns (embeddings, dimensions, count)."""
     if model_id.startswith("amazon.titan-embed"):
         emb = response_data.get("embedding", [])

src/openlayer/lib/integrations/litellm_tracer.py

@@ -18,6 +18,7 @@
 
 from ..tracing import tracer
 from ..tracing import enums as tracer_enums
+from ._openai_embedding_common import build_embedding_step_kwargs
 
 logger = logging.getLogger(__name__)
 
@@ -367,40 +368,38 @@ def handle_embedding(
     try:
         model_name = kwargs.get("model", getattr(response, "model", "unknown"))
         provider = detect_provider_from_response(response, model_name)
-        embeddings, dim, count = _parse_embedding_response(response)
-        usage_data = extract_usage_from_response(response)
         extra_metadata = extract_litellm_metadata(response, model_name)
-        cost = extra_metadata.get("cost", None)
-
-        prompt_tokens = usage_data.get("prompt_tokens") or 0
-        total_tokens = usage_data.get("total_tokens") or prompt_tokens
+        usage_data = extract_usage_from_response(response)
 
-        tracer.add_embedding_step_to_trace(
+        step_kwargs = build_embedding_step_kwargs(
+            response,
+            kwargs,
+            start_time,
+            end_time,
             name="LiteLLM Embedding",
-            end_time=end_time,
-            inputs={"input": kwargs.get("input")},
-            output=embeddings,
-            latency=(end_time - start_time) * 1000,
-            tokens=total_tokens,
-            prompt_tokens=prompt_tokens,
-            model=model_name,
-            model_parameters=_get_embedding_model_parameters(kwargs),
-            embedding_dimensions=dim,
-            embedding_count=count,
-            raw_output=(
-                response.model_dump()
-                if hasattr(response, "model_dump")
-                else str(response)
-            ),
             provider=provider,
-            cost=cost,
-            id=inference_id,
-            metadata={
-                "provider": provider,
-                "litellm_model": model_name,
-                **extra_metadata,
-            },
+            inference_id=inference_id,
         )
+
+        # LiteLLM-specific overlays: usage uses LiteLLM's normalized dict, extra
+        # connection params, response cost, and provider metadata.
+        prompt_tokens = usage_data.get("prompt_tokens") or 0
+        step_kwargs["prompt_tokens"] = prompt_tokens
+        step_kwargs["tokens"] = usage_data.get("total_tokens") or prompt_tokens
+        step_kwargs["model_parameters"] = {
+            **step_kwargs["model_parameters"],
+            "timeout": kwargs.get("timeout"),
+            "api_base": kwargs.get("api_base"),
+            "api_version": kwargs.get("api_version"),
+        }
+        step_kwargs["cost"] = extra_metadata.get("cost", None)
+        step_kwargs["metadata"] = {
+            **step_kwargs["metadata"],
+            "litellm_model": model_name,
+            **extra_metadata,
+        }
+
+        tracer.add_embedding_step_to_trace(**step_kwargs)
     except Exception as e:
         logger.error(
             "Failed to trace the LiteLLM embedding request with Openlayer. %s", e
@@ -409,38 +408,6 @@ def handle_embedding(
     return response
 
 
-def _parse_embedding_response(response: Any) -> tuple:
-    """Returns (embeddings, dimensions, count). Mirrors OpenAI EmbeddingResponse."""
-    try:
-        data = getattr(response, "data", None)
-        if data is None and isinstance(response, dict):
-            data = response.get("data", [])
-        if not data:
-            return [], 0, 0
-        embeddings = [
-            item["embedding"] if isinstance(item, dict) else item.embedding
-            for item in data
-        ]
-        if not embeddings:
-            return [], 0, 0
-        if len(embeddings) == 1:
-            return embeddings[0], len(embeddings[0]), 1
-        return embeddings, len(embeddings[0]), len(embeddings)
-    except Exception:
-        return [], 0, 0
-
-
-def _get_embedding_model_parameters(kwargs: Dict[str, Any]) -> Dict[str, Any]:
-    return {
-        "dimensions": kwargs.get("dimensions"),
-        "encoding_format": kwargs.get("encoding_format"),
-        "user": kwargs.get("user"),
-        "timeout": kwargs.get("timeout"),
-        "api_base": kwargs.get("api_base"),
-        "api_version": kwargs.get("api_version"),
-    }
-
-
 def get_model_parameters(kwargs: Dict[str, Any]) -> Dict[str, Any]:
     """Gets the model parameters from the kwargs."""
     return {

src/openlayer/lib/integrations/openai_tracer.py

@@ -27,6 +27,7 @@
     ImageContent,
     TextContent,
 )
+from ._openai_embedding_common import build_embedding_step_kwargs
 
 logger = logging.getLogger(__name__)
 
@@ -1635,46 +1636,21 @@ def handle_embedding(
     **kwargs,
 ) -> Any:
     """Trace a sync OpenAI client.embeddings.create() call."""
-    from ._openai_embedding_common import (
-        get_embedding_model_parameters as _get_embedding_model_parameters,
-    )
-    from ._openai_embedding_common import (
-        parse_embedding_response as _parse_embedding_response,
-    )
-
     start_time = time.time()
     response = original_func(*args, **kwargs)
     end_time = time.time()
 
     try:
-        model_name = getattr(response, "model", kwargs.get("model", "unknown"))
-        embeddings, dim, count = _parse_embedding_response(response)
-        usage = getattr(response, "usage", None)
-        prompt_tokens = getattr(usage, "prompt_tokens", 0) if usage else 0
-        total_tokens = (
-            getattr(usage, "total_tokens", prompt_tokens) if usage else prompt_tokens
-        )
-
         tracer.add_embedding_step_to_trace(
-            name="OpenAI Embedding",
-            end_time=end_time,
-            inputs={"input": kwargs.get("input")},
-            output=embeddings,
-            latency=(end_time - start_time) * 1000,
-            tokens=total_tokens,
-            prompt_tokens=prompt_tokens,
-            model=model_name,
-            model_parameters=_get_embedding_model_parameters(kwargs),
-            embedding_dimensions=dim,
-            embedding_count=count,
-            raw_output=(
-                response.model_dump()
-                if hasattr(response, "model_dump")
-                else str(response)
-            ),
-            provider="OpenAI",
-            id=inference_id,
-            metadata={"provider": "OpenAI"},
+            **build_embedding_step_kwargs(
+                response,
+                kwargs,
+                start_time,
+                end_time,
+                name="OpenAI Embedding",
+                provider="OpenAI",
+                inference_id=inference_id,
+            )
         )
     except Exception as e:
         logger.error(