ragaai-catalyst 2.1b0__py3-none-any.whl → 2.1b1__py3-none-any.whl

ragaai_catalyst/tracers/agentic_tracing/upload_agentic_traces.py

@@ -0,0 +1,187 @@
+import requests
+import json
+import os
+from datetime import datetime
+
+
+class UploadAgenticTraces:
+    def __init__(self, 
+                 json_file_path,
+                 project_name,
+                 project_id,
+                 dataset_name,
+                 user_detail,
+                 base_url):
+        self.json_file_path = json_file_path
+        self.project_name = project_name
+        self.project_id = project_id
+        self.dataset_name = dataset_name
+        self.user_detail = user_detail
+        self.base_url = base_url
+        self.timeout = 99999
+
+    def _create_dataset_schema_with_trace(self):
+        SCHEMA_MAPPING_NEW = {
+            "trace_id": {"columnType": "traceId"},
+            "trace_uri": {"columnType": "traceUri"},
+            "prompt": {"columnType": "prompt"},
+            "response":{"columnType": "response"},
+            "context": {"columnType": "context"},
+            "llm_model": {"columnType":"pipeline"},
+            "recorded_on": {"columnType": "metadata"},
+            "embed_model": {"columnType":"pipeline"},
+            "log_source": {"columnType": "metadata"},
+            "vector_store":{"columnType":"pipeline"},
+            "feedback": {"columnType":"feedBack"}
+        }
+        def make_request():
+            headers = {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "X-Project-Name": self.project_name,
+            }
+            payload = json.dumps({
+                "datasetName": self.dataset_name,
+                # "schemaMapping": SCHEMA_MAPPING_NEW,
+                "traceFolderUrl": None,
+            })
+            response = requests.request("POST",
+                f"{self.base_url}/v1/llm/dataset/logs",
+                headers=headers,
+                data=payload,
+                timeout=self.timeout
+            )
+
+            return response
+
+        response = make_request()
+
+        if response.status_code == 401:
+            # get_token()  # Fetch a new token and set it in the environment
+            response = make_request()  # Retry the request
+        if response.status_code != 200:
+            return response.status_code
+        return response.status_code
+
+
+    def _get_presigned_url(self):
+        payload = json.dumps({
+                "datasetName": self.dataset_name,
+                "numFiles": 1,
+            })
+        headers = {
+            "Content-Type": "application/json",
+            "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+            "X-Project-Name": self.project_name,
+        }
+
+        try:
+            response = requests.request("GET", 
+                                        f"{self.base_url}/v1/llm/presigned-url", 
+                                        headers=headers, 
+                                        data=payload,
+                                        timeout=self.timeout)
+            if response.status_code == 200:
+                presignedUrls = response.json()["data"]["presignedUrls"][0]
+                return presignedUrls
+        except requests.exceptions.RequestException as e:
+            print(f"Error while getting presigned url: {e}")
+            return None
+
+    def _put_presigned_url(self, presignedUrl, filename):
+        headers = {
+                "Content-Type": "application/json",
+            }
+
+        if "blob.core.windows.net" in presignedUrl:  # Azure
+            headers["x-ms-blob-type"] = "BlockBlob"
+        print(f"Uploading agentic traces...")
+        try:
+            with open(filename) as f:
+                payload = f.read().replace("\n", "").replace("\r", "").encode()
+        except Exception as e:
+            print(f"Error while reading file: {e}")
+            return None
+        try:
+            response = requests.request("PUT", 
+                                        presignedUrl, 
+                                        headers=headers, 
+                                        data=payload,
+                                        timeout=self.timeout)
+            if response.status_code != 200 or response.status_code != 201:
+                return response, response.status_code
+        except requests.exceptions.RequestException as e:
+            print(f"Error while uploading to presigned url: {e}")
+            return None
+
+    def insert_traces(self, presignedUrl):
+        headers = {
+                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "Content-Type": "application/json",
+                "X-Project-Name": self.project_name,
+            }
+        payload = json.dumps({
+                "datasetName": self.dataset_name,
+                "presignedUrl": presignedUrl,
+                "datasetSpans": self._get_dataset_spans(), #Extra key for agentic traces
+            })
+        try:
+            response = requests.request("POST", 
+                                        f"{self.base_url}/v1/llm/insert/trace", 
+                                        headers=headers, 
+                                        data=payload,
+                                        timeout=self.timeout)
+        except requests.exceptions.RequestException as e:
+            print(f"Error while inserting traces: {e}")
+            return None
+
+    def _get_dataset_spans(self):
+        try:
+            with open(self.json_file_path) as f:
+                data = json.load(f)
+        except Exception as e:
+            print(f"Error while reading file: {e}")
+            return None
+        try:
+            spans = data["data"][0]["spans"]
+            datasetSpans = []
+            for span in spans:
+                if span["type"] != "agent":
+                    existing_span = next((s for s in datasetSpans if s["spanHash"] == span["hash_id"]), None)
+                    if existing_span is None:
+                        datasetSpans.append({
+                            "spanId": span["id"],
+                            "spanName": span["name"],
+                            "spanHash": span["hash_id"],
+                            "spanType": span["type"],
+                        })
+                else:
+                    datasetSpans.append({
+                                "spanId": span["id"],
+                                "spanName": span["name"],
+                                "spanHash": span["hash_id"],
+                                "spanType": span["type"],
+                            })
+                    children = span["data"]["children"]
+                    for child in children:
+                        existing_span = next((s for s in datasetSpans if s["spanHash"] == child["hash_id"]), None)
+                        if existing_span is None:
+                            datasetSpans.append({
+                                "spanId": child["id"],
+                                "spanName": child["name"],
+                                "spanHash": child["hash_id"],
+                                "spanType": child["type"],
+                            })
+            return datasetSpans
+        except Exception as e:
+            print(f"Error while reading dataset spans: {e}")
+            return None
+    
+    def upload_agentic_traces(self):
+        self._create_dataset_schema_with_trace()
+        presignedUrl = self._get_presigned_url()
+        if presignedUrl is None:
+            return
+        self._put_presigned_url(presignedUrl, self.json_file_path)
+        self.insert_traces(presignedUrl)
+        print("Agentic Traces uploaded")

