ragaai-catalyst 2.0.7.2b1__py3-none-any.whl → 2.1__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/base.py

@@ -1,25 +1,27 @@
 import json
 import os
 import platform
-import re
 import psutil
 import pkg_resources
 from datetime import datetime
 from pathlib import Path
-from typing import Optional, Dict, Any, List
+from typing import List
 import uuid
 import sys
+import tempfile
 
 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
+    Component, 
 )
 
-from ..upload_traces import UploadTraces
-from ...ragaai_catalyst import RagaAICatalyst
+from .upload_agentic_traces import UploadAgenticTraces
+from .upload_code import upload_code
+
+from .file_name_tracker import TrackName
+from .zip_list_of_unique_files import zip_list_of_unique_files
 
 class TracerJSONEncoder(json.JSONEncoder):
     def default(self, obj):
@@ -50,13 +52,10 @@ class BaseTracer:
         self.project_id = self.user_details['project_id']  # Access the project_id
         
         # Initialize trace data
-        self.trace_id = str(uuid.uuid4())
-        self.start_time = datetime.now().isoformat()
+        self.trace_id = None
+        self.start_time = None 
         self.components: List[Component] = []
-        self.data_key = [{"start_time": self.start_time, 
-                        "end_time": "",
-                        "spans": self.components
-                        }]
+        self.file_tracker = TrackName()
         
     def _get_system_info(self) -> SystemInfo:
         # Get OS info
@@ -146,11 +145,22 @@ class BaseTracer:
             system_info=self._get_system_info(),
             resources=self._get_resources()
         )
