ragaai-catalyst 2.1.5b6__py3-none-any.whl → 2.1.5b8__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -150,8 +150,6 @@ class LLMTracerMixin:
         beta_module = openai_module.beta
 
         # Patch openai.beta.threads
-        import openai
-        openai.api_type = "openai"
         if hasattr(beta_module, "threads"):
             threads_obj = beta_module.threads
             # Patch top-level methods on openai.beta.threads

ragaai_catalyst/tracers/langchain_callback.py

@@ -48,8 +48,9 @@ class LangchainTracer(BaseCallbackHandler):
         self._original_methods = {}
         self.additional_metadata = {}
         self._save_task = None
-        self._current_query = None  # Add this line to track the current query
+        self._current_query = None
         self.filepath = None
+        self.model_names = {}  # Store model names by component instance
         logger.setLevel(log_level)
 
         if not os.path.exists(output_path):
@@ -86,7 +87,7 @@ class LangchainTracer(BaseCallbackHandler):
             "retriever_actions": [],
             "tokens": [],
             "errors": [],
-            "query": self._current_query,  # Add this line to include the query in the trace
+            "query": self._current_query,
             "metadata": {
                 "version": "2.0",
                 "trace_all": self.trace_all,
@@ -172,6 +173,14 @@ class LangchainTracer(BaseCallbackHandler):
                     kwargs_copy['callbacks'] = [self]
                 elif self not in kwargs_copy['callbacks']:
                     kwargs_copy['callbacks'].append(self)
+
+                # Store model name if available
+                if component_name in ["OpenAI", "ChatOpenAI_LangchainOpenAI", "ChatOpenAI_ChatModels", 
+                                    "ChatVertexAI", "ChatGoogleGenerativeAI", "ChatAnthropic", "ChatLiteLLM"]:
+                    instance = args[0] if args else None
+                    model_name = kwargs.get('model_name') or kwargs.get('model')
+                    if instance and model_name:
+                        self.model_names[id(instance)] = model_name
                 
                 # Try different method signatures
                 try:
@@ -201,28 +210,56 @@ class LangchainTracer(BaseCallbackHandler):
 
     def _monkey_patch(self):
         """Enhanced monkey-patching with comprehensive component support"""
-        from langchain.llms import OpenAI
-        # from langchain_groq import ChatGroq
-        # from langchain_google_genai import ChatGoogleGenerativeAI
-        # from langchain_anthropic import ChatAnthropic
-        from langchain_community.chat_models import ChatLiteLLM
-        # from langchain_cohere import ChatCohere
-        from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
-        from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
-        from langchain.chains import create_retrieval_chain, RetrievalQA
-
-        components_to_patch = {
-            "OpenAI": (OpenAI, "__init__"),
-            # "ChatGroq": (ChatGroq, "__init__"),
-            # "ChatGoogleGenerativeAI": (ChatGoogleGenerativeAI, "__init__"),
-            # "ChatAnthropic": (ChatAnthropic, "__init__"),
-            "ChatLiteLLM": (ChatLiteLLM, "__init__"),
-            # "ChatCohere": (ChatCohere, "__init__"),
-            "ChatOpenAI_LangchainOpenAI": (ChatOpenAI_LangchainOpenAI, "__init__"),
-            "ChatOpenAI_ChatModels": (ChatOpenAI_ChatModels, "__init__"),
-            "RetrievalQA": (RetrievalQA, "from_chain_type"),
-            "create_retrieval_chain": (create_retrieval_chain, None),
-        }
+        components_to_patch = {}
+        
+        try:
+            from langchain.llms import OpenAI
+            components_to_patch["OpenAI"] = (OpenAI, "__init__")
+        except ImportError:
+            logger.debug("OpenAI not available for patching")
+            
+        try:
+            from langchain_google_vertexai import ChatVertexAI
+            components_to_patch["ChatVertexAI"] = (ChatVertexAI, "__init__")
+        except ImportError:
+            logger.debug("ChatVertexAI not available for patching")
+            
+        try:
+            from langchain_google_genai import ChatGoogleGenerativeAI
+            components_to_patch["ChatGoogleGenerativeAI"] = (ChatGoogleGenerativeAI, "__init__")
+        except ImportError:
+            logger.debug("ChatGoogleGenerativeAI not available for patching")
+            
+        try:
+            from langchain_anthropic import ChatAnthropic
+            components_to_patch["ChatAnthropic"] = (ChatAnthropic, "__init__")
+        except ImportError:
+            logger.debug("ChatAnthropic not available for patching")
+            
+        try:
+            from langchain_community.chat_models import ChatLiteLLM
+            components_to_patch["ChatLiteLLM"] = (ChatLiteLLM, "__init__")
+        except ImportError:
+            logger.debug("ChatLiteLLM not available for patching")
+            
+        try:
+            from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
+            components_to_patch["ChatOpenAI_LangchainOpenAI"] = (ChatOpenAI_LangchainOpenAI, "__init__")
+        except ImportError:
+            logger.debug("ChatOpenAI_LangchainOpenAI not available for patching")
+            
+        try:
+            from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
+            components_to_patch["ChatOpenAI_ChatModels"] = (ChatOpenAI_ChatModels, "__init__")
+        except ImportError:
+            logger.debug("ChatOpenAI_ChatModels not available for patching")
+            
+        try:
+            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)
+        except ImportError:
+            logger.debug("Langchain chains not available for patching")
 
         for name, (component, method_name) in components_to_patch.items():
             try:
@@ -249,21 +286,45 @@ class LangchainTracer(BaseCallbackHandler):
 
     def _restore_original_methods(self):
         """Restore all original methods and functions with enhanced error handling"""
-        from langchain.llms import OpenAI
-        # from langchain_groq import ChatGroq
-        # from langchain_google_genai import ChatGoogleGenerativeAI
-        # from langchain_anthropic import ChatAnthropic
-        from langchain_community.chat_models import ChatLiteLLM
-        # from langchain_cohere import ChatCohere
-        from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
-        from langchain.chat_models import ChatOpenAI as ChatOpenAI_ChatModels
-        from langchain.chains import create_retrieval_chain, RetrievalQA
-
+        # Dynamically import only what we need based on what was patched
+        imported_components = {}
+        
+        if self._original_inits or self._original_methods:
+            for name in list(self._original_inits.keys()) + list(self._original_methods.keys()):
+                try:
+                    if name == "OpenAI":
+                        from langchain.llms import OpenAI
+                        imported_components[name] = OpenAI
+                    elif name == "ChatVertexAI":
+                        from langchain_google_vertexai import ChatVertexAI
+                        imported_components[name] = ChatVertexAI
+                    elif name == "ChatGoogleGenerativeAI":
+                        from langchain_google_genai import ChatGoogleGenerativeAI
+                        imported_components[name] = ChatGoogleGenerativeAI
+                    elif name == "ChatAnthropic":
+                        from langchain_anthropic import ChatAnthropic
+                        imported_components[name] = ChatAnthropic
+                    elif name == "ChatLiteLLM":
+                        from langchain_community.chat_models import ChatLiteLLM
+                        imported_components[name] = ChatLiteLLM
+                    elif name == "ChatOpenAI_LangchainOpenAI":
+                        from langchain_openai import ChatOpenAI as ChatOpenAI_LangchainOpenAI
+                        imported_components[name] = ChatOpenAI_LangchainOpenAI
+                    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"]:
+                        from langchain.chains import create_retrieval_chain, RetrievalQA
+                        imported_components["RetrievalQA"] = RetrievalQA
+                        imported_components["create_retrieval_chain"] = create_retrieval_chain
+                except ImportError:
+                    logger.debug(f"{name} not available for restoration")
 
         for name, original in self._original_inits.items():
             try:
