ragaai-catalyst 2.1.4.1b0__py3-none-any.whl → 2.1.5__py3-none-any.whl

ragaai_catalyst/tracers/distributed.py

@@ -0,0 +1,300 @@
+"""
+Distributed tracing functionality for RagaAI Catalyst.
+Provides simplified initialization and decorator-based tracing.
+"""
+
+import os
+import threading
+from typing import Optional, Dict, Any, List
+from functools import wraps
+from contextlib import contextmanager
+import uuid
+from .agentic_tracing.utils.unique_decorator import generate_unique_hash_simple
+from datetime import datetime
+import asyncio
+
+from .tracer import Tracer
+from ..ragaai_catalyst import RagaAICatalyst
+
+# Global state
+_global_tracer: Optional[Tracer] = None
+_global_catalyst: Optional[RagaAICatalyst] = None
+_tracer_lock = threading.Lock()
+_active_spans = threading.local()
+
+def get_current_tracer() -> Optional[Tracer]:
+    """Get the current global tracer instance."""
+    return _global_tracer
+
+def get_current_catalyst() -> Optional[RagaAICatalyst]:
+    """Get the current global catalyst instance."""
+    return _global_catalyst
+
+def init_tracing(
+    project_name: str = None,
+    dataset_name: str = None,
+    access_key: str = None,
+    secret_key: str = None,
+    base_url: str = None,
+    tracer: Tracer = None,
+    catalyst: RagaAICatalyst = None, 
+    **kwargs
+) -> None:
+    """Initialize distributed tracing.
+    
+    Args:
+        project_name: Project name for new tracer
+        dataset_name: Dataset name for new tracer
+        access_key: RagaAI Catalyst access key
+        secret_key: RagaAI Catalyst secret key  
+        base_url: RagaAI Catalyst API base URL
+        tracer: Existing Tracer instance
+        catalyst: Existing RagaAICatalyst instance
+        **kwargs: Additional tracer parameters
+    """
+    global _global_tracer, _global_catalyst
+    
+    with _tracer_lock:
+        if tracer and 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:
+            raise ValueError("Both Tracer and Catalyst objects must be provided.")
+
+
+def trace_agent(name: str = None, agent_type: str = "generic", version: str = "1.0.0", **kwargs):
+    """Decorator for tracing agent functions."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        span_name = name or func.__name__
+        # Generate hash based on the decorated function
+        top_level_hash_id = generate_unique_hash_simple(func)
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Set current agent name and store the token
+            name_token = tracer.current_agent_name.set(span_name)
+            
+            try:
+                # Use async agent tracing
+                return await tracer._trace_agent_execution(
+                    func,
+                    span_name,
+                    agent_type,
+                    version,
+                    None,  # capabilities
+                    top_level_hash_id, 
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_agent_name.reset(name_token)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Set current agent name and store the token
+            name_token = tracer.current_agent_name.set(span_name)
+            
+            try:
+                # Use synchronous agent tracing
+                return tracer._trace_sync_agent_execution(
+                    func,
+                    span_name,
+                    agent_type,
+                    version,
+                    None,  # capabilities
+                    top_level_hash_id,   
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_agent_name.reset(name_token)
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+
+def trace_llm(name: str = None, model: str = None, **kwargs):
+    """Decorator for tracing LLM calls."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        span_name = name or func.__name__
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Set current LLM name and store the token
+            name_token = tracer.current_llm_call_name.set(span_name)
+            
+            try:
+                # Just execute the function within the current span
+                result = await func(*args, **kwargs)
+                return result
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_llm_call_name.reset(name_token)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Set current LLM name and store the token
+            name_token = tracer.current_llm_call_name.set(span_name)
+            
+            try:
+                # Just execute the function within the current span
+                result = func(*args, **kwargs)
+                return result
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_llm_call_name.reset(name_token)
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+    
+
+def trace_tool(name: str = None, tool_type: str = "generic", version: str = "1.0.0", **kwargs):
+    """Decorator for tracing tool functions."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        span_name = name or func.__name__
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Set current tool name and store the token
+            name_token = tracer.current_tool_name.set(span_name)
+            
+            try:
+                # Use async tool tracing
+                return await tracer._trace_tool_execution(
+                    func,
+                    span_name,
+                    tool_type,
+                    version,
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_tool_name.reset(name_token)
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Set current tool name and store the token
+            name_token = tracer.current_tool_name.set(span_name)
+            
+            try:
+                # Use synchronous tool tracing
+                return tracer._trace_sync_tool_execution(
+                    func,
+                    span_name,
+                    tool_type,
+                    version,
+                    *args,
+                    **kwargs
+                )
+            finally:
+                # Reset using the stored token
+                if name_token:
+                    tracer.current_tool_name.reset(name_token)
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+
+
+
+def trace_custom(name: str = None, custom_type: str = "generic", version: str = "1.0.0", trace_variables: bool = False, **kwargs):
+    """Decorator for tracing custom functions."""
+    def decorator(func):
+        is_async = asyncio.iscoroutinefunction(func)
+        
+        @wraps(func)
+        async def async_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return await func(*args, **kwargs)
+
+            # Use async tool tracing
+            return await tracer._trace_custom_execution(
+                func,
+                name or func.__name__,
+                custom_type,
+                version,
+                trace_variables,
+                *args,
+                **kwargs
+            )
+
+        @wraps(func)
+        def sync_wrapper(*args, **kwargs):
+            tracer = get_current_tracer()
+            if not tracer:
+                return func(*args, **kwargs)
+
+            # Use synchronous tool tracing
+            return tracer._trace_sync_custom_execution(
+                func,
+                name or func.__name__,
+                custom_type,
+                version,
+                trace_variables,
+                *args,
+                **kwargs
+            )
+
+        return async_wrapper if is_async else sync_wrapper
+    return decorator
+
+
+def current_span():
+    """Get the current active span for adding metrics."""
+    tracer = get_current_tracer()
+    if not tracer:
+        return None
+    
+    # First check for LLM context
+    llm_name = tracer.current_llm_call_name.get()
+    if llm_name:
+        return tracer.span(llm_name)
+    
+    # Then check for tool context
+    tool_name = tracer.current_tool_name.get()
+    if tool_name:
+        return tracer.span(tool_name)
+    
+    # Finally fall back to agent context
+    agent_name = tracer.current_agent_name.get()
+    if not agent_name:
+        raise ValueError("No active span found. Make sure you're calling this within a traced function.")
+    
+    return tracer.span(agent_name)

ragaai_catalyst/tracers/exporters/__init__.py

@@ -1,5 +1,7 @@
 from .file_span_exporter import FileSpanExporter
 from .raga_exporter import RagaExporter
+from .ragaai_trace_exporter import RAGATraceExporter
+from .dynamic_trace_exporter import DynamicTraceExporter
 
 
-__all__ = ["FileSpanExporter", "RagaExporter"]
+__all__ = ["FileSpanExporter", "RagaExporter", "RAGATraceExporter", "DynamicTraceExporter"]

ragaai_catalyst/tracers/exporters/dynamic_trace_exporter.py

@@ -0,0 +1,160 @@
+"""
+Dynamic Trace Exporter - A wrapper for RAGATraceExporter that allows dynamic updates to properties.
+"""
+import logging
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+from ragaai_catalyst.tracers.exporters.ragaai_trace_exporter import RAGATraceExporter
+
+logger = logging.getLogger("RagaAICatalyst")
+
+class DynamicTraceExporter(SpanExporter):
+    """
+    A wrapper around RAGATraceExporter that allows dynamic updates to properties.
+    This exporter forwards all calls to the underlying RAGATraceExporter but allows
+    certain properties to be updated dynamically during execution.
+    """
+    
+    def __init__(self, files_to_zip, project_name, project_id, dataset_name, user_details, base_url, custom_model_cost):
+        """
+        Initialize the DynamicTraceExporter.
+        
+        Args:
+            files_to_zip: List of files to zip
+            project_name: Project name
+            project_id: Project ID
+            dataset_name: Dataset name
+            user_details: User details
+            base_url: Base URL for API
+        """
+        self._exporter = RAGATraceExporter(
+            files_to_zip=files_to_zip,
+            project_name=project_name,
+            project_id=project_id,
+            dataset_name=dataset_name,
+            user_details=user_details,
+            base_url=base_url,
+            custom_model_cost=custom_model_cost
+        )
+        
+        # Store the initial values
+        self._files_to_zip = files_to_zip
+        self._project_name = project_name
+        self._project_id = project_id
+        self._dataset_name = dataset_name
+        self._user_details = user_details
+        self._base_url = base_url
+        self._custom_model_cost = custom_model_cost
+
+    
+    def export(self, spans):
+        """
+        Export spans by forwarding to the underlying exporter.
+        Before exporting, update the exporter's properties with the current values.
+        
+        Args:
+            spans: Spans to export
+            
+        Returns:
+            SpanExportResult: Result of the export operation
+        """
+        try:
+            # Update the exporter's properties
+            self._update_exporter_properties()
+        except Exception as e:
+            raise Exception(f"Error updating exporter properties: {e}")
+
+        try:
+            # Forward the call to the underlying exporter
+            result = self._exporter.export(spans)
+            return result
+        except Exception as e:
+            raise Exception(f"Error exporting trace: {e}")
+            
+
+
+    def shutdown(self):
+        """
+        Shutdown the exporter by forwarding to the underlying exporter.
+        Before shutting down, update the exporter's properties with the current values.
+        """
+        try:
+            # Update the exporter's properties
+            self._update_exporter_properties()
+        except Exception as e:
+            raise Exception(f"Error updating exporter properties: {e}")
+
+        try:
+            # Forward the call to the underlying exporter
+            return self._exporter.shutdown()
+        except Exception as e:
+            raise Exception(f"Error shutting down exporter: {e}")
+    
+    def _update_exporter_properties(self):
+        """
+        Update the underlying exporter's properties with the current values.
+        """
+        self._exporter.files_to_zip = self._files_to_zip
+        self._exporter.project_name = self._project_name
+        self._exporter.project_id = self._project_id
+        self._exporter.dataset_name = self._dataset_name
+        self._exporter.user_details = self._user_details
+        self._exporter.base_url = self._base_url
+        self._exporter.custom_model_cost = self._custom_model_cost
+    
+    # Getter and setter methods for dynamic properties
+    
+    @property
+    def files_to_zip(self):
+        return self._files_to_zip
+    
+    @files_to_zip.setter
+    def files_to_zip(self, value):
+        self._files_to_zip = value
+    
+    @property
+    def project_name(self):
+        return self._project_name
+    
+    @project_name.setter
+    def project_name(self, value):
+        self._project_name = value
+    
+    @property
+    def project_id(self):
+        return self._project_id
+    
+    @project_id.setter
+    def project_id(self, value):
+        self._project_id = value
+    
+    @property
+    def dataset_name(self):
+        return self._dataset_name
+    
+    @dataset_name.setter
+    def dataset_name(self, value):
+        self._dataset_name = value
+    
+    @property
+    def user_details(self):
+        return self._user_details
+    
+    @user_details.setter
+    def user_details(self, value):
+        self._user_details = value
+    
+    @property
+    def base_url(self):
+        return self._base_url
+    
+    @base_url.setter
+    def base_url(self, value):
+        self._base_url = value
+
+    @property
+    def custom_model_cost(self):
+        return self._custom_model_cost
+    
+    @custom_model_cost.setter
+    def custom_model_cost(self, value):
+        self._custom_model_cost = value

ragaai_catalyst/tracers/exporters/ragaai_trace_exporter.py

@@ -0,0 +1,129 @@
+import os
+import json
+import tempfile
+from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
+import logging
+from datetime import datetime
+from dataclasses import asdict
+from ragaai_catalyst.tracers.utils.trace_json_converter import convert_json_format
+from ragaai_catalyst.tracers.agentic_tracing.tracers.base import TracerJSONEncoder
+from ragaai_catalyst.tracers.agentic_tracing.utils.system_monitor import SystemMonitor
+from ragaai_catalyst.tracers.agentic_tracing.upload.trace_uploader import submit_upload_task
+from ragaai_catalyst.tracers.agentic_tracing.utils.zip_list_of_unique_files import zip_list_of_unique_files
+
+
+logger = logging.getLogger("RagaAICatalyst")
+logging_level = (
+    logger.setLevel(logging.DEBUG) if os.getenv("DEBUG") == "1" else logging.INFO
+)
+
+
+class RAGATraceExporter(SpanExporter):
+    def __init__(self, files_to_zip, project_name, project_id, dataset_name, user_details, base_url, custom_model_cost):
+        self.trace_spans = dict()
+        self.tmp_dir = tempfile.gettempdir()
+        self.files_to_zip = files_to_zip
+        self.project_name = project_name
+        self.project_id = project_id
+        self.dataset_name = dataset_name
+        self.user_details = user_details
+        self.base_url = base_url
+        self.custom_model_cost = custom_model_cost
+        self.system_monitor = SystemMonitor(dataset_name)
+
+    def export(self, spans):
+        for span in spans:
+            span_json = json.loads(span.to_json())
+            trace_id = span_json.get("context").get("trace_id")
+            if trace_id is None:
+                raise Exception("Trace ID is None")
+
+            if trace_id not in self.trace_spans:
+                self.trace_spans[trace_id] = list()
+
+            self.trace_spans[trace_id].append(span_json)
+
+            if span_json["parent_id"] is None:
+                trace = self.trace_spans[trace_id]
+                try:
+                    self.process_complete_trace(trace, trace_id)
+                except Exception as e:
+                    raise Exception(f"Error processing complete trace: {e}")
+                try:
+                    del self.trace_spans[trace_id]
+                except Exception as e:
+                    raise Exception(f"Error deleting trace: {e}")
+
+        return SpanExportResult.SUCCESS
+
+    def shutdown(self):
+        # Process any remaining traces during shutdown
+        for trace_id, spans in self.trace_spans.items():
+            self.process_complete_trace(spans, trace_id)
+        self.trace_spans.clear()
+
+    def process_complete_trace(self, spans, trace_id):
+        # Convert the trace to ragaai trace format
+        try:
+            ragaai_trace_details = self.prepare_trace(spans, trace_id)
+        except Exception as e:
+            print(f"Error converting trace {trace_id}: {e}")
+        
+        # Upload the trace if upload_trace function is provided
+        try:
+            self.upload_trace(ragaai_trace_details, trace_id)
+        except Exception as e:
+            print(f"Error uploading trace {trace_id}: {e}")
+
+    def prepare_trace(self, spans, trace_id):
+        try:
+            ragaai_trace = convert_json_format(spans, self.custom_model_cost)            
+            ragaai_trace["workflow"] = []
+
+            # Add source code hash
+            hash_id, zip_path = zip_list_of_unique_files(
+                self.files_to_zip, output_dir=self.tmp_dir
+            )
+
+            ragaai_trace["metadata"]["system_info"] = asdict(self.system_monitor.get_system_info())
+            ragaai_trace["metadata"]["resources"] = asdict(self.system_monitor.get_resources())
+            ragaai_trace["metadata"]["system_info"]["source_code"] = hash_id
+
+            ragaai_trace["data"][0]["start_time"] = ragaai_trace["start_time"]
+            ragaai_trace["data"][0]["end_time"] = ragaai_trace["end_time"]
+
+            ragaai_trace["project_name"] = self.project_name
+            
+            # Save the trace_json 
+            trace_file_path = os.path.join(self.tmp_dir, f"{trace_id}.json")
+            with open(trace_file_path, "w") as file:
+                json.dump(ragaai_trace, file, cls=TracerJSONEncoder, indent=2)
+
+            return {
+                'trace_file_path': trace_file_path,
+                'code_zip_path': zip_path,
+                'hash_id': hash_id
+            }
+        except Exception as e:
+            logger.error(f"Error converting trace {trace_id}: {str(e)}")
+            return None
+
+    def upload_trace(self, ragaai_trace_details, trace_id):
+        filepath = ragaai_trace_details['trace_file_path']
+        hash_id = ragaai_trace_details['hash_id']
+        zip_path = ragaai_trace_details['code_zip_path']
+
+        
+
+        self.upload_task_id = submit_upload_task(
+                filepath=filepath,
+                hash_id=hash_id,
+                zip_path=zip_path,
+                project_name=self.project_name,
+                project_id=self.project_id,
+                dataset_name=self.dataset_name,
+                user_details=self.user_details,
+                base_url=self.base_url
+            )
+
+        logger.info(f"Submitted upload task with ID: {self.upload_task_id}")