ragaai-catalyst 2.1.5b21__py3-none-any.whl → 2.1.5b23__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_local_metric.py

@@ -0,0 +1,72 @@
+import logging
+import os
+import requests
+
+from ragaai_catalyst import RagaAICatalyst
+
+logger = logging.getLogger(__name__)
+logging_level = (
+    logger.setLevel(logging.DEBUG)
+    if os.getenv("DEBUG")
+    else logger.setLevel(logging.INFO)
+)
+
+def calculate_metric(project_id, metric_name, model, provider, prompt, response, context, expected_response=None):
+    user_id = "1"
+    org_domain = "raga"
+
+    headers = {
+        "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+        "X-Project-Id": str(project_id),  # TODO to change it to project_id
+        "Content-Type": "application/json"
+    }
+
+    payload = {
+        "data": [
+            {
+                "metric_name": metric_name,
+                "metric_config": {
+                    "threshold": {
+                        "isEditable": True,
+                        "lte": 0.3
+                    },
+                    "model": model,
+                    "orgDomain": org_domain,
+                    "provider": provider,
+                    "user_id": user_id,
+                    "job_id": 1,
+                    "metric_name": metric_name,
+                    "request_id": 1
+                },
+                "trace_object": {
+                    "Data": {
+                        "DocId": "doc-1",
+                        "Prompt": prompt,
+                        "Response": response,
+                        "Context": context,
+                        "ExpectedResponse": "",
+                        "ExpectedContext": expected_response,
+                        "Chat": "",
+                        "Instructions": "",
+                        "SystemPrompt": "",
+                        "Text": ""
+                    },
+                    "claims": {},
+                    "last_computed_metrics": {
+                        metric_name: {
+                        }
+                    }
+                }
+            }
+        ]
+    }
+
+    try:
+        BASE_URL = RagaAICatalyst.BASE_URL
+        response = requests.post(f"{BASE_URL}/v1/llm/calculate-metric", headers=headers, json=payload, timeout=30)
+        logger.debug(f"Metric calculation response status {response.status_code}")
+        response.raise_for_status()
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        logger.debug(f"Error in calculate-metric api: {e}, payload: {payload}")
+        raise Exception(f"Error in calculate-metric: {e}")

ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py

@@ -1,27 +1,39 @@
+import logging
+
 import requests
 import os
 import json
 from ....ragaai_catalyst import RagaAICatalyst
 from ..utils.get_user_trace_metrics import get_user_trace_metrics
 
+logger = logging.getLogger(__name__)
+logging_level = (
+    logger.setLevel(logging.DEBUG)
+    if os.getenv("DEBUG")
+    else logger.setLevel(logging.INFO)
+)
+
+
 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)
 
         user_trace_metrics = get_user_trace_metrics(project_name, dataset_name)
         if user_trace_metrics:
             user_trace_metrics_list = [metric["displayName"] for metric in user_trace_metrics]
-        
+
         if user_trace_metrics:
             for metric in metrics:
                 if metric["displayName"] in user_trace_metrics_list:
-                    metricConfig = next((user_metric["metricConfig"] for user_metric in user_trace_metrics if user_metric["displayName"] == metric["displayName"]), None)
+                    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}.")
-
+                        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')}",
@@ -31,9 +43,9 @@ def upload_trace_metric(json_file_path, dataset_name, project_name):
             "datasetName": dataset_name,
             "metrics": metrics
         })
-        response = requests.request("POST", 
-                                    f"{RagaAICatalyst.BASE_URL}/v1/llm/trace/metrics", 
-                                    headers=headers, 
+        response = requests.request("POST",
+                                    f"{RagaAICatalyst.BASE_URL}/v1/llm/trace/metrics",
+                                    headers=headers,
                                     data=payload,
                                     timeout=10)
         if response.status_code != 200:
@@ -68,23 +80,28 @@ 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):
     formatted_metrics = []
     for metric in metrics:
-        if any(m["name"] == metric["name"] for m in formatted_metrics):
+        if any(m["name"] == metric.get("displayName") or m['name'] == metric.get("name") for m in formatted_metrics):
             continue
+        metric_display_name = metric["name"]
+        if metric.get("displayName"):
+            metric_display_name = metric['displayName']
         formatted_metrics.append({
-            "name": metric["name"],
-            "displayName": metric["name"],
+            "name": metric_display_name,
+            "displayName": metric_display_name,
             "config": {"source": "user"},
         })
-    return formatted_metrics
+    return formatted_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"
 
@@ -70,6 +71,9 @@ def extract_parameters(kwargs):
     # Remove messages key in parameters (OpenAI message)
     if 'messages' in parameters:
         del parameters['messages']
+        
+    if 'run_manager' in parameters:
+        del parameters['run_manager']
 
     if 'generation_config' in parameters:
         generation_config = parameters['generation_config']
@@ -268,10 +272,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 +385,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 +539,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)