+
+        # Generate a unique trace ID, when trace starts
+        self.trace_id = str(uuid.uuid4()) 
+        
+        # Get the start time
+        self.start_time = datetime.now().isoformat()
+        
+        self.data_key = [{"start_time": datetime.now().isoformat(), 
+                        "end_time": "",
+                        "spans": self.components
+                        }]
         
         self.trace = Trace(
             id=self.trace_id,
             project_name=self.project_name,
-            start_time=self.start_time,
+            start_time=datetime.now().isoformat(),
             end_time="",  # Will be set when trace is stopped
             metadata=metadata,
             data=self.data_key,
@@ -165,29 +175,37 @@ class BaseTracer:
 
             # Change span ids to int
             self.trace = self._change_span_ids_to_int(self.trace)
-            self.trace = self._change_agent_intput_output(self.trace)
+            self.trace = self._change_agent_input_output(self.trace)
             self.trace = self._extract_cost_tokens(self.trace)
             
             # Create traces directory if it doesn't exist
-            self.traces_dir = Path("traces")
-            self.traces_dir.mkdir(exist_ok=True)
+            self.traces_dir = tempfile.gettempdir()
             filename = self.trace.id + ".json"
-            filepath = self.traces_dir / filename
-            
-            # Save to JSON file using custom encoder
+            filepath = f"{self.traces_dir}/{filename}"
+
+            #get unique files and zip it. Generate a unique hash ID for the contents of the files
+            list_of_unique_files = self.file_tracker.get_unique_files()
+            hash_id, zip_path = zip_list_of_unique_files(list_of_unique_files, output_dir=self.traces_dir)
+
+            #replace source code with zip_path
+            self.trace.metadata.system_info.source_code = hash_id
+
+            # Clean up trace_data before saving
+            trace_data = self.trace.__dict__
+            cleaned_trace_data = self._clean_trace(trace_data)
+
             with open(filepath, 'w') as f:
-                json.dump(self.trace.__dict__, f, cls=TracerJSONEncoder, indent=2)
+                json.dump(cleaned_trace_data, f, cls=TracerJSONEncoder, indent=2)
                 
             print(f"Trace saved to {filepath}")
-            # import pdb; pdb.set_trace()
             # Upload traces
             json_file_path = str(filepath)
             project_name = self.project_name
-            project_id = self.project_id  # TODO: Replace with actual project ID
+            project_id = self.project_id 
             dataset_name = self.dataset_name
             user_detail = self.user_details
             base_url = os.getenv('RAGAAI_CATALYST_BASE_URL')
-            upload_traces = UploadTraces(
+            upload_traces = UploadAgenticTraces(
                 json_file_path=json_file_path,
                 project_name=project_name,
                 project_id=project_id,
@@ -195,7 +213,20 @@ class BaseTracer:
                 user_detail=user_detail,
                 base_url=base_url
             )
-            upload_traces.upload_traces()
+            upload_traces.upload_agentic_traces()
+
+            #Upload Codehash
+            response = upload_code(
+                hash_id=hash_id,
+                zip_path=zip_path,
+                project_name=project_name,
+                dataset_name=dataset_name
+            )
+            print(response)
+            
+        # Cleanup
+        self.components = []
+        self.file_tracker.reset()
                 
     def add_component(self, component: Component):
         """Add a component to the trace"""
@@ -209,7 +240,6 @@ class BaseTracer:
         self.stop()
 
     def _change_span_ids_to_int(self, trace):
-        # import pdb; pdb.set_trace()
         id, parent_id = 1, 0
         for span in trace.data[0]["spans"]:
             span.id = id
@@ -222,14 +252,32 @@ class BaseTracer:
                     id += 1
         return trace
 
-    def _change_agent_intput_output(self, trace):
+    def _change_agent_input_output(self, trace):
         for span in trace.data[0]["spans"]:
             if span.type == "agent":
-                # import pdb; pdb.set_trace()
                 childrens = span.data["children"]
-                if childrens != []:
-                    span.data["input"] = childrens[0]["data"]["input"]
-                    span.data["output"] = childrens[-1]["data"]["output"]
+                span.data["input"] = None
+                span.data["output"] = None
+                if childrens:
+                    # Find first non-null input going forward
+                    for child in childrens:
+                        if "data" not in child:
+                            continue
+                        input_data = child["data"].get("input")
+
+                        if input_data:
+                            span.data["input"] = input_data['args'] if hasattr(input_data, 'args') else input_data
+                            break
+                    
+                    # Find first non-null output going backward
+                    for child in reversed(childrens):
+                        if "data" not in child:
+                            continue
+                        output_data = child["data"].get("output")
+                
+                        if output_data and output_data != "" and output_data != "None":
+                            span.data["output"] = output_data
+                            break
         return trace
     
     def _extract_cost_tokens(self, trace):
@@ -251,6 +299,8 @@ class BaseTracer:
                         tokens[key] += value
             if span.type == "agent":
                 for children in span.data["children"]:
+                    if 'type' not in children:
+                        continue
                     if children["type"] != "llm":
                         continue
                     info = children["info"]
@@ -267,4 +317,56 @@ class BaseTracer:
                             tokens[key] += value
         trace.metadata.cost = cost
         trace.metadata.tokens = tokens
+        return trace
+
+    def _clean_trace(self, trace):
+        # Convert span to dict if it has to_dict method
+        def _to_dict_if_needed(obj):
+            if hasattr(obj, 'to_dict'):
+                return obj.to_dict()
+            return obj
+
+        def deduplicate_spans(spans):
+            seen_llm_spans = {}  # Dictionary to track unique LLM spans
+            unique_spans = []
+            
+            for span in spans:
+                # Convert span to dictionary if needed
+                span_dict = _to_dict_if_needed(span)
+                
+                # Skip spans without hash_id
+                if 'hash_id' not in span_dict:
+                    continue
+                
+                if span_dict.get('type') == 'llm':
+                    # Create a unique key based on hash_id, input, and output
+                    span_key = (
+                        span_dict.get('hash_id'),
+                        str(span_dict.get('data', {}).get('input')),
+                        str(span_dict.get('data', {}).get('output'))
+                    )
+                    
+                    if span_key not in seen_llm_spans:
+                        seen_llm_spans[span_key] = True
+                        unique_spans.append(span)
+                else:
+                    # For non-LLM spans, process their children if they exist
+                    if 'data' in span_dict and 'children' in span_dict['data']:
+                        children = span_dict['data']['children']
+                        # Filter and deduplicate children
+                        filtered_children = deduplicate_spans(children)
+                        if isinstance(span, dict):
+                            span['data']['children'] = filtered_children
+                        else:
+                            span.data['children'] = filtered_children
+                    unique_spans.append(span)
+            
+            return unique_spans
+
+        # Remove any spans without hash ids
+        for data in trace.get('data', []):
+            if 'spans' in data:
+                # First filter out spans without hash_ids, then deduplicate
+                data['spans'] = deduplicate_spans(data['spans'])
+        
         return trace

ragaai_catalyst/tracers/agentic_tracing/data_structure.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
-from typing import List, Dict, Optional, Any
+from typing import List, Dict, Optional, Any, Union
 from datetime import datetime
+import uuid
 
 @dataclass
 class OSInfo:
@@ -99,12 +100,20 @@ class NetworkCall:
     request: Dict[str, Any]
     response: Dict[str, Any]
 
-@dataclass
 class Interaction:
-    id: str
-    interaction_type: str
-    content: Optional[str]
-    timestamp: str
+    def __init__(self, id, type: str, content: str, timestamp: str):
+        self.id = id
+        self.type = type
+        self.content = content
+        self.timestamp = timestamp
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "interaction_type": self.type,
+            "content": self.content,
+            "timestamp": self.timestamp
+        }
 
 @dataclass
 class Error:
@@ -150,54 +159,63 @@ class ToolInfo:
     version: str
     memory_used: int
 
-@dataclass
-class LLMComponent:
-    id: str
-    hash_id: str
-    source_hash_id: Optional[str]
-    type: str = "llm"
-    name: str = ""
-    start_time: str = ""
-    end_time: str = ""
-    error: Optional[Error] = None
-    parent_id: Optional[str] = None
-    info: LLMInfo = None
-    data: Dict[str, Any] = None
-    network_calls: List[NetworkCall] = None
-    interactions: List[Interaction] = None
-
-@dataclass
-class AgentComponent:
-    id: str
-    hash_id: str
-    source_hash_id: Optional[str]
-    type: str = "agent"
-    name: str = ""
-    start_time: str = ""
-    end_time: str = ""
-    error: Optional[Error] = None
-    parent_id: Optional[str] = None
-    info: AgentInfo = None
-    data: Dict[str, Any] = None
-    network_calls: List[NetworkCall] = None
-    interactions: List[Interaction] = None
-    # children: List['Component'] = None
-
-@dataclass
-class ToolComponent:
-    id: str
-    hash_id: str
-    source_hash_id: Optional[str]
-    type: str = "tool"
-    name: str = ""
-    start_time: str = ""
-    end_time: str = ""
-    error: Optional[Error] = None
-    parent_id: Optional[str] = None
-    info: ToolInfo = None
-    data: Dict[str, Any] = None
-    network_calls: List[NetworkCall] = None
-    interactions: List[Interaction] = None
+class Component:
+    def __init__(self, id: str, hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
+        self.id = id
+        self.hash_id = hash_id
+        self.type = type
+        self.name = name
+        self.start_time = start_time
+        self.end_time = end_time
+        self.parent_id = parent_id
+        self.info = info
+        self.data = data
+        self.error = error
+        self.network_calls = network_calls or []
+        self.interactions = []
+        self.error = error
+        if interactions:
+            for interaction in interactions:
+                if isinstance(interaction, dict):
+                    self.interactions.append(
+                        Interaction(
+                            id=interaction.get("id", str(uuid.uuid4())),
+                            type=interaction.get("interaction_type", ""),
+                            content=str(interaction.get("content", "")),
+                            timestamp=interaction.get("timestamp", datetime.utcnow().isoformat())
+                        )
+                    )
+                else:
+                    self.interactions.append(interaction)
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "hash_id": self.hash_id,
+            "type": self.type,
+            "name": self.name,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "parent_id": self.parent_id,
+            "info": self.info,
+            "error": self.error,
+            "data": self.data,
+            "error": self.error,
+            "network_calls": [call.to_dict() if hasattr(call, 'to_dict') else call for call in self.network_calls],
+            "interactions": self.interactions
+        }
+
+class LLMComponent(Component):
+    def __init__(self, id: str, hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
+        super().__init__(id, hash_id, type, name, start_time, end_time, parent_id, info, data, network_calls, interactions, error)
+
+class AgentComponent(Component):
+    def __init__(self, id: str, hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
+        super().__init__(id, hash_id, type, name, start_time, end_time, parent_id, info, data, network_calls, interactions, error)
+
+class ToolComponent(Component):
+    def __init__(self, id: str, hash_id: str, type: str, name: str, start_time: str, end_time: str, parent_id: int, info: Dict[str, Any], data: Dict[str, Any], network_calls: Optional[List[NetworkCall]] = None, interactions: Optional[List[Union[Interaction, Dict]]] = None, error: Optional[Dict[str, Any]] = None):
+        super().__init__(id, hash_id, type, name, start_time, end_time, parent_id, info, data, network_calls, interactions, error)
 
 @dataclass
 class ComponentInfo:
@@ -211,29 +229,23 @@ class ComponentInfo:
     token_usage: Optional[Dict[str, int]] = None
     cost: Optional[Dict[str, float]] = None
 
-@dataclass
-class Component:
-    id: str
-    hash_id: str
-    source_hash_id: Optional[str]
-    type: str
-    name: str
-    start_time: str
-    end_time: str
-    error: Optional[Error]
-    parent_id: Optional[str]
-    info: ComponentInfo
-    data: Dict[str, Any]
-    network_calls: List[NetworkCall]
-    interactions: List[Interaction]
-    children: Optional[List['Component']] = None
-
-@dataclass
 class Trace:
-    id: str
-    project_name: str
-    start_time: str
-    end_time: str
-    metadata: Metadata
-    data: List[Dict[str, Any]]
-    replays: Optional[Dict[str, Any]]
+    def __init__(self, id: str, project_name: str, start_time: str, end_time: str, metadata: Optional[Metadata] = None, data: Optional[List[Dict[str, Any]]] = None, replays: Optional[Dict[str, Any]] = None):
+        self.id = id
+        self.project_name = project_name
+        self.start_time = start_time
+        self.end_time = end_time
+        self.metadata = metadata or Metadata()
+        self.data = data or []
+        self.replays = replays
+
+    def to_dict(self):
+        return {
+            "id": self.id,
+            "project_name": self.project_name,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "metadata": self.metadata.to_dict() if self.metadata else None,
+            "data": self.data,
+            "replays": self.replays,
+        }