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

ragaai_catalyst/synthetic_data_generation.py

@@ -8,7 +8,9 @@ import markdown
 import pandas as pd
 import json
 from litellm import completion
+import litellm
 from tqdm import tqdm
+import tiktoken
 # import internal_api_completion
 # import proxy_call
 from .internal_api_completion import api_completion as internal_api_completion
@@ -48,13 +50,18 @@ class SyntheticDataGeneration:
         Raises:
             ValueError: If an invalid provider is specified or API key is missing.
         """
+        text_validity = self.validate_input(text)
+        if text_validity:
+            raise ValueError(text_validity)
+
         BATCH_SIZE = 5  # Optimal batch size for maintaining response quality
         provider = model_config.get("provider")
         model = model_config.get("model")
         api_base = model_config.get("api_base")
+        api_version = model_config.get("api_version")
 
         # Initialize the appropriate client based on provider
-        self._initialize_client(provider, api_key, api_base, internal_llm_proxy=kwargs.get("internal_llm_proxy", None))
+        self._initialize_client(provider, api_key, api_base, api_version, internal_llm_proxy=kwargs.get("internal_llm_proxy", None))
 
         # Initialize progress bar
         pbar = tqdm(total=n, desc="Generating QA pairs")
@@ -88,7 +95,7 @@ class SyntheticDataGeneration:
                     pbar.update(len(batch_df))
                     
             except Exception as e:
-                print(f"Batch generation failed.")
+                print(f"Batch generation failed:{str(e)}")
 
                 if any(error in str(e) for error in FAILURE_CASES):
                     raise Exception(f"{e}")
@@ -139,7 +146,7 @@ class SyntheticDataGeneration:
         
         return final_df
 
-    def _initialize_client(self, provider, api_key, api_base=None, internal_llm_proxy=None):
+    def _initialize_client(self, provider, api_key, api_base=None, api_version=None, internal_llm_proxy=None):
         """Initialize the appropriate client based on provider."""
         if not provider:
             raise ValueError("Model configuration must be provided with a valid provider and model.")
@@ -158,7 +165,17 @@ class SyntheticDataGeneration:
             if api_key is None and os.getenv("OPENAI_API_KEY") is None and internal_llm_proxy is None:
                 raise ValueError("API key must be provided for OpenAI.")
             openai.api_key = api_key or os.getenv("OPENAI_API_KEY")
-        
+
+        elif provider == "azure":
+            if api_key is None and os.getenv("AZURE_API_KEY") is None and internal_llm_proxy is None:
+                raise ValueError("API key must be provided for Azure.")
+            litellm.api_key = api_key or os.getenv("AZURE_API_KEY")
+            if api_base is None and os.getenv("AZURE_API_BASE") is None and internal_llm_proxy is None:
+                raise ValueError("API Base must be provided for Azure.")
+            litellm.api_base = api_base or os.getenv("AZURE_API_BASE")
+            if api_version is None and os.getenv("AZURE_API_VERSION") is None and internal_llm_proxy is None:
+                raise ValueError("API version must be provided for Azure.")
+            litellm.api_version = api_version or os.getenv("AZURE_API_VERSION")
         else:
             raise ValueError(f"Provider is not recognized.")
 
@@ -189,7 +206,15 @@ class SyntheticDataGeneration:
             kwargs=kwargs
         )
 
+    def validate_input(self,text):
 
+        if not text.strip():
+            return 'Empty Text provided for qna generation. Please provide valid text'
+        encoding = tiktoken.encoding_for_model("gpt-4")
+        tokens = encoding.encode(text)
+        if len(tokens)<5:
+            return 'Very Small Text provided for qna generation. Please provide longer text'
+        return False
         
           
     def _get_system_message(self, question_type, n):
@@ -274,10 +299,14 @@ class SyntheticDataGeneration:
         # Add optional parameters if they exist in model_config
         if "api_base" in model_config:
             completion_params["api_base"] = model_config["api_base"]
+        if "api_version" in model_config:
+            completion_params["api_version"] = model_config["api_version"]
         if "max_tokens" in model_config:
             completion_params["max_tokens"] = model_config["max_tokens"]
         if "temperature" in model_config:
             completion_params["temperature"] = model_config["temperature"]
+        if 'provider' in model_config:
+            completion_params['model'] = f'{model_config["provider"]}/{model_config["model"]}'
 
         # Make the API call using LiteLLM
         try:
@@ -318,9 +347,13 @@ class SyntheticDataGeneration:
             list_start_index = data.find('[')  # Find the index of the first '['
             substring_data = data[list_start_index:] if list_start_index != -1 else data  # Slice from the list start
             data = substring_data
-
+        elif provider == "azure":
+            data = response.choices[0].message.content.replace('\n', '')
+            list_start_index = data.find('[')  # Find the index of the first '['
+            substring_data = data[list_start_index:] if list_start_index != -1 else data  # Slice from the list start
+            data = substring_data
         else:
-            raise ValueError("Invalid provider. Choose 'groq', 'gemini', or 'openai'.")
+            raise ValueError("Invalid provider. Choose 'groq', 'gemini', 'azure' or 'openai'.")
         try:
             json_data = json.loads(data)
             return pd.DataFrame(json_data)

ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py

@@ -101,7 +101,10 @@ class AgentTracerMixin:
                 original_init = target.__init__
 
                 def wrapped_init(self, *args, **kwargs):
-                    self.gt = kwargs.get("gt", None) if kwargs else None
+                    gt = kwargs.get("gt") if kwargs else None
+                    if gt is not None:
+                        span = self.span(name)
+                        span.add_gt(gt)
                     # Set agent context before initializing
                     component_id = str(uuid.uuid4())
                     hash_id = top_level_hash_id
@@ -159,7 +162,10 @@ class AgentTracerMixin:
                                 @self.file_tracker.trace_decorator
                                 @functools.wraps(method)
                                 def wrapped_method(self, *args, **kwargs):
-                                    self.gt = kwargs.get("gt", None) if kwargs else None
+                                    gt = kwargs.get("gt") if kwargs else None
+                                    if gt is not None:
+                                        span = tracer.span(name)
+                                        span.add_gt(gt)
                                     # Set this agent as current during method execution
                                     token = tracer.current_agent_id.set(
                                         self._agent_component_id
@@ -247,6 +253,7 @@ class AgentTracerMixin:
                             agent_type,
                             version,
                             capabilities,
+                            top_level_hash_id,
                             *args,
                             **kwargs,
                         )
@@ -256,10 +263,9 @@ class AgentTracerMixin:
         return decorator
 
     def _trace_sync_agent_execution(
-        self, func, name, agent_type, version, capabilities, *args, **kwargs
+        self, func, name, agent_type, version, capabilities, top_level_hash_id, *args, **kwargs
     ):
-        # Generate a unique hash_id for this execution context
-        hash_id = str(uuid.uuid4())
+        hash_id = top_level_hash_id
 
         """Synchronous version of agent tracing"""
         if not self.is_active:
@@ -275,7 +281,10 @@ class AgentTracerMixin:
         component_id = str(uuid.uuid4())
 
         # Extract ground truth if present
-        ground_truth = kwargs.pop("gt", None) if kwargs else None
+        ground_truth = kwargs.pop("gt") if kwargs else None
+        if ground_truth is not None:
+            span = self.span(name)
+            span.add_gt(ground_truth)
 
         # Get parent agent ID if exists
         parent_agent_id = self.current_agent_id.get()
@@ -293,7 +302,7 @@ class AgentTracerMixin:
 
         try:
             # Execute the agent
-            result = func(*args, **kwargs)
+            result = self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_memory = psutil.Process().memory_info().rss
@@ -320,9 +329,6 @@ class AgentTracerMixin:
                 children=children,
                 parent_id=parent_agent_id,
             )
-            # Add ground truth to component data if present
-            if ground_truth is not None:
-                agent_component["data"]["gt"] = ground_truth
 
             # Add this component as a child to parent's children list
             parent_children.append(agent_component)
@@ -398,7 +404,10 @@ class AgentTracerMixin:
         component_id = str(uuid.uuid4())
 
         # Extract ground truth if present
-        ground_truth = kwargs.pop("gt", None) if kwargs else None
+        ground_truth = kwargs.pop("gt") if kwargs else None
+        if ground_truth is not None:
+            span = self.span(name)
+            span.add_gt(ground_truth)
 
         # Get parent agent ID if exists
         parent_agent_id = self.current_agent_id.get()
@@ -414,7 +423,7 @@ class AgentTracerMixin:
 
         try:
             # Execute the agent
-            result = await func(*args, **kwargs)
+            result = await self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_memory = psutil.Process().memory_info().rss
@@ -441,10 +450,6 @@ class AgentTracerMixin:
                 parent_id=parent_agent_id,
             )
 
-            # Add ground truth to component data if present
-            if ground_truth is not None:
-                agent_component["data"]["gt"] = ground_truth
-
             # Add this component as a child to parent's children list
             parent_children.append(agent_component)
             self.agent_children.set(parent_children)
@@ -576,8 +581,13 @@ class AgentTracerMixin:
             "interactions": interactions,
         }
 
-        if self.gt:
-            component["data"]["gt"] = self.gt
+        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
 
         # Reset the SpanAttributes context variable
         self.span_attributes_dict[kwargs["name"]] = SpanAttributes(kwargs["name"])

ragaai_catalyst/tracers/agentic_tracing/tracers/base.py

@@ -83,6 +83,7 @@ class BaseTracer:
         self.tracking_thread = None
         self.tracking = False
         self.system_monitor = None
+        self.gt = None
 
     def _get_system_info(self) -> SystemInfo:
         return self.system_monitor.get_system_info()
@@ -249,7 +250,8 @@ class BaseTracer:
 
             # Format interactions and add to trace
             interactions = self.format_interactions()
-            trace_data["workflow"] = interactions["workflow"]
+            # trace_data["workflow"] = interactions["workflow"]
+            cleaned_trace_data["workflow"] = interactions["workflow"]
 
             with open(filepath, "w") as f:
                 json.dump(cleaned_trace_data, f, cls=TracerJSONEncoder, indent=2)

ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py

@@ -45,7 +45,10 @@ class CustomTracerMixin:
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
                 async_wrapper.metadata = metadata
-                self.gt = kwargs.get('gt', None) if kwargs else None
+                gt = kwargs.get('gt') if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return await self._trace_custom_execution(
                     func, name or func.__name__, custom_type, version, trace_variables, *args, **kwargs
                 )
@@ -54,7 +57,10 @@ class CustomTracerMixin:
             @functools.wraps(func)
             def sync_wrapper(*args, **kwargs):
                 sync_wrapper.metadata = metadata
-                self.gt = kwargs.get('gt', None) if kwargs else None
+                gt = kwargs.get('gt') if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return self._trace_sync_custom_execution(
                     func, name or func.__name__, custom_type, version, trace_variables, *args, **kwargs
                 )
@@ -98,7 +104,7 @@ class CustomTracerMixin:
 
         try:
             # Execute the function
-            result = func(*args, **kwargs)
+            result = self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_time = datetime.now().astimezone().isoformat()
@@ -186,7 +192,7 @@ class CustomTracerMixin:
 
         try:
             # Execute the function
-            result = await func(*args, **kwargs)
+            result = await self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_time = datetime.now().astimezone().isoformat()
@@ -284,9 +290,13 @@ class CustomTracerMixin:
             "interactions": interactions
         }
 
-        if self.gt:
-            component["data"]["gt"] = self.gt
-
+        if kwargs["name"] in self.span_attributes_dict:
+            span_gt = self.span_attributes_dict[kwargs["name"]].gt
+            if span_gt is not None:
+                component["data"]["gt"] = span_gt
+            span_context = self.span_attributes_dict[kwargs["name"]].context
+            if span_context:
+                component["data"]["context"] = span_context
         return component
 
     def start_component(self, component_id):

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -12,6 +12,7 @@ import contextvars
 import traceback
 import importlib
 import sys
+from litellm import model_cost
 
 from ..utils.llm_utils import (
     extract_model_name,
@@ -24,8 +25,7 @@ from ..utils.llm_utils import (
     extract_llm_output,
     num_tokens_from_messages
 )
-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
 from ..utils.file_name_tracker import TrackName
 from ..utils.span_attributes import SpanAttributes
 import logging
@@ -44,7 +44,7 @@ class LLMTracerMixin:
         self.file_tracker = TrackName()
         self.patches = []
         try:
-            self.model_costs = load_model_costs()
+            self.model_costs = model_cost
         except Exception as e:
             self.model_costs = {
                 "default": {"input_cost_per_token": 0.0, "output_cost_per_token": 0.0}
@@ -97,6 +97,10 @@ 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"])
 
         # Register hooks for future imports with availability checks
         if self.check_package_available("vertexai"):
@@ -130,6 +134,15 @@ class LLMTracerMixin:
             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"
+            )
+        if self.check_package_available("langchain_anthropic"):
+            wrapt.register_post_import_hook(
+                self.patch_langchain_anthropic_methods, "langchain_anthropic"
+            )
 
     def instrument_user_interaction_calls(self):
         """Enable user interaction instrumentation for LLM calls"""
@@ -154,6 +167,42 @@ 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"):
+                    self.wrap_langchain_openai_method(client_class, f"{client_class.__name__}.run")
+            if hasattr(module, 'AsyncChatOpenAI'):
+                if hasattr(client_class, "ainvoke"):
+                    self.wrap_langchain_openai_method(client_class, f"{client_class.__name__}.ainvoke")
+                elif hasattr(client_class, "arun"):
+                    self.wrap_langchain_openai_method(client_class, f"{client_class.__name__}.arun")
+        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'):
+                client_class = getattr(module, "ChatAnthropic")
+                if hasattr(client_class, "invoke"):
+                    self.wrap_langchain_anthropic_method(client_class, f"{client_class.__name__}.invoke")
+                if hasattr(client_class, "ainvoke"):
+                    self.wrap_langchain_anthropic_method(client_class, f"{client_class.__name__}.ainvoke")
+            if hasattr(module, 'AsyncChatAnthropic'):
+                async_client_class = getattr(module, "AsyncChatAnthropic")
+                if hasattr(async_client_class, "ainvoke"):
+                    self.wrap_langchain_anthropic_method(async_client_class, f"{async_client_class.__name__}.ainvoke")
+                if hasattr(async_client_class, "arun"):
+                    self.wrap_langchain_anthropic_method(async_client_class, f"{async_client_class.__name__}.arun")
+        except Exception as e:
+            # Log the error but continue execution
+            print(f"Warning: Failed to patch Anthropic methods: {str(e)}")
 
     def patch_openai_beta_methods(self, openai_module):
         """
