mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
6c76f5f9c6
Every ruff-strict rule that sat above its budget limit (FURB188, RUF022, SIM118, UP007, UP032, UP037) is now at zero, LIT001 and LIT006 are back under their ceilings, and the freed headroom is ratcheted out of ruff-strict-budget.json, type-discipline-budget.json, and basedpyright-code-budget.json so the gates take the fast path again
102 lines
4 KiB
Python
102 lines
4 KiB
Python
import json
|
|
from typing import TYPE_CHECKING, Any, Final
|
|
|
|
from litellm.proxy._types import SpanAttributes
|
|
|
|
if TYPE_CHECKING:
|
|
from opentelemetry.trace import Span as _Span
|
|
|
|
Span = _Span | Any
|
|
else:
|
|
Span = Any
|
|
|
|
|
|
class LangtraceAttributes:
|
|
"""
|
|
This class is used to save trace attributes to Langtrace's spans
|
|
"""
|
|
|
|
def set_langtrace_attributes(self, span: Span, kwargs, response_obj):
|
|
"""
|
|
This function is used to log the event to Langtrace
|
|
"""
|
|
|
|
vendor: Final = kwargs.get("litellm_params").get("custom_llm_provider")
|
|
optional_params: Final = kwargs.get("optional_params", {})
|
|
options: Final = {**kwargs, **optional_params}
|
|
self.set_request_attributes(span, options, vendor)
|
|
self.set_response_attributes(span, response_obj)
|
|
self.set_usage_attributes(span, response_obj)
|
|
|
|
def set_request_attributes(self, span: Span, kwargs, vendor):
|
|
"""
|
|
This function is used to get span attributes for the LLM request
|
|
"""
|
|
span_attributes: Final = {
|
|
"gen_ai.operation.name": "chat",
|
|
"langtrace.service.name": vendor,
|
|
SpanAttributes.LLM_REQUEST_MODEL.value: kwargs.get("model"),
|
|
SpanAttributes.LLM_IS_STREAMING.value: kwargs.get("stream"),
|
|
SpanAttributes.LLM_REQUEST_TEMPERATURE.value: kwargs.get("temperature"),
|
|
SpanAttributes.LLM_TOP_K.value: kwargs.get("top_k"),
|
|
SpanAttributes.LLM_REQUEST_TOP_P.value: kwargs.get("top_p"),
|
|
SpanAttributes.LLM_USER.value: kwargs.get("user"),
|
|
SpanAttributes.LLM_REQUEST_MAX_TOKENS.value: kwargs.get("max_tokens"),
|
|
SpanAttributes.LLM_RESPONSE_STOP_REASON.value: kwargs.get("stop"),
|
|
SpanAttributes.LLM_FREQUENCY_PENALTY.value: kwargs.get("frequency_penalty"),
|
|
SpanAttributes.LLM_PRESENCE_PENALTY.value: kwargs.get("presence_penalty"),
|
|
}
|
|
|
|
prompts: Final = kwargs.get("messages")
|
|
|
|
if prompts:
|
|
span.add_event(
|
|
name="gen_ai.content.prompt",
|
|
attributes={SpanAttributes.LLM_PROMPTS.value: json.dumps(prompts)},
|
|
)
|
|
|
|
self.set_span_attributes(span, span_attributes)
|
|
|
|
def set_response_attributes(self, span: Span, response_obj):
|
|
"""
|
|
This function is used to get span attributes for the LLM response
|
|
"""
|
|
response_attributes: Final = {
|
|
"gen_ai.response_id": response_obj.get("id"),
|
|
"gen_ai.system_fingerprint": response_obj.get("system_fingerprint"),
|
|
SpanAttributes.LLM_RESPONSE_MODEL.value: response_obj.get("model"),
|
|
}
|
|
completions: Final = []
|
|
for choice in response_obj.get("choices", []):
|
|
role = choice.get("message").get("role")
|
|
content = choice.get("message").get("content")
|
|
completions.append({"role": role, "content": content})
|
|
|
|
span.add_event(
|
|
name="gen_ai.content.completion",
|
|
attributes={SpanAttributes.LLM_COMPLETIONS: json.dumps(completions)},
|
|
)
|
|
|
|
self.set_span_attributes(span, response_attributes)
|
|
|
|
def set_usage_attributes(self, span: Span, response_obj):
|
|
"""
|
|
This function is used to get span attributes for the LLM usage
|
|
"""
|
|
usage: Final = response_obj.get("usage")
|
|
if usage:
|
|
usage_attributes: Final = {
|
|
SpanAttributes.LLM_USAGE_PROMPT_TOKENS.value: usage.get("prompt_tokens"),
|
|
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS.value: usage.get("completion_tokens"),
|
|
SpanAttributes.LLM_USAGE_TOTAL_TOKENS.value: usage.get("total_tokens"),
|
|
}
|
|
self.set_span_attributes(span, usage_attributes)
|
|
|
|
def set_span_attributes(self, span: Span, attributes):
|
|
"""
|
|
This function is used to set span attributes
|
|
"""
|
|
for key, value in attributes.items():
|
|
if not value:
|
|
continue
|
|
span.set_attribute(key, value)
|