ragaai-catalyst 2.0.7.2b0__py3-none-any.whl → 2.1__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/llm_tracer.py

@@ -1,42 +1,41 @@
 from typing import Optional, Any, Dict, List
 import asyncio
 import psutil
-import json
 import wrapt
 import functools
 from datetime import datetime
 import uuid
-import os
 import contextvars
-import sys
-import gc
+import traceback
 
-from .unique_decorator import mydecorator
-from .utils.trace_utils import calculate_cost, load_model_costs
+from .unique_decorator import generate_unique_hash_simple   
+from .utils.trace_utils import load_model_costs
 from .utils.llm_utils import extract_llm_output
+from .file_name_tracker import TrackName
 
 
 class LLMTracerMixin:
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.file_tracker = TrackName()
         self.patches = []
         try:
             self.model_costs = load_model_costs()
         except Exception as e:
-            # If model costs can't be loaded, use default costs
             self.model_costs = {
+                # TODO: Default cost handling needs to be improved
                 "default": {
-                    "input_cost_per_token": 0.00002,
-                    "output_cost_per_token": 0.00002
+                    "input_cost_per_token": 0.0,
+                    "output_cost_per_token": 0.0
                 }
             }
         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.total_tokens = 0
         self.total_cost = 0.0
-        # Apply decorator to trace_llm_call method
-        self.trace_llm_call = mydecorator(self.trace_llm_call)
+        self.llm_data = {}
 
     def instrument_llm_calls(self):
         # Handle modules that are already imported
@@ -238,7 +237,7 @@ class LLMTracerMixin:
             setattr(obj, method_name, wrapped_method)
             self.patches.append((obj, method_name, original_method))
 
-    def _extract_model_name(self, kwargs):
+    def _extract_model_name(self, args, kwargs, result):
         """Extract model name from kwargs or result"""
         # First try direct model parameter
         model = kwargs.get("model", "")
@@ -254,6 +253,7 @@ class LLMTracerMixin:
                 elif hasattr(instance, "model"):
                     model = instance.model
         
+        # TODO: This way isn't scalable. The necessity for normalising model names needs to be fixed. We shouldn't have to do this
         # Normalize Google model names
         if model and isinstance(model, str):
             model = model.lower()
@@ -263,33 +263,48 @@ class LLMTracerMixin:
                 return "gemini-1.5-pro"
             if "gemini-pro" in model:
                 return "gemini-pro"
+
+        if 'to_dict' in dir(result):
+            result = result.to_dict()
+            if 'model_version' in result:
+                model = result['model_version']
         
         return model or "default"
 
-    def _extract_parameters(self, kwargs, result=None):
-        """Extract parameters from kwargs or result"""
-        params = {
-            "temperature": kwargs.get("temperature", getattr(result, "temperature", 0.7)),
-            "top_p": kwargs.get("top_p", getattr(result, "top_p", 1.0)),
-            "max_tokens": kwargs.get("max_tokens", getattr(result, "max_tokens", 512))
-        }
-        
-        # Add Google AI specific parameters if available
-        if hasattr(kwargs.get("self", None), "generation_config"):
-            gen_config = kwargs["self"].generation_config
-            params.update({
-                "candidate_count": getattr(gen_config, "candidate_count", 1),
-                "stop_sequences": getattr(gen_config, "stop_sequences", []),
-                "top_k": getattr(gen_config, "top_k", 40)
-            })
-        
-        return params
+    def _extract_parameters(self, kwargs):
+        """Extract all non-null parameters from kwargs"""
+        parameters = {k: v for k, v in kwargs.items() if v is not None}
+
+        # Remove contents key in parameters (Google LLM Response)
+        if 'contents' in parameters:
+            del parameters['contents']
+
+        # Remove messages key in parameters (OpenAI message)
+        if 'messages' in parameters:
+            del parameters['messages']
+
+        if 'generation_config' in parameters:
+            generation_config = parameters['generation_config']
+            # If generation_config is already a dict, use it directly
+            if isinstance(generation_config, dict):
+                config_dict = generation_config
+            else:
+                # Convert GenerationConfig to dictionary if it has a to_dict method, otherwise try to get its __dict__
+                config_dict = getattr(generation_config, 'to_dict', lambda: generation_config.__dict__)()
+            parameters.update(config_dict)
+            del parameters['generation_config']
+            
+        return parameters
 
     def _extract_token_usage(self, result):
         """Extract token usage from result"""
         # Handle coroutines
         if asyncio.iscoroutine(result):
