ragaai-catalyst 2.0.7.2__py3-none-any.whl → 2.0.7.2b0__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/__init__.py

@@ -0,0 +1,3 @@
+from .agentic_tracing import AgenticTracing
+
+__all__ = ['AgenticTracing']

ragaai_catalyst/tracers/agentic_tracing/agent_tracer.py

@@ -0,0 +1,311 @@
+import functools
+import uuid
+from datetime import datetime
+import psutil
+from typing import Optional, Any, Dict, List
+from .unique_decorator import mydecorator
+
+import contextvars
+import asyncio
+
+class AgentTracerMixin:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        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._trace_sync_agent_execution = mydecorator(self._trace_sync_agent_execution)
+        self._trace_agent_execution = mydecorator(self._trace_agent_execution)
+
+
+    def trace_agent(self, name: str, agent_type: str = "generic", version: str = "1.0.0", capabilities: List[str] = None):
+        def decorator(func):
+            # Check if the function is async
+            is_async = asyncio.iscoroutinefunction(func)
+
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                return await self._trace_agent_execution(
+                    func, name, agent_type, version, capabilities, *args, **kwargs
+                )
+
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                return self._trace_sync_agent_execution(
+                    func, name, agent_type, version, capabilities, *args, **kwargs
+                )
+
+            return async_wrapper if is_async else sync_wrapper
+
+        return decorator
+
+    def _trace_sync_agent_execution(self, func, name, agent_type, version, capabilities, *args, **kwargs):
+        """Synchronous version of agent tracing"""
+        if not self.is_active:
+            return func(*args, **kwargs)
+
+        start_time = datetime.now()
+        start_memory = psutil.Process().memory_info().rss
+        component_id = str(uuid.uuid4())
+        hash_id = self._trace_sync_agent_execution.hash_id
+
+        # Initialize empty children list for this agent
+        children_token = self.agent_children.set([])
+        
+        # Set the current agent context
+        agent_token = self.current_agent_id.set(component_id)
+        agent_name_token = self.current_agent_name.set(name)
+
+        # Start tracking network calls for this component
+        self.start_component(component_id)
+
+        try:
+            # Execute the agent
+            result = func(*args, **kwargs)
+
+            # Calculate resource usage
+            end_time = datetime.now()
+            end_memory = psutil.Process().memory_info().rss
+            memory_used = max(0, end_memory - start_memory)
+
+            # Get children components collected during execution
+            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)
+
+            # Create agent component with children
+            agent_component = self.create_agent_component(
+                component_id=component_id,
+                hash_id=hash_id,
+                name=name,
+                agent_type=agent_type,
+                version=version,
+                capabilities=capabilities or [],
+                start_time=start_time,
+                end_time=end_time,
+                memory_used=memory_used,
+                input_data=self._sanitize_input(args, kwargs),
+                output_data=self._sanitize_output(result),
+                children=children
+            )
+
+            self.add_component(agent_component)
+            return result
+
+        except Exception as e:
+            error_component = {
+                "code": 500,
+                "type": type(e).__name__,
+                "message": str(e),
+                "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,
+                name=name,
+                agent_type=agent_type,
+                version=version,
+                capabilities=capabilities or [],
+                start_time=start_time,
+                end_time=datetime.now(),
+                memory_used=0,
+                input_data=self._sanitize_input(args, kwargs),
+                output_data=None,
+                error=error_component,
+                children=children
+            )
+
+            self.add_component(agent_component)
+            raise
+        finally:
+            # Reset context variables
+            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, *args, **kwargs):
+        """Asynchronous version of agent tracing"""
+        if not self.is_active:
+            return await func(*args, **kwargs)
+
+        start_time = datetime.now()
+        start_memory = psutil.Process().memory_info().rss
+        component_id = str(uuid.uuid4())
+        hash_id = self._trace_agent_execution.hash_id
+
+        # Initialize empty children list for this agent
+        children_token = self.agent_children.set([])
+        
+        # Set the current agent context
+        agent_token = self.current_agent_id.set(component_id)
+        agent_name_token = self.current_agent_name.set(name)
+
+        # Start tracking network calls for this component
+        self.start_component(component_id)
+
+        try:
+            # Execute the agent
+            result = await func(*args, **kwargs)
+
+            # Calculate resource usage
+            end_time = datetime.now()
+            end_memory = psutil.Process().memory_info().rss
+            memory_used = max(0, end_memory - start_memory)
+
+            # Get children components collected during execution
+            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)
+
+            # Create agent component with children
+            agent_component = self.create_agent_component(
+                component_id=component_id,
+                hash_id=hash_id,
+                name=name,
+                agent_type=agent_type,
+                version=version,
+                capabilities=capabilities or [],
+                start_time=start_time,
+                end_time=end_time,
+                memory_used=memory_used,
+                input_data=self._sanitize_input(args, kwargs),
+                output_data=self._sanitize_output(result),
+                children=children
+            )
+
+            self.add_component(agent_component)
+            return result
+
+        except Exception as e:
+            error_component = {
+                "code": 500,
+                "type": type(e).__name__,
+                "message": str(e),
+                "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,
+                name=name,
+                agent_type=agent_type,
+                version=version,
+                capabilities=capabilities or [],
+                start_time=start_time,
+                end_time=datetime.now(),
+                memory_used=0,
+                input_data=self._sanitize_input(args, kwargs),
+                output_data=None,
+                error=error_component,
+                children=children
+            )
+
+            self.add_component(agent_component)
+            raise
+        finally:
+            # Reset context variables
+            self.current_agent_id.reset(agent_token)
+            self.current_agent_name.reset(agent_name_token)
+            self.agent_children.reset(children_token)
+
+    def create_agent_component(self, **kwargs):
+        """Create an agent component according to the data structure"""
+        start_time = kwargs["start_time"]
+        component = {
+            "id": kwargs["component_id"],
+            "hash_id": kwargs["hash_id"],
+            "source_hash_id": None,
+            "type": "agent",
+            "name": kwargs["name"],
+            "start_time": start_time.isoformat(),
+            "end_time": kwargs["end_time"].isoformat(),
+            "error": kwargs.get("error"),
+            "parent_id": None,
+            "info": {
+                "agent_type": kwargs["agent_type"],
+                "version": kwargs["version"],
+                "capabilities": kwargs["capabilities"],
+                "memory_used": kwargs["memory_used"]
+            },
+            "data": {
+                "input": kwargs["input_data"],
+                "output": kwargs["output_data"],
+                "children": kwargs.get("children", [])
+            },
+            "network_calls": self.component_network_calls.get(kwargs["component_id"], []),
+            "interactions": [{
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "input",
+                "timestamp": start_time.isoformat(),
+                "content": kwargs["input_data"]
+            }, {
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "output",
+                "timestamp": kwargs["end_time"].isoformat(),
+                "content": kwargs["output_data"]
+            }]
+        }
+
+        return component
+
+    def start_component(self, component_id):
+        """Start tracking network calls for a component"""
+        component_network_calls = self.component_network_calls.get()
+        if component_id not in component_network_calls:
+            component_network_calls[component_id] = []
+        self.component_network_calls.set(component_network_calls)
+
+    def end_component(self, component_id):
+        """End tracking network calls for a component"""
+        component_network_calls = self.component_network_calls.get()
+        if component_id in component_network_calls:
+            component_network_calls[component_id] = []
+        self.component_network_calls.set(component_network_calls)
+
+    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_output(self, output: Any) -> Any:
+        """Sanitize and format output data"""
+        if isinstance(output, (int, float, bool, str, list, dict)):
+            return output
+        return str(output)

