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

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)