-                component = eval(name)
-                component.__init__ = original
+                if name in imported_components:
+                    component = imported_components[name]
+                    component.__init__ = original
             except Exception as e:
                 logger.error(f"Error restoring {name}: {e}")
                 self.on_error(e, context=f"restore_{name}")
@@ -272,10 +333,12 @@ class LangchainTracer(BaseCallbackHandler):
             try:
                 if "." in name:
                     module_name, method_name = name.rsplit(".", 1)
-                    module = eval(module_name)
-                    setattr(module, method_name, original)
+                    if module_name in imported_components:
+                        module = imported_components[module_name]
+                        setattr(module, method_name, original)
                 else:
-                    globals()[name] = original
+                    if name in imported_components:
+                        globals()[name] = original
             except Exception as e:
                 logger.error(f"Error restoring {name}: {e}")
                 self.on_error(e, context=f"restore_{name}")
@@ -359,16 +422,92 @@ class LangchainTracer(BaseCallbackHandler):
                 }
             )
 
+            # Calculate latency
             end_time = datetime.now()
-            self.additional_metadata["latency"] = (end_time - self.current_trace["start_time"]).total_seconds()
+            latency = (end_time - self.current_trace["start_time"]).total_seconds()
 
+            # Check if values are there in llm_output
+            model = ""
+            prompt_tokens = 0
+            completion_tokens = 0
+            total_tokens = 0
+            
+            # Try to get model name from llm_output first
             if response and response.llm_output:
-                self.additional_metadata["model_name"] = response.llm_output.get("model_name", "")
-                self.additional_metadata["tokens"] = {}
-                if response.llm_output.get("token_usage"):
-                    self.additional_metadata["tokens"]["total"] = response.llm_output["token_usage"].get("total_tokens", 0)
-                    self.additional_metadata["tokens"]["prompt"] = response.llm_output["token_usage"].get("prompt_tokens", 0)
-                    self.additional_metadata["tokens"]["completion"] = response.llm_output["token_usage"].get("completion_tokens", 0)
+                try:
+                    model = response.llm_output.get("model_name")
+                    if not model:
+                        model = response.llm_output.get("model", "")
+                except Exception as e:
+                    # logger.debug(f"Error getting model name: {e}")
+                    model = ""
+
+            # Add model name
+            if not model:
+                try:
+                    model = response.llm_output.get("model_name")
+                    if not model:
+                        model = response.llm_output.get("model", "")
+                except Exception as e:
+                    # logger.debug(f"Error getting model name: {e}")
+                    model = ""
+
+
+            # Add token usage
+            try:
+                token_usage = response.llm_output.get("token_usage", {})
+                if token_usage=={}:
+                    try:
+                        token_usage = response.llm_output.get("usage")
+                    except Exception as e:
+                        # logger.debug(f"Error getting token usage: {e}")
+                        token_usage = {}
+                    
+                if token_usage !={}:
+                    prompt_tokens = token_usage.get("prompt_tokens", 0)
+                    if prompt_tokens==0:
+                        prompt_tokens = token_usage.get("input_tokens", 0)
+                    completion_tokens = token_usage.get("completion_tokens", 0)
+                    if completion_tokens==0:
+                        completion_tokens = token_usage.get("output_tokens", 0)
+
+                    total_tokens = prompt_tokens + completion_tokens
+            except Exception as e:
+                # logger.debug(f"Error getting token usage: {e}")
+                prompt_tokens = 0
+                completion_tokens = 0
+                total_tokens = 0
+
+            # Check if values are there in 
+            if prompt_tokens == 0 and completion_tokens == 0:
+                try:
+                    usage_data = response.generations[0][0].message.usage_metadata
+                    prompt_tokens = usage_data.get("input_tokens", 0)
+                    completion_tokens = usage_data.get("output_tokens", 0)
+                    total_tokens = prompt_tokens + completion_tokens
+                except Exception as e:
+                    # logger.debug(f"Error getting usage data: {e}")
+                    prompt_tokens = 0
+                    completion_tokens = 0
+                    total_tokens = 0
+
+            # If no model name in llm_output, try to get it from stored model names
+            try:
+                if model == "":
+                    model = list(self.model_names.values())[0]
+            except Exception as e:
+                model=""
+            
+            self.additional_metadata = {
+                'latency': latency,
+                'model_name': model,
+                'tokens': {
+                    'prompt': prompt_tokens,
+                    'completion': completion_tokens,
+                    'total': total_tokens
+                }
+            }
+
         except Exception as e:
             self.on_error(e, context="llm_end")
 