-            result = asyncio.run(result)
+            # Get the current event loop
+            loop = asyncio.get_event_loop()
+            # Run the coroutine in the current event loop
+            result = loop.run_until_complete(result)
+
 
         # Handle standard OpenAI/Anthropic format
         if hasattr(result, "usage"):
@@ -317,79 +332,26 @@ class LLMTracerMixin:
                 # Try to get from raw response
                 total_tokens = getattr(result._raw_response, "token_count", 0)
             return {
+                # TODO: This implementation is incorrect. Vertex AI does provide this breakdown    
                 "prompt_tokens": 0,  # Vertex AI doesn't provide this breakdown
                 "completion_tokens": total_tokens,
                 "total_tokens": total_tokens
             }
         
-        return {
+        return {    # TODO: Passing 0 in case of not recorded is not correct. This needs to be fixes. Discuss before making changes to this
             "prompt_tokens": 0,
             "completion_tokens": 0,
             "total_tokens": 0
         }
 
-    def _extract_input_data(self, kwargs, result):
-        """Extract input data from kwargs and result"""
-
-        # For Vertex AI GenerationResponse
-        if hasattr(result, 'candidates') and hasattr(result, 'usage_metadata'):
-            # Extract generation config
-            generation_config = kwargs.get('generation_config', {})
-            config_dict = {}
-            if hasattr(generation_config, 'temperature'):
-                config_dict['temperature'] = generation_config.temperature
-            if hasattr(generation_config, 'top_p'):
-                config_dict['top_p'] = generation_config.top_p
-            if hasattr(generation_config, 'max_output_tokens'):
-                config_dict['max_tokens'] = generation_config.max_output_tokens
-            if hasattr(generation_config, 'candidate_count'):
-                config_dict['n'] = generation_config.candidate_count
-
-            return {
-                "prompt": kwargs.get('contents', ''),
-                "model": "gemini-1.5-flash-002",  
-                **config_dict
-            }
-
-        # For standard OpenAI format
-        messages = kwargs.get("messages", [])
-        if messages:
-            return {
-                "messages": messages,
-                "model": kwargs.get("model", "unknown"),
-                "temperature": kwargs.get("temperature", 0.7),
-                "max_tokens": kwargs.get("max_tokens", None),
-                "top_p": kwargs.get("top_p", None),
-                "frequency_penalty": kwargs.get("frequency_penalty", None),
-                "presence_penalty": kwargs.get("presence_penalty", None)
-            }
-
-        # For text completion format
-        if "prompt" in kwargs:
-            return {
-                "prompt": kwargs["prompt"],
-                "model": kwargs.get("model", "unknown"),
-                "temperature": kwargs.get("temperature", 0.7),
-                "max_tokens": kwargs.get("max_tokens", None),
-                "top_p": kwargs.get("top_p", None),
-                "frequency_penalty": kwargs.get("frequency_penalty", None),
-                "presence_penalty": kwargs.get("presence_penalty", None)
-            }
-
-        # For any other case, try to extract from kwargs
-        if "contents" in kwargs:
-            return {
-                "prompt": kwargs["contents"],
-                "model": kwargs.get("model", "unknown"),
-                "temperature": kwargs.get("temperature", 0.7),
-                "max_tokens": kwargs.get("max_tokens", None),
-                "top_p": kwargs.get("top_p", None)
-            }
-
-        print("No input data found")
-        return {}
+    def _extract_input_data(self, args, kwargs, result):
+        return {
+            'args': args,
+            'kwargs': kwargs
+        }
 
     def _calculate_cost(self, token_usage, model_name):
+        # TODO: Passing default cost is a faulty logic & implementation and should be fixed
         """Calculate cost based on token usage and model"""
         if not isinstance(token_usage, dict):
             token_usage = {
@@ -398,74 +360,60 @@ class LLMTracerMixin:
                 "total_tokens": token_usage if isinstance(token_usage, (int, float)) else 0
             }
 
+        # TODO: This is a temporary fix. This needs to be fixed
+
         # Get model costs, defaulting to Vertex AI PaLM2 costs if unknown
         model_cost = self.model_costs.get(model_name, {
-            "input_cost_per_token": 0.0005,  # $0.0005 per 1K input tokens
-            "output_cost_per_token": 0.0005  # $0.0005 per 1K output tokens
+            "input_cost_per_token": 0.0,   
+            "output_cost_per_token": 0.0   
         })
 
-        # Calculate costs per 1K tokens
-        input_cost = (token_usage.get("prompt_tokens", 0) / 1000.0) * model_cost.get("input_cost_per_token", 0.0005)
-        output_cost = (token_usage.get("completion_tokens", 0) / 1000.0) * model_cost.get("output_cost_per_token", 0.0005)
+        input_cost = (token_usage.get("prompt_tokens", 0)) * model_cost.get("input_cost_per_token", 0.0)
+        output_cost = (token_usage.get("completion_tokens", 0)) * model_cost.get("output_cost_per_token", 0.0)
         total_cost = input_cost + output_cost
 
+        # TODO: Return the value as it is, no need to round
         return {
-            "input_cost": round(input_cost, 6),
-            "output_cost": round(output_cost, 6),
-            "total_cost": round(total_cost, 6)
+            "input_cost": round(input_cost, 10),
+            "output_cost": round(output_cost, 10),
+            "total_cost": round(total_cost, 10)
         }
 
-    def create_llm_component(self, **kwargs):
-        """Create an LLM component according to the data structure"""
-        start_time = kwargs["start_time"]
-        
-        # Ensure cost and usage are dictionaries
-        cost = kwargs.get("cost", {})
-        if not isinstance(cost, dict):
-            cost = {"total_cost": cost}
-            
-        usage = kwargs.get("usage", {})
-        if not isinstance(usage, dict):
-            usage = {"total_tokens": usage}
-            
+    def create_llm_component(self, component_id, hash_id, name, llm_type, version, memory_used, start_time, end_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)
+
         component = {
-            "id": kwargs["component_id"],
-            "hash_id": kwargs["hash_id"],
+            "id": component_id,
+            "hash_id": hash_id,
             "source_hash_id": None,
             "type": "llm",
-            "name": kwargs["name"],
+            "name": name,
             "start_time": start_time.isoformat(),
-            "end_time": kwargs["end_time"].isoformat(),
-            "error": kwargs.get("error"),
+            "end_time": end_time.isoformat(),
+            "error": error,
             "parent_id": self.current_agent_id.get(),
             "info": {
-                "llm_type": kwargs.get("llm_type", "unknown"),
-                "version": kwargs.get("version", "1.0.0"),
-                "memory_used": kwargs.get("memory_used", 0),
+                "model": llm_type,
+                "version": version,
+                "memory_used": memory_used,
                 "cost": cost,
-                "tokens": usage
+                "tokens": usage,
+                **parameters
             },
             "data": {
-                "input": kwargs.get("input_data"),
-                "output": kwargs.get("output_data"),
-                "memory_used": kwargs.get("memory_used", 0)
+                "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
             },
-            "network_calls": self.component_network_calls.get(kwargs["component_id"], []),
-            "interactions": [
-                {
-                    "id": f"int_{uuid.uuid4()}",
-                    "interaction_type": "input",
-                    "timestamp": start_time.isoformat(),
-                    "content": kwargs.get("input_data")
-                },
-                {
-                    "id": f"int_{uuid.uuid4()}",
-                    "interaction_type": "output",
-                    "timestamp": kwargs["end_time"].isoformat(),
-                    "content": kwargs.get("output_data")
-                }
-            ]
+            "network_calls": self.component_network_calls.get(component_id, []),
+            "interactions": self.component_user_interaction.get(component_id, [])
         }
+
+        if self.gt: 
+            component["data"]["gt"] = self.gt
+
         return component
     
     def start_component(self, component_id):
@@ -481,29 +429,19 @@ class LLMTracerMixin:
     async def trace_llm_call(self, original_func, *args, **kwargs):
         """Trace an LLM API call"""
         if not self.is_active:
-            if asyncio.iscoroutinefunction(original_func):
-                return await original_func(*args, **kwargs)
-            return original_func(*args, **kwargs)
+            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 = self.trace_llm_call.hash_id
+        hash_id = generate_unique_hash_simple(original_func) 
 
         # Start tracking network calls for this component
         self.start_component(component_id)
 
         try:
             # Execute the LLM call
-            result = None
-            if asyncio.iscoroutinefunction(original_func):
-                result = await original_func(*args, **kwargs)
-            else:
-                result = original_func(*args, **kwargs)
-
-            # If result is a coroutine, await it
-            if asyncio.iscoroutine(result):
-                result = await result
+            result = await original_func(*args, **kwargs)
 
             # Calculate resource usage
             end_time = datetime.now().astimezone()
@@ -511,30 +449,40 @@ class LLMTracerMixin:
             memory_used = max(0, end_memory - start_memory)
 
             # Extract token usage and calculate cost
-            token_usage = await self._extract_token_usage(result)
-            model_name = self._extract_model_name(kwargs)
+            token_usage = self._extract_token_usage(result)
+            model_name = self._extract_model_name(args, kwargs, result)
             cost = self._calculate_cost(token_usage, model_name)
+            parameters = self._extract_parameters(kwargs)
 
             # 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__
+            
+            # Create input data with ground truth
+            input_data = self._extract_input_data(args, kwargs, result)
+            
             # Create LLM component
             llm_component = self.create_llm_component(
                 component_id=component_id,
                 hash_id=hash_id,
-                name=self.current_llm_call_name.get(),
+                name=name,
                 llm_type=model_name,
                 version="1.0.0",
                 memory_used=memory_used,
                 start_time=start_time,
                 end_time=end_time,
-                input_data=self._extract_input_data(kwargs, result),
+                input_data=input_data,
                 output_data=extract_llm_output(result),
                 cost=cost,
-                usage=token_usage
+                usage=token_usage,
+                parameters=parameters
             )
-
-            self.add_component(llm_component)
+                
+            # self.add_component(llm_component)
+            self.llm_data = llm_component
             return result
 
         except Exception as e:
@@ -549,67 +497,28 @@ class LLMTracerMixin:
             self.end_component(component_id)
             
             end_time = datetime.now().astimezone()
+
+            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=self.current_llm_call_name.get(),
+                name=name,
                 llm_type="unknown",
                 version="1.0.0",
                 memory_used=0,
                 start_time=start_time,
                 end_time=end_time,
-                input_data=self._extract_input_data(kwargs, None),
+                input_data=self._extract_input_data(args, kwargs, None),
                 output_data=None,
                 error=error_component
             )
