ragaai-catalyst 2.1.4.1b1__py3-none-any.whl → 2.1.5b1__py3-none-any.whl

ragaai_catalyst/__init__.py

@@ -1,6 +1,5 @@
 from .experiment import Experiment
 from .ragaai_catalyst import RagaAICatalyst
-from .tracers import Tracer
 from .utils import response_checker
 from .dataset import Dataset
 from .prompt_manager import PromptManager
@@ -8,6 +7,24 @@ from .evaluation import Evaluation
 from .synthetic_data_generation import SyntheticDataGeneration
 from .guardrails_manager import GuardrailsManager
 from .guard_executor import GuardExecutor
+from .tracers import Tracer, init_tracing, trace_agent, trace_llm, trace_tool, current_span, trace_custom
 
 
-__all__ = ["Experiment", "RagaAICatalyst", "Tracer", "PromptManager", "Evaluation","SyntheticDataGeneration", "GuardrailsManager"]
+
+
+__all__ = [
+    "Experiment", 
+    "RagaAICatalyst", 
+    "Tracer", 
+    "PromptManager", 
+    "Evaluation",
+    "SyntheticDataGeneration", 
+    "GuardrailsManager", 
+    "GuardExecutor",
+    "init_tracing", 
+    "trace_agent", 
+    "trace_llm",
+    "trace_tool",
+    "current_span",
+    "trace_custom"
+]

ragaai_catalyst/dataset.py

@@ -1,4 +1,5 @@
 import os
+import json
 import requests
 from .utils import response_checker
 from typing import Union
@@ -271,3 +272,332 @@ class Dataset:
         except Exception as e:
             logger.error(f"Error in create_from_csv: {e}")
             raise
