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

ragaai_catalyst/tracers/tracer.py

@@ -1,12 +1,23 @@
+from audioop import add
 import os
+import uuid
 import datetime
 import logging
 import asyncio
 import aiohttp
 import requests
+from litellm import model_cost
+
 from contextlib import contextmanager
 from concurrent.futures import ThreadPoolExecutor
-
+from ragaai_catalyst.tracers.langchain_callback import LangchainTracer
+from ragaai_catalyst.tracers.utils.convert_langchain_callbacks_output import convert_langchain_callbacks_output
+
+from ragaai_catalyst.tracers.utils.langchain_tracer_extraction_logic import langchain_tracer_extraction
+from ragaai_catalyst.tracers.upload_traces import UploadTraces
+import tempfile
+import json
+import numpy as np
 from opentelemetry.sdk import trace as trace_sdk
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from ragaai_catalyst.tracers.exporters.file_span_exporter import FileSpanExporter
@@ -18,15 +29,17 @@ 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 import AgenticTracing
 from ragaai_catalyst.tracers.agentic_tracing.tracers.llm_tracer import LLMTracerMixin
-from ragaai_catalyst.tracers.agentic_tracing.utils.trace_utils import load_model_costs, update_model_costs_from_github
+from ragaai_catalyst.tracers.exporters.ragaai_trace_exporter import RAGATraceExporter
+from ragaai_catalyst.tracers.agentic_tracing.utils.file_name_tracker import TrackName
 
 logger = logging.getLogger(__name__)
 
 class Tracer(AgenticTracing):
