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

ragaai_catalyst/__init__.py

@@ -1,13 +1,34 @@
 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
 from .evaluation import Evaluation
 from .synthetic_data_generation import SyntheticDataGeneration
+from .redteaming import RedTeaming
 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
+from .redteaming import RedTeaming
 
 
-__all__ = ["Experiment", "RagaAICatalyst", "Tracer", "PromptManager", "Evaluation","SyntheticDataGeneration", "GuardrailsManager"]
+
+
+__all__ = [
+    "Experiment", 
+    "RagaAICatalyst", 
+    "Tracer", 
+    "PromptManager", 
+    "Evaluation",
+    "SyntheticDataGeneration",
+    "RedTeaming",
+    "GuardrailsManager", 
+    "GuardExecutor",
+    "init_tracing", 
+    "trace_agent", 
+    "trace_llm",
+    "trace_tool",
+    "current_span",
+    "trace_custom"
+    "RedTeaming"
+]

ragaai_catalyst/dataset.py

@@ -1,4 +1,7 @@
 import os
+import csv
+import json
+import tempfile
 import requests
 from .utils import response_checker
 from typing import Union
@@ -8,6 +11,10 @@ import pandas as pd
 logger = logging.getLogger(__name__)
 get_token = RagaAICatalyst.get_token
 
+# Job status constants
+JOB_STATUS_FAILED = "failed"
+JOB_STATUS_IN_PROGRESS = "in_progress"
+JOB_STATUS_COMPLETED = "success"
 
 class Dataset:
     BASE_URL = None
@@ -17,6 +24,7 @@ class Dataset:
         self.project_name = project_name
         self.num_projects = 99999
         Dataset.BASE_URL = RagaAICatalyst.BASE_URL
+        self.jobId = None
         headers = {
             "Authorization": f'Bearer {os.getenv("RAGAAI_CATALYST_TOKEN")}',
         }
@@ -218,7 +226,6 @@ class Dataset:
         try:
 
             put_csv_response = put_csv_to_presignedUrl(url)
-            print(put_csv_response)
             if put_csv_response.status_code not in (200, 201):
                 raise ValueError('Unable to put csv to the presignedUrl')
         except Exception as e:
@@ -268,6 +275,460 @@ class Dataset:
                 raise ValueError('Unable to upload csv')
             else:
                 print(upload_csv_response['message'])
+                self.jobId = upload_csv_response['data']['jobId']
         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 response_data.get('success', False):
