Refactor trace metadata handling for improved flexibility

- Updated `update_current_trace` function to accept dynamic keyword arguments, simplifying the process of updating trace metadata.
- Enhanced `Trace` class to merge new metadata with existing fields, allowing for more efficient metadata management.

This refactor aims to streamline trace updates and improve code maintainability.

Author: viniciusdsmello

src/openlayer/lib/__init__.py

@@ -31,11 +31,6 @@
 update_current_span = tracer.update_current_span
 
 
-# --------------------------------- OCI GenAI -------------------------------- #
-# Alias for backward compatibility
-trace_oci = trace_oci_genai
-
-
 def trace_anthropic(client):
     """Trace Anthropic chat completions."""
     # pylint: disable=import-outside-toplevel
@@ -147,3 +142,8 @@ def trace_oci_genai(client, estimate_tokens: bool = True):
         raise ValueError("Invalid client. Please provide an OCI GenAI client.")
 
     return oci_tracer.trace_oci_genai(client, estimate_tokens=estimate_tokens)
+
+
+# --------------------------------- OCI GenAI -------------------------------- #
+# Alias for backward compatibility
+trace_oci = trace_oci_genai

src/openlayer/lib/tracing/tracer.py

@@ -526,42 +526,24 @@ def log_context(context: List[str]) -> None:
         logger.warning("No current step found to log context.")
 
 
-def update_current_trace(
-    name: Optional[str] = None,
-    tags: Optional[List[str]] = None,
-    metadata: Optional[Dict[str, Any]] = None,
-    thread_id: Optional[str] = None,
-    user_id: Optional[str] = None,
-    input: Optional[Any] = None,
-    output: Optional[Any] = None,
-    feedback: Optional['traces.Feedback'] = None,
-    test_case: Optional['traces.LLMTestCase'] = None,
-) -> None:
+def update_current_trace(**kwargs) -> None:
     """Updates the current trace metadata with the provided values.
     
     This function allows users to set trace-level metadata dynamically
     during execution without having to pass it through function arguments.
     
-    Args:
-        name: Optional trace name
-        tags: Optional list of tags for the trace
-        metadata: Optional dictionary of metadata to merge with existing metadata
-        thread_id: Optional thread identifier
-        user_id: Optional user identifier 
-        input: Optional trace input data
-        output: Optional trace output data
-        feedback: Optional feedback data
-        test_case: Optional LLM test case data
+    All provided key-value pairs will be stored in the trace metadata.
         
     Example:
-        >>> import openlayer
+        >>> from openlayer.lib import trace, update_current_trace
         >>> 
-        >>> @openlayer.trace()
+        >>> @trace()
         >>> def my_function():
         >>>     # Update trace with user context
-        >>>     openlayer.update_current_trace(
+        >>>     update_current_trace(
         >>>         user_id="user123",
-        >>>         metadata={"session_id": "sess456"}
+        >>>         session_id="sess456",
+        >>>         custom_field="any_value"
         >>>     )
         >>>     return "result"
     """
@@ -574,17 +556,7 @@ def update_current_trace(
         )
         return
 
-    current_trace.update_metadata(
-        name=name,
-        tags=tags,
-        metadata=metadata,
-        thread_id=thread_id,
-        user_id=user_id,
-        input=input,
-        output=output,
-        feedback=feedback,
-        test_case=test_case,
-    )
+    current_trace.update_metadata(**kwargs)
     logger.debug("Updated current trace metadata")
 
 
@@ -934,25 +906,9 @@ def post_process_trace(
     }
 
     # Include trace-level metadata if set
-    if trace_obj.name is not None:
-        trace_data["trace_name"] = trace_obj.name
-    if trace_obj.tags is not None:
-        trace_data["tags"] = trace_obj.tags
     if trace_obj.metadata is not None:
         # Merge trace-level metadata (higher precedence than root step metadata)
         trace_data.update(trace_obj.metadata)
-    if trace_obj.thread_id is not None:
-        trace_data["thread_id"] = trace_obj.thread_id
-    if trace_obj.user_id is not None:
-        trace_data["user_id"] = trace_obj.user_id
-    if trace_obj.input is not None:
-        trace_data["trace_input"] = trace_obj.input
-    if trace_obj.output is not None:
-        trace_data["trace_output"] = trace_obj.output
-    if trace_obj.feedback is not None:
-        trace_data["feedback"] = trace_obj.feedback
-    if trace_obj.test_case is not None:
-        trace_data["test_case"] = trace_obj.test_case
 
     if root_step.ground_truth:
         trace_data["groundTruth"] = root_step.ground_truth

src/openlayer/lib/tracing/traces.py

@@ -20,56 +20,32 @@ class Trace:
     def __init__(self):
         self.steps = []
         self.current_step = None
-        
-        # Enhanced trace metadata fields
-        self.name: Optional[str] = None
-        self.tags: Optional[List[str]] = None
         self.metadata: Optional[Dict[str, Any]] = None
-        self.thread_id: Optional[str] = None
-        self.user_id: Optional[str] = None
-        self.input: Optional[Any] = None
-        self.output: Optional[Any] = None
-        self.feedback: Optional[Feedback] = None
-        self.test_case: Optional[LLMTestCase] = None
 
     def add_step(self, step: Step) -> None:
         """Adds a step to the trace."""
         self.steps.append(step)
 
-    def update_metadata(
-        self,
-        name: Optional[str] = None,
-        tags: Optional[List[str]] = None,
-        metadata: Optional[Dict[str, Any]] = None,
-        thread_id: Optional[str] = None,
-        user_id: Optional[str] = None,
-        input: Optional[Any] = None,
-        output: Optional[Any] = None,
-        feedback: Optional[Feedback] = None,
-        test_case: Optional[LLMTestCase] = None,
-    ) -> None:
-        """Updates the trace metadata with the provided values."""
-        if name is not None:
-            self.name = name
-        if tags is not None:
-            self.tags = tags
-        if metadata is not None:
-            # Merge with existing metadata if it exists
-            if self.metadata is None:
-                self.metadata = {}
-            self.metadata.update(metadata)
-        if thread_id is not None:
-            self.thread_id = thread_id
-        if user_id is not None:
-            self.user_id = user_id
-        if input is not None:
-            self.input = input
-        if output is not None:
-            self.output = output
-        if feedback is not None:
-            self.feedback = feedback
-        if test_case is not None:
-            self.test_case = test_case
+    def update_metadata(self, **kwargs) -> None:
+        """Updates the trace metadata with the provided values.
+        
+        All provided key-value pairs will be stored in self.metadata.
+        Special handling for 'metadata' key which gets merged with existing metadata.
+        """
+        # Initialize metadata if it doesn't exist
+        if self.metadata is None:
+            self.metadata = {}
+        
+        # Handle special case for 'metadata' key - merge with existing
+        if 'metadata' in kwargs:
+            metadata_to_merge = kwargs.pop('metadata')
+            if metadata_to_merge is not None:
+                self.metadata.update(metadata_to_merge)
+        
+        # Add all other kwargs to metadata
+        for key, value in kwargs.items():
+            if value is not None:
+                self.metadata[key] = value
 
     def to_dict(self) -> List[Dict[str, Any]]:
         """Dictionary representation of the Trace."""