litellm/integrations/opik/opik.py

"""
Opik Logger that logs LLM events to an Opik server
"""

import asyncio
import json
import traceback
from typing import Dict, List

from litellm._logging import verbose_logger
from litellm.integrations.custom_batch_logger import CustomBatchLogger
from litellm.llms.custom_httpx.http_handler import (
    _get_httpx_client,
    get_async_httpx_client,
    httpxSpecialProvider,
)

from .utils import (
    create_usage_object,
    create_uuid7,
    get_opik_config_variable,
    get_traces_and_spans_from_payload,
)


class OpikLogger(CustomBatchLogger):
    """
    Opik Logger for logging events to an Opik Server
    """

    def __init__(self, **kwargs):
        self.async_httpx_client = get_async_httpx_client(
            llm_provider=httpxSpecialProvider.LoggingCallback
        )
        self.sync_httpx_client = _get_httpx_client()

        self.opik_project_name = get_opik_config_variable(
            "project_name",
            user_value=kwargs.get("project_name", None),
            default_value="Default Project",
        )

        opik_base_url = get_opik_config_variable(
            "url_override",
            user_value=kwargs.get("url", None),
            default_value="https://www.comet.com/opik/api",
        )
        opik_api_key = get_opik_config_variable(
            "api_key", user_value=kwargs.get("api_key", None), default_value=None
        )
        opik_workspace = get_opik_config_variable(
            "workspace", user_value=kwargs.get("workspace", None), default_value=None
        )

        self.trace_url = f"{opik_base_url}/v1/private/traces/batch"
        self.span_url = f"{opik_base_url}/v1/private/spans/batch"

        self.headers = {}
        if opik_workspace:
            self.headers["Comet-Workspace"] = opik_workspace

        if opik_api_key:
            self.headers["authorization"] = opik_api_key

        self.opik_workspace = opik_workspace
        self.opik_api_key = opik_api_key
        try:
            asyncio.create_task(self.periodic_flush())
            self.flush_lock = asyncio.Lock()
        except Exception as e:
            verbose_logger.exception(
                f"OpikLogger - Asynchronous processing not initialized as we are not running in an async context {str(e)}"
            )
            self.flush_lock = None

        super().__init__(**kwargs, flush_lock=self.flush_lock)

    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
        try:
            opik_payload = self._create_opik_payload(
                kwargs=kwargs,
                response_obj=response_obj,
                start_time=start_time,
                end_time=end_time,
            )

            self.log_queue.extend(opik_payload)
            verbose_logger.debug(
                f"OpikLogger added event to log_queue - Will flush in {self.flush_interval} seconds..."
            )

            if len(self.log_queue) >= self.batch_size:
                verbose_logger.debug("OpikLogger - Flushing batch")
                await self.flush_queue()
        except Exception as e:
            verbose_logger.exception(
                f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}"
            )

    def _sync_send(self, url: str, headers: Dict[str, str], batch: Dict):
        try:
            response = self.sync_httpx_client.post(
                url=url, headers=headers, json=batch  # type: ignore
            )
            response.raise_for_status()
            if response.status_code != 204:
                raise Exception(
                    f"Response from opik API status_code: {response.status_code}, text: {response.text}"
                )
        except Exception as e:
            verbose_logger.exception(
                f"OpikLogger failed to send batch - {str(e)}\n{traceback.format_exc()}"
            )

    def log_success_event(self, kwargs, response_obj, start_time, end_time):
        try:
            opik_payload = self._create_opik_payload(
                kwargs=kwargs,
                response_obj=response_obj,
                start_time=start_time,
                end_time=end_time,
            )

            traces, spans = get_traces_and_spans_from_payload(opik_payload)
            if len(traces) > 0:
                self._sync_send(
                    url=self.trace_url, headers=self.headers, batch={"traces": traces}
                )
            if len(spans) > 0:
                self._sync_send(
                    url=self.span_url, headers=self.headers, batch={"spans": spans}
                )
        except Exception as e:
            verbose_logger.exception(
                f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}"
            )

    async def _submit_batch(self, url: str, headers: Dict[str, str], batch: Dict):
        try:
            response = await self.async_httpx_client.post(
                url=url, headers=headers, json=batch  # type: ignore
            )
            response.raise_for_status()

            if response.status_code >= 300:
                verbose_logger.error(
                    f"OpikLogger - Error: {response.status_code} - {response.text}"
                )
            else:
                verbose_logger.info(
                    f"OpikLogger - {len(self.log_queue)} Opik events submitted"
                )
        except Exception as e:
            verbose_logger.exception(f"OpikLogger failed to send batch - {str(e)}")

    def _create_opik_headers(self):
        headers = {}
        if self.opik_workspace:
            headers["Comet-Workspace"] = self.opik_workspace

        if self.opik_api_key:
            headers["authorization"] = self.opik_api_key
        return headers

    async def async_send_batch(self):
        verbose_logger.info("Calling async_send_batch")
        if not self.log_queue:
            return

        # Split the log_queue into traces and spans
        traces, spans = get_traces_and_spans_from_payload(self.log_queue)

        # Send trace batch
        if len(traces) > 0:
            await self._submit_batch(
                url=self.trace_url, headers=self.headers, batch={"traces": traces}
            )
            verbose_logger.info(f"Sent {len(traces)} traces")
        if len(spans) > 0:
            await self._submit_batch(
                url=self.span_url, headers=self.headers, batch={"spans": spans}
            )
            verbose_logger.info(f"Sent {len(spans)} spans")

    def _create_opik_payload(  # noqa: PLR0915
        self, kwargs, response_obj, start_time, end_time
    ) -> List[Dict]:
        # Get metadata
        _litellm_params = kwargs.get("litellm_params", {}) or {}
        litellm_params_metadata = _litellm_params.get("metadata", {}) or {}

        # Extract opik metadata
        litellm_opik_metadata = litellm_params_metadata.get("opik", {})
        verbose_logger.debug(
            f"litellm_opik_metadata - {json.dumps(litellm_opik_metadata, default=str)}"
        )
        project_name = litellm_opik_metadata.get("project_name", self.opik_project_name)

        # Extract trace_id and parent_span_id
        current_span_data = litellm_opik_metadata.get("current_span_data", None)
        if isinstance(current_span_data, dict):
            trace_id = current_span_data.get("trace_id", None)
            parent_span_id = current_span_data.get("id", None)
        elif current_span_data:
            trace_id = current_span_data.trace_id
            parent_span_id = current_span_data.id
        else:
            trace_id = None
            parent_span_id = None
        # Create Opik tags
        opik_tags = litellm_opik_metadata.get("tags", [])
        if kwargs.get("custom_llm_provider"):
            opik_tags.append(kwargs["custom_llm_provider"])

        # Use standard_logging_object to create metadata and input/output data
        standard_logging_object = kwargs.get("standard_logging_object", None)
        if standard_logging_object is None:
            verbose_logger.debug(
                "OpikLogger skipping event; no standard_logging_object found"
            )
            return []

        # Create input and output data
        input_data = standard_logging_object.get("messages", {})
        output_data = standard_logging_object.get("response", {})

        # Create usage object
        usage = create_usage_object(response_obj["usage"])

        # Define span and trace names
        span_name = "%s_%s_%s" % (
            response_obj.get("model", "unknown-model"),
            response_obj.get("object", "unknown-object"),
            response_obj.get("created", 0),
        )
        trace_name = response_obj.get("object", "unknown type")

        # Create metadata object, we add the opik metadata first and then
        # update it with the standard_logging_object metadata
        metadata = litellm_opik_metadata
        if "current_span_data" in metadata:
            del metadata["current_span_data"]
        metadata["created_from"] = "litellm"

        metadata.update(standard_logging_object.get("metadata", {}))
        if "call_type" in standard_logging_object:
            metadata["type"] = standard_logging_object["call_type"]
        if "status" in standard_logging_object:
            metadata["status"] = standard_logging_object["status"]
        if "response_cost" in kwargs:
            metadata["cost"] = {
                "total_tokens": kwargs["response_cost"],
                "currency": "USD",
            }
        if "response_cost_failure_debug_info" in kwargs:
            metadata["response_cost_failure_debug_info"] = kwargs[
                "response_cost_failure_debug_info"
            ]
        if "model_map_information" in standard_logging_object:
            metadata["model_map_information"] = standard_logging_object[
                "model_map_information"
            ]
        if "model" in standard_logging_object:
            metadata["model"] = standard_logging_object["model"]
        if "model_id" in standard_logging_object:
            metadata["model_id"] = standard_logging_object["model_id"]
        if "model_group" in standard_logging_object:
            metadata["model_group"] = standard_logging_object["model_group"]
        if "api_base" in standard_logging_object:
            metadata["api_base"] = standard_logging_object["api_base"]
        if "cache_hit" in standard_logging_object:
            metadata["cache_hit"] = standard_logging_object["cache_hit"]
        if "saved_cache_cost" in standard_logging_object:
            metadata["saved_cache_cost"] = standard_logging_object["saved_cache_cost"]
        if "error_str" in standard_logging_object:
            metadata["error_str"] = standard_logging_object["error_str"]
        if "model_parameters" in standard_logging_object:
            metadata["model_parameters"] = standard_logging_object["model_parameters"]
        if "hidden_params" in standard_logging_object:
            metadata["hidden_params"] = standard_logging_object["hidden_params"]

        payload = []
        if trace_id is None:
            trace_id = create_uuid7()
            verbose_logger.debug(
                f"OpikLogger creating payload for trace with id {trace_id}"
            )

            payload.append(
                {
                    "project_name": project_name,
                    "id": trace_id,
                    "name": trace_name,
                    "start_time": start_time.isoformat() + "Z",
                    "end_time": end_time.isoformat() + "Z",
                    "input": input_data,
                    "output": output_data,
                    "metadata": metadata,
                    "tags": opik_tags,
                }
            )

        span_id = create_uuid7()
        verbose_logger.debug(
            f"OpikLogger creating payload for trace with id {trace_id} and span with id {span_id}"
        )
        payload.append(
            {
                "id": span_id,
                "project_name": project_name,
                "trace_id": trace_id,
                "parent_span_id": parent_span_id,
                "name": span_name,
                "type": "llm",
                "start_time": start_time.isoformat() + "Z",
                "end_time": end_time.isoformat() + "Z",
                "input": input_data,
                "output": output_data,
                "metadata": metadata,
                "tags": opik_tags,
                "usage": usage,
            }
        )
        verbose_logger.debug(f"Payload: {payload}")
        return payload