ragaai-catalyst 2.1.4.1b1__py3-none-any.whl → 2.1.5b1__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -61,6 +61,7 @@ class LLMTracerMixin:
         # Add auto_instrument options
         self.auto_instrument_llm = False
         self.auto_instrument_user_interaction = False
+        self.auto_instrument_file_io = False
         self.auto_instrument_network = False
 
     def instrument_llm_calls(self):
@@ -117,6 +118,10 @@ class LLMTracerMixin:
     def instrument_network_calls(self):
         """Enable network instrumentation for LLM calls"""
         self.auto_instrument_network = True
+        
+    def instrument_file_io_calls(self):
+        """Enable file IO instrumentation for LLM calls"""
+        self.auto_instrument_file_io = True
 
     def patch_openai_methods(self, module):
         try:
@@ -334,7 +339,17 @@ class LLMTracerMixin:
 
         interactions = []
         if self.auto_instrument_user_interaction:
-            interactions = self.component_user_interaction.get(component_id, [])
+            input_output_interactions = []
+            for interaction in self.component_user_interaction.get(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(component_id, []):
+                if interaction["interaction_type"] in ["file_read", "file_write"]:
+                    file_io_interactions.append(interaction)
+            interactions.extend(file_io_interactions)
 
         parameters_to_display = {}
         if "run_manager" in parameters:
@@ -475,7 +490,7 @@ class LLMTracerMixin:
                 parameters=parameters,
             )
 
-            # self.add_component(llm_component)
+            self.add_component(llm_component)
             self.llm_data = llm_component
 
             return result

ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py