@@ -293,7 +342,6 @@ 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
@@ -305,10 +353,39 @@ 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)
+
+        @functools.wraps(original_init)
+        def patched_init(*args, **kwargs):
+            # Check if this is AsyncOpenAI or OpenAI
+            is_async = "AsyncChatOpenAI" in client_class.__name__
+
+            if is_async:
+                return self.trace_llm_call(original_init, *args, **kwargs)
+            else:
+                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:
+                return self.trace_llm_call_sync(original_init, *args, **kwargs)
+
+        setattr(client_class, method_name, patched_init)
 
     def wrap_anthropic_client_methods(self, client_class):
         original_init = client_class.__init__
@@ -475,8 +552,13 @@ class LLMTracerMixin:
             "interactions": interactions,
         }
 
-        if self.gt:
-            component["data"]["gt"] = self.gt
+        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
 
         # Reset the SpanAttributes context variable
         self.span_attributes_dict[name] = SpanAttributes(name)
@@ -503,7 +585,7 @@ class LLMTracerMixin:
         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(original_func, args, kwargs)
 
         # Start tracking network calls for this component
         self.start_component(component_id)
@@ -605,7 +687,7 @@ class LLMTracerMixin:
 
         start_time = datetime.now().astimezone()
         component_id = str(uuid.uuid4())
