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

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

ragaai_catalyst/tracers/llamaindex_instrumentation.py

@@ -0,0 +1,424 @@
+from configparser import InterpolationMissingOptionError
+import json
+from datetime import datetime
+from typing import Any, Optional, Dict, List, ClassVar
+from pydantic import Field
+# from treelib import Tree
+
+from llama_index.core.instrumentation.span import SimpleSpan
+from llama_index.core.instrumentation.span_handlers.base import BaseSpanHandler
+from llama_index.core.instrumentation.events import BaseEvent
+from llama_index.core.instrumentation.event_handlers import BaseEventHandler
+from llama_index.core.instrumentation import get_dispatcher
+from llama_index.core.instrumentation.span_handlers import SimpleSpanHandler
+
+from llama_index.core.instrumentation.events.agent import (
+    AgentChatWithStepStartEvent,
+    AgentChatWithStepEndEvent,
+    AgentRunStepStartEvent,
+    AgentRunStepEndEvent,
+    AgentToolCallEvent,
+)
+from llama_index.core.instrumentation.events.chat_engine import (
+    StreamChatErrorEvent,
+    StreamChatDeltaReceivedEvent,
+)
+from llama_index.core.instrumentation.events.embedding import (
+    EmbeddingStartEvent,
+    EmbeddingEndEvent,
+)
+from llama_index.core.instrumentation.events.llm import (
+    LLMPredictEndEvent,
+    LLMPredictStartEvent,
+    LLMStructuredPredictEndEvent,
+    LLMStructuredPredictStartEvent,
+    LLMCompletionEndEvent,
+    LLMCompletionStartEvent,
+    LLMChatEndEvent,
+    LLMChatStartEvent,
+    LLMChatInProgressEvent,
+)
+from llama_index.core.instrumentation.events.query import (
+    QueryStartEvent,
+    QueryEndEvent,
+)
+from llama_index.core.instrumentation.events.rerank import (
+    ReRankStartEvent,
+    ReRankEndEvent,
+)
+from llama_index.core.instrumentation.events.retrieval import (
+    RetrievalStartEvent,
+    RetrievalEndEvent,
+)
+from llama_index.core.instrumentation.events.span import (
+    SpanDropEvent,
+)
+from llama_index.core.instrumentation.events.synthesis import (
+    SynthesizeStartEvent,
+    SynthesizeEndEvent,
+    GetResponseEndEvent,
+    GetResponseStartEvent,
+)
+
+import uuid
+
+from .utils.extraction_logic_llama_index import extract_llama_index_data
+from .utils.convert_llama_instru_callback import convert_llamaindex_instrumentation_to_callback
+
+class EventHandler(BaseEventHandler):
+    """Example event handler.
+
+    This event handler is an example of how to create a custom event handler.
+
+    In general, logged events are treated as single events in a point in time,
+    that link to a span. The span is a collection of events that are related to
+    a single task. The span is identified by a unique span_id.
+
+    While events are independent, there is some hierarchy.
+    For example, in query_engine.query() call with a reranker attached:
+    - QueryStartEvent
+    - RetrievalStartEvent
+    - EmbeddingStartEvent
+    - EmbeddingEndEvent
+    - RetrievalEndEvent
+    - RerankStartEvent
+    - RerankEndEvent
+    - SynthesizeStartEvent
+    - GetResponseStartEvent
+    - LLMPredictStartEvent
+    - LLMChatStartEvent
+    - LLMChatEndEvent
+    - LLMPredictEndEvent
+    - GetResponseEndEvent
+    - SynthesizeEndEvent
+    - QueryEndEvent
+    """
+
+    events: List[BaseEvent] = []
+    current_trace: List[Dict[str, Any]] = []  # Store events for the current trace
+
+
+    @classmethod
+    def class_name(cls) -> str:
+        """Class name."""
+        return "EventHandler"
+
+    def handle(self, event: BaseEvent) -> None:
+        """Logic for handling event."""
+        # print("-----------------------")
+        # # all events have these attributes
+        # print(event.id_)
+        # print(event.timestamp)
+        # print(event.span_id)
+
+        # Prepare event details dictionary
+        event_details = {
+            "id": event.id_,
+            "timestamp": event.timestamp,
+            "span_id": event.span_id,
+            "event_type": event.class_name(),
+        }
+
+        # event specific attributes
+        # print(f"Event type: {event.class_name()}")
+        if isinstance(event, AgentRunStepStartEvent):
+            event_details.update({
+                "task_id": event.task_id,
+                "step": event.step,
+                "input": event.input,
+            })
+        if isinstance(event, AgentRunStepEndEvent):
+            event_details.update({
+                "step_output": event.step_output,
+            })
+        if isinstance(event, AgentChatWithStepStartEvent):
+            event_details.update({
+                "user_msg": event.user_msg,
+            })
+        if isinstance(event, AgentChatWithStepEndEvent):
+            event_details.update({
+                "response": event.response,
+            })
+        if isinstance(event, AgentToolCallEvent):
+            event_details.update({
+                "arguments": event.arguments,
+                "tool_name": event.tool.name,
+                "tool_description": event.tool.description,
+                "tool_openai": event.tool.to_openai_tool(),
+            })
+        if isinstance(event, StreamChatDeltaReceivedEvent):
+            event_details.update({
+                "delta": event.delta,
+            })
+        if isinstance(event, StreamChatErrorEvent):
+            event_details.update({
+                "exception": event.exception,
+            })
+        if isinstance(event, EmbeddingStartEvent):
+            event_details.update({
+                "model_dict": event.model_dict,
+            })
+        if isinstance(event, EmbeddingEndEvent):
+            event_details.update({
+                "chunks": event.chunks,
+                "embeddings": event.embeddings[0][:5],
+            })
+        if isinstance(event, LLMPredictStartEvent):
+            event_details.update({
+                "template": event.template,
+                "template_args": event.template_args,
+            })
+        if isinstance(event, LLMPredictEndEvent):
+            event_details.update({
+                "output": event.output,
+            })
+        if isinstance(event, LLMStructuredPredictStartEvent):
+            event_details.update({
+                "template": event.template,
+                "template_args": event.template_args,
+                "output_cls": event.output_cls,
+            })
+        if isinstance(event, LLMStructuredPredictEndEvent):
+            event_details.update({
+                "output": event.output,
+            })
+        if isinstance(event, LLMCompletionStartEvent):
+            event_details.update({
+                "model_dict": event.model_dict,
+                "prompt": event.prompt,
+                "additional_kwargs": event.additional_kwargs,
+            })
+        if isinstance(event, LLMCompletionEndEvent):
+            event_details.update({
+                "response": event.response,
+                "prompt": event.prompt,
+            })
+        if isinstance(event, LLMChatInProgressEvent):
+            event_details.update({
+                "messages": event.messages,
+                "response": event.response,
+            })
+        if isinstance(event, LLMChatStartEvent):
+            event_details.update({
+                "messages": event.messages,
+                "additional_kwargs": event.additional_kwargs,
+                "model_dict": event.model_dict,
+            })
+        if isinstance(event, LLMChatEndEvent):
+            event_details.update({
+                "messages": event.messages,
+                "response": event.response,
+            })
+        if isinstance(event, RetrievalStartEvent):
+            event_details.update({
+                "str_or_query_bundle": event.str_or_query_bundle,
+            })
+        if isinstance(event, RetrievalEndEvent):
+            event_details.update({
+                "str_or_query_bundle": event.str_or_query_bundle,
+                "nodes": event.nodes,
+                "text": event.nodes[0].text
+            })
+        if isinstance(event, ReRankStartEvent):
+            event_details.update({
+                "query": event.query,
+                "nodes": event.nodes,
+                "top_n": event.top_n,
+                "model_name": event.model_name,
+            })
+        if isinstance(event, ReRankEndEvent):
+            event_details.update({
+                "nodes": event.nodes,
+            })
+        if isinstance(event, QueryStartEvent):
+            event_details.update({
+                "query": event.query,
+            })
+        if isinstance(event, QueryEndEvent):
+            event_details.update({
+                "response": event.response,
+                "query": event.query,
+            })
+        if isinstance(event, SpanDropEvent):
+            event_details.update({
+                "err_str": event.err_str,
+            })
+        if isinstance(event, SynthesizeStartEvent):
+            event_details.update({
+                "query": event.query,
+            })
+        if isinstance(event, SynthesizeEndEvent):
+            event_details.update({
+                "response": event.response,
+                "query": event.query,
+            })
+        if isinstance(event, GetResponseStartEvent):
+            event_details.update({
+                "query_str": event.query_str,
+            })
+
+        # Append event details to current_trace
+        self.current_trace.append(event_details)
+
+        self.events.append(event)
+
+    def _get_events_by_span(self) -> Dict[str, List[BaseEvent]]:
+        events_by_span: Dict[str, List[BaseEvent]] = {}
+        for event in self.events:
+            if event.span_id in events_by_span:
+                events_by_span[event.span_id].append(event)
+            else:
+                events_by_span[event.span_id] = [event]
+        return events_by_span
+
+    # def _get_event_span_trees(self) -> List[Tree]:
+    #     events_by_span = self._get_events_by_span()
+
+    #     trees = []
+    #     tree = Tree()
+
+    #     for span, sorted_events in events_by_span.items():
+    #         # create root node i.e. span node
+    #         tree.create_node(
+    #             tag=f"{span} (SPAN)",
+    #             identifier=span,
+    #             parent=None,
+    #             data=sorted_events[0].timestamp,
+    #         )
+
+    #         for event in sorted_events:
+    #             tree.create_node(
+    #                 tag=f"{event.class_name()}: {event.id_}",
+    #                 identifier=event.id_,
+    #                 parent=event.span_id,
+    #                 data=event.timestamp,
+    #             )
+
+    #         trees.append(tree)
+    #         tree = Tree()
+    #     return trees
+
+    # def print_event_span_trees(self) -> None:
+    #     """Method for viewing trace trees."""
+    #     trees = self._get_event_span_trees()
+    #     for tree in trees:
+    #         print(
+    #             tree.show(
+    #                 stdout=False, sorting=True, key=lambda node: node.data
+    #             )
+    #         )
+    #         print("")
+
+
+
+class SpanHandler(BaseSpanHandler[SimpleSpan]):
+    # span_dict = {}
+    span_dict: ClassVar[Dict[str, List[SimpleSpan]]] = {}
+
+    @classmethod
+    def class_name(cls) -> str:
+        """Class name."""
+        return "SpanHandler"
+
+    def new_span(
+        self,
+        id_: str,
+        bound_args: Any,
+        instance: Optional[Any] = None,
+        parent_span_id: Optional[str] = None,
+        tags: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> Optional[SimpleSpan]:
+        """Create a span."""
+        # logic for creating a new MyCustomSpan
+        if id_ not in self.span_dict:
+            self.span_dict[id_] = []
+        self.span_dict[id_].append(
+            SimpleSpan(id_=id_, parent_id=parent_span_id)
+        )
+
+    def prepare_to_exit_span(
+        self,
+        id_: str,
+        bound_args: Any,
+        instance: Optional[Any] = None,
+        result: Optional[Any] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Logic for preparing to exit a span."""
+        pass
+        # if id in self.span_dict:
+        #    return self.span_dict[id].pop()
+
+    def prepare_to_drop_span(
+        self,
+        id_: str,
+        bound_args: Any,
+        instance: Optional[Any] = None,
+        err: Optional[BaseException] = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Logic for preparing to drop a span."""
+        pass
+        # if id in self.span_dict:
+        #    return self.span_dict[id].pop()
+
+
+
+class LlamaIndexInstrumentationTracer:
+    def __init__(self, user_detail):
+        """Initialize the LlamaIndexTracer with handlers but don't start tracing yet."""
+        # Initialize the root dispatcher
+        self.root_dispatcher = get_dispatcher()
+
+        # Initialize handlers
+        self.json_event_handler = EventHandler()
+        self.span_handler = SpanHandler()
+        self.simple_span_handler = SimpleSpanHandler()
+
+        self.is_tracing = False  # Flag to check if tracing is active
+
+        self.user_detail = user_detail
+
+    def start(self):
+        """Start tracing by registering handlers."""
+        if self.is_tracing:
+            print("Tracing is already active.")
+            return
+
+        # Register handlers
+        self.root_dispatcher.add_span_handler(self.span_handler)
+        self.root_dispatcher.add_span_handler(self.simple_span_handler)
+        self.root_dispatcher.add_event_handler(self.json_event_handler)
+
+        self.is_tracing = True
+        print("Tracing started.")
+
+    def stop(self):
+        """Stop tracing by unregistering handlers."""
+        if not self.is_tracing:
+            print("Tracing is not active.")
+            return
+
+        # Write current_trace to a JSON file
+        final_traces = {
+            "project_id": self.user_detail["project_id"],
+            "trace_id": str(uuid.uuid4()),
+            "session_id": None,
+            "trace_type": "llamaindex",
+            "metadata": self.user_detail["trace_user_detail"]["metadata"],
+            "pipeline": self.user_detail["trace_user_detail"]["pipeline"],
+            "traces": self.json_event_handler.current_trace,
+
+        }
+
+        with open('new_llamaindex_traces.json', 'w') as f:
+            json.dump([final_traces], f, default=str, indent=4)
+        
+        llamaindex_instrumentation_data = extract_llama_index_data([final_traces])
+        converted_back_to_callback = convert_llamaindex_instrumentation_to_callback(llamaindex_instrumentation_data)
+
+         # Just indicate tracing is stopped
+        self.is_tracing = False
+        print("Tracing stopped.")
+        return converted_back_to_callback