Spaces:
Running
Running
| import json | |
| from typing import TYPE_CHECKING, Any, Union | |
| from litellm.proxy._types import SpanAttributes | |
| if TYPE_CHECKING: | |
| from opentelemetry.trace import Span as _Span | |
| Span = Union[_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 = kwargs.get("litellm_params").get("custom_llm_provider") | |
| optional_params = kwargs.get("optional_params", {}) | |
| options = {**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 = { | |
| "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 = 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 = { | |
| "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 = [] | |
| 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 = response_obj.get("usage") | |
| if usage: | |
| usage_attributes = { | |
| 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) | |