diff --git a/examples/tracing/trace_metadata_updates.py b/examples/tracing/trace_metadata_updates.py index 67660d89..2f37e9a4 100644 --- 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,12 +177,22 @@ 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(), } ) @@ -187,13 +200,6 @@ def batch_processing_example(): 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 index a616eb2b..5a15a243 100644 --- 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 index 2f483fc5..bb5b8e5c 100644 --- 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')