ragaai-catalyst 2.1.5b2__py3-none-any.whl → 2.1.5b4__py3-none-any.whl

ragaai_catalyst/ragaai_catalyst.py

@@ -2,7 +2,7 @@ import os
 import logging
 import requests
 from typing import Dict, Optional, Union
-
+import re
 logger = logging.getLogger("RagaAICatalyst")
 
 
@@ -55,10 +55,11 @@ class RagaAICatalyst:
         self.api_keys = api_keys or {}
 
         if base_url:
-            RagaAICatalyst.BASE_URL = base_url
+            RagaAICatalyst.BASE_URL = self._normalize_base_url(base_url)
             try:
+                #set the os.environ["RAGAAI_CATALYST_BASE_URL"] before getting the token as it is used in the get_token method
+                os.environ["RAGAAI_CATALYST_BASE_URL"] = RagaAICatalyst.BASE_URL
                 self.get_token()
-                os.environ["RAGAAI_CATALYST_BASE_URL"] = base_url
             except requests.exceptions.RequestException:
                 raise ConnectionError(
                     "The provided base_url is not accessible. Please re-check the base_url."
@@ -71,6 +72,14 @@ class RagaAICatalyst:
         if self.api_keys:
             self._upload_keys()
 
+    @staticmethod
+    def _normalize_base_url(url):
+        url = re.sub(r'(?<!:)//+', '/', url)  # Ignore the `://` part of URLs and remove extra // if any
+        url = url.rstrip("/") # To remove trailing slashes
+        if not url.endswith("/api"): # To ensure it ends with /api
+            url = f"{url}/api"
+        return url
+
     def _set_access_key_secret_key(self, access_key, secret_key):
         os.environ["RAGAAI_CATALYST_ACCESS_KEY"] = access_key
         os.environ["RAGAAI_CATALYST_SECRET_KEY"] = secret_key

ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py

@@ -150,6 +150,8 @@ 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

@@ -0,0 +1,568 @@
+from typing import Any, Dict, List, Optional, Union, Sequence
+
+import attr
+from langchain.callbacks.base import BaseCallbackHandler
+from langchain.schema import LLMResult, AgentAction, AgentFinish, BaseMessage
+from datetime import datetime
+import json
+import os
+from uuid import UUID
+from functools import wraps
+import asyncio
+from langchain_core.documents import Document
+import logging
+import tempfile
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class LangchainTracer(BaseCallbackHandler):
+    """
+    An enhanced callback handler for LangChain that traces all actions and saves them to a JSON file.
+    Includes improved error handling, async support, and configuration options.
+    """
+
+    def __init__(
+        self,
+        output_path: str = tempfile.gettempdir(),
+        trace_all: bool = True,
+        save_interval: Optional[int] = None,
+        log_level: int = logging.INFO,
+    ):
+        """
+        Initialize the tracer with enhanced configuration options.
+
+        Args:
+            output_path (str): Directory where trace files will be saved
+            trace_all (bool): Whether to trace all components or only specific ones
+            save_interval (Optional[int]): Interval in seconds to auto-save traces
+            log_level (int): Logging level for the tracer
+        """
+        super().__init__()
+        self.output_path = output_path
+        self.trace_all = trace_all
+        self.save_interval = save_interval
+        self._active = False
+        self._original_inits = {}
+        self._original_methods = {}
+        self.additional_metadata = {}
+        self._save_task = None
+        self._current_query = None  # Add this line to track the current query
+        self.filepath = None
+        logger.setLevel(log_level)
+
+        if not os.path.exists(output_path):
+            os.makedirs(output_path)
+
+        self.reset_trace()
+        
+
+    def __enter__(self):
+        """Context manager entry"""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Context manager exit"""
+        
+        self.stop()
+        if exc_type:
+            logger.error(f"Error in context manager: {exc_val}")
+            return False
+        return True
+
+    def reset_trace(self):
+        """Reset the current trace to initial state with enhanced structure"""
+        self.current_trace: Dict[str, Any] = {
+            "start_time": None,
+            "end_time": None,
+            "actions": [],
+            "llm_calls": [],
+            "chain_starts": [],
+            "chain_ends": [],
+            "agent_actions": [],
+            "chat_model_calls": [],
+            "retriever_actions": [],
+            "tokens": [],
+            "errors": [],
+            "query": self._current_query,  # Add this line to include the query in the trace
+            "metadata": {
+                "version": "2.0",
+                "trace_all": self.trace_all,
+                "save_interval": self.save_interval,
+            },
+        }
+
+    async def _periodic_save(self):
+        """Periodically save traces if save_interval is set"""
+        while self._active and self.save_interval:
+            await asyncio.sleep(self.save_interval)
+            await self._async_save_trace()
+
+    async def _async_save_trace(self, force: bool = False):
+        """Asynchronously save the current trace to a JSON file"""
+        if not self.current_trace["start_time"] and not force:
+            return
+
+        try:
+            self.current_trace["end_time"] = datetime.now()
+            
+            # Use the query from the trace or fallback to a default
+            safe_query = self._current_query or "unknown"
+            
+            # Sanitize the query for filename
+            safe_query = ''.join(c for c in safe_query if c.isalnum() or c.isspace())[:50].strip()
+
+            # Add a timestamp to ensure unique filenames
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"langchain_callback_traces.json"
+            filepath = os.path.join(self.output_path, filename)
+            self.filepath = filepath
+
+            trace_to_save = self.current_trace.copy()
+            trace_to_save["start_time"] = str(trace_to_save["start_time"])
+            trace_to_save["end_time"] = str(trace_to_save["end_time"])
+
+            # Save if there are meaningful events or if force is True
+            if (
+                len(trace_to_save["llm_calls"]) > 0
+                or len(trace_to_save["chain_starts"]) > 0
+                or len(trace_to_save["chain_ends"]) > 0
+                or len(trace_to_save["errors"]) > 0
+                or force
+            ):
+                async with asyncio.Lock():
+                    with open(filepath, "w", encoding="utf-8") as f:
+                        json.dump(trace_to_save, f, indent=2, default=str)
+
+                logger.info(f"Trace saved to: {filepath}")
+                
+                # Reset the current query after saving
+                self._current_query = None
+                
+                # Reset the trace
+                self.reset_trace()
+
+        except Exception as e:
+            logger.error(f"Error saving trace: {e}")
+            self.on_error(e, context="save_trace")
+
+    def _save_trace(self, force: bool = False):
+        """Synchronous version of trace saving"""
+        if asyncio.get_event_loop().is_running():
+            asyncio.create_task(self._async_save_trace(force))
+        else:
+            asyncio.run(self._async_save_trace(force))
+
+    def _create_safe_wrapper(self, original_func, component_name):
+        """Create a safely wrapped version of an original function with enhanced error handling"""
+
+        @wraps(original_func)
+        def wrapped(*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 'callbacks' not in kwargs_copy:
+                    kwargs_copy['callbacks'] = [self]
+                elif self not in kwargs_copy['callbacks']:
+                    kwargs_copy['callbacks'].append(self)
+                
+                # 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)
+
+        return wrapped
+
+
+    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),
+        }
+
+        for name, (component, method_name) in components_to_patch.items():
+            try:
+                if method_name == "__init__":
+                    original = component.__init__
+                    self._original_inits[name] = original
+                    component.__init__ = self._create_safe_wrapper(original, 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)
+                        )
+                    else:
+                        wrapped = self._create_safe_wrapper(original, name)
+                    setattr(component, method_name, wrapped)
+                else:
+                    self._original_methods[name] = component
+                    globals()[name] = self._create_safe_wrapper(component, name)
+            except Exception as e:
+                logger.error(f"Error patching {name}: {e}")
+                self.on_error(e, context=f"patch_{name}")
+
+    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
+
+
+        for name, original in self._original_inits.items():
+            try:
+                component = eval(name)
+                component.__init__ = original
+            except Exception as e:
+                logger.error(f"Error restoring {name}: {e}")
+                self.on_error(e, context=f"restore_{name}")
+
+        for name, original in self._original_methods.items():
+            try:
+                if "." in name:
+                    module_name, method_name = name.rsplit(".", 1)
+                    module = eval(module_name)
+                    setattr(module, method_name, original)
+                else:
+                    globals()[name] = original
+            except Exception as e:
+                logger.error(f"Error restoring {name}: {e}")
+                self.on_error(e, context=f"restore_{name}")
+
+    def start(self):
+        """Start tracing with enhanced error handling and async support"""
+        try:
+            self.reset_trace()
+            self.current_trace["start_time"] = datetime.now()
+            self._active = True
+            self._monkey_patch()
+
+            if self.save_interval:
+                loop = asyncio.get_event_loop()
+                self._save_task = loop.create_task(self._periodic_save())
+
+            logger.info("Tracing started")
+        except Exception as e:
+            logger.error(f"Error starting tracer: {e}")
+            self.on_error(e, context="start")
+            raise
+
+    def stop(self):
+        """Stop tracing with enhanced cleanup"""
+        try:
+            self._active = False
+            if self._save_task:
+                self._save_task.cancel()
+            self._restore_original_methods()
+            # self._save_trace(force=True)
+
+            return self.current_trace.copy(), self.additional_metadata
+
+            logger.info("Tracing stopped")
+        except Exception as e:
+            logger.error(f"Error stopping tracer: {e}")
+            self.on_error(e, context="stop")
+            raise
+        finally:
+            self._original_inits.clear()
+            self._original_methods.clear()
+
+    def force_save(self):
+        """Force save the current trace"""
+        self._save_trace(force=True)
+
+    # Callback methods with enhanced error handling and logging
+    def on_llm_start(
+        self,
+        serialized: Dict[str, Any],
+        prompts: List[str],
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            if not self.current_trace["start_time"]:
+                self.current_trace["start_time"] = datetime.now()
+
+            self.current_trace["llm_calls"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "event": "llm_start",
+                    "serialized": serialized,
+                    "prompts": prompts,
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+            )
+        except Exception as e:
+            self.on_error(e, context="llm_start")
+
+    def on_llm_end(self, response: LLMResult, *, run_id: UUID, **kwargs: Any) -> None:
+        try:
+            self.current_trace["llm_calls"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "event": "llm_end",
+                    "response": response.dict(),
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+            )
+
+            end_time = datetime.now()
+            self.additional_metadata["latency"] = (end_time - self.current_trace["start_time"]).total_seconds()
+
+            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)
+        except Exception as e:
+            self.on_error(e, context="llm_end")
+
+    def on_chat_model_start(
+        self,
+        serialized: Dict[str, Any],
+        messages: List[List[BaseMessage]],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            messages_dict = [
+                [
+                    {
+                        "type": msg.type,
+                        "content": msg.content,
+                        "additional_kwargs": msg.additional_kwargs,
+                    }
+                    for msg in batch
+                ]
+                for batch in messages
+            ]
+
+            self.current_trace["chat_model_calls"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "event": "chat_model_start",
+                    "serialized": serialized,
+                    "messages": messages_dict,
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+            )
+        except Exception as e:
+            self.on_error(e, context="chat_model_start")
+
+    def on_chain_start(
+        self,
+        serialized: Dict[str, Any],
+        inputs: Dict[str, Any],
+        *,
+        run_id: UUID,
+        **kwargs: Any,
+    ) -> None:
+        try:
+            context = ""
+            query = ""
+            if isinstance(inputs, dict):
+                if "context" in inputs:
+                    if isinstance(inputs["context"], Document):
+                        context = inputs["context"].page_content
+                    elif isinstance(inputs["context"], list):
+                        context = "\n".join(
+                            doc.page_content if isinstance(doc, Document) else str(doc)
+                            for doc in inputs["context"]
+                        )
+                    elif isinstance(inputs["context"], str):
+                        context = inputs["context"]
+
+                query = inputs.get("question", inputs.get("input", ""))
+                
+                # Set the current query
+                self._current_query = query
+                
+                chain_event = {
+                    "timestamp": datetime.now(),
+                    "serialized": serialized,
+                    "context": context,
+                    "query": inputs.get("question", inputs.get("input", "")),
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+
+                self.current_trace["chain_starts"].append(chain_event)
+        except Exception as e:
+            self.on_error(e, context="chain_start")
+
+    def on_chain_end(
+        self, outputs: Dict[str, Any], *, run_id: UUID, **kwargs: Any
+    ) -> None:
+        try:
+            self.current_trace["chain_ends"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "outputs": outputs,
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+            )
+        except Exception as e:
+            self.on_error(e, context="chain_end")
+
+    def on_agent_action(self, action: AgentAction, run_id: UUID, **kwargs: Any) -> None:
+        try:
+            self.current_trace["agent_actions"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "action": action.dict(),
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+            )
+        except Exception as e:
+            self.on_error(e, context="agent_action")
+
+    def on_agent_finish(self, finish: AgentFinish, run_id: UUID, **kwargs: Any) -> None:
+        try:
+            self.current_trace["agent_actions"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "event": "agent_finish",
+                    "finish": finish.dict(),
+                    "run_id": str(run_id),
+                    "additional_kwargs": kwargs,
+                }
+            )
+        except Exception as e:
+            self.on_error(e, context="agent_finish")
+
+    def on_retriever_start(
+        self, serialized: Dict[str, Any], query: str, *, run_id: UUID, **kwargs: Any
+    ) -> None:
+        try:
+            retriever_event = {
+                "timestamp": datetime.now(),
+                "event": "retriever_start",
+                "serialized": serialized,
+                "query": query,
+                "run_id": str(run_id),
+                "additional_kwargs": kwargs,
+            }
+
+            self.current_trace["retriever_actions"].append(retriever_event)
+        except Exception as e:
+            self.on_error(e, context="retriever_start")
+
+    def on_retriever_end(
+        self, documents: Sequence[Document], *, run_id: UUID, **kwargs: Any
+    ) -> None:
+        try:
+            processed_documents = [
+                {"page_content": doc.page_content, "metadata": doc.metadata}
+                for doc in documents
+            ]
+
+            retriever_event = {
+                "timestamp": datetime.now(),
+                "event": "retriever_end",
+                "documents": processed_documents,
+                "run_id": str(run_id),
+                "additional_kwargs": kwargs,
+            }
+
+            self.current_trace["retriever_actions"].append(retriever_event)
+        except Exception as e:
+            self.on_error(e, context="retriever_end")
+
+    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        try:
+            self.current_trace["tokens"].append(
+                {
+                    "timestamp": datetime.now(),
+                    "event": "new_token",
+                    "token": token,
+                    "additional_kwargs": kwargs,
+                }
+            )
+        except Exception as e:
+            self.on_error(e, context="llm_new_token")
+
+    def on_error(self, error: Exception, context: str = "", **kwargs: Any) -> None:
+        """Enhanced error handling with context"""
+        try:
+            error_event = {
+                "timestamp": datetime.now(),
+                "error": str(error),
+                "error_type": type(error).__name__,
+                "context": context,
+                "additional_kwargs": kwargs,
+            }
+            self.current_trace["errors"].append(error_event)
+            logger.error(f"Error in {context}: {error}")
+        except Exception as e:
+            logger.critical(f"Error in error handler: {e}")
+
+    def on_chain_error(self, error: Exception, **kwargs: Any) -> None:
+        self.on_error(error, context="chain", **kwargs)
+
+    def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        self.on_error(error, context="llm", **kwargs)
+
+    def on_tool_error(self, error: Exception, **kwargs: Any) -> None:
+        self.on_error(error, context="tool", **kwargs)
+
+    def on_retriever_error(self, error: Exception, **kwargs: Any) -> None:
+        self.on_error(error, context="retriever", **kwargs)

ragaai_catalyst/tracers/tracer.py

@@ -1,4 +1,6 @@
+from audioop import add
 import os
+import uuid
 import datetime
 import logging
 import asyncio
@@ -6,6 +8,13 @@ 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
@@ -118,6 +127,7 @@ class Tracer(AgenticTracing):
         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
@@ -152,11 +162,12 @@ 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
@@ -239,11 +250,12 @@ 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()
@@ -254,17 +266,74 @@ 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
+
+            # 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"]
+            additional_metadata["total_cost"] = additional_metadata["cost"]["total_cost"]
+
+            del additional_metadata["tokens"]
+            del additional_metadata["cost"]
+            
+            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)
             
-            return "Trace upload initiated. Use get_upload_status() to check the status."
+            
+            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()

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,15 @@ class UploadTraces:
             "vector_store":{"columnType":"pipeline"},
             "feedback": {"columnType":"feedBack"}
         }
+
+        if additional_metadata_keys:
+            for key in additional_metadata_keys:
+                SCHEMA_MAPPING_NEW[key] = {"columnType": "metadata"}
+
+        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 +128,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/langchain_tracer_extraction_logic.py

@@ -0,0 +1,81 @@
+import json
+import uuid
+
+def langchain_tracer_extraction(data):
+    trace_aggregate = {}
+    import uuid
+
+    def generate_trace_id():
+        """
+        Generate a random trace ID using UUID4.
+        Returns a string representation of the UUID with no hyphens.
+        """
+        return '0x'+str(uuid.uuid4()).replace('-', '')
+
+    trace_aggregate["tracer_type"] = "langchain"
+    trace_aggregate['trace_id'] = generate_trace_id()
+    trace_aggregate['session_id'] = None
+    trace_aggregate["pipeline"] = {
+        'llm_model': 'gpt-3.5-turbo', 
+        'vector_store': 'faiss',
+        'embed_model': 'text-embedding-ada-002'
+        }
+    trace_aggregate["metadata"] = {
+        'key1': 'value1',
+        'key2': 'value2',
+        'log_source': 'langchain_tracer',
+        'recorded_on': '2024-06-14 08:57:27.324410'
+        }
+    trace_aggregate["prompt_length"] = 0
+    trace_aggregate["data"] = {}
+
+    def get_prompt(data):
+        # if "chain_starts" in data and data["chain_starts"] != []:
+        #     for item in data["chain_starts"]:
+
+        if "chat_model_calls" in data and data["chat_model_calls"] != []:
+            for item in data["chat_model_calls"]:
+                messages = item["messages"][0]
+                for message in messages:
+                    if message["type"]=="human":
+                        human_messages = message["content"].strip()
+                        return human_messages
+        if  "llm_calls" in data and data["llm_calls"] != []:
+            if "llm_start" in data["llm_calls"][0]["event"]:
+                for item in data["llm_calls"]:
+                    prompt = item["prompts"]
+                    return prompt[0].strip()
+
+    def get_response(data):
+        for item in data["llm_calls"]:
+            if item["event"] == "llm_end":
+                # import pdb; pdb.set_trace()
+                llm_end_responses = item["response"]["generations"][0]
+                for llm_end_response in llm_end_responses:
+                    response = llm_end_response["text"]
+                return response.strip()
+
+    def get_context(data):
+        if "retriever_actions" in data and data["retriever_actions"] != []:
+            for item in data["retriever_actions"]:
+                if item["event"] == "retriever_end":
+                    context = item["documents"][0]["page_content"].replace('\n', ' ')
+                    return context
+        if "chat_model_calls" in data and data["chat_model_calls"] != []:
+            for item in data["chat_model_calls"]:
+                messages = item["messages"][0]
+                for message in messages:
+                    if message["type"]=="system":
+                        content = message["content"].strip().replace('\n', ' ')
+                        return content
+
+
+    prompt = get_prompt(data)
+    response = get_response(data)
+    context = get_context(data)
+
+    trace_aggregate["data"]["prompt"]=prompt
+    trace_aggregate["data"]["response"]=response
+    trace_aggregate["data"]["context"]=context
+
+    return trace_aggregate

{ragaai_catalyst-2.1.5b2.dist-info → ragaai_catalyst-2.1.5b4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: ragaai_catalyst
-Version: 2.1.5b2
+Version: 2.1.5b4
 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.5b2.dist-info → ragaai_catalyst-2.1.5b4.dist-info}/RECORD

@@ -8,14 +8,15 @@ ragaai_catalyst/guardrails_manager.py,sha256=DILMOAASK57FH9BLq_8yC1AQzRJ8McMFLwC
 ragaai_catalyst/internal_api_completion.py,sha256=DdICI5yfEudiOAIC8L4oxH0Qz7kX-BZCdo9IWsi2gNo,2965
 ragaai_catalyst/prompt_manager.py,sha256=W8ypramzOprrJ7-22d5vkBXIuIQ8v9XAzKDGxKsTK28,16550
 ragaai_catalyst/proxy_call.py,sha256=CHxldeceZUaLU-to_hs_Kf1z_b2vHMssLS_cOBedu78,5499
-ragaai_catalyst/ragaai_catalyst.py,sha256=FdqMzwuQLqS2-3JJDsTQ8uh2itllOxfPrRUjb8Kwmn0,17428
+ragaai_catalyst/ragaai_catalyst.py,sha256=5nVg3_-lcvhrXjNkPTeGhe3tdUjm_4ZIctOcqWXBkRA,17939
 ragaai_catalyst/synthetic_data_generation.py,sha256=uDV9tNwto2xSkWg5XHXUvjErW-4P34CTrxaJpRfezyA,19250
 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/llamaindex_callback.py,sha256=ZY0BJrrlz-P9Mg2dX-ZkVKG3gSvzwqBtk7JL_05MiYA,14028
-ragaai_catalyst/tracers/tracer.py,sha256=S_ANRm5zSMvQiUyQTRwyUepFci_T3AN26wAOXoURfyc,15648
-ragaai_catalyst/tracers/upload_traces.py,sha256=mT5rverNUL5Rcal9VR5_c75wHBAUrm2pvYetTZqP3ok,4796
+ragaai_catalyst/tracers/tracer.py,sha256=k2HjH6ONaabbPvoX6xJRck-A2l-9GVW7Nueimuu-Ua8,19096
+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
 ragaai_catalyst/tracers/agentic_tracing/data/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -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=91aWXJGb3GDfyDfJyA7Irnk3XSyfkQaQppW_NMORGJQ,31725
+ragaai_catalyst/tracers/agentic_tracing/tracers/llm_tracer.py,sha256=s6BRoBteCRF8XrXGnmZ98ZWPrSONC5RObPXNaq-im3w,31782
 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
@@ -61,9 +62,11 @@ ragaai_catalyst/tracers/instrumentators/langchain.py,sha256=yMN0qVF0pUVk6R5M1vJo
 ragaai_catalyst/tracers/instrumentators/llamaindex.py,sha256=SMrRlR4xM7k9HK43hakE8rkrWHxMlmtmWD-AX6TeByc,416
 ragaai_catalyst/tracers/instrumentators/openai.py,sha256=14R4KW9wQCR1xysLfsP_nxS7cqXrTPoD8En4MBAaZUU,379
 ragaai_catalyst/tracers/utils/__init__.py,sha256=KeMaZtYaTojilpLv65qH08QmpYclfpacDA0U3wg6Ybw,64
+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.5b2.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-ragaai_catalyst-2.1.5b2.dist-info/METADATA,sha256=OtU5W4jpT4K2GVF82w9jYeaVglZOWOIwpPfXJNN_SmM,12764
-ragaai_catalyst-2.1.5b2.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-ragaai_catalyst-2.1.5b2.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
-ragaai_catalyst-2.1.5b2.dist-info/RECORD,,
+ragaai_catalyst-2.1.5b4.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+ragaai_catalyst-2.1.5b4.dist-info/METADATA,sha256=9KXrmDoxY6TV9BL3TXWha0rm4n4wXJy03sTW7dc4ZRU,12764
+ragaai_catalyst-2.1.5b4.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+ragaai_catalyst-2.1.5b4.dist-info/top_level.txt,sha256=HpgsdRgEJMk8nqrU6qdCYk3di7MJkDL0B19lkc7dLfM,16
+ragaai_catalyst-2.1.5b4.dist-info/RECORD,,