Merge pull request #647 from nirga/traceloop-bug

fix: bugs in traceloop integration

docs/my-website/docs/observability/traceloop_integration.md

@@ -10,22 +10,19 @@ It is based on [OpenTelemetry](https://opentelemetry.io), so it can provide full
 
 ## Getting Started
 
-First, sign up to get an API key on the [Traceloop dashboard](https://app.traceloop.com/settings/api-keys).
-
-Then, install the Traceloop SDK:
+Install the Traceloop SDK:
 
 ```
 pip install traceloop-sdk
 ```
 
-Use just 1 line of code, to instantly log your LLM responses:
+Use just 2 lines of code, to instantly log your LLM responses with OpenTelemetry:
 
 ```python
+Traceloop.init(app_name=<YOUR APP NAME>, disable_batch=True)
 litellm.success_callback = ["traceloop"]
 ```
 
-When running your app, make sure to set the `TRACELOOP_API_KEY` environment variable to your API key.
-
 To get better visualizations on how your code behaves, you may want to annotate specific parts of your LLM chain. See [Traceloop docs on decorators](https://traceloop.com/docs/python-sdk/decorators) for more information.
 
 ## Exporting traces to other systems (e.g. Datadog, New Relic, and others)

litellm/integrations/traceloop.py

@@ -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}")

litellm/tests/test_traceloop.py

@@ -0,0 +1,57 @@
+import litellm
+from litellm import completion
+from traceloop.sdk import Traceloop
+
+Traceloop.init(app_name="test_traceloop", disable_batch=True)
+litellm.success_callback = ["traceloop"]
+
+
+def test_traceloop_logging():
+    try:
+        response = completion(
+            model="claude-instant-1.2",
+            messages=[
+                {"role": "user", "content": "Tell me a joke about OpenTelemetry"}
+            ],
+            max_tokens=10,
+            temperature=0.2,
+        )
+        print(response)
+    except Exception as e:
+        print(e)
+
+
+test_traceloop_logging()
+
+
+def test_traceloop_tracing_function_calling():
+    function1 = [
+        {
+            "name": "get_current_weather",
+            "description": "Get the current weather in a given location",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                "required": ["location"],
+            },
+        }
+    ]
+    try:
+        response = completion(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "what's the weather in boston"}],
+            temperature=0.1,
+            functions=function1,
+        )
+        print(response)
+    except Exception as e:
+        print(e)
+
+
+test_traceloop_tracing_function_calling()

litellm/utils.py

@@ -2411,6 +2411,15 @@ def handle_success(args, kwargs, result, start_time, end_time):
                         end_time=end_time,
                         print_verbose=print_verbose,
                     )
+                
+                elif callback == "traceloop":
+                    traceloopLogger.log_event(
+                        kwargs=kwargs,
+                        response_obj=result,
+                        start_time=start_time,
+                        end_time=end_time,
+                        print_verbose=print_verbose,
+                    )
 
                 elif callback == "aispend":
                     print_verbose("reaches aispend for logging!")