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

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

ragaai_catalyst/tracers/tracer.py

@@ -29,6 +29,7 @@ from ragaai_catalyst.tracers.instrumentators import (
 )
 from ragaai_catalyst.tracers.utils import get_unique_key
 # from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
+from ragaai_catalyst.tracers.llamaindex_instrumentation import LlamaIndexInstrumentationTracer
 from ragaai_catalyst import RagaAICatalyst
 from ragaai_catalyst.tracers.agentic_tracing import AgenticTracing, TrackName
 from ragaai_catalyst.tracers.agentic_tracing.tracers.llm_tracer import LLMTracerMixin
@@ -170,7 +171,6 @@ class Tracer(AgenticTracing):
         elif tracer_type == "llamaindex":
             self._upload_task = None
             self.llamaindex_tracer = None
-            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
 
         else:
             self._upload_task = None
@@ -256,8 +256,7 @@ class Tracer(AgenticTracing):
             self.langchain_tracer = LangchainTracer()
             return self.langchain_tracer.start()
         elif self.tracer_type == "llamaindex":
-            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
-            self.llamaindex_tracer = LlamaIndexTracer(self._pass_user_data())
+            self.llamaindex_tracer = LlamaIndexInstrumentationTracer(self._pass_user_data())
             return self.llamaindex_tracer.start()
         else:
             super().start()
@@ -286,30 +285,31 @@ class Tracer(AgenticTracing):
                     if 'tokens' in additional_metadata and all(k in additional_metadata['tokens'] for k in ['prompt', 'completion']):
                         prompt_cost = additional_metadata["tokens"]["prompt"]*model_cost_data["input_cost_per_token"]
                         completion_cost = additional_metadata["tokens"]["completion"]*model_cost_data["output_cost_per_token"]
-                        additional_metadata.setdefault('cost', {})["total_cost"] = prompt_cost + completion_cost 
+                        additional_metadata["cost"] = prompt_cost + completion_cost 
+
+                        additional_metadata["prompt_tokens"] = float(additional_metadata["tokens"].get("prompt", 0.0))
+                        additional_metadata["completion_tokens"] = float(additional_metadata["tokens"].get("completion", 0.0))
+
+                        logger.debug("Metadata added successfully")
                     else:
                         logger.warning("Token information missing in additional_metadata")
+
+                    if 'cost' in additional_metadata:
+                        additional_metadata["cost"] = float(additional_metadata["cost"])
+                    else:
+                        additional_metadata["cost"] = 0.0
+                        logger.warning("Total cost information not available")
+
+
                 except Exception as e:
                     logger.warning(f"Error adding cost: {e}")
             else:
                 logger.debug("Model name not available in additional_metadata, skipping cost calculation")
             
-            # Safely get total tokens and cost
-            if 'tokens' in additional_metadata and 'total' in additional_metadata['tokens']:
-                additional_metadata["total_tokens"] = float(additional_metadata["tokens"]["total"])
-            else:
-                additional_metadata["total_tokens"] = 0.0
-                logger.warning("Total tokens information not available")
-
-            if 'cost' in additional_metadata and 'total_cost' in additional_metadata['cost']:
-                additional_metadata["total_cost"] = float(additional_metadata["cost"]["total_cost"])
-            else:
-                additional_metadata["total_cost"] = 0.0
-                logger.warning("Total cost information not available")
 
             # Safely remove tokens and cost dictionaries if they exist
             additional_metadata.pop("tokens", None)
-            additional_metadata.pop("cost", None)
+            # additional_metadata.pop("cost", None)
             
             # Safely merge metadata
             combined_metadata = {}
@@ -337,7 +337,8 @@ class Tracer(AgenticTracing):
             else:
                 logger.warning("No valid langchain traces found in final_result")
 
-            additional_metadata_keys = list(additional_metadata.keys()) if additional_metadata else None
+            # additional_metadata_keys = list(additional_metadata.keys()) if additional_metadata else None
+            additional_metadata_dict = additional_metadata if additional_metadata else {}
 
             UploadTraces(json_file_path=filepath_3,
                          project_name=self.project_name,
@@ -345,14 +346,30 @@ class Tracer(AgenticTracing):
                          dataset_name=self.dataset_name,
                          user_detail=user_detail,
                          base_url=self.base_url
-                         ).upload_traces(additional_metadata_keys=additional_metadata_keys)
+                         ).upload_traces(additional_metadata_keys=additional_metadata_dict)
             
             return 
 
         elif self.tracer_type == "llamaindex":
             if self.llamaindex_tracer is None:
                 raise ValueError("LlamaIndex tracer was not started")
-            return self.llamaindex_tracer.stop()
+
+            user_detail = self._pass_user_data()
+            converted_back_to_callback = self.llamaindex_tracer.stop()
+
+            filepath_3 = os.path.join(os.getcwd(), "llama_final_result.json")
+            with open(filepath_3, 'w') as f:
+                json.dump(converted_back_to_callback, f, default=str, indent=2)
+
+            if converted_back_to_callback:
+                UploadTraces(json_file_path=filepath_3,
+                             project_name=self.project_name,
+                             project_id=self.project_id,
+                             dataset_name=self.dataset_name,
+                             user_detail=user_detail,
+                             base_url=self.base_url
+                             ).upload_traces()
+            return 
         else:
             super().stop()
 

