ragaai-catalyst 2.1.5b25__py3-none-any.whl → 2.1.5b27__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/utils/zip_list_of_unique_files.py

@@ -8,6 +8,7 @@ import ast
 import importlib.util
 import json
 import ipynbname
+from copy import deepcopy
 
 from pathlib import Path
 from IPython import get_ipython
@@ -211,7 +212,6 @@ def comment_magic_commands(script_content: str) -> str:
 class TraceDependencyTracker:
     def __init__(self, output_dir=None):
         self.tracked_files = set()
-        self.python_imports = set()
         self.notebook_path = None
         self.colab_content = None  
         
@@ -302,7 +302,7 @@ class TraceDependencyTracker:
                     except (UnicodeDecodeError, IOError):
                         pass
 
-    def analyze_python_imports(self, filepath):
+    def analyze_python_imports(self, filepath, ignored_locations):
         try:
             with open(filepath, 'r', encoding='utf-8') as file:
                 tree = ast.parse(file.read(), filename=filepath)
@@ -315,8 +315,10 @@ class TraceDependencyTracker:
                             module_name = name.name.split('.')[0]
                     try:
                         spec = importlib.util.find_spec(module_name)
-                        if spec and spec.origin and not spec.origin.startswith(os.path.dirname(importlib.__file__)):
-                            self.python_imports.add(spec.origin)
+                        if spec and spec.origin:    
+                            if not (any(spec.origin.startswith(location) for location in ignored_locations) or (spec.origin in ['built-in', 'frozen'])):
+                                self.tracked_files.add(spec.origin)
+                                self.analyze_python_imports(spec.origin, ignored_locations)
                     except (ImportError, AttributeError):
                         pass
         except Exception as e:
@@ -332,6 +334,13 @@ class TraceDependencyTracker:
         except ImportError:
             logger.error("Error getting Catalyst location")
             return 'ragaai_catalyst'
+    
+    def should_ignore_path(self, path, main_filepaths):
+        if any(os.path.abspath(path) in os.path.abspath(main_filepath) for main_filepath in main_filepaths):
+            return False
+        if path in ['', os.path.abspath('')]:
+            return False
+        return True
 
     def create_zip(self, filepaths):
         self.track_jupyter_notebook()
@@ -349,21 +358,32 @@ class TraceDependencyTracker:
             # Get current cell content
             self.check_environment_and_save()
 
+        env_location = self.get_env_location()
+        catalyst_location = self.get_catalyst_location()
+
         # Process all files (existing code)
+        ignored_locations = [env_location, catalyst_location] + [path for path in sys.path if self.should_ignore_path(path, filepaths)]
         for filepath in filepaths:
             abs_path = os.path.abspath(filepath)
             self.track_file_access(abs_path)
             try:
-                with open(abs_path, 'r', encoding='utf-8') as file:
+                if filepath.endswith('.py'):
+                    self.analyze_python_imports(abs_path, ignored_locations)
+            except Exception as e:
+                pass
+        
+        curr_tracked_files = deepcopy(self.tracked_files)
+        for filepath in curr_tracked_files:
+            try:
+                with open(filepath, 'r', encoding='utf-8') as file:
                     content = file.read()
                     # Comment out magic commands before processing
                     content = comment_magic_commands(content)
-                self.find_config_files(content, abs_path)
-                if filepath.endswith('.py'):
-                    self.analyze_python_imports(abs_path)
+                self.find_config_files(content, filepath)
             except Exception as e:
                 pass
 
+
         notebook_content_str = None
         if self.notebook_path and os.path.exists(self.notebook_path):
             try:
@@ -386,11 +406,7 @@ class TraceDependencyTracker:
             except Exception as e:
                 pass
 
-        env_location = self.get_env_location()
-        catalyst_location = self.get_catalyst_location()
-
         # Calculate hash and create zip
-        self.tracked_files.update(self.python_imports)
         hash_contents = []
 
         for filepath in sorted(self.tracked_files):
@@ -433,7 +449,11 @@ class TraceDependencyTracker:
                     continue
                 try:
                     relative_path = os.path.relpath(filepath, base_path)
