fix(tracer): handle missing active trace and improve error logging in LiteLLM integration

Author: viniciusdsmello

src/openlayer/lib/integrations/litellm_tracer.py

@@ -1,5 +1,6 @@
 """Module with methods used to trace LiteLLM completions."""
 
+import contextvars
 import json
 import logging
 import time
@@ -153,6 +154,7 @@ 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
@@ -253,14 +255,15 @@ def stream_chunks(
                     **latest_chunk_metadata,  # Add all LiteLLM-specific metadata
                 },
             )
-            add_to_trace(**trace_args)
+            captured_context.run(add_to_trace, **trace_args)
 
         # pylint: disable=broad-except
         except Exception as e:
-            logger.error(
-                "Failed to trace the LiteLLM completion request with Openlayer. %s",
-                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

@@ -1349,6 +1349,14 @@ def _handle_trace_completion(
         logger.debug("Ending the trace...")
         current_trace = get_current_trace()
 
+        if current_trace is None:
+            logger.warning(
+                "Cannot complete trace for step '%s': no active trace found. "
+                "This can happen when OPENLAYER_DISABLE_PUBLISH=true or trace context was lost.",
+                step_name,
+            )
+            return
+
         trace_data, input_variable_names = post_process_trace(current_trace)
 
         config = dict(
@@ -1644,7 +1652,7 @@ async def _invoke_with_context(
 
 
 def post_process_trace(
-    trace_obj: traces.Trace,
+    trace_obj: Optional[traces.Trace],
 ) -> Tuple[Dict[str, Any], List[str]]:
     """Post processing of the trace data before uploading to Openlayer.