mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 19:24:27 +00:00
105 lines
3.1 KiB
Python
105 lines
3.1 KiB
Python
"""
|
|
arize AI is OTEL compatible
|
|
|
|
this file has Arize ai specific helper functions
|
|
"""
|
|
|
|
import os
|
|
from datetime import datetime
|
|
from typing import TYPE_CHECKING, Any, Optional, Union
|
|
|
|
from litellm.integrations.arize import _utils
|
|
from litellm.integrations.opentelemetry import OpenTelemetry
|
|
from litellm.types.integrations.arize import ArizeConfig
|
|
from litellm.types.services import ServiceLoggerPayload
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
|
|
from litellm.types.integrations.arize import Protocol as _Protocol
|
|
|
|
Protocol = _Protocol
|
|
Span = _Span
|
|
else:
|
|
Protocol = Any
|
|
Span = Any
|
|
|
|
|
|
class ArizeLogger(OpenTelemetry):
|
|
|
|
def set_attributes(self, span: Span, kwargs, response_obj: Optional[Any]):
|
|
ArizeLogger.set_arize_attributes(span, kwargs, response_obj)
|
|
return
|
|
|
|
@staticmethod
|
|
def set_arize_attributes(span: Span, kwargs, response_obj):
|
|
_utils.set_attributes(span, kwargs, response_obj)
|
|
return
|
|
|
|
@staticmethod
|
|
def get_arize_config() -> ArizeConfig:
|
|
"""
|
|
Helper function to get Arize configuration.
|
|
|
|
Returns:
|
|
ArizeConfig: A Pydantic model containing Arize configuration.
|
|
|
|
Raises:
|
|
ValueError: If required environment variables are not set.
|
|
"""
|
|
space_key = os.environ.get("ARIZE_SPACE_KEY")
|
|
api_key = os.environ.get("ARIZE_API_KEY")
|
|
|
|
grpc_endpoint = os.environ.get("ARIZE_ENDPOINT")
|
|
http_endpoint = os.environ.get("ARIZE_HTTP_ENDPOINT")
|
|
|
|
endpoint = None
|
|
protocol: Protocol = "otlp_grpc"
|
|
|
|
if grpc_endpoint:
|
|
protocol = "otlp_grpc"
|
|
endpoint = grpc_endpoint
|
|
elif http_endpoint:
|
|
protocol = "otlp_http"
|
|
endpoint = http_endpoint
|
|
else:
|
|
protocol = "otlp_grpc"
|
|
endpoint = "https://otlp.arize.com/v1"
|
|
|
|
return ArizeConfig(
|
|
space_key=space_key,
|
|
api_key=api_key,
|
|
protocol=protocol,
|
|
endpoint=endpoint,
|
|
)
|
|
|
|
async def async_service_success_hook(
|
|
self,
|
|
payload: ServiceLoggerPayload,
|
|
parent_otel_span: Optional[Span] = None,
|
|
start_time: Optional[Union[datetime, float]] = None,
|
|
end_time: Optional[Union[datetime, float]] = None,
|
|
event_metadata: Optional[dict] = None,
|
|
):
|
|
"""Arize is used mainly for LLM I/O tracing, sending router+caching metrics adds bloat to arize logs"""
|
|
pass
|
|
|
|
async def async_service_failure_hook(
|
|
self,
|
|
payload: ServiceLoggerPayload,
|
|
error: Optional[str] = "",
|
|
parent_otel_span: Optional[Span] = None,
|
|
start_time: Optional[Union[datetime, float]] = None,
|
|
end_time: Optional[Union[float, datetime]] = None,
|
|
event_metadata: Optional[dict] = None,
|
|
):
|
|
"""Arize is used mainly for LLM I/O tracing, sending router+caching metrics adds bloat to arize logs"""
|
|
pass
|
|
|
|
def create_litellm_proxy_request_started_span(
|
|
self,
|
|
start_time: datetime,
|
|
headers: dict,
|
|
):
|
|
"""Arize is used mainly for LLM I/O tracing, sending Proxy Server Request adds bloat to arize logs"""
|
|
pass
|