-                    zipf.write(filepath, relative_path)
+                    if relative_path in ['', '.']:
+                        zipf.write(filepath, os.path.basename(filepath))
+                    else:
+                        zipf.write(filepath, relative_path)
+                    
                     logger.debug(f"Added python script to zip: {relative_path}")
                 except Exception as e:
                     pass

ragaai_catalyst/tracers/distributed.py

@@ -37,7 +37,7 @@ def init_tracing(
     secret_key: str = None,
     base_url: str = None,
     tracer: Tracer = None,
-    catalyst: RagaAICatalyst = None,
+    catalyst: RagaAICatalyst = None, 
     **kwargs
 ) -> None:
     """Initialize distributed tracing.
@@ -50,37 +50,20 @@ def init_tracing(
         base_url: RagaAI Catalyst API base URL
         tracer: Existing Tracer instance
         catalyst: Existing RagaAICatalyst instance
-        **kwargs: Additional tracer configuration
+        **kwargs: Additional tracer parameters
     """
     global _global_tracer, _global_catalyst
     
     with _tracer_lock:
         if tracer and catalyst:
-            _global_tracer = tracer
-            _global_catalyst = catalyst
+            if isinstance(tracer, Tracer) and isinstance(catalyst, RagaAICatalyst):
+                _global_tracer = tracer
+                _global_catalyst = catalyst
+            else:
+                raise ValueError("Both Tracer and Catalyst objects must be instances of Tracer and RagaAICatalyst, respectively.")
         else:
-            # Use env vars as fallback
-            access_key = access_key or os.getenv("RAGAAI_CATALYST_ACCESS_KEY")
-            secret_key = secret_key or os.getenv("RAGAAI_CATALYST_SECRET_KEY") 
-            base_url = base_url or os.getenv("RAGAAI_CATALYST_BASE_URL")
+            raise ValueError("Both Tracer and Catalyst objects must be provided.")
 
-            if not all([access_key, secret_key]):
-                raise ValueError(
-                    "Missing required credentials. Either provide access_key and secret_key "
-                    "or set RAGAAI_CATALYST_ACCESS_KEY and RAGAAI_CATALYST_SECRET_KEY environment variables."
-                )
-
-            _global_catalyst = RagaAICatalyst(
-                access_key=access_key,
-                secret_key=secret_key,
-                base_url=base_url
-            )
-
-            _global_tracer = Tracer(
-                project_name=project_name,
-                dataset_name=dataset_name,
-                **kwargs
-            )
 
 def trace_agent(name: str = None, agent_type: str = "generic", version: str = "1.0.0", **kwargs):
     """Decorator for tracing agent functions."""
@@ -162,7 +145,7 @@ def trace_llm(name: str = None, model: str = None, **kwargs):
             
             try:
                 # Just execute the function within the current span
-                result = await tracer.file_tracker.trace_wrapper(func)(*args, **kwargs)
+                result = await func(*args, **kwargs)
                 return result
             finally:
                 # Reset using the stored token
@@ -180,7 +163,7 @@ def trace_llm(name: str = None, model: str = None, **kwargs):
             
             try:
                 # Just execute the function within the current span
-                result = tracer.file_tracker.trace_wrapper(func)(*args, **kwargs)
+                result = func(*args, **kwargs)
                 return result
             finally:
                 # Reset using the stored token

ragaai_catalyst/tracers/langchain_callback.py

@@ -159,7 +159,7 @@ class LangchainTracer(BaseCallbackHandler):
         else:
             asyncio.run(self._async_save_trace(force))
 
-    def _create_safe_wrapper(self, original_func, component_name):
+    def _create_safe_wrapper(self, original_func, component_name, method_name):
         """Create a safely wrapped version of an original function with enhanced error handling"""
 
         @wraps(original_func)
@@ -209,7 +209,59 @@ class LangchainTracer(BaseCallbackHandler):
                 
                 # Fallback to calling the original function without modifications
                 return original_func(*args, **kwargs)
