ragaai-catalyst 2.2.2__py3-none-any.whl → 2.2.3b0__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/upload/trace_uploader.py

@@ -2,17 +2,23 @@
 trace_uploader.py - A dedicated process for handling trace uploads
 """
 
-import argparse
-import atexit
-import concurrent.futures
+import os
+import sys
 import json
+import time
+import signal
 import logging
-import os
+import argparse
 import tempfile
-import time
+from pathlib import Path
+import multiprocessing
+import queue
 from datetime import datetime
+import atexit
+import glob
 from logging.handlers import RotatingFileHandler
-from typing import Any, Dict
+import concurrent.futures
+from typing import Dict, Any, Optional
 
 # Set up logging
 log_dir = os.path.join(tempfile.gettempdir(), "ragaai_logs")
@@ -37,16 +43,11 @@ logging.basicConfig(
 logger = logging.getLogger("trace_uploader")
 
 try:
-    from ragaai_catalyst import RagaAICatalyst
-    from ragaai_catalyst.tracers.agentic_tracing.upload.upload_agentic_traces import (
-        UploadAgenticTraces,
-    )
+    from ragaai_catalyst.tracers.agentic_tracing.upload.upload_agentic_traces import UploadAgenticTraces
     from ragaai_catalyst.tracers.agentic_tracing.upload.upload_code import upload_code
-
     # from ragaai_catalyst.tracers.agentic_tracing.upload.upload_trace_metric import upload_trace_metric
-    from ragaai_catalyst.tracers.agentic_tracing.utils.create_dataset_schema import (
-        create_dataset_schema_with_trace,
-    )
+    from ragaai_catalyst.tracers.agentic_tracing.utils.create_dataset_schema import create_dataset_schema_with_trace
+    from ragaai_catalyst import RagaAICatalyst
     IMPORTS_AVAILABLE = True
 except ImportError:
     logger.warning("RagaAI Catalyst imports not available - running in test mode")

ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py

@@ -1,6 +1,6 @@
 import asyncio
 
-#from litellm import model_cost
+# from litellm import model_cost
 import json
 import logging
 import os
@@ -12,24 +12,26 @@ from ..data.data_structure import LLMCall
 
 logger = logging.getLogger(__name__)
 
+
 def get_model_cost():
-    """Load model costs from a JSON file. 
+    """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={}
+    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 
+        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"""
     # First try direct model parameter
     model = kwargs.get("model", "")
-    
     if not model:
         # Try to get from instance
         instance = kwargs.get("self", None)
@@ -45,24 +47,23 @@ def extract_model_name(args, kwargs, result):
     if not model:
         manager = kwargs.get("run_manager", None)
         if manager:
-            if hasattr(manager, 'metadata'):
+            if hasattr(manager, "metadata"):
                 metadata = manager.metadata
-                model_name = metadata.get('ls_model_name', None)
+                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):
+        if "to_dict" in dir(result):
             result = result.to_dict()
-            if 'model_version' in result:
-                model = result['model_version']  
+            if "model_version" in result:
+                model = result["model_version"]
     try:
         if not model:
             model = result.raw.model
     except Exception:
         pass
-    
-    
+
     # Normalize Google model names
     if model and isinstance(model, str):
         model = model.lower()
@@ -73,10 +74,10 @@ def extract_model_name(args, kwargs, result):
         if "gemini-pro" in model:
             return "gemini-pro"
 
-    if 'response_metadata' in dir(result):
-        if 'model_name' in result.response_metadata:
-            model = result.response_metadata['model_name']
-    
+    if "response_metadata" in dir(result):
+        if "model_name" in result.response_metadata:
+            model = result.response_metadata["model_name"]
+
     return model or "default"
 
 
@@ -85,27 +86,29 @@ def extract_parameters(kwargs):
     parameters = {k: v for k, v in kwargs.items() if v is not None}
 
     # Remove contents key in parameters (Google LLM Response)
-    if 'contents' in parameters:
-        del parameters['contents']
+    if "contents" in parameters:
+        del parameters["contents"]
 
     # 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']