+
+    def add_rows(self, csv_path, dataset_name):
+        """
+        Add rows to an existing dataset from a CSV file.
+
+        Args:
+            csv_path (str): Path to the CSV file to be added
+            dataset_name (str): Name of the existing dataset to add rows to
+
+        Raises:
+            ValueError: If dataset does not exist or columns are incompatible
+        """
+        # Get existing dataset columns
+        existing_columns = self.get_dataset_columns(dataset_name)
+
+        # Read the CSV file to check columns
+        try:
+            import pandas as pd
+            df = pd.read_csv(csv_path)
+            csv_columns = df.columns.tolist()
+        except Exception as e:
+            logger.error(f"Failed to read CSV file: {e}")
+            raise ValueError(f"Unable to read CSV file: {e}")
+
+        # Check column compatibility
+        for column in existing_columns:
+            if column not in csv_columns:
+                df[column] = None  
+
+        # Get presigned URL for the CSV
+        def get_presignedUrl():
+            headers = {
+                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "X-Project-Id": str(self.project_id),
+            }
+            try:
+                response = requests.get(
+                    f"{Dataset.BASE_URL}/v2/llm/dataset/csv/presigned-url",
+                    headers=headers,
+                    timeout=Dataset.TIMEOUT,
+                )
+                response.raise_for_status()
+                return response.json()
+            except requests.exceptions.RequestException as e:
+                logger.error(f"Failed to get presigned URL: {e}")
+                raise
+
+        try:
+            presignedUrl = get_presignedUrl()
+            if presignedUrl['success']:
+                url = presignedUrl['data']['presignedUrl']
+                filename = presignedUrl['data']['fileName']
+            else:
+                raise ValueError('Unable to fetch presignedUrl')
+        except Exception as e:
+            logger.error(f"Error in get_presignedUrl: {e}")
+            raise
+
+        # Upload CSV to presigned URL
+        def put_csv_to_presignedUrl(url):
+            headers = {
+                'Content-Type': 'text/csv',
+                'x-ms-blob-type': 'BlockBlob',
+            }
+            try:
+                with open(csv_path, 'rb') as file:
+                    response = requests.put(
+                        url,
+                        headers=headers,
+                        data=file,
+                        timeout=Dataset.TIMEOUT,
+                    )
+                    response.raise_for_status()
+                    return response
+            except requests.exceptions.RequestException as e:
+                logger.error(f"Failed to put CSV to presigned URL: {e}")
+                raise
+
+        try:
+            put_csv_response = put_csv_to_presignedUrl(url)
+            if put_csv_response.status_code not in (200, 201):
+                raise ValueError('Unable to put csv to the presignedUrl')
+        except Exception as e:
+            logger.error(f"Error in put_csv_to_presignedUrl: {e}")
+            raise
+
+        # Prepare schema mapping (assuming same mapping as original dataset)
+        def generate_schema_mapping(dataset_name):
+            headers = {
+                'Content-Type': 'application/json',
+                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "X-Project-Id": str(self.project_id),
+            }
+            json_data = {
+                "size": 12, 
+                "page": "0", 
+                "projectId": str(self.project_id), 
+                "search": ""
+            }
+            try:
+                # First get dataset details
+                response = requests.post(
+                    f"{Dataset.BASE_URL}/v2/llm/dataset",
+                    headers=headers,
+                    json=json_data,
+                    timeout=Dataset.TIMEOUT,
+                )
+                response.raise_for_status()
+                datasets = response.json()["data"]["content"]
+                dataset_id = [dataset["id"] for dataset in datasets if dataset["name"]==dataset_name][0]
+
+                # Get dataset details to extract schema mapping
+                response = requests.get(
+                    f"{Dataset.BASE_URL}/v2/llm/dataset/{dataset_id}?initialCols=0",
+                    headers=headers,
+                    timeout=Dataset.TIMEOUT,
+                )
+                response.raise_for_status()
+                
+                # Extract schema mapping
+                schema_mapping = {}
+                for col in response.json()["data"]["datasetColumnsResponses"]:
+                    schema_mapping[col["displayName"]] = {"columnType": col["columnType"]}
+                
+                return schema_mapping
+            except requests.exceptions.RequestException as e:
+                logger.error(f"Failed to get schema mapping: {e}")
+                raise
+
+        # Upload CSV to elastic
+        try:
+            schema_mapping = generate_schema_mapping(dataset_name)
+            
+            data = {
+                "projectId": str(self.project_id),
+                "datasetName": dataset_name,
+                "fileName": filename,
+                "schemaMapping": schema_mapping,
+                "opType": "update",  # Use update for adding rows
+                "description": "Adding new rows to dataset"
+            }
+            
+            headers = {
+                'Content-Type': 'application/json',
+                'Authorization': f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "X-Project-Id": str(self.project_id)
+            }
+            
+            response = requests.post(
+                f"{Dataset.BASE_URL}/v2/llm/dataset/csv",
+                headers=headers,
+                json=data,
+                timeout=Dataset.TIMEOUT,
+            )
+            
+            if response.status_code == 400:
+                raise ValueError(response.json().get("message", "Failed to add rows"))
+            
+            response.raise_for_status()
+            
+            # Check response
+            response_data = response.json()
+            if not response_data.get('success', False):
+                raise ValueError(response_data.get('message', 'Unknown error occurred'))
+            
+            print(f"Successfully added rows to dataset {dataset_name}")
+            return response_data
+        
+        except Exception as e:
+            logger.error(f"Error in add_rows_to_dataset: {e}")
+            raise
+
+    def add_columns(self,text_fields,dataset_name, column_name, provider, model,variables={}):
+        """
+        Add a column to a dataset with dynamically fetched model parameters
+        
+        Args:
+            project_id (int): Project ID
+            dataset_id (int): Dataset ID
+            column_name (str): Name of the new column
+            provider (str): Name of the model provider
+            model (str): Name of the model
+        """
+        # First, get model parameters
+
+        # Validate text_fields input
+        if not isinstance(text_fields, list):
+            raise ValueError("text_fields must be a list of dictionaries")
+        
+        for field in text_fields:
+            if not isinstance(field, dict) or 'role' not in field or 'content' not in field:
+                raise ValueError("Each text field must be a dictionary with 'role' and 'content' keys")
+            
+        # First, get the dataset ID
+        headers = {
+            'Content-Type': 'application/json',
+            "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+            "X-Project-Id": str(self.project_id),
+        }
+        json_data = {"size": 12, "page": "0", "projectId": str(self.project_id), "search": ""}
+        
+        try:
+            # Get dataset list
+            response = requests.post(
+                f"{Dataset.BASE_URL}/v2/llm/dataset",
+                headers=headers,
+                json=json_data,
+                timeout=Dataset.TIMEOUT,
+            )
+            response.raise_for_status()
+            datasets = response.json()["data"]["content"]
+            
+            # Find dataset ID
+            dataset_id = next((dataset["id"] for dataset in datasets if dataset["name"] == dataset_name), None)
+            
+            if dataset_id is None:
+                raise ValueError(f"Dataset {dataset_name} not found")
+
+
+
+            parameters_url= f"{Dataset.BASE_URL}/playground/providers/models/parameters/list"
+            
+            headers = {
+                'Content-Type': 'application/json',
+                "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                "X-Project-Id": str(self.project_id),
+            }
+            
+            # Fetch model parameters
+            parameters_payload = {
+                "providerName": provider,
+                "modelName": model
+            }
+        
+            # Get model parameters
+            params_response = requests.post(
+                parameters_url, 
+                headers=headers, 
+                json=parameters_payload, 
+                timeout=30
+            )
+            params_response.raise_for_status()
+            
+            # Extract parameters
+            all_parameters = params_response.json().get('data', [])
+            
+            # Filter and transform parameters for add-column API
+            formatted_parameters = []
+            for param in all_parameters:
+                value = param.get('value')
+                param_type = param.get('type')
+
+                if value is None:
+                    formatted_param = {
+                        "name": param.get('name'),
+                        "value": None,  # Pass None if the value is null
+                        "type": param.get('type')
+                    }
+                else:
+                    # Improved type handling
+                    if param_type == "float":
+                        value = float(value)  # Ensure value is converted to float
+                    elif param_type == "int":
+                        value = int(value)  # Ensure value is converted to int
+                    elif param_type == "bool":
+                        value = bool(value)  # Ensure value is converted to bool
+                    elif param_type == "string":
+                        value = str(value)  # Ensure value is converted to string
+                    else:
+                        raise ValueError(f"Unsupported parameter type: {param_type}")  # Handle unsupported types
+
+                    formatted_param = {
+                        "name": param.get('name'),
+                        "value": value,
+                        "type": param.get('type')
+                    }
+                formatted_parameters.append(formatted_param)
+            dataset_id = next((dataset["id"] for dataset in datasets if dataset["name"] == dataset_name), None)
+
+            # Prepare payload for add column API
+            add_column_payload = {
+                "rowFilterList": [],
+                "columnName": column_name,
+                "datasetId": dataset_id,
+                "variables": variables,
+                "promptTemplate": {
+                    "textFields": text_fields,
+                    "modelSpecs": {
+                        "model": f"{provider}/{model}",
+                        "parameters": formatted_parameters
+                    }
+                }
+            }
+            if variables:
+                variable_specs = []
+                for key, values in variables.items():
+                    variable_specs.append({
+                        "name": key,
+                        "type": "string",
+                        "schema": values
+                    })
+                add_column_payload["promptTemplate"]["variableSpecs"] = variable_specs
+            
+            # Make API call to add column
+            add_column_url = f"{Dataset.BASE_URL}/v2/llm/dataset/add-column"
+            
+            response = requests.post(
+                add_column_url, 
+                headers={
+                    'Content-Type': 'application/json',
+                    'Authorization': f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+                    "X-Project-Id": str(self.project_id)
+                }, 
+                json=add_column_payload,
+                timeout=30
+            )
+            
+            # Check response
+            response.raise_for_status()
+            response_data = response.json()
+            
+            print("Column added successfully:")
+            print(json.dumps(response_data, indent=2))
+            return response_data
+        
+        except requests.exceptions.RequestException as e:
+            print(f"Error adding column: {e}")
+            raise
+

