ragaai-catalyst 2.1.5b20__py3-none-any.whl → 2.1.5b22__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py

@@ -258,7 +258,10 @@ class ToolTracerMixin:
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
                 async_wrapper.metadata = metadata
-                self.gt = kwargs.get("gt", None) if kwargs else None
+                gt = kwargs.get("gt") if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return await self._trace_tool_execution(
                     func, name, tool_type, version, *args, **kwargs
                 )
@@ -267,7 +270,10 @@ class ToolTracerMixin:
             @functools.wraps(func)
             def sync_wrapper(*args, **kwargs):
                 sync_wrapper.metadata = metadata
-                self.gt = kwargs.get("gt", None) if kwargs else None
+                gt = kwargs.get("gt") if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return self._trace_sync_tool_execution(
                     func, name, tool_type, version, *args, **kwargs
                 )
@@ -302,7 +308,7 @@ class ToolTracerMixin:
 
         try:
             # Execute the tool
-            result = func(*args, **kwargs)
+            result = self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_memory = psutil.Process().memory_info().rss
@@ -384,7 +390,7 @@ class ToolTracerMixin:
         self.start_component(component_id)
         try:
             # Execute the tool
-            result = await func(*args, **kwargs)
+            result = await self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_memory = psutil.Process().memory_info().rss
@@ -504,8 +510,13 @@ class ToolTracerMixin:
             "interactions": interactions,
         }
 
-        if self.gt:
-            component["data"]["gt"] = self.gt
+        if name in self.span_attributes_dict:
+            span_gt = self.span_attributes_dict[name].gt
+            if span_gt is not None:
+                component["data"]["gt"] = span_gt
+            span_context = self.span_attributes_dict[name].context
+            if span_context:
+                component["data"]["context"] = span_context
 
         # Reset the SpanAttributes context variable
         self.span_attributes_dict[kwargs["name"]] = SpanAttributes(kwargs["name"])

ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py

@@ -8,6 +8,7 @@ def upload_trace_metric(json_file_path, dataset_name, project_name):
     try:
         with open(json_file_path, "r") as f:
             traces = json.load(f)
+                    
         metrics = get_trace_metrics_from_trace(traces)
         metrics = _change_metrics_format_for_payload(metrics)
 
@@ -21,7 +22,6 @@ def upload_trace_metric(json_file_path, dataset_name, project_name):
                     metricConfig = next((user_metric["metricConfig"] for user_metric in user_trace_metrics if user_metric["displayName"] == metric["displayName"]), None)
                     if not metricConfig or metricConfig.get("Metric Source", {}).get("value") != "user":
                         raise ValueError(f"Metrics {metric['displayName']} already exist in dataset {dataset_name} of project {project_name}.")
-
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
@@ -68,13 +68,14 @@ def get_trace_metrics_from_trace(traces):
     # get span level metrics
     for span in traces["data"][0]["spans"]:
         if span["type"] == "agent":
+            # Add children metrics of agent
             children_metric = _get_children_metrics_of_agent(span["data"]["children"])
             if children_metric:
                 metrics.extend(children_metric)
-        else:
-            metric = span.get("metrics", [])
-            if metric:
-                metrics.extend(metric)
+
+        metric = span.get("metrics", [])
+        if metric:
+            metrics.extend(metric)
     return metrics
 
 def _change_metrics_format_for_payload(metrics):

ragaai_catalyst/tracers/agentic_tracing/utils/file_name_tracker.py

@@ -8,13 +8,32 @@ class TrackName:
     def trace_decorator(self, func):
         @wraps(func)
         def wrapper(*args, **kwargs):
-            file_name = self._get_file_name()
+            file_name = self._get_decorated_file_name()
             self.files.add(file_name)
 
             return func(*args, **kwargs)
         return wrapper
+    
+    def trace_wrapper(self, func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            file_name = self._get_wrapped_file_name()
+            self.files.add(file_name)
+            return func(*args, **kwargs)
+        return wrapper
+    
+    def _get_wrapped_file_name(self):
+        try:
+            from IPython import get_ipython
+            if 'IPKernelApp' in get_ipython().config:
+                return self._get_notebook_name()
+        except Exception:
+            pass
+
+        frame = inspect.stack()[4]
+        return frame.filename
 
-    def _get_file_name(self):
+    def _get_decorated_file_name(self):
         # Check if running in a Jupyter notebook
         try:
             from IPython import get_ipython

ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py

@@ -2,9 +2,9 @@ from ..data.data_structure import LLMCall
 from .trace_utils import (
     calculate_cost,
     convert_usage_to_dict,
-    load_model_costs,
 )
 from importlib import resources
+from litellm import model_cost
 import json
 import os
 import asyncio
@@ -38,7 +38,13 @@ def extract_model_name(args, kwargs, result):
                 metadata = manager.metadata
                 model_name = metadata.get('ls_model_name', None)
                 if model_name:
-                    model = model_name         
+                    model = model_name       
+                    
+    if not model:
+        if 'to_dict' in dir(result):
+            result = result.to_dict()
+            if 'model_version' in result:
+                model = result['model_version']  
     
     
     # Normalize Google model names
@@ -50,11 +56,6 @@ def extract_model_name(args, kwargs, result):
             return "gemini-1.5-pro"
         if "gemini-pro" in model:
             return "gemini-pro"
-
-    if 'to_dict' in dir(result):
-        result = result.to_dict()
-        if 'model_version' in result:
-            model = result['model_version']
     
     return model or "default"
 
@@ -268,10 +269,21 @@ def num_tokens_from_messages(model="gpt-4o-mini-2024-07-18", prompt_messages=Non
     }
 
 def extract_input_data(args, kwargs, result):
-    """Extract input data from function call"""
+    """Sanitize and format input data, including handling of nested lists and dictionaries."""
+
+    def sanitize_value(value):
+        if isinstance(value, (int, float, bool, str)):
+            return value
+        elif isinstance(value, list):
+            return [sanitize_value(item) for item in value]
+        elif isinstance(value, dict):
+            return {key: sanitize_value(val) for key, val in value.items()}
+        else:
+            return str(value)  # Convert non-standard types to string
+
     return {
-        'args': args,
-        'kwargs': kwargs
+        "args": [sanitize_value(arg) for arg in args],
+        "kwargs": {key: sanitize_value(val) for key, val in kwargs.items()},
     }
 
 
@@ -370,6 +382,13 @@ def extract_llm_output(result):
                         })
         return OutputResponse(output)
     
+    # Handle AIMessage Format
+    if hasattr(result, "content"):
+        return OutputResponse([{
+            "content": result.content,
+            "role": getattr(result, "role", "assistant")
+        }])
+    
     # Handle Vertex AI format
     # format1
     if hasattr(result, "text"):
@@ -517,7 +536,7 @@ def extract_llm_data(args, kwargs, result):
     token_usage = extract_token_usage(result)
 
     # Load model costs
-    model_costs = load_model_costs()
+    model_costs = model_cost
 
     # Calculate cost
     cost = calculate_llm_cost(token_usage, model_name, model_costs)