-
+    
             self.add_component(llm_component)
             raise
 
-    def _extract_token_usage_sync(self, result):
-        """Sync version of extract token usage"""
-        # Handle coroutines
-        if asyncio.iscoroutine(result):
-            result = asyncio.run(result)
-
-        # Handle standard OpenAI/Anthropic format
-        if hasattr(result, "usage"):
-            usage = result.usage
-            return {
-                "prompt_tokens": getattr(usage, "prompt_tokens", 0),
-                "completion_tokens": getattr(usage, "completion_tokens", 0),
-                "total_tokens": getattr(usage, "total_tokens", 0)
-            }
-        
-        # Handle Google GenerativeAI format with usage_metadata
-        if hasattr(result, "usage_metadata"):
-            metadata = result.usage_metadata
-            return {
-                "prompt_tokens": getattr(metadata, "prompt_token_count", 0),
-                "completion_tokens": getattr(metadata, "candidates_token_count", 0),
-                "total_tokens": getattr(metadata, "total_token_count", 0)
-            }
-        
-        # Handle Vertex AI format
-        if hasattr(result, "text"):
-            # For LangChain ChatVertexAI
-            total_tokens = getattr(result, "token_count", 0)
-            if not total_tokens and hasattr(result, "_raw_response"):
-                # Try to get from raw response
-                total_tokens = getattr(result._raw_response, "token_count", 0)
-            return {
-                "prompt_tokens": 0,  # Vertex AI doesn't provide this breakdown
-                "completion_tokens": total_tokens,
-                "total_tokens": total_tokens
-            }
-        
-        return {
-            "prompt_tokens": 0,
-            "completion_tokens": 0,
-            "total_tokens": 0
-        }
-
     def trace_llm_call_sync(self, original_func, *args, **kwargs):
         """Sync version of trace_llm_call"""
         if not self.is_active:
@@ -618,54 +527,59 @@ class LLMTracerMixin:
             return original_func(*args, **kwargs)
 
         start_time = datetime.now().astimezone()
-        start_memory = psutil.Process().memory_info().rss
         component_id = str(uuid.uuid4())
-        hash_id = self.trace_llm_call.hash_id
+        hash_id = generate_unique_hash_simple(original_func)
 
         # Start tracking network calls for this component
         self.start_component(component_id)
 
+        # Calculate resource usage
+        end_time = datetime.now().astimezone()
+        start_memory = psutil.Process().memory_info().rss
+
         try:
-            # Execute the LLM call
-            result = None
+            # Execute the function
             if asyncio.iscoroutinefunction(original_func):
                 result = asyncio.run(original_func(*args, **kwargs))
             else:
                 result = original_func(*args, **kwargs)
 
-            # If result is a coroutine, run it
-            if asyncio.iscoroutine(result):
-                result = asyncio.run(result)
-
-            # Calculate resource usage
-            end_time = datetime.now().astimezone()
             end_memory = psutil.Process().memory_info().rss
             memory_used = max(0, end_memory - start_memory)
 
             # Extract token usage and calculate cost
-            token_usage = self._extract_token_usage_sync(result)
-            model_name = self._extract_model_name(kwargs)
+            token_usage = self._extract_token_usage(result)
+            model_name = self._extract_model_name(args, kwargs, result)
             cost = self._calculate_cost(token_usage, model_name)