-        hash_id = generate_unique_hash_simple(original_func)
+        hash_id = generate_unique_hash(original_func, args, kwargs)
 
         # Start tracking network calls for this component
         self.start_component(component_id)
@@ -745,12 +827,14 @@ class LLMTracerMixin:
             @self.file_tracker.trace_decorator
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
-                self.gt = kwargs.get("gt", None) if kwargs else None
+                gt = kwargs.get("gt") if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 self.current_llm_call_name.set(name)
                 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)
@@ -777,8 +861,13 @@ class LLMTracerMixin:
                     if (name is not None) or (name != ""):
                         llm_component["name"] = name
 
-                    if self.gt:
-                        llm_component["data"]["gt"] = self.gt
+                    if name in self.span_attributes_dict:
+                        span_gt = self.span_attributes_dict[name].gt
+                        if span_gt is not None:
+                            llm_component["data"]["gt"] = span_gt
+                        span_context = self.span_attributes_dict[name].context
+                        if span_context:
+                            llm_component["data"]["context"] = span_context
 
                     if error_info:
                         llm_component["error"] = error_info["error"]
@@ -811,13 +900,14 @@ class LLMTracerMixin:
             @self.file_tracker.trace_decorator
             @functools.wraps(func)
             def sync_wrapper(*args, **kwargs):
