phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix opik types

Browse source
This commit is contained in:
Ishaan Jaff 2024-10-10 18:37:53 +05:30
parent 4064bfc6dd
commit fbf756806e
1 changed files with 102 additions and 84 deletions
Download patch file Download diff file Expand all files Collapse all files

158
litellm/integrations/opik/opik.py
View file

@ -2,27 +2,26 @@
Opik Logger that logs LLM events to an Opik server Opik Logger that logs LLM events to an Opik server
""" """
from typing import Dict, List import asyncio
import json import json
import traceback
from typing import Dict, List
from litellm._logging import verbose_logger from litellm._logging import verbose_logger
import traceback from litellm.integrations.custom_batch_logger import CustomBatchLogger
from .utils import (
get_opik_config_variable,
create_uuid7,
create_usage_object,
get_traces_and_spans_from_payload
)
import asyncio
from litellm.llms.custom_httpx.http_handler import ( from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
_get_httpx_client, _get_httpx_client,
get_async_httpx_client,
httpxSpecialProvider, httpxSpecialProvider,
) )
from litellm.integrations.custom_batch_logger import CustomBatchLogger from .utils import (
create_usage_object,
create_uuid7,
get_opik_config_variable,
get_traces_and_spans_from_payload,
)
class OpikLogger(CustomBatchLogger): class OpikLogger(CustomBatchLogger):
""" """
@ -38,23 +37,19 @@ class OpikLogger(CustomBatchLogger):
self.opik_project_name = get_opik_config_variable( self.opik_project_name = get_opik_config_variable(
"project_name", "project_name",
user_value=kwargs.get("project_name", None), user_value=kwargs.get("project_name", None),
default_value="Default Project" default_value="Default Project",
) )
opik_base_url = get_opik_config_variable( opik_base_url = get_opik_config_variable(
"url_override", "url_override",
user_value=kwargs.get("url", None), user_value=kwargs.get("url", None),
default_value="https://www.comet.com/opik/api" default_value="https://www.comet.com/opik/api",
) )
opik_api_key = get_opik_config_variable( opik_api_key = get_opik_config_variable(
"api_key", "api_key", user_value=kwargs.get("api_key", None), default_value=None
user_value=kwargs.get("api_key", None),
default_value=None
) )
opik_workspace = get_opik_config_variable( opik_workspace = get_opik_config_variable(
"workspace", "workspace", user_value=kwargs.get("workspace", None), default_value=None
user_value=kwargs.get("workspace", None),
default_value=None
) )
self.trace_url = f"{opik_base_url}/v1/private/traces/batch" self.trace_url = f"{opik_base_url}/v1/private/traces/batch"
@ -67,12 +62,14 @@ class OpikLogger(CustomBatchLogger):
if opik_api_key: if opik_api_key:
self.headers["authorization"] = opik_api_key self.headers["authorization"] = opik_api_key
self.opik_workspace = opik_workspace
self.opik_api_key = opik_api_key
try: try:
asyncio.create_task(self.periodic_flush()) asyncio.create_task(self.periodic_flush())
self.flush_lock = asyncio.Lock() self.flush_lock = asyncio.Lock()
except Exception as e: except Exception as e:
verbose_logger.debug( verbose_logger.exception(
f"OpikLogger - Asynchronous processing not initialized as we are not running in an async context" f"OpikLogger - Asynchronous processing not initialized as we are not running in an async context {str(e)}"
) )
self.flush_lock = None self.flush_lock = None
@ -84,11 +81,13 @@ class OpikLogger(CustomBatchLogger):
kwargs=kwargs, kwargs=kwargs,
response_obj=response_obj, response_obj=response_obj,
start_time=start_time, start_time=start_time,
end_time=end_time end_time=end_time,
) )
self.log_queue.extend(opik_payload) self.log_queue.extend(opik_payload)
verbose_logger.debug(f"OpikLogger added event to log_queue - Will flush in {self.flush_interval} seconds...") 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: if len(self.log_queue) >= self.batch_size:
verbose_logger.debug("OpikLogger - Flushing batch") verbose_logger.debug("OpikLogger - Flushing batch")
@ -98,12 +97,10 @@ class OpikLogger(CustomBatchLogger):
f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}" f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}"
) )
def _sync_send(self, url: str, headers: Dict[str, str], batch: List[Dict]): def _sync_send(self, url: str, headers: Dict[str, str], batch: Dict):
try: try:
response = self.sync_httpx_client.post( response = self.sync_httpx_client.post(
url=url, url=url, headers=headers, json=batch # type: ignore
headers=headers,
json=batch
) )
response.raise_for_status() response.raise_for_status()
if response.status_code != 204: if response.status_code != 204:
@ -121,25 +118,27 @@ class OpikLogger(CustomBatchLogger):
kwargs=kwargs, kwargs=kwargs,
response_obj=response_obj, response_obj=response_obj,
start_time=start_time, start_time=start_time,
end_time=end_time end_time=end_time,
) )
traces, spans = get_traces_and_spans_from_payload(opik_payload) traces, spans = get_traces_and_spans_from_payload(opik_payload)
if len(traces) > 0: if len(traces) > 0:
self._sync_send(self.trace_url, self.headers, {"traces": traces}) self._sync_send(
url=self.trace_url, headers=self.headers, batch={"traces": traces}
)
if len(spans) > 0: if len(spans) > 0:
self._sync_send(self.span_url, self.headers, {"spans": spans}) self._sync_send(
url=self.span_url, headers=self.headers, batch={"spans": spans}
)
except Exception as e: except Exception as e:
verbose_logger.exception( verbose_logger.exception(
f"OpikLogger failed to log success event - {str(e)}\n{traceback.format_exc()}" 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: List[Dict]): async def _submit_batch(self, url: str, headers: Dict[str, str], batch: Dict):
try: try:
response = await self.async_httpx_client.post( response = await self.async_httpx_client.post(
url=url, url=url, headers=headers, json=batch # type: ignore
headers=headers,
json=batch
) )
response.raise_for_status() response.raise_for_status()
@ -165,7 +164,6 @@ class OpikLogger(CustomBatchLogger):
headers["authorization"] = self.opik_api_key headers["authorization"] = self.opik_api_key
return headers return headers
async def async_send_batch(self): async def async_send_batch(self):
verbose_logger.exception("Calling async_send_batch") verbose_logger.exception("Calling async_send_batch")
if not self.log_queue: if not self.log_queue:
@ -176,13 +174,17 @@ class OpikLogger(CustomBatchLogger):
# Send trace batch # Send trace batch
if len(traces) > 0: if len(traces) > 0:
await self._submit_batch(self.trace_url, self.headers, {"traces": traces}) await self._submit_batch(
url=self.trace_url, headers=self.headers, batch={"traces": traces}
)
if len(spans) > 0: if len(spans) > 0:
await self._submit_batch(self.span_url, self.headers, {"spans": spans}) await self._submit_batch(
url=self.span_url, headers=self.headers, batch={"spans": spans}
)
def _create_opik_payload(self, kwargs, response_obj, start_time, end_time) -> List[Dict]:
def _create_opik_payload(
self, kwargs, response_obj, start_time, end_time
) -> List[Dict]:
# Get metadata # Get metadata
_litellm_params = kwargs.get("litellm_params", {}) or {} _litellm_params = kwargs.get("litellm_params", {}) or {}
@ -190,7 +192,9 @@ class OpikLogger(CustomBatchLogger):
# Extract opik metadata # Extract opik metadata
litellm_opik_metadata = litellm_params_metadata.get("opik", {}) litellm_opik_metadata = litellm_params_metadata.get("opik", {})
verbose_logger.debug(f"litellm_opik_metadata - {json.dumps(litellm_opik_metadata, default=str)}") 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) project_name = litellm_opik_metadata.get("project_name", self.opik_project_name)
# Extract trace_id and parent_span_id # Extract trace_id and parent_span_id
@ -212,7 +216,9 @@ class OpikLogger(CustomBatchLogger):
# Use standard_logging_object to create metadata and input/output data # Use standard_logging_object to create metadata and input/output data
standard_logging_object = kwargs.get("standard_logging_object", None) standard_logging_object = kwargs.get("standard_logging_object", None)
if standard_logging_object is None: if standard_logging_object is None:
verbose_logger.debug("OpikLogger skipping event; no standard_logging_object found") verbose_logger.debug(
"OpikLogger skipping event; no standard_logging_object found"
)
return [] return []
# Create input and output data # Create input and output data
@ -245,12 +251,16 @@ class OpikLogger(CustomBatchLogger):
if "response_cost" in kwargs: if "response_cost" in kwargs:
metadata["cost"] = { metadata["cost"] = {
"total_tokens": kwargs["response_cost"], "total_tokens": kwargs["response_cost"],
"currency": "USD" "currency": "USD",
} }
if "response_cost_failure_debug_info" in kwargs: if "response_cost_failure_debug_info" in kwargs:
metadata["response_cost_failure_debug_info"] = kwargs["response_cost_failure_debug_info"] metadata["response_cost_failure_debug_info"] = kwargs[
"response_cost_failure_debug_info"
]
if "model_map_information" in standard_logging_object: if "model_map_information" in standard_logging_object:
metadata["model_map_information"] = standard_logging_object["model_map_information"] metadata["model_map_information"] = standard_logging_object[
"model_map_information"
]
if "model" in standard_logging_object: if "model" in standard_logging_object:
metadata["model"] = standard_logging_object["model"] metadata["model"] = standard_logging_object["model"]
if "model_id" in standard_logging_object: if "model_id" in standard_logging_object:
@ -273,36 +283,44 @@ class OpikLogger(CustomBatchLogger):
payload = [] payload = []
if trace_id is None: if trace_id is None:
trace_id = create_uuid7() trace_id = create_uuid7()
verbose_logger.debug(f"OpikLogger creating payload for trace with id {trace_id}") verbose_logger.debug(
f"OpikLogger creating payload for trace with id {trace_id}"
)
payload.append({ 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, "project_name": project_name,
"id": trace_id, "trace_id": trace_id,
"name": trace_name, "parent_span_id": parent_span_id,
"name": span_name,
"type": "llm",
"start_time": start_time.isoformat() + "Z", "start_time": start_time.isoformat() + "Z",
"end_time": end_time.isoformat() + "Z", "end_time": end_time.isoformat() + "Z",
"input": input_data, "input": input_data,
"output": output_data, "output": output_data,
"metadata": metadata, "metadata": metadata,
"tags": opik_tags, "tags": opik_tags,
}) "usage": usage,
}
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}") verbose_logger.debug(f"Payload: {payload}")
return payload return payload