Spaces:
Running
Running
| #### What this does #### | |
| # On success + failure, log events to Logfire | |
| import os | |
| import traceback | |
| import uuid | |
| from enum import Enum | |
| from typing import Any, Dict, NamedTuple | |
| from typing_extensions import LiteralString | |
| from litellm._logging import print_verbose, verbose_logger | |
| from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info | |
| class SpanConfig(NamedTuple): | |
| message_template: LiteralString | |
| span_data: Dict[str, Any] | |
| class LogfireLevel(str, Enum): | |
| INFO = "info" | |
| ERROR = "error" | |
| class LogfireLogger: | |
| # Class variables or attributes | |
| def __init__(self): | |
| try: | |
| verbose_logger.debug("in init logfire logger") | |
| import logfire | |
| # only setting up logfire if we are sending to logfire | |
| # in testing, we don't want to send to logfire | |
| if logfire.DEFAULT_LOGFIRE_INSTANCE.config.send_to_logfire: | |
| logfire.configure(token=os.getenv("LOGFIRE_TOKEN")) | |
| except Exception as e: | |
| print_verbose(f"Got exception on init logfire client {str(e)}") | |
| raise e | |
| def _get_span_config(self, payload) -> SpanConfig: | |
| if ( | |
| payload["call_type"] == "completion" | |
| or payload["call_type"] == "acompletion" | |
| ): | |
| return SpanConfig( | |
| message_template="Chat Completion with {request_data[model]!r}", | |
| span_data={"request_data": payload}, | |
| ) | |
| elif ( | |
| payload["call_type"] == "embedding" or payload["call_type"] == "aembedding" | |
| ): | |
| return SpanConfig( | |
| message_template="Embedding Creation with {request_data[model]!r}", | |
| span_data={"request_data": payload}, | |
| ) | |
| elif ( | |
| payload["call_type"] == "image_generation" | |
| or payload["call_type"] == "aimage_generation" | |
| ): | |
| return SpanConfig( | |
| message_template="Image Generation with {request_data[model]!r}", | |
| span_data={"request_data": payload}, | |
| ) | |
| else: | |
| return SpanConfig( | |
| message_template="Litellm Call with {request_data[model]!r}", | |
| span_data={"request_data": payload}, | |
| ) | |
| async def _async_log_event( | |
| self, | |
| kwargs, | |
| response_obj, | |
| start_time, | |
| end_time, | |
| print_verbose, | |
| level: LogfireLevel, | |
| ): | |
| self.log_event( | |
| kwargs=kwargs, | |
| response_obj=response_obj, | |
| start_time=start_time, | |
| end_time=end_time, | |
| print_verbose=print_verbose, | |
| level=level, | |
| ) | |
| def log_event( | |
| self, | |
| kwargs, | |
| start_time, | |
| end_time, | |
| print_verbose, | |
| level: LogfireLevel, | |
| response_obj, | |
| ): | |
| try: | |
| import logfire | |
| verbose_logger.debug( | |
| f"logfire Logging - Enters logging function for model {kwargs}" | |
| ) | |
| if not response_obj: | |
| response_obj = {} | |
| litellm_params = kwargs.get("litellm_params", {}) | |
| metadata = ( | |
| litellm_params.get("metadata", {}) or {} | |
| ) # if litellm_params['metadata'] == None | |
| messages = kwargs.get("messages") | |
| optional_params = kwargs.get("optional_params", {}) | |
| call_type = kwargs.get("call_type", "completion") | |
| cache_hit = kwargs.get("cache_hit", False) | |
| usage = response_obj.get("usage", {}) | |
| id = response_obj.get("id", str(uuid.uuid4())) | |
| try: | |
| response_time = (end_time - start_time).total_seconds() | |
| except Exception: | |
| response_time = None | |
| # Clean Metadata before logging - never log raw metadata | |
| # the raw metadata can contain circular references which leads to infinite recursion | |
| # we clean out all extra litellm metadata params before logging | |
| clean_metadata = {} | |
| if isinstance(metadata, dict): | |
| for key, value in metadata.items(): | |
| # clean litellm metadata before logging | |
| if key in [ | |
| "endpoint", | |
| "caching_groups", | |
| "previous_models", | |
| ]: | |
| continue | |
| else: | |
| clean_metadata[key] = value | |
| clean_metadata = redact_user_api_key_info(metadata=clean_metadata) | |
| # Build the initial payload | |
| payload = { | |
| "id": id, | |
| "call_type": call_type, | |
| "cache_hit": cache_hit, | |
| "startTime": start_time, | |
| "endTime": end_time, | |
| "responseTime (seconds)": response_time, | |
| "model": kwargs.get("model", ""), | |
| "user": kwargs.get("user", ""), | |
| "modelParameters": optional_params, | |
| "spend": kwargs.get("response_cost", 0), | |
| "messages": messages, | |
| "response": response_obj, | |
| "usage": usage, | |
| "metadata": clean_metadata, | |
| } | |
| logfire_openai = logfire.with_settings(custom_scope_suffix="openai") | |
| message_template, span_data = self._get_span_config(payload) | |
| if level == LogfireLevel.INFO: | |
| logfire_openai.info( | |
| message_template, | |
| **span_data, | |
| ) | |
| elif level == LogfireLevel.ERROR: | |
| logfire_openai.error( | |
| message_template, | |
| **span_data, | |
| _exc_info=True, | |
| ) | |
| print_verbose(f"\ndd Logger - Logging payload = {payload}") | |
| print_verbose( | |
| f"Logfire Layer Logging - final response object: {response_obj}" | |
| ) | |
| except Exception as e: | |
| verbose_logger.debug( | |
| f"Logfire Layer Error - {str(e)}\n{traceback.format_exc()}" | |
| ) | |
| pass | |