ragaai-catalyst 2.1.4.1b0__py3-none-any.whl → 2.1.5__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py

@@ -25,6 +25,7 @@ class CustomTracerMixin:
         self.auto_instrument_custom = False
         self.auto_instrument_user_interaction = False
         self.auto_instrument_network = False
+        self.auto_instrument_file_io = False
 
     def trace_custom(self, name: str = None, custom_type: str = "generic", version: str = "1.0.0", trace_variables: bool = True):
         def decorator(func):
@@ -40,20 +41,24 @@ class CustomTracerMixin:
             # Check if the function is async
             is_async = asyncio.iscoroutinefunction(func)
 
-            @self.file_tracker.trace_decorator
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
                 async_wrapper.metadata = metadata
-                self.gt = kwargs.get('gt', None) if kwargs else None
+                gt = kwargs.get('gt') if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return await self._trace_custom_execution(
                     func, name or func.__name__, custom_type, version, trace_variables, *args, **kwargs
                 )
 
-            @self.file_tracker.trace_decorator
             @functools.wraps(func)
             def sync_wrapper(*args, **kwargs):
                 sync_wrapper.metadata = metadata
-                self.gt = kwargs.get('gt', None) if kwargs else None
+                gt = kwargs.get('gt') if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return self._trace_sync_custom_execution(
                     func, name or func.__name__, custom_type, version, trace_variables, *args, **kwargs
                 )
@@ -153,7 +158,7 @@ class CustomTracerMixin:
                 error=error_component
             )
 
-            self.add_component(custom_component)
+            self.add_component(custom_component, is_error=True)
             raise
 
     async def _trace_custom_execution(self, func, name, custom_type, version, trace_variables, *args, **kwargs):
@@ -233,7 +238,7 @@ class CustomTracerMixin:
                 output_data=None,
                 error=error_component
             )
-            self.add_component(custom_component)
+            self.add_component(custom_component, is_error=True)
             raise
 
     def create_custom_component(self, **kwargs):
@@ -246,8 +251,18 @@ class CustomTracerMixin:
             
         interactions = []
         if self.auto_instrument_user_interaction:
-            interactions = self.component_user_interaction.get(kwargs["component_id"], [])
-
+            input_output_interactions = []
+            for interaction in self.component_user_interaction.get(kwargs["component_id"], []):
+                if interaction["interaction_type"] in ["input", "output"]:
+                    input_output_interactions.append(interaction)
+            interactions.extend(input_output_interactions) 
+        if self.auto_instrument_file_io:
+            file_io_interactions = []
+            for interaction in self.component_user_interaction.get(kwargs["component_id"], []):
+                if interaction["interaction_type"] in ["file_read", "file_write"]:
+                    file_io_interactions.append(interaction)
+            interactions.extend(file_io_interactions)
+            
         component = {
             "id": kwargs["component_id"],
             "hash_id": kwargs["hash_id"],
@@ -273,9 +288,13 @@ class CustomTracerMixin:
             "interactions": interactions
         }
 
-        if self.gt:
-            component["data"]["gt"] = self.gt
-
+        if kwargs["name"] in self.span_attributes_dict:
+            span_gt = self.span_attributes_dict[kwargs["name"]].gt
+            if span_gt is not None:
+                component["data"]["gt"] = span_gt
+            span_context = self.span_attributes_dict[kwargs["name"]].context
+            if span_context:
+                component["data"]["context"] = span_context
         return component
 
     def start_component(self, component_id):
@@ -286,15 +305,23 @@ class CustomTracerMixin:
         """End tracking network calls for a component"""
         pass
 
-    def _sanitize_input(self, args: tuple, kwargs: dict) -> Dict:
-        """Sanitize and format input data"""
-        return {
-            "args": [str(arg) if not isinstance(arg, (int, float, bool, str, list, dict)) else arg for arg in args],
-            "kwargs": {
-                k: str(v) if not isinstance(v, (int, float, bool, str, list, dict)) else v 
-                for k, v in kwargs.items()
+    def _sanitize_input(self, args: tuple, kwargs: dict) -> dict:
+            """Sanitize and format input data, including handling of nested lists and dictionaries."""
+
+            def sanitize_value(value):
+                if isinstance(value, (int, float, bool, str)):
+                    return value
+                elif isinstance(value, list):
+                    return [sanitize_value(item) for item in value]
+                elif isinstance(value, dict):
+                    return {key: sanitize_value(val) for key, val in value.items()}
+                else:
+                    return str(value)  # Convert non-standard types to string
+
+            return {
+                "args": [sanitize_value(arg) for arg in args],
+                "kwargs": {key: sanitize_value(val) for key, val in kwargs.items()},
             }
-        }
 
     def _sanitize_output(self, output: Any) -> Any:
         """Sanitize and format output data"""
@@ -314,3 +341,7 @@ class CustomTracerMixin:
     def instrument_network_calls(self):
         """Enable auto-instrumentation for network calls"""
         self.auto_instrument_network = True
+
+    def instrument_file_io_calls(self):
+        """Enable auto-instrumentation for file IO calls"""
+        self.auto_instrument_file_io = True