ragaai_catalyst/tracers/__init__.py

@@ -1,3 +1,19 @@
 from .tracer import Tracer
+from .distributed import (
+    init_tracing,
+    trace_agent,
+    trace_llm,
+    trace_tool,
+    current_span,
+    trace_custom,
+)
 
-__all__ = ["Tracer"] 
+__all__ = [
+    "Tracer",
+    "init_tracing",
+    "trace_agent", 
+    "trace_llm",
+    "trace_tool",
+    "current_span",
+    "trace_custom"
+]

ragaai_catalyst/tracers/agentic_tracing/tracers/agent_tracer.py

@@ -44,6 +44,7 @@ class AgentTracerMixin:
         # Add auto instrument flags
         self.auto_instrument_agent = False
         self.auto_instrument_user_interaction = False
+        self.auto_instrument_file_io = False
         self.auto_instrument_network = False
 
     def trace_agent(
@@ -512,10 +513,22 @@ class AgentTracerMixin:
             network_calls = self.component_network_calls.get(kwargs["component_id"], [])
         interactions = []
         if self.auto_instrument_user_interaction:
-            interactions = self.component_user_interaction.get(
-                kwargs["component_id"], []
-            )
-        start_time = kwargs["start_time"]
+            input_output_interactions = []
+            for interaction in self.component_user_interaction.get(kwargs["component_id"], []):
+                if interaction["interaction_type"] in ["input", "output"]:
+                    input_output_interactions.append(interaction)
+            interactions.extend(input_output_interactions) 
+        if self.auto_instrument_file_io:
+            file_io_interactions = []
+            for interaction in self.component_user_interaction.get(kwargs["component_id"], []):
+                if interaction["interaction_type"] in ["file_read", "file_write"]:
+                    file_io_interactions.append(interaction)
+            interactions.extend(file_io_interactions)
+
+        # Get start time
+        start_time = None
+        if "start_time" in kwargs:
+            start_time = kwargs["start_time"]
 
         # Get tags, metrics
         name = kwargs["name"]
@@ -621,3 +634,6 @@ class AgentTracerMixin:
 
     def instrument_network_calls(self):
         self.auto_instrument_network = True
+        
+    def instrument_file_io_calls(self):
+        self.auto_instrument_file_io = True

ragaai_catalyst/tracers/agentic_tracing/tracers/base.py

@@ -308,15 +308,14 @@ class BaseTracer:
 
             # Add metrics to trace before saving
             trace_data = self.trace.to_dict()
-
             trace_data["metrics"] = self.trace_metrics
-
+            
             # Clean up trace_data before saving
             cleaned_trace_data = self._clean_trace(trace_data)
 
             # Format interactions and add to trace
             interactions = self.format_interactions()
-            self.trace.workflow = interactions["workflow"]
+            trace_data["workflow"] = interactions["workflow"]
 
             with open(filepath, "w") as f:
                 json.dump(cleaned_trace_data, f, cls=TracerJSONEncoder, indent=2)
@@ -434,38 +433,44 @@ class BaseTracer:
     def _extract_cost_tokens(self, trace):
         cost = {}
         tokens = {}
-        for span in trace.data[0]["spans"]:
-            if span.type == "llm":
-                info = span.info
-                if isinstance(info, dict):
-                    cost_info = info.get("cost", {})
-                    for key, value in cost_info.items():
-                        if key not in cost:
-                            cost[key] = 0
-                        cost[key] += value
-                    token_info = info.get("tokens", {})
-                    for key, value in token_info.items():
-                        if key not in tokens:
-                            tokens[key] = 0
-                        tokens[key] += value
-            if span.type == "agent":
-                for children in span.data["children"]:
-                    if "type" not in children:
-                        continue
-                    if children["type"] != "llm":
-                        continue
-                    info = children["info"]
-                    if isinstance(info, dict):
-                        cost_info = info.get("cost", {})
-                        for key, value in cost_info.items():
-                            if key not in cost:
-                                cost[key] = 0
-                            cost[key] += value
-                        token_info = info.get("tokens", {})
-                        for key, value in token_info.items():
-                            if key not in tokens:
-                                tokens[key] = 0
-                            tokens[key] += value
+
+        def process_span_info(info):
+            if not isinstance(info, dict):
+                return
+            cost_info = info.get("cost", {})
+            for key, value in cost_info.items():
+                if key not in cost:
+                    cost[key] = 0
+                cost[key] += value
+            token_info = info.get("tokens", {})
+            for key, value in token_info.items():
+                if key not in tokens:
+                    tokens[key] = 0
+                tokens[key] += value
+
+        def process_spans(spans):
+            for span in spans:
+                # Get span type, handling both span objects and dictionaries
+                span_type = span.type if hasattr(span, 'type') else span.get('type')
+                span_info = span.info if hasattr(span, 'info') else span.get('info', {})
+                span_data = span.data if hasattr(span, 'data') else span.get('data', {})
+
+                # Process direct LLM spans
+                if span_type == "llm":
+                    process_span_info(span_info)
+                # Process agent spans recursively
+                elif span_type == "agent":
+                    # Process LLM children in the current agent span
+                    children = span_data.get("children", [])
+                    for child in children:
+                        child_type = child.get("type")
+                        if child_type == "llm":
+                            process_span_info(child.get("info", {}))
+                        # Recursively process nested agent spans
+                        elif child_type == "agent":
+                            process_spans([child])
+
+        process_spans(trace.data[0]["spans"])
         trace.metadata.cost = cost
         trace.metadata.tokens = tokens
         return trace
@@ -665,7 +670,7 @@ class BaseTracer:
                 {
                     "id": str(interaction_id),
                     "span_id": child.get("id"),
-                    "interaction_type": child_type,
+                    "interaction_type": f"{child_type}_call_start",
                     "name": child.get("name"),
                     "content": child.get("data", {}),
                     "timestamp": child.get("start_time"),
@@ -673,15 +678,29 @@ class BaseTracer:
                 }
             )
             interaction_id += 1
