ragaai-catalyst 2.1.5b7__py3-none-any.whl → 2.1.5b9__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -150,8 +150,6 @@ class LLMTracerMixin:
         beta_module = openai_module.beta
 
         # Patch openai.beta.threads
-        import openai
-        openai.api_type = "openai"
         if hasattr(beta_module, "threads"):
             threads_obj = beta_module.threads
             # Patch top-level methods on openai.beta.threads

ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json

@@ -14,7 +14,8 @@
         "supports_audio_output": true,
         "supports_prompt_caching": true,
         "supports_response_schema": true,
-        "supports_system_messages": true
+        "supports_system_messages": true,
+        "deprecation_date": "date when the model becomes deprecated in the format YYYY-MM-DD"
     },
     "omni-moderation-latest": {
         "max_tokens": 32768,
@@ -441,7 +442,8 @@
         "mode": "chat",
         "supports_function_calling": true,
         "supports_prompt_caching": true,
-        "supports_system_messages": true
+        "supports_system_messages": true,
+        "deprecation_date": "2025-06-06"
     },
     "gpt-4-32k": {
         "max_tokens": 4096,
@@ -540,7 +542,8 @@
         "mode": "chat",
         "supports_vision": true,
         "supports_prompt_caching": true,
-        "supports_system_messages": true
+        "supports_system_messages": true,
+        "deprecation_date": "2024-12-06"
     },
     "gpt-4-1106-vision-preview": {
         "max_tokens": 4096,
@@ -552,7 +555,8 @@
         "mode": "chat",
         "supports_vision": true,
         "supports_prompt_caching": true,
-        "supports_system_messages": true
+        "supports_system_messages": true,
+        "deprecation_date": "2024-12-06"
     },
     "gpt-3.5-turbo": {
         "max_tokens": 4097,
@@ -1223,7 +1227,8 @@
         "litellm_provider": "azure",
         "mode": "chat",
         "supports_function_calling": true,
-        "supports_parallel_function_calling": true
+        "supports_parallel_function_calling": true,
+        "deprecation_date": "2025-03-31"
     },
     "azure/gpt-35-turbo-0613": {
         "max_tokens": 4097,
@@ -1234,7 +1239,8 @@
         "litellm_provider": "azure",
         "mode": "chat",
         "supports_function_calling": true,
-        "supports_parallel_function_calling": true
+        "supports_parallel_function_calling": true,
+        "deprecation_date": "2025-02-13"
     },
     "azure/gpt-35-turbo-0301": {
         "max_tokens": 4097,
@@ -1245,7 +1251,8 @@
         "litellm_provider": "azure",
         "mode": "chat",
         "supports_function_calling": true,
-        "supports_parallel_function_calling": true
+        "supports_parallel_function_calling": true,
+        "deprecation_date": "2025-02-13"
     },
     "azure/gpt-35-turbo-0125": {
         "max_tokens": 4096,
@@ -1256,7 +1263,8 @@
         "litellm_provider": "azure",
         "mode": "chat",
         "supports_function_calling": true,
-        "supports_parallel_function_calling": true
+        "supports_parallel_function_calling": true,
+        "deprecation_date": "2025-03-31"
     },
     "azure/gpt-35-turbo-16k": {
         "max_tokens": 4096,
@@ -2042,6 +2050,84 @@
         "supports_function_calling": true,
         "supports_vision": true
     },
