Using Local Datasets with Maxim SDK for Test Runs

pip install maxim-py

from typing import Dict, Optional, List, Any
from maxim import Maxim
import csv
import json

from maxim.evaluators import BaseEvaluator
from maxim.models import (
    LocalEvaluatorResultParameter,
    LocalEvaluatorReturn,
    ManualData,
    PassFailCriteria,
    QueryBuilder
)

from maxim.models.evaluator import (
    PassFailCriteriaForTestrunOverall,
    PassFailCriteriaOnEachEntry,
)

# For Google Colab users
from google.colab import userdata

API_KEY: str = userdata.get("MAXIM_API_KEY") or ""
WORKSPACE_ID: str = userdata.get("MAXIM_WORKSPACE_ID") or ""
PROMPT_ID: str = userdata.get("PROMPT_ID") or ""

# For VS Code users, use environment variables:
# import os
# API_KEY = os.getenv("MAXIM_API_KEY")
# WORKSPACE_ID = os.getenv("MAXIM_WORKSPACE_ID")
# PROMPT_ID = os.getenv("PROMPT_ID")

maxim = Maxim({
    "api_key": API_KEY, 
    "prompt_management": True
})

dataStructure = {
    "Input": "INPUT",                           # Main input text (required, only one allowed)
    "Expected_Output": "EXPECTED_OUTPUT",       # Expected response (optional, only one allowed)
    # "contextColumn": "CONTEXT_TO_EVALUATE",  # Context for evaluation (optional, only one allowed)
    # "additionalDataColumn": "VARIABLE"       # Additional data columns (multiple allowed)
}

class AIQualityEvaluator(BaseEvaluator):
    """
    Evaluates response quality using AI judgment.
    Scores between 1-5 based on how well the response answers the prompt.
    """

    def evaluate(self, result: LocalEvaluatorResultParameter, data: ManualData) -> Dict[str, LocalEvaluatorReturn]:
        prompt = data["Input"]
        response = result.output

        prompt_quality = self._get_quality_evaluator_prompt()
        
        response = prompt_quality.run(
            f"prompt: {prompt} \n output: {response}"
        )

        content = json.loads(response.choices[0].message.content)

        return {
            "qualityScore": LocalEvaluatorReturn(
                score=content['score'],
                reasoning=content['reasoning']
            )
        }

    def _get_quality_evaluator_prompt(self):
        env = "prod"
        tenantId = 222

        rule = (QueryBuilder()
        .and_()
        .deployment_var("env", env)
        .deployment_var("tenant", tenantId)
        .build()
        )

        return maxim.get_prompt("your_quality_evaluator_prompt_id", rule)

class AISafetyEvaluator(BaseEvaluator):
    """
    Evaluates if the response contains any unsafe content.
    Returns True if safe, False if unsafe.
    """

    def evaluate(self, result: LocalEvaluatorResultParameter, data: ManualData) -> Dict[str, LocalEvaluatorReturn]:
        response = result.output

        prompt_safety = self._get_safety_evaluator_prompt()
        evaluation_response = prompt_safety.run(response)
        
        content = json.loads(evaluation_response.choices[0].message.content)
        safe = content['safe'] == 1

        return {
            "safetyCheck": LocalEvaluatorReturn(
                score=safe,
                reasoning=content['reasoning']
            )
        }

    def _get_safety_evaluator_prompt(self):
        env = "prod-2"
        tenantId = 111

        rule = (QueryBuilder()
        .and_()
        .deployment_var("env", env)
        .deployment_var("tenant", tenantId)
        .build()
        )

        return maxim.get_prompt("your_safety_evaluator_prompt_id", rule)

class KeywordPresenceEvaluator(BaseEvaluator):
    """
    Checks if required keywords are present in the response.
    """

    def __init__(self, required_keywords: list):
        super().__init__()
        self.required_keywords = required_keywords

    def evaluate(self, result: LocalEvaluatorResultParameter, data: ManualData) -> Dict[str, LocalEvaluatorReturn]:
        response = result.outputs.get("response", "").lower()
        missing_keywords = [
            kw for kw in self.required_keywords
            if kw.lower() not in response
        ]

        all_present = len(missing_keywords) == 0

        return {
            "isKeywordPresent": LocalEvaluatorReturn(
                score=all_present,
                reasoning="All keywords present" if all_present
                         else f"Missing keywords: {', '.join(missing_keywords)}"
            )
        }

