fix logfire

This commit is contained in:
Ishaan Jaff 2024-06-17 19:21:35 -07:00
parent 253b504199
commit 4a75e57b6f
3 changed files with 37 additions and 6 deletions

View file

@ -34,7 +34,7 @@ input_callback: List[Union[str, Callable]] = []
success_callback: List[Union[str, Callable]] = [] success_callback: List[Union[str, Callable]] = []
failure_callback: List[Union[str, Callable]] = [] failure_callback: List[Union[str, Callable]] = []
service_callback: List[Union[str, Callable]] = [] service_callback: List[Union[str, Callable]] = []
_custom_logger_compatible_callbacks_literal = Literal["lago", "openmeter"] _custom_logger_compatible_callbacks_literal = Literal["lago", "openmeter", "logfire"]
callbacks: List[Union[Callable, _custom_logger_compatible_callbacks_literal]] = [] callbacks: List[Union[Callable, _custom_logger_compatible_callbacks_literal]] = []
_langfuse_default_tags: Optional[ _langfuse_default_tags: Optional[
List[ List[

View file

@ -10,7 +10,7 @@ import sys
import time import time
import traceback import traceback
import uuid import uuid
from typing import Callable, Optional from typing import Any, Callable, Dict, List, Optional
import litellm import litellm
from litellm import ( from litellm import (
@ -73,6 +73,8 @@ from ..integrations.supabase import Supabase
from ..integrations.traceloop import TraceloopLogger from ..integrations.traceloop import TraceloopLogger
from ..integrations.weights_biases import WeightsBiasesLogger from ..integrations.weights_biases import WeightsBiasesLogger
_in_memory_loggers: List[Any] = []
class Logging: class Logging:
global supabaseClient, liteDebuggerClient, promptLayerLogger, weightsBiasesLogger, langsmithLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app global supabaseClient, liteDebuggerClient, promptLayerLogger, weightsBiasesLogger, langsmithLogger, logfireLogger, capture_exception, add_breadcrumb, lunaryLogger, logfireLogger, prometheusLogger, slack_app
@ -1741,13 +1743,33 @@ def _init_custom_logger_compatible_class(
logging_integration: litellm._custom_logger_compatible_callbacks_literal, logging_integration: litellm._custom_logger_compatible_callbacks_literal,
) -> Callable: ) -> Callable:
if logging_integration == "lago": if logging_integration == "lago":
return LagoLogger() # type: ignore for callback in _in_memory_loggers:
if isinstance(callback, LagoLogger):
return callback # type: ignore
lago_logger = LagoLogger()
_in_memory_loggers.append(lago_logger)
return lago_logger # type: ignore
elif logging_integration == "openmeter": elif logging_integration == "openmeter":
return OpenMeterLogger() # type: ignore for callback in _in_memory_loggers:
if isinstance(callback, OpenMeterLogger):
return callback # type: ignore
_openmeter_logger = OpenMeterLogger()
_in_memory_loggers.append(_openmeter_logger)
return _openmeter_logger # type: ignore
elif logging_integration == "logfire": elif logging_integration == "logfire":
if "LOGFIRE_TOKEN" not in os.environ:
raise ValueError("LOGFIRE_TOKEN not found in environment variables")
otel_config = OpenTelemetryConfig( otel_config = OpenTelemetryConfig(
exporter="otlp_http", exporter="otlp_http",
endpoint="https://logfire-api.pydantic.dev/v1/traces", endpoint="https://logfire-api.pydantic.dev/v1/traces",
headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}", headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}",
) )
return OpenTelemetry(config=otel_config) # type: ignore for callback in _in_memory_loggers:
if isinstance(callback, OpenTelemetry):
return callback # type: ignore
_otel_logger = OpenTelemetry(config=otel_config)
_in_memory_loggers.append(_otel_logger)
return _otel_logger # type: ignore

View file

@ -340,6 +340,15 @@ def function_setup(
try: try:
global callback_list, add_breadcrumb, user_logger_fn, Logging global callback_list, add_breadcrumb, user_logger_fn, Logging
function_id = kwargs["id"] if "id" in kwargs else None function_id = kwargs["id"] if "id" in kwargs else None
# Note: maintain backwards compatibility with logfire. some users use litellm.success_callback=["logfire"]
if (
litellm.success_callback is not None
and "logfire" in litellm.success_callback
):
litellm.success_callback.remove("logfire")
litellm.callbacks.append("logfire")
if len(litellm.callbacks) > 0: if len(litellm.callbacks) > 0:
for callback in litellm.callbacks: for callback in litellm.callbacks:
# check if callback is a string - e.g. "lago", "openmeter" # check if callback is a string - e.g. "lago", "openmeter"