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

(feat) Arize - Allow using Arize HTTP endpoint (#6364)

Browse source 
* arize use helper for get_arize_opentelemetry_config

* use helper to get Arize OTEL config

* arize add helpers for arize

* docs allow using arize http endpoint

* fix importing OTEL for Arize

* use static methods for ArizeLogger

* fix ArizeLogger tests
This commit is contained in:
Ishaan Jaff 2024-10-23 09:38:35 +05:30 • committed by GitHub
parent f943410e32
commit b75019c1a5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 257 additions and 124 deletions
Download patch file Download diff file Expand all files Collapse all files

3
docs/my-website/docs/observability/arize_integration.md
View file

@ -62,7 +62,8 @@ litellm_settings:
environment_variables: environment_variables:
ARIZE_SPACE_KEY: "d0*****" ARIZE_SPACE_KEY: "d0*****"
ARIZE_API_KEY: "141a****" ARIZE_API_KEY: "141a****"
ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize api endpoint ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize GRPC api endpoint
ARIZE_HTTP_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize HTTP api endpoint. Set either this or ARIZE_ENDPOINT
``` ```
## Support & Talk to Founders ## Support & Talk to Founders

3
docs/my-website/docs/proxy/logging.md
View file

@ -1279,7 +1279,8 @@ litellm_settings:
environment_variables: environment_variables:
ARIZE_SPACE_KEY: "d0*****" ARIZE_SPACE_KEY: "d0*****"
ARIZE_API_KEY: "141a****" ARIZE_API_KEY: "141a****"
ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize api endpoint ARIZE_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize GRPC api endpoint
ARIZE_HTTP_ENDPOINT: "https://otlp.arize.com/v1" # OPTIONAL - your custom arize HTTP api endpoint. Set either this or ARIZE_ENDPOINT
``` ```
2. Start Proxy 2. Start Proxy

287
litellm/integrations/arize_ai.py
View file

@ -7,135 +7,208 @@ this file has Arize ai specific helper functions
import json import json
from typing import TYPE_CHECKING, Any, Optional, Union from typing import TYPE_CHECKING, Any, Optional, Union
from litellm._logging import verbose_proxy_logger from litellm._logging import verbose_logger
if TYPE_CHECKING: if TYPE_CHECKING:
from opentelemetry.trace import Span as _Span from opentelemetry.trace import Span as _Span
from .opentelemetry import OpenTelemetryConfig as _OpenTelemetryConfig
Span = _Span Span = _Span
OpenTelemetryConfig = _OpenTelemetryConfig
else: else:
Span = Any Span = Any
OpenTelemetryConfig = Any
import os
from litellm.types.integrations.arize import *
def make_json_serializable(payload: dict) -> dict: class ArizeLogger:
for key, value in payload.items(): @staticmethod
def set_arize_ai_attributes(span: Span, kwargs, response_obj):
from litellm.integrations._types.open_inference import (
MessageAttributes,
MessageContentAttributes,
OpenInferenceSpanKindValues,
SpanAttributes,
)
try: try:
if isinstance(value, dict):
# recursively sanitize dicts optional_params = kwargs.get("optional_params", {})
payload[key] = make_json_serializable(value.copy()) # litellm_params = kwargs.get("litellm_params", {}) or {}
elif not isinstance(value, (str, int, float, bool, type(None))):
# everything else becomes a string #############################################
payload[key] = str(value) ############ LLM CALL METADATA ##############
except Exception: #############################################
# non blocking if it can't cast to a str # commented out for now - looks like Arize AI could not log this
# metadata = litellm_params.get("metadata", {}) or {}
# span.set_attribute(SpanAttributes.METADATA, str(metadata))
#############################################
########## LLM Request Attributes ###########
#############################################
# The name of the LLM a request is being made to
if kwargs.get("model"):
span.set_attribute(SpanAttributes.LLM_MODEL_NAME, kwargs.get("model"))
span.set_attribute(
SpanAttributes.OPENINFERENCE_SPAN_KIND,
OpenInferenceSpanKindValues.LLM.value,
)
messages = kwargs.get("messages")
# for /chat/completions
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
if messages:
span.set_attribute(
SpanAttributes.INPUT_VALUE,
messages[-1].get("content", ""), # get the last message for input
)
# LLM_INPUT_MESSAGES shows up under `input_messages` tab on the span page
for idx, msg in enumerate(messages):
# Set the role per message
span.set_attribute(
f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}.{MessageAttributes.MESSAGE_ROLE}",
msg["role"],
)
# Set the content per message
span.set_attribute(
f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}.{MessageAttributes.MESSAGE_CONTENT}",
msg.get("content", ""),
)
# The Generative AI Provider: Azure, OpenAI, etc.
_optional_params = ArizeLogger.make_json_serializable(optional_params)
_json_optional_params = json.dumps(_optional_params)
span.set_attribute(
SpanAttributes.LLM_INVOCATION_PARAMETERS, _json_optional_params
)
if optional_params.get("user"):
span.set_attribute(SpanAttributes.USER_ID, optional_params.get("user"))
#############################################
########## LLM Response Attributes ##########
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
#############################################
for choice in response_obj.get("choices"):
response_message = choice.get("message", {})
span.set_attribute(
SpanAttributes.OUTPUT_VALUE, response_message.get("content", "")
)
# This shows up under `output_messages` tab on the span page
# This code assumes a single response
span.set_attribute(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}",
response_message["role"],
)
span.set_attribute(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}",
response_message.get("content", ""),
)
usage = response_obj.get("usage")
if usage:
span.set_attribute(
SpanAttributes.LLM_TOKEN_COUNT_TOTAL,
usage.get("total_tokens"),
)
# The number of tokens used in the LLM response (completion).
span.set_attribute(
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION,
usage.get("completion_tokens"),
)
# The number of tokens used in the LLM prompt.
span.set_attribute(
SpanAttributes.LLM_TOKEN_COUNT_PROMPT,
usage.get("prompt_tokens"),
)
pass pass
return payload except Exception as e:
verbose_logger.error(f"Error setting arize attributes: {e}")
###################### Helper functions ######################
def set_arize_ai_attributes(span: Span, kwargs, response_obj): @staticmethod
from litellm.integrations._types.open_inference import ( def _get_arize_config() -> ArizeConfig:
MessageAttributes, """
MessageContentAttributes, Helper function to get Arize configuration.
OpenInferenceSpanKindValues,
SpanAttributes,
)
try: Returns:
ArizeConfig: A Pydantic model containing Arize configuration.
optional_params = kwargs.get("optional_params", {}) Raises:
# litellm_params = kwargs.get("litellm_params", {}) or {} ValueError: If required environment variables are not set.
"""
space_key = os.environ.get("ARIZE_SPACE_KEY")
api_key = os.environ.get("ARIZE_API_KEY")
############################################# if not space_key:
############ LLM CALL METADATA ############## raise ValueError("ARIZE_SPACE_KEY not found in environment variables")
############################################# if not api_key:
# commented out for now - looks like Arize AI could not log this raise ValueError("ARIZE_API_KEY not found in environment variables")
# metadata = litellm_params.get("metadata", {}) or {}
# span.set_attribute(SpanAttributes.METADATA, str(metadata))
############################################# grpc_endpoint = os.environ.get("ARIZE_ENDPOINT")
########## LLM Request Attributes ########### http_endpoint = os.environ.get("ARIZE_HTTP_ENDPOINT")
############################################# if grpc_endpoint is None and http_endpoint is None:
# use default arize grpc endpoint
# The name of the LLM a request is being made to verbose_logger.debug(
if kwargs.get("model"): "No ARIZE_ENDPOINT or ARIZE_HTTP_ENDPOINT found, using default endpoint: https://otlp.arize.com/v1"
span.set_attribute(SpanAttributes.LLM_MODEL_NAME, kwargs.get("model"))
span.set_attribute(
SpanAttributes.OPENINFERENCE_SPAN_KIND,
OpenInferenceSpanKindValues.LLM.value,
)
messages = kwargs.get("messages")
# for /chat/completions
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
if messages:
span.set_attribute(
SpanAttributes.INPUT_VALUE,
messages[-1].get("content", ""), # get the last message for input
) )
grpc_endpoint = "https://otlp.arize.com/v1"
# LLM_INPUT_MESSAGES shows up under `input_messages` tab on the span page return ArizeConfig(
for idx, msg in enumerate(messages): space_key=space_key,
# Set the role per message api_key=api_key,
span.set_attribute( grpc_endpoint=grpc_endpoint,
f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}.{MessageAttributes.MESSAGE_ROLE}", http_endpoint=http_endpoint,
msg["role"],
)
# Set the content per message
span.set_attribute(
f"{SpanAttributes.LLM_INPUT_MESSAGES}.{idx}.{MessageAttributes.MESSAGE_CONTENT}",
msg.get("content", ""),
)
# The Generative AI Provider: Azure, OpenAI, etc.
_optional_params = make_json_serializable(optional_params)
_json_optional_params = json.dumps(_optional_params)
span.set_attribute(
SpanAttributes.LLM_INVOCATION_PARAMETERS, _json_optional_params
) )
if optional_params.get("user"): @staticmethod
span.set_attribute(SpanAttributes.USER_ID, optional_params.get("user")) def get_arize_opentelemetry_config() -> Optional[OpenTelemetryConfig]:
"""
Helper function to get OpenTelemetry configuration for Arize.
############################################# Args:
########## LLM Response Attributes ########## arize_config (ArizeConfig): Arize configuration object.
# https://docs.arize.com/arize/large-language-models/tracing/semantic-conventions
############################################# Returns:
for choice in response_obj.get("choices"): OpenTelemetryConfig: Configuration for OpenTelemetry.
response_message = choice.get("message", {}) """
span.set_attribute( from .opentelemetry import OpenTelemetryConfig
SpanAttributes.OUTPUT_VALUE, response_message.get("content", "")
arize_config = ArizeLogger._get_arize_config()
if arize_config.http_endpoint:
return OpenTelemetryConfig(
exporter="otlp_http",
endpoint=arize_config.http_endpoint,
) )
# This shows up under `output_messages` tab on the span page # use default arize grpc endpoint
# This code assumes a single response return OpenTelemetryConfig(
span.set_attribute( exporter="otlp_grpc",
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_ROLE}", endpoint=arize_config.grpc_endpoint,
response_message["role"], )
)
span.set_attribute(
f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.0.{MessageAttributes.MESSAGE_CONTENT}",
response_message.get("content", ""),
)
usage = response_obj.get("usage") @staticmethod
if usage: def make_json_serializable(payload: dict) -> dict:
span.set_attribute( for key, value in payload.items():
SpanAttributes.LLM_TOKEN_COUNT_TOTAL, try:
usage.get("total_tokens"), if isinstance(value, dict):
) # recursively sanitize dicts
payload[key] = ArizeLogger.make_json_serializable(value.copy())
# The number of tokens used in the LLM response (completion). elif not isinstance(value, (str, int, float, bool, type(None))):
span.set_attribute( # everything else becomes a string
SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, payload[key] = str(value)
usage.get("completion_tokens"), except Exception:
) # non blocking if it can't cast to a str
pass
# The number of tokens used in the LLM prompt. return payload
span.set_attribute(
SpanAttributes.LLM_TOKEN_COUNT_PROMPT,
usage.get("prompt_tokens"),
)
pass
except Exception as e:
verbose_proxy_logger.error(f"Error setting arize attributes: {e}")

