ragaai-catalyst 2.1.3b0__py3-none-any.whl → 2.1.4b1__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -22,8 +22,17 @@ from ..utils.llm_utils import (
     extract_llm_output,
 )
 from ..utils.trace_utils import load_model_costs
-from ..utils.unique_decorator import generate_unique_hash_simple   
+from ..utils.unique_decorator import generate_unique_hash_simple
 from ..utils.file_name_tracker import TrackName
+from ..utils.span_attributes import SpanAttributes
+import logging
+
+logger = logging.getLogger(__name__)
+logging_level = (
+    logger.setLevel(logging.DEBUG)
+    if os.getenv("DEBUG")
+    else logger.setLevel(logging.INFO)
+)
 
 
 class LLMTracerMixin:
@@ -36,29 +45,36 @@ class LLMTracerMixin:
         except Exception as e:
             self.model_costs = {
                 # TODO: Default cost handling needs to be improved
-                "default": {
-                    "input_cost_per_token": 0.0,
-                    "output_cost_per_token": 0.0
-                }
+                "default": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}
             }
         self.MAX_PARAMETERS_TO_DISPLAY = 10
-        self.current_llm_call_name = contextvars.ContextVar("llm_call_name", default=None)
-        self.component_network_calls = {}  
+        self.current_llm_call_name = contextvars.ContextVar(
+            "llm_call_name", default=None
+        )
+        self.component_network_calls = {}
         self.component_user_interaction = {}
-        self.current_component_id = None  
+        self.current_component_id = None
         self.total_tokens = 0
         self.total_cost = 0.0
         self.llm_data = {}
 
+        # Add auto_instrument options
+        self.auto_instrument_llm = False
+        self.auto_instrument_user_interaction = False
+        self.auto_instrument_network = False
+
     def instrument_llm_calls(self):
+        """Enable LLM instrumentation"""
+        self.auto_instrument_llm = True
+
         # Handle modules that are already imported
         import sys
-        
+
         if "vertexai" in sys.modules:
             self.patch_vertex_ai_methods(sys.modules["vertexai"])
         if "vertexai.generative_models" in sys.modules:
             self.patch_vertex_ai_methods(sys.modules["vertexai.generative_models"])
-            
+
         if "openai" in sys.modules:
             self.patch_openai_methods(sys.modules["openai"])
         if "litellm" in sys.modules:
@@ -68,21 +84,39 @@ class LLMTracerMixin:
         if "google.generativeai" in sys.modules:
             self.patch_google_genai_methods(sys.modules["google.generativeai"])
         if "langchain_google_vertexai" in sys.modules:
-            self.patch_langchain_google_methods(sys.modules["langchain_google_vertexai"])
+            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"])
 
         # Register hooks for future imports
         wrapt.register_post_import_hook(self.patch_vertex_ai_methods, "vertexai")
-        wrapt.register_post_import_hook(self.patch_vertex_ai_methods, "vertexai.generative_models")
+        wrapt.register_post_import_hook(
+            self.patch_vertex_ai_methods, "vertexai.generative_models"
+        )
         wrapt.register_post_import_hook(self.patch_openai_methods, "openai")
         wrapt.register_post_import_hook(self.patch_litellm_methods, "litellm")
         wrapt.register_post_import_hook(self.patch_anthropic_methods, "anthropic")
-        wrapt.register_post_import_hook(self.patch_google_genai_methods, "google.generativeai")
-        
+        wrapt.register_post_import_hook(
+            self.patch_google_genai_methods, "google.generativeai"
+        )
+
         # Add hooks for LangChain integrations
-        wrapt.register_post_import_hook(self.patch_langchain_google_methods, "langchain_google_vertexai")
-        wrapt.register_post_import_hook(self.patch_langchain_google_methods, "langchain_google_genai")
+        wrapt.register_post_import_hook(
+            self.patch_langchain_google_methods, "langchain_google_vertexai"
+        )
+        wrapt.register_post_import_hook(
+            self.patch_langchain_google_methods, "langchain_google_genai"
+        )
+
+    def instrument_user_interaction_calls(self):
+        """Enable user interaction instrumentation for LLM calls"""
+        self.auto_instrument_user_interaction = True
+
+    def instrument_network_calls(self):
+        """Enable network instrumentation for LLM calls"""
+        self.auto_instrument_network = True
 
     def patch_openai_methods(self, module):
         try:
@@ -106,20 +140,20 @@ class LLMTracerMixin:
         if hasattr(module, "GenerativeModel"):
             model_class = getattr(module, "GenerativeModel")
             self.wrap_genai_model_methods(model_class)
-        
+
         # Patch LangChain integration
         if hasattr(module, "ChatGoogleGenerativeAI"):
             chat_class = getattr(module, "ChatGoogleGenerativeAI")
             # Wrap invoke method to capture messages
             original_invoke = chat_class.invoke
-            
+
             def patched_invoke(self, messages, *args, **kwargs):
                 # Store messages in the instance for later use
                 self._last_messages = messages
                 return original_invoke(self, messages, *args, **kwargs)
-            
+
             chat_class.invoke = patched_invoke
-            
+
             # LangChain v0.2+ uses invoke/ainvoke
             self.wrap_method(chat_class, "_generate")
             if hasattr(chat_class, "_agenerate"):
@@ -137,7 +171,7 @@ class LLMTracerMixin:
             if hasattr(gen_models, "GenerativeModel"):
                 model_class = getattr(gen_models, "GenerativeModel")
                 self.wrap_vertex_model_methods(model_class)
-        
+
         # Also patch the class directly if available
         if hasattr(module, "GenerativeModel"):
             model_class = getattr(module, "GenerativeModel")
@@ -186,23 +220,31 @@ class LLMTracerMixin:
         def patched_init(client_self, *args, **kwargs):
             original_init(client_self, *args, **kwargs)
             # Check if this is AsyncOpenAI or OpenAI
-            is_async = 'AsyncOpenAI' in client_class.__name__
-            
+            is_async = "AsyncOpenAI" in client_class.__name__
+
             if is_async:
                 # Patch async methods for AsyncOpenAI
                 if hasattr(client_self.chat.completions, "create"):
                     original_create = client_self.chat.completions.create
+
                     @functools.wraps(original_create)
                     async def wrapped_create(*args, **kwargs):
-                        return await self.trace_llm_call(original_create, *args, **kwargs)
+                        return await self.trace_llm_call(
+                            original_create, *args, **kwargs
+                        )
+
                     client_self.chat.completions.create = wrapped_create
             else:
                 # Patch sync methods for OpenAI
                 if hasattr(client_self.chat.completions, "create"):
                     original_create = client_self.chat.completions.create
+
                     @functools.wraps(original_create)
                     def wrapped_create(*args, **kwargs):
-                        return self.trace_llm_call_sync(original_create, *args, **kwargs)
+                        return self.trace_llm_call_sync(
+                            original_create, *args, **kwargs
+                        )
+
                     client_self.chat.completions.create = wrapped_create
 
         setattr(client_class, "__init__", patched_init)
@@ -240,41 +282,64 @@ class LLMTracerMixin:
         if isinstance(obj, type):
             # Store the original class method
             original_method = getattr(obj, method_name)
-            
+
             @wrapt.decorator
             def wrapper(wrapped, instance, args, kwargs):
                 if asyncio.iscoroutinefunction(wrapped):
                     return self.trace_llm_call(wrapped, *args, **kwargs)
                 return self.trace_llm_call_sync(wrapped, *args, **kwargs)
-            
+
             # Wrap the class method
             wrapped_method = wrapper(original_method)
             setattr(obj, method_name, wrapped_method)
             self.patches.append((obj, method_name, original_method))
-            
+
         else:
             # For instance methods
             original_method = getattr(obj, method_name)
-            
+
             @wrapt.decorator
             def wrapper(wrapped, instance, args, kwargs):
                 if asyncio.iscoroutinefunction(wrapped):
                     return self.trace_llm_call(wrapped, *args, **kwargs)
                 return self.trace_llm_call_sync(wrapped, *args, **kwargs)
