diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index 16e5bf562d4..d0002cad450 100644 --- a/litellm/integrations/arize/_utils.py +++ b/litellm/integrations/arize/_utils.py @@ -96,132 +96,147 @@ def _set_response_attributes(span: "Span", response_obj): SpanAttributes, AudioAttributes, EmbeddingAttributes, - RerankerAttributes, ) if not hasattr(response_obj, "get"): return - # set chat completion attributes + _set_choice_outputs(span, response_obj, MessageAttributes, SpanAttributes) + _set_image_outputs(span, response_obj, ImageAttributes, SpanAttributes) + _set_audio_outputs(span, response_obj, AudioAttributes, SpanAttributes) + _set_embedding_outputs(span, response_obj, EmbeddingAttributes, SpanAttributes) + _set_structured_outputs(span, response_obj, MessageAttributes, SpanAttributes) + _set_usage_outputs(span, response_obj, SpanAttributes) + + +def _set_choice_outputs(span: "Span", response_obj, msg_attrs, span_attrs): for idx, choice in enumerate(response_obj.get("choices", [])): response_message = choice.get("message", {}) safe_set_attribute( span, - SpanAttributes.OUTPUT_VALUE, + span_attrs.OUTPUT_VALUE, response_message.get("content", ""), ) - prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{idx}" + prefix = f"{span_attrs.LLM_OUTPUT_MESSAGES}.{idx}" safe_set_attribute( span, - f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", + f"{prefix}.{msg_attrs.MESSAGE_ROLE}", response_message.get("role"), ) safe_set_attribute( span, - f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}", + f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", response_message.get("content", ""), ) - # set image generation attributes + +def _set_image_outputs(span: "Span", response_obj, image_attrs, span_attrs): images = response_obj.get("data", []) - if images: - for i, image in enumerate(images): - img_url = image.get("url") - if img_url is None and image.get("b64_json"): - img_url = f"data:image/png;base64,{image.get('b64_json')}" + for i, image in enumerate(images): + img_url = image.get("url") + if img_url is None and image.get("b64_json"): + img_url = f"data:image/png;base64,{image.get('b64_json')}" - if img_url: - if i == 0: - safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, img_url) + if not img_url: + continue - # ste attribute per-image url - safe_set_attribute(span, f"{ImageAttributes.IMAGE_URL}.{i}", img_url) + if i == 0: + safe_set_attribute(span, span_attrs.OUTPUT_VALUE, img_url) - # set audio generation attr + safe_set_attribute(span, f"{image_attrs.IMAGE_URL}.{i}", img_url) + + +def _set_audio_outputs(span: "Span", response_obj, audio_attrs, span_attrs): audio = response_obj.get("audio", []) - if audio: - for i, audio_item in enumerate(audio): - audio_url = audio_item.get("url") - if audio_url is None and audio_item.get("b64_json"): - audio_url = f"data:audio/wav;base64,{audio_item.get('b64_json')}" + for i, audio_item in enumerate(audio): + audio_url = audio_item.get("url") + if audio_url is None and audio_item.get("b64_json"): + audio_url = f"data:audio/wav;base64,{audio_item.get('b64_json')}" - if audio_url: - if i == 0: - safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, audio_url) + if audio_url: + if i == 0: + safe_set_attribute(span, span_attrs.OUTPUT_VALUE, audio_url) + safe_set_attribute(span, f"{audio_attrs.AUDIO_URL}.{i}", audio_url) - safe_set_attribute(span, f"{AudioAttributes.AUDIO_URL}.{i}", audio_url) + audio_mime = audio_item.get("mime_type") + if audio_mime: + safe_set_attribute(span, f"{audio_attrs.AUDIO_MIME_TYPE}.{i}", audio_mime) - audio_mime = audio_item.get("mime_type") - if audio_mime: - safe_set_attribute(span, f"{AudioAttributes.AUDIO_MIME_TYPE}.{i}", audio_mime) + audio_transcript = audio_item.get("transcript") + if audio_transcript: + safe_set_attribute(span, f"{audio_attrs.AUDIO_TRANSCRIPT}.{i}", audio_transcript) - audio_transcript = audio_item.get("transcript") - if audio_transcript: - safe_set_attribute(span, f"{AudioAttributes.AUDIO_TRANSCRIPT}.{i}", audio_transcript) +def _set_embedding_outputs(span: "Span", response_obj, embedding_attrs, span_attrs): embeddings = response_obj.get("data", []) - if embeddings: - for i, embedding_item in enumerate(embeddings): - embedding_vector = embedding_item.get("embedding") - if embedding_vector: - if i == 0: - safe_set_attribute( - span, - SpanAttributes.OUTPUT_VALUE, - str(embedding_vector), - ) - + for i, embedding_item in enumerate(embeddings): + embedding_vector = embedding_item.get("embedding") + if embedding_vector: + if i == 0: safe_set_attribute( span, - f"{EmbeddingAttributes.EMBEDDING_VECTOR}.{i}", + span_attrs.OUTPUT_VALUE, str(embedding_vector), ) - embedding_text = embedding_item.get("text") - if embedding_text: - safe_set_attribute( - span, - f"{EmbeddingAttributes.EMBEDDING_TEXT}.{i}", - str(embedding_text), - ) + safe_set_attribute( + span, + f"{embedding_attrs.EMBEDDING_VECTOR}.{i}", + str(embedding_vector), + ) + embedding_text = embedding_item.get("text") + if embedding_text: + safe_set_attribute( + span, + f"{embedding_attrs.EMBEDDING_TEXT}.{i}", + str(embedding_text), + ) + + +def _set_structured_outputs(span: "Span", response_obj, msg_attrs, span_attrs): output_items = response_obj.get("output", []) - if output_items: - for i, item in enumerate(output_items): - prefix = f"{SpanAttributes.LLM_OUTPUT_MESSAGES}.{i}" - if hasattr(item, "type"): - item_type = item.type - if item_type == "reasoning" and hasattr(item, "summary"): - for summary in item.summary: - if hasattr(summary, "text"): - safe_set_attribute( - span, - f"{prefix}.{MessageAttributes.MESSAGE_REASONING_SUMMARY}", - summary.text, - ) - elif item_type == "message" and hasattr(item, "content"): - message_content = "" - content_list = item.content - if content_list and len(content_list) > 0: - first_content = content_list[0] - message_content = getattr(first_content, "text", "") - message_role = getattr(item, "role", "assistant") - safe_set_attribute(span, SpanAttributes.OUTPUT_VALUE, message_content) - safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_CONTENT}", message_content) - safe_set_attribute(span, f"{prefix}.{MessageAttributes.MESSAGE_ROLE}", message_role) + for i, item in enumerate(output_items): + prefix = f"{span_attrs.LLM_OUTPUT_MESSAGES}.{i}" + if not hasattr(item, "type"): + continue + item_type = item.type + if item_type == "reasoning" and hasattr(item, "summary"): + for summary in item.summary: + if hasattr(summary, "text"): + safe_set_attribute( + span, + f"{prefix}.{msg_attrs.MESSAGE_REASONING_SUMMARY}", + summary.text, + ) + elif item_type == "message" and hasattr(item, "content"): + message_content = "" + content_list = item.content + if content_list and len(content_list) > 0: + first_content = content_list[0] + message_content = getattr(first_content, "text", "") + message_role = getattr(item, "role", "assistant") + safe_set_attribute(span, span_attrs.OUTPUT_VALUE, message_content) + safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_CONTENT}", message_content) + safe_set_attribute(span, f"{prefix}.{msg_attrs.MESSAGE_ROLE}", message_role) + + +def _set_usage_outputs(span: "Span", response_obj, span_attrs): usage = response_obj and response_obj.get("usage") - if usage: - safe_set_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens")) - completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") - if completion_tokens: - safe_set_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION, completion_tokens) - prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") - if prompt_tokens: - safe_set_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_PROMPT, prompt_tokens) - reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens") - if reasoning_tokens: - safe_set_attribute(span, SpanAttributes.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, reasoning_tokens) + if not usage: + return + + safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, usage.get("total_tokens")) + completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") + if completion_tokens: + safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens) + prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") + if prompt_tokens: + safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens) + reasoning_tokens = usage.get("output_tokens_details", {}).get("reasoning_tokens") + if reasoning_tokens: + safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION_DETAILS_REASONING, reasoning_tokens) def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: @@ -250,7 +265,7 @@ def _infer_open_inference_span_kind(call_type: Optional[str]) -> str: if lowered == "call_mcp_tool" or lowered == "mcp" or lowered.endswith("tool"): return OpenInferenceSpanKindValues.TOOL.value - if "assistant" in lowered: + if "asend_message" in lowered or "a2a" in lowered or "assistant" in lowered: return OpenInferenceSpanKindValues.AGENT.value if any( @@ -336,81 +351,41 @@ def set_attributes( ) try: - # Remove secret_fields to prevent leaking sensitive data (e.g., authorization headers) - optional_params = kwargs.get("optional_params", {}) - if isinstance(optional_params, dict): - optional_params.pop("secret_fields", None) - litellm_params = kwargs.get("litellm_params", {}) + optional_params = _sanitize_optional_params(kwargs.get("optional_params")) + litellm_params = kwargs.get("litellm_params", {}) or {} standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object" ) if standard_logging_payload is None: raise ValueError("standard_logging_object not found in kwargs") - metadata = ( - standard_logging_payload.get("metadata") - if standard_logging_payload - else None - ) - if metadata is not None: - safe_set_attribute(span, SpanAttributes.METADATA, safe_dumps(metadata)) + metadata = standard_logging_payload.get("metadata") if standard_logging_payload else None + _set_metadata_attributes(span, metadata, SpanAttributes) - metadata_tools: Optional[list] = None - if isinstance(metadata, dict): - llm_obj = metadata.get("llm") - if isinstance(llm_obj, dict): - metadata_tools = llm_obj.get("tools") - - optional_tools = None - if isinstance(optional_params, dict): - optional_tools = optional_params.get("tools") + metadata_tools = _extract_metadata_tools(metadata) + optional_tools = _extract_optional_tools(optional_params) call_type = standard_logging_payload.get("call_type") - - if kwargs.get("model"): - safe_set_attribute(span, SpanAttributes.LLM_MODEL_NAME, kwargs.get("model")) - - safe_set_attribute(span, "llm.request.type", standard_logging_payload["call_type"]) - safe_set_attribute(span, SpanAttributes.LLM_PROVIDER, litellm_params.get("custom_llm_provider", "Unknown")) - - if optional_params.get("max_tokens"): - safe_set_attribute(span, "llm.request.max_tokens", optional_params.get("max_tokens")) - if optional_params.get("temperature"): - safe_set_attribute(span, "llm.request.temperature", optional_params.get("temperature")) - if optional_params.get("top_p"): - safe_set_attribute(span, "llm.request.top_p", optional_params.get("top_p")) - - safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False))) - - if optional_params.get("user"): - safe_set_attribute(span, "llm.user", optional_params.get("user")) - - if response_obj and response_obj.get("id"): - safe_set_attribute(span, "llm.response.id", response_obj.get("id")) - if response_obj and response_obj.get("model"): - safe_set_attribute(span, "llm.response.model", response_obj.get("model")) + _set_request_attributes( + span=span, + kwargs=kwargs, + standard_logging_payload=standard_logging_payload, + optional_params=optional_params, + litellm_params=litellm_params, + response_obj=response_obj, + span_attrs=SpanAttributes, + ) span_kind = _infer_open_inference_span_kind(call_type=call_type) - _set_tool_attributes(span, optional_tools, metadata_tools) - if (optional_tools or metadata_tools) and span_kind != OpenInferenceSpanKindValues.TOOL.value: span_kind = OpenInferenceSpanKindValues.TOOL.value safe_set_attribute(span, SpanAttributes.OPENINFERENCE_SPAN_KIND, span_kind) attributes.set_messages(span, kwargs) - model_params = ( - standard_logging_payload.get("model_parameters") - if standard_logging_payload - else None - ) - if model_params: - safe_set_attribute(span, SpanAttributes.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)) - if model_params.get("user"): - user_id = model_params.get("user") - if user_id is not None: - safe_set_attribute(span, SpanAttributes.USER_ID, user_id) + model_params = standard_logging_payload.get("model_parameters") if standard_logging_payload else None + _set_model_params(span, model_params, SpanAttributes) _set_response_attributes(span=span, response_obj=response_obj) @@ -420,3 +395,72 @@ def set_attributes( ) if hasattr(span, "record_exception"): span.record_exception(e) + + +def _sanitize_optional_params(optional_params: Optional[dict]) -> dict: + if not isinstance(optional_params, dict): + return {} + optional_params.pop("secret_fields", None) + return optional_params + + +def _set_metadata_attributes(span: "Span", metadata: Optional[dict], span_attrs) -> None: + if metadata is not None: + safe_set_attribute(span, span_attrs.METADATA, safe_dumps(metadata)) + + +def _extract_metadata_tools(metadata: Optional[dict]) -> Optional[list]: + if not isinstance(metadata, dict): + return None + llm_obj = metadata.get("llm") + if isinstance(llm_obj, dict): + return llm_obj.get("tools") + return None + + +def _extract_optional_tools(optional_params: dict) -> Optional[list]: + return optional_params.get("tools") if isinstance(optional_params, dict) else None + + +def _set_request_attributes( + span: "Span", + kwargs, + standard_logging_payload: StandardLoggingPayload, + optional_params: dict, + litellm_params: dict, + response_obj, + span_attrs, +): + if kwargs.get("model"): + safe_set_attribute(span, span_attrs.LLM_MODEL_NAME, kwargs.get("model")) + + safe_set_attribute(span, "llm.request.type", standard_logging_payload.get("call_type")) + safe_set_attribute(span, span_attrs.LLM_PROVIDER, litellm_params.get("custom_llm_provider", "Unknown")) + + if optional_params.get("max_tokens"): + safe_set_attribute(span, "llm.request.max_tokens", optional_params.get("max_tokens")) + if optional_params.get("temperature"): + safe_set_attribute(span, "llm.request.temperature", optional_params.get("temperature")) + if optional_params.get("top_p"): + safe_set_attribute(span, "llm.request.top_p", optional_params.get("top_p")) + + safe_set_attribute(span, "llm.is_streaming", str(optional_params.get("stream", False))) + + if optional_params.get("user"): + safe_set_attribute(span, "llm.user", optional_params.get("user")) + + if response_obj and response_obj.get("id"): + safe_set_attribute(span, "llm.response.id", response_obj.get("id")) + if response_obj and response_obj.get("model"): + safe_set_attribute(span, "llm.response.model", response_obj.get("model")) + + +def _set_model_params(span: "Span", model_params: Optional[dict], span_attrs) -> None: + if not model_params: + return + + safe_set_attribute(span, span_attrs.LLM_INVOCATION_PARAMETERS, safe_dumps(model_params)) + if model_params.get("user"): + user_id = model_params.get("user") + if user_id is not None: + safe_set_attribute(span, span_attrs.USER_ID, user_id)