+            
+            interactions.append(
+                {
+                    "id": str(interaction_id),
+                    "span_id": child.get("id"),
+                    "interaction_type": f"{child_type}_call_end",
+                    "name": child.get("name"),
+                    "content": child.get("data", {}),
+                    "timestamp": child.get("end_time"),
+                    "error": child.get("error"),
+                }
+            )
+            interaction_id += 1
 
         # Process additional interactions and network calls
         if "interactions" in child:
             for interaction in child["interactions"]:
-                interaction["id"] = str(interaction_id)
-                interaction["span_id"] = child.get("id")
-                interaction["error"] = None
-                interactions.append(interaction)
-                interaction_id += 1
+                if interaction!=[]:
+                    interaction["id"] = str(interaction_id)
+                    interaction["span_id"] = child.get("id")
+                    interaction["error"] = None
+                    interactions.append(interaction)
+                    interaction_id += 1
 
         if "network_calls" in child:
             for child_network_call in child["network_calls"]:
@@ -833,7 +852,7 @@ class BaseTracer:
                     {
                         "id": str(interaction_id),
                         "span_id": span.id,
-                        "interaction_type": span.type,
+                        "interaction_type": f"{span.type}_call_start",
                         "name": span.name,
                         "content": span.data,
                         "timestamp": span.start_time,
@@ -841,19 +860,33 @@ class BaseTracer:
                     }
                 )
                 interaction_id += 1