+    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"]
         # If generation_config is already a dict, use it directly
         if isinstance(generation_config, dict):
             config_dict = generation_config
         else:
             # Convert GenerationConfig to dictionary if it has a to_dict method, otherwise try to get its __dict__
-            config_dict = getattr(generation_config, 'to_dict', lambda: generation_config.__dict__)()
+            config_dict = getattr(
+                generation_config, "to_dict", lambda: generation_config.__dict__
+            )()
         parameters.update(config_dict)
-        del parameters['generation_config']
-        
+        del parameters["generation_config"]
+
     return parameters
 
 
@@ -123,13 +126,14 @@ def extract_token_usage(result):
         # First try parsing as JSON for OpenAI responses
         try:
             import json
+
             json_data = json.loads(result.text)
             if isinstance(json_data, dict) and "usage" in json_data:
                 usage = json_data["usage"]
                 return {
                     "prompt_tokens": usage.get("prompt_tokens", 0),
                     "completion_tokens": usage.get("completion_tokens", 0),
-                    "total_tokens": usage.get("total_tokens", 0)
+                    "total_tokens": usage.get("total_tokens", 0),
                 }
         except (json.JSONDecodeError, AttributeError):
             pass
@@ -142,7 +146,7 @@ def extract_token_usage(result):
             return {
                 "prompt_tokens": 0,  # Vertex AI doesn't provide this breakdown
                 "completion_tokens": total_tokens,
-                "total_tokens": total_tokens
+                "total_tokens": total_tokens,
             }
 
     # Handle Claude 3 message format
@@ -152,15 +156,15 @@ def extract_token_usage(result):
             return {
                 "prompt_tokens": usage.input_tokens,
                 "completion_tokens": usage.output_tokens,
-                "total_tokens": usage.input_tokens + usage.output_tokens
+                "total_tokens": usage.input_tokens + usage.output_tokens,
             }
         # Handle standard OpenAI/Anthropic format
         return {
             "prompt_tokens": getattr(usage, "prompt_tokens", 0),
             "completion_tokens": getattr(usage, "completion_tokens", 0),
-            "total_tokens": getattr(usage, "total_tokens", 0)
+            "total_tokens": getattr(usage, "total_tokens", 0),
         }
-    
+
     # Handle Google GenerativeAI format with usage_metadata
     if hasattr(result, "usage_metadata"):
         metadata = result.usage_metadata
@@ -168,37 +172,35 @@ def extract_token_usage(result):
             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)
+                "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)
+                "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"]
+                "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(usage, "prompt_tokens", 0),
             "completion_tokens": getattr(usage, "completion_tokens", 0),
-            "total_tokens": getattr(usage, "total_tokens", 0)
+            "total_tokens": getattr(usage, "total_tokens", 0),
         }
-    
+
     # Handle ChatResult format with generations
     if hasattr(result, "generations") and result.generations:
         # Get the first generation
         generation = result.generations[0]
-        
+
         # Try to get usage from generation_info
         if hasattr(generation, "generation_info"):
             metadata = generation.generation_info.get("usage_metadata", {})
@@ -206,46 +208,47 @@ def extract_token_usage(result):
                 return {
                     "prompt_tokens": metadata.get("prompt_token_count", 0),
                     "completion_tokens": metadata.get("candidates_token_count", 0),
-                    "total_tokens": metadata.get("total_token_count", 0)
+                    "total_tokens": metadata.get("total_token_count", 0),
                 }
-        
+
         # Try to get usage from message's usage_metadata
-        if hasattr(generation, "message") and hasattr(generation.message, "usage_metadata"):
+        if hasattr(generation, "message") and hasattr(
+            generation.message, "usage_metadata"
+        ):
             metadata = generation.message.usage_metadata
             return {
                 "prompt_tokens": metadata.get("input_tokens", 0),
                 "completion_tokens": metadata.get("output_tokens", 0),
-                "total_tokens": metadata.get("total_tokens", 0)
+                "total_tokens": metadata.get("total_tokens", 0),
             }