-            
+
             wrapped_method = wrapper(original_method)
             setattr(obj, method_name, wrapped_method)
             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={}):
+    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={},
+    ):
         # Update total metrics
         self.total_tokens += usage.get("total_tokens", 0)
         self.total_cost += cost.get("total_cost", 0)
 
+        network_calls = []
+        if self.auto_instrument_network:
+            network_calls = self.component_network_calls.get(component_id, [])
+
+        interactions = []
+        if self.auto_instrument_user_interaction:
+            interactions = self.component_user_interaction.get(component_id, [])
+
         parameters_to_display = {}
-        if 'run_manager' in parameters:
-            parameters_obj = parameters['run_manager']
-            if hasattr(parameters_obj, 'metadata'):
+        if "run_manager" in parameters:
+            parameters_obj = parameters["run_manager"]
+            if hasattr(parameters_obj, "metadata"):
                 metadata = parameters_obj.metadata
                 # parameters = {'metadata': metadata}
                 parameters_to_display.update(metadata)
@@ -283,9 +348,11 @@ class LLMTracerMixin:
         for key, value in parameters.items():
             if isinstance(value, (str, int, float, bool)):
                 parameters_to_display[key] = value
-        
+
         # Limit the number of parameters to display
-        parameters_to_display = dict(list(parameters_to_display.items())[:self.MAX_PARAMETERS_TO_DISPLAY])
+        parameters_to_display = dict(
+            list(parameters_to_display.items())[: self.MAX_PARAMETERS_TO_DISPLAY]
+        )
 
         component = {
             "id": component_id,
@@ -303,23 +370,30 @@ class LLMTracerMixin:
                 "memory_used": memory_used,
                 "cost": cost,
                 "tokens": usage,
-                **parameters_to_display
+                "tags": self.span_attributes_dict[name].tags or [],
+                **parameters_to_display,
             },
             "extra_info": parameters,
             "data": {
-                "input": input_data['args'] if hasattr(input_data, 'args') else input_data,
+                "input": (
+                    input_data["args"] if hasattr(input_data, "args") else input_data
+                ),
                 "output": output_data.output_response if output_data else None,
-                "memory_used": memory_used
+                "memory_used": memory_used,
             },
-            "network_calls": self.component_network_calls.get(component_id, []),
-            "interactions": self.component_user_interaction.get(component_id, [])
+            "metrics": self.span_attributes_dict[name].metrics or [],
+            "network_calls": network_calls,
+            "interactions": interactions,
         }
 
-        if self.gt: 
+        if self.gt:
             component["data"]["gt"] = self.gt
 
+        # Reset the SpanAttributes context variable
+        self.span_attributes_dict[name] = SpanAttributes(name)
+
         return component
-    
+
     def start_component(self, component_id):
         """Start tracking network calls for a component"""
         self.component_network_calls[component_id] = []
@@ -329,16 +403,18 @@ class LLMTracerMixin:
         """Stop tracking network calls for a component"""
         self.current_component_id = None
 
-
     async def trace_llm_call(self, original_func, *args, **kwargs):
         """Trace an LLM API call"""
         if not self.is_active:
             return await original_func(*args, **kwargs)
 
+        if not self.auto_instrument_llm:
+            return await original_func(*args, **kwargs)
+
         start_time = datetime.now().astimezone()
         start_memory = psutil.Process().memory_info().rss
         component_id = str(uuid.uuid4())
-        hash_id = generate_unique_hash_simple(original_func) 
+        hash_id = generate_unique_hash_simple(original_func)
 
         # Start tracking network calls for this component
         self.start_component(component_id)
@@ -364,25 +440,26 @@ class LLMTracerMixin:
             name = self.current_llm_call_name.get()
             if name is None:
                 name = original_func.__name__