ragaai_catalyst/tracers/agentic_tracing/upload_code.py

@@ -0,0 +1,115 @@
+from aiohttp import payload
+import requests
+import json
+import os
+import logging
+logger = logging.getLogger(__name__)
+
+def upload_code(hash_id, zip_path, project_name, dataset_name):
+    code_hashes_list = _fetch_dataset_code_hashes(project_name, dataset_name)
+
+    if hash_id not in code_hashes_list:
+        presigned_url = _fetch_presigned_url(project_name, dataset_name)
+        _put_zip_presigned_url(project_name, presigned_url, zip_path)
+
+        response = _insert_code(dataset_name, hash_id, presigned_url, project_name)
+        return response
+    else:
+        return "Code already exists"
+
+def _fetch_dataset_code_hashes(project_name, dataset_name):
+    payload = {}
+    headers = {
+        "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+        "X-Project-Name": project_name,
+    }
+
+    try:
+        response = requests.request("GET", 
+                                    f"{os.getenv('RAGAAI_CATALYST_BASE_URL')}/v2/llm/dataset/code?datasetName={dataset_name}", 
+                                    headers=headers, 
+                                    data=payload,
+                                    timeout=99999)
+
+        if response.status_code == 200:
+            return response.json()["data"]["codeHashes"]
+        else:
+            raise Exception(f"Failed to fetch code hashes: {response.json()['message']}")
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to list datasets: {e}")
+        raise 
+
+def _fetch_presigned_url(project_name, dataset_name):
+    payload = json.dumps({
+            "datasetName": dataset_name,
+            "numFiles": 1,
+            "contentType": "application/zip"
+            })
+
+    headers = {
+        "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+        "Content-Type": "application/json",
+        "X-Project-Name": project_name,
+    }
+
+    try:
+        response = requests.request("GET", 
+                                    f"{os.getenv('RAGAAI_CATALYST_BASE_URL')}/v1/llm/presigned-url", 
+                                    headers=headers, 
+                                    data=payload,
+                                    timeout=99999)
+
+        if response.status_code == 200:
+            return response.json()["data"]["presignedUrls"][0]
+        else:
+            raise Exception(f"Failed to fetch code hashes: {response.json()['message']}")
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to list datasets: {e}")
+        raise
+
+def _put_zip_presigned_url(project_name, presignedUrl, filename):
+    headers = {
+            "X-Project-Name": project_name,
+            "Content-Type": "application/zip",
+        }
+
+    if "blob.core.windows.net" in presignedUrl:  # Azure
+        headers["x-ms-blob-type"] = "BlockBlob"
+    print(f"Uploading code...")
+    with open(filename, 'rb') as f:
+        payload = f.read()
+
+    response = requests.request("PUT", 
+                                presignedUrl, 
+                                headers=headers, 
+                                data=payload,
+                                timeout=99999)
+    if response.status_code != 200 or response.status_code != 201:
+        return response, response.status_code
+
+def _insert_code(dataset_name, hash_id, presigned_url, project_name):
+    payload = json.dumps({
+        "datasetName": dataset_name,
+        "codeHash": hash_id,
+        "presignedUrl": presigned_url
+        })
+    
+    headers = {
+        'X-Project-Name': project_name,
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {os.getenv("RAGAAI_CATALYST_TOKEN")}'
+        }
+    
+    try:
+        response = requests.request("POST", 
+                                    f"{os.getenv('RAGAAI_CATALYST_BASE_URL')}/v2/llm/dataset/code", 
+                                    headers=headers, 
+                                    data=payload,
+                                    timeout=99999)
+        if response.status_code == 200:
+            return response.json()["message"]
+        else:
+            raise Exception(f"Failed to insert code: {response.json()['message']}")
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to insert code: {e}")
+        raise