ragaai_catalyst/tracers/upload_traces.py

@@ -37,7 +37,10 @@ class UploadTraces:
 
         if additional_metadata_keys:
             for key in additional_metadata_keys:
-                SCHEMA_MAPPING_NEW[key] = {"columnType": "metadata"}
+                if key == "model_name":
+                    SCHEMA_MAPPING_NEW['response']["modelName"] = additional_metadata_keys[key]
+                else:
+                    SCHEMA_MAPPING_NEW[key] = {"columnType": key, "parentColumn": "response"}
 
         if additional_pipeline_keys:
             for key in additional_pipeline_keys:

ragaai_catalyst/tracers/utils/convert_llama_instru_callback.py

@@ -0,0 +1,69 @@
+def convert_llamaindex_instrumentation_to_callback(data):
+    data = data[0]
+    initial_struc = [{
+        "trace_id": data["trace_id"],
+        "project_id": data["project_id"],
+        "session_id": data["session_id"],
+        "trace_type": data["trace_type"],
+        "metadata" : data["metadata"],
+        "pipeline" : data["pipeline"],
+        "traces" : []
+    }]
+
+    traces_data = []
+
+    prompt = data["data"]["prompt"]
+    response = data["data"]["response"]
+    context = data["data"]["context"]
+    system_prompt = data["data"]["system_prompt"]
+
+    prompt_structured_data = {
+        "event_type": "query",
+        "payload": {
+            "query_str": prompt
+        }
+    }
+    traces_data.append(prompt_structured_data)
+
+    response_structured_data = {
+        "event_type": "llm",
+        "payload": {
+            "response": {
+                "message": {
+                    "content": response,
+                }
+            }
+        }
+    }
+    traces_data.append(response_structured_data)
+
+    context_structured_data = {
+        "event_type": "retrieve",
+        "payload": {
+            "nodes": [
+                {
+                    "node": {
+                        "text": context
+                    }
+                }
+            ]
+        }
+    }
+    traces_data.append(context_structured_data)
+
+    system_prompt_structured_data = {
+        "event_type": "llm",
+        "payload": {
+            "messages": [
+                {
+                    "role": "system",
+                    "content": system_prompt
+                }
+            ]
+        }
+    }
+    traces_data.append(system_prompt_structured_data)
+
+    initial_struc[0]["traces"] = traces_data    
+
+    return initial_struc

ragaai_catalyst/tracers/utils/extraction_logic_llama_index.py

@@ -0,0 +1,74 @@
+import json
+from typing import Dict, Any, Optional
+
+
+def extract_llama_index_data(data):
+    """
+    Transform llama_index trace data into standardized format
+    """
+    data = data[0]
+
+    # Extract top-level metadata
+    trace_data = {
+        "project_id": data.get("project_id"),
+        "trace_id": data.get("trace_id"),
+        "session_id": data.get("session_id"),
+        "trace_type": data.get("trace_type"),
+        "pipeline": data.get("pipeline"),
+        "metadata":data.get("metadata") ,
+        "prompt_length": 0,  
+        "data": {
+            "prompt": None,
+            "context": None,
+            "response": None,
+            "system_prompt": None
+        }
+    }
+
+    def get_prompt(data):
+        for span in data:
+            if span["event_type"]=="QueryStartEvent":
+                prompt = span.get("query", "")
+                return prompt
+            if span["event_type"]=="QueryEndEvent":
+                prompt = span.get("query", "")
+                return prompt
+
+
+    def get_context(data):
+        for span in data:
+            if span["event_type"]=="RetrievalEndEvent":
+                context = span.get("text", "")
+                return context
+    
+    def get_response(data):
+        for span in data:
+            if span["event_type"]=="QueryEndEvent":
+                response = span.get("response", "")
+                return response
+            # if span["event_type"]=="LLMPredictEndEvent":
+            #     response = span.get("output", "")
+            #     return response
+            # if span["event_type"]=="SynthesizeEndEvent":
+            #     response = span.get("response", "")
+            #     return response
+
+    def get_system_prompt(data):
+        for span in data:
+            if span["event_type"]=="LLMChatStartEvent":
+                response = span.get("messages", "")
+                response = response[0]
+                return response
+
+    # Process traces
+    if "traces" in data:
+        prompt = get_prompt(data["traces"])
+        context = get_context(data["traces"])
+        response = get_response(data["traces"])
+        system_prompt = get_system_prompt(data["traces"])
+
+    trace_data["data"]["prompt"] = prompt
+    trace_data["data"]["context"] = context
+    trace_data["data"]["response"] = response
+    trace_data["data"]["system_prompt"] = system_prompt
+    return [trace_data]