ragaai-catalyst 2.1.4b0__py3-none-any.whl → 2.1.4b2__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py

@@ -1,3 +1,4 @@
+import os
 import functools
 import uuid
 from datetime import datetime
@@ -7,6 +8,15 @@ from ..utils.unique_decorator import mydecorator, generate_unique_hash_simple
 import contextvars
 import asyncio
 from ..utils.file_name_tracker import TrackName
+from ..utils.span_attributes import SpanAttributes
+import logging
+
+logger = logging.getLogger(__name__)
+logging_level = (
+    logger.setLevel(logging.DEBUG)
+    if os.getenv("DEBUG")
+    else logger.setLevel(logging.INFO)
+)
 
 
 class AgentTracerMixin:
@@ -16,42 +26,89 @@ class AgentTracerMixin:
         self.current_agent_id = contextvars.ContextVar("agent_id", default=None)
         self.current_agent_name = contextvars.ContextVar("agent_name", default=None)
         self.agent_children = contextvars.ContextVar("agent_children", default=[])
-        self.component_network_calls = contextvars.ContextVar("component_network_calls", default={})
-        self.component_user_interaction = contextvars.ContextVar("component_user_interaction", default={})
-        self.version = contextvars.ContextVar("version", default="1.0.0")
+        self.component_network_calls = contextvars.ContextVar(
+            "component_network_calls", default={}
+        )
+        self.component_user_interaction = contextvars.ContextVar(
+            "component_user_interaction", default={}
+        )
+        self.version = contextvars.ContextVar("version", default=None)
         self.agent_type = contextvars.ContextVar("agent_type", default="generic")
         self.capabilities = contextvars.ContextVar("capabilities", default=[])
         self.start_time = contextvars.ContextVar("start_time", default=None)
         self.input_data = contextvars.ContextVar("input_data", default=None)
         self.gt = None
 
+        self.span_attributes_dict = {}
+
+        # Add auto instrument flags
+        self.auto_instrument_agent = False
+        self.auto_instrument_user_interaction = False
+        self.auto_instrument_network = False
+
+    def trace_agent(
+        self,
+        name: str,
+        agent_type: str = None,
+        version: str = None,
+        capabilities: List[str] = None,
+        tags: List[str] = [],
+        metadata: Dict[str, Any] = {},
+        metrics: List[Dict[str, Any]] = [],
+        feedback: Optional[Any] = None,
+    ):
+        if name not in self.span_attributes_dict:
+            self.span_attributes_dict[name] = SpanAttributes(name)
+        if tags:
+            self.span(name).add_tags(tags)
+        if metadata:
+            self.span(name).add_metadata(metadata)
+        if metrics:
+            if isinstance(metrics, dict):
+                metrics = [metrics]
+            for metric in metrics:
+                try:
+                    self.span(name).add_metrics(
+                        name=metric["name"],
+                        score=metric["score"],
+                        reasoning=metric.get("reasoning", ""),
+                        cost=metric.get("cost", None),
+                        latency=metric.get("latency", None),
+                        metadata=metric.get("metadata", {}),
+                        config=metric.get("config", {}),
+                    )
+                except KeyError as e:
+                    logger.error(f"Error adding metric: {e}")
+        if feedback:
+            self.span(name).add_feedback(feedback)
 
-    def trace_agent(self, name: str, agent_type: str = None, version: str = None, capabilities: List[str] = None):
         def decorator(target):
             # Check if target is a class
             is_class = isinstance(target, type)
             tracer = self  # Store reference to tracer instance
-            top_level_hash_id = generate_unique_hash_simple(target)   # Generate hash based on the decorated target code
+            top_level_hash_id = generate_unique_hash_simple(
+                target
+            )  # Generate hash based on the decorated target code
             self.version.set(version)
             self.agent_type.set(agent_type)
             self.capabilities.set(capabilities)
-            
+
             if is_class:
                 # Store original __init__
                 original_init = target.__init__
-                
+
                 def wrapped_init(self, *args, **kwargs):
