phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 19:24:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix - otel fixes

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-17 19:15:14 -07:00
parent a4b1b9a3ac
commit b343911d59
1 changed files with 30 additions and 24 deletions
Download patch file Download diff file Expand all files Collapse all files

54
litellm/integrations/opentelemetry.py
View file

@ -1,20 +1,21 @@
import os import os
from dataclasses import dataclass from dataclasses import dataclass
from datetime import datetime from datetime import datetime
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm._logging import verbose_logger
from litellm.types.services import ServiceLoggerPayload
from functools import wraps from functools import wraps
from typing import Union, Optional, TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any, Optional, Union
import litellm
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.services import ServiceLoggerPayload
if TYPE_CHECKING: if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span from opentelemetry.trace import Span as _Span
from litellm.proxy.proxy_server import UserAPIKeyAuth as _UserAPIKeyAuth
from litellm.proxy._types import ( from litellm.proxy._types import (
ManagementEndpointLoggingPayload as _ManagementEndpointLoggingPayload, ManagementEndpointLoggingPayload as _ManagementEndpointLoggingPayload,
) )
from litellm.proxy.proxy_server import UserAPIKeyAuth as _UserAPIKeyAuth
Span = _Span Span = _Span
UserAPIKeyAuth = _UserAPIKeyAuth UserAPIKeyAuth = _UserAPIKeyAuth
@ -107,8 +108,9 @@ class OpenTelemetry(CustomLogger):
start_time: Optional[datetime] = None, start_time: Optional[datetime] = None,
end_time: Optional[datetime] = None, end_time: Optional[datetime] = None,
): ):
from opentelemetry import trace
from datetime import datetime from datetime import datetime
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode from opentelemetry.trace import Status, StatusCode
_start_time_ns = start_time _start_time_ns = start_time
@ -145,8 +147,9 @@ class OpenTelemetry(CustomLogger):
start_time: Optional[datetime] = None, start_time: Optional[datetime] = None,
end_time: Optional[datetime] = None, end_time: Optional[datetime] = None,
): ):
from opentelemetry import trace
from datetime import datetime from datetime import datetime
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode from opentelemetry.trace import Status, StatusCode
_start_time_ns = start_time _start_time_ns = start_time
@ -179,8 +182,8 @@ class OpenTelemetry(CustomLogger):
async def async_post_call_failure_hook( async def async_post_call_failure_hook(
self, original_exception: Exception, user_api_key_dict: UserAPIKeyAuth self, original_exception: Exception, user_api_key_dict: UserAPIKeyAuth
): ):
from opentelemetry.trace import Status, StatusCode
from opentelemetry import trace from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
parent_otel_span = user_api_key_dict.parent_otel_span parent_otel_span = user_api_key_dict.parent_otel_span
if parent_otel_span is not None: if parent_otel_span is not None:
@ -202,8 +205,8 @@ class OpenTelemetry(CustomLogger):
parent_otel_span.end(end_time=self._to_ns(datetime.now())) parent_otel_span.end(end_time=self._to_ns(datetime.now()))
def _handle_sucess(self, kwargs, response_obj, start_time, end_time): def _handle_sucess(self, kwargs, response_obj, start_time, end_time):
from opentelemetry.trace import Status, StatusCode
from opentelemetry import trace from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode
verbose_logger.debug( verbose_logger.debug(
"OpenTelemetry Logger: Logging kwargs: %s, OTEL config settings=%s", "OpenTelemetry Logger: Logging kwargs: %s, OTEL config settings=%s",
@ -253,9 +256,10 @@ class OpenTelemetry(CustomLogger):
span.end(end_time=self._to_ns(end_time)) span.end(end_time=self._to_ns(end_time))
def set_tools_attributes(self, span: Span, tools): def set_tools_attributes(self, span: Span, tools):
from litellm.proxy._types import SpanAttributes
import json import json
from litellm.proxy._types import SpanAttributes
if not tools: if not tools:
return return
@ -320,7 +324,7 @@ class OpenTelemetry(CustomLogger):
) )
span.set_attribute( span.set_attribute(
SpanAttributes.LLM_IS_STREAMING, optional_params.get("stream", False) SpanAttributes.LLM_IS_STREAMING, str(optional_params.get("stream", False))
) )
if optional_params.get("tools"): if optional_params.get("tools"):
@ -439,7 +443,7 @@ class OpenTelemetry(CustomLogger):
############################################# #############################################
########## LLM Response Attributes ########## ########## LLM Response Attributes ##########
############################################# #############################################
if _raw_response: if _raw_response and isinstance(_raw_response, str):
# cast sr -> dict # cast sr -> dict
import json import json
@ -478,10 +482,10 @@ class OpenTelemetry(CustomLogger):
return _parent_context return _parent_context
def _get_span_context(self, kwargs): def _get_span_context(self, kwargs):
from opentelemetry import trace
from opentelemetry.trace.propagation.tracecontext import ( from opentelemetry.trace.propagation.tracecontext import (
TraceContextTextMapPropagator, TraceContextTextMapPropagator,
) )
from opentelemetry import trace
litellm_params = kwargs.get("litellm_params", {}) or {} litellm_params = kwargs.get("litellm_params", {}) or {}
proxy_server_request = litellm_params.get("proxy_server_request", {}) or {} proxy_server_request = litellm_params.get("proxy_server_request", {}) or {}
@ -505,17 +509,17 @@ class OpenTelemetry(CustomLogger):
return TraceContextTextMapPropagator().extract(carrier=carrier), None return TraceContextTextMapPropagator().extract(carrier=carrier), None
def _get_span_processor(self): def _get_span_processor(self):
from opentelemetry.sdk.trace.export import ( from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import (
SpanExporter, OTLPSpanExporter as OTLPSpanExporterGRPC,
SimpleSpanProcessor,
BatchSpanProcessor,
ConsoleSpanExporter,
) )
from opentelemetry.exporter.otlp.proto.http.trace_exporter import ( from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
OTLPSpanExporter as OTLPSpanExporterHTTP, OTLPSpanExporter as OTLPSpanExporterHTTP,
) )
from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( from opentelemetry.sdk.trace.export import (
OTLPSpanExporter as OTLPSpanExporterGRPC, BatchSpanProcessor,
ConsoleSpanExporter,
SimpleSpanProcessor,
SpanExporter,
) )
verbose_logger.debug( verbose_logger.debug(
@ -574,8 +578,9 @@ class OpenTelemetry(CustomLogger):
logging_payload: ManagementEndpointLoggingPayload, logging_payload: ManagementEndpointLoggingPayload,
parent_otel_span: Optional[Span] = None, parent_otel_span: Optional[Span] = None,
): ):
from opentelemetry import trace
from datetime import datetime from datetime import datetime
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode from opentelemetry.trace import Status, StatusCode
_start_time_ns = logging_payload.start_time _start_time_ns = logging_payload.start_time
@ -619,8 +624,9 @@ class OpenTelemetry(CustomLogger):
logging_payload: ManagementEndpointLoggingPayload, logging_payload: ManagementEndpointLoggingPayload,
parent_otel_span: Optional[Span] = None, parent_otel_span: Optional[Span] = None,
): ):
from opentelemetry import trace
from datetime import datetime from datetime import datetime
from opentelemetry import trace
from opentelemetry.trace import Status, StatusCode from opentelemetry.trace import Status, StatusCode
_start_time_ns = logging_payload.start_time _start_time_ns = logging_payload.start_time