ragaai_catalyst/tracers/agentic_tracing/user_interaction_tracer.py

@@ -0,0 +1,43 @@
+import builtins
+from datetime import datetime
+import contextvars
+import inspect
+import uuid
+
+class UserInteractionTracer:
+    def __init__(self, *args, **kwargs):
+        self.project_id = contextvars.ContextVar("project_id", default=None)
+        self.trace_id = contextvars.ContextVar("trace_id", default=None)
+        self.tracer = contextvars.ContextVar("tracer", default=None)
+        self.component_id = contextvars.ContextVar("component_id", default=None)
+        self.original_input = builtins.input
+        self.original_print = builtins.print
+        self.interactions = []
+
+    def traced_input(self, prompt=""):
+        # Get caller information
+        if prompt:
+            self.traced_print(prompt, end="")
+        try:
+            content = self.original_input()
+        except EOFError:
+            content = ""  # Return empty string on EOF
+            
+        self.interactions.append({
+            "id": str(uuid.uuid4()),
+            "interaction_type": "input",
+            "content": content,
+            "timestamp": datetime.now().isoformat()
+        })
+        return content
+
+    def traced_print(self, *args, **kwargs):
+        content = " ".join(str(arg) for arg in args)
+        
+        self.interactions.append({
+            "id": str(uuid.uuid4()),
+            "interaction_type": "output",
+            "content": content,
+            "timestamp": datetime.now().isoformat()
+        })
+        return self.original_print(*args, **kwargs)

ragaai_catalyst/tracers/agentic_tracing/utils/__init__.py

@@ -0,0 +1,3 @@
+from .generic import get_db_path
+
+__all__ = ["get_db_path"]

ragaai_catalyst/tracers/agentic_tracing/utils/api_utils.py

@@ -0,0 +1,18 @@
+import requests
+
+def fetch_analysis_trace(base_url, trace_id):
+    """
+    Fetches the analysis trace data from the server.
+
+    :param base_url: The base URL of the server (e.g., "http://localhost:3000").
+    :param trace_id: The ID of the trace to fetch.
+    :return: The JSON response from the server if successful, otherwise None.
+    """
+    try:
+        url = f"{base_url}/api/analysis_traces/{trace_id}"
+        response = requests.get(url)
+        response.raise_for_status()  # Raise an error for bad responses (4xx, 5xx)
+        return response.json()
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching analysis trace: {e}")
+        return None

ragaai_catalyst/tracers/agentic_tracing/utils/data_classes.py