4
litellm/integrations/opentelemetry.py
View file

@ -396,9 +396,9 @@ class OpenTelemetry(CustomLogger):
def set_attributes(self, span: Span, kwargs, response_obj): # noqa: PLR0915 def set_attributes(self, span: Span, kwargs, response_obj): # noqa: PLR0915
try: try:
if self.callback_name == "arize": if self.callback_name == "arize":
from litellm.integrations.arize_ai import set_arize_ai_attributes from litellm.integrations.arize_ai import ArizeLogger
set_arize_ai_attributes(span, kwargs, response_obj) ArizeLogger.set_arize_ai_attributes(span, kwargs, response_obj)
return return
elif self.callback_name == "langtrace": elif self.callback_name == "langtrace":
from litellm.integrations.langtrace import LangtraceAttributes from litellm.integrations.langtrace import LangtraceAttributes

18
litellm/litellm_core_utils/litellm_logging.py
View file

@ -60,6 +60,7 @@ from litellm.utils import (
from ..integrations.aispend import AISpendLogger from ..integrations.aispend import AISpendLogger
from ..integrations.argilla import ArgillaLogger from ..integrations.argilla import ArgillaLogger
from ..integrations.arize_ai import ArizeLogger
from ..integrations.athina import AthinaLogger from ..integrations.athina import AthinaLogger
from ..integrations.braintrust_logging import BraintrustLogger from ..integrations.braintrust_logging import BraintrustLogger
from ..integrations.datadog.datadog import DataDogLogger from ..integrations.datadog.datadog import DataDogLogger
@ -2323,22 +2324,16 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_opik_logger) _in_memory_loggers.append(_opik_logger)
return _opik_logger # type: ignore return _opik_logger # type: ignore
elif logging_integration == "arize": elif logging_integration == "arize":
if "ARIZE_SPACE_KEY" not in os.environ:
raise ValueError("ARIZE_SPACE_KEY not found in environment variables")
if "ARIZE_API_KEY" not in os.environ:
raise ValueError("ARIZE_API_KEY not found in environment variables")
from litellm.integrations.opentelemetry import ( from litellm.integrations.opentelemetry import (
OpenTelemetry, OpenTelemetry,
OpenTelemetryConfig, OpenTelemetryConfig,
) )
arize_endpoint = ( otel_config = ArizeLogger.get_arize_opentelemetry_config()
os.environ.get("ARIZE_ENDPOINT", None) or "https://otlp.arize.com/v1" if otel_config is None:
) raise ValueError(
otel_config = OpenTelemetryConfig( "No valid endpoint found for Arize, please set 'ARIZE_ENDPOINT' to your GRPC endpoint or 'ARIZE_HTTP_ENDPOINT' to your HTTP endpoint"
exporter="otlp_grpc", )
endpoint=arize_endpoint,
)
os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = ( os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
f"space_key={os.getenv('ARIZE_SPACE_KEY')},api_key={os.getenv('ARIZE_API_KEY')}" f"space_key={os.getenv('ARIZE_SPACE_KEY')},api_key={os.getenv('ARIZE_API_KEY')}"
) )
@ -2351,7 +2346,6 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_otel_logger = OpenTelemetry(config=otel_config, callback_name="arize") _otel_logger = OpenTelemetry(config=otel_config, callback_name="arize")
_in_memory_loggers.append(_otel_logger) _in_memory_loggers.append(_otel_logger)
return _otel_logger # type: ignore return _otel_logger # type: ignore
elif logging_integration == "otel": elif logging_integration == "otel":
from litellm.integrations.opentelemetry import OpenTelemetry from litellm.integrations.opentelemetry import OpenTelemetry

