lmnr 0.6.16__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/opentelemetry_lib/litellm/__init__.py

@@ -3,16 +3,28 @@
 import json
 from datetime import datetime
 
+from opentelemetry.semconv._incubating.attributes.gen_ai_attributes import GEN_AI_PROMPT
 from opentelemetry.trace import SpanKind, Status, StatusCode, Tracer
-from lmnr.opentelemetry_lib.litellm.utils import model_as_dict, set_span_attribute
+from lmnr.opentelemetry_lib.litellm.utils import (
+    get_tool_definition,
+    is_validator_iterator,
+    model_as_dict,
+    set_span_attribute,
+)
 from lmnr.opentelemetry_lib.tracing import TracerWrapper
 
+from lmnr.opentelemetry_lib.tracing.context import (
+    get_current_context,
+    get_event_attributes_from_context,
+)
+from lmnr.opentelemetry_lib.tracing.attributes import ASSOCIATION_PROPERTIES
 from lmnr.opentelemetry_lib.utils.package_check import is_package_installed
 from lmnr.sdk.log import get_default_logger
+from lmnr.sdk.utils import json_dumps
 
 logger = get_default_logger(__name__)
 
-SUPPORTED_CALL_TYPES = ["completion", "acompletion"]
+SUPPORTED_CALL_TYPES = ["completion", "acompletion", "responses", "aresponses"]
 
 # Try to import the necessary LiteLLM components and gracefully handle ImportError
 try:
@@ -35,11 +47,29 @@ try:
             litellm.callbacks = [LaminarLiteLLMCallback()]
         """
 
+        logged_openai_responses: set[str]
+
         def __init__(self, **kwargs):
             super().__init__(**kwargs)
             if not hasattr(TracerWrapper, "instance") or TracerWrapper.instance is None:
                 raise ValueError("Laminar must be initialized before LiteLLM callback")
 
+            self.logged_openai_responses = set()
+            if is_package_installed("openai"):
+                from lmnr.opentelemetry_lib.opentelemetry.instrumentation.openai import (
+                    OpenAIInstrumentor,
+                )
+
+                openai_instrumentor = OpenAIInstrumentor()
+                if (
+                    openai_instrumentor
+                    and openai_instrumentor.is_instrumented_by_opentelemetry
+                ):
+                    logger.debug(
+                        "Disabling OpenTelemetry instrumentation for OpenAI to avoid double-instrumentation of LiteLLM."
+                    )
+                    openai_instrumentor.uninstrument()
+
         def _get_tracer(self) -> Tracer:
             if not hasattr(TracerWrapper, "instance") or TracerWrapper.instance is None:
                 raise ValueError("Laminar must be initialized before LiteLLM callback")
@@ -50,6 +80,14 @@ try:
         ):
             if kwargs.get("call_type") not in SUPPORTED_CALL_TYPES:
                 return
+            if kwargs.get("call_type") in ["responses", "aresponses"]:
+                # responses API may be called multiple times with the same response_obj
+                response_id = getattr(response_obj, "id", None)
+                if response_id in self.logged_openai_responses:
+                    return
+                if response_id:
+                    self.logged_openai_responses.add(response_id)
+            self.logged_openai_responses.add(response_obj.id)
             try:
                 self._create_span(
                     kwargs, response_obj, start_time, end_time, is_success=True
@@ -88,12 +126,18 @@ try:
             is_success: bool,
         ):
             """Create an OpenTelemetry span for the LiteLLM call"""
-            span_name = "litellm.completion"
+            call_type = kwargs.get("call_type", "completion")
+            if call_type == "aresponses":
+                call_type = "responses"
+            if call_type == "acompletion":
+                call_type = "completion"
+            span_name = f"litellm.{call_type}"
             try:
                 tracer = self._get_tracer()
             except Exception as e:
                 logger.error(f"Error getting tracer: {e}")
                 return
+
             span = tracer.start_span(
                 span_name,
                 kind=SpanKind.CLIENT,
@@ -101,6 +145,7 @@ try:
                 attributes={
                     "lmnr.internal.provider": "litellm",
                 },
+                context=get_current_context(),
             )
             try:
                 model = kwargs.get("model", "unknown")
@@ -129,6 +174,52 @@ try:
                 if "top_p" in kwargs:
                     set_span_attribute(span, "gen_ai.request.top_p", kwargs["top_p"])
 
+                metadata = (
+                    kwargs.get("litellm_params").get(
+                        "metadata", kwargs.get("metadata", {})
+                    )
+                    or {}
+                )
+                tags = metadata.get("tags", [])
+                if isinstance(tags, str):
+                    try:
+                        tags = json.loads(tags)
+                    except Exception:
+                        pass
+                if (
+                    tags
+                    and isinstance(tags, (list, tuple, set))
+                    and all(isinstance(tag, str) for tag in tags)
+                ):
+                    span.set_attribute(f"{ASSOCIATION_PROPERTIES}.tags", tags)
+
+                user_id = metadata.get("user_id")
+                if user_id:
+                    span.set_attribute(f"{ASSOCIATION_PROPERTIES}.user_id", user_id)
+
+                session_id = metadata.get("session_id")
+                if session_id:
+                    span.set_attribute(
+                        f"{ASSOCIATION_PROPERTIES}.session_id", session_id
+                    )
+
+                optional_params = kwargs.get("optional_params") or {}
+                if not optional_params:
+                    hidden_params = metadata.get("hidden_params") or {}
+                    optional_params = hidden_params.get("optional_params") or {}
+                response_format = optional_params.get("response_format")
+                if (
+                    response_format
+                    and isinstance(response_format, dict)
+                    and response_format.get("type") == "json_schema"
+                ):
+                    schema = (response_format.get("json_schema") or {}).get("schema")
+                    if schema:
+                        span.set_attribute(
+                            "gen_ai.request.structured_output_schema",
+                            json_dumps(schema),
+                        )
+
                 if is_success:
                     span.set_status(Status(StatusCode.OK))
                     if kwargs.get("complete_streaming_response"):
@@ -141,10 +232,12 @@ try:
                 else:
                     span.set_status(Status(StatusCode.ERROR))
                     if isinstance(response_obj, Exception):
-                        span.record_exception(response_obj)
+                        attributes = get_event_attributes_from_context()
+                        span.record_exception(response_obj, attributes=attributes)
 
             except Exception as e:
-                span.record_exception(e)
+                attributes = get_event_attributes_from_context()
+                span.record_exception(e, attributes=attributes)
                 logger.error(f"Error in Laminar LiteLLM instrumentation: {e}")
             finally:
                 span.end(int(end_time.timestamp() * 1e9))
@@ -154,35 +247,107 @@ try:
             if not isinstance(messages, list):
                 return
 
-            for i, message in enumerate(messages):
-                message_dict = model_as_dict(message)
-                role = message_dict.get("role", "unknown")
-                set_span_attribute(span, f"gen_ai.prompt.{i}.role", role)
-
-                tool_calls = message_dict.get("tool_calls", [])
-                self._process_tool_calls(span, tool_calls, i, is_response=False)
+            prompt_index = 0
+            for item in messages:
+                block_dict = model_as_dict(item)
+                if block_dict.get("type", "message") == "message":
+                    tool_calls = block_dict.get("tool_calls", [])
+                    self._process_tool_calls(
+                        span, tool_calls, prompt_index, is_response=False
+                    )
+                    content = block_dict.get("content")
+                    if is_validator_iterator(content):
+                        # Have not been able to catch this in the wild, but keeping
+                        # just in case, as raw OpenAI responses do that
+                        content = [self._process_content_part(part) for part in content]
+                    try:
+                        stringified_content = (
+                            content if isinstance(content, str) else json_dumps(content)
+                        )
+                    except Exception:
+                        stringified_content = (
+                            str(content) if content is not None else ""
+                        )
+                    set_span_attribute(
+                        span,
+                        f"{GEN_AI_PROMPT}.{prompt_index}.content",
+                        stringified_content,
+                    )
+                    set_span_attribute(
+                        span,
+                        f"{GEN_AI_PROMPT}.{prompt_index}.role",
+                        block_dict.get("role"),
+                    )
+                    prompt_index += 1
 
-                content = message_dict.get("content", "")
-                if content is None:
-                    continue
-                if isinstance(content, str):
-                    set_span_attribute(span, f"gen_ai.prompt.{i}.content", content)
-                elif isinstance(content, list):
+                elif block_dict.get("type") == "computer_call_output":
                     set_span_attribute(
-                        span, f"gen_ai.prompt.{i}.content", json.dumps(content)
+                        span,
+                        f"{GEN_AI_PROMPT}.{prompt_index}.role",
+                        "computer_call_output",
+                    )
+                    output_image_url = block_dict.get("output", {}).get("image_url")
+                    if output_image_url:
+                        set_span_attribute(
+                            span,
+                            f"{GEN_AI_PROMPT}.{prompt_index}.content",
+                            json.dumps(
+                                [
+                                    {
+                                        "type": "image_url",
+                                        "image_url": {"url": output_image_url},
+                                    }
+                                ]
+                            ),
+                        )
+                    prompt_index += 1
+                elif block_dict.get("type") == "computer_call":
+                    set_span_attribute(
+                        span, f"{GEN_AI_PROMPT}.{prompt_index}.role", "assistant"
+                    )
+                    call_content = {}
+                    if block_dict.get("id"):
+                        call_content["id"] = block_dict.get("id")
+                    if block_dict.get("action"):
+                        call_content["action"] = block_dict.get("action")
+                    set_span_attribute(
+                        span,
+                        f"{GEN_AI_PROMPT}.{prompt_index}.tool_calls.0.arguments",
+                        json.dumps(call_content),
                     )
-                else:
                     set_span_attribute(
                         span,
-                        f"gen_ai.prompt.{i}.content",
-                        json.dumps(model_as_dict(content)),
+                        f"{GEN_AI_PROMPT}.{prompt_index}.tool_calls.0.id",
+                        block_dict.get("call_id"),
                     )
-                if role == "tool":
                     set_span_attribute(
                         span,
-                        f"gen_ai.prompt.{i}.tool_call_id",
-                        message_dict.get("tool_call_id"),
+                        f"{GEN_AI_PROMPT}.{prompt_index}.tool_calls.0.name",
+                        "computer_call",
                     )
+                    prompt_index += 1
+                elif block_dict.get("type") == "reasoning":
+                    reasoning_summary = block_dict.get("summary")
+                    if reasoning_summary and isinstance(reasoning_summary, list):
+                        processed_chunks = [
+                            {"type": "text", "text": chunk.get("text")}
+                            for chunk in reasoning_summary
+                            if isinstance(chunk, dict)
+                            and chunk.get("type") == "summary_text"
+                        ]
+                        set_span_attribute(
+                            span,
+                            f"{GEN_AI_PROMPT}.{prompt_index}.reasoning",
+                            json_dumps(processed_chunks),
+                        )
+                        set_span_attribute(
+                            span,
+                            f"{GEN_AI_PROMPT}.{prompt_index}.role",
+                            "assistant",
+                        )
+                    # reasoning is followed by other content parts in the same messge,
+                    # so we don't increment the prompt index
+                # TODO: handle other block types
 
         def _process_request_tool_definitions(self, span, tools):
             """Process and set tool definitions attributes on the span"""
@@ -191,14 +356,10 @@ try:
 
             for i, tool in enumerate(tools):
                 tool_dict = model_as_dict(tool)
-                if tool_dict.get("type") != "function":
-                    # TODO: parse other tool types
-                    continue
-
-                function_dict = tool_dict.get("function", {})
-                function_name = function_dict.get("name", "")
-                function_description = function_dict.get("description", "")
-                function_parameters = function_dict.get("parameters", {})
+                tool_definition = get_tool_definition(tool_dict)
+                function_name = tool_definition.get("name")
+                function_description = tool_definition.get("description")
+                function_parameters = tool_definition.get("parameters")
                 set_span_attribute(
                     span,
                     f"llm.request.functions.{i}.name",
@@ -245,7 +406,15 @@ try:
                         details.get("cached_tokens"),
                     )
                 # TODO: add audio/image/text token details
-            # TODO: add completion tokens details (reasoning tokens)
+            if usage_dict.get("completion_tokens_details"):
+                details = usage_dict.get("completion_tokens_details", {})
+                details = model_as_dict(details)
+                if details.get("reasoning_tokens"):
+                    set_span_attribute(
+                        span,
+                        "gen_ai.usage.reasoning_tokens",
+                        details.get("reasoning_tokens"),
+                    )
 
         def _process_tool_calls(self, span, tool_calls, choice_index, is_response=True):
             """Process and set tool call attributes on the span"""
@@ -306,19 +475,160 @@ try:
                 content = message.get("content", "")
                 if content is None:
                     continue
+                reasoning_content = message.get("reasoning_content")
+                if reasoning_content:
+                    if isinstance(reasoning_content, str):
+                        reasoning_content = [
+                            {
+                                "type": "text",
+                                "text": reasoning_content,
+                            }
+                        ]
+                    elif not isinstance(reasoning_content, list):
+                        reasoning_content = [
+                            {
+                                "type": "text",
+                                "text": str(reasoning_content),
+                            }
+                        ]
+                else:
+                    reasoning_content = []
                 if isinstance(content, str):
-                    set_span_attribute(span, f"gen_ai.completion.{i}.content", content)
+                    if reasoning_content:
+                        set_span_attribute(
+                            span,
+                            f"gen_ai.completion.{i}.content",
+                            json.dumps(
+                                reasoning_content
+                                + [
+                                    {
+                                        "type": "text",
+                                        "text": content,
+                                    }
+                                ]
+                            ),
+                        )
+                    else:
+                        set_span_attribute(
+                            span,
+                            f"gen_ai.completion.{i}.content",
+                            content,
+                        )
                 elif isinstance(content, list):
                     set_span_attribute(
-                        span, f"gen_ai.completion.{i}.content", json.dumps(content)
+                        span,
+                        f"gen_ai.completion.{i}.content",
+                        json.dumps(reasoning_content + content),
                     )
                 else:
                     set_span_attribute(
                         span,
                         f"gen_ai.completion.{i}.content",
-                        json.dumps(model_as_dict(content)),
+                        json.dumps(reasoning_content + [model_as_dict(content)]),
                     )
 
+        def _process_content_part(self, content_part: dict) -> dict:
+            content_part_dict = model_as_dict(content_part)
+            if content_part_dict.get("type") == "output_text":
+                return {"type": "text", "text": content_part_dict.get("text")}
+            return content_part_dict
+
+        def _process_response_output(self, span, output):
+            """Response of OpenAI Responses API"""
+            if not isinstance(output, list):
+                return
+            set_span_attribute(span, "gen_ai.completion.0.role", "assistant")
+            tool_call_index = 0
+            for block in output:
+                block_dict = model_as_dict(block)
+                if block_dict.get("type") == "message":
+                    content = block_dict.get("content")
+                    if content is None:
+                        continue
+                    if isinstance(content, str):
+                        set_span_attribute(span, "gen_ai.completion.0.content", content)
+                    elif isinstance(content, list):
+                        set_span_attribute(
+                            span,
+                            "gen_ai.completion.0.content",
+                            json_dumps(
+                                [self._process_content_part(part) for part in content]
+                            ),
+                        )
+                if block_dict.get("type") == "function_call":
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.id",
+                        block_dict.get("id"),
+                    )
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.name",
+                        block_dict.get("name"),
+                    )
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.arguments",
+                        block_dict.get("arguments"),
+                    )
+                    tool_call_index += 1
+                elif block_dict.get("type") == "file_search_call":
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.id",
+                        block_dict.get("id"),
+                    )
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.name",
+                        "file_search_call",
+                    )
+                    tool_call_index += 1
+                elif block_dict.get("type") == "web_search_call":
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.id",
+                        block_dict.get("id"),
+                    )
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.name",
+                        "web_search_call",
+                    )
+                    tool_call_index += 1
+                elif block_dict.get("type") == "computer_call":
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.id",
+                        block_dict.get("call_id"),
+                    )
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.name",
+                        "computer_call",
+                    )
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.0.tool_calls.{tool_call_index}.arguments",
+                        json_dumps(block_dict.get("action")),
+                    )
+                    tool_call_index += 1
+                elif block_dict.get("type") == "reasoning":
+                    reasoning_summary = block_dict.get("summary")
+                    if reasoning_summary and isinstance(reasoning_summary, list):
+                        processed_chunks = [
+                            {"type": "text", "text": chunk.get("text")}
+                            for chunk in reasoning_summary
+                            if isinstance(chunk, dict)
+                            and chunk.get("type") == "summary_text"
+                        ]
+                        set_span_attribute(
+                            span,
+                            "gen_ai.completion.0.reasoning",
+                            json_dumps(processed_chunks),
+                        )
+                # TODO: handle other block types, in particular other calls
+
         def _process_success_response(self, span, response_obj):
             """Process successful response attributes"""
             response_dict = model_as_dict(response_obj)
@@ -327,7 +637,9 @@ try:
                 span, "gen_ai.response.model", response_dict.get("model")
             )
 
-            if response_dict.get("usage"):
+            if getattr(response_obj, "usage", None):
+                self._process_response_usage(span, getattr(response_obj, "usage", None))
+            elif response_dict.get("usage"):
                 self._process_response_usage(span, response_dict.get("usage"))
 
             if response_dict.get("cache_creation_input_tokens"):
@@ -345,6 +657,8 @@ try:
 
             if response_dict.get("choices"):
                 self._process_response_choices(span, response_dict.get("choices"))
+            elif response_dict.get("output"):
+                self._process_response_output(span, response_dict.get("output"))
 
 except ImportError as e:
     logger.debug(f"LiteLLM callback unavailable: {e}")

lmnr/opentelemetry_lib/litellm/utils.py

@@ -1,6 +1,14 @@
+import re
 from pydantic import BaseModel
 from opentelemetry.sdk.trace import Span
 from opentelemetry.util.types import AttributeValue
+from typing_extensions import TypedDict
+
+
+class ToolDefinition(TypedDict):
+    name: str | None
+    description: str | None
+    parameters: dict | None
 
 
 def model_as_dict(model: BaseModel | dict) -> dict:
@@ -16,3 +24,77 @@ def set_span_attribute(span: Span, key: str, value: AttributeValue | None):
     if value is None or value == "":
         return
     span.set_attribute(key, value)
+
+
+def get_tool_definition(tool: dict) -> ToolDefinition:
+    parameters = None
+    description = None
+    name = (tool.get("function") or {}).get("name") or tool.get("name")
+    if tool.get("type") == "function":
+        function = tool.get("function") or {}
+        parameters = function.get("parameters") or tool.get("parameters")
+        description = function.get("description") or tool.get("description")
+    elif isinstance(tool.get("type"), str) and tool.get("type").startswith("computer"):
+        # Anthropic beta computer tools
+        # https://docs.anthropic.com/en/docs/agents-and-tools/tool-use/computer-use-tool
+
+        # OpenAI computer use API
+        # https://platform.openai.com/docs/guides/tools-computer-use
+        if not name:
+            name = tool.get("type")
+
+        parameters = {}
+        tool_parameters = (tool.get("function") or {}).get("parameters") or {}
+        # Anthropic
+        display_width_px = tool_parameters.get("display_width_px") or tool.get(
+            "display_width_px"
+        )
+        display_height_px = tool_parameters.get("display_height_px") or tool.get(
+            "display_height_px"
+        )
+        display_number = tool_parameters.get("display_number") or tool.get(
+            "display_number"
+        )
+        if display_width_px:
+            parameters["display_width_px"] = display_width_px
+        if display_height_px:
+            parameters["display_height_px"] = display_height_px
+        if display_number:
+            parameters["display_number"] = display_number
+        # OpenAI
+        display_width = tool_parameters.get("display_width") or tool.get(
+            "display_width"
+        )
+        display_height = tool_parameters.get("display_height") or tool.get(
+            "display_height"
+        )
+        environment = tool_parameters.get("environment") or tool.get("environment")
+        if display_width:
+            parameters["display_width"] = display_width
+        if display_height:
+            parameters["display_height"] = tool.get("display_height")
+        if environment:  # Literal['browser', 'mac', 'windows', 'ubuntu']
+            parameters["environment"] = environment
+    else:
+        # Some versions of LiteLLM (around 1.69.0) flatten the tool definition in
+        # anthropic style, not the OpenAI style as they do with other tool types.
+        function = tool.get("function") or tool
+        parameters = function.get("parameters") or function.get("input_schema")
+        description = function.get("description") or ""
+
+    return ToolDefinition(
+        name=name,
+        description=description,
+        parameters=parameters,
+    )
+
+
+def is_validator_iterator(content):
+    """
+    Some OpenAI objects contain fields typed as Iterable, which pydantic
+    internally converts to a ValidatorIterator, and they cannot be trivially
+    serialized without consuming the iterator to, for example, a list.
+
+    See: https://github.com/pydantic/pydantic/issues/9541#issuecomment-2189045051
+    """
+    return re.search(r"pydantic.*ValidatorIterator'>$", str(type(content)))