lmnr 0.4.53.dev0__py3-none-any.whl → 0.7.26__py3-none-any.whl

lmnr/opentelemetry_lib/litellm/__init__.py

@@ -0,0 +1,685 @@
+"""LiteLLM callback logger for Laminar"""
+
+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 (
+    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", "responses", "aresponses"]
+
+# Try to import the necessary LiteLLM components and gracefully handle ImportError
+try:
+    if not is_package_installed("litellm"):
+        raise ImportError("LiteLLM is not installed")
+
+    from litellm.integrations.custom_batch_logger import CustomBatchLogger
+
+    class LaminarLiteLLMCallback(CustomBatchLogger):
+        """Custom LiteLLM logger that sends logs to Laminar via OpenTelemetry spans
+
+        Usage:
+            import litellm
+            from lmnr import Laminar, LaminarLiteLLMCallback
+
+            # make sure this comes first
+            Laminar.initialize()
+
+            # Add the logger to LiteLLM callbacks
+            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")
+            return TracerWrapper().get_tracer()
+
+        def log_success_event(
+            self, kwargs, response_obj, start_time: datetime, end_time: datetime
+        ):
+            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
+                )
+            except Exception as e:
+                logger.error(f"Error in log_success_event: {e}")
+
+        def log_failure_event(
+            self, kwargs, response_obj, start_time: datetime, end_time: datetime
+        ):
+            if kwargs.get("call_type") not in SUPPORTED_CALL_TYPES:
+                return
+            try:
+                self._create_span(
+                    kwargs, response_obj, start_time, end_time, is_success=False
+                )
+            except Exception as e:
+                logger.error(f"Error in log_failure_event: {e}")
+
+        async def async_log_success_event(
+            self, kwargs, response_obj, start_time: datetime, end_time: datetime
+        ):
+            self.log_success_event(kwargs, response_obj, start_time, end_time)
+
+        async def async_log_failure_event(
+            self, kwargs, response_obj, start_time: datetime, end_time: datetime
+        ):
+            self.log_failure_event(kwargs, response_obj, start_time, end_time)
+
+        def _create_span(
+            self,
+            kwargs,
+            response_obj,
+            start_time: datetime,
+            end_time: datetime,
+            is_success: bool,
+        ):
+            """Create an OpenTelemetry span for the LiteLLM call"""
+            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,
+                start_time=int(start_time.timestamp() * 1e9),
+                attributes={
+                    "lmnr.internal.provider": "litellm",
+                },
+                context=get_current_context(),
+            )
+            try:
+                model = kwargs.get("model", "unknown")
+                if kwargs.get("custom_llm_provider"):
+                    set_span_attribute(
+                        span, "gen_ai.system", kwargs["custom_llm_provider"]
+                    )
+
+                messages = kwargs.get("messages", [])
+                self._process_input_messages(span, messages)
+
+                tools = kwargs.get("tools", [])
+                self._process_request_tool_definitions(span, tools)
+
+                set_span_attribute(span, "gen_ai.request.model", model)
+
+                # Add more attributes from kwargs
+                if "temperature" in kwargs:
+                    set_span_attribute(
+                        span, "gen_ai.request.temperature", kwargs["temperature"]
+                    )
+                if "max_tokens" in kwargs:
+                    set_span_attribute(
+                        span, "gen_ai.request.max_tokens", kwargs["max_tokens"]
+                    )
+                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"):
+                        self._process_success_response(
+                            span,
+                            kwargs.get("complete_streaming_response"),
+                        )
+                    else:
+                        self._process_success_response(span, response_obj)
+                else:
+                    span.set_status(Status(StatusCode.ERROR))
+                    if isinstance(response_obj, Exception):
+                        attributes = get_event_attributes_from_context()
+                        span.record_exception(response_obj, attributes=attributes)
+
+            except Exception as 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))
+
+        def _process_input_messages(self, span, messages):
+            """Process and set message attributes on the span"""
+            if not isinstance(messages, list):
+                return
+
+            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
+
+                elif block_dict.get("type") == "computer_call_output":
+                    set_span_attribute(
+                        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),
+                    )
+                    set_span_attribute(
+                        span,
+                        f"{GEN_AI_PROMPT}.{prompt_index}.tool_calls.0.id",
+                        block_dict.get("call_id"),
+                    )
+                    set_span_attribute(
+                        span,
+                        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"""
+            if not isinstance(tools, list):
+                return
+
+            for i, tool in enumerate(tools):
+                tool_dict = model_as_dict(tool)
+                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",
+                    function_name,
+                )
+                set_span_attribute(
+                    span,
+                    f"llm.request.functions.{i}.description",
+                    function_description,
+                )
+                set_span_attribute(
+                    span,
+                    f"llm.request.functions.{i}.parameters",
+                    json.dumps(function_parameters),
+                )
+
+        def _process_response_usage(self, span, usage):
+            """Process and set usage attributes on the span"""
+            usage_dict = model_as_dict(usage)
+            if (
+                not usage_dict.get("prompt_tokens")
+                and not usage_dict.get("completion_tokens")
+                and not usage_dict.get("total_tokens")
+            ):
+                return
+
+            set_span_attribute(
+                span, "gen_ai.usage.input_tokens", usage_dict.get("prompt_tokens")
+            )
+            set_span_attribute(
+                span, "gen_ai.usage.output_tokens", usage_dict.get("completion_tokens")
+            )
+            set_span_attribute(
+                span, "llm.usage.total_tokens", usage_dict.get("total_tokens")
+            )
+
+            if usage_dict.get("prompt_tokens_details"):
+                details = usage_dict.get("prompt_tokens_details", {})
+                details = model_as_dict(details)
+                if details.get("cached_tokens"):
+                    set_span_attribute(
+                        span,
+                        "gen_ai.usage.cache_read_input_tokens",
+                        details.get("cached_tokens"),
+                    )
+                # TODO: add audio/image/text token details
+            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"""
+            attr_prefix = "completion" if is_response else "prompt"
+            if not isinstance(tool_calls, list):
+                return
+
+            for j, tool_call in enumerate(tool_calls):
+                tool_call_dict = model_as_dict(tool_call)
+
+                tool_name = tool_call_dict.get(
+                    "name", tool_call_dict.get("function", {}).get("name", "")
+                )
+                set_span_attribute(
+                    span,
+                    f"gen_ai.{attr_prefix}.{choice_index}.tool_calls.{j}.name",
+                    tool_name,
+                )
+
+                call_id = tool_call_dict.get("id", "")
+                set_span_attribute(
+                    span,
+                    f"gen_ai.{attr_prefix}.{choice_index}.tool_calls.{j}.id",
+                    call_id,
+                )
+
+                tool_arguments = tool_call_dict.get(
+                    "arguments", tool_call_dict.get("function", {}).get("arguments", "")
+                )
+                if isinstance(tool_arguments, str):
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.{attr_prefix}.{choice_index}.tool_calls.{j}.arguments",
+                        tool_arguments,
+                    )
+                else:
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.{attr_prefix}.{choice_index}.tool_calls.{j}.arguments",
+                        json.dumps(model_as_dict(tool_arguments)),
+                    )
+
+        def _process_response_choices(self, span, choices):
+            """Process and set choice attributes on the span"""
+            if not isinstance(choices, list):
+                return
+
+            for i, choice in enumerate(choices):
+                choice_dict = model_as_dict(choice)
+                message = choice_dict.get("message", choice_dict)
+
+                role = message.get("role", "unknown")
+                set_span_attribute(span, f"gen_ai.completion.{i}.role", role)
+
+                tool_calls = message.get("tool_calls", [])
+                self._process_tool_calls(span, tool_calls, i, is_response=True)
+
+                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):
+                    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(reasoning_content + content),
+                    )
+                else:
+                    set_span_attribute(
+                        span,
+                        f"gen_ai.completion.{i}.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)
+            set_span_attribute(span, "gen_ai.response.id", response_dict.get("id"))
+            set_span_attribute(
+                span, "gen_ai.response.model", response_dict.get("model")
+            )
+
+            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"):
+                set_span_attribute(
+                    span,
+                    "gen_ai.usage.cache_creation_input_tokens",
+                    response_dict.get("cache_creation_input_tokens"),
+                )
+            if response_dict.get("cache_read_input_tokens"):
+                set_span_attribute(
+                    span,
+                    "gen_ai.usage.cache_read_input_tokens",
+                    response_dict.get("cache_read_input_tokens"),
+                )
+
+            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}")
+
+    # Create a no-op logger when LiteLLM is not available
+    class LaminarLiteLLMCallback:
+        """No-op logger when LiteLLM is not available"""
+
+        def __init__(self, **kwargs):
+            logger.warning(
+                "LiteLLM is not installed. Install with: pip install litellm"
+            )
+
+        def log_success_event(self, *args, **kwargs):
+            pass
+
+        def log_failure_event(self, *args, **kwargs):
+            pass
+
+        async def async_log_success_event(self, *args, **kwargs):
+            pass
+
+        async def async_log_failure_event(self, *args, **kwargs):
+            pass