10
litellm/types/integrations/arize.py Normal file
View file

@ -0,0 +1,10 @@
from typing import Optional
from pydantic import BaseModel
class ArizeConfig(BaseModel):
space_key: str
api_key: str
grpc_endpoint: Optional[str] = None
http_endpoint: Optional[str] = None

56
tests/local_testing/test_arize_ai.py
View file

@ -10,9 +10,9 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE
import litellm import litellm
from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._logging import verbose_logger, verbose_proxy_logger
from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig from litellm.integrations.opentelemetry import OpenTelemetry, OpenTelemetryConfig
from litellm.integrations.arize_ai import ArizeConfig, ArizeLogger
load_dotenv() load_dotenv()
import logging
@pytest.mark.asyncio() @pytest.mark.asyncio()
@ -32,3 +32,57 @@ async def test_async_otel_callback():
) )
await asyncio.sleep(2) await asyncio.sleep(2)
@pytest.fixture
def mock_env_vars(monkeypatch):
monkeypatch.setenv("ARIZE_SPACE_KEY", "test_space_key")
monkeypatch.setenv("ARIZE_API_KEY", "test_api_key")
def test_get_arize_config(mock_env_vars):
"""
Use Arize default endpoint when no endpoints are provided
"""
config = ArizeLogger._get_arize_config()
assert isinstance(config, ArizeConfig)
assert config.space_key == "test_space_key"
assert config.api_key == "test_api_key"
assert config.grpc_endpoint == "https://otlp.arize.com/v1"
assert config.http_endpoint is None
def test_get_arize_config_with_endpoints(mock_env_vars, monkeypatch):
"""
Use provided endpoints when they are set
"""
monkeypatch.setenv("ARIZE_ENDPOINT", "grpc://test.endpoint")
monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "http://test.endpoint")
config = ArizeLogger._get_arize_config()
assert config.grpc_endpoint == "grpc://test.endpoint"
assert config.http_endpoint == "http://test.endpoint"
def test_get_arize_opentelemetry_config_grpc(mock_env_vars, monkeypatch):
"""
Use provided GRPC endpoint when it is set
"""
monkeypatch.setenv("ARIZE_ENDPOINT", "grpc://test.endpoint")
config = ArizeLogger.get_arize_opentelemetry_config()
assert isinstance(config, OpenTelemetryConfig)
assert config.exporter == "otlp_grpc"
assert config.endpoint == "grpc://test.endpoint"
def test_get_arize_opentelemetry_config_http(mock_env_vars, monkeypatch):
"""
Use provided HTTP endpoint when it is set
"""
monkeypatch.setenv("ARIZE_HTTP_ENDPOINT", "http://test.endpoint")
config = ArizeLogger.get_arize_opentelemetry_config()
assert isinstance(config, OpenTelemetryConfig)
assert config.exporter == "otlp_http"
assert config.endpoint == "http://test.endpoint"