-                    self.gt = kwargs.get('gt', None) if kwargs else None
+                    self.gt = kwargs.get("gt", None) if kwargs else None
                     # Set agent context before initializing
                     component_id = str(uuid.uuid4())
                     hash_id = top_level_hash_id
-                    
+
                     # Store the component ID in the instance
                     self._agent_component_id = component_id
-                    
+
                     # Get parent agent ID if exists
                     parent_agent_id = tracer.current_agent_id.get()
-                    
+
                     # Create agent component
                     agent_component = tracer.create_agent_component(
                         component_id=component_id,
@@ -65,14 +122,14 @@ class AgentTracerMixin:
                         input_data=tracer._sanitize_input(args, kwargs),
                         output_data=None,
                         children=[],
-                        parent_id=parent_agent_id
+                        parent_id=parent_agent_id,
                     )
-                      
+
                     # Store component for later updates
-                    if not hasattr(tracer, '_agent_components'):
+                    if not hasattr(tracer, "_agent_components"):
                         tracer._agent_components = {}
                     tracer._agent_components[component_id] = agent_component
-                    
+
                     # If this is a nested agent, add it to parent's children
                     if parent_agent_id:
                         parent_children = tracer.agent_children.get()
@@ -81,61 +138,82 @@ class AgentTracerMixin:
                     else:
                         # Only add to root components if no parent
                         tracer.add_component(agent_component)
-                    
+
                     # Call original __init__ with this agent as current
                     token = tracer.current_agent_id.set(component_id)
                     try:
                         original_init(self, *args, **kwargs)
                     finally:
                         tracer.current_agent_id.reset(token)
-                
+
                 # Wrap all public methods to track execution
                 for attr_name in dir(target):
-                    if not attr_name.startswith('_'):
+                    if not attr_name.startswith("_"):
                         attr_value = getattr(target, attr_name)
                         if callable(attr_value):
+
                             def wrap_method(method):
                                 @self.file_tracker.trace_decorator
                                 @functools.wraps(method)
                                 def wrapped_method(self, *args, **kwargs):
-                                    self.gt = kwargs.get('gt', None) if kwargs else None
+                                    self.gt = kwargs.get("gt", None) if kwargs else None
                                     # Set this agent as current during method execution
-                                    token = tracer.current_agent_id.set(self._agent_component_id)
-                                    
+                                    token = tracer.current_agent_id.set(
+                                        self._agent_component_id
+                                    )
+
                                     # Store parent's children before setting new empty list
                                     parent_children = tracer.agent_children.get()
                                     children_token = tracer.agent_children.set([])
-                                    
+
                                     try:
                                         start_time = datetime.now()
                                         result = method(self, *args, **kwargs)
-                                        
+
                                         # Update agent component with method result
-                                        if hasattr(tracer, '_agent_components'):
-                                            component = tracer._agent_components.get(self._agent_component_id)
+                                        if hasattr(tracer, "_agent_components"):
+                                            component = tracer._agent_components.get(
+                                                self._agent_component_id
+                                            )
                                             if component:
-                                                component['data']['output'] = tracer._sanitize_output(result)
-                                                component['data']['input'] = tracer._sanitize_input(args, kwargs)
-                                                component['start_time'] = start_time.isoformat()
-                                                
+                                                component["data"]["output"] = (
+                                                    tracer._sanitize_output(result)
+                                                )
+                                                component["data"]["input"] = (
+                                                    tracer._sanitize_input(args, kwargs)
+                                                )
+                                                component["start_time"] = (
+                                                    start_time.isoformat()
+                                                )
+
                                                 # Get children accumulated during method execution
                                                 children = tracer.agent_children.get()
                                                 if children:
-                                                    if 'children' not in component['data']:
-                                                        component['data']['children'] = []
-                                                    component['data']['children'].extend(children)
-                                                    
+                                                    if (
+                                                        "children"
+                                                        not in component["data"]
+                                                    ):
+                                                        component["data"][
+                                                            "children"
+                                                        ] = []
+                                                    component["data"][
+                                                        "children"
+                                                    ].extend(children)
+
                                                     # Add this component as a child to parent's children list
                                                     parent_children.append(component)