@@ -0,0 +1,61 @@
+from dataclasses import dataclass, field
+from typing import Dict, List, Any, Optional
+
+
+@dataclass
+class ProjectInfo:
+    project_name: str
+    start_time: float
+    end_time: float = field(default=0)
+    duration: float = field(default=0)
+    total_cost: float = field(default=0)
+    total_tokens: int = field(default=0)
+
+
+@dataclass
+class SystemInfo:
+    project_id: int
+    os_name: str
+    os_version: str
+    python_version: str
+    cpu_info: str
+    memory_total: float
+    installed_packages: str
+
+
+@dataclass
+class LLMCall:
+    name: str
+    model_name: str
+    input_prompt: str
+    output_response: str
+    tool_call: Dict
+    token_usage: Dict[str, int]
+    cost: Dict[str, float]
+    start_time: float = field(default=0)
+    end_time: float = field(default=0)
+    duration: float = field(default=0)
+
+
+@dataclass
+class ToolCall:
+    name: str
+    input_parameters: Dict[str, Any]
+    output: Any
+    start_time: float
+    end_time: float
+    duration: float
+    errors: Optional[str] = None
+
+
+@dataclass
+class AgentCall:
+    name: str
+    input_parameters: Dict[str, Any]
+    output: Any
+    start_time: float
+    end_time: float
+    duration: float
+    tool_calls: List[Dict[str, Any]]
+    llm_calls: List[Dict[str, Any]]
+    errors: Optional[str] = None

ragaai_catalyst/tracers/agentic_tracing/utils/generic.py

@@ -0,0 +1,32 @@
+import os
+import logging
+
+
+def get_db_path():
+    db_filename = "trace_data.db"
+
+    # First, try the package directory
+    package_dir = os.path.dirname(os.path.abspath(__file__))
+    public_dir = os.path.join(package_dir, "..", "ui", "dist")
+    package_db_path = os.path.join(public_dir, db_filename)
+
+    # Ensure the directory exists
+    os.makedirs(os.path.dirname(package_db_path), exist_ok=True)
+
+    if os.path.exists(os.path.dirname(package_db_path)):
+        logging.debug(f"Using package database: {package_db_path}")
+        return f"sqlite:///{package_db_path}"
+
+    # Then, try the local directory
+    local_db_path = os.path.join(os.getcwd(), "agentneo", "ui", "dist", db_filename)
+    if os.path.exists(os.path.dirname(local_db_path)):
+        logging.debug(f"Using local database: {local_db_path}")
+        return f"sqlite:///{local_db_path}"
+
+    # Finally, try the local "/dist" directory
+    local_dist_path = os.path.join(os.getcwd(), "dist", db_filename)
+    if os.path.exists(os.path.dirname(local_dist_path)):
+        logging.debug(f"Using local database: {local_dist_path}")
+        return f"sqlite:///{local_dist_path}"
+
+    return f"sqlite:///{package_db_path}"

ragaai_catalyst/tracers/agentic_tracing/utils/llm_utils.py

