diff --git a/agent-evaluation/SKILL.md b/agent-evaluation/SKILL.md index e007d5f..f4662a7 100644 --- a/agent-evaluation/SKILL.md +++ b/agent-evaluation/SKILL.md @@ -1,6 +1,6 @@ --- name: agent-evaluation -description: Use this when you need to IMPROVE or OPTIMIZE an existing LLM agent's performance - including improving tool selection accuracy, answer quality, reducing costs, or fixing issues where the agent gives wrong/incomplete responses. Evaluates agents systematically using MLflow evaluation with datasets, scorers, and tracing. Covers end-to-end evaluation workflow or individual components (tracing setup, dataset creation, scorer definition, evaluation execution). +description: Use this when you need to EVALUATE OR IMPROVE or OPTIMIZE an existing LLM agent's output quality - including improving tool selection accuracy, answer quality, reducing costs, or fixing issues where the agent gives wrong/incomplete responses. Evaluates agents systematically using MLflow evaluation with datasets, scorers, and tracing. Covers end-to-end evaluation workflow or individual components (tracing setup, dataset creation, scorer definition, evaluation execution). allowed-tools: Read, Write, Bash, Grep, Glob, WebFetch --- @@ -8,6 +8,19 @@ allowed-tools: Read, Write, Bash, Grep, Glob, WebFetch Comprehensive guide for evaluating GenAI agents with MLflow. Use this skill for the complete evaluation workflow or individual components - tracing setup, environment configuration, dataset creation, scorer definition, or evaluation execution. Each section can be used independently based on your needs. +## ⛔ CRITICAL: Must Use MLflow APIs + +**DO NOT create custom evaluation frameworks.** You MUST use MLflow's native APIs: + +- **Datasets**: Use `mlflow.genai.datasets.create_dataset()` - NOT custom test case files +- **Scorers**: Use `mlflow.genai.scorers` and `mlflow.genai.judges.make_judge()` - NOT custom scorer functions +- **Evaluation**: Use `mlflow.genai.evaluate()` - NOT custom evaluation loops +- **Scripts**: Use the provided `scripts/` directory templates - NOT custom `evaluation/` directories + +**Why?** MLflow tracks everything (datasets, scorers, traces, results) in the experiment. Custom frameworks bypass this and lose all observability. + +If you're tempted to create `evaluation/eval_dataset.py` or similar custom files, STOP. Use `scripts/create_dataset_template.py` instead. + ## Table of Contents 1. [Quick Start](#quick-start) @@ -18,12 +31,14 @@ Comprehensive guide for evaluating GenAI agents with MLflow. Use this skill for ## Quick Start +**⚠️ REMINDER: Use MLflow APIs from this skill. Do not create custom evaluation frameworks.** + **Setup (prerequisite)**: Install MLflow 3.8+, configure environment, integrate tracing -**Evaluation workflow in 4 steps**: +**Evaluation workflow in 4 steps** (each uses MLflow APIs): 1. **Understand**: Run agent, inspect traces, understand purpose -2. **Define**: Select/create scorers for quality criteria +2. **Scorers**: Select and register scorers for quality criteria 3. **Dataset**: ALWAYS discover existing datasets first, only create new if needed 4. **Evaluate**: Run agent on dataset, apply scorers, analyze results @@ -44,26 +59,10 @@ This ensures commands run in the correct environment with proper dependencies. When saving CLI command output to files for parsing (JSON, CSV, etc.), always redirect stderr separately to avoid mixing logs with structured data: ```bash -# WRONG - mixes progress bars and logs with JSON output -uv run mlflow traces evaluate ... --output json > results.json - -# CORRECT - separates stderr from JSON output -uv run mlflow traces evaluate ... --output json 2>/dev/null > results.json - -# ALTERNATIVE - save both separately for debugging +# Save both separately for debugging uv run mlflow traces evaluate ... --output json > results.json 2> evaluation.log ``` -**When to separate streams:** -- Any command with `--output json` flag -- Commands that output structured data (CSV, JSON, XML) -- When piping output to parsing tools (`jq`, `grep`, etc.) - -**When NOT to separate:** -- Interactive commands where you want to see progress -- Debugging scenarios where logs provide context -- Commands that only output unstructured text - ## Documentation Access Protocol **All MLflow documentation must be accessed through llms.txt:** @@ -79,17 +78,6 @@ uv run mlflow traces evaluate ... --output json > results.json 2> evaluation.log - Scorer registration (check MLflow docs for scorer APIs) - Evaluation execution (understand mlflow.genai.evaluate API) -## Pre-Flight Validation - -Validate environment before starting: - -```bash -uv run mlflow --version # Should be >=3.8.0 -uv run python -c "import mlflow; print(f'MLflow {mlflow.__version__} installed')" -``` - -If MLflow is missing or version is <3.8.0, see Setup Overview below. - ## Discovering Agent Structure **Each project has unique structure.** Use dynamic exploration instead of assumptions: @@ -108,18 +96,6 @@ grep -r "@app\.(get|post)" . --include="*.py" # FastAPI/Flask grep -r "def.*route" . --include="*.py" ``` -### Find Tracing Integration -```bash -# Find autolog calls -grep -r "mlflow.*autolog" . --include="*.py" - -# Find trace decorators -grep -r "@mlflow.trace" . --include="*.py" - -# Check imports -grep -r "import mlflow" . --include="*.py" -``` - ### Understand Project Structure ```bash # Check entry points in package config @@ -158,8 +134,6 @@ uv run python scripts/validate_environment.py # Check MLflow install, env vars, uv run python scripts/validate_auth.py # Test authentication before expensive operations ``` -**For complete setup instructions:** See `references/setup-guide.md` - ## Evaluation Workflow ### Step 1: Understand Agent Purpose @@ -171,21 +145,40 @@ uv run python scripts/validate_auth.py # Test authentication before expe ### Step 2: Define Quality Scorers -1. **Discover built-in scorers using documentation protocol:** - - Query `https://mlflow.org/docs/latest/llms.txt` for "What built-in LLM judges or scorers are available?" - - Read scorer documentation to understand their purpose and requirements - - Note: Do NOT use `mlflow scorers list -b` - use documentation instead for accurate information - -2. **Check registered scorers in your experiment:** +1. **Check registered scorers in your experiment:** ```bash uv run mlflow scorers list -x $MLFLOW_EXPERIMENT_ID ``` -3. Identify quality dimensions for your agent and select appropriate scorers -4. Register scorers and test on sample trace before full evaluation +**IMPORTANT: if there are registered scorers in the experiment then they must be used for evaluation.** + +2. **Select additional built-in scorers that apply to the agent** + +See `references/scorers.md` for the built-in scorers. Select any that are useful for assessing the agent's quality and that are not already registered. -**For scorer selection and registration:** See `references/scorers.md` -**For CLI constraints (yes/no format, template variables):** See `references/scorers-constraints.md` +3. **Create additional custom scorers as needed** + +If needed, create additional scorers using the `make_judge()` API. See `references/scorers.md` on how to create custom scorers and `references/scorers-constraints.md` for best practices. + +4. **REQUIRED: Register new scorers before evaluation** using Python API: + + ```python + from mlflow.genai.judges import make_judge + from mlflow.genai.scorers import BuiltinScorerName + import os + + scorer = make_judge(...) # Or, scorer = BuiltinScorerName() + scorer.register() + ``` + +** IMPORTANT: See `references/scorers.md` → "Model Selection for Scorers" to configure the `model` parameter of scorers before registration. + +⚠️ **Scorers MUST be registered before evaluation.** Inline scorers that aren't registered won't appear in `mlflow scorers list` and won't be reusable. + +5. **Verify registration:** + ```bash + uv run mlflow scorers list -x $MLFLOW_EXPERIMENT_ID # Should show your scorers + ``` ### Step 3: Prepare Evaluation Dataset @@ -222,32 +215,50 @@ uv run python scripts/validate_auth.py # Test authentication before expe **For complete dataset guide:** See `references/dataset-preparation.md` +**Checkpoint - verify before proceeding:** + +- [ ] Scorers have been registered +- [ ] Dataset has been created + ### Step 4: Run Evaluation -1. Generate traces: +1. Generate and run evaluation script: ```bash - # Generates evaluation script (auto-detects agent module, entry point, dataset) - uv run python scripts/run_evaluation_template.py - uv run python scripts/run_evaluation_template.py --help # Override auto-detection + # Generate evaluation script (specify module and entry point) + uv run python scripts/run_evaluation_template.py \ + --module mlflow_agent.agent \ + --entry-point run_agent + + # Review the generated script, then execute it + uv run python run_agent_evaluation.py ``` - Generated script uses `mlflow.genai.evaluate()` - review and execute it. + The generated script creates a wrapper function that: + - Accepts keyword arguments matching the dataset's input keys + - Provides any additional arguments the agent needs (like `llm_provider`) + - Runs `mlflow.genai.evaluate(data=df, predict_fn=wrapper, scorers=registered_scorers)` + - Saves results to `evaluation_results.csv` -2. Apply scorers: +⚠️ **CRITICAL: wrapper Signature Must Match Dataset Input Keys** - ```bash - # IMPORTANT: Redirect stderr to avoid mixing logs with JSON output - uv run mlflow traces evaluate \ - --trace-ids \ - --scorers ,,... \ - --output json 2>/dev/null > evaluation_results.json - ``` +MLflow calls `predict_fn(**inputs)` - it unpacks the inputs dict as keyword arguments. + +| Dataset Record | MLflow Calls | predict_fn Must Be | +|----------------|--------------|-------------------| +| `{"inputs": {"query": "..."}}` | `predict_fn(query="...")` | `def wrapper(query):` | +| `{"inputs": {"question": "...", "context": "..."}}` | `predict_fn(question="...", context="...")` | `def wrapper(question, context):` | + +**Common Mistake (WRONG):** +```python +def wrapper(inputs): # ❌ WRONG - inputs is NOT a dict + return agent(inputs["query"]) +``` -3. Analyze results: +2. Analyze results: ```bash # Pattern detection, failure analysis, recommendations - uv run python scripts/analyze_results.py evaluation_results.json + uv run python scripts/analyze_results.py evaluation_results.csv ``` Generates `evaluation_report.md` with pass rates and improvement suggestions. diff --git a/agent-evaluation/references/dataset-preparation.md b/agent-evaluation/references/dataset-preparation.md index 2830c3c..90b8ba2 100644 --- a/agent-evaluation/references/dataset-preparation.md +++ b/agent-evaluation/references/dataset-preparation.md @@ -186,8 +186,8 @@ uv run python scripts/create_dataset_template.py \ 1. Detect your environment 2. Guide you through naming conventions: - **OSS**: Simple names like `mlflow-agent-eval-v1` - - **Databricks**: UC table names like `main.default.mlflow_agent_eval_v1` -3. Help create 10+ diverse sample queries interactively + - **Databricks**: UC table names like `{catalog}.{schema}.mlflow_agent_eval_v1` +3. Help create a few diverse sample queries interactively 4. Generate a complete Python script using correct APIs 5. Optionally execute the script to create the dataset @@ -217,7 +217,6 @@ When using Databricks as your tracking URI, special considerations apply. **1. Fully-Qualified Table Name** - Format: `..` -- Example: `main.default.mlflow_agent_eval_v1` - Cannot use simple names like `my_dataset` **2. Tags Not Supported** @@ -254,17 +253,7 @@ databricks schemas list ``` **Option 3: Use Default** -Suggest the default location: - -``` -main.default.mlflow_agent_eval_v1 -``` - -Where: - -- `main`: Default catalog -- `default`: Default schema -- `mlflow_agent_eval_v1`: Your table name (include version) +Suggest `main.default` for {catalog}.{schema} but *ONLY* if it exists. ### Code Pattern @@ -318,8 +307,7 @@ See generated script output for more examples. ### Sample Size -- **Minimum**: 10 queries (for initial testing) -- **Recommended**: 20-50 queries (for comprehensive evaluation) +- **Minimum**: 5 queries (prefer this for initial setup) - **Balance**: Coverage vs execution time/cost More queries = better coverage but longer evaluation time and higher LLM costs. @@ -340,7 +328,7 @@ More queries = better coverage but longer evaluation time and higher LLM costs. ### Iteration -1. **Start small**: 10-15 queries for initial evaluation +1. **Start small**: 5 queries for initial evaluation 2. **Analyze results**: See what fails, what's missing 3. **Expand**: Add queries to cover gaps 4. **Refine**: Improve existing queries based on agent behavior diff --git a/agent-evaluation/references/scorers.md b/agent-evaluation/references/scorers.md index 6557043..f4dd9b0 100644 --- a/agent-evaluation/references/scorers.md +++ b/agent-evaluation/references/scorers.md @@ -6,9 +6,10 @@ Complete guide for selecting and creating scorers to evaluate agent quality. 1. [Understanding Scorers](#understanding-scorers) 2. [Built-in Scorers](#built-in-scorers) -3. [Custom Scorer Design](#custom-scorer-design) -4. [Scorer Registration](#scorer-registration) -5. [Testing Scorers](#testing-scorers) +3. [Model Selection for Scorers](#model-selection-for-scorers) +4. [Custom Scorer Design](#custom-scorer-design) +5. [Scorer Registration](#scorer-registration) +6. [Testing Scorers](#testing-scorers) ## Understanding Scorers @@ -18,7 +19,7 @@ Scorers (also called "judges" or "LLM-as-a-judge") are evaluation criteria that - Take as input single-turn inputs&outputs, or multi-turn conversations, or traces - Apply quality criteria (relevance, accuracy, completeness, etc.) -- Return a score or pass/fail judgment +- Return a structured assessment (pass/fail or numeric) along with a rationale - Can be built-in (provided by MLflow) or custom (defined by you) ### Types of Scorers @@ -26,86 +27,23 @@ Scorers (also called "judges" or "LLM-as-a-judge") are evaluation criteria that **1. Reference-Free Scorers** - Don't require ground truth or expected outputs -- Judge quality based on their inputs alone -- Examples: Relevance, Completeness, Clarity - **Easiest to use** - work with any dataset **2. Ground-Truth Scorers** - Require expectations in the dataset - Compare agent response to ground truth -- Examples: Factual Accuracy, Answer Correctness - Require datasets with `expectations` field -### LLM-as-a-Judge Pattern - -Modern scorers use an LLM to judge quality: - -1. Scorer receives information on agent's execution (input&output, or trace, or conversation) -2. LLM is given evaluation instructions -3. LLM judges whether criteria is met -4. Returns a structured assessment (pass/fail or numeric) along with a rationale - ## Built-in Scorers MLflow provides several built-in scorers for common evaluation criteria. -### Discovering Built-in Scorers - -**IMPORTANT: Use the documentation protocol to discover built-in scorers.** - -Do NOT use `mlflow scorers list -b` - it may be incomplete or unavailable in some environments. Instead: - -1. Query MLflow documentation via llms.txt: - ``` - WebFetch https://mlflow.org/docs/latest/llms.txt with prompt: - "What built-in LLM judges or scorers are available in MLflow for evaluating GenAI agents?" - ``` - -2. Read scorer documentation pages referenced in llms.txt to understand: - - Scorer names and how to import them - - What each scorer evaluates - - Required inputs (trace structure, expected_response, etc.) - - When to use each scorer - -3. Verify scorer availability by attempting import: - ```python - from mlflow.genai.scorers import Correctness, RelevanceToQuery - ``` - -### Checking Registered Scorers - -List scorers registered in your experiment: - -```bash -uv run mlflow scorers list -x $MLFLOW_EXPERIMENT_ID +```python +uv run mlflow scorers list -b ``` -Output shows: -- Scorer names -- Whether they're built-in or custom -- Registration details - -**IMPORTANT: if there are registered scorers in the experiment then they must be used for evaluation.** - -### Understanding Built-in Scorers - -After querying the documentation, you'll typically find scorers in these categories: - -**Reference-free scorers** (judge without ground truth): - -- Relevance, Completeness, Coherence, Clarity -- Use for: All agents, no expected outputs needed - -**Ground-truth scorers** (require expected outputs): - -- Answer Correctness, Faithfulness, Accuracy -- Use for: When you have known correct answers in dataset - -**Context-based scorers** (require context/documents): - -- Groundedness, Citation Quality -- Use for: RAG systems, knowledge base agents +Detailed documentation for builtin-scorers can be found at `https://mlflow.org/docs/latest/api_reference/python_api/mlflow.genai.html` ### Important: Trace Structure Assumptions @@ -117,39 +55,49 @@ Before using a built-in scorer: 2. **Check trace structure** matches expectations 3. **Verify it works** with a test trace before full evaluation -**Example issue**: +Read `troubleshooting.md` for some common issues with built-in scorers. -- Scorer expects `context` field in trace -- Your agent doesn't provide `context` -- Scorer fails or returns null +## Model Selection for Scorers -**Solution**: +Scorers require an LLM to judge quality. The `model=` parameter accepts MLflow model URIs. -- Read scorer docs carefully -- Test on single trace first -- Create custom scorer if built-in doesn't match your structure +### Default Model Behavior -### Using Built-in Scorers +When `model=None` (default), MLflow selects based on tracking URI: +- **Databricks tracking URI**: Uses `"databricks"` managed judge (no API key needed) +- **Local/other tracking URI**: Uses `"openai:/gpt-4.1-mini"` (requires `OPENAI_API_KEY`) -After discovering scorers via documentation, register them to your experiment: +#### Testing whether the default model works ```python -import os -from mlflow.genai.scorers import Correctness, RelevanceToQuery - -# Note: Import exact class names from documentation -# Common mistake: trying to import "Relevance" when it's actually "RelevanceToQuery" +from mlflow.genai.judges import make_judge -# Register built-in scorer to experiment -scorer = Correctness() -scorer.register(experiment_id=os.getenv("MLFLOW_EXPERIMENT_ID")) +test_judge = make_judge(name="test", instructions="Test whether {{inputs}} is in english") +test_judge(inputs="hello") # Should return a valid Feedback object if default model works ``` -**Benefits of registration**: +### If default model does not work + +#### MLflow Model URI Format + +MLflow uses the format `:/`. Examples: + +| URI | Description | +|-----|-------------| +| `databricks` | Databricks managed judge (good default if there is a databricks connection) | +| `databricks:/` | Databricks serving endpoint | +| `openai:/` | OpenAI (e.g., `openai:/gpt-4`) | +| `anthropic:/` | Anthropic | + +#### Reusing the Agent's LLM -- Shows up in `mlflow scorers list -x ` -- Keeps all evaluation criteria in one place -- Makes it clear what scorers are being used for the experiment +If the agent uses a working LLM, reuse it for scorers to avoid credential issues: + +1. **Find agent's model config** - Look for settings/config module +2. **Check environment** - `LLM_MODEL` or provider-specific env vars +3. **Convert to MLflow URI** - If agent uses LiteLLM format, convert: + - `"gpt-4o"` → `"openai:/gpt-4o"` + - `"databricks/endpoint"` → `"databricks:/endpoint"` ## Custom Scorer Design @@ -160,7 +108,7 @@ Create custom scorers when: - Your agent has unique requirements - Trace structure doesn't match built-in assumptions -## MLflow Judge Constraints +### MLflow Judge Constraints ⚠️ **The MLflow CLI has specific requirements for custom scorers.** @@ -268,7 +216,20 @@ Output: yes/no ## Scorer Registration -### Check CLI Help First +*IMPORTANT: All scorers (built-in and custom) must be registered so that they can be reused in future runs* + +### Registering built-in scorers + +```python +from mlflow.genai.scorers import + +scorer=(...) +scorer.register() +``` + +### Registering custom scorers + +#### Check CLI Help First Run `--help` to verify parameter names: @@ -276,9 +237,9 @@ Run `--help` to verify parameter names: uv run mlflow scorers register-llm-judge --help ``` -### Correct CLI Parameters +#### Correct CLI Parameters -### Registration Example - All Requirements Met +#### Registration Example - All Requirements Met ```bash # ✅ CORRECT - Has variable, uses yes/no, correct parameters @@ -296,9 +257,7 @@ uv run mlflow scorers register-llm-judge \ -i "Examine the trace {{ trace }}. Did the agent use appropriate tools for the query? Return 'yes' if appropriate, 'no' if not." ``` -### Using make_judge() Function - -**Programmatic registration** for advanced use cases: +#### Programmatic registration (for advanced use cases): ```python from mlflow.genai.judges import make_judge @@ -319,16 +278,16 @@ scorer = make_judge( ) # Register the scorer -registered_scorer = scorer.register(experiment_id="your_experiment_id") +registered_scorer = scorer.register() ``` -**When to use make_judge()**: +**When to use**: - `mlflow scorers register-llm-judge` fails with an obscure error - Need programmatic control - Integration with existing code -**Important**: The `make_judge()` API follows the same constraints documented in the CRITICAL CONSTRAINTS section above. Use `Literal["yes", "no"]` for `feedback_value_type` for binary scorers. +**Important**: The `make_judge()` API follows the same constraints documented in the CRITICAL CONSTRAINTS section above. ### Best Practices diff --git a/agent-evaluation/references/setup-guide.md b/agent-evaluation/references/setup-guide.md index 235dff0..3f99fb3 100644 --- a/agent-evaluation/references/setup-guide.md +++ b/agent-evaluation/references/setup-guide.md @@ -32,61 +32,6 @@ uv pip install mlflow>=3.8.0 **Auto-detects Databricks or local MLflow server:** -Run these commands to auto-configure MLflow: - -```bash -# 1. Detect tracking server type -if databricks current-user me &> /dev/null; then - # Databricks detected - export MLFLOW_TRACKING_URI="databricks" - export DB_USER=$(databricks current-user me --output json | grep -o '"value":"[^"]*"' | head -1 | cut -d'"' -f4) - export PROJECT_NAME=$(basename $(pwd)) - export EXP_NAME="/Users/$DB_USER/${PROJECT_NAME}-evaluation" - echo "✓ Detected Databricks" - echo " User: $DB_USER" - echo " Experiment: $EXP_NAME" -else - # Local or other server - export MLFLOW_TRACKING_URI="http://127.0.0.1:5000" - export PROJECT_NAME=$(basename $(pwd)) - export EXP_NAME="${PROJECT_NAME}-evaluation" - echo "✓ Using local MLflow server" - echo " URI: $MLFLOW_TRACKING_URI" - echo " Experiment: $EXP_NAME" - echo "" - echo " Note: If MLflow server isn't running, start it with:" - echo " mlflow server --host 127.0.0.1 --port 5000 &" -fi - -# 2. Find existing or create new experiment -export EXP_ID=$(uv run python -c " -import mlflow -mlflow.set_tracking_uri('$MLFLOW_TRACKING_URI') -experiments = mlflow.search_experiments( - filter_string=\"name = '$EXP_NAME'\", - max_results=1 -) -if experiments: - print(experiments[0].experiment_id) -else: - print(mlflow.create_experiment('$EXP_NAME')) -") - -# 3. Display configuration -echo "" -echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -echo "✓ MLflow Configuration Complete" -echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -echo "Tracking URI: $MLFLOW_TRACKING_URI" -echo "Experiment ID: $EXP_ID" -echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━" -echo "" - -# Export for use in subsequent steps -export MLFLOW_EXPERIMENT_ID="$EXP_ID" -``` - -**Alternative**: Use the setup script with auto-detection: ```bash uv run python scripts/setup_mlflow.py @@ -94,7 +39,9 @@ uv run python scripts/setup_mlflow.py # Outputs: export MLFLOW_TRACKING_URI="..." and export MLFLOW_EXPERIMENT_ID="..." ``` -**After running the above commands**, automatically detect and update the agent's configuration: +⚠️ **Do NOT change the detected MLFLOW_TRACKING_URI** - Trust the detection script unless there is a very serious reason/failure. + +**After running the above command**, automatically detect and update the agent's configuration: 1. **Detect configuration mechanism** by checking for: - `.env` file (most common) diff --git a/agent-evaluation/references/tracing-integration.md b/agent-evaluation/references/tracing-integration.md index 594fd35..605e124 100644 --- a/agent-evaluation/references/tracing-integration.md +++ b/agent-evaluation/references/tracing-integration.md @@ -1,41 +1,24 @@ # MLflow Tracing Integration -Complete guide for integrating MLflow tracing with your agent. This is the authoritative source for all tracing implementation - follow these instructions after completing environment setup in setup-guide.md. +Complete guide for integrating MLflow tracing with your agent. **Prerequisites**: Complete setup-guide.md Steps 1-2 (MLflow install + environment configuration) ## Quick Start -Three steps to integrate tracing: +Three steps to integrate tracing in the agent's code: -1. **Enable autolog** - Add `mlflow..autolog()` BEFORE importing agent code +1. **Enable autolog** - Add `mlflow..autolog()` BEFORE initializing the agent 2. **Decorate entry points** - Add `@mlflow.trace` to agent's main functions 3. **Verify** - Run test query and check trace is captured -**Minimum implementation:** -```python -import mlflow -mlflow.langchain.autolog() # Before imports - -from my_agent import agent - -@mlflow.trace -def run_agent(query: str) -> str: - return agent.run(query) -``` - See sections below for detailed instructions and verification steps. ## Documentation Access Protocol **MANDATORY: Follow the documentation protocol to read MLflow documentation before implementing:** -```bash -# Query llms.txt for tracing documentation -curl https://mlflow.org/docs/latest/llms.txt | grep -A 20 "tracing" -``` - -Or use WebFetch: +WebFetch: - Start: `https://mlflow.org/docs/latest/llms.txt` - Query for: "MLflow tracing documentation", "autolog setup", "trace decorators" - Follow referenced URLs for detailed guides @@ -45,6 +28,7 @@ Or use WebFetch: 1. **Enable Autolog FIRST** - Call `mlflow.{library}.autolog()` before importing agent code - Captures internal library calls automatically - Supported: `langchain`, `langgraph`, `openai`, `anthropic`, etc. + - Prefer autologging for authoring frameworks (e.g., `langchain`, `openai`) over gateways (e.g., `litellm`). 2. **Add @mlflow.trace to Entry Points** - Decorate agent's main functions - Creates top-level span in trace hierarchy @@ -68,23 +52,20 @@ Or use WebFetch: ```python # step 1: Enable autolog BEFORE imports -import mlflow -mlflow.langchain.autolog() # Or langgraph, openai, etc. Use the documentation protocol to find the integration for different libraries. - -# step 2: Import agent code -from my_agent import agent ++ import mlflow ++ mlflow.langchain.autolog() # Or langgraph, openai, etc. Use the documentation protocol to find the integration for different libraries. -# step 3: Add @mlflow.trace decorator -@mlflow.trace +# step 2: Add @mlflow.trace decorator to all existing entry points ++ @mlflow.trace def run_agent(query: str, session_id: str = None) -> str: - """Agent entry point with tracing.""" - result = agent.run(query) - - # step 4 (optional): Track session for multi-turn - if session_id: - trace_id = mlflow.get_last_active_trace_id() - if trace_id: - mlflow.set_trace_tag(trace_id, "session_id", session_id) + """Agent entry point.""" + ... # Agent logic + ++ # step 3 (optional): Track session for multi-turn ++ if session_id: ++ trace_id = mlflow.get_last_active_trace_id() ++ if trace_id: ++ mlflow.set_trace_tag(trace_id, "session_id", session_id) return result ``` diff --git a/agent-evaluation/references/troubleshooting.md b/agent-evaluation/references/troubleshooting.md index e7ef55f..260b935 100644 --- a/agent-evaluation/references/troubleshooting.md +++ b/agent-evaluation/references/troubleshooting.md @@ -25,6 +25,28 @@ Common errors and solutions for agent evaluation with MLflow. 3. Check virtual environment is activated 4. For command line: Add MLflow to PATH +### Wrong API Imports + +**Error/Symptoms**: `ImportError`, `AttributeError`, or methods not found + +**WRONG:** +```python +# These don't exist in MLflow 3 GenAI +from mlflow.evaluate import evaluate +from mlflow.metrics import genai +import mlflow.llm +``` + +**CORRECT:** +```python +import mlflow.genai +from mlflow.genai.scorers import Guidelines, Safety, Correctness, scorer +from mlflow.genai.judges import meets_guidelines, is_correct, make_judge +from mlflow.entities import Feedback, Trace +``` + +**Why**: MLflow 3 GenAI has a completely new API structure. Old MLflow 2 imports don't exist. + ### Databricks Profile Not Authenticated **Error**: `Profile X not authenticated` or `Invalid credentials` @@ -156,6 +178,44 @@ Common errors and solutions for agent evaluation with MLflow. 4. Grep for decorators: `grep -B 2 "def run_agent" src/*/agent/*.py` +### Wrong Trace Search Syntax + +**Error/Symptoms**: Empty results, syntax errors, or unexpected trace filtering behavior + +**WRONG:** +```python +# Missing prefix +mlflow.search_traces("status = 'OK'") + +# Using double quotes +mlflow.search_traces('attributes.status = "OK"') + +# Missing backticks for dotted names +mlflow.search_traces("tags.mlflow.traceName = 'my_app'") + +# Using OR (not supported) +mlflow.search_traces("attributes.status = 'OK' OR attributes.status = 'ERROR'") +``` + +**CORRECT:** +```python +# Use prefix and single quotes +mlflow.search_traces("attributes.status = 'OK'") + +# Backticks for dotted names +mlflow.search_traces("tags.`mlflow.traceName` = 'my_app'") + +# AND is supported +mlflow.search_traces("attributes.status = 'OK' AND tags.env = 'prod'") + +# Time in milliseconds +import time +cutoff = int((time.time() - 3600) * 1000) # 1 hour ago +mlflow.search_traces(f"attributes.timestamp_ms > {cutoff}") +``` + +**Why**: MLflow trace search has specific syntax rules - attributes need prefixes, string values use single quotes, dotted tag names need backticks, and OR is not supported. + ### Session ID Not Captured **Symptoms**: Trace exists but no session_id in tags @@ -205,6 +265,59 @@ Common errors and solutions for agent evaluation with MLflow. ## Dataset Creation Issues +### Wrong Data Format + +**Error/Symptoms**: Schema validation error, missing fields, or evaluation doesn't use your data correctly + +**WRONG:** +```python +# Flat data structure - missing nested structure +eval_data = [ + {"query": "What is X?", "expected": "X is..."} +] +``` + +**CORRECT:** +```python +# Must have 'inputs' key, optionally 'expectations' +eval_data = [ + { + "inputs": {"query": "What is X?"}, + "expectations": {"expected_response": "X is..."} + } +] +``` + +**Why**: MLflow GenAI expects a specific nested structure with `inputs` containing the data passed to predict_fn and `expectations` containing ground truth for scorers. + +### Wrong Dataset Creation - Spark Session + +**Error/Symptoms**: `No Spark session available` or Unity Catalog dataset creation fails + +**WRONG:** +```python +# Need Spark for MLflow-managed datasets +import mlflow.genai.datasets + +dataset = mlflow.genai.datasets.create_dataset( + uc_table_name="catalog.schema.my_dataset" +) +# Error: No Spark session available +``` + +**CORRECT:** +```python +from databricks.connect import DatabricksSession + +spark = DatabricksSession.builder.remote(serverless=True).getOrCreate() + +dataset = mlflow.genai.datasets.create_dataset( + uc_table_name="catalog.schema.my_dataset" +) +``` + +**Why**: Unity Catalog dataset creation requires an active Spark session for table operations. + ### Databricks Dataset APIs Not Supported **Error**: `"Evaluation dataset APIs is not supported in Databricks environments"` @@ -298,6 +411,99 @@ records = [ ## Evaluation Execution Issues +### Wrong Evaluate Function + +**Error/Symptoms**: Wrong API, unexpected results, or `model_type` parameter errors + +**WRONG:** +```python +# This is the old API for classic ML +results = mlflow.evaluate( + model=my_model, + data=eval_data, + model_type="text" +) +``` + +**CORRECT:** +```python +# MLflow 3 GenAI evaluation +results = mlflow.genai.evaluate( + data=eval_dataset, + predict_fn=my_app, + scorers=[Guidelines(name="test", guidelines="...")] +) +``` + +**Why**: `mlflow.evaluate()` is for classic ML models. GenAI agents use `mlflow.genai.evaluate()` with different parameters. + +### Wrong predict_fn Signature + +**Error/Symptoms**: `TypeError` about unexpected arguments, or predict_fn receives wrong data + +**WRONG:** +```python +# predict_fn receives **unpacked inputs, not a dict +def my_app(inputs): # Receives dict + query = inputs["query"] + return {"response": "..."} +``` + +**CORRECT:** +```python +# inputs are unpacked as kwargs +def my_app(query, context=None): # Receives individual keys + return {"response": f"Answer to {query}"} + +# If inputs = {"query": "What is X?", "context": "..."} +# Then my_app is called as: my_app(query="What is X?", context="...") +``` + +**Why**: `mlflow.genai.evaluate()` unpacks the `inputs` dict as keyword arguments to your predict_fn. + +### Using Model Serving Endpoints for Development + +**Error/Symptoms**: Slow iteration, hard to debug, no stack traces + +**WRONG:** +```python +# Don't use model serving endpoints during development +from databricks.sdk import WorkspaceClient + +w = WorkspaceClient() +client = w.serving_endpoints.get_open_ai_client() + +def predict_fn(messages): + response = client.chat.completions.create( + model="my-agent-endpoint", # Deployed endpoint + messages=messages + ) + return {"response": response.choices[0].message.content} +``` + +**CORRECT:** +```python +# Import agent directly for fast iteration +from plan_execute_agent import AGENT # Your local agent module + +def predict_fn(messages): + result = AGENT.predict({"messages": messages}) + # Extract response from ResponsesAgent format + if isinstance(result, dict) and "messages" in result: + for msg in reversed(result["messages"]): + if msg.get("role") == "assistant": + return {"response": msg.get("content", "")} + return {"response": str(result)} +``` + +**Why**: +- Local testing enables faster iteration (no deployment needed) +- Full stack traces for debugging +- No serving endpoint costs +- Direct access to agent internals + +**When to use endpoints**: Only for production monitoring, load testing, or A/B testing deployed versions. + ### Agent Import Errors **Error**: Cannot import agent module or entry point @@ -383,6 +589,300 @@ records = [ ## Scorer Issues +### Wrong Scorer Decorator Usage + +**Error/Symptoms**: Scorer not recognized, or function not treated as a scorer + +**WRONG:** +```python +# Missing @scorer decorator +def my_scorer(outputs): + return True +``` + +**CORRECT:** +```python +from mlflow.genai.scorers import scorer + +@scorer +def my_scorer(outputs): + return True +``` + +**Why**: The `@scorer` decorator is required for MLflow to recognize and use your function as a scorer. + +### Wrong Feedback Return + +**Error/Symptoms**: `TypeError`, serialization errors, or scorer results not recorded + +**WRONG:** +```python +@scorer +def bad_scorer(outputs): + # Can't return dict + return {"score": 0.5, "reason": "..."} + + # Can't return tuple + return (True, "rationale") +``` + +**CORRECT:** +```python +from mlflow.entities import Feedback + +@scorer +def good_scorer(outputs): + # Return primitive + return True + return 0.85 + return "yes" + + # Return Feedback object + return Feedback( + value=True, + rationale="Explanation" + ) + + # Return list of Feedbacks + return [ + Feedback(name="metric_1", value=True), + Feedback(name="metric_2", value=0.9) + ] +``` + +**Why**: Scorers must return primitives (bool, float, str) or Feedback objects. Dicts and tuples are not supported. + +### Wrong Guidelines Scorer Setup + +**Error/Symptoms**: `TypeError` about missing required argument, or scorer doesn't work + +**WRONG:** +```python +# Missing 'name' parameter +scorer = Guidelines(guidelines="Must be professional") +``` + +**CORRECT:** +```python +scorer = Guidelines( + name="professional_tone", # REQUIRED + guidelines="The response must be professional" # REQUIRED +) +``` + +**Why**: `Guidelines` requires both `name` and `guidelines` parameters. + +### Wrong Custom Scorer Imports + +**Error/Symptoms**: Serialization fails when registering scorer for production monitoring + +**WRONG:** +```python +# External import at module level - fails for production monitoring +import my_custom_library + +@scorer +def production_scorer(outputs): + return my_custom_library.process(outputs) +``` + +**CORRECT:** +```python +# Import inside function for serialization +@scorer +def production_scorer(outputs): + import json # Import inside for production monitoring + return len(json.dumps(outputs)) > 100 +``` + +**Why**: Production monitoring scorers are serialized. External imports at module level break serialization. Put imports inside the function. + +### Wrong Multiple Feedback Names + +**Error/Symptoms**: Feedback values overwrite each other, missing metrics in results + +**WRONG:** +```python +@scorer +def bad_multi_scorer(outputs): + # Feedbacks will conflict - no unique names + return [ + Feedback(value=True), + Feedback(value=0.8) + ] +``` + +**CORRECT:** +```python +@scorer +def good_multi_scorer(outputs): + # Each has unique name + return [ + Feedback(name="check_1", value=True), + Feedback(name="check_2", value=0.8) + ] +``` + +**Why**: When returning multiple Feedbacks, each must have a unique `name` to avoid conflicts. + +### Wrong Guidelines Context Reference + +**Error/Symptoms**: Guidelines don't see your data, or judge evaluates wrong content + +**WRONG:** +```python +# Guidelines use 'request' and 'response', not custom keys +Guidelines( + name="check", + guidelines="The output must address the query" # 'output' and 'query' not available +) +``` + +**CORRECT:** +```python +# These are auto-extracted +Guidelines( + name="check", + guidelines="The response must address the request" +) +``` + +**Why**: Guidelines scorers automatically extract `request` (input) and `response` (output) from traces. Use these variable names in your guidelines text. + +### Wrong Custom Judge Model Format + +**Error/Symptoms**: Model not found, or wrong LLM used for judging + +**WRONG:** +```python +# Missing provider prefix +Guidelines(name="test", guidelines="...", model="gpt-4o") + +# Wrong separator +Guidelines(name="test", guidelines="...", model="databricks:gpt-4o") +``` + +**CORRECT:** +```python +# Use :/ separator +Guidelines(name="test", guidelines="...", model="databricks:/my-endpoint") +Guidelines(name="test", guidelines="...", model="openai:/gpt-4o") +``` + +**Why**: Model specification requires the format `provider:/model-name` with `:/` as the separator. + +### Wrong Aggregation Values + +**Error/Symptoms**: `ValueError` about invalid aggregation, or aggregations silently ignored + +**WRONG:** +```python +# p50, p99, sum are not valid +@scorer(aggregations=["mean", "p50", "p99", "sum"]) +def my_scorer(outputs) -> float: + return 0.5 +``` + +**CORRECT:** +```python +# Only these 6 are valid +@scorer(aggregations=["min", "max", "mean", "median", "variance", "p90"]) +def my_scorer(outputs) -> float: + return 0.5 +``` + +**Valid aggregations:** +- `min` - minimum value +- `max` - maximum value +- `mean` - average value +- `median` - 50th percentile (NOT `p50`) +- `variance` - statistical variance +- `p90` - 90th percentile (only p90, NOT p50 or p99) + +### Built-in Scorer Not Working + +**Symptoms**: Built-in scorer errors or returns unexpected or null results + +**Causes**: + +1. Trace structure doesn't match scorer assumptions +2. Required fields missing +3. Scorer expects ground truth but dataset doesn't have it +4. MLflow version incompatibility + +**Solutions**: + +1. Read scorer documentation for requirements: + + - Required trace fields + - Expected structure + - Ground truth needs + +2. Verify trace has expected fields/structure + +3. Check dataset has `expectations` if scorer needs ground truth + +4. Consider custom scorer if built-in doesn't match your structure + +5. Test with verbose output to see what scorer received + +#### Correctness Scorer Requires Expectations + +**WRONG:** +```python +# Correctness requires expected_facts or expected_response +eval_data = [ + {"inputs": {"query": "What is X?"}} +] +results = mlflow.genai.evaluate( + data=eval_data, + predict_fn=my_app, + scorers=[Correctness()] # Will fail - no ground truth! +) +``` + +**CORRECT:** +```python +eval_data = [ + { + "inputs": {"query": "What is X?"}, + "expectations": { + "expected_facts": ["X is a platform", "X is open-source"] + } + } +] +``` + +**Why**: `Correctness` scorer compares agent output to expected facts/response. Without expectations, it has nothing to compare against. + +#### RetrievalGroundedness Requires RETRIEVER Span + +**WRONG:** +```python +# App has no RETRIEVER span type +@mlflow.trace +def my_rag_app(query): + docs = get_documents(query) # Not marked as retriever + return generate_response(docs, query) + +# RetrievalGroundedness will fail - can't find retriever spans +``` + +**CORRECT:** +```python +# Use span_type="RETRIEVER" +@mlflow.trace(span_type="RETRIEVER") +def retrieve_documents(query): + return [doc1, doc2] + +@mlflow.trace +def my_rag_app(query): + docs = retrieve_documents(query) # Now has RETRIEVER span + return generate_response(docs, query) +``` + +**Why**: `RetrievalGroundedness` scorer looks for spans with type `RETRIEVER` to find retrieved documents. Without proper span typing, it can't locate the retrieval step. + ### Scorer Returns Null **Symptoms**: Scorer output is null, empty, or missing @@ -448,33 +948,6 @@ records = [ 5. Add examples to scorer instructions -### Built-in Scorer Not Working - -**Symptoms**: Built-in scorer errors or returns unexpected results - -**Causes**: - -1. Trace structure doesn't match scorer assumptions -2. Required fields missing -3. Scorer expects ground truth but dataset doesn't have it -4. MLflow version incompatibility - -**Solutions**: - -1. Read scorer documentation for requirements: - - - Required trace fields - - Expected structure - - Ground truth needs - -2. Verify trace has expected fields/structure - -3. Check dataset has `expectations` if scorer needs ground truth - -4. Consider custom scorer if built-in doesn't match your structure - -5. Test with verbose output to see what scorer received - ### Scorer Registration Fails **Error**: Error during `mlflow scorers register-llm-judge` diff --git a/agent-evaluation/scripts/create_dataset_template.py b/agent-evaluation/scripts/create_dataset_template.py index 446a135..7df185d 100644 --- a/agent-evaluation/scripts/create_dataset_template.py +++ b/agent-evaluation/scripts/create_dataset_template.py @@ -31,7 +31,7 @@ def parse_arguments(): parser.add_argument( "--test-cases-file", required=True, - help="File with test cases (one per line, minimum 10)", + help="File with test cases (one per line, minimum 5)", ) parser.add_argument( "--dataset-name", @@ -44,19 +44,50 @@ def parse_arguments(): return parser.parse_args() -def load_test_cases_from_file(file_path: str) -> list[str]: - """Load test cases from file (one per line).""" +def load_test_cases_from_file(file_path: str) -> list[dict]: + """Load test cases from file. + + Supports two formats: + 1. JSON file: Array of input dictionaries + [{"query": "question 1"}, {"query": "question 2"}] + + 2. Plain text: One query per line (assumes 'query' key) + question 1 + question 2 + + The input dictionaries should match the entry point function signature. + """ + import json + try: with open(file_path) as f: - test_cases = [line.strip() for line in f if line.strip()] + content = f.read().strip() - if len(test_cases) < 10: - print(f"✗ File has only {len(test_cases)} test cases (minimum: 10)") - print(" Please add more test cases to the file") - sys.exit(1) + # Try to parse as JSON first + try: + test_cases = json.loads(content) + if not isinstance(test_cases, list): + print("✗ JSON file must contain an array of input dictionaries") + sys.exit(1) - print(f"✓ Loaded {len(test_cases)} test cases from {file_path}") - return test_cases + # Validate each item is a dict + for i, item in enumerate(test_cases): + if not isinstance(item, dict): + print(f"✗ Test case {i+1} must be a dictionary, got {type(item).__name__}") + sys.exit(1) + + print(f"✓ Loaded {len(test_cases)} test cases from JSON file") + if test_cases: + print(f" Input keys: {list(test_cases[0].keys())}") + return test_cases + + except json.JSONDecodeError: + # Fall back to plain text (one query per line) + lines = [line.strip() for line in content.split('\n') if line.strip()] + test_cases = [{"query": line} for line in lines] + print(f"✓ Loaded {len(test_cases)} test cases from plain text file") + print(" Note: Using 'query' as the input key") + return test_cases except FileNotFoundError: print(f"✗ File not found: {file_path}") @@ -102,15 +133,16 @@ def generate_dataset_creation_code( tracking_uri: str, experiment_id: str, dataset_name: str, - test_cases: list[str], + test_cases: list[dict], catalog: str = None, schema: str = None, table: str = None, ) -> str: """Generate Python code for creating the dataset.""" + import json - # Escape test cases for Python code - test_cases_repr = repr(test_cases) + # Format test cases as Python code (using json for proper escaping) + test_cases_json = json.dumps(test_cases, indent=4) if catalog and schema and table: # Databricks Unity Catalog version @@ -123,7 +155,7 @@ def generate_dataset_creation_code( """ import mlflow -from mlflow.genai.datasets import create_dataset +from mlflow.genai.datasets import create_dataset, get_dataset # Set tracking URI mlflow.set_tracking_uri("{tracking_uri}") @@ -132,37 +164,56 @@ def generate_dataset_creation_code( DATASET_NAME = "{uc_name}" EXPERIMENT_ID = "{experiment_id}" -# Test cases -TEST_CASES = {test_cases_repr} +# Test cases - each dict should match the entry point function signature +# e.g., if entry point is run_agent(query), test cases should be [{{"query": "..."}}, ...] +TEST_CASES = {test_cases_json} print("=" * 60) print("Creating MLflow Evaluation Dataset") print("=" * 60) print(f"Dataset: {{DATASET_NAME}}") print(f"Test cases: {{len(TEST_CASES)}}") +if TEST_CASES: + print(f"Input keys: {{list(TEST_CASES[0].keys())}}") print() -# Create dataset +# Create or get existing dataset print("Creating dataset...") +dataset = None try: + # Try to create new dataset dataset = create_dataset( name=DATASET_NAME, - source={{ - "inputs": [ - {{"query": query}} - for query in TEST_CASES - ] - }}, - targets=[ - {{"expected_output": "TODO: Add expected output"}} - for _ in TEST_CASES - ], experiment_id=[EXPERIMENT_ID] ) + print(f"✓ Dataset created: {{dataset.dataset_id}}") +except Exception as e: + error_str = str(e) + if "TABLE_ALREADY_EXISTS" in error_str or "already exists" in error_str.lower(): + print(f"Dataset already exists, loading existing dataset...") + try: + dataset = get_dataset(name=DATASET_NAME) + print(f"✓ Loaded existing dataset: {{dataset.dataset_id}}") + except Exception as e2: + print(f"✗ Failed to load existing dataset: {{e2}}") + import traceback + traceback.print_exc() + exit(1) + else: + print(f"✗ Error creating dataset: {{e}}") + import traceback + traceback.print_exc() + exit(1) + +# Add records using merge_records +print("Adding records...") +try: + records = [{{"inputs": tc}} for tc in TEST_CASES] + dataset.merge_records(records) - print(f"✓ Dataset created: {{DATASET_NAME}}") - print(f" Location: Unity Catalog table") - print(f" Queries: {{len(TEST_CASES)}}") + # Verify records were added + df = dataset.to_df() + print(f"✓ Added/merged {{len(df)}} records to dataset") print() print("=" * 60) print("Next Steps") @@ -171,9 +222,8 @@ def generate_dataset_creation_code( print("1. Verify dataset in Databricks Unity Catalog") print(f"2. Use in evaluation: python scripts/run_evaluation_template.py --dataset-name {uc_name}") print() - except Exception as e: - print(f"✗ Error creating dataset: {{e}}") + print(f"✗ Error adding records: {{e}}") import traceback traceback.print_exc() exit(1) @@ -197,36 +247,38 @@ def generate_dataset_creation_code( DATASET_NAME = "{dataset_name}" EXPERIMENT_ID = "{experiment_id}" -# Test cases -TEST_CASES = {test_cases_repr} +# Test cases - each dict should match the entry point function signature +# e.g., if entry point is run_agent(query), test cases should be [{{"query": "..."}}, ...] +TEST_CASES = {test_cases_json} print("=" * 60) print("Creating MLflow Evaluation Dataset") print("=" * 60) print(f"Dataset: {{DATASET_NAME}}") print(f"Test cases: {{len(TEST_CASES)}}") +if TEST_CASES: + print(f"Input keys: {{list(TEST_CASES[0].keys())}}") print() -# Create dataset +# Create dataset (two-step process: create empty, then add records) print("Creating dataset...") try: + # Step 1: Create empty dataset dataset = create_dataset( name=DATASET_NAME, - source={{ - "inputs": [ - {{"query": query}} - for query in TEST_CASES - ] - }}, - targets=[ - {{"expected_output": "TODO: Add expected output"}} - for _ in TEST_CASES - ], experiment_id=[EXPERIMENT_ID] ) + print(f"✓ Dataset created: {{dataset.dataset_id}}") + + # Step 2: Add records using merge_records + # Each test case dict is used directly as the "inputs" field + print("Adding records...") + records = [{{"inputs": tc}} for tc in TEST_CASES] + dataset.merge_records(records) - print(f"✓ Dataset created: {{DATASET_NAME}}") - print(f" Queries: {{len(TEST_CASES)}}") + # Verify records were added + df = dataset.to_df() + print(f"✓ Added {{len(df)}} records to dataset") print() print("=" * 60) print("Next Steps") diff --git a/agent-evaluation/scripts/run_evaluation_template.py b/agent-evaluation/scripts/run_evaluation_template.py index fa24553..5e38bb2 100644 --- a/agent-evaluation/scripts/run_evaluation_template.py +++ b/agent-evaluation/scripts/run_evaluation_template.py @@ -43,145 +43,140 @@ def list_datasets() -> list[str]: def generate_evaluation_code( tracking_uri: str, experiment_id: str, dataset_name: str, agent_module: str, entry_point: str ) -> str: - """Generate Python code for running evaluation.""" + """Generate Python code for running evaluation using mlflow.genai.evaluate().""" return f'''#!/usr/bin/env python3 """ -Run agent on evaluation dataset and collect traces. +Run agent evaluation using mlflow.genai.evaluate(). Generated by run_evaluation_template.py + +IMPORTANT: The entry point function signature must match the dataset's input keys. +For example, if your dataset has {{"inputs": {{"query": "..."}}}}, your entry point +must accept `query` as a keyword argument: def run_agent(query): ... """ import os import sys import mlflow from mlflow.genai.datasets import get_dataset +from mlflow.genai.scorers import list_scorers # Set environment variables os.environ["MLFLOW_TRACKING_URI"] = "{tracking_uri}" os.environ["MLFLOW_EXPERIMENT_ID"] = "{experiment_id}" -# Import agent -from {agent_module} import {entry_point} - -# Configuration -DATASET_NAME = "{dataset_name}" - print("=" * 60) -print("Running Agent on Evaluation Dataset") +print("MLflow Agent Evaluation") print("=" * 60) print() +# Configuration +DATASET_NAME = "{dataset_name}" +EXPERIMENT_ID = "{experiment_id}" + # Load dataset -# IMPORTANT: Do not modify this section. It uses the official MLflow API. -# Spark or databricks-sdk approaches are NOT recommended. print("Loading evaluation dataset...") try: dataset = get_dataset(DATASET_NAME) df = dataset.to_df() print(f" Dataset: {{DATASET_NAME}}") - print(f" Total queries: {{len(df)}}") + print(f" Records: {{len(df)}}") + + # Show sample input structure + if len(df) > 0: + sample_inputs = df.iloc[0]['inputs'] + print(f" Input keys: {{list(sample_inputs.keys())}}") print() except Exception as e: print(f"✗ Failed to load dataset: {{e}}") - print() - print("Common issues:") - print(" 1. Dataset name incorrect - check with: mlflow datasets list") - print(" 2. Not authenticated - run: databricks auth login") - print(" 3. Wrong experiment - verify MLFLOW_EXPERIMENT_ID") sys.exit(1) -# TODO: Configure your agent's LLM provider or other dependencies here -# Example: -# from your_agent.llm import LLMConfig, LLMProvider -# llm_config = LLMConfig(model="gpt-4", temperature=0.0) -# llm_provider = LLMProvider(config=llm_config) - -print("⚠ IMPORTANT: Configure your agent's dependencies above before running!") -print(" Update the TODO section with your agent's setup code") -print() - -# Run agent on each query -trace_ids = [] -successful = 0 -failed = 0 - -print("Running agent on dataset queries...") -print() - -for index, row in df.iterrows(): - inputs = row['inputs'] - - # Extract query from inputs - query = inputs.get('query', inputs.get('question', str(inputs))) - - print(f"[{{index + 1}}/{{len(df)}}] Query: {{query[:80]}}{{'...' if len(query) > 80 else ''}}") - - try: - # TODO: Adjust the function call to match your agent's signature - # Examples: - # response = {entry_point}(query, llm_provider) - # response = {entry_point}(query) - # response = {entry_point}(**inputs) - - response = {entry_point}(query) # <-- UPDATE THIS LINE - - # Capture trace ID - trace_id = mlflow.get_last_active_trace_id() +# Get registered scorers from the experiment +print("Loading registered scorers...") +try: + registered_scorers = list_scorers(experiment_id=EXPERIMENT_ID) - if trace_id: - trace_ids.append(trace_id) - successful += 1 - print(f" ✓ Success (trace: {{trace_id}})") - else: - print(f" ✗ No trace captured") - failed += 1 + if not registered_scorers: + print(" ⚠ No registered scorers found in experiment") + print(" Register scorers first with: scorer.register()") + sys.exit(1) - except Exception as e: - print(f" ✗ Error: {{str(e)[:100]}}") - failed += 1 + print(f" Found {{len(registered_scorers)}} registered scorer(s):") + for scorer in registered_scorers: + print(f" - {{scorer.name}}") + print() +except Exception as e: + print(f"✗ Failed to load scorers: {{e}}") + sys.exit(1) +# Import the agent entry point +# IMPORTANT: The function signature must match the dataset's input keys +print("Importing agent entry point...") +try: + from {agent_module} import {entry_point} + print(f" ✓ Imported {entry_point} from {agent_module}") print() +except ImportError as e: + print(f"✗ Failed to import agent: {{e}}") + sys.exit(1) -# Summary +# Run evaluation print("=" * 60) -print("Execution Summary") +print("Running Evaluation") print("=" * 60) -print(f" Total queries: {{len(df)}}") -print(f" Successful: {{successful}}") -print(f" Failed: {{failed}}") -print(f" Traces collected: {{len(trace_ids)}}") +print() +print(f" Dataset: {{DATASET_NAME}} ({{len(df)}} records)") +print(f" Scorers: {{[s.name for s in registered_scorers]}}") +print(f" Entry point: {agent_module}.{entry_point}") print() -# Save trace IDs -if trace_ids: - traces_file = "evaluation_trace_ids.txt" - with open(traces_file, 'w') as f: - f.write(','.join(trace_ids)) +try: + results = mlflow.genai.evaluate( + data=df, + predict_fn={entry_point}, + scorers=registered_scorers, + ) - print(f"Trace IDs saved to: {{traces_file}}") print() - - # Print evaluation command print("=" * 60) - print("Next Step: Evaluate Traces with Scorers") + print("Evaluation Results") print("=" * 60) print() - print("Run the following command to evaluate all traces:") - print() - print(f" mlflow traces evaluate \\\\") - print(f" --trace-ids {{','.join(trace_ids[:3])}}{{',...' if len(trace_ids) > 3 else ''}} \\\\") - print(f" --scorers ,,... \\\\") - print(f" --output json") - print() - print("Replace ,,... with your registered scorers") - print(" Example: RelevanceToQuery,Completeness,ToolUsageAppropriate") - print() -else: - print("✗ No traces were collected. Please check for errors above.") + + # Display aggregate metrics + if hasattr(results, 'metrics') and results.metrics: + print("Aggregate Metrics:") + for metric_name, value in results.metrics.items(): + if isinstance(value, float): + print(f" {{metric_name}}: {{value:.3f}}") + else: + print(f" {{metric_name}}: {{value}}") + print() + + # Save detailed results + if hasattr(results, 'eval_results_table'): + output_file = "evaluation_results.csv" + results.eval_results_table.to_csv(output_file, index=False) + print(f"✓ Detailed results saved to: {{output_file}}") + print() + print("=" * 60) + print("Evaluation Complete") + print("=" * 60) -print("=" * 60) +except Exception as e: + print(f"✗ Evaluation failed: {{e}}") + print() + print("Common issues:") + print(" 1. Function signature mismatch - entry point params must match dataset input keys") + print(f" Dataset input keys: {{list(sample_inputs.keys()) if 'sample_inputs' in dir() else 'unknown'}}") + print(" Entry point must accept these as keyword arguments") + print(" 2. Missing dependencies - check agent imports") + print(" 3. Authentication issues - verify OPENAI_API_KEY or other credentials") + import traceback + traceback.print_exc() + sys.exit(1) ''' @@ -297,10 +292,13 @@ def main(): print("=" * 60) print() print(f"1. Review the generated script: {output_file}") - print("2. Update the TODO sections with your agent's setup code") - print("3. Update the agent call to match your signature") - print(f"4. Execute it: python {output_file}") - print("5. Use the trace IDs to run evaluation with scorers") + print("2. Ensure your entry point function signature matches dataset input keys") + print(f"3. Execute: uv run python {output_file}") + print() + print("The script will:") + print(" - Load the dataset and registered scorers") + print(" - Run mlflow.genai.evaluate() with your agent") + print(" - Save results to evaluation_results.csv") print() print("=" * 60) diff --git a/agent-evaluation/scripts/setup_mlflow.py b/agent-evaluation/scripts/setup_mlflow.py index 5650254..948134d 100644 --- a/agent-evaluation/scripts/setup_mlflow.py +++ b/agent-evaluation/scripts/setup_mlflow.py @@ -6,19 +6,15 @@ Features: - Auto-detects Databricks profiles or local SQLite -- Search experiments by name (post-processes `mlflow experiments list` output) +- Search experiments by name using MLflow Python API - Single command instead of multiple CLI calls - Creates experiments if they don't exist - -Note: Uses MLflow CLI commands underneath (`mlflow experiments list`, `mlflow experiments create`). -For direct CLI usage, see MLflow documentation. """ import argparse import os import subprocess import sys -from pathlib import Path def parse_arguments(): @@ -41,76 +37,45 @@ def parse_arguments(): def check_mlflow_installed() -> bool: """Check if MLflow >=3.6.0 is installed.""" try: - result = subprocess.run(["mlflow", "--version"], capture_output=True, text=True, check=True) - version = result.stdout.strip().split()[-1] + import mlflow + + version = mlflow.__version__ print(f"✓ MLflow {version} is installed") return True - except (subprocess.CalledProcessError, FileNotFoundError): + except ImportError: print("✗ MLflow is not installed") print(" Install with: uv pip install mlflow") return False def detect_databricks_profiles() -> list[str]: - """Detect available Databricks profiles.""" + """Detect available and valid Databricks profiles. + + Returns: + List of profile names that have Valid=YES + """ try: result = subprocess.run( ["databricks", "auth", "profiles"], capture_output=True, text=True, check=True ) lines = result.stdout.strip().split("\n") - # Skip first line (header: "Name Host Valid") - # and filter empty lines - return [line.strip() for line in lines[1:] if line.strip()] + # Skip header line: "Name Host Valid" + profiles = [] + for line in lines[1:]: + if not line.strip(): + continue + # Parse columns: Name, Host, Valid + parts = line.split() + if len(parts) >= 3: + name = parts[0] + valid = parts[-1] # Last column is Valid (YES/NO) + if valid.upper() == "YES": + profiles.append(name) + return profiles except (subprocess.CalledProcessError, FileNotFoundError): return [] -def check_databricks_auth(profile: str) -> bool: - """Check if a Databricks profile is authenticated.""" - try: - # Try a simple API call to check auth - result = subprocess.run( - ["databricks", "auth", "env", "-p", profile], capture_output=True, text=True, check=True - ) - return "DATABRICKS_TOKEN" in result.stdout or "DATABRICKS_HOST" in result.stdout - except subprocess.CalledProcessError: - return False - - -def start_local_mlflow_server(port: int = 5050) -> bool: - """Start local MLflow server in the background.""" - print(f"\nStarting local MLflow server on port {port}...") - - try: - # Create mlruns directory if it doesn't exist - Path("./mlruns").mkdir(exist_ok=True) - - # Start server in background - cmd = [ - "mlflow", - "server", - "--port", - str(port), - "--backend-store-uri", - "sqlite:///mlflow.db", - "--default-artifact-root", - "./mlruns", - ] - - print(f" Command: {' '.join(cmd)}") - print(" Running in background...") - - # Note: In production, you might want to use nohup or subprocess.Popen with proper detachment - print("\n To start the server manually, run:") - print(f" {' '.join(cmd)} &") - print(f"\n Server will be available at: http://127.0.0.1:{port}") - - return True - except Exception as e: - print(f"✗ Error starting server: {e}") - return False - - def auto_detect_tracking_uri() -> str: """Auto-detect best tracking URI. @@ -171,61 +136,46 @@ def configure_tracking_uri(args_uri: str | None = None) -> str: return auto_detect_tracking_uri() -def list_experiments(tracking_uri: str) -> list[dict]: - """List available experiments.""" - try: - env = os.environ.copy() - env["MLFLOW_TRACKING_URI"] = tracking_uri +def list_experiments(tracking_uri: str, max_results: int = 100) -> list[dict]: + """List available experiments using MLflow Python API. - result = subprocess.run( - ["mlflow", "experiments", "list"], capture_output=True, text=True, check=True, env=env - ) + Args: + tracking_uri: MLflow tracking URI + max_results: Maximum number of experiments to return - # Parse output (simplified) - lines = result.stdout.strip().split("\n") - experiments = [] + Returns: + List of dicts with 'id' and 'name' keys + """ + try: + import mlflow - for line in lines[2:]: # Skip header - if line.strip(): - parts = [p.strip() for p in line.split("|") if p.strip()] - if len(parts) >= 2: - exp_id = parts[0] - name = parts[1] - experiments.append({"id": exp_id, "name": name}) + mlflow.set_tracking_uri(tracking_uri) + experiments = mlflow.search_experiments(max_results=max_results) - return experiments + return [{"id": exp.experiment_id, "name": exp.name} for exp in experiments] except Exception as e: print(f"✗ Error listing experiments: {e}") return [] def create_experiment(tracking_uri: str, name: str) -> str | None: - """Create a new experiment.""" - try: - env = os.environ.copy() - env["MLFLOW_TRACKING_URI"] = tracking_uri + """Create a new experiment using MLflow Python API. - result = subprocess.run( - ["mlflow", "experiments", "create", "-n", name], - capture_output=True, - text=True, - check=True, - env=env, - ) + Args: + tracking_uri: MLflow tracking URI + name: Name for the new experiment - # Extract experiment ID from output - for line in result.stdout.split("\n"): - if "Experiment" in line and "created" in line: - # Try to extract ID - words = line.split() - for i, word in enumerate(words): - if word.lower() == "id" and i + 1 < len(words): - return words[i + 1].strip() + Returns: + Experiment ID if created successfully, None otherwise + """ + try: + import mlflow - # If can't parse, return None (but experiment was created) - return None - except subprocess.CalledProcessError as e: - print(f"✗ Error creating experiment: {e.stderr}") + mlflow.set_tracking_uri(tracking_uri) + exp_id = mlflow.create_experiment(name) + return str(exp_id) + except Exception as e: + print(f"✗ Error creating experiment: {e}") return None @@ -265,7 +215,7 @@ def configure_experiment_id( # Priority 3: Create new experiment if --create and --experiment-name provided if create_new and args_exp_name: - print(f"✓ Creating experiment: {args_exp_name}") + print(f"Creating experiment: {args_exp_name}") exp_id = create_experiment(tracking_uri, args_exp_name) if exp_id: print(f"✓ Experiment created with ID: {exp_id}") @@ -282,7 +232,7 @@ def configure_experiment_id( # Priority 4: Search for experiment by name if provided if args_exp_name: - print(f"✓ Searching for experiment: {args_exp_name}") + print(f"Searching for experiment: {args_exp_name}") experiments = list_experiments(tracking_uri) for exp in experiments: if exp["name"] == args_exp_name: @@ -291,7 +241,7 @@ def configure_experiment_id( # Not found - fail with clear message print(f"✗ Experiment '{args_exp_name}' not found") - print(" Use --create flag to create it: --experiment-name '{args_exp_name}' --create") + print(f" Use --create flag to create it: --experiment-name '{args_exp_name}' --create") sys.exit(1) # Priority 5: Auto-select first available experiment diff --git a/tests/test_agent_evaluation.py b/tests/test_agent_evaluation.py new file mode 100755 index 0000000..276fd94 --- /dev/null +++ b/tests/test_agent_evaluation.py @@ -0,0 +1,803 @@ +#!/usr/bin/env python3 +""" +Test script for the agent-evaluation skill. + +This script: +1. Clones a fresh checkout of mlflow-agent +2. Starts a local MLflow tracking server (SQLite backend) +3. Installs the agent-evaluation skill in .claude/skills/ +4. Tests Claude Code headless mode with a simple query +5. Runs Claude Code evaluation (skill should be auto-discovered) +6. Verifies that evaluation artifacts were created +7. Copies Claude session logs for inspection +8. Cleans up (stops server, removes temp files unless --keep-workdir) + +Usage: + python tests/test_agent_evaluation.py [OPTIONS] [extra_prompt] + +Options: + --skill-dir PATH Path to agent-evaluation skill + --timeout SECONDS Claude Code execution timeout (default: 900) + --mlflow-port PORT Port for local MLflow server (default: 5000) + --keep-workdir Keep the working directory after completion + --tracking-uri URI Use external MLflow server instead of local + --test-runs-dir PATH Directory to store test run data (default: /tmp) + +Environment variables: + OPENAI_API_KEY Passed to Claude Code for MLflow's default judge +""" + +from __future__ import annotations + +import argparse +import atexit +import json +import os +import shutil +import signal +import socket +import subprocess +import sys +import tempfile +import time +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Optional + + +# Exit codes +EXIT_SUCCESS = 0 +EXIT_SETUP_FAILED = 1 +EXIT_EXECUTION_FAILED = 2 +EXIT_VERIFICATION_FAILED = 3 + + +@dataclass +class Config: + """Test configuration.""" + + skill_dir: Path + test_runs_dir: Path = field(default_factory=lambda: Path("/tmp")) + timeout_seconds: int = 900 + mlflow_port: int = 5000 + keep_workdir: bool = True + extra_prompt: str = "" + tracking_uri: Optional[str] = None + openai_api_key: Optional[str] = None + + # Runtime state + work_dir: Optional[Path] = None + experiment_name: Optional[str] = None + experiment_id: Optional[str] = None + log_file: Optional[Path] = None + mlflow_server_pid: Optional[int] = None + use_external_server: bool = False + + # Constants + mlflow_agent_repo: str = "https://github.com/alkispoly-db/mlflow-agent" + + +class Logger: + """Simple logger with colored output.""" + + @staticmethod + def info(msg: str) -> None: + print(f"[INFO] {msg}") + + @staticmethod + def error(msg: str) -> None: + print(f"[ERROR] {msg}", file=sys.stderr) + + @staticmethod + def success(msg: str) -> None: + print(f"[SUCCESS] {msg}") + + @staticmethod + def section(msg: str) -> None: + print() + print("=" * 40) + print(msg) + print("=" * 40) + + +log = Logger() + + +def run_command( + cmd: list[str], + cwd: Optional[Path] = None, + capture_output: bool = True, + check: bool = True, + timeout: Optional[int] = None, + env: Optional[dict] = None, +) -> subprocess.CompletedProcess: + """Run a command and return the result.""" + merged_env = os.environ.copy() + if env: + merged_env.update(env) + + return subprocess.run( + cmd, + cwd=cwd, + capture_output=capture_output, + text=True, + check=check, + timeout=timeout, + env=merged_env, + ) + + +def command_exists(cmd: str) -> bool: + """Check if a command exists in PATH.""" + return shutil.which(cmd) is not None + + +def is_port_available(port: int) -> bool: + """Check if a port is available.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.bind(("127.0.0.1", port)) + return True + except OSError: + return False + + +def cleanup(config: Config) -> None: + """Cleanup function called on exit.""" + log.section("Cleanup") + + # Copy Claude session logs to work directory + if config.work_dir and config.work_dir.exists(): + claude_projects_dir = Path.home() / ".claude" / "projects" + project_path_encoded = str(config.work_dir / "mlflow-agent").replace("/", "-") + session_dir = claude_projects_dir / project_path_encoded + + if session_dir.exists(): + log.info("Copying Claude session logs...") + dest_dir = config.work_dir / "claude-sessions" + dest_dir.mkdir(exist_ok=True) + try: + for item in session_dir.iterdir(): + if item.is_file(): + shutil.copy2(item, dest_dir) + else: + shutil.copytree(item, dest_dir / item.name, dirs_exist_ok=True) + log.info(f"Claude session logs copied to: {dest_dir}") + except Exception as e: + log.error(f"Failed to copy session logs: {e}") + + # Stop MLflow server if we started one + if config.mlflow_server_pid and not config.use_external_server: + log.info(f"Stopping MLflow server (PID: {config.mlflow_server_pid})...") + try: + os.kill(config.mlflow_server_pid, signal.SIGTERM) + # Wait for process to terminate + for _ in range(10): + try: + os.kill(config.mlflow_server_pid, 0) + time.sleep(0.5) + except OSError: + break + log.info("MLflow server stopped") + except OSError: + pass + + # Remove or keep working directory + if config.work_dir and config.work_dir.exists(): + if not config.keep_workdir: + log.info(f"Removing working directory: {config.work_dir}") + shutil.rmtree(config.work_dir) + else: + log.info(f"Keeping working directory: {config.work_dir}") + log.info(f" Claude Code output log: {config.log_file or 'N/A'}") + log.info(f" Claude session logs: {config.work_dir}/claude-sessions/") + if config.use_external_server: + log.info(f" MLflow tracking URI: {config.tracking_uri}") + else: + log.info(f" MLflow data: {config.work_dir}/mlflow-data") + log.info(f" Evaluation experiment ID: {config.experiment_id or 'N/A'}") + + +def check_prerequisites(config: Config) -> bool: + """Check that all prerequisites are met.""" + log.section("Checking Prerequisites") + + missing_deps = [] + + # Check required commands + if not command_exists("claude"): + missing_deps.append("claude (Claude Code CLI)") + if not command_exists("git"): + missing_deps.append("git") + if not command_exists("uv"): + missing_deps.append("uv (Python package manager)") + + if missing_deps: + log.error("Missing required commands:") + for dep in missing_deps: + log.error(f" - {dep}") + return False + + log.info("All required commands available: claude, git, uv") + + # Check skill directory + if not config.skill_dir.exists(): + log.error(f"Skill directory not found: {config.skill_dir}") + return False + + if not (config.skill_dir / "SKILL.md").exists(): + log.error(f"SKILL.md not found in: {config.skill_dir}") + return False + + log.info(f"Skill directory found: {config.skill_dir}") + + # Check external server or port availability + if config.tracking_uri: + config.use_external_server = True + log.info(f"Using external MLflow server: {config.tracking_uri}") + else: + if not is_port_available(config.mlflow_port): + log.error(f"Port {config.mlflow_port} is already in use") + log.error("Set --mlflow-port to use a different port") + return False + log.info(f"Port {config.mlflow_port} is available") + + log.success("All prerequisites satisfied") + return True + + +def start_mlflow_server(config: Config) -> bool: + """Start a local MLflow server.""" + log.section("Starting Local MLflow Server") + + mlflow_data_dir = config.work_dir / "mlflow-data" + mlflow_data_dir.mkdir(parents=True, exist_ok=True) + (mlflow_data_dir / "artifacts").mkdir(exist_ok=True) + + backend_store = f"sqlite:///{mlflow_data_dir}/mlflow.db" + artifact_root = str(mlflow_data_dir / "artifacts") + + log.info(f"Backend store: {backend_store}") + log.info(f"Artifact root: {artifact_root}") + log.info(f"Starting server on port {config.mlflow_port}...") + + # Start MLflow server in background + log_file = config.work_dir / "mlflow-server.log" + with open(log_file, "w") as f: + process = subprocess.Popen( + [ + "uv", + "run", + "python", + "-m", + "mlflow", + "server", + "--host", + "127.0.0.1", + "--port", + str(config.mlflow_port), + "--backend-store-uri", + backend_store, + "--default-artifact-root", + artifact_root, + ], + stdout=f, + stderr=subprocess.STDOUT, + cwd=config.work_dir / "mlflow-agent", + ) + + config.mlflow_server_pid = process.pid + + # Wait for server to be ready + max_attempts = 30 + for attempt in range(max_attempts): + try: + result = subprocess.run( + ["curl", "-s", f"http://127.0.0.1:{config.mlflow_port}/health"], + capture_output=True, + timeout=2, + ) + if result.returncode == 0: + break + except (subprocess.TimeoutExpired, Exception): + pass + + # Check if process died + if process.poll() is not None: + log.error("MLflow server process died") + log.error(f"Check log: {log_file}") + with open(log_file) as f: + print(f.read(), file=sys.stderr) + return False + + time.sleep(1) + else: + log.error("MLflow server failed to start (timeout)") + log.error(f"Check log: {log_file}") + with open(log_file) as f: + print(f.read(), file=sys.stderr) + return False + + log.info(f"MLflow server started (PID: {config.mlflow_server_pid})") + + # Set tracking URI + tracking_uri = f"http://127.0.0.1:{config.mlflow_port}" + os.environ["MLFLOW_TRACKING_URI"] = tracking_uri + log.info(f"MLFLOW_TRACKING_URI set to: {tracking_uri}") + + log.success("MLflow server is ready") + return True + + +def setup_phase(config: Config) -> bool: + """Set up the test environment.""" + log.section("Setup Phase") + + # Create working directory + config.test_runs_dir.mkdir(exist_ok=True) + config.work_dir = Path(tempfile.mkdtemp(prefix="agent-eval-test-", dir=config.test_runs_dir)) + log.info(f"Created working directory: {config.work_dir}") + + # Clone mlflow-agent repository + log.info("Cloning mlflow-agent repository...") + try: + run_command( + ["git", "clone", "--depth", "1", config.mlflow_agent_repo, str(config.work_dir / "mlflow-agent")] + ) + log.info("Repository cloned successfully") + except subprocess.CalledProcessError as e: + log.error(f"Failed to clone repository: {e}") + return False + + # Copy skill to .claude/skills directory + log.info("Setting up agent-evaluation skill in .claude/skills/...") + skills_dir = config.work_dir / "mlflow-agent" / ".claude" / "skills" + skills_dir.mkdir(parents=True, exist_ok=True) + shutil.copytree(config.skill_dir, skills_dir / "agent-evaluation") + log.info(f"Skill installed at: {skills_dir / 'agent-evaluation'}") + + project_dir = config.work_dir / "mlflow-agent" + + # Install dependencies + log.info("Installing Python dependencies with uv sync...") + try: + run_command(["uv", "sync"], cwd=project_dir) + log.info("Dependencies installed successfully") + except subprocess.CalledProcessError as e: + log.error(f"Failed to install dependencies: {e}") + return False + + # Add mlflow package + log.info("Adding mlflow package...") + try: + run_command(["uv", "add", "mlflow"], cwd=project_dir) + log.info("MLflow package added") + except subprocess.CalledProcessError as e: + log.error(f"Failed to add mlflow package: {e}") + return False + + # Start local MLflow server or use external one + if config.use_external_server: + log.section("Using External MLflow Server") + os.environ["MLFLOW_TRACKING_URI"] = config.tracking_uri + log.info(f"MLFLOW_TRACKING_URI set to: {config.tracking_uri}") + + # Install Databricks packages if needed + if config.tracking_uri.startswith("databricks://"): + log.info("Installing Databricks packages for Unity Catalog dataset support...") + try: + run_command(["uv", "add", "databricks-agents", "databricks-connect"], cwd=project_dir) + log.info("Databricks packages installed (databricks-agents, databricks-connect)") + except subprocess.CalledProcessError: + log.error("Failed to install Databricks packages") + log.error("Dataset operations may fail without databricks-agents and databricks-connect") + # Don't fail - evaluation can still work with local DataFrames + + log.success("External MLflow server configured") + else: + if not start_mlflow_server(config): + return False + + # Create test experiment + timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") + base_experiment_name = f"agent-eval-test-{timestamp}" + + # For Databricks, experiment names must be absolute workspace paths + if config.use_external_server and config.tracking_uri.startswith("databricks://"): + db_profile = config.tracking_uri.replace("databricks://", "") + log.info(f"Detecting Databricks workspace user for profile: {db_profile}") + + try: + result = run_command(["databricks", "current-user", "me", "-p", db_profile]) + user_data = json.loads(result.stdout) + db_user = user_data.get("userName", "") + except Exception: + db_user = "" + + if not db_user: + log.error(f"Failed to get Databricks user. Make sure 'databricks auth login -p {db_profile}' has been run.") + return False + + config.experiment_name = f"/Users/{db_user}/{base_experiment_name}" + log.info("Using Databricks workspace path for experiment") + else: + config.experiment_name = base_experiment_name + + log.info(f"Creating test experiment: {config.experiment_name}") + + # Create experiment + try: + result = run_command( + [ + "uv", + "run", + "python", + "-c", + f"import mlflow; print(mlflow.create_experiment('{config.experiment_name}'))", + ], + cwd=project_dir, + ) + config.experiment_id = result.stdout.strip() + + if not config.experiment_id or "Error" in config.experiment_id or "Traceback" in config.experiment_id: + log.error(f"Failed to create test experiment: {config.experiment_id}") + return False + + log.info(f"Created evaluation experiment with ID: {config.experiment_id}") + except subprocess.CalledProcessError as e: + log.error(f"Failed to create experiment: {e}") + return False + + # Set experiment ID + os.environ["MLFLOW_EXPERIMENT_ID"] = config.experiment_id + + # Set OpenAI API key if provided + if config.openai_api_key: + os.environ["OPENAI_API_KEY"] = config.openai_api_key + log.info("OPENAI_API_KEY set (for MLflow scorers)") + + log.success("Setup phase completed") + return True + + +def test_claude_headless(config: Config) -> bool: + """Test that Claude Code headless mode works.""" + log.section("Testing Claude Code Headless Mode") + + project_dir = config.work_dir / "mlflow-agent" + + log.info("Running simple Claude Code test query...") + env = os.environ.copy() + try: + result = subprocess.run( + ["claude", "-p", "Say hello world", "--allowedTools", ""], + cwd=project_dir, + capture_output=True, + text=True, + timeout=120, + stdin=subprocess.DEVNULL, + env=env, + ) + output = result.stdout + result.stderr + except subprocess.TimeoutExpired: + log.error("Claude Code headless mode test timed out") + return False + except Exception as e: + log.error(f"Claude Code headless mode test failed: {e}") + return False + + if not output: + log.error("Claude Code headless mode test failed - no output") + return False + + if "error" in output.lower() and "Error" in output: + log.error("Claude Code headless mode test failed") + log.error(f"Output: {output}") + return False + + log.info(f"Claude Code responded: {output[:100]}...") + log.success("Claude Code headless mode is working") + return True + + +def run_claude_code(config: Config) -> bool: + """Run Claude Code with the evaluation prompt.""" + log.section("Running Claude Code Evaluation") + + config.log_file = config.work_dir / "claude_output.log" + project_dir = config.work_dir / "mlflow-agent" + + log.info("Changing to mlflow-agent directory...") + + # Build the prompt + base_prompt = "Evaluate the output quality of my agent. Do not ask for input." + full_prompt = base_prompt + if config.extra_prompt: + full_prompt = f"{base_prompt} {config.extra_prompt}" + log.info(f"Extra prompt: {config.extra_prompt}") + + log.info("Executing Claude Code with agent-evaluation skill...") + log.info(f"Prompt: {full_prompt}") + log.info(f"Timeout: {config.timeout_seconds} seconds") + log.info(f"Log file: {config.log_file}") + log.info(f"MLFLOW_TRACKING_URI: {os.environ.get('MLFLOW_TRACKING_URI', 'not set')}") + log.info(f"MLFLOW_EXPERIMENT_ID: {os.environ.get('MLFLOW_EXPERIMENT_ID', 'not set')}") + + # Run Claude Code with explicit environment + env = os.environ.copy() + try: + with open(config.log_file, "w") as f: + result = subprocess.run( + [ + "claude", + "-p", + full_prompt, + "--dangerously-skip-permissions", + "--allowedTools", + "Bash,Read,Write,Edit,Grep,Glob,WebFetch", + ], + cwd=project_dir, + stdout=f, + stderr=subprocess.STDOUT, + timeout=config.timeout_seconds, + stdin=subprocess.DEVNULL, + env=env, + ) + exit_code = result.returncode + except subprocess.TimeoutExpired: + log.info(f"Claude Code execution timed out after {config.timeout_seconds} seconds") + log.info("Will still verify if artifacts were created before timeout") + # Don't return - continue to verification + return True + except Exception as e: + log.error(f"Claude Code execution failed: {e}") + return False + + if exit_code != 0: + log.error(f"Claude Code exited with code: {exit_code}") + log.error(f"Check log file for details: {config.log_file}") + return False + + log.success("Claude Code execution completed") + return True + + +def verify_results(config: Config) -> bool: + """Verify that evaluation artifacts were created.""" + log.section("Verification Phase") + + project_dir = config.work_dir / "mlflow-agent" + + log.info("Running verification checks...") + + verification_code = f""" +import sys +import json +import subprocess +from mlflow import MlflowClient + +experiment_id = '{config.experiment_id}' +client = MlflowClient() + +results = {{ + 'datasets': {{'found': 0, 'pass': False}}, + 'scorers': {{'found': 0, 'pass': False}}, + 'traces': {{'found': 0, 'with_assessments': 0, 'pass': False}} +}} + +# Check datasets +try: + datasets = client.search_datasets(experiment_ids=[experiment_id]) + results['datasets']['found'] = len(datasets) + results['datasets']['pass'] = len(datasets) >= 1 +except Exception as e: + print(f'Warning: Error checking datasets: {{e}}', file=sys.stderr) + +# Check traces +try: + traces = client.search_traces(experiment_ids=[experiment_id]) + results['traces']['found'] = len(traces) + results['traces']['pass'] = len(traces) >= 1 +except Exception as e: + print(f'Warning: Error checking traces: {{e}}', file=sys.stderr) + +# Check scorers +try: + scorer_result = subprocess.run( + ['uv', 'run', 'python', '-m', 'mlflow', 'scorers', 'list', '-x', experiment_id], + capture_output=True, + text=True + ) + lines = [l for l in scorer_result.stdout.strip().split('\\n') if l and not l.startswith('---') and 'Name' not in l] + results['scorers']['found'] = len(lines) + results['scorers']['pass'] = len(lines) >= 1 +except Exception as e: + print(f'Warning: Error checking scorers: {{e}}', file=sys.stderr) + +print(json.dumps(results)) +""" + + try: + result = run_command( + ["uv", "run", "python", "-c", verification_code], + cwd=project_dir, + check=False, + ) + verification_result = result.stdout.strip() + except Exception as e: + log.error(f"Verification script failed: {e}") + return False + + if not verification_result: + log.error("Verification script failed to produce output") + return False + + try: + results = json.loads(verification_result) + except json.JSONDecodeError as e: + log.error(f"Failed to parse verification results: {e}") + log.error(f"Raw output: {verification_result}") + return False + + log.info("Verification Results:") + print() + + all_pass = True + + # Datasets check + datasets = results.get("datasets", {}) + if datasets.get("pass"): + print(f" [PASS] Datasets: {datasets.get('found', 0)} created") + else: + print(f" [FAIL] Datasets: {datasets.get('found', 0)} created (expected >= 1)") + all_pass = False + + # Scorers check + scorers = results.get("scorers", {}) + if scorers.get("pass"): + print(f" [PASS] Scorers: {scorers.get('found', 0)} registered") + else: + print(f" [FAIL] Scorers: {scorers.get('found', 0)} registered (expected >= 1)") + all_pass = False + + # Traces check + traces = results.get("traces", {}) + if traces.get("pass"): + print(f" [PASS] Evaluation runs: {traces.get('found', 0)} traces created") + else: + print(f" [FAIL] Evaluation runs: {traces.get('found', 0)} traces created (expected >= 1)") + all_pass = False + + print() + + if all_pass: + log.success("All verification checks passed") + return True + else: + log.error("Some verification checks failed") + return False + + +def main() -> int: + """Main entry point.""" + parser = argparse.ArgumentParser( + description="Test script for the agent-evaluation skill", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=__doc__, + ) + + # Determine default skill directory + script_dir = Path(__file__).parent.resolve() + repo_root = script_dir.parent + default_skill_dir = repo_root / "agent-evaluation" + + parser.add_argument( + "extra_prompt", + nargs="?", + default="", + help="Optional text to append to the Claude Code prompt", + ) + parser.add_argument( + "--skill-dir", + type=Path, + default=default_skill_dir, + help=f"Path to agent-evaluation skill (default: {default_skill_dir})", + ) + parser.add_argument( + "--timeout", + type=int, + default=900, + dest="timeout_seconds", + help="Claude Code execution timeout in seconds (default: 900)", + ) + parser.add_argument( + "--mlflow-port", + type=int, + default=5000, + help="Port for local MLflow server (default: 5000)", + ) + parser.add_argument( + "--keep-workdir", + action="store_true", + default=True, + help="Keep the working directory after completion (default: True)", + ) + parser.add_argument( + "--no-keep-workdir", + action="store_false", + dest="keep_workdir", + help="Remove the working directory after completion", + ) + parser.add_argument( + "--tracking-uri", + default=os.environ.get("MLFLOW_TRACKING_URI"), + help="Use external MLflow server instead of local (can also set MLFLOW_TRACKING_URI env var)", + ) + parser.add_argument( + "--test-runs-dir", + type=Path, + default=Path("/tmp"), + help="Directory to store test run data (default: /tmp)", + ) + + args = parser.parse_args() + + config = Config( + skill_dir=args.skill_dir.resolve(), + test_runs_dir=args.test_runs_dir.resolve(), + timeout_seconds=args.timeout_seconds, + mlflow_port=args.mlflow_port, + keep_workdir=args.keep_workdir, + extra_prompt=args.extra_prompt, + tracking_uri=args.tracking_uri, + openai_api_key=os.environ.get("OPENAI_API_KEY"), + ) + + # Register cleanup handler + atexit.register(cleanup, config) + + log.section("Agent Evaluation Skill Test") + log.info(f"Starting test at {datetime.now()}") + if config.extra_prompt: + log.info(f"Extra prompt argument: {config.extra_prompt}") + if config.tracking_uri: + log.info(f"External MLflow URI: {config.tracking_uri}") + if config.openai_api_key: + log.info("OPENAI_API_KEY: provided (for MLflow scorers)") + + # Phase 1: Check prerequisites + if not check_prerequisites(config): + log.error("Prerequisites check failed") + return EXIT_SETUP_FAILED + + # Phase 2: Setup + if not setup_phase(config): + log.error("Setup phase failed") + return EXIT_SETUP_FAILED + + # Phase 3: Test Claude Code headless mode + if not test_claude_headless(config): + log.error("Claude Code headless mode test failed") + return EXIT_SETUP_FAILED + + # Phase 4: Run Claude Code evaluation + if not run_claude_code(config): + log.error("Claude Code execution failed") + log.error(f"Check log file: {config.log_file}") + return EXIT_EXECUTION_FAILED + + # Phase 5: Verify results + if not verify_results(config): + log.error("Verification failed") + return EXIT_VERIFICATION_FAILED + + log.section("Test Completed Successfully") + log.info(f"Experiment: {config.experiment_name} (ID: {config.experiment_id})") + log.info("All evaluation artifacts were created as expected") + + return EXIT_SUCCESS + + +if __name__ == "__main__": + sys.exit(main())