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

ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Dict, Any
+from typing import List, Dict, Any, Optional
 import logging
 
 logger = logging.getLogger(__name__)
@@ -11,13 +11,17 @@ logging_level = (
 
 
 class SpanAttributes:
-    def __init__(self, name):
+    def __init__(self, name, project_id: Optional[int] = None):
         self.name = name
         self.tags = []
         self.metadata = {}
         self.metrics = []
+        self.local_metrics = []
         self.feedback = None
+        self.project_id = project_id
         self.trace_attributes = ["tags", "metadata", "metrics"]
+        self.gt = None
+        self.context = None
 
     def add_tags(self, tags: str | List[str]):
         if isinstance(tags, str):
@@ -30,14 +34,14 @@ class SpanAttributes:
         logger.debug(f"Added metadata: {metadata}")
 
     def add_metrics(
-        self,
-        name: str,
-        score: float | int,
-        reasoning: str = "",
-        cost: float = None,
-        latency: float = None,
-        metadata: Dict[str, Any] = {},
-        config: Dict[str, Any] = {},
+            self,
+            name: str,
+            score: float | int,
+            reasoning: str = "",
+            cost: float = None,
+            latency: float = None,
+            metadata: Dict[str, Any] = {},
+            config: Dict[str, Any] = {},
     ):
         self.metrics.append(
             {
@@ -57,3 +61,70 @@ class SpanAttributes:
     def add_feedback(self, feedback: Any):
         self.feedback = feedback
         logger.debug(f"Added feedback: {self.feedback}")
+
+    # TODO: Add validation to check if all the required parameters are present
+    def execute_metrics(self, **kwargs: Any):
+        name = kwargs.get("name")
+        model = kwargs.get("model")
+        provider = kwargs.get("provider")
+        display_name = kwargs.get("display_name", None)
+        mapping = kwargs.get("mapping", None)
+
+        if isinstance(name, str):
+            metrics = [{
+                "name": name
+            }]
+        else:
+            metrics = name if isinstance(name, list) else [name] if isinstance(name, dict) else []
+
+        for metric in metrics:
+            if not isinstance(metric, dict):
+                raise ValueError(f"Expected dict, got {type(metric)}")
+
+            if "name" not in metric:
+                raise ValueError("Metric must contain 'name'")
+
+            metric_name = metric["name"]
+            if metric_name in self.local_metrics:
+                count = sum(1 for m in self.local_metrics if m.startswith(metric_name))
+                metric_name = f"{metric_name}_{count + 1}"
+
+            prompt =None
+            context = None
+            response = None
+            # if mapping is not None:
+            #     prompt = mapping['prompt']
+            #     context = mapping['context']
+            #     response = mapping['response']
+            new_metric = {
+                "name": metric_name,
+                "model": model,
+                "provider": provider,
+                "project_id": self.project_id,
+                # "prompt": prompt,
+                # "context": context,
+                # "response": response,
+                "displayName": display_name,
+                "mapping": mapping
+            }
+            self.local_metrics.append(new_metric)
+
+    def add_gt(self, gt: Any):
+        if not isinstance(gt, (str, int, float, bool, list, dict)):
+            raise TypeError(f"Unsupported type for gt: {type(gt)}")
+        if self.gt:
+            logger.warning(f"GT already exists: {self.gt} \n Overwriting...")
+        self.gt = gt
+        logger.debug(f"Added gt: {self.gt}")
+
+    def add_context(self, context: Any):
+        if isinstance(context, str):
+            if not context.strip():
+                logger.warning("Empty or whitespace-only context string provided")
+            self.context = str(context)
+        else:
+            try:
+                self.context = str(context)
+            except Exception as e:
+                logger.warning('Cannot cast the context to string... Skipping')
+        logger.debug(f"Added context: {self.context}")

ragaai_catalyst/tracers/agentic_tracing/utils/supported_llm_provider.toml

@@ -0,0 +1,34 @@
+# List of all supported LLM method calls
+
+supported_llm_calls = [
+    # OpenAI
+    "OpenAI.chat.completions.create()",
+    "AsyncOpenAI.chat.completions.create()",
+    
+    # OpenAI Beta
+    "OpenAI.beta.threads.create()",
+    "OpenAI.beta.threads.messages.create()",
+    "OpenAI.beta.threads.runs.create()",
+    
+    # Anthropic
+    "Anthropic.messages.create()",
+    "Anthropic.messages.acreate()",
+    
+    # Google VertexAI/PaLM
+    "GenerativeModel.generate_content()",
+    "GenerativeModel.generate_content_async()",
+    "ChatVertexAI._generate()",
+    "ChatVertexAI._agenerate()",
+    "ChatVertexAI.complete()",
+    "ChatVertexAI.acomplete()",
+    
+    # Google GenerativeAI
+    "ChatGoogleGenerativeAI._generate()",
+    "ChatGoogleGenerativeAI._agenerate()",
+    "ChatGoogleGenerativeAI.complete()",
+    "ChatGoogleGenerativeAI.acomplete()",
+    
+    # LiteLLM
+    "litellm.completion()",
+    "litellm.acompletion()"
+]

ragaai_catalyst/tracers/agentic_tracing/utils/system_monitor.py

@@ -0,0 +1,215 @@
+import platform
+import psutil
+import sys
+import pkg_resources
+import logging
+from typing import Dict, List, Optional
+from ..data.data_structure import (
+    SystemInfo,
+    OSInfo,
+    EnvironmentInfo,
+    Resources,
+    CPUResource,
+    MemoryResource,
+    DiskResource,
+    NetworkResource,
+    ResourceInfo,
+    MemoryInfo,
+    DiskInfo,
+    NetworkInfo,
+)
+
+logger = logging.getLogger(__name__)
+
+class SystemMonitor:
+    def __init__(self, trace_id: str):
+        self.trace_id = trace_id
+
+    def get_system_info(self) -> SystemInfo:
+        # Initialize with None values
+        os_info = OSInfo(
+            name=None,
+            version=None,
+            platform=None,
+            kernel_version=None,
+        )
+        env_info = EnvironmentInfo(
+            name=None,
+            version=None,
+            packages=[],
+            env_path=None,
+            command_to_run=None,
+        )
+
+        try:
+            # Get OS info
+            os_info = OSInfo(
+                name=platform.system(),
+                version=platform.version(),
+                platform=platform.machine(),
+                kernel_version=platform.release(),
+            )
+        except Exception as e:
+            logger.warning(f"Failed to get OS info: {str(e)}")
+             
+        try:
+            # Get Python environment info
+            installed_packages = [
+                f"{pkg.key}=={pkg.version}" for pkg in pkg_resources.working_set
+            ]
+            env_info = EnvironmentInfo(
+                name="Python",
+                version=platform.python_version(),
+                packages=installed_packages,
+                env_path=sys.prefix,
+                command_to_run=f"python {sys.argv[0]}",
+            )
+        except Exception as e:
+            logger.warning(f"Failed to get environment info: {str(e)}")
+             
+
+        # Always return a valid SystemInfo object
+        return SystemInfo(
+            id=f"sys_{self.trace_id}",
+            os=os_info,
+            environment=env_info,
+            source_code="",
+        )
+
+    def get_resources(self) -> Resources:
+        # Initialize with None values
+        cpu_info = ResourceInfo(
+            name=None,
+            cores=None,
+            threads=None,
+        )
+        cpu = CPUResource(info=cpu_info, interval="5s", values=[])
+
+        mem_info = MemoryInfo(
+            total=None,
+            free=None,
+        )
+        mem = MemoryResource(info=mem_info, interval="5s", values=[])
+
+        disk_info = DiskInfo(
+            total=None,
+            free=None,
+        )
+        disk_resource = DiskResource(
+            info=disk_info,
+            interval="5s",
+            read=[],
+            write=[],
+        )
+
+        net_info = NetworkInfo(
+            upload_speed=None,
+            download_speed=None,
+        )
+        net = NetworkResource(
+            info=net_info,
+            interval="5s",
+            uploads=[],
+            downloads=[],
+        )
+
+        try:
+            # CPU info
+            cpu_info = ResourceInfo(
+                name=platform.processor(),
+                cores=psutil.cpu_count(logical=False),
+                threads=psutil.cpu_count(logical=True),
+            )
+            cpu = CPUResource(info=cpu_info, interval="5s", values=[psutil.cpu_percent()])
+        except Exception as e:
+            logger.warning(f"Failed to get CPU info: {str(e)}")
+             
+
+        try:
+            # Memory info
+            memory = psutil.virtual_memory()
+            mem_info = MemoryInfo(
+                total=memory.total / (1024**3),  # Convert to GB
+                free=memory.available / (1024**3),
+            )
+            mem = MemoryResource(info=mem_info, interval="5s", values=[memory.percent])
+        except Exception as e:
+            logger.warning(f"Failed to get memory info: {str(e)}")
+             
+
+        try:
+            # Disk info
+            disk = psutil.disk_usage("/")
+            disk_info = DiskInfo(total=disk.total / (1024**3), free=disk.free / (1024**3))
+            disk_io = psutil.disk_io_counters()
+            disk_resource = DiskResource(
+                info=disk_info,
+                interval="5s",
+                read=[disk_io.read_bytes / (1024**2)],  # MB
+                write=[disk_io.write_bytes / (1024**2)],
+            )
+        except Exception as e:
+            logger.warning(f"Failed to get disk info: {str(e)}")
+             
+        try:
+            # Network info
+            net_io = psutil.net_io_counters()
+            net_info = NetworkInfo(
+                upload_speed=net_io.bytes_sent / (1024**2),  # MB
+                download_speed=net_io.bytes_recv / (1024**2),
+            )
+            net = NetworkResource(
+                info=net_info,
+                interval="5s",
+                uploads=[net_io.bytes_sent / (1024**2)],
+                downloads=[net_io.bytes_recv / (1024**2)],
+            )
+        except Exception as e:
+            logger.warning(f"Failed to get network info: {str(e)}")
+            
+
+        # Always return a valid Resources object
+        return Resources(cpu=cpu, memory=mem, disk=disk_resource, network=net)
+
+    def track_memory_usage(self) -> Optional[float]:
+        """Track memory usage in MB"""
+        try:
+            memory_usage = psutil.Process().memory_info().rss
+            return memory_usage / (1024 * 1024)  # Convert to MB
+        except Exception as e:
+            logger.warning(f"Failed to track memory usage: {str(e)}")
+            return None  
+
+    def track_cpu_usage(self, interval: float) -> Optional[float]:
+        """Track CPU usage percentage"""
+        try:
+            return psutil.cpu_percent(interval=interval)
+        except Exception as e:
+            logger.warning(f"Failed to track CPU usage: {str(e)}")
+            return None  
+
+    def track_disk_usage(self) -> Dict[str, Optional[float]]:
+        """Track disk I/O in MB"""
+        default_response = {'disk_read': None, 'disk_write': None}
+        try:
+            disk_io = psutil.disk_io_counters()
+            return {
+                'disk_read': disk_io.read_bytes / (1024 * 1024),  # Convert to MB
+                'disk_write': disk_io.write_bytes / (1024 * 1024)  # Convert to MB
+            }
+        except Exception as e:
+            logger.warning(f"Failed to track disk usage: {str(e)}")
+            return default_response 
+
+    def track_network_usage(self) -> Dict[str, Optional[float]]:
+        """Track network I/O in MB"""
+        default_response = {'uploads': None, 'downloads': None}
+        try:
+            net_io = psutil.net_io_counters()
+            return {
+                'uploads': net_io.bytes_sent / (1024 * 1024),  # Convert to MB
+                'downloads': net_io.bytes_recv / (1024 * 1024)  # Convert to MB
+            }
+        except Exception as e:
+            logger.warning(f"Failed to track network usage: {str(e)}")
+            return default_response 

ragaai_catalyst/tracers/agentic_tracing/utils/trace_utils.py

@@ -59,38 +59,6 @@ def calculate_cost(
         "total": total_cost,
     }
 
-
-def load_model_costs():
-    try:
-        current_dir = os.path.dirname(os.path.abspath(__file__))
-        model_costs_path = os.path.join(current_dir, "model_costs.json")
-        with open(model_costs_path, "r") as file:
-            return json.load(file)
-    except FileNotFoundError:
-        with resources.open_text("utils", "model_costs.json") as file:
-            return json.load(file)
-
-
-def update_model_costs_from_github():
-    """Updates the model_costs.json file with latest costs from GitHub."""
-    try:
-        logger.debug("loading the latest model costs.")
-        response = requests.get(
-            "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
-        )
-        if response.status_code == 200:
-            current_dir = os.path.dirname(os.path.abspath(__file__))
-            model_costs_path = os.path.join(current_dir, "model_costs.json")
-            with open(model_costs_path, "w") as file:
-                json.dump(response.json(), file, indent=4)
-            logger.debug("Model costs updated successfully.")
-            return True
-        return False
-    except Exception as e:
-        logger.error(f"Failed to update model costs from GitHub: {e}")
-        return False
-
-
 def log_event(event_data, log_file_path):
     event_data = asdict(event_data)
     with open(log_file_path, "a") as f:

ragaai_catalyst/tracers/agentic_tracing/utils/unique_decorator.py

@@ -56,7 +56,9 @@ def generate_unique_hash(func, *args, **kwargs):
                 return '_'.join(f"{normalize_arg(k)}:{normalize_arg(v)}" 
                               for k, v in sorted(arg.items()))
             elif callable(arg):
-                return arg.__name__
+                if hasattr(arg, "__name__"):
+                    return arg.__name__
+                return str(type(arg).__name__)
             else:
                 return str(type(arg).__name__)
 

ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py

@@ -1,13 +1,14 @@
 import os
+import sys
+import importlib
 import hashlib
 import zipfile
 import re
 import ast
 import importlib.util
 import json
-import astor
 import ipynbname
-import sys
+from copy import deepcopy
 
 from pathlib import Path
 from IPython import get_ipython
@@ -23,7 +24,7 @@ logger = logging.getLogger(__name__)
 logging_level = logger.setLevel(logging.DEBUG) if os.getenv("DEBUG") == "1" else logging.INFO
 
 
-# Define the PackageUsageRemover class
+# PackageUsageRemover class
 class PackageUsageRemover(ast.NodeTransformer):
     def __init__(self, package_name):
         self.package_name = package_name
@@ -49,7 +50,12 @@ class PackageUsageRemover(ast.NodeTransformer):
         return node
     
     def visit_Assign(self, node):
-        if self._uses_package(node.value):
+        if isinstance(node.value, ast.Expr):  
+            node_value = node.value.body   
+        else:
+            node_value = node.value
+
+        if self._uses_package(node_value):
             return None
         return node
     
@@ -60,8 +66,10 @@ class PackageUsageRemover(ast.NodeTransformer):
             if isinstance(node.func.value, ast.Name) and node.func.value.id in self.imported_names:
                 return None
         return node
-    
+
     def _uses_package(self, node):
+        if isinstance(node, ast.Expr):
+            return self._uses_package(node.body)
         if isinstance(node, ast.Name) and node.id in self.imported_names:
             return True
         if isinstance(node, ast.Call):
@@ -70,16 +78,19 @@ class PackageUsageRemover(ast.NodeTransformer):
             return self._uses_package(node.value)
         return False
 
-# Define the function to remove package code from a source code string
+
+# Remove package code from a source code string
 def remove_package_code(source_code: str, package_name: str) -> str:
     try:
         tree = ast.parse(source_code)
-        transformer = PackageUsageRemover(package_name)
-        modified_tree = transformer.visit(tree)
-        modified_code = astor.to_source(modified_tree)
+        # remover = PackageUsageRemover(package_name)
+        # modified_tree = remover.visit(tree)
+        modified_code = ast.unparse(tree)
+
         return modified_code
     except Exception as e:
-        raise Exception(f"Error processing source code: {str(e)}")
+        logger.error(f"Error in remove_package_code: {e}")
+        return source_code
 
 class JupyterNotebookHandler:
     @staticmethod
@@ -118,25 +129,12 @@ class JupyterNotebookHandler:
             # Check if running in Colab
             if JupyterNotebookHandler.is_running_in_colab():
                 try:
-                    from google.colab import drive
-                    if not os.path.exists('/content/drive'):
-                        drive.mount('/content/drive')
-                        # logger.info("Google Drive mounted successfully")
-
                     # Look for notebooks in /content first
                     ipynb_files = list(Path('/content').glob('*.ipynb'))
                     if ipynb_files:
                         current_nb = max(ipynb_files, key=os.path.getmtime)
                         # logger.info(f"Found current Colab notebook: {current_nb}")
                         return str(current_nb)
-                        
-                    # Then check Drive if mounted
-                    if os.path.exists('/content/drive'):
-                        drive_ipynb_files = list(Path('/content/drive').rglob('*.ipynb'))
-                        if drive_ipynb_files:
-                            current_nb = max(drive_ipynb_files, key=os.path.getmtime)
-                            # logger.info(f"Found Colab notebook in Drive: {current_nb}")
-                            return str(current_nb)
                 except Exception as e:
                     logger.warning(f"Error in Colab notebook detection: {str(e)}")
 
@@ -201,7 +199,6 @@ def comment_magic_commands(script_content: str) -> str:
 class TraceDependencyTracker:
     def __init__(self, output_dir=None):
         self.tracked_files = set()
-        self.python_imports = set()
         self.notebook_path = None
         self.colab_content = None  
         
@@ -292,7 +289,7 @@ class TraceDependencyTracker:
                     except (UnicodeDecodeError, IOError):
                         pass
 
-    def analyze_python_imports(self, filepath):
+    def analyze_python_imports(self, filepath, ignored_locations):
         try:
             with open(filepath, 'r', encoding='utf-8') as file:
                 tree = ast.parse(file.read(), filename=filepath)
@@ -305,48 +302,75 @@ class TraceDependencyTracker:
                             module_name = name.name.split('.')[0]
                     try:
                         spec = importlib.util.find_spec(module_name)
-                        if spec and spec.origin and not spec.origin.startswith(os.path.dirname(importlib.__file__)):
-                            self.python_imports.add(spec.origin)
+                        if spec and spec.origin:    
+                            if not (any(spec.origin.startswith(location) for location in ignored_locations) or (spec.origin in ['built-in', 'frozen'])):
+                                self.tracked_files.add(spec.origin)
+                                self.analyze_python_imports(spec.origin, ignored_locations)
                     except (ImportError, AttributeError):
                         pass
         except Exception as e:
             pass
 
+    def get_env_location(self):
+        return sys.prefix
+    
+    def get_catalyst_location(self):
+        try:
+            imported_module = importlib.import_module("ragaai_catalyst")
+            return os.path.dirname(os.path.abspath(imported_module.__file__))
+        except ImportError:
+            logger.error("Error getting Catalyst location")
+            return 'ragaai_catalyst'
+    
+    def should_ignore_path(self, path, main_filepaths):
+        if any(os.path.abspath(path) in os.path.abspath(main_filepath) for main_filepath in main_filepaths):
+            return False
+        if path in ['', os.path.abspath('')]:
+            return False
+        return True
+
     def create_zip(self, filepaths):
         self.track_jupyter_notebook()
-        # logger.info("Tracked Jupyter notebook and its dependencies")
 
         # Ensure output directory exists
         os.makedirs(self.output_dir, exist_ok=True)
-        # logger.info(f"Using output directory: {self.output_dir}")
 
         # Special handling for Colab
         if self.jupyter_handler.is_running_in_colab():
-            # logger.info("Running in Google Colab environment")
-            # Try to get the Colab notebook path
+            # Get the Colab notebook path
             colab_notebook = self.jupyter_handler.get_notebook_path()
             if colab_notebook:
                 self.tracked_files.add(os.path.abspath(colab_notebook))
-                # logger.info(f"Added Colab notebook to tracked files: {colab_notebook}")
 
             # Get current cell content
             self.check_environment_and_save()
 
+        env_location = self.get_env_location()
+        catalyst_location = self.get_catalyst_location()
+
         # Process all files (existing code)
+        ignored_locations = [env_location, catalyst_location] + [path for path in sys.path if self.should_ignore_path(path, filepaths)]
         for filepath in filepaths:
             abs_path = os.path.abspath(filepath)
             self.track_file_access(abs_path)
             try:
-                with open(abs_path, 'r', encoding='utf-8') as file:
+                if filepath.endswith('.py'):
+                    self.analyze_python_imports(abs_path, ignored_locations)
+            except Exception as e:
+                pass
+        
+        curr_tracked_files = deepcopy(self.tracked_files)
+        for filepath in curr_tracked_files:
+            try:
+                with open(filepath, 'r', encoding='utf-8') as file:
                     content = file.read()
                     # Comment out magic commands before processing
                     content = comment_magic_commands(content)
-                self.find_config_files(content, abs_path)
-                if filepath.endswith('.py'):
-                    self.analyze_python_imports(abs_path)
+                self.find_config_files(content, filepath)
             except Exception as e:
                 pass
 
+
         notebook_content_str = None
         if self.notebook_path and os.path.exists(self.notebook_path):
             try:
@@ -370,13 +394,12 @@ class TraceDependencyTracker:
                 pass
 
         # Calculate hash and create zip
-        self.tracked_files.update(self.python_imports)
         hash_contents = []
 
         for filepath in sorted(self.tracked_files):
             if not filepath.endswith('.py'):
                 continue
-            elif '/envs' in filepath or '__init__' in filepath:
+            elif env_location in filepath or '__init__' in filepath:
                 continue
             try:
                 with open(filepath, 'rb') as file:
@@ -409,11 +432,15 @@ class TraceDependencyTracker:
 
         with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
             for filepath in sorted(self.tracked_files):
-                if 'env' in filepath or 'ragaai_catalyst' in filepath:
+                if env_location in filepath or catalyst_location in filepath:
                     continue
                 try:
                     relative_path = os.path.relpath(filepath, base_path)
-                    zipf.write(filepath, relative_path)
+                    if relative_path in ['', '.']:
+                        zipf.write(filepath, os.path.basename(filepath))
+                    else:
+                        zipf.write(filepath, relative_path)
+                    
                     logger.debug(f"Added python script to zip: {relative_path}")
                 except Exception as e:
                     pass
@@ -446,10 +473,9 @@ def zip_list_of_unique_files(filepaths, output_dir=None):
     return tracker.create_zip(filepaths)
 
 
-# Example usage
-if __name__ == "__main__":
-    filepaths = ["script1.py", "script2.py"]
-    hash_id, zip_path = zip_list_of_unique_files(filepaths)
-    print(f"Created zip file: {zip_path}")
-    print(f"Hash ID: {hash_id}")
-
+# # Example usage
+# if __name__ == "__main__":
+#     filepaths = ["script1.py", "script2.py"]
+#     hash_id, zip_path = zip_list_of_unique_files(filepaths)
+#     print(f"Created zip file: {zip_path}")
+#     print(f"Hash ID: {hash_id}")