-                self.gt = kwargs.get("gt", None) if kwargs else None
+                gt = kwargs.get("gt") if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 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)

ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py

@@ -258,7 +258,10 @@ class ToolTracerMixin:
             @functools.wraps(func)
             async def async_wrapper(*args, **kwargs):
                 async_wrapper.metadata = metadata
-                self.gt = kwargs.get("gt", None) if kwargs else None
+                gt = kwargs.get("gt") if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return await self._trace_tool_execution(
                     func, name, tool_type, version, *args, **kwargs
                 )
@@ -267,7 +270,10 @@ class ToolTracerMixin:
             @functools.wraps(func)
             def sync_wrapper(*args, **kwargs):
                 sync_wrapper.metadata = metadata
-                self.gt = kwargs.get("gt", None) if kwargs else None
+                gt = kwargs.get("gt") if kwargs else None
+                if gt is not None:
+                    span = self.span(name)
+                    span.add_gt(gt)
                 return self._trace_sync_tool_execution(
                     func, name, tool_type, version, *args, **kwargs
                 )
@@ -302,7 +308,7 @@ class ToolTracerMixin:
 
         try:
             # Execute the tool
-            result = func(*args, **kwargs)
+            result = self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_memory = psutil.Process().memory_info().rss
@@ -384,7 +390,7 @@ class ToolTracerMixin:
         self.start_component(component_id)
         try:
             # Execute the tool