+    "xai/grok-2-vision-1212": {
+        "max_tokens": 32768,
+        "max_input_tokens": 32768,
+        "max_output_tokens": 32768,
+        "input_cost_per_token": 2e-06,
+        "input_cost_per_image": 2e-06,
+        "output_cost_per_token": 1e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_vision": true
+    },
+    "xai/grok-2-vision-latest": {
+        "max_tokens": 32768,
+        "max_input_tokens": 32768,
+        "max_output_tokens": 32768,
+        "input_cost_per_token": 2e-06,
+        "input_cost_per_image": 2e-06,
+        "output_cost_per_token": 1e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_vision": true
+    },
+    "xai/grok-2-vision": {
+        "max_tokens": 32768,
+        "max_input_tokens": 32768,
+        "max_output_tokens": 32768,
+        "input_cost_per_token": 2e-06,
+        "input_cost_per_image": 2e-06,
+        "output_cost_per_token": 1e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_vision": true
+    },
+    "xai/grok-vision-beta": {
+        "max_tokens": 8192,
+        "max_input_tokens": 8192,
+        "max_output_tokens": 8192,
+        "input_cost_per_token": 5e-06,
+        "input_cost_per_image": 5e-06,
+        "output_cost_per_token": 1.5e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true,
+        "supports_vision": true
+    },
+    "xai/grok-2-1212": {
+        "max_tokens": 131072,
+        "max_input_tokens": 131072,
+        "max_output_tokens": 131072,
+        "input_cost_per_token": 2e-06,
+        "output_cost_per_token": 1e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
+    "xai/grok-2": {
+        "max_tokens": 131072,
+        "max_input_tokens": 131072,
+        "max_output_tokens": 131072,
+        "input_cost_per_token": 2e-06,
+        "output_cost_per_token": 1e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
+    "xai/grok-2-latest": {
+        "max_tokens": 131072,
+        "max_input_tokens": 131072,
+        "max_output_tokens": 131072,
+        "input_cost_per_token": 2e-06,
+        "output_cost_per_token": 1e-05,
+        "litellm_provider": "xai",
+        "mode": "chat",
+        "supports_function_calling": true
+    },
     "deepseek/deepseek-coder": {
         "max_tokens": 4096,
         "max_input_tokens": 128000,
@@ -2353,7 +2439,8 @@
         "tool_use_system_prompt_tokens": 264,
         "supports_assistant_prefill": true,
         "supports_prompt_caching": true,
-        "supports_response_schema": true
+        "supports_response_schema": true,
+        "deprecation_date": "2025-03-01"
     },
     "claude-3-5-haiku-20241022": {
         "max_tokens": 8192,
@@ -2369,7 +2456,8 @@
         "tool_use_system_prompt_tokens": 264,
         "supports_assistant_prefill": true,
         "supports_prompt_caching": true,
-        "supports_response_schema": true
+        "supports_response_schema": true,
+        "deprecation_date": "2025-10-01"
     },
     "claude-3-opus-20240229": {
         "max_tokens": 4096,
@@ -2386,7 +2474,8 @@
         "tool_use_system_prompt_tokens": 395,
         "supports_assistant_prefill": true,
         "supports_prompt_caching": true,
-        "supports_response_schema": true
+        "supports_response_schema": true,
+        "deprecation_date": "2025-03-01"
     },
     "claude-3-sonnet-20240229": {
         "max_tokens": 4096,
@@ -2401,7 +2490,8 @@
         "tool_use_system_prompt_tokens": 159,
         "supports_assistant_prefill": true,
         "supports_prompt_caching": true,
-        "supports_response_schema": true
+        "supports_response_schema": true,
+        "deprecation_date": "2025-07-21"
     },
     "claude-3-5-sonnet-20240620": {
         "max_tokens": 8192,
@@ -2418,7 +2508,8 @@
         "tool_use_system_prompt_tokens": 159,
         "supports_assistant_prefill": true,
         "supports_prompt_caching": true,
-        "supports_response_schema": true
+        "supports_response_schema": true,
+        "deprecation_date": "2025-06-01"
     },
     "claude-3-5-sonnet-20241022": {
         "max_tokens": 8192,
@@ -2436,7 +2527,8 @@
         "supports_assistant_prefill": true,
         "supports_pdf_input": true,
         "supports_prompt_caching": true,
-        "supports_response_schema": true
+        "supports_response_schema": true,
+        "deprecation_date": "2025-10-01"
     },
     "text-bison": {
         "max_tokens": 2048,
@@ -2546,7 +2638,8 @@
         "output_cost_per_character": 5e-07,
         "litellm_provider": "vertex_ai-chat-models",
         "mode": "chat",
-        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
+        "deprecation_date": "2025-04-09"
     },
     "chat-bison-32k": {
         "max_tokens": 8192,
@@ -2787,7 +2880,8 @@
         "litellm_provider": "vertex_ai-language-models",
         "mode": "chat",
         "supports_function_calling": true,
-        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
+        "deprecation_date": "2025-04-09"
     },
     "gemini-1.0-ultra": {
         "max_tokens": 8192,
@@ -2832,7 +2926,8 @@
         "litellm_provider": "vertex_ai-language-models",
         "mode": "chat",
         "supports_function_calling": true,
-        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
+        "deprecation_date": "2025-04-09"
     },
     "gemini-1.5-pro": {
         "max_tokens": 8192,
@@ -2914,7 +3009,8 @@
         "supports_function_calling": true,
         "supports_tool_choice": true,
         "supports_response_schema": true,
-        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
+        "deprecation_date": "2025-05-24"
     },
     "gemini-1.5-pro-preview-0514": {
         "max_tokens": 8192,
@@ -3119,7 +3215,8 @@
         "supports_function_calling": true,
         "supports_vision": true,
         "supports_response_schema": true,
-        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
+        "deprecation_date": "2025-05-24"
     },
     "gemini-1.5-flash-preview-0514": {
         "max_tokens": 8192,
@@ -3223,7 +3320,8 @@
         "mode": "chat",
         "supports_function_calling": true,
         "supports_vision": true,
-        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models"
+        "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models",
+        "deprecation_date": "2025-04-09"
     },
     "medlm-medium": {
         "max_tokens": 8192,
@@ -3929,7 +4027,8 @@
         "supports_prompt_caching": true,
         "tpm": 4000000,
         "rpm": 2000,
-        "source": "https://ai.google.dev/pricing"
+        "source": "https://ai.google.dev/pricing",
+        "deprecation_date": "2025-05-24"
     },
     "gemini/gemini-1.5-flash": {
         "max_tokens": 8192,
@@ -4205,7 +4304,8 @@
         "supports_prompt_caching": true,
         "tpm": 4000000,
         "rpm": 1000,
-        "source": "https://ai.google.dev/pricing"
+        "source": "https://ai.google.dev/pricing",
+        "deprecation_date": "2025-05-24"
     },
     "gemini/gemini-1.5-pro-exp-0801": {
         "max_tokens": 8192,
@@ -7341,7 +7441,8 @@
         "input_cost_per_token": 5e-06,
         "output_cost_per_token": 5e-06,
         "litellm_provider": "perplexity",
-        "mode": "chat"
+        "mode": "chat",
+        "deprecation_date": "2025-02-22"
     },
     "perplexity/llama-3.1-sonar-large-128k-online": {
         "max_tokens": 127072,
@@ -7350,7 +7451,8 @@
         "input_cost_per_token": 1e-06,
         "output_cost_per_token": 1e-06,
         "litellm_provider": "perplexity",
-        "mode": "chat"
+        "mode": "chat",
+        "deprecation_date": "2025-02-22"
     },
     "perplexity/llama-3.1-sonar-large-128k-chat": {
         "max_tokens": 131072,
@@ -7359,7 +7461,8 @@
         "input_cost_per_token": 1e-06,
         "output_cost_per_token": 1e-06,
         "litellm_provider": "perplexity",
-        "mode": "chat"
+        "mode": "chat",
+        "deprecation_date": "2025-02-22"
     },
     "perplexity/llama-3.1-sonar-small-128k-chat": {
         "max_tokens": 131072,
@@ -7368,7 +7471,8 @@
         "input_cost_per_token": 2e-07,
         "output_cost_per_token": 2e-07,
         "litellm_provider": "perplexity",
-        "mode": "chat"
+        "mode": "chat",
+        "deprecation_date": "2025-02-22"
     },
     "perplexity/llama-3.1-sonar-small-128k-online": {
         "max_tokens": 127072,
@@ -7377,7 +7481,8 @@
         "input_cost_per_token": 2e-07,
         "output_cost_per_token": 2e-07,
         "litellm_provider": "perplexity",
-        "mode": "chat"
+        "mode": "chat",
+        "deprecation_date": "2025-02-22"
     },
     "perplexity/pplx-7b-chat": {
         "max_tokens": 8192,

ragaai_catalyst/tracers/langchain_callback.py

@@ -48,8 +48,9 @@ class LangchainTracer(BaseCallbackHandler):
         self._original_methods = {}
         self.additional_metadata = {}
         self._save_task = None
-        self._current_query = None  # Add this line to track the current query
+        self._current_query = None
         self.filepath = None
+        self.model_names = {}  # Store model names by component instance
         logger.setLevel(log_level)
 
         if not os.path.exists(output_path):
@@ -86,7 +87,7 @@ class LangchainTracer(BaseCallbackHandler):
             "retriever_actions": [],
             "tokens": [],
             "errors": [],
-            "query": self._current_query,  # Add this line to include the query in the trace
+            "query": self._current_query,
             "metadata": {
                 "version": "2.0",
                 "trace_all": self.trace_all,
@@ -172,6 +173,15 @@ class LangchainTracer(BaseCallbackHandler):
                     kwargs_copy['callbacks'] = [self]
                 elif self not in kwargs_copy['callbacks']:
                     kwargs_copy['callbacks'].append(self)
+
+                # Store model name if available
+                if component_name in ["OpenAI", "ChatOpenAI_LangchainOpenAI", "ChatOpenAI_ChatModels",
+                                    "ChatVertexAI", "VertexAI", "ChatGoogleGenerativeAI", "ChatAnthropic", 
+                                    "ChatLiteLLM"]:
+                    instance = args[0] if args else None
+                    model_name = kwargs.get('model_name') or kwargs.get('model')
+                    if instance and model_name:
+                        self.model_names[id(instance)] = model_name
                 
                 # Try different method signatures
                 try:
@@ -201,28 +211,57 @@ class LangchainTracer(BaseCallbackHandler):
 
     def _monkey_patch(self):
         """Enhanced monkey-patching with comprehensive component support"""
-        from langchain.llms import OpenAI
-        # from langchain_groq import ChatGroq
-        # from langchain_google_genai import ChatGoogleGenerativeAI
-        # from langchain_anthropic import ChatAnthropic
-        from langchain_community.chat_models import ChatLiteLLM
-        # from langchain_cohere import ChatCohere
-        from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
-        from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
-        from langchain.chains import create_retrieval_chain, RetrievalQA
-
-        components_to_patch = {
-            "OpenAI": (OpenAI, "__init__"),
-            # "ChatGroq": (ChatGroq, "__init__"),
-            # "ChatGoogleGenerativeAI": (ChatGoogleGenerativeAI, "__init__"),
-            # "ChatAnthropic": (ChatAnthropic, "__init__"),
-            "ChatLiteLLM": (ChatLiteLLM, "__init__"),
-            # "ChatCohere": (ChatCohere, "__init__"),
-            "ChatOpenAI_LangchainOpenAI": (ChatOpenAI_LangchainOpenAI, "__init__"),
-            "ChatOpenAI_ChatModels": (ChatOpenAI_ChatModels, "__init__"),
-            "RetrievalQA": (RetrievalQA, "from_chain_type"),
-            "create_retrieval_chain": (create_retrieval_chain, None),
-        }
+        components_to_patch = {}
+        
+        try:
+            from langchain.llms import OpenAI
+            components_to_patch["OpenAI"] = (OpenAI, "__init__")
+        except ImportError:
+            logger.debug("OpenAI not available for patching")
+            
+        try:
+            from langchain_google_vertexai import ChatVertexAI, VertexAI
+            components_to_patch["ChatVertexAI"] = (ChatVertexAI, "__init__")
+            components_to_patch["VertexAI"] = (VertexAI, "__init__")
+        except ImportError:
+            logger.debug("ChatVertexAI/VertexAI not available for patching")
+            
+        try:
+            from langchain_google_genai import ChatGoogleGenerativeAI
+            components_to_patch["ChatGoogleGenerativeAI"] = (ChatGoogleGenerativeAI, "__init__")
+        except ImportError:
+            logger.debug("ChatGoogleGenerativeAI not available for patching")
+            
+        try:
+            from langchain_anthropic import ChatAnthropic
+            components_to_patch["ChatAnthropic"] = (ChatAnthropic, "__init__")
+        except ImportError:
+            logger.debug("ChatAnthropic not available for patching")
+            
+        try:
+            from langchain_community.chat_models import ChatLiteLLM
+            components_to_patch["ChatLiteLLM"] = (ChatLiteLLM, "__init__")
+        except ImportError:
+            logger.debug("ChatLiteLLM not available for patching")
+            
+        try:
+            from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
+            components_to_patch["ChatOpenAI_LangchainOpenAI"] = (ChatOpenAI_LangchainOpenAI, "__init__")
+        except ImportError:
+            logger.debug("ChatOpenAI_LangchainOpenAI not available for patching")
+            
+        try:
+            from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
+            components_to_patch["ChatOpenAI_ChatModels"] = (ChatOpenAI_ChatModels, "__init__")
+        except ImportError:
+            logger.debug("ChatOpenAI_ChatModels not available for patching")
+            
+        try:
+            from langchain.chains import create_retrieval_chain, RetrievalQA
+            components_to_patch["RetrievalQA"] = (RetrievalQA, "from_chain_type")
+            components_to_patch["create_retrieval_chain"] = (create_retrieval_chain, None)
+        except ImportError:
+            logger.debug("Langchain chains not available for patching")
 
         for name, (component, method_name) in components_to_patch.items():
             try:
@@ -249,21 +288,48 @@ class LangchainTracer(BaseCallbackHandler):
 
     def _restore_original_methods(self):
         """Restore all original methods and functions with enhanced error handling"""
-        from langchain.llms import OpenAI
-        # from langchain_groq import ChatGroq
-        # from langchain_google_genai import ChatGoogleGenerativeAI
-        # from langchain_anthropic import ChatAnthropic
-        from langchain_community.chat_models import ChatLiteLLM
-        # from langchain_cohere import ChatCohere
-        from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
-        from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
-        from langchain.chains import create_retrieval_chain, RetrievalQA
-
+        # Dynamically import only what we need based on what was patched
+        imported_components = {}
+        
+        if self._original_inits or self._original_methods:
+            for name in list(self._original_inits.keys()) + list(self._original_methods.keys()):
+                try:
+                    if name == "OpenAI":
+                        from langchain.llms import OpenAI
+                        imported_components[name] = OpenAI
+                    elif name == "ChatVertexAI":
+                        from langchain_google_vertexai import ChatVertexAI
+                        imported_components[name] = ChatVertexAI
+                    elif name == "VertexAI":
+                        from langchain_google_vertexai import VertexAI
+                        imported_components[name] = VertexAI
+                    elif name == "ChatGoogleGenerativeAI":
+                        from langchain_google_genai import ChatGoogleGenerativeAI
+                        imported_components[name] = ChatGoogleGenerativeAI
+                    elif name == "ChatAnthropic":
+                        from langchain_anthropic import ChatAnthropic
+                        imported_components[name] = ChatAnthropic
+                    elif name == "ChatLiteLLM":
+                        from langchain_community.chat_models import ChatLiteLLM
+                        imported_components[name] = ChatLiteLLM
+                    elif name == "ChatOpenAI_LangchainOpenAI":
+                        from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
+                        imported_components[name] = ChatOpenAI_LangchainOpenAI
+                    elif name == "ChatOpenAI_ChatModels":
+                        from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
+                        imported_components[name] = ChatOpenAI_ChatModels
+                    elif name in ["RetrievalQA", "create_retrieval_chain"]:
+                        from langchain.chains import create_retrieval_chain, RetrievalQA
+                        imported_components["RetrievalQA"] = RetrievalQA
+                        imported_components["create_retrieval_chain"] = create_retrieval_chain
+                except ImportError:
+                    logger.debug(f"{name} not available for restoration")
 
         for name, original in self._original_inits.items():
             try:
-                component = eval(name)
-                component.__init__ = original
+                if name in imported_components:
+                    component = imported_components[name]
+                    component.__init__ = original
             except Exception as e:
                 logger.error(f"Error restoring {name}: {e}")
                 self.on_error(e, context=f"restore_{name}")
@@ -272,10 +338,12 @@ class LangchainTracer(BaseCallbackHandler):
             try:
                 if "." in name:
                     module_name, method_name = name.rsplit(".", 1)
-                    module = eval(module_name)
-                    setattr(module, method_name, original)
+                    if module_name in imported_components:
+                        module = imported_components[module_name]
+                        setattr(module, method_name, original)
                 else:
-                    globals()[name] = original
+                    if name in imported_components:
+                        globals()[name] = original
             except Exception as e:
                 logger.error(f"Error restoring {name}: {e}")
                 self.on_error(e, context=f"restore_{name}")
@@ -359,16 +427,99 @@ class LangchainTracer(BaseCallbackHandler):
                 }
             )
 
+            # Calculate latency
             end_time = datetime.now()
-            self.additional_metadata["latency"] = (end_time - self.current_trace["start_time"]).total_seconds()
+            latency = (end_time - self.current_trace["start_time"]).total_seconds()
 
+            # Check if values are there in llm_output
+            model = ""
+            prompt_tokens = 0
+            completion_tokens = 0
+            total_tokens = 0
+            
+            # Try to get model name from llm_output first
             if response and response.llm_output:
-                self.additional_metadata["model_name"] = response.llm_output.get("model_name", "")
-                self.additional_metadata["tokens"] = {}
-                if response.llm_output.get("token_usage"):
-                    self.additional_metadata["tokens"]["total"] = response.llm_output["token_usage"].get("total_tokens", 0)
-                    self.additional_metadata["tokens"]["prompt"] = response.llm_output["token_usage"].get("prompt_tokens", 0)
-                    self.additional_metadata["tokens"]["completion"] = response.llm_output["token_usage"].get("completion_tokens", 0)
+                try:
+                    model = response.llm_output.get("model_name")
+                    if not model:
+                        model = response.llm_output.get("model", "")
+                except Exception as e:
+                    # logger.debug(f"Error getting model name: {e}")
+                    model = ""
+
+            # Add model name
+            if not model:
+                try:
+                    model = response.llm_output.get("model_name")
+                    if not model:
+                        model = response.llm_output.get("model", "")
+                except Exception as e:
+                    # logger.debug(f"Error getting model name: {e}")
+                    model = ""
+
+
+            # Add token usage
+            try:
+                token_usage = response.llm_output.get("token_usage", {})
+                if token_usage=={}:
+                    try:
+                        token_usage = response.llm_output.get("usage")
+                    except Exception as e:
+                        # logger.debug(f"Error getting token usage: {e}")
+                        token_usage = {}
+                    
+                if token_usage !={}:
+                    prompt_tokens = token_usage.get("prompt_tokens", 0)
+                    if prompt_tokens==0:
+                        prompt_tokens = token_usage.get("input_tokens", 0)
+                    completion_tokens = token_usage.get("completion_tokens", 0)
+                    if completion_tokens==0:
+                        completion_tokens = token_usage.get("output_tokens", 0)
+
+                    total_tokens = prompt_tokens + completion_tokens
+            except Exception as e:
+                # logger.debug(f"Error getting token usage: {e}")
+                prompt_tokens = 0
+                completion_tokens = 0
+                total_tokens = 0
+
+            # Check if values are there in 
+            if prompt_tokens == 0 and completion_tokens == 0:
+                try:
+                    usage_data = response.generations[0][0].message.usage_metadata
+                    prompt_tokens = usage_data.get("input_tokens", 0)
+                    completion_tokens = usage_data.get("output_tokens", 0)
+                    total_tokens = prompt_tokens + completion_tokens
+                except Exception as e:
+                    # logger.debug(f"Error getting usage data: {e}")
+                    try:
+                        usage_data = response.generations[0][0].generation_info['usage_metadata']
+                        prompt_tokens = usage_data.get("prompt_token_count", 0)
+                        completion_tokens = usage_data.get("candidates_token_count", 0)
+                        total_tokens = prompt_tokens + completion_tokens
+                    except Exception as e:
+                        # logger.debug(f"Error getting token usage: {e}")
+                        prompt_tokens = 0
+                        completion_tokens = 0
+                        total_tokens = 0
+
+            # If no model name in llm_output, try to get it from stored model names
+            try:
+                if model == "":
+                    model = list(self.model_names.values())[0]
+            except Exception as e:
+                model=""
+            
+            self.additional_metadata = {
+                'latency': latency,
+                'model_name': model,
+                'tokens': {
+                    'prompt': prompt_tokens,
+                    'completion': completion_tokens,
+                    'total': total_tokens
+                }
+            }
+
         except Exception as e:
             self.on_error(e, context="llm_end")
 

ragaai_catalyst/tracers/tracer.py

@@ -283,64 +283,48 @@ class Tracer(AgenticTracing):
             data, additional_metadata = self.langchain_tracer.stop()
 
             # Add cost if possible
-            if additional_metadata.get('model_name'):
+            # import pdb; pdb.set_trace()
+            if additional_metadata['model_name']:
                 try:
                     model_cost_data = self.model_cost_dict[additional_metadata['model_name']]
-                    if 'tokens' in additional_metadata and all(k in additional_metadata['tokens'] for k in ['prompt', 'completion']):
-                        prompt_cost = additional_metadata["tokens"]["prompt"]*model_cost_data["input_cost_per_token"]
-                        completion_cost = additional_metadata["tokens"]["completion"]*model_cost_data["output_cost_per_token"]
-                        additional_metadata.setdefault('cost', {})["total_cost"] = prompt_cost + completion_cost 
-                    else:
-                        logger.warning("Token information missing in additional_metadata")
+                    prompt_cost = additional_metadata["tokens"]["prompt"]*model_cost_data["input_cost_per_token"]
+                    completion_cost = additional_metadata["tokens"]["completion"]*model_cost_data["output_cost_per_token"]
+                    # additional_metadata.setdefault('cost', {})["prompt_cost"] = prompt_cost
+                    # additional_metadata.setdefault('cost', {})["completion_cost"] = completion_cost
+                    additional_metadata.setdefault('cost', {})["total_cost"] = prompt_cost + completion_cost 
                 except Exception as e:
                     logger.warning(f"Error adding cost: {e}")
-            else:
-                logger.debug("Model name not available in additional_metadata, skipping cost calculation")
-            
-            # Safely get total tokens and cost
-            if 'tokens' in additional_metadata and 'total' in additional_metadata['tokens']:
-                additional_metadata["total_tokens"] = additional_metadata["tokens"]["total"]
-            else:
-                additional_metadata["total_tokens"] = 0
-                logger.warning("Total tokens information not available")
 
-            if 'cost' in additional_metadata and 'total_cost' in additional_metadata['cost']:
+            # with open(filepath, 'r') as f:
+            #     data = json.load(f)
+            additional_metadata["total_tokens"] = additional_metadata["tokens"]["total"]
+            del additional_metadata["tokens"]
+            if "cost" in additional_metadata:
                 additional_metadata["total_cost"] = additional_metadata["cost"]["total_cost"]
+                del additional_metadata["cost"]
             else:
                 additional_metadata["total_cost"] = 0.0
-                logger.warning("Total cost information not available")
-
-            # Safely remove tokens and cost dictionaries if they exist
-            additional_metadata.pop("tokens", None)
-            additional_metadata.pop("cost", None)
             
-            # Safely merge metadata
-            combined_metadata = {}
-            if user_detail.get('trace_user_detail', {}).get('metadata'):
-                combined_metadata.update(user_detail['trace_user_detail']['metadata'])
-            if additional_metadata:
-                combined_metadata.update(additional_metadata)
+            combined_metadata = user_detail['trace_user_detail']['metadata'].copy()
+            combined_metadata.update(additional_metadata)
+            combined_metadata
 
             langchain_traces = langchain_tracer_extraction(data)
             final_result = convert_langchain_callbacks_output(langchain_traces)
+            final_result[0]['project_name'] = user_detail['project_name']
+            final_result[0]['trace_id'] = str(uuid.uuid4())
+            final_result[0]['session_id'] = None
+            final_result[0]['metadata'] = combined_metadata
+            final_result[0]['pipeline'] = user_detail['trace_user_detail']['pipeline']
+
+            filepath_3 = os.path.join(os.getcwd(), "final_result.json")
+            with open(filepath_3, 'w') as f:
+                json.dump(final_result, f, indent=2)
             
-            # Safely set required fields in final_result
-            if final_result and isinstance(final_result, list) and len(final_result) > 0:
-                final_result[0]['project_name'] = user_detail.get('project_name', '')
-                final_result[0]['trace_id'] = str(uuid.uuid4())
-                final_result[0]['session_id'] = None
-                final_result[0]['metadata'] = combined_metadata
-                final_result[0]['pipeline'] = user_detail.get('trace_user_detail', {}).get('pipeline')
-
-                filepath_3 = os.path.join(os.getcwd(), "final_result.json")
-                with open(filepath_3, 'w') as f:
-                    json.dump(final_result, f, indent=2)
-                
-                print(filepath_3)
-            else:
-                logger.warning("No valid langchain traces found in final_result")
+            
+            print(filepath_3)
 
-            additional_metadata_keys = list(additional_metadata.keys()) if additional_metadata else None
+            additional_metadata_keys = additional_metadata.keys() if additional_metadata else None
 
             UploadTraces(json_file_path=filepath_3,
                          project_name=self.project_name,

{ragaai_catalyst-2.1.5b7.dist-info → ragaai_catalyst-2.1.5b9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ragaai_catalyst
-Version: 2.1.5b7
+Version: 2.1.5b9
 Summary: RAGA AI CATALYST
 Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>
 Requires-Python: <3.13,>=3.9

{ragaai_catalyst-2.1.5b7.dist-info → ragaai_catalyst-2.1.5b9.dist-info}/RECORD

@@ -13,9 +13,9 @@ ragaai_catalyst/synthetic_data_generation.py,sha256=uDV9tNwto2xSkWg5XHXUvjErW-4P
 ragaai_catalyst/utils.py,sha256=TlhEFwLyRU690HvANbyoRycR3nQ67lxVUQoUOfTPYQ0,3772
 ragaai_catalyst/tracers/__init__.py,sha256=LfgTes-nHpazssbGKnn8kyLZNr49kIPrlkrqqoTFTfc,301
 ragaai_catalyst/tracers/distributed.py,sha256=AIRvS5Ur4jbFDXsUkYuCTmtGoHHx3LOG4n5tWOh610U,10330
-ragaai_catalyst/tracers/langchain_callback.py,sha256=LvMBhgvAX8ftyBQ9Naeui46EoDa2nHQZq48Ra6nL-Qg,21991
+ragaai_catalyst/tracers/langchain_callback.py,sha256=yR6TdzcLzkqerRwIkT7qvtb6v5HLVFiSHvwyvyI8wnE,28940
 ragaai_catalyst/tracers/llamaindex_callback.py,sha256=ZY0BJrrlz-P9Mg2dX-ZkVKG3gSvzwqBtk7JL_05MiYA,14028
-ragaai_catalyst/tracers/tracer.py,sha256=bLgO3lQmoumo-JtqZFi4DUqhu9itM5GaLFNY7hmsI1g,20267
+ragaai_catalyst/tracers/tracer.py,sha256=Yq2HhgT4785t9573kksJ7ngM3qCLPgZbZ0IpgOHdTTo,19223
 ragaai_catalyst/tracers/upload_traces.py,sha256=2TWdRTN6FMaX-dqDv8BJWQS0xrCGYKkXEYOi2kK3Z3Y,5487
 ragaai_catalyst/tracers/agentic_tracing/README.md,sha256=X4QwLb7-Jg7GQMIXj-SerZIgDETfw-7VgYlczOR8ZeQ,4508
 ragaai_catalyst/tracers/agentic_tracing/__init__.py,sha256=yf6SKvOPSpH-9LiKaoLKXwqj5sez8F_5wkOb91yp0oE,260
@@ -32,7 +32,7 @@ ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py,sha256=--wvhOJ-J
 ragaai_catalyst/tracers/agentic_tracing/tracers/base.py,sha256=88rX7OkOGEyVNECUrc4bYqODyulXve_-99d9ku5hBeQ,37373
 ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py,sha256=l3x3uFO5ov93I7UUrUX1M06WVGy2ug2jEZ1G7o315z4,13075
 ragaai_catalyst/tracers/agentic_tracing/tracers/langgraph_tracer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=s6BRoBteCRF8XrXGnmZ98ZWPrSONC5RObPXNaq-im3w,31782
+ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=91aWXJGb3GDfyDfJyA7Irnk3XSyfkQaQppW_NMORGJQ,31725
 ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py,sha256=6hsg-Yw11v4qeELI1CWrdX8BXf-wJrTF5smBI5prgoo,15873
 ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py,sha256=m8CxYkl7iMiFya_lNwN1ykBc3Pmo-2pR_2HmpptwHWQ,10352
 ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py,sha256=4rWL7fIJE5wN0nwh6fMWyh3OrrenZHJkNzyQXikyzQI,13771
@@ -48,7 +48,7 @@ ragaai_catalyst/tracers/agentic_tracing/utils/file_name_tracker.py,sha256=515NND
 ragaai_catalyst/tracers/agentic_tracing/utils/generic.py,sha256=WwXT01xmp8MSr7KinuDCSK9a1ifpLcT7ajFkvYviG_A,1190
 ragaai_catalyst/tracers/agentic_tracing/utils/get_user_trace_metrics.py,sha256=vPZ4dn4EHFW0kqd1GyRpsYXbfrRrd0DXCmh-pzsDBNE,1109
 ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py,sha256=wlXCuaRe81s-7FWdJ_MquXFGRZZfNrZxLIIxl-Ohbqk,15541
-ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json,sha256=E_uKa1SSrigaorCiAShZr4inKNMc54jcEy4B_7pT4DA,295002
+ragaai_catalyst/tracers/agentic_tracing/utils/model_costs.json,sha256=kQwC8AYTfJCqPm1F_heR7FoEhIpEZgBRWvkHRncfhzU,298689
 ragaai_catalyst/tracers/agentic_tracing/utils/span_attributes.py,sha256=MqeRNGxzeuh9qTK0NbYMftl9V9Z0V7gMgBoHkrXP56k,1592
 ragaai_catalyst/tracers/agentic_tracing/utils/system_monitor.py,sha256=H8WNsk4v_5T6OUw4TFOzlDLjQhJwjh1nAMyMAoqMEi4,6946
 ragaai_catalyst/tracers/agentic_tracing/utils/trace_utils.py,sha256=RciiDdo2riibEoM8X0FKHaXi78y3bWwNkV8U0leqigk,3508
@@ -65,8 +65,8 @@ ragaai_catalyst/tracers/utils/__init__.py,sha256=KeMaZtYaTojilpLv65qH08QmpYclfpa
 ragaai_catalyst/tracers/utils/convert_langchain_callbacks_output.py,sha256=ofrNrxf2b1hpjDh_zeaxiYq86azn1MF3kW8-ViYPEg0,1641
 ragaai_catalyst/tracers/utils/langchain_tracer_extraction_logic.py,sha256=cghjCuUe8w-2MZdh9xgtRGe3y219u26GGzpnuY4Wt6Q,3047
 ragaai_catalyst/tracers/utils/utils.py,sha256=ViygfJ7vZ7U0CTSA1lbxVloHp4NSlmfDzBRNCJuMhis,2374
-ragaai_catalyst-2.1.5b7.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ragaai_catalyst-2.1.5b7.dist-info/METADATA,sha256=cVHEchxtHjkR_9AKHeDlqTpZqYNMSoF5gKR37D9icAY,12764
-ragaai_catalyst-2.1.5b7.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ragaai_catalyst-2.1.5b7.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
-ragaai_catalyst-2.1.5b7.dist-info/RECORD,,
+ragaai_catalyst-2.1.5b9.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ragaai_catalyst-2.1.5b9.dist-info/METADATA,sha256=r88NroNJJuoi1skQURXDpKdMc4fo1FciVpxFtruaYuo,12764
+ragaai_catalyst-2.1.5b9.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ragaai_catalyst-2.1.5b9.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
+ragaai_catalyst-2.1.5b9.dist-info/RECORD,,