chore: initial public snapshot for github upload
This commit is contained in:
Your Name
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
Opik Logger that logs LLM events to an Opik server
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
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 . import opik_payload_builder, utils
|
||||
|
||||
try:
|
||||
from opik.api_objects import opik_client
|
||||
except Exception:
|
||||
opik_client = None
|
||||
|
||||
|
||||
def _should_skip_event(kwargs: Dict[str, Any]) -> bool:
|
||||
"""Check if event should be skipped due to missing standard_logging_object."""
|
||||
if kwargs.get("standard_logging_object") is None:
|
||||
verbose_logger.debug(
|
||||
"OpikLogger skipping event; no standard_logging_object found"
|
||||
)
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
class OpikLogger(CustomBatchLogger):
|
||||
"""
|
||||
Opik Logger for logging events to an Opik Server
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self.async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=httpxSpecialProvider.LoggingCallback
|
||||
)
|
||||
self.sync_httpx_client = _get_httpx_client()
|
||||
|
||||
self.opik_project_name: str = (
|
||||
utils.get_opik_config_variable(
|
||||
"project_name",
|
||||
user_value=kwargs.get("project_name", None),
|
||||
default_value="Default Project",
|
||||
)
|
||||
or "Default Project"
|
||||
)
|
||||
|
||||
opik_base_url: str = (
|
||||
utils.get_opik_config_variable(
|
||||
"url_override",
|
||||
user_value=kwargs.get("url", None),
|
||||
default_value="https://www.comet.com/opik/api",
|
||||
)
|
||||
or "https://www.comet.com/opik/api"
|
||||
)
|
||||
opik_api_key: Optional[str] = utils.get_opik_config_variable(
|
||||
"api_key", user_value=kwargs.get("api_key", None), default_value=None
|
||||
)
|
||||
opik_workspace: Optional[str] = utils.get_opik_config_variable(
|
||||
"workspace", user_value=kwargs.get("workspace", None), default_value=None
|
||||
)
|
||||
|
||||
self.trace_url: str = f"{opik_base_url}/v1/private/traces/batch"
|
||||
self.span_url: str = f"{opik_base_url}/v1/private/spans/batch"
|
||||
|
||||
self.headers: Dict[str, str] = {}
|
||||
if opik_workspace:
|
||||
self.headers["Comet-Workspace"] = opik_workspace
|
||||
|
||||
if opik_api_key:
|
||||
self.headers["authorization"] = opik_api_key
|
||||
|
||||
self.opik_workspace: Optional[str] = opik_workspace
|
||||
self.opik_api_key: Optional[str] = opik_api_key
|
||||
try:
|
||||
asyncio.create_task(self.periodic_flush())
|
||||
self.flush_lock: Optional[asyncio.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
|
||||
|
||||
# Initialize _opik_client attribute
|
||||
if opik_client is not None:
|
||||
self._opik_client = opik_client.get_client_cached()
|
||||
else:
|
||||
self._opik_client = None
|
||||
|
||||
super().__init__(**kwargs, flush_lock=self.flush_lock)
|
||||
|
||||
async def async_log_success_event(
|
||||
self,
|
||||
kwargs: Dict[str, Any],
|
||||
response_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
try:
|
||||
if _should_skip_event(kwargs):
|
||||
return
|
||||
|
||||
# Build payload using the payload builder
|
||||
trace_payload, span_payload = opik_payload_builder.build_opik_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
project_name=self.opik_project_name,
|
||||
)
|
||||
|
||||
if self._opik_client is not None:
|
||||
# Opik native client is available, use it to send data
|
||||
if trace_payload is not None:
|
||||
self._opik_client.trace(
|
||||
id=trace_payload.id,
|
||||
name=trace_payload.name,
|
||||
start_time=datetime.fromisoformat(trace_payload.start_time),
|
||||
end_time=datetime.fromisoformat(trace_payload.end_time),
|
||||
input=trace_payload.input,
|
||||
output=trace_payload.output,
|
||||
metadata=trace_payload.metadata,
|
||||
tags=trace_payload.tags,
|
||||
thread_id=trace_payload.thread_id,
|
||||
project_name=trace_payload.project_name,
|
||||
)
|
||||
|
||||
self._opik_client.span(
|
||||
id=span_payload.id,
|
||||
trace_id=span_payload.trace_id,
|
||||
parent_span_id=span_payload.parent_span_id,
|
||||
name=span_payload.name,
|
||||
type=span_payload.type,
|
||||
model=span_payload.model,
|
||||
start_time=datetime.fromisoformat(span_payload.start_time),
|
||||
end_time=datetime.fromisoformat(span_payload.end_time),
|
||||
input=span_payload.input,
|
||||
output=span_payload.output,
|
||||
metadata=span_payload.metadata,
|
||||
tags=span_payload.tags,
|
||||
usage=span_payload.usage,
|
||||
project_name=span_payload.project_name,
|
||||
provider=span_payload.provider,
|
||||
total_cost=span_payload.total_cost,
|
||||
)
|
||||
else:
|
||||
# Add payloads to LiteLLM queue
|
||||
if trace_payload is not None:
|
||||
self.log_queue.append(trace_payload.__dict__)
|
||||
self.log_queue.append(span_payload.__dict__)
|
||||
|
||||
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[str, Any]
|
||||
) -> None:
|
||||
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: Dict[str, Any],
|
||||
response_obj: Any,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
) -> None:
|
||||
try:
|
||||
if _should_skip_event(kwargs):
|
||||
return
|
||||
|
||||
# Build payload using the payload builder
|
||||
trace_payload, span_payload = opik_payload_builder.build_opik_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
project_name=self.opik_project_name,
|
||||
)
|
||||
if self._opik_client is not None:
|
||||
# Opik native client is available, use it to send data
|
||||
if trace_payload is not None:
|
||||
self._opik_client.trace(
|
||||
id=trace_payload.id,
|
||||
name=trace_payload.name,
|
||||
start_time=datetime.fromisoformat(trace_payload.start_time),
|
||||
end_time=datetime.fromisoformat(trace_payload.end_time),
|
||||
input=trace_payload.input,
|
||||
output=trace_payload.output,
|
||||
metadata=trace_payload.metadata,
|
||||
tags=trace_payload.tags,
|
||||
thread_id=trace_payload.thread_id,
|
||||
project_name=trace_payload.project_name,
|
||||
)
|
||||
|
||||
self._opik_client.span(
|
||||
id=span_payload.id,
|
||||
trace_id=span_payload.trace_id,
|
||||
parent_span_id=span_payload.parent_span_id,
|
||||
name=span_payload.name,
|
||||
type=span_payload.type,
|
||||
model=span_payload.model,
|
||||
start_time=datetime.fromisoformat(span_payload.start_time),
|
||||
end_time=datetime.fromisoformat(span_payload.end_time),
|
||||
input=span_payload.input,
|
||||
output=span_payload.output,
|
||||
metadata=span_payload.metadata,
|
||||
tags=span_payload.tags,
|
||||
usage=span_payload.usage,
|
||||
project_name=span_payload.project_name,
|
||||
provider=span_payload.provider,
|
||||
total_cost=span_payload.total_cost,
|
||||
)
|
||||
else:
|
||||
# Opik native client is not available, use LiteLLM queue to send data
|
||||
if trace_payload is not None:
|
||||
self._sync_send(
|
||||
url=self.trace_url,
|
||||
headers=self.headers,
|
||||
batch={"traces": [trace_payload.__dict__]},
|
||||
)
|
||||
|
||||
# Always send span
|
||||
self._sync_send(
|
||||
url=self.span_url,
|
||||
headers=self.headers,
|
||||
batch={"spans": [span_payload.__dict__]},
|
||||
)
|
||||
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[str, Any]
|
||||
) -> None:
|
||||
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) -> Dict[str, str]:
|
||||
headers: Dict[str, str] = {}
|
||||
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) -> None:
|
||||
verbose_logger.info("Calling async_send_batch")
|
||||
if not self.log_queue:
|
||||
return
|
||||
|
||||
# Split the log_queue into traces and spans
|
||||
traces, spans = utils.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")
|
||||
Reference in New Issue
Block a user