-    
-    return {
-        "prompt_tokens": 0,
-        "completion_tokens": 0,
-        "total_tokens": 0
-    }
 
-def num_tokens_from_messages(model="gpt-4o-mini-2024-07-18", prompt_messages=None, response_message=None):
+    return {"prompt_tokens": 0, "completion_tokens": 0, "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()
+    # 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",
@@ -253,31 +256,51 @@ def num_tokens_from_messages(model="gpt-4o-mini-2024-07-18", prompt_messages=Non
         "gpt-4-0613",
         "gpt-4-32k-0613",
         "gpt-4o-2024-08-06",
-        "gpt-4o-mini-2024-07-18"
-        }:
+        "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)
+        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)
+        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)
+        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)
+        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)
@@ -286,36 +309,39 @@ def num_tokens_from_messages(model="gpt-4o-mini-2024-07-18", prompt_messages=Non
             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
+            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
+        "total_tokens": total_tokens,
     }
 
+
 def extract_input_data(args, kwargs, result):
     """Sanitize and format input data, including handling of nested lists and dictionaries."""
 
@@ -336,6 +362,7 @@ def extract_input_data(args, kwargs, result):
 
 
 def calculate_llm_cost(token_usage, model_name, model_costs, model_custom_cost=None):
+    model_name = extract_model_name([], {"model": model_name}, None)
     """Calculate cost based on token usage and model"""
     if model_custom_cost is None:
         model_custom_cost = {}
@@ -344,40 +371,51 @@ def calculate_llm_cost(token_usage, model_name, model_costs, model_custom_cost=N
         token_usage = {
             "prompt_tokens": 0,
             "completion_tokens": 0,
-            "total_tokens": token_usage if isinstance(token_usage, (int, float)) else 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,   
-        "output_cost_per_token": 0.0   
-    })
-    if model_cost['input_cost_per_token'] == 0.0 and model_cost['output_cost_per_token'] == 0.0:
-        provide_name = model_name.split('-')[0]
-        if provide_name == 'azure':
-            model_name = os.path.join('azure', '-'.join(model_name.split('-')[1:]))
-
-            model_cost = model_costs.get(model_name, {
-                "input_cost_per_token": 0.0,   
-                "output_cost_per_token": 0.0   
-            })
-
-    input_cost = (token_usage.get("prompt_tokens", 0)) * model_cost.get("input_cost_per_token", 0.0)
-    output_cost = (token_usage.get("completion_tokens", 0)) * model_cost.get("output_cost_per_token", 0.0)
+    model_cost = model_cost = model_costs.get(
+        model_name, {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}
+    )
+    if (
+        model_cost["input_cost_per_token"] == 0.0
+        and model_cost["output_cost_per_token"] == 0.0
+    ):
+        provide_name = model_name.split("-")[0]
+        if provide_name == "azure":
+            model_name = os.path.join("azure", "-".join(model_name.split("-")[1:]))
+
+            model_cost = model_costs.get(
+                model_name, {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}
+            )
+
+    input_cost = (token_usage.get("prompt_tokens", 0)) * model_cost.get(
+        "input_cost_per_token", 0.0
+    )
+    output_cost = (token_usage.get("completion_tokens", 0)) * model_cost.get(
+        "output_cost_per_token", 0.0
+    )
     total_cost = input_cost + output_cost
 
     return {
         "input_cost": round(input_cost, 10),
         "output_cost": round(output_cost, 10),
-        "total_cost": round(total_cost, 10)
+        "total_cost": round(total_cost, 10),
     }
 
 
 def sanitize_api_keys(data):
     """Remove sensitive information from data"""
     if isinstance(data, dict):
-        return {k: sanitize_api_keys(v) for k, v in data.items() 
-                if not any(sensitive in k.lower() for sensitive in ['key', 'token', 'secret', 'password'])}
+        return {
+            k: sanitize_api_keys(v)
+            for k, v in data.items()
+            if not any(
+                sensitive in k.lower()
+                for sensitive in ["key", "token", "secret", "password"]
+            )
+        }
     elif isinstance(data, list):
         return [sanitize_api_keys(item) for item in data]
     elif isinstance(data, tuple):
