fix(sdk): manual logging example + fix span ended error (#3352)

Author: nirga

packages/sample-app/sample_app/manual_logging_example.py

@@ -0,0 +1,109 @@
+"""
+Manual logging example using OpenLLMetry tracing.
+
+This example demonstrates how to manually report LLM calls using the
+track_llm_call context manager from traceloop.sdk.tracing.manual.
+"""
+
+import os
+import openai
+from dotenv import load_dotenv
+from traceloop.sdk import Traceloop
+from traceloop.sdk.tracing.manual import LLMMessage, LLMUsage, track_llm_call
+
+
+def main():
+    # Load environment variables
+    load_dotenv()
+
+    # Initialize Traceloop
+    Traceloop.init(
+        app_name="manual-logging-example",
+        disable_batch=True,
+        instruments={}
+    )
+
+    # Initialize OpenAI client
+    openai_client = openai.OpenAI(
+        api_key=os.getenv("OPENAI_API_KEY")
+    )
+
+    print("Starting manual logging example...")
+
+    # Example 1: Basic LLM call with manual reporting
+    print("\n=== Example 1: Basic LLM Call ===")
+    with track_llm_call(vendor="openai", type="chat") as span:
+        # Report the request
+        messages = [
+            LLMMessage(role="user", content="Tell me a joke about opentelemetry")
+        ]
+        span.report_request(
+            model="gpt-4o",
+            messages=messages,
+        )
+
+        # Make the actual API call
+        response = openai_client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "user", "content": "Tell me a joke about opentelemetry"}
+            ],
+        )
+
+        # Report the response
+        response_messages = [choice.message.content for choice in response.choices]
+        span.report_response(response.model, response_messages)
+
+        # Report usage metrics
+        if response.usage:
+            span.report_usage(
+                LLMUsage(
+                    prompt_tokens=response.usage.prompt_tokens,
+                    completion_tokens=response.usage.completion_tokens,
+                    total_tokens=response.usage.total_tokens,
+                )
+            )
+
+        print(f"Response: {response_messages[0]}")
+
+    # Example 2: Second independent LLM call
+    print("\n=== Example 2: Second LLM Call ===")
+    with track_llm_call(vendor="openai", type="chat") as span:
+        # Report the request
+        messages = [
+            LLMMessage(role="user", content="What is machine learning?")
+        ]
+        span.report_request(
+            model="gpt-4o",
+            messages=messages,
+        )
+
+        # Make the actual API call
+        response = openai_client.chat.completions.create(
+            model="gpt-4o",
+            messages=[
+                {"role": "user", "content": "What is machine learning?"}
+            ],
+        )
+
+        # Report the response
+        response_messages = [choice.message.content for choice in response.choices]
+        span.report_response(response.model, response_messages)
+
+        # Report usage metrics
+        if response.usage:
+            span.report_usage(
+                LLMUsage(
+                    prompt_tokens=response.usage.prompt_tokens,
+                    completion_tokens=response.usage.completion_tokens,
+                    total_tokens=response.usage.total_tokens,
+                )
+            )
+
+        print(f"Response: {response_messages[0]}")
+
+    print("\nManual logging example completed!")
+
+
+if __name__ == "__main__":
+    main()

packages/traceloop-sdk/traceloop/sdk/tracing/manual.py

@@ -1,7 +1,9 @@
 from contextlib import contextmanager
+from opentelemetry import context
 from opentelemetry.semconv_ai import SpanAttributes
-from opentelemetry.trace import Span
+from opentelemetry.trace import Span, set_span_in_context
 from pydantic import BaseModel
+from typing import Optional
 from traceloop.sdk.tracing.context_manager import get_tracer
 
 
@@ -14,8 +16,8 @@ class LLMUsage(BaseModel):
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
-    cache_creation_input_tokens: int
-    cache_read_input_tokens: int
+    cache_creation_input_tokens: Optional[int] = None
+    cache_read_input_tokens: Optional[int] = None
 
 
 class LLMSpan:
@@ -55,24 +57,28 @@ def report_usage(self, usage: LLMUsage):
         self._span.set_attribute(
             SpanAttributes.LLM_USAGE_TOTAL_TOKENS, usage.total_tokens
         )
-        self._span.set_attribute(
-            SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS,
-            usage.cache_creation_input_tokens,
-        )
-        self._span.set_attribute(
-            SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
-            usage.cache_read_input_tokens,
-        )
+        if usage.cache_creation_input_tokens is not None:
+            self._span.set_attribute(
+                SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS,
+                usage.cache_creation_input_tokens,
+            )
+        if usage.cache_read_input_tokens is not None:
+            self._span.set_attribute(
+                SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS,
+                usage.cache_read_input_tokens,
+            )
 
 
 @contextmanager
 def track_llm_call(vendor: str, type: str):
     with get_tracer() as tracer:
-        with tracer.start_as_current_span(name=f"{vendor}.{type}") as span:
-            span.set_attribute(SpanAttributes.LLM_SYSTEM, vendor)
-            span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, type)
-            llm_span = LLMSpan(span)
-            try:
-                yield llm_span
-            finally:
-                span.end()
+        span = tracer.start_span(name=f"{vendor}.{type}")
+        span.set_attribute(SpanAttributes.LLM_SYSTEM, vendor)
+        span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, type)
+        ctx = set_span_in_context(span)
+        token = context.attach(ctx)
+        try:
+            yield LLMSpan(span)
+        finally:
+            context.detach(token)
+            span.end()