-            
+
             # Create LLM component
             llm_component = self.create_llm_component(
                 component_id=component_id,
                 hash_id=hash_id,
                 name=name,
                 llm_type=model_name,
-                version="1.0.0",
+                version=None,
                 memory_used=memory_used,
                 start_time=start_time,
                 input_data=input_data,
                 output_data=extract_llm_output(result),
                 cost=cost,
                 usage=token_usage,
-                parameters=parameters
+                parameters=parameters,
             )
-                
+
             # self.add_component(llm_component)
             self.llm_data = llm_component
+
             return result
 
         except Exception as e:
@@ -390,30 +467,31 @@ class LLMTracerMixin:
                 "code": 500,
                 "type": type(e).__name__,
                 "message": str(e),
-                "details": {}
+                "details": {},
             }
-            
+
             # End tracking network calls for this component
             self.end_component(component_id)
 
             name = self.current_llm_call_name.get()
             if name is None:
                 name = original_func.__name__
-            
+
             llm_component = self.create_llm_component(
                 component_id=component_id,
                 hash_id=hash_id,
                 name=name,
                 llm_type="unknown",
-                version="1.0.0",
+                version=None,
                 memory_used=0,
                 start_time=start_time,
                 input_data=extract_input_data(args, kwargs, None),
                 output_data=None,
-                error=error_component
+                error=error_component,
             )
-    
+
             self.add_component(llm_component)
+
             raise
 
     def trace_llm_call_sync(self, original_func, *args, **kwargs):
@@ -423,6 +501,9 @@ class LLMTracerMixin:
                 return asyncio.run(original_func(*args, **kwargs))
             return original_func(*args, **kwargs)
 
+        if not self.auto_instrument_llm:
+            return original_func(*args, **kwargs)
+
         start_time = datetime.now().astimezone()
         component_id = str(uuid.uuid4())
         hash_id = generate_unique_hash_simple(original_func)
@@ -456,24 +537,25 @@ class LLMTracerMixin:
             name = self.current_llm_call_name.get()
             if name is None:
                 name = original_func.__name__
-            
+
             # Create LLM component
             llm_component = self.create_llm_component(
                 component_id=component_id,
                 hash_id=hash_id,
                 name=name,
                 llm_type=model_name,
-                version="1.0.0",
+                version=None,
                 memory_used=memory_used,
                 start_time=start_time,
                 input_data=input_data,
                 output_data=extract_llm_output(result),
                 cost=cost,
                 usage=token_usage,
-                parameters=parameters
+                parameters=parameters,
             )
             self.llm_data = llm_component
             self.add_component(llm_component)
+
             return result
 
         except Exception as e:
@@ -481,9 +563,9 @@ class LLMTracerMixin:
                 "code": 500,
                 "type": type(e).__name__,
                 "message": str(e),
-                "details": {}
+                "details": {},
             }
-            
+
             # End tracking network calls for this component
             self.end_component(component_id)
 
@@ -493,41 +575,76 @@ class LLMTracerMixin:
 
             end_memory = psutil.Process().memory_info().rss
             memory_used = max(0, end_memory - start_memory)
-            
+
             llm_component = self.create_llm_component(
                 component_id=component_id,
                 hash_id=hash_id,
                 name=name,
                 llm_type="unknown",
-                version="1.0.0",
+                version=None,
                 memory_used=memory_used,
                 start_time=start_time,
                 input_data=extract_input_data(args, kwargs, None),
                 output_data=None,
-                error=error_component
+                error=error_component,
             )
             self.llm_data = llm_component
             self.add_component(llm_component, is_error=True)
+
             raise
 
