feat(tracing): enable custom inferenceId for trace decorator

viniciusdsmello · whoseoyster · commit 1f01ca5e39dc · 2025-09-08T21:22:57.000-07:00
- Add dedicated inference_id field to Trace class instead of storing in metadata
- Allow users to set custom inference ID via update_current_trace(inferenceId=...)
- Maintain backward compatibility with auto-generated UUID fallback
- Simplify post_process_trace logic by removing special metadata filtering
- Update examples to demonstrate custom inference ID usage patterns

BREAKING CHANGE: None - fully backward compatible

Closes: #OPEN-7312
diff --git a/examples/tracing/trace_metadata_updates.py b/examples/tracing/trace_metadata_updates.py
@@ -41,8 +41,11 @@ def handle_user_request(self, request_text: str, session_token: str) -> str:
         # Get user session (this info isn't available as function arguments)
         user_session = self.get_user_session(session_token)
         
-        # Set trace-level metadata with user context
+        # Set trace-level metadata with user context and custom inference ID
+        custom_inference_id = f"chat_{user_session.user_id}_{user_session.interaction_count}_{int(datetime.now().timestamp())}"
+        
         update_current_trace(
+            inferenceId=custom_inference_id,
             name=f"chat_request_{user_session.user_id}",
             user_id=user_session.user_id,
             tags=["chat", "user_request", user_session.preferences.get("tier", "free")],
@@ -174,26 +177,29 @@ def make_formal(self, text: str) -> str:
 def batch_processing_example():
     """Example showing batch processing with trace metadata updates."""
     
-    # Set trace metadata for batch job
+    # Process multiple requests
+    test_requests = [
+        ("Hello there!", "premium_session_123"),
+        ("What's the weather like?", "free_session_456"), 
+        ("Help me with coding", "premium_session_789")
+    ]
+    
+    # Set trace metadata for batch job with custom batch ID
+    batch_inference_id = f"batch_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{len(test_requests)}_items"
+    
     update_current_trace(
+        inferenceId=batch_inference_id,
         name="batch_user_requests",
         tags=["batch", "processing", "multiple_users"],
         metadata={
-            "batch_size": 3,
+            "batch_size": len(test_requests),
             "processing_start": datetime.now().isoformat(),
         }
     )
     
     app = ChatApplication()
     results = []
     
-    # Process multiple requests
-    test_requests = [
-        ("Hello there!", "premium_session_123"),
-        ("What's the weather like?", "free_session_456"), 
-        ("Help me with coding", "premium_session_789")
-    ]
-    
     for i, (request, session) in enumerate(test_requests):
         result = app.handle_user_request(request, session)
         results.append(result)
diff --git a/src/openlayer/lib/tracing/tracer.py b/src/openlayer/lib/tracing/tracer.py
@@ -779,6 +779,7 @@ def update_current_trace(**kwargs) -> None:
         >>> def my_function():
         >>>     # Update trace with user context
         >>>     update_current_trace(
+        >>>         inferenceId="custom_inference_id",
         >>>         user_id="user123",
         >>>         session_id="sess456",
         >>>         custom_field="any_value"
@@ -1170,7 +1171,7 @@ def post_process_trace(
 
     trace_data = {
         "inferenceTimestamp": root_step.start_time,
-        "inferenceId": str(root_step.id),
+        "inferenceId": trace_obj.inference_id or str(root_step.id),
         "output": root_step.output,
         "latency": root_step.latency,
         "cost": processed_steps[0].get("cost", 0),
diff --git a/src/openlayer/lib/tracing/traces.py b/src/openlayer/lib/tracing/traces.py
@@ -16,6 +16,7 @@ def __init__(self):
         self.steps = []
         self.current_step = None
         self.metadata: Optional[Dict[str, Any]] = None
+        self.inference_id: Optional[str] = None
 
     def add_step(self, step: Step) -> None:
         """Adds a step to the trace."""
@@ -26,11 +27,16 @@ def update_metadata(self, **kwargs) -> None:
         
         All provided key-value pairs will be stored in self.metadata.
         Special handling for 'metadata' key which gets merged with existing metadata.
+        Special handling for 'inferenceId' which gets stored in dedicated field.
         """
         # Initialize metadata if it doesn't exist
         if self.metadata is None:
             self.metadata = {}
         
+        # Handle special case for inferenceId - store in dedicated field
+        if 'inferenceId' in kwargs:
+            self.inference_id = kwargs.pop('inferenceId')
+        
         # Handle special case for 'metadata' key - merge with existing
         if 'metadata' in kwargs:
             metadata_to_merge = kwargs.pop('metadata')
