mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
79 lines
2.5 KiB
Python
79 lines
2.5 KiB
Python
"""
|
|
arize AI is OTEL compatible
|
|
|
|
this file has Arize ai specific helper functions
|
|
"""
|
|
|
|
from typing import TYPE_CHECKING, Any, Optional, Union
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
|
|
Span = _Span
|
|
else:
|
|
Span = Any
|
|
|
|
|
|
def set_arize_ai_attributes(span: Span, kwargs, response_obj):
|
|
from litellm.integrations._types.open_inference import (
|
|
MessageAttributes,
|
|
MessageContentAttributes,
|
|
SpanAttributes,
|
|
)
|
|
|
|
optional_params = kwargs.get("optional_params", {})
|
|
litellm_params = kwargs.get("litellm_params", {}) or {}
|
|
|
|
#############################################
|
|
############ LLM CALL METADATA ##############
|
|
#############################################
|
|
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,
|
|
f"litellm-{str(kwargs.get('call_type', None))}",
|
|
)
|
|
span.set_attribute(SpanAttributes.LLM_INPUT_MESSAGES, str(kwargs.get("messages")))
|
|
|
|
# The Generative AI Provider: Azure, OpenAI, etc.
|
|
span.set_attribute(SpanAttributes.LLM_INVOCATION_PARAMETERS, str(optional_params))
|
|
|
|
if optional_params.get("user"):
|
|
span.set_attribute(SpanAttributes.USER_ID, optional_params.get("user"))
|
|
|
|
#############################################
|
|
########## LLM Response Attributes ##########
|
|
#############################################
|
|
llm_output_messages = []
|
|
for choice in response_obj.get("choices"):
|
|
llm_output_messages.append(choice.get("message"))
|
|
|
|
span.set_attribute(SpanAttributes.LLM_OUTPUT_MESSAGES, str(llm_output_messages))
|
|
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
|