-    NUM_PROJECTS = 100
+    NUM_PROJECTS = 99999
     TIMEOUT = 10
     def __init__(
         self,
@@ -77,7 +90,17 @@ class Tracer(AgenticTracing):
 
         # take care of auto_instrumentation
         if isinstance(auto_instrumentation, bool):
-            if auto_instrumentation:
+            if tracer_type == "agentic/llamaindex":
+                auto_instrumentation = {
+                    "llm": False,
+                    "tool": False,
+                    "agent": False,
+                    "user_interaction": False,
+                    "file_io": False,
+                    "network": False,
+                    "custom": False
+                }
+            elif auto_instrumentation:
                 auto_instrumentation = {
                     "llm": True,
                     "tool": True,
@@ -98,11 +121,11 @@ class Tracer(AgenticTracing):
                     "custom": False
                 }
         elif isinstance(auto_instrumentation, dict):
-            auto_instrumentation = {k: v for k, v in auto_instrumentation.items() if v}
+            auto_instrumentation = {k: v for k, v in auto_instrumentation.items()}
             for key in ["llm", "tool", "agent", "user_interaction", "file_io", "network", "custom"]:
                 if key not in auto_instrumentation:
-                    auto_instrumentation[key] = False
-        
+                    auto_instrumentation[key] = True
+        self.model_custom_cost = {}
         super().__init__(user_detail=user_detail, auto_instrumentation=auto_instrumentation)
 
         self.project_name = project_name
@@ -116,12 +139,11 @@ class Tracer(AgenticTracing):
         self.upload_timeout = upload_timeout
         self.base_url = f"{RagaAICatalyst.BASE_URL}"
         self.timeout = 30
-        self.num_projects = 100
+        self.num_projects = 99999
         self.start_time = datetime.datetime.now().astimezone().isoformat()
-
-        if update_llm_cost:
-            # First update the model costs file from GitHub
-            update_model_costs_from_github()
+        self.model_cost_dict = model_cost
+        self.user_context = ""  # Initialize user_context to store context from add_context
+        self.file_tracker = TrackName()
         
         try:
             response = requests.get(
@@ -152,43 +174,112 @@ class Tracer(AgenticTracing):
             raise
 
         if tracer_type == "langchain":
-            self.raga_client = RagaExporter(project_name=self.project_name, dataset_name=self.dataset_name)
+            # self.raga_client = RagaExporter(project_name=self.project_name, dataset_name=self.dataset_name)
 
-            self._tracer_provider = self._setup_provider()
-            self._instrumentor = self._setup_instrumentor(tracer_type)
-            self.is_instrumented = False
+            # self._tracer_provider = self._setup_provider()
+            # self._instrumentor = self._setup_instrumentor(tracer_type)
+            # self.is_instrumented = False
+            # self._upload_task = None
             self._upload_task = None
         elif tracer_type == "llamaindex":
             self._upload_task = None
-            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
-
+            self.llamaindex_tracer = None
+        elif tracer_type == "agentic/llamaindex":
+            from opentelemetry.sdk import trace as trace_sdk
+            from opentelemetry.sdk.trace.export import SimpleSpanProcessor 
+            from openinference.instrumentation.llama_index import LlamaIndexInstrumentor 
+            from ragaai_catalyst.tracers.exporters.dynamic_trace_exporter import DynamicTraceExporter
+
+            # Get the code_files
+            self.file_tracker.trace_main_file()
+            list_of_unique_files = self.file_tracker.get_unique_files()
+
+            # Create a dynamic exporter that allows property updates
+            self.dynamic_exporter = DynamicTraceExporter(
+                files_to_zip=list_of_unique_files,
+                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,
+                custom_model_cost=self.model_custom_cost
+            )
+            
+            tracer_provider = trace_sdk.TracerProvider()
+            tracer_provider.add_span_processor(SimpleSpanProcessor(self.dynamic_exporter))
+            LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
         else:
             self._upload_task = None
             # raise ValueError (f"Currently supported tracer types are 'langchain' and 'llamaindex'.")
 
-        
+    def set_model_cost(self, cost_config):
+        """
+        Set custom cost values for a specific model.
+
+        Args:
+            cost_config (dict): Dictionary containing model cost configuration with keys:
+                - model_name (str): Name of the model
+                - input_cost_per_token (float): Cost per input token
+                - output_cost_per_token (float): Cost per output token
+
+        Example:
+            tracer.set_model_cost({
+                "model_name": "gpt-4",
+                "input_cost_per_million_token": 6,
+                "output_cost_per_million_token": 2.40
+            })
+        """
+        if not isinstance(cost_config, dict):
+            raise TypeError("cost_config must be a dictionary")
+
+        required_keys = {"model_name", "input_cost_per_million_token", "output_cost_per_million_token"}
+        if not all(key in cost_config for key in required_keys):
+            raise ValueError(f"cost_config must contain all required keys: {required_keys}")
+
+        model_name = cost_config["model_name"]
+        self.model_custom_cost[model_name] = {
+            "input_cost_per_token": float(cost_config["input_cost_per_million_token"])/ 1000000,
+            "output_cost_per_token": float(cost_config["output_cost_per_million_token"]) /1000000
+        }
+
     def set_dataset_name(self, dataset_name):
         """
         Reinitialize the Tracer with a new dataset name while keeping all other parameters the same.
+        If using agentic/llamaindex tracer with dynamic exporter, update the exporter's dataset_name property.
         
         Args:
             dataset_name (str): The new dataset name to set
         """
-        # Store current parameters
-        current_params = {
-            'project_name': self.project_name,
-            'tracer_type': self.tracer_type,
-            'pipeline': self.pipeline,
-            'metadata': self.metadata,
-            'description': self.description,
-            'upload_timeout': self.upload_timeout
-        }
-        
-        # Reinitialize self with new dataset_name and stored parameters
-        self.__init__(
-            dataset_name=dataset_name,
-            **current_params
-        )
+        # If we have a dynamic exporter, update its dataset_name property
+        if self.tracer_type == "agentic/llamaindex" and hasattr(self, "dynamic_exporter"):
+            # Update the dataset name in the dynamic exporter
+            self.dynamic_exporter.dataset_name = dataset_name
+            logger.debug(f"Updated dynamic exporter's dataset_name to {dataset_name}")
+            
+            # Update the instance variable
+            self.dataset_name = dataset_name
+            
+            # Update user_details with new dataset_name
+            self.user_details = self._pass_user_data()
+            
+            # Also update the user_details in the dynamic exporter
+            self.dynamic_exporter.user_details = self.user_details
+        else:
+            # Store current parameters
+            current_params = {
+                'project_name': self.project_name,
+                'tracer_type': self.tracer_type,
+                'pipeline': self.pipeline,
+                'metadata': self.metadata,
+                'description': self.description,
+                'upload_timeout': self.upload_timeout
+            }
+            
+            # Reinitialize self with new dataset_name and stored parameters
+            self.__init__(
+                dataset_name=dataset_name,
+                **current_params
+            )
 
     def _improve_metadata(self, metadata, tracer_type):
         if metadata is None:
@@ -239,14 +330,15 @@ class Tracer(AgenticTracing):
     def start(self):
         """Start the tracer."""
         if self.tracer_type == "langchain":
-            if not self.is_instrumented:
-                self._instrumentor().instrument(tracer_provider=self._tracer_provider)
-                self.is_instrumented = True
-            print(f"Tracer started for project: {self.project_name}")
-            return self
+            # if not self.is_instrumented:
+            #     self._instrumentor().instrument(tracer_provider=self._tracer_provider)
+            #     self.is_instrumented = True
+            # print(f"Tracer started for project: {self.project_name}")
+            self.langchain_tracer = LangchainTracer()
+            return self.langchain_tracer.start()
         elif self.tracer_type == "llamaindex":
-            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
-            return LlamaIndexTracer(self._pass_user_data()).start()
+            self.llamaindex_tracer = LlamaIndexInstrumentationTracer(self._pass_user_data())
+            return self.llamaindex_tracer.start()
         else:
             super().start()
             return self
@@ -254,20 +346,111 @@ class Tracer(AgenticTracing):
     def stop(self):
         """Stop the tracer and initiate trace upload."""
         if self.tracer_type == "langchain":
-            if not self.is_instrumented:
-                logger.warning("Tracer was not started. No traces to upload.")
-                return "No traces to upload"
-
-            print("Stopping tracer and initiating trace upload...")
-            self._cleanup()
-            self._upload_task = self._run_async(self._upload_traces())
-            self.is_active = False
-            self.dataset_name = None
+            # if not self.is_instrumented:
+            #     logger.warning("Tracer was not started. No traces to upload.")
+            #     return "No traces to upload"
+
+            # print("Stopping tracer and initiating trace upload...")
+            # self._cleanup()
+            # self._upload_task = self._run_async(self._upload_traces())
+            # self.is_active = False
+            # self.dataset_name = None
             
-            return "Trace upload initiated. Use get_upload_status() to check the status."
+            user_detail = self._pass_user_data()
+            data, additional_metadata = self.langchain_tracer.stop()
+
+            # Add cost if possible
+            if additional_metadata.get('model_name'):
+                try:
+                    model_cost_data = self.model_cost_dict[additional_metadata['model_name']]
+                    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["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 remove tokens and cost dictionaries if they exist
+            additional_metadata.pop("tokens", None)
+            # additional_metadata.pop("cost", None)
+            
+            # Safely merge metadata
+            combined_metadata = {}
+            if user_detail.get('trace_user_detail', {}).get('metadata'):
+                combined_metadata.update(user_detail['trace_user_detail']['metadata'])
+            if additional_metadata:
+                combined_metadata.update(additional_metadata)
+
+            langchain_traces = langchain_tracer_extraction(data, self.user_context)
+            final_result = convert_langchain_callbacks_output(langchain_traces)
+            
+            # Safely set required fields in final_result
+            if final_result and isinstance(final_result, list) and len(final_result) > 0:
+                final_result[0]['project_name'] = user_detail.get('project_name', '')
+                final_result[0]['trace_id'] = str(uuid.uuid4())
+                final_result[0]['session_id'] = None
+                final_result[0]['metadata'] = combined_metadata
+                final_result[0]['pipeline'] = user_detail.get('trace_user_detail', {}).get('pipeline')
+
+                filepath_3 = os.path.join(os.getcwd(), "final_result.json")
+                with open(filepath_3, 'w') as f:
+                    json.dump(final_result, f, indent=2)
+                
+                print(filepath_3)
+            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_dict = additional_metadata if additional_metadata else {}
+
+            UploadTraces(json_file_path=filepath_3,
+                         project_name=self.project_name,
+                         project_id=self.project_id,
+                         dataset_name=self.dataset_name,
+                         user_detail=self._pass_user_data(),
+                         base_url=self.base_url
+                         ).upload_traces(additional_metadata_keys=additional_metadata_dict)
+            
+            return 
+
         elif self.tracer_type == "llamaindex":
-            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
-            return LlamaIndexTracer(self._pass_user_data()).stop()
+            if self.llamaindex_tracer is None:
+                raise ValueError("LlamaIndex tracer was not started")
+
+            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()
 
@@ -379,4 +562,67 @@ class Tracer(AgenticTracing):
                     }
                 }
             }
