refactor(litellm_tracer): streamline chunk processing and enhance tracing with structured logging

Author: viniciusdsmello

src/openlayer/lib/integrations/litellm_tracer.py

@@ -17,6 +17,7 @@
     import litellm
 
 from ..tracing import tracer
+from ..tracing import enums as tracer_enums
 
 logger = logging.getLogger(__name__)
 
@@ -154,121 +155,120 @@ def stream_chunks(
     latest_usage_data = {"total_tokens": None, "prompt_tokens": None, "completion_tokens": None}
     provider = "unknown"
     latest_chunk_metadata = {}
-    captured_context = contextvars.copy_context()
 
-    try:
-        i = 0
-        for i, chunk in enumerate(chunks):
-            raw_outputs.append(chunk.model_dump() if hasattr(chunk, 'model_dump') else str(chunk))
-            
-            if i == 0:
-                first_token_time = time.time()
-                # Try to detect provider from the first chunk
-                provider = detect_provider_from_chunk(chunk, model_name)
-            
-            # Extract usage data from this chunk if available (usually in final chunks)
-            chunk_usage = extract_usage_from_chunk(chunk)
-            if any(v is not None for v in chunk_usage.values()):
-                latest_usage_data = chunk_usage
+    # Create step immediately so it's added to parent trace before parent publishes
+    with tracer.create_step(
+        name="LiteLLM Chat Completion",
+        step_type=tracer_enums.StepType.CHAT_COMPLETION,
+        inputs={"prompt": kwargs.get("messages", [])},
+    ) as step:
+        try:
+            i = 0
+            for i, chunk in enumerate(chunks):
+                raw_outputs.append(chunk.model_dump() if hasattr(chunk, 'model_dump') else str(chunk))
 
-            # Always update metadata from latest chunk (for cost, headers, etc.)
-            chunk_metadata = extract_litellm_metadata(chunk, model_name)
-            if chunk_metadata:
-                latest_chunk_metadata.update(chunk_metadata)
+                if i == 0:
+                    first_token_time = time.time()
+                    # Try to detect provider from the first chunk
+                    provider = detect_provider_from_chunk(chunk, model_name)
 
-            if i > 0:
-                num_of_completion_tokens = i + 1
-
-            # Handle different chunk formats based on provider
-            delta = get_delta_from_chunk(chunk)
-
-            if delta and hasattr(delta, 'content') and delta.content:
-                collected_output_data.append(delta.content)
-            elif delta and hasattr(delta, 'function_call') and delta.function_call:
-                if delta.function_call.name:
-                    collected_function_call["name"] += delta.function_call.name
-                if delta.function_call.arguments:
-                    collected_function_call["arguments"] += delta.function_call.arguments
-            elif delta and hasattr(delta, 'tool_calls') and delta.tool_calls:
-                if delta.tool_calls[0].function.name:
-                    collected_function_call["name"] += delta.tool_calls[0].function.name
-                if delta.tool_calls[0].function.arguments:
-                    collected_function_call["arguments"] += delta.tool_calls[0].function.arguments
-
-            yield chunk
-            
-        end_time = time.time()
-        latency = (end_time - start_time) * 1000
-        
-    # pylint: disable=broad-except
-    except Exception as e:
-        logger.error("Failed to yield chunk. %s", e)
-    finally:
-        # #region agent log - Debug: trace finally block execution
-        _parent = tracer.get_current_step()
-        _trace = tracer.get_current_trace()
-        print(f"[OPENLAYER_DEBUG] litellm_tracer.py:finally | has_parent_step={_parent is not None} | parent_step_name={_parent.name if _parent else None} | has_trace={_trace is not None} | trace_steps_count={len(_trace.steps) if _trace else 0}", flush=True)
-        # #endregion
-        # Try to add step to the trace
-        try:
-            collected_output_data = [message for message in collected_output_data if message is not None]
-            if collected_output_data:
-                output_data = "".join(collected_output_data)
-            else:
-                if collected_function_call["arguments"]:
-                    try:
-                        collected_function_call["arguments"] = json.loads(collected_function_call["arguments"])
-                    except json.JSONDecodeError:
-                        pass
-                output_data = collected_function_call
-
-            # Post-streaming calculations (after streaming is finished)
-            completion_tokens_calculated, prompt_tokens_calculated, total_tokens_calculated, cost_calculated = calculate_streaming_usage_and_cost(
-                chunks=raw_outputs,
-                messages=kwargs.get("messages", []),
-                output_content=output_data,
-                model_name=model_name,
-                latest_usage_data=latest_usage_data,
-                latest_chunk_metadata=latest_chunk_metadata
-            )
-            
-            # Use calculated values (fall back to extracted data if calculation fails)
-            usage_data = latest_usage_data if any(v is not None for v in latest_usage_data.values()) else {}
-            
-            final_prompt_tokens = prompt_tokens_calculated if prompt_tokens_calculated is not None else usage_data.get("prompt_tokens", 0)
-            final_completion_tokens = completion_tokens_calculated if completion_tokens_calculated is not None else usage_data.get("completion_tokens", num_of_completion_tokens)
-            final_total_tokens = total_tokens_calculated if total_tokens_calculated is not None else usage_data.get("total_tokens", final_prompt_tokens + final_completion_tokens)
-            final_cost = cost_calculated if cost_calculated is not None else latest_chunk_metadata.get('cost', None)
+                # Extract usage data from this chunk if available (usually in final chunks)
+                chunk_usage = extract_usage_from_chunk(chunk)
+                if any(v is not None for v in chunk_usage.values()):
+                    latest_usage_data = chunk_usage
+                    
+                # Always update metadata from latest chunk (for cost, headers, etc.)
+                chunk_metadata = extract_litellm_metadata(chunk, model_name)
+                if chunk_metadata:
+                    latest_chunk_metadata.update(chunk_metadata)
+                    
+                if i > 0:
+                    num_of_completion_tokens = i + 1
+
+                # Handle different chunk formats based on provider
+                delta = get_delta_from_chunk(chunk)
+
+                if delta and hasattr(delta, 'content') and delta.content:
+                    collected_output_data.append(delta.content)
+                elif delta and hasattr(delta, 'function_call') and delta.function_call:
+                    if delta.function_call.name:
+                        collected_function_call["name"] += delta.function_call.name
+                    if delta.function_call.arguments:
+                        collected_function_call["arguments"] += delta.function_call.arguments
+                elif delta and hasattr(delta, 'tool_calls') and delta.tool_calls:
+                    if delta.tool_calls[0].function.name:
+                        collected_function_call["name"] += delta.tool_calls[0].function.name
+                    if delta.tool_calls[0].function.arguments:
+                        collected_function_call["arguments"] += delta.tool_calls[0].function.arguments
+
+                yield chunk
+                
+            end_time = time.time()
+            latency = (end_time - start_time) * 1000
 
-            trace_args = create_trace_args(
-                end_time=end_time,
-                inputs={"prompt": kwargs.get("messages", [])},
-                output=output_data,
-                latency=latency,
-                tokens=final_total_tokens,
-                prompt_tokens=final_prompt_tokens,
-                completion_tokens=final_completion_tokens,
-                model=model_name,
-                model_parameters=get_model_parameters(kwargs),
-                raw_output=raw_outputs,
-                id=inference_id,
-                cost=final_cost,  # Use calculated cost
-                metadata={
-                    "timeToFirstToken": ((first_token_time - start_time) * 1000 if first_token_time else None),
-                    "provider": provider,
-                    "litellm_model": model_name,
-                    **latest_chunk_metadata,  # Add all LiteLLM-specific metadata
-                },
-            )
-            captured_context.run(add_to_trace, **trace_args)
-
         # pylint: disable=broad-except
         except Exception as e:
-            if logger is not None:
-                logger.error(
-                    "Failed to trace the LiteLLM completion request with Openlayer. %s",
-                    e,
+            logger.error("Failed to yield chunk. %s", e)
+        finally:
+            # Update step with final data before context manager exits
+            try:
+                collected_output_data = [message for message in collected_output_data if message is not None]
+                if collected_output_data:
+                    output_data = "".join(collected_output_data)
+                else:
+                    if collected_function_call["arguments"]:
+                        try:
+                            collected_function_call["arguments"] = json.loads(collected_function_call["arguments"])
+                        except json.JSONDecodeError:
+                            pass
+                    output_data = collected_function_call
+
+                # Post-streaming calculations (after streaming is finished)
+                completion_tokens_calculated, prompt_tokens_calculated, total_tokens_calculated, cost_calculated = calculate_streaming_usage_and_cost(
+                    chunks=raw_outputs,
+                    messages=kwargs.get("messages", []),
+                    output_content=output_data,
+                    model_name=model_name,
+                    latest_usage_data=latest_usage_data,
+                    latest_chunk_metadata=latest_chunk_metadata
                 )
+                
+                # Use calculated values (fall back to extracted data if calculation fails)
+                usage_data = latest_usage_data if any(v is not None for v in latest_usage_data.values()) else {}
+                
+                final_prompt_tokens = prompt_tokens_calculated if prompt_tokens_calculated is not None else usage_data.get("prompt_tokens", 0)
+                final_completion_tokens = completion_tokens_calculated if completion_tokens_calculated is not None else usage_data.get("completion_tokens", num_of_completion_tokens)
+                final_total_tokens = total_tokens_calculated if total_tokens_calculated is not None else usage_data.get("total_tokens", final_prompt_tokens + final_completion_tokens)
+                final_cost = cost_calculated if cost_calculated is not None else latest_chunk_metadata.get('cost', None)
+                
+                # Update the step with final trace data
+                step.log(
+                    output=output_data,
+                    latency=latency,
+                    tokens=final_total_tokens,
+                    prompt_tokens=final_prompt_tokens,
+                    completion_tokens=final_completion_tokens,
+                    model=model_name,
+                    model_parameters=get_model_parameters(kwargs),
+                    raw_output=raw_outputs,
+                    id=inference_id,
+                    cost=final_cost,
+                    provider=provider,
+                    metadata={
+                        "timeToFirstToken": ((first_token_time - start_time) * 1000 if first_token_time else None),
+                        "provider": provider,
+                        "litellm_model": model_name,
+                        **latest_chunk_metadata,
+                    },
+                )
+
+            # pylint: disable=broad-except
+            except Exception as e:
+                if logger is not None:
+                    logger.error(
+                        "Failed to trace the LiteLLM completion request with Openlayer. %s",
+                        e,
+                    )
 
 
 def handle_non_streaming_completion(

src/openlayer/lib/tracing/tracer.py

@@ -1322,10 +1322,6 @@ def _create_and_initialize_step(
     parent_step = get_current_step()
     is_root_step = parent_step is None
 
-    # #region agent log - Debug: step creation
-    print(f"[OPENLAYER_DEBUG] tracer.py:_create_and_initialize_step | step_name={step_name} | step_type={step_type} | is_root_step={is_root_step} | parent_step_name={parent_step.name if parent_step else None}", flush=True)
-    # #endregion
-
     if parent_step is None:
         logger.debug("Starting a new trace...")
         current_trace = traces.Trace()
@@ -1349,16 +1345,6 @@ def _handle_trace_completion(
     on_flush_failure: Optional[OnFlushFailureCallback] = None,
 ) -> None:
     """Handle trace completion and data streaming."""
-    # #region agent log - Debug: trace completion
-    _trace = get_current_trace()
-    _steps = [s.name for s in _trace.steps] if _trace and _trace.steps else []
-    _nested = []
-    if _trace and _trace.steps:
-        for s in _trace.steps:
-            if hasattr(s, 'steps') and s.steps:
-                _nested.extend([ns.name for ns in s.steps])
-    print(f"[OPENLAYER_DEBUG] tracer.py:_handle_trace_completion | step_name={step_name} | is_root_step={is_root_step} | has_trace={_trace is not None} | root_steps={_steps} | nested_steps={_nested}", flush=True)
-    # #endregion
     if is_root_step:
         logger.debug("Ending the trace...")
         current_trace = get_current_trace()