+            
+        @wraps(original_func)
+        def wrapped_invoke(*args, **kwargs):
+            if not self._active:
+                return original_func(*args, **kwargs)
+            
+            try:
+                # Deep copy kwargs to avoid modifying the original
+                kwargs_copy = kwargs.copy() if kwargs is not None else {}
+
+                # Handle different calling conventions
+                if 'config' not in kwargs_copy:
+                    kwargs_copy['config'] = {'callbacks': [self]}
+                elif 'callbacks' not in kwargs_copy['config']:
+                    kwargs_copy['config']['callbacks'] = [self]
+                elif self not in kwargs_copy['config']['callbacks']:
+                    kwargs_copy['config']['callbacks'].append(self)
+
+                # Store model name if available
+                if component_name in ["OpenAI", "ChatOpenAI_LangchainOpenAI", "ChatOpenAI_ChatModels",
+                                    "ChatVertexAI", "VertexAI", "ChatGoogleGenerativeAI", "ChatAnthropic", 
+                                    "ChatLiteLLM", "ChatBedrock", "AzureChatOpenAI", "ChatAnthropicVertex"]:
+                    instance = args[0] if args else None
+                    model_name = kwargs.get('model_name') or kwargs.get('model') or kwargs.get('model_id')
 
+                    if instance and model_name:
+                        self.model_names[id(instance)] = model_name
+                
+                # Try different method signatures
+                try:
+                    # First, try calling with modified kwargs
+                    return original_func(*args, **kwargs_copy)
+                except TypeError:
+                    # If that fails, try without kwargs
+                    try:
+                        return original_func(*args)
+                    except Exception as e:
+                        # If all else fails, use original call
+                        logger.error(f"Failed to invoke {component_name} with modified callbacks: {e}")
+                        return original_func(*args, **kwargs)
+            
+            except Exception as e:
+                # Log any errors that occur during the function call
+                logger.error(f"Error in {component_name} wrapper: {e}")
+                
+                # Record the error using the tracer's error handling method
+                self.on_error(e, context=f"wrapper_{component_name}")
+                
+                # Fallback to calling the original function without modifications
+                return original_func(*args, **kwargs)
+        
+        if method_name == 'invoke':
+            return wrapped_invoke
         return wrapped
 
 
@@ -287,6 +339,7 @@ class LangchainTracer(BaseCallbackHandler):
             from langchain.chains import create_retrieval_chain, RetrievalQA
             components_to_patch["RetrievalQA"] = (RetrievalQA, "from_chain_type")
             components_to_patch["create_retrieval_chain"] = (create_retrieval_chain, None)
+            components_to_patch['RetrievalQA.invoke'] = (RetrievalQA, 'invoke')
         except ImportError:
             logger.debug("Langchain chains not available for patching")
 
@@ -295,20 +348,20 @@ class LangchainTracer(BaseCallbackHandler):
                 if method_name == "__init__":
                     original = component.__init__
                     self._original_inits[name] = original
-                    component.__init__ = self._create_safe_wrapper(original, name)
+                    component.__init__ = self._create_safe_wrapper(original, name, method_name)
                 elif method_name:
                     original = getattr(component, method_name)
                     self._original_methods[name] = original
                     if isinstance(original, classmethod):
                         wrapped = classmethod(
-                            self._create_safe_wrapper(original.__func__, name)
+                            self._create_safe_wrapper(original.__func__, name, method_name)
                         )
                     else:
-                        wrapped = self._create_safe_wrapper(original, name)
+                        wrapped = self._create_safe_wrapper(original, name, method_name)
                     setattr(component, method_name, wrapped)
                 else:
                     self._original_methods[name] = component
-                    globals()[name] = self._create_safe_wrapper(component, name)
+                    globals()[name] = self._create_safe_wrapper(component, name, method_name)
             except Exception as e:
                 logger.error(f"Error patching {name}: {e}")
                 self.on_error(e, context=f"patch_{name}")
@@ -354,7 +407,7 @@ class LangchainTracer(BaseCallbackHandler):
                     elif name == "ChatOpenAI_ChatModels":
                         from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
                         imported_components[name] = ChatOpenAI_ChatModels
-                    elif name in ["RetrievalQA", "create_retrieval_chain"]:
+                    elif name in ["RetrievalQA", "create_retrieval_chain", 'RetrievalQA.invoke']:
                         from langchain.chains import create_retrieval_chain, RetrievalQA
                         imported_components["RetrievalQA"] = RetrievalQA
                         imported_components["create_retrieval_chain"] = create_retrieval_chain