ragaai_catalyst/tracers/agentic_tracing/agentic_tracing.py

@@ -0,0 +1,212 @@
+import contextvars
+from typing import Optional, Dict
+import json
+from datetime import datetime
+import uuid
+import os
+from pathlib import Path
+
+from .base import BaseTracer
+from .llm_tracer import LLMTracerMixin
+from .tool_tracer import ToolTracerMixin
+from .agent_tracer import AgentTracerMixin
+from .network_tracer import NetworkTracer
+
+from .data_structure import (
+    Trace, Metadata, SystemInfo, OSInfo, EnvironmentInfo,
+    Resources, CPUResource, MemoryResource, DiskResource, NetworkResource,
+    ResourceInfo, MemoryInfo, DiskInfo, NetworkInfo,
+    Component, LLMComponent, AgentComponent, ToolComponent,
+    NetworkCall, Interaction, Error
+)
+
+
+from ...ragaai_catalyst import RagaAICatalyst
+from ..upload_traces import UploadTraces
+
+class AgenticTracing(BaseTracer, LLMTracerMixin, ToolTracerMixin, AgentTracerMixin):
+    def __init__(self, user_detail, auto_instrument_llm: bool = True):
+        # Initialize all parent classes
+        LLMTracerMixin.__init__(self)
+        ToolTracerMixin.__init__(self)
+        AgentTracerMixin.__init__(self)
+        
+        self.project_name = user_detail["project_name"]
+        self.project_id = user_detail["project_id"]
+        self.dataset_name = user_detail["dataset_name"]
+        self.trace_user_detail = user_detail["trace_user_detail"]
+        self.base_url = f"{RagaAICatalyst.BASE_URL}"
+        self.timeout = 10
+
+        BaseTracer.__init__(self, user_detail)
+        
+        self.auto_instrument_llm = auto_instrument_llm
+        self.tools: Dict[str, Tool] = {}
+        self.call_depth = contextvars.ContextVar("call_depth", default=0)
+        self.network_tracer = NetworkTracer()
+        self.is_active = False
+        self.current_agent_id = contextvars.ContextVar("current_agent_id", default=None)
+        self.agent_children = contextvars.ContextVar("agent_children", default=[])
+        self.component_network_calls = {}  # Store network calls per component
+        
+        # Create output directory if it doesn't exist
+        self.output_dir = Path("./traces")  # Using default traces directory
+        self.output_dir.mkdir(exist_ok=True)
+
+    def start_component(self, component_id: str):
+        """Start tracking network calls for a component"""
+        self.component_network_calls[component_id] = []
+        self.network_tracer.network_calls = []  # Reset network calls
+
+    def end_component(self, component_id: str):
+        """End tracking network calls for a component"""
+        self.component_network_calls[component_id] = self.network_tracer.network_calls.copy()
+        self.network_tracer.network_calls = []  # Reset for next component
+
+    def start(self):
+        """Start tracing"""
+        # Start base tracer (includes system info and resource monitoring)
+        super().start()
+        self.is_active = True
+        
+        # Activate network tracing
+        self.network_tracer.activate_patches()
+        
+        # Instrument calls from mixins
+        if self.auto_instrument_llm:
+            self.instrument_llm_calls()
+
+    def stop(self):
+        """Stop tracing and save results"""
+        if self.is_active:
+            # Calculate final metrics before stopping
+            self._calculate_final_metrics()
+            
+            # Deactivate network tracing
+            self.network_tracer.deactivate_patches()
+            
+            # Stop base tracer (includes saving to file)
+            super().stop()
+            
+            # Cleanup
+            self.unpatch_llm_calls()
+            self.is_active = False
+
+
+    def _calculate_final_metrics(self):
+        """Calculate total cost and tokens from all components"""
+        total_cost = 0.0
+        total_tokens = 0
+        
+        def process_component(component):
+            nonlocal total_cost, total_tokens
+            # Convert component to dict if it's an object
+            comp_dict = component.__dict__ if hasattr(component, '__dict__') else component
+            
+            if comp_dict.get('type') == "llm":
+                info = comp_dict.get('info', {})
+                if isinstance(info, dict):
+                    # Extract cost
+                    cost_info = info.get('cost', {})
+                    if isinstance(cost_info, dict):
+                        total_cost += cost_info.get('total_cost', 0)
+                    
+                    # Extract tokens
+                    token_info = info.get('tokens', {})
+                    if isinstance(token_info, dict):
+                        total_tokens += token_info.get('total_tokens', 0)
+                    else:
+                        token_info = info.get('token_usage', {})
+                        if isinstance(token_info, dict):
+                            total_tokens += token_info.get('total_tokens', 0)
+            
+            # Process children if they exist
+            data = comp_dict.get('data', {})
+            if isinstance(data, dict):
+                children = data.get('children', [])
+                if children:
+                    for child in children:
+                        process_component(child)
+        
+        # Process all root components
+        for component in self.components:
+            process_component(component)
+        
+        # Update metadata in trace
+        if hasattr(self, 'trace'):
+            if isinstance(self.trace.metadata, dict):
+                self.trace.metadata['total_cost'] = total_cost
+                self.trace.metadata['total_tokens'] = total_tokens
+            else:
+                self.trace.metadata.total_cost = total_cost
+                self.trace.metadata.total_tokens = total_tokens
+
+    def add_component(self, component_data: dict):
+        """Add a component to the trace data"""
+        component_id = component_data["id"]
+        
+        # Convert dict to appropriate Component type
+        if component_data["type"] == "llm":
+            component = LLMComponent(**component_data)
+            # Add network calls specific to this LLM component
+            component_data["network_calls"] = self.component_network_calls.get(component_id, [])
+            # Add user interaction for LLM calls
+            component_data["interactions"] = [{
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "input",
+                "content": component_data["data"]["input"]
+            }, {
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "output",
+                "content": component_data["data"]["output"]
+            }]
+        elif component_data["type"] == "agent":
+            component = AgentComponent(**component_data)
+            # Add network calls specific to this agent component
+            component_data["network_calls"] = self.component_network_calls.get(component_id, [])
+            # Add user interaction for agent
+            component_data["interactions"] = [{
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "input",
+                "content": component_data["data"]["input"]
+            }, {
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "output",
+                "content": component_data["data"]["output"]
+            }]
+        elif component_data["type"] == "tool":
+            component = ToolComponent(**component_data)
+            # Add network calls specific to this tool component
+            component_data["network_calls"] = self.component_network_calls.get(component_id, [])
+            # Add user interaction for tool
+            component_data["interactions"] = [{
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "input",
+                "content": component_data["data"]["input"]
+            }, {
+                "id": f"int_{uuid.uuid4()}",
+                "interaction_type": "output",
+                "content": component_data["data"]["output"]
+            }]
+        else:
+            component = Component(**component_data)
+        
+        # Check if there's an active agent context
+        current_agent_id = self.current_agent_id.get()
+        if current_agent_id and component_data["type"] in ["llm", "tool"]:
+            # Add this component as a child of the current agent
+            current_children = self.agent_children.get()
+            current_children.append(component_data)
+            self.agent_children.set(current_children)
+        else:
+            # Add component to the main trace
+            super().add_component(component)
+
+    def __enter__(self):
+        """Context manager entry"""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Context manager exit"""
+        self.stop()