-                                                    tracer.agent_children.set(parent_children)
+                                                    tracer.agent_children.set(
+                                                        parent_children
+                                                    )
                                         return result
                                     finally:
                                         tracer.current_agent_id.reset(token)
                                         tracer.agent_children.reset(children_token)
+
                                 return wrapped_method
-                            
+
                             setattr(target, attr_name, wrap_method(attr_value))
-                
+
                 # Replace __init__ with wrapped version
                 target.__init__ = wrapped_init
                 return target
@@ -143,17 +221,40 @@ class AgentTracerMixin:
                 # For function decorators, use existing sync/async tracing
                 is_async = asyncio.iscoroutinefunction(target)
                 if is_async:
+
                     async def wrapper(*args, **kwargs):
-                        return await self._trace_agent_execution(target, name, agent_type, version, capabilities, top_level_hash_id, *args, **kwargs)
+                        return await self._trace_agent_execution(
+                            target,
+                            name,
+                            agent_type,
+                            version,
+                            capabilities,
+                            top_level_hash_id,
+                            *args,
+                            **kwargs,
+                        )
+
                     return wrapper
                 else:
+
                     def wrapper(*args, **kwargs):
-                        return self._trace_sync_agent_execution(target, name, agent_type, version, capabilities, *args, **kwargs)
+                        return self._trace_sync_agent_execution(
+                            target,
+                            name,
+                            agent_type,
+                            version,
+                            capabilities,
+                            *args,
+                            **kwargs,
+                        )
+
                     return wrapper
 
         return decorator
 
-    def _trace_sync_agent_execution(self, func, name, agent_type, version, capabilities, *args, **kwargs):
+    def _trace_sync_agent_execution(
+        self, func, name, agent_type, version, capabilities, *args, **kwargs
+    ):
         # Generate a unique hash_id for this execution context
         hash_id = str(uuid.uuid4())
 
@@ -161,6 +262,9 @@ class AgentTracerMixin:
         if not self.is_active:
             return func(*args, **kwargs)
 
+        if not self.auto_instrument_agent:
+            return func(*args, **kwargs)
+
         start_time = datetime.now()
         self.start_time = start_time
         self.input_data = self._sanitize_input(args, kwargs)
@@ -168,11 +272,11 @@ class AgentTracerMixin:
         component_id = str(uuid.uuid4())
 
         # Extract ground truth if present
-        ground_truth = kwargs.pop('gt', None) if kwargs else None
+        ground_truth = kwargs.pop("gt", None) if kwargs else None
 
         # Get parent agent ID if exists
         parent_agent_id = self.current_agent_id.get()
-        
+
         # Set the current agent context
         agent_token = self.current_agent_id.set(component_id)
         agent_name_token = self.current_agent_name.set(name)
@@ -180,7 +284,7 @@ class AgentTracerMixin:
         # Initialize empty children list for this agent
         parent_children = self.agent_children.get()
         children_token = self.agent_children.set([])
-        
+
         # Start tracking network calls for this component
         self.start_component(component_id)
 
@@ -211,12 +315,12 @@ class AgentTracerMixin:
                 input_data=self.input_data,
                 output_data=self._sanitize_output(result),
                 children=children,
-                parent_id=parent_agent_id
+                parent_id=parent_agent_id,
             )
             # Add ground truth to component data if present
             if ground_truth is not None:
                 agent_component["data"]["gt"] = ground_truth
-                        
+
             # Add this component as a child to parent's children list
             parent_children.append(agent_component)
             self.agent_children.set(parent_children)
@@ -231,19 +335,19 @@ class AgentTracerMixin:
                 "code": 500,
                 "type": type(e).__name__,
                 "message": str(e),
-                "details": {}
+                "details": {},
             }
-            
+
             # Get children even in case of error
             children = self.agent_children.get()