+                print(f"{response_data['message']}")
+                self.jobId = response_data['data']['jobId']
+            else:
+                raise ValueError(response_data.get('message', 'Failed to add rows'))
+        
+        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": "query"
+                    })
+                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()
+            
+            if response_data.get('success', False):
+                print(f"Column '{column_name}' added successfully to dataset '{dataset_name}'")
+                self.jobId = response_data['data']['jobId']
+            else:
+                raise ValueError(response_data.get('message', 'Failed to add column'))
+        
+        except requests.exceptions.RequestException as e:
+            print(f"Error adding column: {e}")
+            raise
+
+    def get_status(self):
+        headers = {
+            'Content-Type': 'application/json',
+            "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+            'X-Project-Id': str(self.project_id),
+        }
+        try:
+            response = requests.get(
+                f'{Dataset.BASE_URL}/job/status', 
+                headers=headers, 
+                timeout=30)
+            response.raise_for_status()
+            if response.json()["success"]:
+
+                status_json = [item["status"] for item in response.json()["data"]["content"] if item["id"]==self.jobId]
+                status_json = status_json[0]
+                if status_json == "Failed":
+                    print("Job failed. No results to fetch.")
+                    return JOB_STATUS_FAILED
+                elif status_json == "In Progress":
+                    print(f"Job in progress. Please wait while the job completes.\nVisit Job Status: {Dataset.BASE_URL.removesuffix('/api')}/projects/job-status?projectId={self.project_id} to track")
+                    return JOB_STATUS_IN_PROGRESS
+                elif status_json == "Completed":
+                    print(f"Job completed. Fetching results.\nVisit Job Status: {Dataset.BASE_URL.removesuffix('/api')}/projects/job-status?projectId={self.project_id} to check")
+                    return JOB_STATUS_COMPLETED
+                else:
+                    logger.error(f"Unknown status received: {status_json}")
+                    return JOB_STATUS_FAILED
+            else:
+                logger.error("Request was not successful")
+                return JOB_STATUS_FAILED
+        except requests.exceptions.HTTPError as http_err:
+            logger.error(f"HTTP error occurred: {http_err}")
+            return JOB_STATUS_FAILED
+        except requests.exceptions.ConnectionError as conn_err:
+            logger.error(f"Connection error occurred: {conn_err}")
+            return JOB_STATUS_FAILED
+        except requests.exceptions.Timeout as timeout_err:
+            logger.error(f"Timeout error occurred: {timeout_err}")
+            return JOB_STATUS_FAILED
+        except requests.exceptions.RequestException as req_err:
+            logger.error(f"An error occurred: {req_err}")
+            return JOB_STATUS_FAILED
+        except Exception as e:
+            logger.error(f"An unexpected error occurred: {e}")
+            return JOB_STATUS_FAILED
+
+    def _jsonl_to_csv(self, jsonl_file, csv_file):
+        """Convert a JSONL file to a CSV file."""
+        with open(jsonl_file, 'r', encoding='utf-8') as infile:
+            data = [json.loads(line) for line in infile]
+        
+        if not data:
+            print("Empty JSONL file.")
+            return
+        
+        with open(csv_file, 'w', newline='', encoding='utf-8') as outfile:
+            writer = csv.DictWriter(outfile, fieldnames=data[0].keys())
+            writer.writeheader()
+            writer.writerows(data)
+        
+        print(f"Converted {jsonl_file} to {csv_file}")
+
+    def create_from_jsonl(self, jsonl_path, dataset_name, schema_mapping):
+        tmp_csv_path = os.path.join(tempfile.gettempdir(), f"{dataset_name}.csv")
+        try:
+            self._jsonl_to_csv(jsonl_path, tmp_csv_path)
+            self.create_from_csv(tmp_csv_path, dataset_name, schema_mapping)
+        except (IOError, UnicodeError) as e:
+            logger.error(f"Error converting JSONL to CSV: {e}")
+            raise
+        finally:
+            if os.path.exists(tmp_csv_path):
+                try:
+                    os.remove(tmp_csv_path)
+                except Exception as e:
+                    logger.error(f"Error removing temporary CSV file: {e}")
+
+    def add_rows_from_jsonl(self, jsonl_path, dataset_name):
+        tmp_csv_path = os.path.join(tempfile.gettempdir(), f"{dataset_name}.csv")
+        try:
+            self._jsonl_to_csv(jsonl_path, tmp_csv_path)
+            self.add_rows(tmp_csv_path, dataset_name)
+        except (IOError, UnicodeError) as e:
+            logger.error(f"Error converting JSONL to CSV: {e}")
+            raise
+        finally:
+            if os.path.exists(tmp_csv_path):
+                try:
+                    os.remove(tmp_csv_path)
+                except Exception as e:
+                    logger.error(f"Error removing temporary CSV file: {e}")
+
+    def create_from_df(self, df, dataset_name, schema_mapping):
+        tmp_csv_path = os.path.join(tempfile.gettempdir(), f"{dataset_name}.csv")
+        try:
+            df.to_csv(tmp_csv_path, index=False)
+            self.create_from_csv(tmp_csv_path, dataset_name, schema_mapping)
+        except (IOError, UnicodeError) as e:
+            logger.error(f"Error converting DataFrame to CSV: {e}")
+            raise
+        finally:
+            if os.path.exists(tmp_csv_path):
+                try:
+                    os.remove(tmp_csv_path)
+                except Exception as e:
+                    logger.error(f"Error removing temporary CSV file: {e}")
+
+    def add_rows_from_df(self, df, dataset_name):
+        tmp_csv_path = os.path.join(tempfile.gettempdir(), f"{dataset_name}.csv")
+        try:
+            df.to_csv(tmp_csv_path, index=False)
+            self.add_rows(tmp_csv_path, dataset_name)
+        except (IOError, UnicodeError) as e:
+            logger.error(f"Error converting DataFrame to CSV: {e}")
+            raise
+        finally:
+            if os.path.exists(tmp_csv_path):
+                try:
+                    os.remove(tmp_csv_path)
+                except Exception as e:
+                    logger.error(f"Error removing temporary CSV file: {e}")

ragaai_catalyst/evaluation.py

@@ -4,16 +4,22 @@ import pandas as pd
 import io
 from .ragaai_catalyst import RagaAICatalyst
 import logging
+import json
 
 logger = logging.getLogger(__name__)
 
+# Job status constants
+JOB_STATUS_FAILED = "failed"
+JOB_STATUS_IN_PROGRESS = "in_progress"
+JOB_STATUS_COMPLETED = "success"
+
 class Evaluation:
 
     def __init__(self, project_name, dataset_name):
         self.project_name = project_name
         self.dataset_name = dataset_name
         self.base_url = f"{RagaAICatalyst.BASE_URL}"
-        self.timeout = 10
+        self.timeout = 20
         self.jobId = None
         self.num_projects=99999
 