manual_data = [
    {
        "Input": "Doctor: Hi, what brings you in today?\nPatient: I've had a sore throat and mild fever since yesterday.\nDoctor: Any cough or difficulty swallowing?\nPatient: Some coughing, but no trouble swallowing.",
        "Expected_Output": "Chief complaint: Sore throat and mild fever x1 day.\nHistory: Mild cough, no dysphagia.\nAssessment: Likely viral pharyngitis.\nPlan: Symptomatic treatment, hydration, follow-up if worsens."
    },
    {
        "Input": "Doctor: Good morning! How's the blood pressure?\nPatient: A bit high yesterday—140/95.\nDoctor: Any dizziness, headache?\nPatient: Slight headache in the morning.",
        "Expected_Output": "Chief complaint: Elevated BP noted.\nHistory: Headache AM, BP 140/95.\nAssessment: Mild hypertension.\nPlan: Monitor BP, reinforce lifestyle, follow-up in 1 week."
    }
    # Add more entries as needed
]

def load_csv_data(filepath: str) -> List[Dict[str, Any]]:
    """
    Load data from CSV file and return as list of dictionaries
    """
    with open(filepath, newline='', encoding='utf-8') as f:
        dialect = csv.Sniffer().sniff(f.read(2048))
        f.seek(0)
        return list(csv.DictReader(f, dialect=dialect))

# Load your CSV data
db = load_csv_data("/path/to/your/dataset.csv")

Input,Expected_Output
"Doctor: Hi, what brings you in today?...","Chief complaint: Sore throat..."
"Doctor: Good morning! How's the blood pressure?...","Chief complaint: Elevated BP..."

def load_database_data():
    """
    Example function to load data from a database
    """
    # Your database connection and query logic here
    # Return list of dictionaries matching your data structure
    pass

def load_excel_data(filepath: str):
    """
    Example function to load data from Excel
    """
    import pandas as pd
    df = pd.read_excel(filepath)
    return df.to_dict('records')

quality_criteria = PassFailCriteria(
    on_each_entry_pass_if=PassFailCriteriaOnEachEntry(
        score_should_be=">",
        value=2  # Quality score must be > 2
    ),
    for_testrun_overall_pass_if=PassFailCriteriaForTestrunOverall(
        overall_should_be=">=",
        value=80,  # 80% of entries must pass
        for_result="percentageOfPassedResults"
    )
)

safety_criteria = PassFailCriteria(
    on_each_entry_pass_if=PassFailCriteriaOnEachEntry(
        score_should_be="=",
        value=True  # Must be safe
    ),
    for_testrun_overall_pass_if=PassFailCriteriaForTestrunOverall(
        overall_should_be=">=",
        value=100,  # 100% must be safe
        for_result="percentageOfPassedResults"
    )
)

# Create and trigger test run
test_run = maxim.create_test_run(
    name="Local Dataset Comprehensive Evaluation",
    in_workspace_id=WORKSPACE_ID
).with_data_structure(
    dataStructure
).with_data(
    db  # Use 'manual_data' for manual data option
).with_concurrency(1
).with_evaluators(
    # Built-in evaluator from Maxim store
    "Bias",
    
    # Custom AI evaluators
    AIQualityEvaluator(
        pass_fail_criteria={
            "qualityScore": quality_criteria
        }
    ),
    
    AISafetyEvaluator(
        pass_fail_criteria={
            "safetyCheck": safety_criteria
        }
    ),
    
    # Uncomment to add keyword evaluator
    # KeywordPresenceEvaluator(
    #     required_keywords=["assessment", "plan"]
    # )
).with_prompt_version_id(
    PROMPT_ID
).run()

print("Test run triggered successfully!")
print(f"Status: {test_run.status}")

# ❌ Wrong - column names don't match
dataStructure = {"input": "INPUT"}  # lowercase 'input'
data = [{"Input": "..."}]           # uppercase 'Input'

# ✅ Correct - matching column names
dataStructure = {"Input": "INPUT"}
data = [{"Input": "..."}]

# ❌ Wrong - missing INPUT type
dataStructure = {"Output": "EXPECTED_OUTPUT"}

# ✅ Correct - includes INPUT type
dataStructure = {
    "Input": "INPUT",
    "Output": "EXPECTED_OUTPUT"
}