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

ragaai_catalyst/__init__.py

@@ -1,6 +1,5 @@
 from .experiment import Experiment
 from .ragaai_catalyst import RagaAICatalyst
-from .tracers import Tracer
 from .utils import response_checker
 from .dataset import Dataset
 from .prompt_manager import PromptManager
@@ -8,6 +7,24 @@ from .evaluation import Evaluation
 from .synthetic_data_generation import SyntheticDataGeneration
 from .guardrails_manager import GuardrailsManager
 from .guard_executor import GuardExecutor
+from .tracers import Tracer, init_tracing, trace_agent, trace_llm, trace_tool, current_span, trace_custom
 
 
-__all__ = ["Experiment", "RagaAICatalyst", "Tracer", "PromptManager", "Evaluation","SyntheticDataGeneration", "GuardrailsManager"]
+
+
+__all__ = [
+    "Experiment", 
+    "RagaAICatalyst", 
+    "Tracer", 
+    "PromptManager", 
+    "Evaluation",
+    "SyntheticDataGeneration", 
+    "GuardrailsManager", 
+    "GuardExecutor",
+    "init_tracing", 
+    "trace_agent", 
+    "trace_llm",
+    "trace_tool",
+    "current_span",
+    "trace_custom"
+]

ragaai_catalyst/tracers/__init__.py

@@ -1,3 +1,19 @@
 from .tracer import Tracer
+from .distributed import (
+    init_tracing,
+    trace_agent,
+    trace_llm,
+    trace_tool,
+    current_span,
+    trace_custom,
+)
 
-__all__ = ["Tracer"] 
+__all__ = [
+    "Tracer",
+    "init_tracing",
+    "trace_agent", 
+    "trace_llm",
+    "trace_tool",
+    "current_span",
+    "trace_custom"
+]

ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py

@@ -271,7 +271,7 @@ class ComponentInfo:
     cost: Optional[Dict[str, float]] = None
 
 class Trace:
-    def __init__(self, id: str, trace_name: 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):
+    def __init__(self, id: str, trace_name: 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, metrics: Optional[List[Dict[str, Any]]] = None):
         self.id = id
         self.trace_name = trace_name
         self.project_name = project_name
@@ -280,6 +280,7 @@ class Trace:
         self.metadata = metadata or Metadata()
         self.data = data or []
         self.replays = replays
+        self.metrics = metrics or []
 
     def to_dict(self):
         return {
@@ -288,7 +289,8 @@ class Trace:
             "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,
+            "metadata": self.metadata,
             "data": self.data,
             "replays": self.replays,
+            "metrics": self.metrics
         }

ragaai_catalyst/tracers/agentic_tracing/tracers/base.py

@@ -5,7 +5,7 @@ import psutil
 import pkg_resources
 from datetime import datetime
 from pathlib import Path
-from typing import List, Any
+from typing import List, Any, Dict
 import uuid
 import sys
 import tempfile
@@ -81,6 +81,7 @@ class BaseTracer:
         self.project_id = self.user_details["project_id"]  # Access the project_id
         self.trace_name = self.user_details["trace_name"]  # Access the trace_name
         self.visited_metrics = []
+        self.trace_metrics = []  # Store metrics here
 
         # Initialize trace data
         self.trace_id = None
@@ -211,6 +212,10 @@ class BaseTracer:
         threading.Thread(target=self._track_disk_usage).start()
         threading.Thread(target=self._track_network_usage).start()
 
+        # Reset metrics
+        self.visited_metrics = []
+        self.trace_metrics = []
+
         metadata = Metadata(
             cost={},
             tokens={},
@@ -241,6 +246,7 @@ class BaseTracer:
             metadata=metadata,
             data=self.data_key,
             replays={"source": None},
+            metrics=[]  # Initialize empty metrics list
         )
 
     def stop(self):
@@ -300,13 +306,16 @@ class BaseTracer:
             # replace source code with zip_path
             self.trace.metadata.system_info.source_code = hash_id
 
+            # Add metrics to trace before saving
+            trace_data = self.trace.to_dict()
+            trace_data["metrics"] = self.trace_metrics
+            
             # Clean up trace_data before saving