@@ -0,0 +1,177 @@
+from .data_classes import LLMCall
+from .trace_utils import (
+    calculate_cost,
+    convert_usage_to_dict,
+    load_model_costs,
+)
+from importlib import resources
+import json
+import os
+
+
+# Load the Json configuration
+try:
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    model_costs_path = os.path.join(current_dir, "model_costs.json")
+    with open(model_costs_path, "r") as file:
+        config = json.load(file)
+except FileNotFoundError:
+    from importlib.resources import files
+    with (files("") / "model_costs.json").open("r") as file:
+        config = json.load(file)
+
+
+
+def extract_llm_output(result):
+    # Initialize variables
+    model_name = None
+    output_response = ""
+    function_call = None
+    tool_call = None
+    token_usage = {}
+    cost = {}
+
+    # Try to get model_name from result or result.content
+    model_name = None
+    if hasattr(result, "model"):
+        model_name = result.model
+    elif hasattr(result, "content"):
+        try:
+            content_dict = json.loads(result.content)
+            model_name = content_dict.get("model", None)
+        except (json.JSONDecodeError, TypeError):
+            model_name = None
+
+    # Try to get choices from result or result.content
+    choices = None
+    if hasattr(result, "choices"):
+        choices = result.choices
+    elif hasattr(result, "content"):
+        try:
+            content_dict = json.loads(result.content)
+            choices = content_dict.get("choices", None)
+        except (json.JSONDecodeError, TypeError):
+            choices = None
+
+    if choices and len(choices) > 0:
+        first_choice = choices[0]
+
+        # Get message or text
+        message = None
+        if hasattr(first_choice, "message"):
+            message = first_choice.message
+        elif isinstance(first_choice, dict) and "message" in first_choice:
+            message = first_choice["message"]
+
+        if message:
+            # For chat completion
+            # Get output_response
+            if hasattr(message, "content"):
+                output_response = message.content
+            elif isinstance(message, dict) and "content" in message:
+                output_response = message["content"]
+
+            # Get function_call
+            if hasattr(message, "function_call"):
+                function_call = message.function_call
+            elif isinstance(message, dict) and "function_call" in message:
+                function_call = message["function_call"]
+
+            # Get tool_calls (if any)
+            if hasattr(message, "tool_calls"):
+                tool_call = message.tool_calls
+            elif isinstance(message, dict) and "tool_calls" in message:
+                tool_call = message["tool_calls"]
+        else:
+            # For completion
+            # Get output_response
+            if hasattr(first_choice, "text"):
+                output_response = first_choice.text
+            elif isinstance(first_choice, dict) and "text" in first_choice:
+                output_response = first_choice["text"]
+            else:
+                output_response = ""
+
+            # No message, so no function_call or tool_call
+            function_call = None
+            tool_call = None
+    else:
+        output_response = ""
+        function_call = None
+        tool_call = None
+
+    # Set tool_call to function_call if tool_call is None
+    if not tool_call:
+        tool_call = function_call
+
+    # Parse tool_call
+    parsed_tool_call = None
+    if tool_call:
+        if isinstance(tool_call, dict):
+            arguments = tool_call.get("arguments", "{}")
+            name = tool_call.get("name", "")
+        else:
+            # Maybe it's an object with attributes
+            arguments = getattr(tool_call, "arguments", "{}")
+            name = getattr(tool_call, "name", "")
+        try:
+            if isinstance(arguments, str):
+                arguments = json.loads(arguments)
+            else:
+                arguments = arguments  # If already a dict
+        except json.JSONDecodeError:
+            arguments = {}
+        parsed_tool_call = {"arguments": arguments, "name": name}
+
+    # Try to get token_usage from result.usage or result.content
+    usage = None
+    if hasattr(result, "usage"):
+        usage = result.usage
+    elif hasattr(result, "content"):
+        try:
+            content_dict = json.loads(result.content)
+            usage = content_dict.get("usage", {})
+        except (json.JSONDecodeError, TypeError):
+            usage = {}
+    else:
+        usage = {}
+
+    token_usage = convert_usage_to_dict(usage)
+
+    # Load model costs
+    model_costs = load_model_costs()
+
+    # Calculate cost
+    if model_name in model_costs:
+        model_config = model_costs[model_name]
+        input_cost_per_token = model_config.get("input_cost_per_token", 0.0)
+        output_cost_per_token = model_config.get("output_cost_per_token", 0.0)
+        reasoning_cost_per_token = model_config.get(
+            "reasoning_cost_per_token", output_cost_per_token
+        )
+    else:
+        # Default costs or log a warning
+        print(
+            f"Warning: Model '{model_name}' not found in config. Using default costs."
+        )
+        input_cost_per_token = 0.0
+        output_cost_per_token = 0.0
+        reasoning_cost_per_token = 0.0
+
+    cost = calculate_cost(
+        token_usage,
+        input_cost_per_token=input_cost_per_token,
+        output_cost_per_token=output_cost_per_token,
+        reasoning_cost_per_token=reasoning_cost_per_token,
+    )
+
+    llm_data = LLMCall(
+        name="",
+        model_name=model_name,
+        input_prompt="",  # Not available here
+        output_response=output_response,
+        token_usage=token_usage,
+        cost=cost,
+        tool_call=parsed_tool_call,
+    )
+    return llm_data