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

ragaai_catalyst/tracers/agentic_tracing/upload/upload_local_metric.py

@@ -0,0 +1,74 @@
+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, **kwargs):
+    user_id = "1"
+    org_domain = "raga"
+
+    headers = {
+        "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+        "X-Project-Id": str(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
+                },
+                "variable_mapping": kwargs,
+                "trace_object": {
+                    "Data": {
+                        "DocId": "doc-1",
+                        "Prompt": kwargs.get("prompt"),
+                        "Response": kwargs.get("response"),
+                        "Context": kwargs.get("context"),
+                        "ExpectedResponse": kwargs.get("expected_response"),
+                        "ExpectedContext": kwargs.get("expected_context"),
+                        "Chat": kwargs.get("chat"),
+                        "Instructions": kwargs.get("instructions"),
+                        "SystemPrompt": kwargs.get("system_prompt"),
+                        "Text": kwargs.get("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,40 @@
+import logging
+
 import requests
 import os
 import json
+import time
 from ....ragaai_catalyst import RagaAICatalyst
 from ..utils.get_user_trace_metrics import get_user_trace_metrics
 
-def upload_trace_metric(json_file_path, dataset_name, project_name):
+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, base_url=None):
     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,11 +44,17 @@ 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, 
+        url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
+        start_time = time.time()
+        endpoint = f"{url_base}/v1/llm/trace/metrics"
+        response = requests.request("POST",
+                                    endpoint,
+                                    headers=headers,
                                     data=payload,
                                     timeout=10)
+        elapsed_ms = (time.time() - start_time) * 1000
+        logger.debug(
+            f"API Call: [POST] {endpoint} | Status: {response.status_code} | Time: {elapsed_ms:.2f}ms")
         if response.status_code != 200:
             raise ValueError(f"Error inserting agentic trace metrics")
     except requests.exceptions.RequestException as e:
@@ -59,25 +78,37 @@ 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":
+            # 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/create_dataset_schema.py

@@ -4,7 +4,7 @@ import re
 import requests
 from ragaai_catalyst.tracers.agentic_tracing.tracers.base import RagaAICatalyst
 
-def create_dataset_schema_with_trace(project_name, dataset_name):
+def create_dataset_schema_with_trace(project_name, dataset_name, base_url=None):
     def make_request():
         headers = {
             "Content-Type": "application/json",
@@ -15,8 +15,10 @@ def create_dataset_schema_with_trace(project_name, dataset_name):
             "datasetName": dataset_name,
             "traceFolderUrl": None,
         })
+        # Use provided base_url or fall back to default
+        url_base = base_url if base_url is not None else RagaAICatalyst.BASE_URL
         response = requests.request("POST",
-            f"{RagaAICatalyst.BASE_URL}/v1/llm/dataset/logs",
+            f"{url_base}/v1/llm/dataset/logs",
             headers=headers,
             data=payload,
             timeout=10

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
@@ -43,4 +62,8 @@ class TrackName:
         
     def reset(self):
         """Reset the file tracker by clearing all tracked files."""
-        self.files.clear()
+        self.files.clear()
+    
+    def trace_main_file(self):
+        frame = inspect.stack()[-1]
+        self.files.add(frame.filename)

ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py

@@ -2,14 +2,30 @@ 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
 import psutil
+import tiktoken
+import logging
 
+logger = logging.getLogger(__name__)
+
+def get_model_cost():
+    """Load model costs from a JSON file. 
+    Note: This file should be updated periodically or whenever a new package is created to ensure accurate cost calculations.
+    To Do: Implement to do this automatically.
+    """
+    file="model_prices_and_context_window_backup.json"
+    d={}
+    with resources.open_text("ragaai_catalyst.tracers.utils", file) as f:
+        d= json.load(f)
+    return d 
+
+model_cost = get_model_cost()
 
 def extract_model_name(args, kwargs, result):
     """Extract model name from kwargs or result"""
@@ -35,7 +51,18 @@ 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']  
+    try:
+        if not model:
+            model = result.raw.model
+    except Exception as e:
+        pass
     
     
     # Normalize Google model names
@@ -48,10 +75,9 @@ def extract_model_name(args, kwargs, result):
         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']
+    if 'response_metadata' in dir(result):
+        if 'model_name' in result.response_metadata:
+            model = result.response_metadata['model_name']
     
     return model or "default"
 
@@ -67,6 +93,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']
@@ -91,8 +120,8 @@ def extract_token_usage(result):
         # Run the coroutine in the current event loop
         result = loop.run_until_complete(result)
 
-    # Handle text attribute responses (JSON string or Vertex AI)
-    if hasattr(result, "text"):
+    # Handle text attribute responses (JSON string for Vertex AI)
+    if hasattr(result, "text") and isinstance(result.text, (str, bytes, bytearray)):
         # First try parsing as JSON for OpenAI responses
         try:
             import json
@@ -137,10 +166,34 @@ def extract_token_usage(result):
     # Handle Google GenerativeAI format with usage_metadata
     if hasattr(result, "usage_metadata"):
         metadata = result.usage_metadata
+        if hasattr(metadata, "prompt_token_count"):
+            return {
+                "prompt_tokens": getattr(metadata, "prompt_token_count", 0),
+                "completion_tokens": getattr(metadata, "candidates_token_count", 0),
+                "total_tokens": getattr(metadata, "total_token_count", 0)
+            }
+        elif hasattr(metadata, "input_tokens"):
+            return {
+                "prompt_tokens": getattr(metadata, "input_tokens", 0),
+                "completion_tokens": getattr(metadata, "output_tokens", 0),
+                "total_tokens": getattr(metadata, "total_tokens", 0)
+            }
+        elif "input_tokens" in metadata:
+            return {
+                "prompt_tokens": metadata["input_tokens"],
+                "completion_tokens": metadata["output_tokens"],
+                "total_tokens": metadata["total_tokens"]
+            }
+
+
+    
+    # Handle ChatResponse format with raw usuage
+    if hasattr(result, "raw") and hasattr(result.raw, "usage"):
+        usage = result.raw.usage
         return {
-            "prompt_tokens": getattr(metadata, "prompt_token_count", 0),
-            "completion_tokens": getattr(metadata, "candidates_token_count", 0),
-            "total_tokens": getattr(metadata, "total_token_count", 0)
+            "prompt_tokens": getattr(usage, "prompt_tokens", 0),
+            "completion_tokens": getattr(usage, "completion_tokens", 0),
+            "total_tokens": getattr(usage, "total_tokens", 0)
         }
     
     # Handle ChatResult format with generations
@@ -173,24 +226,129 @@ def extract_token_usage(result):
         "total_tokens": 0
     }
 
+def num_tokens_from_messages(model="gpt-4o-mini-2024-07-18", prompt_messages=None, response_message=None):
+    """Calculate the number of tokens used by messages.
+    
+    Args:
+        messages: Optional list of messages (deprecated, use prompt_messages and response_message instead)
+        model: The model name to use for token calculation
+        prompt_messages: List of prompt messages
+        response_message: Response message from the assistant
+    
+    Returns:
+        dict: A dictionary containing:
+            - prompt_tokens: Number of tokens in the prompt
+            - completion_tokens: Number of tokens in the completion
+            - total_tokens: Total number of tokens
+    """
+    #import pdb; pdb.set_trace()
+    try:
+        encoding = tiktoken.encoding_for_model(model)
+    except KeyError:
+        logging.warning("Warning: model not found. Using o200k_base encoding.")
+        encoding = tiktoken.get_encoding("o200k_base")
+    
+    if model in {
+        "gpt-3.5-turbo-0125",
+        "gpt-4-0314",
+        "gpt-4-32k-0314",
+        "gpt-4-0613",
+        "gpt-4-32k-0613",
+        "gpt-4o-2024-08-06",
+        "gpt-4o-mini-2024-07-18"
+        }:
+        tokens_per_message = 3
+        tokens_per_name = 1
+    elif "gpt-3.5-turbo" in model:
+        logging.warning("Warning: gpt-3.5-turbo may update over time. Returning num tokens assuming gpt-3.5-turbo-0125.")
+        return num_tokens_from_messages(model="gpt-3.5-turbo-0125", 
+                                     prompt_messages=prompt_messages, response_message=response_message)
+    elif "gpt-4o-mini" in model:
+        logging.warning("Warning: gpt-4o-mini may update over time. Returning num tokens assuming gpt-4o-mini-2024-07-18.")
+        return num_tokens_from_messages(model="gpt-4o-mini-2024-07-18",
+                                     prompt_messages=prompt_messages, response_message=response_message)
+    elif "gpt-4o" in model:
+        logging.warning("Warning: gpt-4o and gpt-4o-mini may update over time. Returning num tokens assuming gpt-4o-2024-08-06.")
+        return num_tokens_from_messages(model="gpt-4o-2024-08-06",
+                                     prompt_messages=prompt_messages, response_message=response_message)
+    elif "gpt-4" in model:
+        logging.warning("Warning: gpt-4 may update over time. Returning num tokens assuming gpt-4-0613.")
+        return num_tokens_from_messages(model="gpt-4-0613",
+                                     prompt_messages=prompt_messages, response_message=response_message)
+    else:
+        raise NotImplementedError(
+            f"""num_tokens_from_messages() is not implemented for model {model}."""
+        )
+    
+    all_messages = []
+    if prompt_messages:
+        all_messages.extend(prompt_messages)
+    if response_message:
+        if isinstance(response_message, dict):
+            all_messages.append(response_message)
+        else:
+            all_messages.append({"role": "assistant", "content": response_message})
+    
+    prompt_tokens = 0
+    completion_tokens = 0
+    
+    for message in all_messages:
+        num_tokens = tokens_per_message
+        for key, value in message.items():
+            token_count = len(encoding.encode(str(value)))  # Convert value to string for safety
+            num_tokens += token_count
+            if key == "name":
+                num_tokens += tokens_per_name
+        
+        # Add tokens to prompt or completion based on role
+        if message.get("role") == "assistant":
+            completion_tokens += num_tokens
+        else:
+            prompt_tokens += num_tokens
+    
+    # Add the assistant message prefix tokens to completion tokens if we have a response
+    if completion_tokens > 0:
+        completion_tokens += 3  # <|start|>assistant<|message|>
+    
+    total_tokens = prompt_tokens + completion_tokens
+    
+    return {
+        "prompt_tokens": prompt_tokens,
+        "completion_tokens": completion_tokens,
+        "total_tokens": total_tokens
+    }
 
 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()},
     }
 
 
-def calculate_llm_cost(token_usage, model_name, model_costs):
+def calculate_llm_cost(token_usage, model_name, model_costs, model_custom_cost=None):
     """Calculate cost based on token usage and model"""
+    if model_custom_cost is None:
+        model_custom_cost = {}
+    model_costs.update(model_custom_cost)
     if not isinstance(token_usage, dict):
         token_usage = {
             "prompt_tokens": 0,
             "completion_tokens": 0,
             "total_tokens": token_usage if isinstance(token_usage, (int, float)) else 0
         }
-
+    
     # Get model costs, defaulting to default costs if unknown
     model_cost = model_cost = model_costs.get(model_name, {
         "input_cost_per_token": 0.0,   
@@ -277,6 +435,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"):
@@ -424,7 +589,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)