+            parameters = self._extract_parameters(kwargs)
 
             # 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__
+
+            # Create input data with ground truth
+            input_data = self._extract_input_data(args, kwargs, result)
+            
             # Create LLM component
             llm_component = self.create_llm_component(
                 component_id=component_id,
                 hash_id=hash_id,
-                name=self.current_llm_call_name.get(),
+                name=name,
                 llm_type=model_name,
                 version="1.0.0",
                 memory_used=memory_used,
                 start_time=start_time,
                 end_time=end_time,
-                input_data=self._extract_input_data(kwargs, result),
+                input_data=input_data,
                 output_data=extract_llm_output(result),
                 cost=cost,
-                usage=token_usage
+                usage=token_usage,
+                parameters=parameters
             )
-
+            
             self.add_component(llm_component)
             return result
 
@@ -681,105 +595,139 @@ class LLMTracerMixin:
             self.end_component(component_id)
             
             end_time = datetime.now().astimezone()
+
+            name = self.current_llm_call_name.get()
+            if name is None:
+                name = original_func.__name__
+
+            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=self.current_llm_call_name.get(),
+                name=name,
                 llm_type="unknown",
                 version="1.0.0",
-                memory_used=0,
+                memory_used=memory_used,
                 start_time=start_time,
                 end_time=end_time,
-                input_data=self._extract_input_data(kwargs, None),
+                input_data=self._extract_input_data(args, kwargs, None),
                 output_data=None,
                 error=error_component
             )
-
+    
             self.add_component(llm_component)
             raise
 
-    async def _extract_token_usage(self, result):
-        """Extract token usage from result"""
-        # Handle coroutines
-        if asyncio.iscoroutine(result):
-            result = await result
-
-        # Handle standard OpenAI/Anthropic format
-        if hasattr(result, "usage"):
-            usage = result.usage
-            return {
-                "prompt_tokens": getattr(usage, "prompt_tokens", 0),
-                "completion_tokens": getattr(usage, "completion_tokens", 0),
-                "total_tokens": getattr(usage, "total_tokens", 0)
-            }
-        
-        # Handle Google GenerativeAI format with usage_metadata
-        if hasattr(result, "usage_metadata"):
-            metadata = result.usage_metadata
-            return {
-                "prompt_tokens": getattr(metadata, "prompt_token_count", 0),
-                "completion_tokens": getattr(metadata, "candidates_token_count", 0),
-                "total_tokens": getattr(metadata, "total_token_count", 0)
-            }
-        
-        # Handle Vertex AI format
-        if hasattr(result, "text"):
-            # For LangChain ChatVertexAI
-            total_tokens = getattr(result, "token_count", 0)
-            if not total_tokens and hasattr(result, "_raw_response"):
-                # Try to get from raw response
-                total_tokens = getattr(result._raw_response, "token_count", 0)
-            return {
-                "prompt_tokens": 0,  # Vertex AI doesn't provide this breakdown
-                "completion_tokens": total_tokens,
-                "total_tokens": total_tokens
-            }
-        
-        return {
-            "prompt_tokens": 0,
-            "completion_tokens": 0,
-            "total_tokens": 0
-        }
-
-    def trace_llm(self, name: str, tool_type: str = "llm", version: str = "1.0.0"):
-        def decorator(func_or_class):
-            if isinstance(func_or_class, type):
-                for attr_name, attr_value in func_or_class.__dict__.items():
-                    if callable(attr_value) and not attr_name.startswith("__"):
-                        setattr(
-                            func_or_class,
-                            attr_name,
-                            self.trace_llm(f"{name}.{attr_name}", tool_type, version)(attr_value),
-                        )
-                return func_or_class
-            else:
-                @functools.wraps(func_or_class)
-                async def async_wrapper(*args, **kwargs):
-                    token = self.current_llm_call_name.set(name)
-                    try:
-                        return await func_or_class(*args, **kwargs)
-                    finally:
-                        self.current_llm_call_name.reset(token)
-
-                @functools.wraps(func_or_class)
-                def sync_wrapper(*args, **kwargs):
-                    token = self.current_llm_call_name.set(name)
-                    try:
-                        return func_or_class(*args, **kwargs)
-                    finally:
-                        self.current_llm_call_name.reset(token)
-
-                return async_wrapper if asyncio.iscoroutinefunction(func_or_class) else sync_wrapper
-
+    def trace_llm(self, name: str = None):
+        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
+                if not self.is_active:
+                    return await 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 = await func(*args, **kwargs)
+                    return result
+                except Exception as e:
+                    error_info = {
+                        "error": {
+                            "type": type(e).__name__,
+                            "message": str(e),
+                            "traceback": traceback.format_exc(),
+                            "timestamp": datetime.now().isoformat()
+                        }
+                    }
+                    raise
+                finally:
+
+                    llm_component = self.llm_data
+                    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)
+
+            @self.file_tracker.trace_decorator
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                self.gt = kwargs.get('gt', None) if kwargs else None
+                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
+                except Exception as e:
+                    error_info = {
+                        "error": {
+                            "type": type(e).__name__,
+                            "message": str(e),
+                            "traceback": traceback.format_exc(),
+                            "timestamp": datetime.now().isoformat()
+                        }
+                    }
+                    raise
+                finally:
+
+                    llm_component = self.llm_data
+
+                    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)
+
+            return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper
         return decorator
 
     def unpatch_llm_calls(self):