-        return user_detail
+        return user_detail
+
+    def update_dynamic_exporter(self, **kwargs):
+        """
+        Update the dynamic exporter's properties.
+        
+        Args:
+            **kwargs: Keyword arguments to update. Can include any of the following:
+                - 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
+                
+        Raises:
+            AttributeError: If the tracer_type is not 'agentic/llamaindex' or if the dynamic_exporter is not initialized.
+        """
+        if self.tracer_type != "agentic/llamaindex" or not hasattr(self, "dynamic_exporter"):
+            raise AttributeError("Dynamic exporter is only available for 'agentic/llamaindex' tracer type")
+            
+        for key, value in kwargs.items():
+            if hasattr(self.dynamic_exporter, key):
+                setattr(self.dynamic_exporter, key, value)
+                logger.debug(f"Updated dynamic exporter's {key} to {value}")
+            else:
+                logger.warning(f"Dynamic exporter has no attribute '{key}'")
+                
+    def update_file_list(self):
+        """
+        Update the file list in the dynamic exporter with the latest tracked files.
+        This is useful when new files are added to the project during execution.
+        
+        Raises:
+            AttributeError: If the tracer_type is not 'agentic/llamaindex' or if the dynamic_exporter is not initialized.
+        """
+        if self.tracer_type != "agentic/llamaindex" or not hasattr(self, "dynamic_exporter"):
+            raise AttributeError("Dynamic exporter is only available for 'agentic/llamaindex' tracer type")
+            
+        # Get the latest list of unique files
+        list_of_unique_files = self.file_tracker.get_unique_files()
+        
+        # Update the dynamic exporter's files_to_zip property
+        self.dynamic_exporter.files_to_zip = list_of_unique_files
+        logger.debug(f"Updated dynamic exporter's files_to_zip with {len(list_of_unique_files)} files")
+    
+    def add_context(self, context):
+        """
+        Add context information to the trace. This method is only supported for 'langchain' and 'llamaindex' tracer types.
+
+        Args:
+            context: Additional context information to be added to the trace. Can be a string.
+
+        Raises:
+            ValueError: If tracer_type is not 'langchain' or 'llamaindex'.
+        """
+        if self.tracer_type not in ["langchain", "llamaindex"]:
+            raise ValueError("add_context is only supported for 'langchain' and 'llamaindex' tracer types")
+        
+        # Convert string context to string if needed
+        if isinstance(context, str):
+            self.user_context = context
+        else:
+            raise TypeError("context must be a string")