+                
+                interactions.append(
+                    {
+                        "id": str(interaction_id),
+                        "span_id": span.id,
+                        "interaction_type": f"{span.type}_call_end",
+                        "name": span.name,
+                        "content": span.data,
+                        "timestamp": span.end_time,
+                        "error": span.error,
+                    }
+                )
+                interaction_id += 1
 
             # Process interactions from span.data if they exist
             if span.interactions:
                 for span_interaction in span.interactions:
-                    interaction = {}
-                    interaction["id"] = str(interaction_id)
-                    interaction["span_id"] = span.id
-                    interaction["interaction_type"] = span_interaction.type
-                    interaction["content"] = span_interaction.content
-                    interaction["timestamp"] = span_interaction.timestamp
-                    interaction["error"] = span.error
-                    interactions.append(interaction)
-                    interaction_id += 1
+                    if span_interaction != []:
+                        interaction = {}
+                        interaction["id"] = str(interaction_id)
+                        interaction["span_id"] = span.id
+                        interaction["interaction_type"] = span_interaction.type
+                        interaction["content"] = span_interaction.content
+                        interaction["timestamp"] = span_interaction.timestamp
+                        interaction["error"] = span.error
+                        interactions.append(interaction)
+                        interaction_id += 1
 
             if span.network_calls:
                 for span_network_call in span.network_calls:
@@ -954,7 +987,7 @@ class BaseTracer:
                 self.visited_metrics.append(metric_name)
 
                 formatted_metric = {
-                    "name": metric_name,  # Use potentially modified name
+                    "name": metric_name,  
                     "score": metric["score"],
                     "reason": metric.get("reasoning", ""),
                     "source": "user",

ragaai_catalyst/tracers/agentic_tracing/tracers/custom_tracer.py

@@ -25,6 +25,7 @@ class CustomTracerMixin:
         self.auto_instrument_custom = False
         self.auto_instrument_user_interaction = False
         self.auto_instrument_network = False
+        self.auto_instrument_file_io = False
 
     def trace_custom(self, name: str = None, custom_type: str = "generic", version: str = "1.0.0", trace_variables: bool = True):
         def decorator(func):
@@ -246,8 +247,18 @@ class CustomTracerMixin:
             
         interactions = []
         if self.auto_instrument_user_interaction:
-            interactions = self.component_user_interaction.get(kwargs["component_id"], [])
-
+            input_output_interactions = []
+            for interaction in self.component_user_interaction.get(kwargs["component_id"], []):
+                if interaction["interaction_type"] in ["input", "output"]:
+                    input_output_interactions.append(interaction)
+            interactions.extend(input_output_interactions) 
+        if self.auto_instrument_file_io:
+            file_io_interactions = []
+            for interaction in self.component_user_interaction.get(kwargs["component_id"], []):
+                if interaction["interaction_type"] in ["file_read", "file_write"]:
+                    file_io_interactions.append(interaction)
+            interactions.extend(file_io_interactions)
+            
         component = {
             "id": kwargs["component_id"],
             "hash_id": kwargs["hash_id"],
@@ -314,3 +325,7 @@ class CustomTracerMixin:
     def instrument_network_calls(self):
         """Enable auto-instrumentation for network calls"""
         self.auto_instrument_network = True
+
+    def instrument_file_io_calls(self):
+        """Enable auto-instrumentation for file IO calls"""
+        self.auto_instrument_file_io = True