@@ -84,27 +84,23 @@ class AgenticTracing(
             self.auto_instrument_custom = True
         else:
             # Set global active state
-            self.is_active = (
-                any(auto_instrumentation.values())
-                if isinstance(auto_instrumentation, dict)
-                else bool(auto_instrumentation)
-            )
+            self.is_active = True
 
             # Set individual components
             if isinstance(auto_instrumentation, dict):
-                self.auto_instrument_llm = auto_instrumentation.get("llm", False)
-                self.auto_instrument_tool = auto_instrumentation.get("tool", False)
-                self.auto_instrument_agent = auto_instrumentation.get("agent", False)
+                self.auto_instrument_llm = auto_instrumentation.get("llm", True)
+                self.auto_instrument_tool = auto_instrumentation.get("tool", True)
+                self.auto_instrument_agent = auto_instrumentation.get("agent", True)
                 self.auto_instrument_user_interaction = auto_instrumentation.get(
-                    "user_interaction", False
+                    "user_interaction", True
                 )
                 self.auto_instrument_file_io = auto_instrumentation.get(
-                    "file_io", False
+                    "file_io", True
                 )
                 self.auto_instrument_network = auto_instrumentation.get(
-                    "network", False
+                    "network", True
                 )
-                self.auto_instrument_custom = auto_instrumentation.get("custom", False)
+                self.auto_instrument_custom = auto_instrumentation.get("custom", True)
             else:
                 # If boolean provided, apply to all components
                 self.auto_instrument_llm = bool(auto_instrumentation)
@@ -137,13 +133,28 @@ class AgenticTracing(
             self.network_tracer.network_calls.copy()
         )
         self.network_tracer.network_calls = []  # Reset for next component
-        # self.component_user_interaction[component_id] = [interaction for interaction in self.user_interaction_tracer.interactions if interaction.get('component_id') == component_id]
+
+        # Store user interactions for the component
         for interaction in self.user_interaction_tracer.interactions:
             interaction_component_id = interaction.get("component_id")
             if interaction_component_id not in self.component_user_interaction:
                 self.component_user_interaction[interaction_component_id] = []
             if interaction not in self.component_user_interaction[interaction_component_id]:
                 self.component_user_interaction[interaction_component_id].append(interaction)
+                
+        # Only reset component_id if it matches the current one
+        # This ensures we don't reset a parent's component_id when a child component ends
+        if self.current_component_id.get() == component_id:
+            # Get the parent agent's component_id if it exists
+            parent_agent_id = self.current_agent_id.get()
+            # If there's a parent agent, set the component_id back to the parent's
+            if parent_agent_id:
+                self.current_component_id.set(parent_agent_id)
+                self.user_interaction_tracer.component_id.set(parent_agent_id)
+            else:
+                # Only reset to None if there's no parent
+                self.current_component_id.set(None)
+                self.user_interaction_tracer.component_id.set(None)
 
     def start(self):
         """Start tracing"""
@@ -155,9 +166,6 @@ class AgenticTracing(
         self.user_interaction_tracer.trace_id.set(self.trace_id)
         self.user_interaction_tracer.tracer = self
         self.user_interaction_tracer.component_id.set(self.current_component_id.get())
-        builtins.print = self.user_interaction_tracer.traced_print
-        builtins.input = self.user_interaction_tracer.traced_input
-        builtins.open = self.user_interaction_tracer.traced_open
 
         # Start base tracer (includes system info and resource monitoring)
         super().start()
@@ -179,11 +187,12 @@ class AgenticTracing(
             self.instrument_custom_calls()
 
         if self.auto_instrument_user_interaction:
-
             ToolTracerMixin.instrument_user_interaction_calls(self)
             LLMTracerMixin.instrument_user_interaction_calls(self)
             AgentTracerMixin.instrument_user_interaction_calls(self)
             CustomTracerMixin.instrument_user_interaction_calls(self)
+            builtins.print = self.user_interaction_tracer.traced_print
+            builtins.input = self.user_interaction_tracer.traced_input
 
         if self.auto_instrument_network:
             ToolTracerMixin.instrument_network_calls(self)
@@ -191,9 +200,12 @@ class AgenticTracing(
             AgentTracerMixin.instrument_network_calls(self)
             CustomTracerMixin.instrument_network_calls(self)
 
-        # These will be implemented later
-        # if self.auto_instrument_file_io:
-        #     self.instrument_file_io_calls()
+        if self.auto_instrument_file_io:
+            ToolTracerMixin.instrument_file_io_calls(self)
+            LLMTracerMixin.instrument_file_io_calls(self)
+            AgentTracerMixin.instrument_file_io_calls(self)
+            CustomTracerMixin.instrument_file_io_calls(self)
+            builtins.open = self.user_interaction_tracer.traced_open
 
     def stop(self):
         """Stop tracing and save results"""

ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py

@@ -41,7 +41,7 @@ class NetworkTracer:
         
         # Extract protocol from URL
         protocol = "https" if url.startswith("https") else "http"
-        
+            
         self.network_calls.append(
             {
                 "url": url,
@@ -57,11 +57,11 @@ class NetworkTracer:
                 "parent_id": None,  # Will be set by the component
                 "request": {
                     "headers": request_headers,
-                    "body": request_body[:1000] if request_body else None,  # Limit to 1000 chars
+                    "body": request_body if request_body else None,
                 },
                 "response": {
                     "headers": response_headers,
-                    "body": response_body[:1000] if response_body else None,  # Limit to 1000 chars
+                    "body": response_body if response_body else None,
                 },
                 "error": str(error) if error else None,
             }

ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py

@@ -32,6 +32,7 @@ class ToolTracerMixin:
         # add auto_instrument option
         self.auto_instrument_tool = False
         self.auto_instrument_user_interaction = False
+        self.auto_instrument_file_io = False
         self.auto_instrument_network = False
 
     # take care of auto_instrument
@@ -40,6 +41,9 @@ class ToolTracerMixin:
 
     def instrument_user_interaction_calls(self):
         self.auto_instrument_user_interaction = True
+        
+    def instrument_file_io_calls(self):
+        self.auto_instrument_file_io = True
 
     def instrument_network_calls(self):
         self.auto_instrument_network = True
@@ -133,6 +137,10 @@ class ToolTracerMixin:
         component_id = str(uuid.uuid4())
         hash_id = generate_unique_hash_simple(func)
 
+        # Set current tool name and store the token
+        name_token = self.current_tool_name.set(name)
+        id_token = self.current_tool_id.set(component_id)
+
         # Start tracking network calls for this component
         self.start_component(component_id)
 
@@ -191,6 +199,12 @@ class ToolTracerMixin:
             self.add_component(tool_component)
 
             raise
+        finally:
+            # Reset the tool name and id context
+            if name_token:
+                self.current_tool_name.reset(name_token)
+            if id_token:
+                self.current_tool_id.reset(id_token)
 
     async def _trace_tool_execution(
         self, func, name, tool_type, version, *args, **kwargs
@@ -207,6 +221,10 @@ class ToolTracerMixin:
         component_id = str(uuid.uuid4())
         hash_id = generate_unique_hash_simple(func)
 
+        # Set current tool name and store the token
+        name_token = self.current_tool_name.set(name)
+        id_token = self.current_tool_id.set(component_id)
+
         self.start_component(component_id)
         try:
             # Execute the tool
@@ -256,6 +274,12 @@ class ToolTracerMixin:
             self.add_component(tool_component)
 
             raise
+        finally:
+            # Reset the tool name and id context
+            if name_token:
+                self.current_tool_name.reset(name_token)
+            if id_token:
+                self.current_tool_id.reset(id_token)
 
     def create_tool_component(self, **kwargs):
         """Create a tool component according to the data structure"""
@@ -264,9 +288,19 @@ class ToolTracerMixin:
             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"], []
-            )
+            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)
+            if input_output_interactions!=[]:
+                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)
+            if file_io_interactions!=[]:
+                interactions.extend(file_io_interactions)
 
         # Get tags, metrics
         name = kwargs["name"]

ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json

@@ -2063,7 +2063,9 @@
         "input_cost_per_token": 5.9e-07,
         "output_cost_per_token": 7.9e-07,
         "litellm_provider": "groq",
-        "mode": "chat"
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_response_schema": true
     },
     "groq/llama-3.3-70b-specdec": {
         "max_tokens": 8192,
@@ -4605,6 +4607,20 @@
         "litellm_provider": "replicate",
         "mode": "chat"
     },
+    "openrouter/deepseek/deepseek-r1": {
+        "max_tokens": 8192,
+        "max_input_tokens": 64000,
+        "max_output_tokens": 8192,
+        "input_cost_per_token": 5.5e-07,
+        "input_cost_per_token_cache_hit": 1.4e-07,
+        "output_cost_per_token": 2.19e-06,
+        "litellm_provider": "openrouter",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_assistant_prefill": true,
+        "supports_tool_choice": true,
+        "supports_prompt_caching": true
+    },
     "openrouter/deepseek/deepseek-chat": {
         "max_tokens": 8192,
         "max_input_tokens": 66000,

ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py

@@ -5,7 +5,6 @@ import re
 import ast
 import importlib.util
 import json
-import astor
 import ipynbname
 import sys
 
@@ -74,12 +73,14 @@ class PackageUsageRemover(ast.NodeTransformer):
 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(modified_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
@@ -452,4 +453,3 @@ if __name__ == "__main__":
     hash_id, zip_path = zip_list_of_unique_files(filepaths)
     print(f"Created zip file: {zip_path}")
     print(f"Hash ID: {hash_id}")
-

ragaai_catalyst/tracers/distributed.py

@@ -0,0 +1,313 @@
+"""
+Distributed tracing functionality for RagaAI Catalyst.
+Provides simplified initialization and decorator-based tracing.
+"""
+
+import os
+import threading
+from typing import Optional, Dict, Any, List
+from functools import wraps
+from contextlib import contextmanager
+import uuid
+from datetime import datetime
+import asyncio
+
+from .tracer import Tracer
+from ..ragaai_catalyst import RagaAICatalyst
+
+# Global state
+_global_tracer: Optional[Tracer] = None
+_global_catalyst: Optional[RagaAICatalyst] = None
+_tracer_lock = threading.Lock()
+_active_spans = threading.local()
+
+def get_current_tracer() -> Optional[Tracer]:
+    """Get the current global tracer instance."""
+    return _global_tracer
+
+def get_current_catalyst() -> Optional[RagaAICatalyst]:
+    """Get the current global catalyst instance."""
+    return _global_catalyst
+
+def init_tracing(
+    project_name: str = None,
+    dataset_name: str = None,
+    access_key: str = None,
+    secret_key: str = None,
+    base_url: str = None,
+    tracer: Tracer = None,
+    catalyst: RagaAICatalyst = None,
+    **kwargs
+) -> None:
+    """Initialize distributed tracing.
+    
+    Args:
+        project_name: Project name for new tracer
+        dataset_name: Dataset name for new tracer
+        access_key: RagaAI Catalyst access key
+        secret_key: RagaAI Catalyst secret key  
+        base_url: RagaAI Catalyst API base URL
+        tracer: Existing Tracer instance
+        catalyst: Existing RagaAICatalyst instance
+        **kwargs: Additional tracer configuration
+    """
+    global _global_tracer, _global_catalyst
+    
+    with _tracer_lock:
+        if tracer and catalyst:
+            _global_tracer = tracer
+            _global_catalyst = catalyst
+        else:
+            # Use env vars as fallback
+            access_key = access_key or os.getenv("RAGAAI_CATALYST_ACCESS_KEY")
+            secret_key = secret_key or os.getenv("RAGAAI_CATALYST_SECRET_KEY") 
+            base_url = base_url or os.getenv("RAGAAI_CATALYST_BASE_URL")
+
+            if not all([access_key, secret_key]):
+                raise ValueError(
+                    "Missing required credentials. Either provide access_key and secret_key "
+                    "or set RAGAAI_CATALYST_ACCESS_KEY and RAGAAI_CATALYST_SECRET_KEY environment variables."
+                )
+
+            _global_catalyst = RagaAICatalyst(
+                access_key=access_key,
+                secret_key=secret_key,
+                base_url=base_url
+            )
+
+            _global_tracer = Tracer(
+                project_name=project_name,
+                dataset_name=dataset_name,
+                **kwargs
+            )
+
+def trace_agent(name: str = None, agent_type: str = "generic", version: str = "1.0.0", **kwargs):
+    """Decorator for tracing agent functions."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        span_name = name or func.__name__
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Set current agent name and store the token
+            name_token = tracer.current_agent_name.set(span_name)
+            
+            try:
+                # Use async agent tracing
+                return await tracer._trace_agent_execution(
+                    func,
+                    span_name,
+                    agent_type,
+                    version,
+                    None,  # capabilities
+                    str(uuid.uuid4()),  # hash_id
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_agent_name.reset(name_token)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Set current agent name and store the token
+            name_token = tracer.current_agent_name.set(span_name)
+            
+            try:
+                # Use synchronous agent tracing
+                return tracer._trace_sync_agent_execution(
+                    func,
+                    span_name,
+                    agent_type,
+                    version,
+                    None,  # capabilities
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_agent_name.reset(name_token)
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+
+def trace_llm(name: str = None, model: str = None, **kwargs):
+    """Decorator for tracing LLM calls."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        span_name = name or func.__name__
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Set current LLM name and store the token
+            name_token = tracer.current_llm_call_name.set(span_name)
+            
+            try:
+                # Just execute the function within the current span
+                result = await func(*args, **kwargs)
+                return result
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_llm_call_name.reset(name_token)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Set current LLM name and store the token
+            name_token = tracer.current_llm_call_name.set(span_name)
+            
+            try:
+                # Just execute the function within the current span
+                result = func(*args, **kwargs)
+                return result
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_llm_call_name.reset(name_token)
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+    
+
+def trace_tool(name: str = None, tool_type: str = "generic", version: str = "1.0.0", **kwargs):
+    """Decorator for tracing tool functions."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        span_name = name or func.__name__
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Set current tool name and store the token
+            name_token = tracer.current_tool_name.set(span_name)
+            
+            try:
+                # Use async tool tracing
+                return await tracer._trace_tool_execution(
+                    func,
+                    span_name,
+                    tool_type,
+                    version,
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_tool_name.reset(name_token)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Set current tool name and store the token
+            name_token = tracer.current_tool_name.set(span_name)
+            
+            try:
+                # Use synchronous tool tracing
+                return tracer._trace_sync_tool_execution(
+                    func,
+                    span_name,
+                    tool_type,
+                    version,
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_tool_name.reset(name_token)
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+
+
+
+def trace_custom(name: str = None, custom_type: str = "generic", version: str = "1.0.0", trace_variables: bool = False, **kwargs):
+    """Decorator for tracing custom functions."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Use async tool tracing
+            return await tracer._trace_custom_execution(
+                func,
+                name or func.__name__,
+                custom_type,
+                version,
+                trace_variables,
+                *args,
+                **kwargs
+            )
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Use synchronous tool tracing
+            return tracer._trace_sync_custom_execution(
+                func,
+                name or func.__name__,
+                custom_type,
+                version,
+                trace_variables,
+                *args,
+                **kwargs
+            )
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+
+
+def current_span():
+    """Get the current active span for adding metrics."""
+    tracer = get_current_tracer()
+    if not tracer:
+        return None
+    
+    # First check for LLM context
+    llm_name = tracer.current_llm_call_name.get()
+    if llm_name:
+        return tracer.span(llm_name)
+    
+    # Then check for tool context
+    tool_name = tracer.current_tool_name.get()
+    if tool_name:
+        return tracer.span(tool_name)
+    
+    # Finally fall back to agent context
+    agent_name = tracer.current_agent_name.get()
+    if not agent_name:
+        raise ValueError("No active span found. Make sure you're calling this within a traced function.")
+    
+    return tracer.span(agent_name)

ragaai_catalyst/tracers/tracer.py

@@ -98,10 +98,10 @@ class Tracer(AgenticTracing):
                     "custom": False
                 }
         elif isinstance(auto_instrumentation, dict):
-            auto_instrumentation = {k: v for k, v in auto_instrumentation.items() if v}
+            auto_instrumentation = {k: v for k, v in auto_instrumentation.items()}
             for key in ["llm", "tool", "agent", "user_interaction", "file_io", "network", "custom"]:
                 if key not in auto_instrumentation:
-                    auto_instrumentation[key] = False
+                    auto_instrumentation[key] = True
         
         super().__init__(user_detail=user_detail, auto_instrumentation=auto_instrumentation)
 

{ragaai_catalyst-2.1.4.1b1.dist-info → ragaai_catalyst-2.1.5b1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ragaai_catalyst
-Version: 2.1.4.1b1
+Version: 2.1.5b1
 Summary: RAGA AI CATALYST
 Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>
 Requires-Python: <3.13,>=3.9
@@ -34,7 +34,6 @@ Requires-Dist: psutil~=6.0.0
 Requires-Dist: py-cpuinfo~=9.0.0
 Requires-Dist: requests~=2.32.3
 Requires-Dist: GPUtil~=1.4.0
-Requires-Dist: astor>=0.8.1
 Requires-Dist: ipynbname
 Provides-Extra: dev
 Requires-Dist: pytest; extra == "dev"