-    def trace_llm(self, name: str = None):
+    def trace_llm(
+        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)
+        if tags:
+            self.span(name).add_tags(tags)
+        if metadata:
+            self.span(name).add_metadata(metadata)
+        if metrics:
+            if isinstance(metrics, dict):
+                metrics = [metrics]
+            for metric in metrics:
+                try:
+                    self.span(name).add_metrics(
+                        name=metric["name"],
+                        score=metric["score"],
+                        reasoning=metric.get("reasoning", ""),
+                        cost=metric.get("cost", None),
+                        latency=metric.get("latency", None),
+                        metadata=metric.get("metadata", {}),
+                        config=metric.get("config", {}),
+                    )
+                except KeyError as e:
+                    logger.error(f"Error adding metric: {e}")
+        if feedback:
+            self.span(name).add_feedback(feedback)
+
+        self.current_llm_call_name.set(name)
+
         def decorator(func):
             @self.file_tracker.trace_decorator
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
-                self.gt = kwargs.get('gt', None) if kwargs else None
+                self.gt = kwargs.get("gt", None) if kwargs else None
                 self.current_llm_call_name.set(name)
                 if not self.is_active:
                     return await func(*args, **kwargs)
-                
-                hash_id = generate_unique_hash_simple(func)                
+
+                hash_id = generate_unique_hash_simple(func)
                 component_id = str(uuid.uuid4())
                 parent_agent_id = self.current_agent_id.get()
                 self.start_component(component_id)
-                
+
                 error_info = None
                 result = None
-                
+
                 try:
                     result = await func(*args, **kwargs)
                     return result
@@ -537,7 +654,7 @@ class LLMTracerMixin:
                             "type": type(e).__name__,
                             "message": str(e),
                             "traceback": traceback.format_exc(),
-                            "timestamp": datetime.now().isoformat()
+                            "timestamp": datetime.now().isoformat(),
                         }
                     }
                     raise
@@ -545,41 +662,45 @@ class LLMTracerMixin:
 
                     llm_component = self.llm_data
                     if (name is not None) or (name != ""):
-                        llm_component['name'] = name 
+                        llm_component["name"] = name
 
                     if self.gt:
                         llm_component["data"]["gt"] = self.gt
 
                     if error_info:
                         llm_component["error"] = error_info["error"]
-                    
+
                     if parent_agent_id:
                         children = self.agent_children.get()
                         children.append(llm_component)
                         self.agent_children.set(children)
                     else:
                         self.add_component(llm_component)
-                    
+
                     self.end_component(component_id)
+                    llm_component["interactions"] = self.component_user_interaction.get(
+                        component_id, []
+                    )
+                    self.add_component(llm_component)
 
             @self.file_tracker.trace_decorator
             @functools.wraps(func)
             def sync_wrapper(*args, **kwargs):
-                self.gt = kwargs.get('gt', None) if kwargs else None
+                self.gt = kwargs.get("gt", None) if kwargs else None
                 self.current_llm_call_name.set(name)
                 if not self.is_active:
                     return func(*args, **kwargs)
-                
+
                 hash_id = generate_unique_hash_simple(func)
 
                 component_id = str(uuid.uuid4())
                 parent_agent_id = self.current_agent_id.get()
                 self.start_component(component_id)
-                
+
                 start_time = datetime.now()
                 error_info = None
                 result = None
-                
+
                 try:
                     result = func(*args, **kwargs)
                     return result
@@ -589,7 +710,7 @@ class LLMTracerMixin:
                             "type": type(e).__name__,
                             "message": str(e),
                             "traceback": traceback.format_exc(),
-                            "timestamp": datetime.now().isoformat()
+                            "timestamp": datetime.now().isoformat(),
                         }
                     }
                     raise
@@ -597,21 +718,26 @@ class LLMTracerMixin:
 
                     llm_component = self.llm_data
                     if (name is not None) or (name != ""):
-                        llm_component['name'] = name 
+                        llm_component["name"] = name
 
                     if error_info:
                         llm_component["error"] = error_info["error"]
-                    
+
                     if parent_agent_id:
                         children = self.agent_children.get()
                         children.append(llm_component)
                         self.agent_children.set(children)
                     else:
                         self.add_component(llm_component)
-                    
+
                     self.end_component(component_id)
+                    llm_component["interactions"] = self.component_user_interaction.get(
+                        component_id, []
+                    )
+                    self.add_component(llm_component)
 
             return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
+
         return decorator
 
     def unpatch_llm_calls(self):