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

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -13,7 +13,9 @@ import traceback
 import importlib
 import sys
 from litellm import model_cost
+from llama_index.core.base.llms.types import ChatResponse
 
+from .base import BaseTracer
 from ..utils.llm_utils import (
     extract_model_name,
     extract_parameters,
@@ -80,7 +82,6 @@ class LLMTracerMixin:
     def instrument_llm_calls(self):
         """Enable LLM instrumentation"""
         self.auto_instrument_llm = True
-
         # Check currently loaded modules
         if "vertexai" in sys.modules:
             self.patch_vertex_ai_methods(sys.modules["vertexai"])
@@ -97,28 +98,32 @@ class LLMTracerMixin:
             self.patch_langchain_google_methods(sys.modules["langchain_google_vertexai"])
         if "langchain_google_genai" in sys.modules:
             self.patch_langchain_google_methods(sys.modules["langchain_google_genai"])
+
         if "langchain_openai" in sys.modules:
             self.patch_langchain_openai_methods(sys.modules["langchain_openai"])
         if "langchain_anthropic" in sys.modules:
             self.patch_langchain_anthropic_methods(sys.modules["langchain_anthropic"])
 
+        if "llama_index" in sys.modules:
+            self.patch_llama_index_methods(sys.modules["llama_index"])
+
         # Register hooks for future imports with availability checks
         if self.check_package_available("vertexai"):
             wrapt.register_post_import_hook(self.patch_vertex_ai_methods, "vertexai")
             wrapt.register_post_import_hook(
                 self.patch_vertex_ai_methods, "vertexai.generative_models"
             )
-        
+
         if self.check_package_available("openai") and self.validate_openai_key():
             wrapt.register_post_import_hook(self.patch_openai_methods, "openai")
             wrapt.register_post_import_hook(self.patch_openai_beta_methods, "openai")
-        
+
         if self.check_package_available("litellm"):
             wrapt.register_post_import_hook(self.patch_litellm_methods, "litellm")
-        
+
         if self.check_package_available("anthropic"):
             wrapt.register_post_import_hook(self.patch_anthropic_methods, "anthropic")
-        
+
         if self.check_package_available("google.generativeai"):
             wrapt.register_post_import_hook(
                 self.patch_google_genai_methods, "google.generativeai"
@@ -129,12 +134,16 @@ class LLMTracerMixin:
             wrapt.register_post_import_hook(
                 self.patch_langchain_google_methods, "langchain_google_vertexai"
             )
+
+
+        # Add hooks for llama-index
+        wrapt.register_post_import_hook(self.patch_llama_index_methods, "llama_index")
         
         if self.check_package_available("langchain_google_genai"):
             wrapt.register_post_import_hook(
                 self.patch_langchain_google_methods, "langchain_google_genai"
             )
-            
+
         if self.check_package_available("langchain_openai"):
             wrapt.register_post_import_hook(
                 self.patch_langchain_openai_methods, "langchain_openai"
@@ -151,11 +160,88 @@ class LLMTracerMixin:
     def instrument_network_calls(self):
         """Enable network instrumentation for LLM calls"""
         self.auto_instrument_network = True
-        
+
     def instrument_file_io_calls(self):
         """Enable file IO instrumentation for LLM calls"""
         self.auto_instrument_file_io = True
 
+    def patch_llama_index_methods(self, module):
+        """Patch llama-index LLM methods"""
+        try:
+            # Handle OpenAI LLM from llama-index
+            if hasattr(module, "llms"):
+                # OpenAI
+                if hasattr(module.llms, "openai"):
+                    openai_module = module.llms.openai
+                    if hasattr(openai_module, "OpenAI"):
+                        llm_class = getattr(openai_module, "OpenAI")
+                        self.wrap_method(llm_class, "complete")
+                        self.wrap_method(llm_class, "acomplete")
+                        self.wrap_method(llm_class, "chat")
+                        self.wrap_method(llm_class, "achat")
+                        self.wrap_method(llm_class, "stream_chat")
+                        # self.wrap_method(llm_class, "stream_achat")
+                        self.wrap_method(llm_class, "stream_complete")
+                        # self.wrap_method(llm_class, "stream_acomplete")
+
+                # Anthropic
+                if hasattr(module.llms, "anthropic"):
+                    anthropic_module = module.llms.anthropic
+                    if hasattr(anthropic_module, "Anthropic"):
+                        llm_class = getattr(anthropic_module, "Anthropic")
+                        self.wrap_method(llm_class, "complete")
+                        self.wrap_method(llm_class, "acomplete")
+                        self.wrap_method(llm_class, "chat")
+                        self.wrap_method(llm_class, "achat")
+                        self.wrap_method(llm_class, "stream_chat")
+                        # self.wrap_method(llm_class, "stream_achat")
+
+                # Azure OpenAI
+                if hasattr(module.llms, "azure_openai"):
+                    azure_module = module.llms.azure_openai
+                    if hasattr(azure_module, "AzureOpenAI"):
+                        llm_class = getattr(azure_module, "AzureOpenAI")
+                        self.wrap_method(llm_class, "complete")
+                        self.wrap_method(llm_class, "acomplete")
+                        self.wrap_method(llm_class, "chat")
+                        self.wrap_method(llm_class, "achat")
+                        self.wrap_method(llm_class, "stream_chat")
+                        # self.wrap_method(llm_class, "stream_achat")
+
+                # LiteLLM
+                if hasattr(module.llms, "litellm"):
+                    litellm_module = module.llms.litellm
+                    if hasattr(litellm_module, "LiteLLM"):
+                        llm_class = getattr(litellm_module, "LiteLLM")
+                        self.wrap_method(llm_class, "complete")
+                        self.wrap_method(llm_class, "acomplete")
+                        self.wrap_method(llm_class, "chat")
+                        self.wrap_method(llm_class, "achat")
+
+                # Vertex AI
+                if hasattr(module.llms, "vertex"):
+                    vertex_module = module.llms.vertex
+                    if hasattr(vertex_module, "Vertex"):
+                        llm_class = getattr(vertex_module, "Vertex")
+                        self.wrap_method(llm_class, "complete")
+                        self.wrap_method(llm_class, "acomplete")
+                        self.wrap_method(llm_class, "chat")
+                        self.wrap_method(llm_class, "achat")
+
+                # Gemini
+                if hasattr(module.llms, "gemini"):
+                    gemini_module = module.llms.gemini
+                    if hasattr(gemini_module, "Gemini"):
+                        llm_class = getattr(gemini_module, "Gemini")
+                        self.wrap_method(llm_class, "complete")
+                        self.wrap_method(llm_class, "acomplete")
+                        self.wrap_method(llm_class, "chat")
+                        self.wrap_method(llm_class, "achat")
+
+        except Exception as e:
+            # Log the error but continue execution
+            print(f"Warning: Failed to patch llama-index methods: {str(e)}")
+
     def patch_openai_methods(self, module):
         try:
             if hasattr(module, "OpenAI"):
@@ -167,12 +253,12 @@ class LLMTracerMixin:
         except Exception as e:
             # Log the error but continue execution
             print(f"Warning: Failed to patch OpenAI methods: {str(e)}")
-            
+
     def patch_langchain_openai_methods(self, module):
         try:
             if hasattr(module, 'ChatOpenAI'):
                 client_class = getattr(module, "ChatOpenAI")
-                
+
                 if hasattr(client_class, "invoke"):
                     self.wrap_langchain_openai_method(client_class, f"{client_class.__name__}.invoke")
                 elif hasattr(client_class, "run"):
@@ -185,7 +271,7 @@ class LLMTracerMixin:
         except Exception as e:
             # Log the error but continue execution
             print(f"Warning: Failed to patch OpenAI methods: {str(e)}")
-            
+
     def patch_langchain_anthropic_methods(self, module):
         try:
             if hasattr(module, 'ChatAnthropic'):
@@ -211,6 +297,7 @@ class LLMTracerMixin:
         openai.beta.threads.runs.create(...) are automatically traced.
         """
         # Make sure openai_module has a 'beta' attribute
+        openai_module.api_type = "openai"
         if not hasattr(openai_module, "beta"):
             return
 
@@ -239,7 +326,6 @@ class LLMTracerMixin:
                     if hasattr(runs_obj, method_name):
                         self.wrap_method(runs_obj, method_name)
 
-
     def patch_anthropic_methods(self, module):
         if hasattr(module, "Anthropic"):
             client_class = getattr(module, "Anthropic")
@@ -342,6 +428,7 @@ class LLMTracerMixin:
                         return await self.trace_llm_call(
                             original_create, *args, **kwargs
                         )
+
                     client_self.chat.completions.create = wrapped_create
             else:
                 # Patch sync methods for OpenAI
@@ -353,10 +440,11 @@ class LLMTracerMixin:
                         return self.trace_llm_call_sync(
                             original_create, *args, **kwargs
                         )
+
                     client_self.chat.completions.create = wrapped_create
 
         setattr(client_class, "__init__", patched_init)
-        
+
     def wrap_langchain_openai_method(self, client_class, method_name):
         method = method_name.split(".")[-1]
         original_init = getattr(client_class, method)
@@ -372,14 +460,14 @@ class LLMTracerMixin:
                 return self.trace_llm_call_sync(original_init, *args, **kwargs)
 
         setattr(client_class, method, patched_init)
-        
+
     def wrap_langchain_anthropic_method(self, client_class, method_name):
         original_init = getattr(client_class, method_name)
 
         @functools.wraps(original_init)
         def patched_init(*args, **kwargs):
             is_async = "AsyncChatAnthropic" in client_class.__name__
-            
+
             if is_async:
                 return self.trace_llm_call(original_init, *args, **kwargs)
             else:
@@ -447,20 +535,20 @@ class LLMTracerMixin:
             self.patches.append((obj, method_name, original_method))
 
     def create_llm_component(
-        self,
-        component_id,
-        hash_id,
-        name,
-        llm_type,
-        version,
-        memory_used,
-        start_time,
-        input_data,
-        output_data,
-        cost={},
-        usage={},
-        error=None,
-        parameters={},
+            self,
+            component_id,
+            hash_id,
+            name,
+            llm_type,
+            version,
+            memory_used,
+            start_time,
+            input_data,
+            output_data,
+            cost={},
+            usage={},
+            error=None,
+            parameters={},
     ):
         # Update total metrics
         self.total_tokens += usage.get("total_tokens", 0)
@@ -476,7 +564,7 @@ class LLMTracerMixin:
             for interaction in self.component_user_interaction.get(component_id, []):
                 if interaction["interaction_type"] in ["input", "output"]:
                     input_output_interactions.append(interaction)
-            interactions.extend(input_output_interactions) 
+            interactions.extend(input_output_interactions)
         if self.auto_instrument_file_io:
             file_io_interactions = []
             for interaction in self.component_user_interaction.get(component_id, []):
@@ -502,13 +590,24 @@ class LLMTracerMixin:
             list(parameters_to_display.items())[: self.MAX_PARAMETERS_TO_DISPLAY]
         )
 
-        # Get tags, metrics
-        # tags
+        # Set the Context and GT
+        span_gt = None
+        span_context = None
+        if name in self.span_attributes_dict:
+            span_gt = self.span_attributes_dict[name].gt
+            span_context = self.span_attributes_dict[name].context
+
+            logger.debug(f"span context {span_context}, span_gt {span_gt}")
+
+        # Tags
         tags = []
         if name in self.span_attributes_dict:
             tags = self.span_attributes_dict[name].tags or []
 
-        # metrics
+        # Get End Time
+        end_time = datetime.now().astimezone().isoformat()
+
+        # Metrics
         metrics = []
         if name in self.span_attributes_dict:
             raw_metrics = self.span_attributes_dict[name].metrics or []
@@ -517,17 +616,32 @@ class LLMTracerMixin:
                 counter = sum(1 for x in self.visited_metrics if x.startswith(base_metric_name))
                 metric_name = f'{base_metric_name}_{counter}' if counter > 0 else base_metric_name
                 self.visited_metrics.append(metric_name)
-                metric["name"] = metric_name  
+                metric["name"] = metric_name
                 metrics.append(metric)
 
+        # TODO TO check i/p and o/p is according or not
+        input = input_data["args"] if hasattr(input_data, "args") else input_data
+        output = output_data.output_response if output_data else None
+        #print("Prompt input:",input)
+        prompt = self.convert_to_content(input)
+        #print("Prompt Output: ",prompt)
+        #print("Response input: ",output)
+        response = self.convert_to_content(output)
+        #print("Response output: ",response)
+
+        # TODO: Execute & Add the User requested metrics here
+        formatted_metrics = BaseTracer.get_formatted_metric(self.span_attributes_dict, self.project_id, name)
+        if formatted_metrics:
+            metrics.extend(formatted_metrics)
+
         component = {
             "id": component_id,
             "hash_id": hash_id,
             "source_hash_id": None,
             "type": "llm",
             "name": name,
-            "start_time": start_time.isoformat(),
-            "end_time": datetime.now().astimezone().isoformat(),
+            "start_time": start_time,
+            "end_time": end_time,
             "error": error,
             "parent_id": self.current_agent_id.get(),
             "info": {
@@ -541,10 +655,8 @@ class LLMTracerMixin:
             },
             "extra_info": parameters,
             "data": {
-                "input": (
-                    input_data["args"] if hasattr(input_data, "args") else input_data
-                ),
-                "output": output_data.output_response if output_data else None,
+                "input": input,
+                "output": output,
                 "memory_used": memory_used,
             },
             "metrics": metrics,
@@ -552,19 +664,77 @@ class LLMTracerMixin:
             "interactions": interactions,
         }
 
-        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
+        # Assign context and gt if available
+        component["data"]["gt"] = span_gt
+        component["data"]["context"] = span_context
 
         # Reset the SpanAttributes context variable
         self.span_attributes_dict[name] = SpanAttributes(name)
 
         return component
 
+    # def convert_to_content(self, input_data):
+    #     if isinstance(input_data, dict):
+    #         messages = input_data.get("kwargs", {}).get("messages", [])
+    #     elif isinstance(input_data, list):
+    #         messages = input_data
+    #     else:
+    #         return ""
+        # return "\n".join(process_content(msg.get("content", "")) for msg in messages if msg.get("content"))
+  
+    def convert_to_content(self, input_data):
+        try:
+            if isinstance(input_data, dict):
+                messages = input_data.get("kwargs", {}).get("messages", [])
+            elif isinstance(input_data, list):
+                if len(input_data)>0 and isinstance(input_data[0]['content'],ChatResponse):
+                    extracted_messages = []
+
+                    for item in input_data:
+                        chat_response = item.get('content')
+                        if hasattr(chat_response, 'message') and hasattr(chat_response.message, 'blocks'):
+                            for block in chat_response.message.blocks:
+                                if hasattr(block, 'text'):
+                                    extracted_messages.append(block.text)
+                    messages=extracted_messages
+                    if isinstance(messages,list):
+                        return "\n".join(messages)
+                    
+                    #messages=[msg["content"] for msg in input_data if isinstance(msg, dict) and "content" in msg]
+                    #messages = [msg["content"].message for msg in input_data if isinstance(msg, dict) and "content" in msg and isinstance(msg["content"], ChatResponse)]
+                else:
+                    messages = input_data
+            elif isinstance(input_data,ChatResponse):
+                messages=input_data['content']
+            else:
+                return ""
+            res=""
+            # try:
+            res="\n".join(msg.get("content", "").strip() for msg in messages if msg.get("content"))
+        except Exception as e:
+            res=str(messages)
+        return res
+
+    def process_content(content):
+        if isinstance(content, str):
+            return content.strip()
+        elif isinstance(content, list):
+            # Handle list of content blocks
+            text_parts = []
+            for block in content:
+                if hasattr(block, 'text'):
+                    # Handle TextBlock-like objects
+                    text_parts.append(block.text.strip())
+                elif isinstance(block, dict) and 'text' in block:
+                    # Handle dictionary with text field
+                    text_parts.append(block['text'].strip())
+            return " ".join(text_parts)
+        elif isinstance(content, dict):
+            # Handle dictionary content
+            return content.get('text', '').strip()
+        return ""
+
+    
     def start_component(self, component_id):
         """Start tracking network calls for a component"""
         self.component_network_calls[component_id] = []
@@ -582,7 +752,7 @@ class LLMTracerMixin:
         if not self.auto_instrument_llm:
             return await original_func(*args, **kwargs)
 
-        start_time = datetime.now().astimezone()
+        start_time = datetime.now().astimezone().isoformat()
         start_memory = psutil.Process().memory_info().rss
         component_id = str(uuid.uuid4())
         hash_id = generate_unique_hash(original_func, args, kwargs)
@@ -605,7 +775,8 @@ class LLMTracerMixin:
                 if stream:
                     prompt_messages = kwargs['messages']
                     # Create response message for streaming case
-                    response_message = {"role": "assistant", "content": result} if result else {"role": "assistant", "content": ""}
+                    response_message = {"role": "assistant", "content": result} if result else {"role": "assistant",
+                                                                                                "content": ""}
                     token_usage = num_tokens_from_messages(model_name, prompt_messages, response_message)
                 else:
                     token_usage = extract_token_usage(result)
@@ -685,7 +856,7 @@ class LLMTracerMixin:
         if not self.auto_instrument_llm:
             return original_func(*args, **kwargs)
 
-        start_time = datetime.now().astimezone()
+        start_time = datetime.now().astimezone().isoformat()
         component_id = str(uuid.uuid4())
         hash_id = generate_unique_hash(original_func, args, kwargs)
 
@@ -707,13 +878,14 @@ class LLMTracerMixin:
 
             # Extract token usage and calculate cost
             model_name = extract_model_name(args, kwargs, result)
-            
+
             if 'stream' in kwargs:
                 stream = kwargs['stream']
                 if stream:
                     prompt_messages = kwargs['messages']
                     # Create response message for streaming case
-                    response_message = {"role": "assistant", "content": result} if result else {"role": "assistant", "content": ""}
+                    response_message = {"role": "assistant", "content": result} if result else {"role": "assistant",
+                                                                                                "content": ""}
                     token_usage = num_tokens_from_messages(model_name, prompt_messages, response_message)
                 else:
                     token_usage = extract_token_usage(result)
@@ -786,12 +958,12 @@ class LLMTracerMixin:
             raise
 
     def trace_llm(
-        self,
-        name: str = None,
-        tags: List[str] = [],
-        metadata: Dict[str, Any] = {},
-        metrics: List[Dict[str, Any]] = [],
-        feedback: Optional[Any] = None,
+            self,
+            name: str = None,
+            tags: List[str] = [],
+            metadata: Dict[str, Any] = {},
+            metrics: List[Dict[str, Any]] = [],
+            feedback: Optional[Any] = None,
     ):
         if name not in self.span_attributes_dict:
             self.span_attributes_dict[name] = SpanAttributes(name)
@@ -817,7 +989,7 @@ class LLMTracerMixin:
                 logger.error(f"Validation Error: {e}")
             except Exception as e:
                 logger.error(f"Error adding metric: {e}")
-                
+
         if feedback:
             self.span(name).add_feedback(feedback)
 
@@ -871,8 +1043,9 @@ class LLMTracerMixin:
 
                     if error_info:
                         llm_component["error"] = error_info["error"]
-                    
+
                     self.end_component(component_id)
+
                     # metrics
                     metrics = []
                     if name in self.span_attributes_dict:
@@ -882,7 +1055,7 @@ class LLMTracerMixin:
                             counter = sum(1 for x in self.visited_metrics if x.startswith(base_metric_name))
                             metric_name = f'{base_metric_name}_{counter}' if counter > 0 else base_metric_name
                             self.visited_metrics.append(metric_name)
-                            metric["name"] = metric_name  
+                            metric["name"] = metric_name
                             metrics.append(metric)
                     llm_component["metrics"] = metrics
                     if parent_agent_id:
@@ -946,9 +1119,9 @@ class LLMTracerMixin:
                             counter = sum(1 for x in self.visited_metrics if x.startswith(base_metric_name))
                             metric_name = f'{base_metric_name}_{counter}' if counter > 0 else base_metric_name
                             self.visited_metrics.append(metric_name)
-                            metric["name"] = metric_name  
+                            metric["name"] = metric_name
                             metrics.append(metric)
-                    llm_component["metrics"] = metrics  
+                    llm_component["metrics"] = metrics
                     if parent_agent_id:
                         children = self.agent_children.get()
                         children.append(llm_component)

ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py

@@ -119,9 +119,6 @@ class AgenticTracing(
         self.component_network_calls = {}  # Store network calls per component
         self.component_user_interaction = {}
 
-        # Create output directory if it doesn't exist
-        self.output_dir = Path("./traces")  # Using default traces directory
-        self.output_dir.mkdir(exist_ok=True)
 
     def start_component(self, component_id: str):
         """Start tracking network calls for a component"""
@@ -364,7 +361,7 @@ class AgenticTracing(
 
         # Check if there's an active agent context
         current_agent_id = self.current_agent_id.get()
-        if current_agent_id and component_data["type"] in ["llm", "tool"]:
+        if current_agent_id and component_data["type"] in ["llm", "tool", "custom"]:
             # Add this component as a child of the current agent
             current_children = self.agent_children.get()
             current_children.append(component_data)

ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py

@@ -7,6 +7,7 @@ import functools
 from typing import Optional, Any, Dict, List
 
 from pydantic import tools
+from .base import BaseTracer
 from ..utils.unique_decorator import generate_unique_hash_simple
 import contextvars
 import asyncio
@@ -483,6 +484,10 @@ class ToolTracerMixin:
                 metric["name"] = metric_name  
                 metrics.append(metric)
 
+        formatted_metrics = BaseTracer.get_formatted_metric(self.span_attributes_dict, self.project_id, name)
+        if formatted_metrics:
+            metrics.extend(formatted_metrics)
+
         start_time = kwargs["start_time"]
         component = {
             "id": kwargs["component_id"],

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, **kwargs):
+    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": 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}")