-        """Remove all patches"""
+        # Remove all patches
         for obj, method_name, original_method in self.patches:
-            if hasattr(obj, method_name):
+            try:
                 setattr(obj, method_name, original_method)
-        self.patches.clear()
+            except Exception as e:
+                print(f"Error unpatching {method_name}: {str(e)}")
+        self.patches = []
 
     def _sanitize_api_keys(self, data):
         """Remove sensitive information from data"""
@@ -792,62 +740,20 @@ class LLMTracerMixin:
             return tuple(self._sanitize_api_keys(item) for item in data)
         return data
 
-    def _create_llm_component(self, component_id, hash_id, name, llm_type, version, memory_used, start_time, end_time, input_data, output_data, usage=None, error=None):
-        cost = None
-        tokens = None
+    def _sanitize_input(self, args, kwargs):
+        """Convert input arguments to text format.
         
-        if usage:
-            tokens = {
-                "prompt_tokens": usage.get("prompt_tokens", 0),
-                "completion_tokens": usage.get("completion_tokens", 0),
-                "total_tokens": usage.get("prompt_tokens", 0) + usage.get("completion_tokens", 0)
-            }
-            cost = calculate_cost(usage)
+        Args:
+            args: Input arguments that may contain nested dictionaries
             
-            # Update total metrics
-            self.total_tokens += tokens["total_tokens"]
-            self.total_cost += cost["total"]
-
-        component = {
-            "id": component_id,
-            "hash_id": hash_id,
-            "source_hash_id": None,
-            "type": "llm",
-            "name": name,
-            "start_time": start_time.isoformat(),
-            "end_time": end_time.isoformat(),
-            "error": error,
-            "parent_id": self.current_agent_id.get(),
-            "info": {
-                "llm_type": llm_type,
-                "version": version,
-                "memory_used": memory_used,
-                "cost": cost,
-                "tokens": tokens
-            },
-            "data": {
-                "input": input_data,
-                "output": output_data.output_response if output_data else None,
-                "memory_used": memory_used
-            },
-            "network_calls": self.component_network_calls.get(component_id, []),
-            "interactions": [
-                {
-                    "id": f"int_{uuid.uuid4()}",
-                    "interaction_type": "input",
-                    "timestamp": start_time.isoformat(),
-                    "content": input_data
-                },
-                {
-                    "id": f"int_{uuid.uuid4()}",
-                    "interaction_type": "output",
-                    "timestamp": end_time.isoformat(),
-                    "content": output_data.output_response if output_data else None
-                }
-            ]
-        }
-
-        return component
+        Returns:
+            str: Text representation of the input arguments
+        """
+        if isinstance(args, dict):
+            return str({k: self._sanitize_input(v, {}) for k, v in args.items()})
+        elif isinstance(args, (list, tuple)):
+            return str([self._sanitize_input(item, {}) for item in args])
+        return str(args)
 
 def extract_llm_output(result):
     """Extract output from LLM response"""