-            trace_data = self.trace.__dict__
             cleaned_trace_data = self._clean_trace(trace_data)
 
             # Format interactions and add to trace
             interactions = self.format_interactions()
-            self.trace.workflow = interactions["workflow"]
+            trace_data["workflow"] = interactions["workflow"]
 
             with open(filepath, "w") as f:
                 json.dump(cleaned_trace_data, f, cls=TracerJSONEncoder, indent=2)
@@ -424,38 +433,44 @@ class BaseTracer:
     def _extract_cost_tokens(self, trace):
         cost = {}
         tokens = {}
-        for span in trace.data[0]["spans"]:
-            if span.type == "llm":
-                info = span.info
-                if isinstance(info, dict):
-                    cost_info = info.get("cost", {})
-                    for key, value in cost_info.items():
-                        if key not in cost:
-                            cost[key] = 0
-                        cost[key] += value
-                    token_info = info.get("tokens", {})
-                    for key, value in token_info.items():
-                        if key not in tokens:
-                            tokens[key] = 0
-                        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"]
-                    if isinstance(info, dict):
-                        cost_info = info.get("cost", {})
-                        for key, value in cost_info.items():
-                            if key not in cost:
-                                cost[key] = 0
-                            cost[key] += value
-                        token_info = info.get("tokens", {})
-                        for key, value in token_info.items():
-                            if key not in tokens:
-                                tokens[key] = 0
-                            tokens[key] += value
+
+        def process_span_info(info):
+            if not isinstance(info, dict):
+                return
+            cost_info = info.get("cost", {})
+            for key, value in cost_info.items():
+                if key not in cost:
+                    cost[key] = 0
+                cost[key] += value
+            token_info = info.get("tokens", {})
+            for key, value in token_info.items():
+                if key not in tokens:
+                    tokens[key] = 0
+                tokens[key] += value
+
+        def process_spans(spans):
+            for span in spans:
+                # Get span type, handling both span objects and dictionaries
+                span_type = span.type if hasattr(span, 'type') else span.get('type')
+                span_info = span.info if hasattr(span, 'info') else span.get('info', {})
+                span_data = span.data if hasattr(span, 'data') else span.get('data', {})
+
+                # Process direct LLM spans
+                if span_type == "llm":
+                    process_span_info(span_info)
+                # Process agent spans recursively
+                elif span_type == "agent":
+                    # Process LLM children in the current agent span
+                    children = span_data.get("children", [])
+                    for child in children:
+                        child_type = child.get("type")
+                        if child_type == "llm":
+                            process_span_info(child.get("info", {}))
+                        # Recursively process nested agent spans
+                        elif child_type == "agent":
+                            process_spans([child])
+
+        process_spans(trace.data[0]["spans"])
         trace.metadata.cost = cost
         trace.metadata.tokens = tokens
         return trace
@@ -609,9 +624,7 @@ class BaseTracer:
                     "span_id": child.get("id"),
                     "interaction_type": "llm_call_end",
                     "name": child.get("name"),
-                    "content": {
-                        "response": child.get("data", {}).get("output")
-                    },
+                    "content": {"response": child.get("data", {}).get("output")},
                     "timestamp": child.get("end_time"),
                     "error": child.get("error"),
                 }
@@ -657,7 +670,7 @@ class BaseTracer:
                 {
                     "id": str(interaction_id),
                     "span_id": child.get("id"),
-                    "interaction_type": child_type,
+                    "interaction_type": f"{child_type}_call_start",
                     "name": child.get("name"),
                     "content": child.get("data", {}),
                     "timestamp": child.get("start_time"),
@@ -665,6 +678,19 @@ class BaseTracer:
                 }
             )
             interaction_id += 1
+            
+            interactions.append(
+                {
+                    "id": str(interaction_id),
+                    "span_id": child.get("id"),
+                    "interaction_type": f"{child_type}_call_end",
+                    "name": child.get("name"),
+                    "content": child.get("data", {}),
+                    "timestamp": child.get("end_time"),
+                    "error": child.get("error"),
+                }
+            )
+            interaction_id += 1
 
         # Process additional interactions and network calls
         if "interactions" in child:
@@ -825,7 +851,7 @@ class BaseTracer:
                     {
                         "id": str(interaction_id),
                         "span_id": span.id,
-                        "interaction_type": span.type,
+                        "interaction_type": f"{span.type}_call_start",
                         "name": span.name,
                         "content": span.data,
                         "timestamp": span.start_time,
@@ -833,6 +859,19 @@ class BaseTracer:
                     }
                 )
                 interaction_id += 1
+                
+                interactions.append(
+                    {
+                        "id": str(interaction_id),
+                        "span_id": span.id,
+                        "interaction_type": f"{span.type}_call_end",
+                        "name": span.name,
+                        "content": span.data,
+                        "timestamp": span.end_time,
+                        "error": span.error,
+                    }
+                )
+                interaction_id += 1
 
             # Process interactions from span.data if they exist
             if span.interactions:
@@ -882,8 +921,90 @@ class BaseTracer:
 
         return {"workflow": sorted_interactions}
 
+    def add_metrics(
+        self,
+        name: str | List[Dict[str, Any]] | Dict[str, Any] = None,
+        score: float | int = None,
+        reasoning: str = "",
+        cost: float = None,
+        latency: float = None,
+        metadata: Dict[str, Any] = None,
+        config: Dict[str, Any] = None,
+    ):
+        """Add metrics at the trace level.
+
+        Can be called in two ways:
+        1. With individual parameters:
+           tracer.add_metrics(name="metric_name", score=0.9, reasoning="Good performance")
+           
+        2. With a dictionary or list of dictionaries:
+           tracer.add_metrics({"name": "metric_name", "score": 0.9})
+           tracer.add_metrics([{"name": "metric1", "score": 0.9}, {"name": "metric2", "score": 0.8}])
+
+        Args:
+            name: Either the metric name (str) or a metric dictionary/list of dictionaries
+            score: Score value (float or int) when using individual parameters
+            reasoning: Optional explanation for the score
+            cost: Optional cost associated with the metric
+            latency: Optional latency measurement
+            metadata: Optional additional metadata as key-value pairs
+            config: Optional configuration parameters
+        """
+        if not hasattr(self, 'trace'):
+            logger.warning("Cannot add metrics before trace is initialized. Call start() first.")
+            return
+
+        # Convert individual parameters to metric dict if needed
+        if isinstance(name, str):
+            metrics = [{
+                "name": name,
+                "score": score,
+                "reasoning": reasoning,
+                "cost": cost,
+                "latency": latency,
+                "metadata": metadata or {},
+                "config": config or {}
+            }]
+        else:
+            # Handle dict or list input
+            metrics = name if isinstance(name, list) else [name] if isinstance(name, dict) else []
+
+        try:
+            for metric in metrics:
+                if not isinstance(metric, dict):
+                    raise ValueError(f"Expected dict, got {type(metric)}")
+                
+                if "name" not in metric or "score" not in metric:
+                    raise ValueError("Metric must contain 'name' and 'score' fields")
+
+                # Handle duplicate metric names
+                metric_name = metric["name"]
+                if metric_name in self.visited_metrics:
+                    count = sum(1 for m in self.visited_metrics if m.startswith(metric_name))
+                    metric_name = f"{metric_name}_{count + 1}"
+                self.visited_metrics.append(metric_name)
+
+                formatted_metric = {
+                    "name": metric_name,  
+                    "score": metric["score"],
+                    "reason": metric.get("reasoning", ""),
+                    "source": "user",
+                    "cost": metric.get("cost"),
+                    "latency": metric.get("latency"),
+                    "metadata": metric.get("metadata", {}),
+                    "mappings": [],
+                    "config": metric.get("config", {})
+                }
+                
+                self.trace_metrics.append(formatted_metric)
+                logger.debug(f"Added trace-level metric: {formatted_metric}")
+
+        except ValueError as e:
+            logger.error(f"Validation Error: {e}")
+        except Exception as e:
+            logger.error(f"Error adding metric: {e}")
+    
     def span(self, span_name):
         if span_name not in self.span_attributes_dict:
             self.span_attributes_dict[span_name] = SpanAttributes(span_name)