ragaai_catalyst/tracers/tracer.py

@@ -0,0 +1,453 @@
+from audioop import add
+import os
+import uuid
+import datetime
+import logging
+import asyncio
+import aiohttp
+import requests
+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
+
+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
+from ragaai_catalyst.tracers.exporters.raga_exporter import RagaExporter
+from ragaai_catalyst.tracers.instrumentators import (
+    LangchainInstrumentor,
+    OpenAIInstrumentor,
+    LlamaIndexInstrumentor,
+)
+from ragaai_catalyst.tracers.utils import get_unique_key
+# from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
+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
+from ragaai_catalyst.tracers.agentic_tracing.utils.trace_utils import load_model_costs, update_model_costs_from_github
+
+logger = logging.getLogger(__name__)
+
+class Tracer(AgenticTracing):
+    NUM_PROJECTS = 100
+    TIMEOUT = 10
+    def __init__(
+        self,
+        project_name,
+        dataset_name,
+        trace_name=None,
+        tracer_type=None,
+        pipeline=None,
+        metadata=None,
+        description=None,
+        upload_timeout=30,  # Default timeout of 30 seconds
+        update_llm_cost=True,  # Parameter to control model cost updates
+        auto_instrumentation={ # to control automatic instrumentation of different components
+            'llm':True,
+            'tool':True,
+            'agent':True,
+            'user_interaction':True,
+            'file_io':True,
+            'network':True,
+            'custom':True
+        },
+        interval_time=2,
+        # auto_instrumentation=True/False  # to control automatic instrumentation of everything
+
+    ):
+        """
+        Initializes a Tracer object. 
+
+        Args:
+            project_name (str): The name of the project.
+            dataset_name (str): The name of the dataset.
+            tracer_type (str, optional): The type of tracer. Defaults to None.
+            pipeline (dict, optional): The pipeline configuration. Defaults to None.
+            metadata (dict, optional): The metadata. Defaults to None.
+            description (str, optional): The description. Defaults to None.
+            upload_timeout (int, optional): The upload timeout in seconds. Defaults to 30.
+            update_llm_cost (bool, optional): Whether to update model costs from GitHub. Defaults to True.
+        """
+
+        user_detail = {
+            "project_name": project_name,
+            "project_id": None,  # Will be set after project validation
+            "dataset_name": dataset_name,
+            "interval_time": interval_time,
+            "trace_name": trace_name if trace_name else f"trace_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}",
+            "trace_user_detail": {"metadata": metadata} if metadata else {}
+        }
+
+        # take care of auto_instrumentation
+        if isinstance(auto_instrumentation, bool):
+            if auto_instrumentation:
+                auto_instrumentation = {
+                    "llm": True,
+                    "tool": True,
+                    "agent": True,
+                    "user_interaction": True,
+                    "file_io": True,
+                    "network": True,
+                    "custom": True
+                }
+            else:
+                auto_instrumentation = {
+                    "llm": False,
+                    "tool": False,
+                    "agent": False,
+                    "user_interaction": False,
+                    "file_io": False,
+                    "network": False,
+                    "custom": False
+                }
+        elif isinstance(auto_instrumentation, dict):
+            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] = True
+        
+        super().__init__(user_detail=user_detail, auto_instrumentation=auto_instrumentation)
+
+        self.project_name = project_name
+        self.dataset_name = dataset_name
+        self.tracer_type = tracer_type
+        self.metadata = self._improve_metadata(metadata, tracer_type)
+        # self.metadata["total_cost"] = 0.0
+        # self.metadata["total_tokens"] = 0
+        self.pipeline = pipeline
+        self.description = description
+        self.upload_timeout = upload_timeout
+        self.base_url = f"{RagaAICatalyst.BASE_URL}"
+        self.timeout = 30
+        self.num_projects = 100
+        self.start_time = datetime.datetime.now().astimezone().isoformat()
+        self.model_cost_dict = load_model_costs()
+
+        if update_llm_cost:
+            # First update the model costs file from GitHub
+            update_model_costs_from_github()
+        
+        try:
+            response = requests.get(
+                f"{self.base_url}/v2/llm/projects?size={self.num_projects}",
+                headers={
+                    "Authorization": f'Bearer {os.getenv("RAGAAI_CATALYST_TOKEN")}',
+                },
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+            logger.debug("Projects list retrieved successfully")
+
+            project_list = [
+                project["name"] for project in response.json()["data"]["content"]
+            ]
+            if project_name not in project_list:
+                raise ValueError("Project not found. Please enter a valid project name")
+            
+            self.project_id = [
+                project["id"] for project in response.json()["data"]["content"] if project["name"] == project_name
+            ][0]
+            # super().__init__(user_detail=self._pass_user_data())
+            # self.file_tracker = TrackName()
+            self._pass_user_data()
+
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Failed to retrieve projects list: {e}")
+            raise
+
+        if tracer_type == "langchain":
+            # 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._upload_task = None
+            self._upload_task = None
+        elif tracer_type == "llamaindex":
+            self._upload_task = None
+            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
+
+        else:
+            self._upload_task = None
+            # raise ValueError (f"Currently supported tracer types are 'langchain' and 'llamaindex'.")
+
+        
+    def set_dataset_name(self, dataset_name):
+        """
+        Reinitialize the Tracer with a new dataset name while keeping all other parameters the same.
+        
+        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
+        )
+
+    def _improve_metadata(self, metadata, tracer_type):
+        if metadata is None:
+            metadata = {}
+        metadata.setdefault("log_source", f"{tracer_type}_tracer")
+        metadata.setdefault("recorded_on", str(datetime.datetime.now()))
+        return metadata
+
+    def _add_unique_key(self, data, key_name):
+        data[key_name] = get_unique_key(data)
+        return data
+
+    def _setup_provider(self):
+        self.filespanx = FileSpanExporter(
+            project_name=self.project_name,
+            metadata=self.metadata,
+            pipeline=self.pipeline,
+            raga_client=self.raga_client,
+        )
+        tracer_provider = trace_sdk.TracerProvider()
+        tracer_provider.add_span_processor(SimpleSpanProcessor(self.filespanx))
+        return tracer_provider
+
+    def _setup_instrumentor(self, tracer_type):
+        instrumentors = {
+            "langchain": LangchainInstrumentor,
+            "openai": OpenAIInstrumentor,
+            "llama_index": LlamaIndexInstrumentor,
+        }
+        if tracer_type not in instrumentors:
+            raise ValueError(f"Invalid tracer type: {tracer_type}")
+        return instrumentors[tracer_type]().get()
+
+    @contextmanager
+    def trace(self):
+        """
+        Synchronous context manager for tracing.
+        Usage:
+            with tracer.trace():
+                # Your code here
+        """
+        self.start()
+        try:
+            yield self
+        finally:
+            self.stop()
+
+    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}")
+            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()
+        else:
+            super().start()
+            return self
+
+    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
+
+            # filename = f"langchain_callback_traces.json"
+            # filepath = os.path.join(tempfile.gettempdir(), filename) 
+            
+            user_detail = self._pass_user_data()
+            data, additional_metadata = self.langchain_tracer.stop()
+
+            # Add cost if possible
+            # import pdb; pdb.set_trace()
+            if additional_metadata['model_name']:
+                try:
+                    model_cost_data = self.model_cost_dict[additional_metadata['model_name']]
+                    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', {})["prompt_cost"] = prompt_cost
+                    # additional_metadata.setdefault('cost', {})["completion_cost"] = completion_cost
+                    additional_metadata.setdefault('cost', {})["total_cost"] = prompt_cost + completion_cost 
+                except Exception as e:
+                    logger.warning(f"Error adding cost: {e}")
+
+            # with open(filepath, 'r') as f:
+            #     data = json.load(f)
+            additional_metadata["total_tokens"] = additional_metadata["tokens"]["total"]
+            del additional_metadata["tokens"]
+            if "cost" in additional_metadata:
+                additional_metadata["total_cost"] = additional_metadata["cost"]["total_cost"]
+                del additional_metadata["cost"]
+            else:
+                additional_metadata["total_cost"] = 0.0
+            
+            combined_metadata = user_detail['trace_user_detail']['metadata'].copy()
+            combined_metadata.update(additional_metadata)
+            combined_metadata
+
+            langchain_traces = langchain_tracer_extraction(data)
+            final_result = convert_langchain_callbacks_output(langchain_traces)
+            final_result[0]['project_name'] = user_detail['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['trace_user_detail']['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)
+
+            additional_metadata_keys = additional_metadata.keys() if additional_metadata else None
+
+            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(additional_metadata_keys=additional_metadata_keys)
+            
+            return 
+
+        elif self.tracer_type == "llamaindex":
+            from ragaai_catalyst.tracers.llamaindex_callback import LlamaIndexTracer
+            return LlamaIndexTracer(self._pass_user_data()).stop()
+        else:
+            super().stop()
+
+    def get_upload_status(self):
+        """Check the status of the trace upload."""
+        if self.tracer_type == "langchain":
+            if self._upload_task is None:
+                return "No upload task in progress."
+            if self._upload_task.done():
+                try:
+                    result = self._upload_task.result()
+                    return f"Upload completed: {result}"
+                except Exception as e:
+                    return f"Upload failed: {str(e)}"
+            return "Upload in progress..."
+
+    def _run_async(self, coroutine):
+        """Run an asynchronous coroutine in a separate thread."""
+        loop = asyncio.new_event_loop()
+        with ThreadPoolExecutor() as executor:
+            future = executor.submit(lambda: loop.run_until_complete(coroutine))
+        return future
+
+    async def _upload_traces(self):
+        """
+        Asynchronously uploads traces to the RagaAICatalyst server.
+
+        This function uploads the traces generated by the RagaAICatalyst client to the RagaAICatalyst server. It uses the `aiohttp` library to make an asynchronous HTTP request to the server. The function first checks if the `RAGAAI_CATALYST_TOKEN` environment variable is set. If not, it raises a `ValueError` with the message "RAGAAI_CATALYST_TOKEN not found. Cannot upload traces.".
+
+        The function then uses the `asyncio.wait_for` function to wait for the `check_and_upload_files` method of the `raga_client` object to complete. The `check_and_upload_files` method is called with the `session` object and a list of file paths to be uploaded. The `timeout` parameter is set to the value of the `upload_timeout` attribute of the `Tracer` object.
+
+        If the upload is successful, the function returns the string "Files uploaded successfully" if the `upload_stat` variable is truthy, otherwise it returns the string "No files to upload".
+
+        If the upload times out, the function returns a string with the message "Upload timed out after {self.upload_timeout} seconds".
+
+        If any other exception occurs during the upload, the function returns a string with the message "Upload failed: {str(e)}", where `{str(e)}` is the string representation of the exception.
+
+        Parameters:
+            None
+
+        Returns:
+            A string indicating the status of the upload.
+        """
+        async with aiohttp.ClientSession() as session:
+            if not os.getenv("RAGAAI_CATALYST_TOKEN"):
+                raise ValueError(
+                    "RAGAAI_CATALYST_TOKEN not found. Cannot upload traces."
+                )
+
+            try:
+                upload_stat = await asyncio.wait_for(
+                    self.raga_client.check_and_upload_files(
+                        session=session,
+                        file_paths=[self.filespanx.sync_file],
+                    ),
+                    timeout=self.upload_timeout,
+                )
+                return (
+                    "Files uploaded successfully"
+                    if upload_stat
+                    else "No files to upload"
+                )
+            except asyncio.TimeoutError:
+                return f"Upload timed out after {self.upload_timeout} seconds"
+            except Exception as e:
+                return f"Upload failed: {str(e)}"
+
+    def _cleanup(self):
+        """
+        Cleans up the tracer by uninstrumenting the instrumentor, shutting down the tracer provider,
+        and resetting the instrumentation flag. This function is called when the tracer is no longer
+        needed.
+
+        Parameters:
+            self (Tracer): The Tracer instance.
+
+        Returns:
+            None
+        """
+        if self.is_instrumented:
+            try:
+                self._instrumentor().uninstrument()
+                self._tracer_provider.shutdown()
+                self.is_instrumented = False
+                print("Tracer provider shut down successfully")
+            except Exception as e:
+                logger.error(f"Error during tracer shutdown: {str(e)}")
+
+        # Reset instrumentation flag
+        self.is_instrumented = False
+        # Note: We're not resetting all attributes here to allow for upload status checking
+
+    def _pass_user_data(self):
+        user_detail = {
+            "project_name":self.project_name, 
+            "project_id": self.project_id,
+            "dataset_name":self.dataset_name, 
+            "trace_user_detail" : {
+                "project_id": self.project_id,
+                "trace_id": "",
+                "session_id": None,
+                "trace_type": self.tracer_type,
+                "traces": [],
+                "metadata": self.metadata,
+                "pipeline": {
+                    "llm_model": (getattr(self, "pipeline", {}) or {}).get("llm_model", ""),
+                    "vector_store": (getattr(self, "pipeline", {}) or {}).get("vector_store", ""),
+                    "embed_model": (getattr(self, "pipeline", {}) or {}).get("embed_model", "")
+                    }
+                }
+            }
+        return user_detail

{ragaai_catalyst-2.1.5b6.dist-info → ragaai_catalyst-2.1.5b8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ragaai_catalyst
-Version: 2.1.5b6
+Version: 2.1.5b8
 Summary: RAGA AI CATALYST
 Author-email: Kiran Scaria <kiran.scaria@raga.ai>, Kedar Gaikwad <kedar.gaikwad@raga.ai>, Dushyant Mahajan <dushyant.mahajan@raga.ai>, Siddhartha Kosti <siddhartha.kosti@raga.ai>, Ritika Goel <ritika.goel@raga.ai>, Vijay Chaurasia <vijay.chaurasia@raga.ai>
 Requires-Python: <3.13,>=3.9

{ragaai_catalyst-2.1.5b6.dist-info → ragaai_catalyst-2.1.5b8.dist-info}/RECORD

@@ -13,8 +13,9 @@ ragaai_catalyst/synthetic_data_generation.py,sha256=uDV9tNwto2xSkWg5XHXUvjErW-4P
 ragaai_catalyst/utils.py,sha256=TlhEFwLyRU690HvANbyoRycR3nQ67lxVUQoUOfTPYQ0,3772
 ragaai_catalyst/tracers/__init__.py,sha256=LfgTes-nHpazssbGKnn8kyLZNr49kIPrlkrqqoTFTfc,301
 ragaai_catalyst/tracers/distributed.py,sha256=AIRvS5Ur4jbFDXsUkYuCTmtGoHHx3LOG4n5tWOh610U,10330
-ragaai_catalyst/tracers/langchain_callback.py,sha256=LvMBhgvAX8ftyBQ9Naeui46EoDa2nHQZq48Ra6nL-Qg,21991
+ragaai_catalyst/tracers/langchain_callback.py,sha256=lLeED0Eg2kT4-_O9IUw3pAyi_Hm4AaX57VfeSiOwaUw,28134
 ragaai_catalyst/tracers/llamaindex_callback.py,sha256=ZY0BJrrlz-P9Mg2dX-ZkVKG3gSvzwqBtk7JL_05MiYA,14028
+ragaai_catalyst/tracers/tracer.py,sha256=Yq2HhgT4785t9573kksJ7ngM3qCLPgZbZ0IpgOHdTTo,19223
 ragaai_catalyst/tracers/upload_traces.py,sha256=2TWdRTN6FMaX-dqDv8BJWQS0xrCGYKkXEYOi2kK3Z3Y,5487
 ragaai_catalyst/tracers/agentic_tracing/README.md,sha256=X4QwLb7-Jg7GQMIXj-SerZIgDETfw-7VgYlczOR8ZeQ,4508
 ragaai_catalyst/tracers/agentic_tracing/__init__.py,sha256=yf6SKvOPSpH-9LiKaoLKXwqj5sez8F_5wkOb91yp0oE,260
@@ -31,7 +32,7 @@ ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py,sha256=--wvhOJ-J
 ragaai_catalyst/tracers/agentic_tracing/tracers/base.py,sha256=88rX7OkOGEyVNECUrc4bYqODyulXve_-99d9ku5hBeQ,37373
 ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py,sha256=l3x3uFO5ov93I7UUrUX1M06WVGy2ug2jEZ1G7o315z4,13075
 ragaai_catalyst/tracers/agentic_tracing/tracers/langgraph_tracer.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=s6BRoBteCRF8XrXGnmZ98ZWPrSONC5RObPXNaq-im3w,31782
+ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=91aWXJGb3GDfyDfJyA7Irnk3XSyfkQaQppW_NMORGJQ,31725
 ragaai_catalyst/tracers/agentic_tracing/tracers/main_tracer.py,sha256=6hsg-Yw11v4qeELI1CWrdX8BXf-wJrTF5smBI5prgoo,15873
 ragaai_catalyst/tracers/agentic_tracing/tracers/network_tracer.py,sha256=m8CxYkl7iMiFya_lNwN1ykBc3Pmo-2pR_2HmpptwHWQ,10352
 ragaai_catalyst/tracers/agentic_tracing/tracers/tool_tracer.py,sha256=4rWL7fIJE5wN0nwh6fMWyh3OrrenZHJkNzyQXikyzQI,13771
@@ -64,8 +65,8 @@ ragaai_catalyst/tracers/utils/__init__.py,sha256=KeMaZtYaTojilpLv65qH08QmpYclfpa
 ragaai_catalyst/tracers/utils/convert_langchain_callbacks_output.py,sha256=ofrNrxf2b1hpjDh_zeaxiYq86azn1MF3kW8-ViYPEg0,1641
 ragaai_catalyst/tracers/utils/langchain_tracer_extraction_logic.py,sha256=cghjCuUe8w-2MZdh9xgtRGe3y219u26GGzpnuY4Wt6Q,3047
 ragaai_catalyst/tracers/utils/utils.py,sha256=ViygfJ7vZ7U0CTSA1lbxVloHp4NSlmfDzBRNCJuMhis,2374
-ragaai_catalyst-2.1.5b6.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ragaai_catalyst-2.1.5b6.dist-info/METADATA,sha256=9wEYmF1UNiH1Py1WJeM_pf5qL71wfyzoXWvUsHtyj6Y,12764
-ragaai_catalyst-2.1.5b6.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ragaai_catalyst-2.1.5b6.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
-ragaai_catalyst-2.1.5b6.dist-info/RECORD,,
+ragaai_catalyst-2.1.5b8.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ragaai_catalyst-2.1.5b8.dist-info/METADATA,sha256=OaiEW7uA1wnQO562QbKGgtlZuue1PTGTjK9-AW5gkLQ,12764
+ragaai_catalyst-2.1.5b8.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ragaai_catalyst-2.1.5b8.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
+ragaai_catalyst-2.1.5b8.dist-info/RECORD,,