@@ -387,10 +425,10 @@ def sanitize_api_keys(data):
 
 def sanitize_input(args, kwargs):
     """Convert input arguments to text format.
-    
+
     Args:
         args: Input arguments that may contain nested dictionaries
-        
+
     Returns:
         str: Text representation of the input arguments
     """
@@ -403,6 +441,7 @@ def sanitize_input(args, kwargs):
 
 def extract_llm_output(result):
     """Extract output from LLM response"""
+
     class OutputResponse:
         def __init__(self, output_response):
             self.output_response = output_response
@@ -415,7 +454,9 @@ def extract_llm_output(result):
         else:
             # We're in an async context, but this function is called synchronously
             # Return a placeholder and let the caller handle the coroutine
-            return OutputResponse([{'content': "Coroutine result pending", "role": "assistant"}])
+            return OutputResponse(
+                [{"content": "Coroutine result pending", "role": "assistant"}]
+            )
 
     # Handle Google GenerativeAI format
     if hasattr(result, "result"):
@@ -426,56 +467,52 @@ def extract_llm_output(result):
             if content and hasattr(content, "parts"):
                 for part in content.parts:
                     if hasattr(part, "text"):
-                        output.append({
-                            "content": part.text,
-                            "role": getattr(content, "role", "assistant"),
-                            "finish_reason": getattr(candidate, "finish_reason", None)
-                        })
+                        output.append(
+                            {
+                                "content": part.text,
+                                "role": getattr(content, "role", "assistant"),
+                                "finish_reason": getattr(
+                                    candidate, "finish_reason", None
+                                ),
+                            }
+                        )
         return OutputResponse(output)
-    
+
     # Handle AIMessage Format
     if hasattr(result, "content"):
-        return OutputResponse([{
-            "content": result.content,
-            "role": getattr(result, "role", "assistant")
-        }])
-    
+        return OutputResponse(
+            [{"content": result.content, "role": getattr(result, "role", "assistant")}]
+        )
+
     # Handle Vertex AI format
     # format1
     if hasattr(result, "text"):
-        return OutputResponse([{
-            "content": result.text,
-            "role": "assistant"
-        }])
-
+        return OutputResponse([{"content": result.text, "role": "assistant"}])
 
     # format2
     if hasattr(result, "generations"):
         output = []
         for generation in result.generations:
-            output.append({
-                "content": generation.text,
-                "role": "assistant"
-            })
+            output.append({"content": generation.text, "role": "assistant"})
         return OutputResponse(output)
-    
+
     # Handle OpenAI format
     if hasattr(result, "choices"):
-        return OutputResponse([{
-            "content": choice.message.content,
-            "role": choice.message.role
-        } for choice in result.choices])
-
+        return OutputResponse(
+            [
+                {"content": choice.message.content, "role": choice.message.role}
+                for choice in result.choices
+            ]
+        )
 
     # Handle Anthropic format
     if hasattr(result, "content"):
-        return OutputResponse([{
-            "content": result.content[0].text,
-            "role": "assistant"
-        }])
-    
+        return OutputResponse(
+            [{"content": result.content[0].text, "role": "assistant"}]
+        )
+
     # Default case
-    return OutputResponse([{'content': result, 'role': 'assistant'}])
+    return OutputResponse([{"content": result, "role": "assistant"}])
 
 
 def extract_llm_data(args, kwargs, result):
@@ -608,12 +645,12 @@ def count_tokens(input_str: str) -> int:
     # Use tiktoken to count tokens
     try:
         import tiktoken
-        
+
         # Use GPT-4o model's encoding (cl100k_base)
         encoding = tiktoken.get_encoding("cl100k_base")
-        
+
         # Count tokens
         tokens = encoding.encode(input_str)
         return len(tokens)
     except Exception:
-        raise Exception("Failed to count tokens")
+        raise Exception("Failed to count tokens")