-        return self.span_attributes_dict[span_name]
-    
+        return self.span_attributes_dict[span_name]

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -475,7 +475,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

@@ -137,13 +137,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"""

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/upload/upload_trace_metric.py

@@ -59,6 +59,13 @@ def _get_children_metrics_of_agent(children_traces):
 
 def get_trace_metrics_from_trace(traces):
     metrics = []
+
+    # get trace level metrics
+    if "metrics" in traces.keys():
+        if len(traces["metrics"]) > 0:
+            metrics.extend(traces["metrics"])
+
+    # get span level metrics
     for span in traces["data"][0]["spans"]:
         if span["type"] == "agent":
             children_metric = _get_children_metrics_of_agent(span["data"]["children"])

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,286 @@
+"""
+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)
+        
+        @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_tool_execution(
+                func,
+                name or func.__name__,
+                tool_type,
+                version,
+                *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_tool_execution(
+                func,
+                name or func.__name__,
+                tool_type,
+                version,
+                *args,
+                **kwargs
+            )
+
+        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
+    
+    # Get the current agent name from 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-2.1.4.1b0.dist-info → ragaai_catalyst-2.1.5b0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ragaai_catalyst
-Version: 2.1.4.1b0
+Version: 2.1.5b0
 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"

{ragaai_catalyst-2.1.4.1b0.dist-info → ragaai_catalyst-2.1.5b0.dist-info}/RECORD

@@ -1,4 +1,4 @@
-ragaai_catalyst/__init__.py,sha256=BdIJ_UUre0uEnRTsLw_hE0C0muWk6XWNZqdVOel22R4,537
+ragaai_catalyst/__init__.py,sha256=tvESV8UuVtth14E89wQxgf0WvQZSApWfDeLiIdmMhkE,791
 ragaai_catalyst/_version.py,sha256=JKt9KaVNOMVeGs8ojO6LvIZr7ZkMzNN-gCcvryy4x8E,460
 ragaai_catalyst/dataset.py,sha256=aTRvZicAXmrC0KdmmsoJH_rtEJrxbqYmf1P806c1Wg0,10521
 ragaai_catalyst/evaluation.py,sha256=34H2bYZNSrcu0jMQgDZw1OLVbQU80PaVLo2avju8POM,20311
@@ -11,14 +11,15 @@ ragaai_catalyst/proxy_call.py,sha256=CHxldeceZUaLU-to_hs_Kf1z_b2vHMssLS_cOBedu78
 ragaai_catalyst/ragaai_catalyst.py,sha256=FdqMzwuQLqS2-3JJDsTQ8uh2itllOxfPrRUjb8Kwmn0,17428
 ragaai_catalyst/synthetic_data_generation.py,sha256=uDV9tNwto2xSkWg5XHXUvjErW-4P34CTrxaJpRfezyA,19250
 ragaai_catalyst/utils.py,sha256=TlhEFwLyRU690HvANbyoRycR3nQ67lxVUQoUOfTPYQ0,3772
-ragaai_catalyst/tracers/__init__.py,sha256=yxepo7iVjTNI_wFdk3Z6Ghu64SazVyszCPEHYrX5WQk,50
+ragaai_catalyst/tracers/__init__.py,sha256=LfgTes-nHpazssbGKnn8kyLZNr49kIPrlkrqqoTFTfc,301
+ragaai_catalyst/tracers/distributed.py,sha256=_KO7swUDoPErbfahxNSRpxQk22rstQItfPGP8e4l7E4,9334
 ragaai_catalyst/tracers/llamaindex_callback.py,sha256=ZY0BJrrlz-P9Mg2dX-ZkVKG3gSvzwqBtk7JL_05MiYA,14028
 ragaai_catalyst/tracers/tracer.py,sha256=UX-01NYWcH2y4UW1W287Cn-jy760rgaFqu8llJbeMdg,15654
 ragaai_catalyst/tracers/upload_traces.py,sha256=mT5rverNUL5Rcal9VR5_c75wHBAUrm2pvYetTZqP3ok,4796
 ragaai_catalyst/tracers/agentic_tracing/README.md,sha256=X4QwLb7-Jg7GQMIXj-SerZIgDETfw-7VgYlczOR8ZeQ,4508
 ragaai_catalyst/tracers/agentic_tracing/__init__.py,sha256=yf6SKvOPSpH-9LiKaoLKXwqj5sez8F_5wkOb91yp0oE,260
 ragaai_catalyst/tracers/agentic_tracing/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py,sha256=nFnwqL1-Uznwndi2ugDnhziUbIASlcBYnM6Dyq7pPt8,9243
+ragaai_catalyst/tracers/agentic_tracing/data/data_structure.py,sha256=icAtNzKN_I0YtfuJ3RF8BdZJK3ohqxkVZIdvM5_YugY,9327
 ragaai_catalyst/tracers/agentic_tracing/tests/FinancialAnalysisSystem.ipynb,sha256=0qZxjWqYCTAVvdo3Tsp544D8Am48wfeMQ9RKpKgAL8g,34291
 ragaai_catalyst/tracers/agentic_tracing/tests/GameActivityEventPlanner.ipynb,sha256=QCMFJYbGX0fd9eMW4PqyQLZjyWuTXo7n1nqO_hMLf0s,4225
 ragaai_catalyst/tracers/agentic_tracing/tests/TravelPlanner.ipynb,sha256=fU3inXoemJbdTkGAQl_N1UwVEZ10LrKv4gCEpbQ4ISg,43481
@@ -27,18 +28,18 @@ ragaai_catalyst/tracers/agentic_tracing/tests/ai_travel_agent.py,sha256=S4rCcKzU
 ragaai_catalyst/tracers/agentic_tracing/tests/unique_decorator_test.py,sha256=Xk1cLzs-2A3dgyBwRRnCWs7Eubki40FVonwd433hPN8,4805
 ragaai_catalyst/tracers/agentic_tracing/tracers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py,sha256=aLiq5nPie5TT61QYtvAtvErsxjPFYiUxjayn5aCX1_k,25543
-ragaai_catalyst/tracers/agentic_tracing/tracers/base.py,sha256=u2DX_BsMGsuJkWqN6ucxboMO8GnvOcArKJe5L08y-YI,35014
+ragaai_catalyst/tracers/agentic_tracing/tracers/base.py,sha256=XZH5qi8eMIHG_qjedRHa3lGYp0NiAtCUxN9P_uW_NMo,39981
 ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py,sha256=uay8lU7T-CKsVu8KvWX31qfMqufK9S3Ive7XKo2Ksmk,12252
 ragaai_catalyst/tracers/agentic_tracing/tracers/langgraph_tracer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=fn5qxb365GmQkJy_yZAY5TiiWRMFKPNJdYk8KFr8uWA,29343
-ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py,sha256=nxUOQSyWBTnbsZfxmr1lje2OggqNf9fwtGUb-sBo6mI,15215
-ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py,sha256=CviGiAg0W0krJxORMBDTosQytIoJDQ5RwU6xt_U_mOg,10408
+ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=hU97LWEdqFL-tnIb_Qlk6VJdS6yVPH0QVzVzg1yuCB0,29341
+ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py,sha256=9K1swf0IOEEzgpkO2A_k_Yxk6ybLRXDcVoQAzm-_fWs,15889
+ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py,sha256=m8CxYkl7iMiFya_lNwN1ykBc3Pmo-2pR_2HmpptwHWQ,10352
 ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py,sha256=4h6epgRHPVRfTrCznwc6KPK6xevwBaZwE-q0Zj-4Kkk,12160
 ragaai_catalyst/tracers/agentic_tracing/tracers/user_interaction_tracer.py,sha256=bhSUhNQCuJXKjgJAXhjKEYjnHMpYN90FSZdR84fNIKU,4614
 ragaai_catalyst/tracers/agentic_tracing/upload/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ragaai_catalyst/tracers/agentic_tracing/upload/upload_agentic_traces.py,sha256=1MDKXAAPzOEdxFKWWQrRgrmM3kz--DGXSywGXQmR3lQ,6041
 ragaai_catalyst/tracers/agentic_tracing/upload/upload_code.py,sha256=HgpMgI-JTWZrizcM7GGUIaAgaZF4aRT3D0dJXVEkblY,4271
-ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py,sha256=iFagpPltlg6aKvdyAFvXsuxyjUUjcHAMmvXlevL-uYM,3312
+ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py,sha256=id66gfx-XYj_zsAmicBKojBOqJQ__FJLSoZ0db56aes,3493
 ragaai_catalyst/tracers/agentic_tracing/utils/__init__.py,sha256=XdB3X_ufe4RVvGorxSqAiB9dYv4UD7Hvvuw3bsDUppY,60
 ragaai_catalyst/tracers/agentic_tracing/utils/api_utils.py,sha256=JyNCbfpW-w4O9CjtemTqmor2Rh1WGpQwhRaDSRmBxw8,689
 ragaai_catalyst/tracers/agentic_tracing/utils/create_dataset_schema.py,sha256=lgvJL0cakJrX8WGsnU05YGvotequSj6HgSohyR4OJNE,804
@@ -50,7 +51,7 @@ ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json,sha256=weznWpRd6N
 ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py,sha256=MqeRNGxzeuh9qTK0NbYMftl9V9Z0V7gMgBoHkrXP56k,1592
 ragaai_catalyst/tracers/agentic_tracing/utils/trace_utils.py,sha256=RciiDdo2riibEoM8X0FKHaXi78y3bWwNkV8U0leqigk,3508
 ragaai_catalyst/tracers/agentic_tracing/utils/unique_decorator.py,sha256=DQHjcEuqEKsNSWaNs7SoOaq50yK4Jsl966S7mBnV-zA,5723
-ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py,sha256=fvL7SQSO1-FPQ3zT1ErAcj8cLCKn4rUSG0Ug-nL-By0,18824
+ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py,sha256=8SAN4pjA3XKBWlUZiDGXGh26ve1qRH_fv0VwewsfCFc,18818
 ragaai_catalyst/tracers/exporters/__init__.py,sha256=kVA8zp05h3phu4e-iHSlnznp_PzMRczB7LphSsZgUjg,138
 ragaai_catalyst/tracers/exporters/file_span_exporter.py,sha256=RgGteu-NVGprXKkynvyIO5yOjpbtA41R3W_NzCjnkwE,6445
 ragaai_catalyst/tracers/exporters/raga_exporter.py,sha256=6xvjWXyh8XPkHKSLLmAZUQSvwuyY17ov8pv2VdfI0qA,17875
@@ -60,8 +61,8 @@ ragaai_catalyst/tracers/instrumentators/llamaindex.py,sha256=SMrRlR4xM7k9HK43hak
 ragaai_catalyst/tracers/instrumentators/openai.py,sha256=14R4KW9wQCR1xysLfsP_nxS7cqXrTPoD8En4MBAaZUU,379
 ragaai_catalyst/tracers/utils/__init__.py,sha256=KeMaZtYaTojilpLv65qH08QmpYclfpacDA0U3wg6Ybw,64
 ragaai_catalyst/tracers/utils/utils.py,sha256=ViygfJ7vZ7U0CTSA1lbxVloHp4NSlmfDzBRNCJuMhis,2374
-ragaai_catalyst-2.1.4.1b0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ragaai_catalyst-2.1.4.1b0.dist-info/METADATA,sha256=V5fmS2GcC5Ykd7yakB0_gDvqRPKHxzTtkp7TCt6T3O4,12794
-ragaai_catalyst-2.1.4.1b0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ragaai_catalyst-2.1.4.1b0.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
-ragaai_catalyst-2.1.4.1b0.dist-info/RECORD,,
+ragaai_catalyst-2.1.5b0.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ragaai_catalyst-2.1.5b0.dist-info/METADATA,sha256=R0r0eV49WTeIO0O1NlkFk99eIbAPnvqf0urcA-GpF48,12764
+ragaai_catalyst-2.1.5b0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ragaai_catalyst-2.1.5b0.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
+ragaai_catalyst-2.1.5b0.dist-info/RECORD,,