-            
+
             # Set parent_id for all children
             for child in children:
                 child["parent_id"] = component_id
-            
+
             # End tracking network calls for this component
             self.end_component(component_id)
-            
+
             agent_component = self.create_agent_component(
                 component_id=component_id,
                 hash_id=hash_id,
@@ -257,39 +361,45 @@ class AgentTracerMixin:
                 output_data=None,
                 error=error_component,
                 children=children,
-                parent_id=parent_agent_id  # Add parent ID if exists
+                parent_id=parent_agent_id,  # Add parent ID if exists
             )
             # If this is a nested agent, add it to parent's children
             if parent_agent_id:
                 parent_component = self._agent_components.get(parent_agent_id)
                 if parent_component:
-                    if 'children' not in parent_component['data']:
-                        parent_component['data']['children'] = []
-                    parent_component['data']['children'].append(agent_component)
+                    if "children" not in parent_component["data"]:
+                        parent_component["data"]["children"] = []
+                    parent_component["data"]["children"].append(agent_component)
             else:
                 # Only add to root components if no parent
                 self.add_component(agent_component)
+
             raise
         finally:
             self.current_agent_id.reset(agent_token)
             self.current_agent_name.reset(agent_name_token)
             self.agent_children.reset(children_token)
 
-    async def _trace_agent_execution(self, func, name, agent_type, version, capabilities, hash_id, *args, **kwargs):
+    async def _trace_agent_execution(
+        self, func, name, agent_type, version, capabilities, hash_id, *args, **kwargs
+    ):
         """Asynchronous version of agent tracing"""
         if not self.is_active:
             return await func(*args, **kwargs)
 
+        if not self.auto_instrument_agent:
+            return await func(*args, **kwargs)
+
         start_time = datetime.now()
         start_memory = psutil.Process().memory_info().rss
         component_id = str(uuid.uuid4())
 
         # Extract ground truth if present
-        ground_truth = kwargs.pop('gt', None) if kwargs else None
+        ground_truth = kwargs.pop("gt", None) if kwargs else None
 
         # Get parent agent ID if exists
         parent_agent_id = self.current_agent_id.get()
-        
+
         # Set the current agent context
         agent_token = self.current_agent_id.set(component_id)
         agent_name_token = self.current_agent_name.set(name)
@@ -297,6 +407,7 @@ class AgentTracerMixin:
         # Initialize empty children list for this agent
         parent_children = self.agent_children.get()
         children_token = self.agent_children.set([])
+        self.start_component(component_id)
 
         try:
             # Execute the agent
@@ -309,6 +420,8 @@ class AgentTracerMixin:
             # Get children components collected during execution
             children = self.agent_children.get()
 
+            self.end_component(component_id)
+
             # Create agent component with children and parent if exists
             agent_component = self.create_agent_component(
                 component_id=component_id,
@@ -322,7 +435,7 @@ class AgentTracerMixin:
                 input_data=self._sanitize_input(args, kwargs),
                 output_data=self._sanitize_output(result),
                 children=children,
-                parent_id=parent_agent_id
+                parent_id=parent_agent_id,
             )
 
             # Add ground truth to component data if present
@@ -343,19 +456,19 @@ class AgentTracerMixin:
                 "code": 500,
                 "type": type(e).__name__,
                 "message": str(e),
-                "details": {}
+                "details": {},
             }
-            
+
             # Get children even in case of error
             children = self.agent_children.get()
-            
+
             # Set parent_id for all children
             for child in children:
                 child["parent_id"] = component_id
-            
+
             # End tracking network calls for this component
             self.end_component(component_id)
-            
+
             agent_component = self.create_agent_component(
                 component_id=component_id,
                 hash_id=hash_id,
@@ -369,19 +482,20 @@ class AgentTracerMixin:
                 output_data=None,
                 error=error_component,
                 children=children,
-                parent_id=parent_agent_id  # Add parent ID if exists
+                parent_id=parent_agent_id,  # Add parent ID if exists
             )