-            result = await func(*args, **kwargs)
+            result = await self.file_tracker.trace_wrapper(func)(*args, **kwargs)
 
             # Calculate resource usage
             end_memory = psutil.Process().memory_info().rss
@@ -504,8 +510,13 @@ class ToolTracerMixin:
             "interactions": interactions,
         }
 
-        if self.gt:
-            component["data"]["gt"] = self.gt
+        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
 
         # Reset the SpanAttributes context variable
         self.span_attributes_dict[kwargs["name"]] = SpanAttributes(kwargs["name"])

ragaai_catalyst/tracers/agentic_tracing/upload/upload_trace_metric.py

@@ -8,6 +8,7 @@ def upload_trace_metric(json_file_path, dataset_name, project_name):
     try:
         with open(json_file_path, "r") as f:
             traces = json.load(f)
+                    
         metrics = get_trace_metrics_from_trace(traces)
         metrics = _change_metrics_format_for_payload(metrics)
 
@@ -21,7 +22,6 @@ def upload_trace_metric(json_file_path, dataset_name, project_name):
                     metricConfig = next((user_metric["metricConfig"] for user_metric in user_trace_metrics if user_metric["displayName"] == metric["displayName"]), None)
                     if not metricConfig or metricConfig.get("Metric Source", {}).get("value") != "user":
                         raise ValueError(f"Metrics {metric['displayName']} already exist in dataset {dataset_name} of project {project_name}.")
-
         headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
@@ -68,13 +68,14 @@ def get_trace_metrics_from_trace(traces):
     # get span level metrics
     for span in traces["data"][0]["spans"]:
         if span["type"] == "agent":
+            # Add children metrics of agent
             children_metric = _get_children_metrics_of_agent(span["data"]["children"])
             if children_metric:
                 metrics.extend(children_metric)
-        else:
-            metric = span.get("metrics", [])
-            if metric:
-                metrics.extend(metric)
+
+        metric = span.get("metrics", [])
+        if metric:
+            metrics.extend(metric)
     return metrics
 
 def _change_metrics_format_for_payload(metrics):

ragaai_catalyst/tracers/agentic_tracing/utils/file_name_tracker.py

@@ -8,13 +8,32 @@ class TrackName:
     def trace_decorator(self, func):
         @wraps(func)
         def wrapper(*args, **kwargs):
-            file_name = self._get_file_name()
+            file_name = self._get_decorated_file_name()
             self.files.add(file_name)
 
             return func(*args, **kwargs)
         return wrapper
+    
+    def trace_wrapper(self, func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            file_name = self._get_wrapped_file_name()
+            self.files.add(file_name)
+            return func(*args, **kwargs)
+        return wrapper
+    
+    def _get_wrapped_file_name(self):
+        try:
+            from IPython import get_ipython
+            if 'IPKernelApp' in get_ipython().config:
+                return self._get_notebook_name()
+        except Exception:
+            pass
+
+        frame = inspect.stack()[4]
+        return frame.filename
 
-    def _get_file_name(self):
+    def _get_decorated_file_name(self):
         # Check if running in a Jupyter notebook
         try:
             from IPython import get_ipython