@@ -352,6 +358,52 @@ class Evaluation:
         except Exception as e:
             logger.error(f"An unexpected error occurred: {e}")
 
+    def append_metrics(self, display_name):
+        if not isinstance(display_name, str):
+            raise ValueError("display_name should be a string")
+        
+        headers = {
+            "Authorization": f"Bearer {os.getenv('RAGAAI_CATALYST_TOKEN')}",
+            'X-Project-Id': str(self.project_id),
+            'Content-Type': 'application/json',
+        }
+
+        payload = json.dumps({
+        "datasetId": self.dataset_id,
+        "metricParams": [
+            {
+            "metricSpec": {
+                "displayName": display_name
+            }
+            }
+        ]
+        })
+        
+        try:
+            response = requests.request(
+                "POST", 
+                f'{self.base_url}/v2/llm/metric-evaluation-rerun', 
+                headers=headers, 
+                data=payload, 
+                timeout=self.timeout)
+            if response.status_code == 400:
+                raise ValueError(response.json()["message"])
+            response.raise_for_status()
+            if response.json()["success"]:
+                print(response.json()["message"])
+                self.jobId = response.json()["data"]["jobId"]
+
+        except requests.exceptions.HTTPError as http_err:
+            logger.error(f"HTTP error occurred: {http_err}")
+        except requests.exceptions.ConnectionError as conn_err:
+            logger.error(f"Connection error occurred: {conn_err}")
+        except requests.exceptions.Timeout as timeout_err:
+            logger.error(f"Timeout error occurred: {timeout_err}")
+        except requests.exceptions.RequestException as req_err:
+            logger.error(f"An error occurred: {req_err}")
+        except Exception as e:
+            logger.error(f"An unexpected error occurred: {e}")
+
     def get_status(self):
         headers = {
             'Content-Type': 'application/json',
@@ -366,22 +418,36 @@ class Evaluation:
             response.raise_for_status()
             if response.json()["success"]:
                 status_json = [item["status"] for item in response.json()["data"]["content"] if item["id"]==self.jobId][0]
-            if status_json == "Failed":
-                return print("Job failed. No results to fetch.")
-            elif status_json == "In Progress":
-                return print(f"Job in progress. Please wait while the job completes.\nVisit Job Status: {self.base_url.removesuffix('/api')}/projects/job-status?projectId={self.project_id} to track")
-            elif status_json == "Completed":
-                print(f"Job completed. Fetching results.\nVisit Job Status: {self.base_url.removesuffix('/api')}/projects/job-status?projectId={self.project_id} to check")
+                if status_json == "Failed":
+                    print("Job failed. No results to fetch.")
+                    return JOB_STATUS_FAILED
+                elif status_json == "In Progress":
+                    print(f"Job in progress. Please wait while the job completes.\nVisit Job Status: {self.base_url.removesuffix('/api')}/projects/job-status?projectId={self.project_id} to track")
+                    return JOB_STATUS_IN_PROGRESS
+                elif status_json == "Completed":
+                    print(f"Job completed. Fetching results.\nVisit Job Status: {self.base_url.removesuffix('/api')}/projects/job-status?projectId={self.project_id} to check")
+                    return JOB_STATUS_COMPLETED
+                else:
+                    logger.error(f"Unknown status received: {status_json}")
+                    return JOB_STATUS_FAILED
+            else:
+                logger.error("Request was not successful")
+                return JOB_STATUS_FAILED
         except requests.exceptions.HTTPError as http_err:
             logger.error(f"HTTP error occurred: {http_err}")
+            return JOB_STATUS_FAILED
         except requests.exceptions.ConnectionError as conn_err:
             logger.error(f"Connection error occurred: {conn_err}")
+            return JOB_STATUS_FAILED
         except requests.exceptions.Timeout as timeout_err:
             logger.error(f"Timeout error occurred: {timeout_err}")
+            return JOB_STATUS_FAILED
         except requests.exceptions.RequestException as req_err:
             logger.error(f"An error occurred: {req_err}")
+            return JOB_STATUS_FAILED
         except Exception as e:
             logger.error(f"An unexpected error occurred: {e}")
+            return JOB_STATUS_FAILED
 
     def get_results(self):
 
@@ -444,8 +510,11 @@ class Evaluation:
             df = pd.read_csv(io.StringIO(response_text))
 
             column_list = df.columns.to_list()
+            # Remove unwanted columns
             column_list = [col for col in column_list if not col.startswith('_')]
             column_list = [col for col in column_list if '.' not in col]
+            # Remove _claims_ columns
+            column_list = [col for col in column_list if '_claims_' not in col]
             return df[column_list]
         else:
             return pd.DataFrame()