-            
+
             # If this is a nested agent, add it to parent's children
             if parent_agent_id:
                 parent_component = self._agent_components.get(parent_agent_id)
                 if parent_component:
-                    if 'children' not in parent_component['data']:
-                        parent_component['data']['children'] = []
-                    parent_component['data']['children'].append(agent_component)
+                    if "children" not in parent_component["data"]:
+                        parent_component["data"]["children"] = []
+                    parent_component["data"]["children"].append(agent_component)
             else:
                 # Only add to root components if no parent
                 self.add_component(agent_component)
+
             raise
         finally:
             # Reset context variables
@@ -391,7 +505,28 @@ class AgentTracerMixin:
 
     def create_agent_component(self, **kwargs):
         """Create an agent component according to the data structure"""
+        network_calls = []
+        if self.auto_instrument_network:
+            network_calls = self.component_network_calls.get(kwargs["component_id"], [])
+        interactions = []
+        if self.auto_instrument_user_interaction:
+            interactions = self.component_user_interaction.get(
+                kwargs["component_id"], []
+            )
         start_time = kwargs["start_time"]
+
+        # Get tags, metrics
+        name = kwargs["name"]
+        # tags
+        tags = []
+        if name in self.span_attributes_dict:
+            tags = self.span_attributes_dict[name].tags or []
+
+        # metrics
+        metrics = []
+        if name in self.span_attributes_dict:
+            metrics = self.span_attributes_dict[name].metrics or []
+
         component = {
             "id": kwargs["component_id"],
             "hash_id": kwargs["hash_id"],
@@ -406,20 +541,25 @@ class AgentTracerMixin:
                 "agent_type": kwargs["agent_type"],
                 "version": kwargs["version"],
                 "capabilities": kwargs["capabilities"],
-                "memory_used": kwargs["memory_used"]
+                "memory_used": kwargs["memory_used"],
+                "tags": tags,
             },
             "data": {
                 "input": kwargs["input_data"],
                 "output": kwargs["output_data"],
-                "children": kwargs.get("children", [])
+                "children": kwargs.get("children", []),
             },
-            "network_calls": self.component_network_calls.get(kwargs["component_id"], []),
-            "interactions": self.component_user_interaction.get(kwargs["component_id"], [])
+            "metrics": metrics,
+            "network_calls": network_calls,
+            "interactions": interactions,
         }
 
-        if self.gt: 
+        if self.gt:
             component["data"]["gt"] = self.gt
 
+        # Reset the SpanAttributes context variable
+        self.span_attributes_dict[kwargs["name"]] = SpanAttributes(kwargs["name"])
+
         return component
 
     def start_component(self, component_id):
@@ -438,21 +578,22 @@ class AgentTracerMixin:
 
     def _sanitize_input(self, args: tuple, kwargs: dict) -> str:
         """Convert input arguments to text format.
-        
+
         Args:
             args: Input arguments tuple
             kwargs: Input keyword arguments dict
-            
+
         Returns:
             str: Text representation of the input arguments
         """
+
         def _sanitize_value(value):
             if isinstance(value, dict):
                 return str({k: _sanitize_value(v) for k, v in value.items()})
             elif isinstance(value, (list, tuple)):
                 return str([_sanitize_value(item) for item in value])
             return str(value)
-        
+
         sanitized_args = [_sanitize_value(arg) for arg in args]
         sanitized_kwargs = {k: _sanitize_value(v) for k, v in kwargs.items()}
         return str({"args": sanitized_args, "kwargs": sanitized_kwargs})
@@ -461,4 +602,13 @@ class AgentTracerMixin:
         """Sanitize and format output data"""
         if isinstance(output, (int, float, bool, str, list, dict)):
             return output
-        return str(output)
+        return str(output)
+
+    def instrument_agent_calls(self):
+        self.auto_instrument_agent = True
+
+    def instrument_user_interaction_calls(self):
+        self.auto_instrument_user_interaction = True
+
+    def instrument_network_calls(self):
+        self.auto_instrument_network = True