ragaai_catalyst/tracers/upload_traces.py

@@ -20,7 +20,7 @@ class UploadTraces:
         self.base_url = base_url
         self.timeout = 10
 
-    def _create_dataset_schema_with_trace(self):
+    def _create_dataset_schema_with_trace(self, additional_metadata_keys=None, additional_pipeline_keys=None):
         SCHEMA_MAPPING_NEW = {
             "trace_id": {"columnType": "traceId"},
             "trace_uri": {"columnType": "traceUri"},
@@ -34,6 +34,18 @@ class UploadTraces:
             "vector_store":{"columnType":"pipeline"},
             "feedback": {"columnType":"feedBack"}
         }
+
+        if additional_metadata_keys:
+            for key in additional_metadata_keys:
+                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:
+                SCHEMA_MAPPING_NEW[key] = {"columnType": "pipeline"}
+                
         def make_request():
             headers = {
                 "Content-Type": "application/json",
@@ -119,9 +131,14 @@ class UploadTraces:
                                     data=payload,
                                     timeout=self.timeout)
 
-    def upload_traces(self):
-        self._create_dataset_schema_with_trace()
-        presignedUrl = self._get_presigned_url()
-        self._put_presigned_url(presignedUrl, self.json_file_path)
-        self._insert_traces(presignedUrl)
-        print("Traces uploaded")
+    def upload_traces(self, additional_metadata_keys=None, additional_pipeline_keys=None):
+        try:
+            self._create_dataset_schema_with_trace(additional_metadata_keys, additional_pipeline_keys)
+            presignedUrl = self._get_presigned_url()
+            if presignedUrl is None:
+                return
+            self._put_presigned_url(presignedUrl, self.json_file_path)
+            self._insert_traces(presignedUrl)
+            print("Traces uploaded")
+        except Exception as e:
+            print(f"Error while uploading agentic traces: {e}")

ragaai_catalyst/tracers/utils/convert_langchain_callbacks_output.py

@@ -0,0 +1,61 @@
+import json
+
+def convert_langchain_callbacks_output(result, project_name="", metadata="", pipeline=""):
+    initial_struc = [{
+        "project_name": project_name,
+        "trace_id": "NA",
+        "session_id": "NA",
+        "metadata" : metadata,
+        "pipeline" : pipeline,
+        "traces" : []
+    }]
+    traces_data = []
+
+    prompt = result["data"]["prompt"]
+    response = result["data"]["response"]
+    context = result["data"]["context"]
+    final_prompt = ""
+
+    prompt_structured_data = {
+        "traceloop.entity.input": json.dumps({
+            "kwargs": {
+                "input": prompt,
+            }
+        })
+    }    
+    prompt_data = {
+        "name": "retrieve_documents.langchain.workflow",
+        "attributes": prompt_structured_data,
+    }
+
+    traces_data.append(prompt_data)
+
+    context_structured_data = {
+        "traceloop.entity.input": json.dumps({
+            "kwargs": {
+                "context": context
+            }
+        }),
+        "traceloop.entity.output": json.dumps({
+            "kwargs": {
+                "text": prompt
+            }
+        })
+    }
+    context_data = {
+        "name": "PromptTemplate.langchain.task",
+        "attributes": context_structured_data,
+    }
+    traces_data.append(context_data)
+
+    response_structured_data = {"gen_ai.completion.0.content": response,
+                                "gen_ai.prompt.0.content": prompt}
+    response_data = {
+        "name": "ChatOpenAI.langchain.task",
+        "attributes" : response_structured_data
+    }
+    traces_data.append(response_data)
+
+    initial_struc[0]["traces"] = traces_data
+
+    return initial_struc

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