phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 10:44:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix: bugs in traceloop integration

Browse source
This commit is contained in:
Nir Gazit 2023-10-19 17:23:51 +02:00
parent 00ac43c19a
commit c21e86faec
4 changed files with 144 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

80
litellm/integrations/traceloop.py
View file

@ -1,8 +1,78 @@
import sys
class TraceloopLogger:
def __init__(self):
from traceloop.tracing import Tracer
from traceloop.sdk.tracing.tracing import TracerWrapper
self.tracer = Tracer.init(app_name=sys.argv[0], disable_batch=True)
self.tracer_wrapper = TracerWrapper()
def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose):
from opentelemetry.trace import SpanKind
from opentelemetry.semconv.ai import SpanAttributes
try:
tracer = self.tracer_wrapper.get_tracer()
model = kwargs.get("model")
# LiteLLM uses the standard OpenAI library, so it's already handled by Traceloop SDK
if "gpt" in model:
return
with tracer.start_as_current_span(
"litellm.completion",
kind=SpanKind.CLIENT,
) as span:
if span.is_recording():
span.set_attribute(
SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model")
)
span.set_attribute(
SpanAttributes.LLM_REQUEST_MAX_TOKENS, kwargs.get("max_tokens")
)
span.set_attribute(
SpanAttributes.LLM_TEMPERATURE, kwargs.get("temperature")
)
for idx, prompt in enumerate(kwargs.get("messages")):
span.set_attribute(
f"{SpanAttributes.LLM_PROMPTS}.{idx}.role",
prompt.get("role"),
)
span.set_attribute(
f"{SpanAttributes.LLM_PROMPTS}.{idx}.content",
prompt.get("content"),
)
span.set_attribute(
SpanAttributes.LLM_RESPONSE_MODEL, response_obj.get("model")
)
usage = response_obj.get("usage")
if usage:
span.set_attribute(
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
usage.get("total_tokens"),
)
span.set_attribute(
SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
usage.get("completion_tokens"),
)
span.set_attribute(
SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
usage.get("prompt_tokens"),
)
for idx, choice in enumerate(response_obj.get("choices")):
span.set_attribute(
f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.finish_reason",
choice.get("finish_reason"),
)
span.set_attribute(
f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.role",
choice.get("message").get("role"),
)
span.set_attribute(
f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.content",
choice.get("message").get("content"),
)
except Exception as e:
print_verbose(f"Traceloop Layer Error - {e}")