diff --git a/plugins/claude-code/README.md b/plugins/claude-code/README.md index a7c39e4..0423e56 100644 --- a/plugins/claude-code/README.md +++ b/plugins/claude-code/README.md @@ -1,20 +1,22 @@ # Supercompact — Claude Code Plugin -Entity-preservation conversation compaction for Claude Code. Replaces the built-in LLM-based `/compact` with EITF scoring — **~400x faster** and **2x better entity retention**. +Entity-preservation conversation compaction for Claude Code. **~400x faster** and **2x better entity retention** than the built-in LLM-based `/compact`. ## Quick Install ```bash -git clone https://github.com/yourusername/supercompact.git +git clone https://github.com/heiervang-technologies/supercompact.git cd supercompact/plugins/claude-code ./install.sh ``` -**Prerequisites:** Python 3.11+, [uv](https://github.com/astral-sh/uv) +The installer automatically registers the plugin in `~/.claude/settings.json`. Restart Claude Code, then use `/supercompact`. + +**Prerequisites:** Python 3.11+, [uv](https://github.com/astral-sh/uv), jq ## What It Does -When Claude Code compacts your conversation (either automatically or via `/compact`), it normally calls an LLM to summarize — slow (~30s) and lossy. Supercompact replaces this with **EITF** (Entity-frequency Inverse Turn Frequency), a zero-model algorithm that: +When Claude Code compacts your conversation (either automatically or via `/compact`), it normally calls an LLM to summarize — slow (~30s) and lossy. Supercompact uses **EITF** (Entity-frequency Inverse Turn Frequency), a zero-model algorithm that: 1. Extracts structured entities (file paths, errors, functions, URLs, etc.) 2. Scores each conversation turn by entity importance × rarity @@ -25,13 +27,13 @@ Result: compaction in **~0.2 seconds** with **2x better retention** of file path ## How It Works -The installer sets up three integration points: +The plugin provides three integration points: -1. **cli.js patch** — Replaces the LLM API call in Claude Code's main compaction function with a subprocess call to supercompact. Falls back to the original LLM on error. +1. **`/supercompact` command** — On-demand compaction. Replaces the session with a compacted version and restarts. This is the primary interface. -2. **PreCompact hook** — Backs up the full transcript before any compaction runs, and produces a supercompact alternative alongside Claude's built-in result. +2. **PreCompact hook** — When Claude's built-in compaction triggers, the hook backs up the full transcript before it's lost. The backup is saved as `*.pre-compact-full` alongside the session JSONL. -3. **`/supercompact` command** — Manual on-demand compaction with configurable method and budget. +3. **cli.js patch** *(npm installations only)* — Replaces the LLM API call in Claude Code's compaction function with supercompact. Falls back to the original LLM on error. Not available on standalone binary installations. ## Configuration @@ -70,6 +72,16 @@ Manual compaction. Examples: ./install.sh --patch-only # Patch cli.js only (plugin must be installed first) ``` +## Update + +```bash +cd supercompact +git pull +./plugins/claude-code/install.sh +``` + +Re-running the installer is safe — it replaces all files and is fully idempotent. + ## Uninstall ```bash @@ -95,9 +107,10 @@ Manual compaction. Examples: ├── hooks/ │ └── hooks.json # PreCompact hook registration ├── hooks-handlers/ - │ └── supercompact-precompact.sh + │ └── supercompact-precompact.sh # Backup-only hook └── scripts/ - ├── patcher.py # cli.js patching logic + ├── compact-session.sh # Main compaction script + ├── patcher.py # cli.js patching logic └── patch-compaction.sh ``` @@ -105,6 +118,15 @@ Manual compaction. Examples: Hook activity is logged to `~/.cache/supercompact/hook.log`. +## Standalone Binary Installation + +If Claude Code is installed as a standalone binary (not via npm), the cli.js patch cannot be applied. The installer detects this automatically, skips patching, and configures `settings.json` for you. + +In standalone mode: +- **`/supercompact`** — Works fully. This is the primary way to compact. +- **`/compact`** — Still uses Claude's built-in LLM compaction (cannot be replaced without cli.js patch). +- **PreCompact hook** — Backs up the full transcript before Claude's built-in compaction runs. + ## Troubleshooting **Compaction not working after Claude Code update:** @@ -118,7 +140,7 @@ Check `~/.cache/supercompact/hook.log` for errors. Common causes: - Python/uv not in PATH during compaction - Supercompact directory removed or corrupted -**Verify patch status:** +**Verify patch status (npm installations only):** ```bash grep -c "SUPERCOMPACT_EITF" "$(readlink -f "$(which claude)" | sed 's|[^/]*$|cli.js|')" # 1 = patched, 0 = not patched diff --git a/plugins/claude-code/commands/supercompact.md b/plugins/claude-code/commands/supercompact.md index 5785357..01d6154 100644 --- a/plugins/claude-code/commands/supercompact.md +++ b/plugins/claude-code/commands/supercompact.md @@ -1,73 +1,26 @@ --- description: EITF entity-preservation compaction (~400x faster than /compact, 2x better entity retention) argument-hint: "[budget] [--method eitf|setcover|dedup]" -allowed-tools: Bash(cd *), Bash(uv *), Bash(PROJECT_DIR*), Bash(JSONL_FILE*), Bash(ls *), Bash(wc *), Bash(cp *), Bash(mv *), Bash(restart-claude*), Bash(echo *) +allowed-tools: Bash(*/compact-session.sh*), Bash(restart-claude*) --- # Supercompact — Entity-Preservation Compaction -**CRITICAL: Do NOT use the built-in /compact command. You must follow the exact steps below using Bash tool calls.** - -You are running the supercompact algorithm. This is completely separate from Claude Code's built-in /compact. You must execute the bash commands below, not delegate to any built-in compaction. - -## Configuration - -Settings come from environment variables (set via plugin config): -- `PLUGIN_SETTING_METHOD` — scoring method (default: `eitf`). Options: `eitf`, `setcover`, `dedup` -- `PLUGIN_SETTING_BUDGET` — token budget (default: `80000`) - -The user can override these via arguments: `/supercompact 120000 --method setcover` - -## Step 1: Find the conversation JSONL and supercompact directory +Run the compaction script. It will find the session JSONL automatically, compact it, and report results. ```bash -PROJECT_DIR=$(echo "$PWD" | sed 's|/|-|g; s|^|'"$HOME"'/.claude/projects/|') -JSONL_FILE=$(ls -t "$PROJECT_DIR"/*.jsonl 2>/dev/null | head -1) -echo "JSONL: $JSONL_FILE" -wc -l "$JSONL_FILE" +SCRIPT="${CLAUDE_PLUGIN_ROOT:-${HOME}/.local/share/supercompact/claude-code/plugin}/scripts/compact-session.sh" +"$SCRIPT" $ARGUMENTS ``` -## Step 2: Run compaction - -Parse $ARGUMENTS for an optional numeric budget and `--method `. Fall back to env vars, then defaults. +If the script succeeds and reports compaction was performed (not "already within budget"), restart to load the compacted context: ```bash -METHOD="${PLUGIN_SETTING_METHOD:-eitf}" -BUDGET="${PLUGIN_SETTING_BUDGET:-80000}" -# Override from arguments if provided (e.g. "/supercompact 120000 --method setcover") -for arg in $ARGUMENTS; do - if [[ "$arg" =~ ^[0-9]+$ ]]; then BUDGET="$arg"; fi - if [[ "$prev" == "--method" ]]; then METHOD="$arg"; fi - prev="$arg" -done -# Find supercompact installation -SUPERCOMPACT_DIR="$HOME/.local/share/supercompact/claude-code/supercompact" -if [[ ! -f "$SUPERCOMPACT_DIR/compact.py" ]]; then - echo "ERROR: supercompact not found at $SUPERCOMPACT_DIR. Run install.sh first." - exit 1 -fi -echo "Method: $METHOD, Budget: $BUDGET" -cd "$SUPERCOMPACT_DIR" && uv run python compact.py "$JSONL_FILE" --method "$METHOD" --budget "$BUDGET" --output /tmp/supercompact-output.jsonl --verbose +restart-claude "Session compacted with supercompact. Restarting to load compacted context." ``` -## Step 3: Replace the session JSONL - -```bash -cp "$JSONL_FILE" "${JSONL_FILE}.pre-supercompact" -mv /tmp/supercompact-output.jsonl "$JSONL_FILE" -echo "Replaced session JSONL (backup: ${JSONL_FILE}.pre-supercompact)" -``` - -## Step 4: Report results briefly - -Report: method used, turns kept vs dropped, compression ratio, wall clock time. - -## Step 5: Restart to reload compacted context - -The JSONL on disk is now compacted, but the live session still has old context in memory. Restart to load the compacted version: +If `restart-claude` is not available, tell the user: "Run `/quit` then `claude --resume` to load the compacted context." -```bash -restart-claude "Session compacted with supercompact ($METHOD). Restarting to load compacted context." -``` +If the script reports "already within budget", tell the user and do NOT restart. -If `restart-claude` is not available, tell the user: "Run `/quit` then `claude --resume` to load the compacted context." +If the script fails, show the error output to the user and do not restart. diff --git a/plugins/claude-code/hooks-handlers/supercompact-precompact.sh b/plugins/claude-code/hooks-handlers/supercompact-precompact.sh index d0c6dd4..f4ff67b 100755 --- a/plugins/claude-code/hooks-handlers/supercompact-precompact.sh +++ b/plugins/claude-code/hooks-handlers/supercompact-precompact.sh @@ -1,42 +1,18 @@ #!/usr/bin/env bash -# supercompact-precompact.sh - PreCompact hook for entity-preservation compaction +# supercompact-precompact.sh - PreCompact hook (backup-only) # -# Triggered when Claude Code is about to compact the conversation. # The PreCompact hook CANNOT block or replace Claude's built-in compaction — -# it is notification-only. So we use it to: +# it is notification-only. Running supercompact here is wasted work since +# Claude's LLM compaction overwrites the result anyway. # +# Instead, we just: # 1. Back up the full transcript before Claude's summarization loses detail -# 2. Run compaction (configurable method) to produce a superior alternative -# 3. The user can later resume from the supercompact version instead of Claude's -# -# Configuration via environment variables: -# PLUGIN_SETTING_METHOD Scoring method (default: eitf) -# PLUGIN_SETTING_BUDGET Token budget (default: 80000) +# 2. Log the event +# 3. Clean up old backups set -euo pipefail -# Resolve supercompact installation root -# Layout: ~/.local/share/supercompact/claude-code/plugin/hooks-handlers/THIS_SCRIPT -# ~/.local/share/supercompact/claude-code/supercompact/compact.py -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -PLUGIN_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" -INSTALL_ROOT="$(cd "${PLUGIN_ROOT}/.." && pwd)" -SUPERCOMPACT_DIR="${INSTALL_ROOT}/supercompact" - -if [[ ! -f "${SUPERCOMPACT_DIR}/compact.py" ]]; then - # Fallback: check if supercompact is bundled inside the plugin dir (dev mode) - if [[ -f "${PLUGIN_ROOT}/supercompact/compact.py" ]]; then - SUPERCOMPACT_DIR="${PLUGIN_ROOT}/supercompact" - else - echo "$(date -Iseconds) ERROR: supercompact not found at ${SUPERCOMPACT_DIR}" >> "${HOME}/.cache/supercompact/hook.log" 2>/dev/null - exit 0 - fi -fi - -METHOD="${PLUGIN_SETTING_METHOD:-eitf}" -BUDGET="${PLUGIN_SETTING_BUDGET:-80000}" LOG_DIR="${HOME}/.cache/supercompact" - mkdir -p "${LOG_DIR}" # Read hook input from stdin (JSON with transcript_path, session_id, trigger, etc.) @@ -48,42 +24,22 @@ JSONL_FILE=$(echo "${HOOK_INPUT}" | jq -r '.transcript_path // empty') echo "$(date -Iseconds) PreCompact hook triggered (trigger=${TRIGGER})" >> "${LOG_DIR}/hook.log" if [[ -z "${JSONL_FILE}" || ! -f "${JSONL_FILE}" ]]; then - echo "$(date -Iseconds) ERROR: No transcript_path in hook input or file missing" >> "${LOG_DIR}/hook.log" - exit 0 + echo "$(date -Iseconds) ERROR: No transcript_path in hook input or file missing" >> "${LOG_DIR}/hook.log" + exit 0 fi JSONL_SIZE=$(wc -l < "${JSONL_FILE}") echo "$(date -Iseconds) Transcript: ${JSONL_FILE} (${JSONL_SIZE} lines)" >> "${LOG_DIR}/hook.log" -# 1. Back up the full transcript before Claude's compaction destroys detail +# Back up the full transcript before Claude's compaction destroys detail BACKUP_FILE="${JSONL_FILE}.pre-compact-full" cp "${JSONL_FILE}" "${BACKUP_FILE}" echo "$(date -Iseconds) Full backup saved: ${BACKUP_FILE}" >> "${LOG_DIR}/hook.log" -# 2. Run supercompact to produce a superior alternative -SC_OUTPUT="${JSONL_FILE}.supercompact" - -echo "$(date -Iseconds) Running supercompact (method=${METHOD}, budget=${BUDGET})" >> "${LOG_DIR}/hook.log" - -cd "${SUPERCOMPACT_DIR}" -if uv run python compact.py "${JSONL_FILE}" \ - --method "${METHOD}" \ - --budget "${BUDGET}" \ - --output "${SC_OUTPUT}" 2>> "${LOG_DIR}/hook.log"; then +# Clean up old backups (keep last 3) +ls -t "${JSONL_FILE}.pre-compact-full"* 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true +ls -t "${JSONL_FILE}.pre-supercompact"* 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true - SC_SIZE=$(wc -l < "${SC_OUTPUT}") - echo "$(date -Iseconds) Supercompact (${METHOD}): ${JSONL_SIZE} -> ${SC_SIZE} lines (saved as .supercompact)" >> "${LOG_DIR}/hook.log" - - # Clean up old backups (keep last 3 of each type) - ls -t "${JSONL_FILE}.pre-compact-full"* 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true - ls -t "${JSONL_FILE}.supercompact"* 2>/dev/null | tail -n +4 | xargs rm -f 2>/dev/null || true - - echo "$(date -Iseconds) SUCCESS: Supercompact alternative ready at ${SC_OUTPUT}" >> "${LOG_DIR}/hook.log" - echo "$(date -Iseconds) NOTE: Claude's built-in compaction will still run (hook cannot block it)" >> "${LOG_DIR}/hook.log" - echo "$(date -Iseconds) To use supercompact version: cp '${SC_OUTPUT}' '${JSONL_FILE}'" >> "${LOG_DIR}/hook.log" -else - echo "$(date -Iseconds) ERROR: Supercompact (${METHOD}) failed (Claude's compaction will proceed)" >> "${LOG_DIR}/hook.log" - rm -f "${SC_OUTPUT}" 2>/dev/null || true -fi +echo "$(date -Iseconds) Backup-only hook complete (use /supercompact for manual compaction)" >> "${LOG_DIR}/hook.log" exit 0 diff --git a/plugins/claude-code/install.sh b/plugins/claude-code/install.sh index a361c7f..4369801 100755 --- a/plugins/claude-code/install.sh +++ b/plugins/claude-code/install.sh @@ -168,19 +168,39 @@ if [[ "$DO_INSTALL" == true ]]; then ok "Plugin installed to ${INSTALL_DIR}" - # Print plugin-dir usage - echo "" - info "To load the plugin, use one of:" - echo " claude --plugin-dir ${PLUGIN_DEST}" - echo "" - echo " Or add to ~/.claude/settings.json:" - echo " { \"pluginDirs\": [\"${PLUGIN_DEST}\"] }" - echo "" + # Auto-configure settings.json to load the plugin + SETTINGS_FILE="${HOME}/.claude/settings.json" + info "Configuring Claude Code to load plugin..." + mkdir -p "$(dirname "${SETTINGS_FILE}")" + + if [[ ! -f "${SETTINGS_FILE}" ]]; then + # Create settings.json with pluginDirs + echo '{"pluginDirs":["'"${PLUGIN_DEST}"'"]}' | jq . > "${SETTINGS_FILE}" + ok "Created ${SETTINGS_FILE} with pluginDirs" + elif jq -e '.pluginDirs' "${SETTINGS_FILE}" >/dev/null 2>&1; then + # pluginDirs exists — check if our path is already there + if jq -e --arg p "${PLUGIN_DEST}" '.pluginDirs | index($p)' "${SETTINGS_FILE}" >/dev/null 2>&1; then + ok "Plugin already registered in settings.json" + else + # Add our path to existing pluginDirs array + jq --arg p "${PLUGIN_DEST}" '.pluginDirs += [$p]' "${SETTINGS_FILE}" > "${SETTINGS_FILE}.tmp" \ + && mv "${SETTINGS_FILE}.tmp" "${SETTINGS_FILE}" + ok "Added plugin to existing pluginDirs in settings.json" + fi + else + # settings.json exists but no pluginDirs key — add it + jq --arg p "${PLUGIN_DEST}" '. + {pluginDirs: [$p]}' "${SETTINGS_FILE}" > "${SETTINGS_FILE}.tmp" \ + && mv "${SETTINGS_FILE}.tmp" "${SETTINGS_FILE}" + ok "Added pluginDirs to settings.json" + fi fi # ------------------------------------------------------------------ # Patch cli.js # ------------------------------------------------------------------ +PATCH_APPLIED=false +STANDALONE_BINARY=false + if [[ "$DO_PATCH" == true ]]; then SUPERCOMPACT_DEST="${INSTALL_DIR}/supercompact" @@ -188,16 +208,33 @@ if [[ "$DO_PATCH" == true ]]; then fatal "Supercompact not installed at ${SUPERCOMPACT_DEST}. Run install first (without --patch-only)." fi - echo "" - info "Patching Claude Code cli.js..." - bash "${INSTALL_DIR}/plugin/scripts/patch-compaction.sh" "${SUPERCOMPACT_DEST}" - EXIT_CODE=$? + # Detect standalone binary vs npm installation + CLAUDE_BIN="${CLAUDE_BIN:-$(which claude 2>/dev/null || echo "")}" + CLAUDE_REAL="" + if [[ -n "$CLAUDE_BIN" ]]; then + CLAUDE_REAL="$(readlink -f "$CLAUDE_BIN" 2>/dev/null || echo "$CLAUDE_BIN")" + fi - if [[ $EXIT_CODE -eq 0 ]]; then - ok "cli.js patched — compaction now uses supercompact" + if [[ -n "$CLAUDE_REAL" ]] && head -c 4 "$CLAUDE_REAL" 2>/dev/null | grep -q "ELF\|MZ"; then + STANDALONE_BINARY=true + echo "" + warn "Claude Code is installed as a standalone binary (not via npm)" + warn "cli.js patching is not available for standalone installations" + info "The /supercompact slash command and PreCompact hook will still work" + info "Use '/supercompact' for on-demand compaction" else - err "Patching failed (exit code $EXIT_CODE)" - exit $EXIT_CODE + echo "" + info "Patching Claude Code cli.js..." + bash "${INSTALL_DIR}/plugin/scripts/patch-compaction.sh" "${SUPERCOMPACT_DEST}" + EXIT_CODE=$? + + if [[ $EXIT_CODE -eq 0 ]]; then + ok "cli.js patched — compaction now uses supercompact" + PATCH_APPLIED=true + else + warn "cli.js patching failed (exit code $EXIT_CODE)" + warn "The /supercompact slash command and PreCompact hook will still work" + fi fi fi @@ -210,15 +247,19 @@ echo "" echo "What's installed:" echo " • Supercompact library at ${INSTALL_DIR}/supercompact/" echo " • Plugin at ${INSTALL_DIR}/plugin/" -if [[ "$DO_PATCH" == true ]]; then +echo " • Plugin registered in ~/.claude/settings.json" +if [[ "$PATCH_APPLIED" == true ]]; then echo " • cli.js patched for automatic compaction replacement" fi echo "" -echo "Configuration (via environment variables or plugin settings):" -echo " PLUGIN_SETTING_METHOD=eitf # eitf, setcover, dedup" -echo " PLUGIN_SETTING_BUDGET=80000 # token budget" -echo " PLUGIN_SETTING_FALLBACK_TO_BUILTIN=true # fall back to LLM on error" -echo "" -if [[ "$DO_PATCH" == true ]]; then - echo "Restart Claude Code to activate the patch." +echo "Usage:" +if [[ "$PATCH_APPLIED" == true ]]; then + echo " /compact and /supercompact both use supercompact now." + echo " Restart Claude Code to activate." +else + echo " /supercompact # On-demand entity-preservation compaction" + echo " /supercompact 120000 # Custom token budget" fi +echo "" +echo "To update later: git pull && ./install.sh" +echo "To uninstall: ./uninstall.sh" diff --git a/plugins/claude-code/scripts/compact-session.sh b/plugins/claude-code/scripts/compact-session.sh new file mode 100755 index 0000000..8432872 --- /dev/null +++ b/plugins/claude-code/scripts/compact-session.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# compact-session.sh - Self-contained supercompact session compaction +# +# Finds the current Claude Code session JSONL, runs supercompact, +# backs up the original, and replaces it with the compacted version. +# +# Usage: compact-session.sh [budget] [--method name] +# +# Environment: +# CLAUDE_PROJECT_DIR Project dir set by Claude Code (preferred for JSONL lookup) +# PLUGIN_SETTING_METHOD Scoring method (default: eitf) +# PLUGIN_SETTING_BUDGET Token budget (default: 80000) + +set -euo pipefail + +# --- Resolve supercompact installation --- +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PLUGIN_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +INSTALL_ROOT="$(cd "${PLUGIN_ROOT}/.." && pwd)" +SUPERCOMPACT_DIR="${INSTALL_ROOT}/supercompact" + +if [[ ! -f "${SUPERCOMPACT_DIR}/compact.py" ]]; then + # Dev mode: repo layout + if [[ -f "${PLUGIN_ROOT}/../../supercompact/compact.py" ]]; then + SUPERCOMPACT_DIR="$(cd "${PLUGIN_ROOT}/../../supercompact" && pwd)" + else + echo "ERROR: supercompact not found at ${SUPERCOMPACT_DIR}" + exit 1 + fi +fi + +# --- Parse arguments --- +METHOD="${PLUGIN_SETTING_METHOD:-eitf}" +BUDGET="${PLUGIN_SETTING_BUDGET:-80000}" +prev="" +for arg in "$@"; do + if [[ "$arg" =~ ^[0-9]+$ ]]; then BUDGET="$arg"; fi + if [[ "$prev" == "--method" ]]; then METHOD="$arg"; fi + prev="$arg" +done + +# --- Find the session JSONL --- +find_project_dir() { + # Prefer CLAUDE_PROJECT_DIR if set + if [[ -n "${CLAUDE_PROJECT_DIR:-}" && -d "${CLAUDE_PROJECT_DIR}" ]]; then + echo "${CLAUDE_PROJECT_DIR}" + return + fi + # Fallback: derive from PWD (same logic Claude Code uses) + local derived + derived="${HOME}/.claude/projects/$(echo "${PWD}" | sed 's|/|-|g')" + if [[ -d "${derived}" ]]; then + echo "${derived}" + return + fi + echo "" +} + +PROJECT_DIR="$(find_project_dir)" +if [[ -z "${PROJECT_DIR}" ]]; then + echo "ERROR: Could not find Claude project directory" + echo " Tried CLAUDE_PROJECT_DIR=${CLAUDE_PROJECT_DIR:-}" + echo " Tried PWD-derived=${HOME}/.claude/projects/$(echo "${PWD}" | sed 's|/|-|g')" + exit 1 +fi + +JSONL_FILE="$(ls -t "${PROJECT_DIR}"/*.jsonl 2>/dev/null | head -1)" +if [[ -z "${JSONL_FILE}" || ! -f "${JSONL_FILE}" ]]; then + echo "ERROR: No .jsonl files found in ${PROJECT_DIR}" + exit 1 +fi + +LINES_BEFORE=$(wc -l < "${JSONL_FILE}") +echo "Session JSONL: ${JSONL_FILE}" +echo "Lines before: ${LINES_BEFORE}" +echo "Method: ${METHOD}" +echo "Budget: ${BUDGET}" +echo "" + +# --- Run supercompact --- +SC_OUTPUT="/tmp/supercompact-output-$$.jsonl" +trap 'rm -f "${SC_OUTPUT}"' EXIT + +START_TIME=$(date +%s%N) + +cd "${SUPERCOMPACT_DIR}" +SC_STDOUT=$(uv run python compact.py compact "${JSONL_FILE}" \ + --method "${METHOD}" \ + --budget "${BUDGET}" \ + --output "${SC_OUTPUT}" \ + --verbose 2>&1) || { + echo "" + echo "${SC_STDOUT}" + echo "" + echo "ERROR: supercompact failed" + exit 1 +} + +END_TIME=$(date +%s%N) +ELAPSED_MS=$(( (END_TIME - START_TIME) / 1000000 )) + +echo "${SC_STDOUT}" + +if [[ ! -f "${SC_OUTPUT}" ]]; then + # Already within budget — not an error + echo "" + echo "Session is already within budget. No compaction needed." + exit 0 +fi + +# --- Backup and replace --- +BACKUP_FILE="${JSONL_FILE}.pre-supercompact" +cp "${JSONL_FILE}" "${BACKUP_FILE}" +mv "${SC_OUTPUT}" "${JSONL_FILE}" +trap - EXIT # output file moved, no cleanup needed + +LINES_AFTER=$(wc -l < "${JSONL_FILE}") + +# --- Report --- +if [[ ${LINES_BEFORE} -gt 0 ]]; then + REDUCTION=$(( (LINES_BEFORE - LINES_AFTER) * 100 / LINES_BEFORE )) +else + REDUCTION=0 +fi + +echo "" +echo "=== Compaction Complete ===" +echo "Lines before: ${LINES_BEFORE}" +echo "Lines after: ${LINES_AFTER}" +echo "Reduction: ${REDUCTION}%" +echo "Time: ${ELAPSED_MS}ms" +echo "Backup: ${BACKUP_FILE}" diff --git a/plugins/claude-code/uninstall.sh b/plugins/claude-code/uninstall.sh index e065ea2..fde814f 100755 --- a/plugins/claude-code/uninstall.sh +++ b/plugins/claude-code/uninstall.sh @@ -131,6 +131,18 @@ if [[ "$DO_REMOVE" == true ]]; then ok "No plugin files found at ${INSTALL_DIR}" fi + # Remove pluginDirs entry from settings.json + SETTINGS_FILE="${HOME}/.claude/settings.json" + PLUGIN_PATH="${INSTALL_DIR}/plugin" + if [[ -f "${SETTINGS_FILE}" ]] && command -v jq &>/dev/null; then + if jq -e --arg p "${PLUGIN_PATH}" '.pluginDirs | index($p)' "${SETTINGS_FILE}" >/dev/null 2>&1; then + info "Removing plugin from settings.json..." + jq --arg p "${PLUGIN_PATH}" '.pluginDirs = [.pluginDirs[] | select(. != $p)]' "${SETTINGS_FILE}" > "${SETTINGS_FILE}.tmp" \ + && mv "${SETTINGS_FILE}.tmp" "${SETTINGS_FILE}" + ok "Plugin removed from settings.json" + fi + fi + # Remove log directory LOG_DIR="${HOME}/.cache/supercompact" if [[ -d "$LOG_DIR" ]]; then diff --git a/tests/test_dedup.py b/tests/test_dedup.py new file mode 100644 index 0000000..2d30bac --- /dev/null +++ b/tests/test_dedup.py @@ -0,0 +1,165 @@ +"""Tests for lib/dedup.py — SuffixAutomaton, _turn_unique_ratio, dedup_scores.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.dedup import ( + SuffixAutomaton, + _build_automaton, + _turn_unique_ratio, + dedup_scores, +) +from lib.parser import Turn +from lib.types import ScoredTurn + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(kind: str, index: int, text: str) -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +# --------------------------------------------------------------------------- +# SuffixAutomaton — basic invariants +# --------------------------------------------------------------------------- + +class TestSuffixAutomaton: + def test_initial_state_count(self): + sa = SuffixAutomaton() + assert len(sa.states) == 1 + + def test_extend_adds_states(self): + sa = SuffixAutomaton() + sa.extend("a", 0) + assert len(sa.states) >= 2 + + def test_extend_single_char(self): + sa = SuffixAutomaton() + sa.extend("x", 0) + sa.propagate_counts() + # Should not crash and last state should be set + assert sa.last >= 0 + + def test_match_repeated_length_all_zeros_for_unique(self): + sa = SuffixAutomaton() + for i, c in enumerate("abcdef"): + sa.extend(c, i) + sa.propagate_counts() + lengths = sa.match_repeated_length("abcdef") + # Each char in "abcdef" appears only once → no repeated substrings + assert all(l == 0 for l in lengths) + + def test_match_repeated_length_for_repeated_text(self): + # Build automaton over "aaa" — substring "aa" appears twice + sa = SuffixAutomaton() + text = "aaaa" + for i, c in enumerate(text): + sa.extend(c, i) + sa.propagate_counts() + lengths = sa.match_repeated_length("aa") + # "aa" occurs in "aaaa", so matching "aa" should find repeated lengths > 0 + assert any(l > 0 for l in lengths) + + def test_match_repeated_length_returns_one_per_char(self): + sa = SuffixAutomaton() + for i, c in enumerate("hello"): + sa.extend(c, i) + sa.propagate_counts() + lengths = sa.match_repeated_length("hel") + assert len(lengths) == 3 + + +# --------------------------------------------------------------------------- +# _turn_unique_ratio +# --------------------------------------------------------------------------- + +class TestTurnUniqueRatio: + def _build_sa(self, corpus: str) -> SuffixAutomaton: + sa = SuffixAutomaton() + for i, c in enumerate(corpus): + sa.extend(c, i) + sa.propagate_counts() + return sa + + def test_empty_text_returns_one(self): + sa = self._build_sa("hello world") + assert _turn_unique_ratio(sa, "") == pytest.approx(1.0) + + def test_unique_text_returns_high_ratio(self): + corpus = "the quick brown fox jumps over the lazy dog" + sa = self._build_sa(corpus) + # Text with no long repeated substrings should have high unique ratio + ratio = _turn_unique_ratio(sa, "xyz123abc", min_repeat_len=4) + assert ratio == pytest.approx(1.0) + + def test_highly_repeated_content_returns_lower_ratio(self): + # Build over repeated content + repeated = "ABCD" * 50 + sa = self._build_sa(repeated) + ratio = _turn_unique_ratio(sa, "ABCD" * 20, min_repeat_len=4) + assert ratio < 1.0 + + def test_ratio_between_zero_and_one(self): + sa = self._build_sa("hello world hello world") + ratio = _turn_unique_ratio(sa, "hello", min_repeat_len=3) + assert 0.0 <= ratio <= 1.0 + + +# --------------------------------------------------------------------------- +# dedup_scores +# --------------------------------------------------------------------------- + +class TestDedupScores: + def test_empty_system_turns_returns_empty(self): + turns = [_turn("user", 0, "hello")] + result = dedup_scores(turns, [], {0: 10}) + assert result == [] + + def test_returns_scored_turns(self): + t = _turn("system", 1, "unique content here") + result = dedup_scores([t], [t], {1: 50}) + assert len(result) == 1 + assert isinstance(result[0], ScoredTurn) + + def test_scores_between_zero_and_one(self): + turns = [ + _turn("user", 0, "question"), + _turn("system", 1, "some response"), + _turn("system", 2, "another response"), + ] + system = [turns[1], turns[2]] + token_counts = {0: 10, 1: 20, 2: 20} + result = dedup_scores(turns, system, token_counts) + for st in result: + assert 0.0 <= st.score <= 1.0 + + def test_tokens_set_from_token_counts(self): + t = _turn("system", 5, "content") + result = dedup_scores([t], [t], {5: 999}) + assert result[0].tokens == 999 + + def test_one_result_per_system_turn(self): + turns = [_turn("system", i, f"text{i}") for i in range(4)] + result = dedup_scores(turns, turns, {i: 10 for i in range(4)}) + assert len(result) == 4 + + def test_highly_repeated_turn_scores_lower(self): + # Same text repeated many times — turn has low unique ratio + repeated = "ABCDE" * 40 + all_turns = [_turn("system", i, repeated) for i in range(3)] + system = all_turns[:] + token_counts = {i: 200 for i in range(3)} + result = dedup_scores(all_turns, system, token_counts, min_repeat_len=5) + # At least one turn should have lower score (repeated content) + scores = [st.score for st in result] + assert min(scores) < 1.0 diff --git a/tests/test_eitf.py b/tests/test_eitf.py new file mode 100644 index 0000000..7d1ba12 --- /dev/null +++ b/tests/test_eitf.py @@ -0,0 +1,117 @@ +"""Tests for lib/eitf.py — eitf_scores.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.eitf import eitf_scores +from lib.parser import Turn +from lib.types import ScoredTurn + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(kind: str, index: int, text: str) -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +# --------------------------------------------------------------------------- +# eitf_scores +# --------------------------------------------------------------------------- + +class TestEitfScoresEmpty: + def test_no_system_turns_returns_empty(self): + turns = [_turn("user", 0, "hello")] + result = eitf_scores(turns, [], {0: 10}) + assert result == [] + + def test_empty_all_turns_returns_empty(self): + result = eitf_scores([], [], {}) + assert result == [] + + +class TestEitfScoresBasic: + def test_returns_scored_turns(self): + t = _turn("system", 1, "Error: ValueError on line 42") + all_turns = [_turn("user", 0, "question"), t] + result = eitf_scores(all_turns, [t], {0: 10, 1: 50}) + assert len(result) == 1 + assert isinstance(result[0], ScoredTurn) + + def test_scores_normalized_to_zero_one(self): + turns = [ + _turn("user", 0, "what is the error?"), + _turn("system", 1, "ValueError on line 42 in /src/main.py"), + _turn("system", 2, "ModuleNotFoundError for numpy package"), + ] + system = turns[1:] + token_counts = {0: 10, 1: 80, 2: 60} + result = eitf_scores(turns, system, token_counts) + for st in result: + assert 0.0 <= st.score <= 1.0 + + def test_one_result_per_system_turn(self): + turns = [_turn("system", i, f"content{i}") for i in range(4)] + result = eitf_scores(turns, turns, {i: 50 for i in range(4)}) + assert len(result) == 4 + + def test_max_score_is_one(self): + turns = [ + _turn("user", 0, "question"), + _turn("system", 1, "ValueError at /path/to/file.py:42"), + _turn("system", 2, "plain content"), + ] + system = turns[1:] + token_counts = {0: 10, 1: 80, 2: 40} + result = eitf_scores(turns, system, token_counts) + scores = [st.score for st in result] + assert max(scores) == pytest.approx(1.0) + + def test_tokens_set_from_token_counts(self): + t = _turn("system", 3, "some content with ValueError") + all_turns = [t] + result = eitf_scores(all_turns, [t], {3: 777}) + assert result[0].tokens == 777 + + def test_turn_with_rare_entities_scores_higher(self): + # Turn 1 has a unique entity (rare), turn 2 has a common entity + # that appears in all turns + common_err = "ValueError" + unique_path = "/home/user/special_project/module.py" + + turns = [ + _turn("user", 0, f"question about {common_err}"), + _turn("system", 1, f"Error in {unique_path}: {common_err}"), + _turn("system", 2, f"{common_err} appears here"), + _turn("system", 3, f"{common_err} again"), + ] + system = turns[1:] + token_counts = {i: 50 for i in range(4)} + result = eitf_scores(turns, system, token_counts) + # turn 1 (index 1) has the unique path entity, should score higher + # than turn 3 which only has the common entity + score_map = {st.turn.index: st.score for st in result} + assert score_map[1] >= score_map[3] + + def test_plain_content_no_entities_scores_zero(self): + # A turn with no extractable entities should score 0 (before normalization) + # but with normalization it can be 0 if all turns have no entities + turns = [ + _turn("system", 0, "aaaaa bbbbb ccccc"), + _turn("system", 1, "zzzzz yyyyy xxxxx"), + ] + system = turns[:] + result = eitf_scores(turns, system, {0: 50, 1: 50}) + scores = [st.score for st in result] + # All should be in [0,1] + for s in scores: + assert 0.0 <= s <= 1.0 diff --git a/tests/test_eval_aggregate.py b/tests/test_eval_aggregate.py new file mode 100644 index 0000000..20114fd --- /dev/null +++ b/tests/test_eval_aggregate.py @@ -0,0 +1,272 @@ +"""Tests for lib/eval/aggregate.py — _dcg, DimensionScore, AggregateResult, aggregate.""" + +from __future__ import annotations + +import math +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.eval.aggregate import ( + DIFFICULTY_WEIGHTS, + AggregateResult, + DimensionScore, + _dcg, + aggregate, +) +from lib.eval.judge import ProbeAnswer +from lib.eval.probes import Probe, ProbeSet + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _probe(id="p1", dimension="progress", difficulty="medium") -> Probe: + return Probe( + id=id, + dimension=dimension, + tier="factual", + question="Q?", + gold_answer="A", + difficulty=difficulty, + ) + + +def _answer(probe_id="p1", model_key="cheap", score=2) -> ProbeAnswer: + return ProbeAnswer( + probe_id=probe_id, + model_key=model_key, + model_label="haiku", + answer="x", + score=score, + ) + + +def _probe_set(*probes) -> ProbeSet: + return ProbeSet(probes=list(probes)) + + +# --------------------------------------------------------------------------- +# DIFFICULTY_WEIGHTS constant +# --------------------------------------------------------------------------- + +class TestDifficultyWeights: + def test_easy_has_weight_one(self): + assert DIFFICULTY_WEIGHTS["easy"] == 1.0 + + def test_medium_has_weight_two(self): + assert DIFFICULTY_WEIGHTS["medium"] == 2.0 + + def test_hard_has_weight_three(self): + assert DIFFICULTY_WEIGHTS["hard"] == 3.0 + + +# --------------------------------------------------------------------------- +# DimensionScore dataclass +# --------------------------------------------------------------------------- + +class TestDimensionScore: + def test_fields_stored(self): + ds = DimensionScore( + dimension="progress", + weight=0.25, + mean_score=0.8, + probe_count=5, + raw_scores=[2, 3, 2, 3, 2], + ) + assert ds.dimension == "progress" + assert ds.weight == 0.25 + assert ds.mean_score == pytest.approx(0.8) + assert ds.probe_count == 5 + assert ds.raw_scores == [2, 3, 2, 3, 2] + + def test_default_raw_scores_empty(self): + ds = DimensionScore(dimension="noise", weight=0.05, mean_score=0.0, probe_count=0) + assert ds.raw_scores == [] + + +# --------------------------------------------------------------------------- +# AggregateResult dataclass +# --------------------------------------------------------------------------- + +class TestAggregateResult: + def test_fields_stored(self): + ar = AggregateResult( + method="dedup", budget=80_000, model_key="cheap", model_label="haiku" + ) + assert ar.method == "dedup" + assert ar.budget == 80_000 + assert ar.model_key == "cheap" + assert ar.model_label == "haiku" + + def test_default_composite_zero(self): + ar = AggregateResult(method="x", budget=0, model_key="k", model_label="l") + assert ar.composite == 0.0 + + def test_default_ndcg_zero(self): + ar = AggregateResult(method="x", budget=0, model_key="k", model_label="l") + assert ar.ndcg == 0.0 + + def test_default_dimensions_empty(self): + ar = AggregateResult(method="x", budget=0, model_key="k", model_label="l") + assert ar.dimensions == [] + + def test_dimension_map_property(self): + ds = DimensionScore(dimension="progress", weight=0.25, mean_score=0.7, probe_count=3) + ar = AggregateResult( + method="dedup", budget=80_000, model_key="cheap", model_label="haiku", + dimensions=[ds], + ) + dm = ar.dimension_map + assert "progress" in dm + assert dm["progress"] is ds + + def test_dimension_map_empty_when_no_dimensions(self): + ar = AggregateResult(method="x", budget=0, model_key="k", model_label="l") + assert ar.dimension_map == {} + + +# --------------------------------------------------------------------------- +# _dcg +# --------------------------------------------------------------------------- + +class TestDcg: + def test_empty_list_returns_zero(self): + assert _dcg([]) == pytest.approx(0.0) + + def test_single_item(self): + # score=3, weight=1.0 → 3*1 / log2(0+2) = 3/1 = 3 + result = _dcg([(3, 1.0)]) + assert result == pytest.approx(3.0 / math.log2(2)) + + def test_higher_weight_item_placed_first(self): + # Items sorted by weight desc: hard(3.0) before easy(1.0) + # So hard item gets position 0, easy gets position 1 + hard = (2, 3.0) + easy = (2, 1.0) + result = _dcg([easy, hard]) # out of order by weight + expected = (2 * 3.0) / math.log2(2) + (2 * 1.0) / math.log2(3) + assert result == pytest.approx(expected) + + def test_all_zero_scores_returns_zero(self): + items = [(0, 1.0), (0, 2.0), (0, 3.0)] + assert _dcg(items) == pytest.approx(0.0) + + def test_higher_position_discounts_more(self): + # Same score+weight at different positions + result_first = _dcg([(3, 2.0)]) + result_second_only = [(0, 3.0), (3, 2.0)] + # Single item should be larger than if it were at position 1 + result_two = _dcg(result_second_only) + assert result_first > result_two - _dcg([(0, 3.0)]) + + +# --------------------------------------------------------------------------- +# aggregate — basic paths +# --------------------------------------------------------------------------- + +class TestAggregateEmpty: + def test_no_answers_returns_empty(self): + ps = _probe_set(_probe()) + assert aggregate([], ps, "dedup", 80_000) == [] + + def test_no_probes_no_answers_returns_empty(self): + ps = _probe_set() + assert aggregate([], ps, "dedup", 80_000) == [] + + +class TestAggregateSingleModel: + def test_returns_one_result_per_model(self): + ps = _probe_set(_probe("p1", "progress")) + answers = [_answer("p1", "cheap", score=3)] + results = aggregate(answers, ps, "dedup", 80_000) + assert len(results) == 1 + + def test_result_has_correct_method_and_budget(self): + ps = _probe_set(_probe("p1", "progress")) + answers = [_answer("p1", "cheap", score=2)] + result = aggregate(answers, ps, "eitf", 40_000)[0] + assert result.method == "eitf" + assert result.budget == 40_000 + + def test_result_has_model_key(self): + ps = _probe_set(_probe("p1", "progress")) + answers = [_answer("p1", "cheap", score=2)] + result = aggregate(answers, ps, "dedup", 80_000)[0] + assert result.model_key == "cheap" + + def test_score_normalized_to_zero_one(self): + ps = _probe_set(_probe("p1", "progress")) + answers = [_answer("p1", "cheap", score=3)] + result = aggregate(answers, ps, "dedup", 80_000)[0] + progress_dim = result.dimension_map["progress"] + assert progress_dim.mean_score == pytest.approx(1.0) # 3/3 + + def test_score_zero_normalizes_to_zero(self): + ps = _probe_set(_probe("p1", "progress")) + answers = [_answer("p1", "cheap", score=0)] + result = aggregate(answers, ps, "dedup", 80_000)[0] + assert result.dimension_map["progress"].mean_score == pytest.approx(0.0) + + def test_probe_not_in_probe_set_skipped(self): + ps = _probe_set() # empty + answers = [_answer("ghost", "cheap", score=3)] + results = aggregate(answers, ps, "dedup", 80_000) + # No probes matched → model group exists but no dims scored + if results: + # All dimension scores should have probe_count=0 + for ds in results[0].dimensions: + assert ds.probe_count == 0 + + def test_dimension_with_no_probes_has_zero_mean(self): + ps = _probe_set(_probe("p1", "progress")) # only progress + answers = [_answer("p1", "cheap", score=3)] + result = aggregate(answers, ps, "dedup", 80_000)[0] + # error_solution has no probes + esr_dim = result.dimension_map.get("error_solution") + if esr_dim: + assert esr_dim.mean_score == pytest.approx(0.0) + assert esr_dim.probe_count == 0 + + +class TestAggregateMultipleModels: + def test_two_models_give_two_results(self): + ps = _probe_set(_probe("p1", "progress")) + answers = [ + _answer("p1", "cheap", score=2), + _answer("p1", "capable", score=3), + ] + results = aggregate(answers, ps, "dedup", 80_000) + assert len(results) == 2 + model_keys = {r.model_key for r in results} + assert model_keys == {"cheap", "capable"} + + +class TestAggregateNdcg: + def test_perfect_scores_give_ndcg_one(self): + ps = _probe_set(_probe("p1", "progress", difficulty="medium")) + answers = [_answer("p1", "cheap", score=3)] + result = aggregate(answers, ps, "dedup", 80_000)[0] + assert result.ndcg == pytest.approx(1.0) + + def test_zero_scores_give_ndcg_zero(self): + ps = _probe_set(_probe("p1", "progress", difficulty="medium")) + answers = [_answer("p1", "cheap", score=0)] + result = aggregate(answers, ps, "dedup", 80_000)[0] + assert result.ndcg == pytest.approx(0.0) + + def test_ndcg_between_zero_and_one(self): + ps = _probe_set( + _probe("p1", "progress", difficulty="easy"), + _probe("p2", "instruction", difficulty="hard"), + ) + answers = [ + _answer("p1", "cheap", score=1), + _answer("p2", "cheap", score=2), + ] + result = aggregate(answers, ps, "dedup", 80_000)[0] + assert 0.0 <= result.ndcg <= 1.0 diff --git a/tests/test_eval_entity_coverage.py b/tests/test_eval_entity_coverage.py new file mode 100644 index 0000000..fe65e42 --- /dev/null +++ b/tests/test_eval_entity_coverage.py @@ -0,0 +1,256 @@ +"""Tests for lib/eval/entity_coverage.py — EntitySet, extract_entities, compute_coverage.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.eval.entity_coverage import ( + ENTITY_TYPES, + EntitySet, + compute_coverage, + extract_entities, +) + + +# --------------------------------------------------------------------------- +# EntitySet +# --------------------------------------------------------------------------- + +class TestEntitySet: + def test_total_count_empty(self): + es = EntitySet() + assert es.total_count == 0 + + def test_total_count_single_type(self): + es = EntitySet(entities={"file_path": {"/foo/bar/baz", "/a/b/c"}}) + assert es.total_count == 2 + + def test_total_count_multiple_types(self): + es = EntitySet(entities={ + "file_path": {"/foo/bar"}, + "exception": {"ValueError", "TypeError"}, + }) + assert es.total_count == 3 + + def test_all_entities_empty(self): + es = EntitySet() + assert es.all_entities() == set() + + def test_all_entities_returns_type_value_pairs(self): + es = EntitySet(entities={"exception": {"ValueError"}}) + assert ("exception", "ValueError") in es.all_entities() + + def test_all_entities_multiple_types(self): + es = EntitySet(entities={ + "exception": {"ValueError"}, + "port": {"8080"}, + }) + pairs = es.all_entities() + assert ("exception", "ValueError") in pairs + assert ("port", "8080") in pairs + + def test_default_entities_is_empty_dict(self): + es = EntitySet() + assert es.entities == {} + + +# --------------------------------------------------------------------------- +# ENTITY_TYPES constant +# --------------------------------------------------------------------------- + +class TestEntityTypes: + def test_has_file_path(self): + assert "file_path" in ENTITY_TYPES + + def test_has_exception(self): + assert "exception" in ENTITY_TYPES + + def test_has_url(self): + assert "url" in ENTITY_TYPES + + def test_weights_are_positive(self): + for k, v in ENTITY_TYPES.items(): + assert v > 0, f"Weight for {k} should be positive" + + def test_file_path_and_error_are_highest_weight(self): + assert ENTITY_TYPES["file_path"] >= 1.0 + assert ENTITY_TYPES["error"] >= 1.0 + + +# --------------------------------------------------------------------------- +# extract_entities — individual types +# --------------------------------------------------------------------------- + +class TestExtractEntitiesExceptions: + def test_extracts_value_error(self): + es = extract_entities("Traceback: ValueError: invalid literal") + assert "exception" in es.entities + assert "valueerror" in es.entities["exception"] + + def test_extracts_module_not_found(self): + es = extract_entities("ModuleNotFoundError: No module named 'requests'") + assert "modulenotfounderror" in es.entities.get("exception", set()) + + def test_extracts_type_error(self): + es = extract_entities("Got TypeError when calling the function") + assert "typeerror" in es.entities.get("exception", set()) + + +class TestExtractEntitiesUrls: + def test_extracts_https_url(self): + es = extract_entities("Visit https://example.com/api for docs") + assert "url" in es.entities + assert any("example.com" in u for u in es.entities["url"]) + + def test_extracts_http_url(self): + es = extract_entities("Server running at http://localhost:8080/api") + assert "url" in es.entities + assert any("localhost" in u for u in es.entities["url"]) + + def test_no_url_in_plain_text(self): + es = extract_entities("Just some plain text without any links") + assert "url" not in es.entities or len(es.entities.get("url", set())) == 0 + + +class TestExtractEntitiesPorts: + def test_extracts_colon_port(self): + es = extract_entities("Server running on port :8080 now") + assert "port" in es.entities + assert "8080" in es.entities["port"] + + def test_extracts_port_keyword(self): + es = extract_entities("PORT=3000 must be set") + assert "port" in es.entities + assert "3000" in es.entities["port"] + + def test_does_not_extract_low_port_numbers(self): + # Ports 100-999 should be filtered as false positives + es = extract_entities("error code :500 status") + # Port 500 should not appear (it's in the filtered range 100-999) + ports = es.entities.get("port", set()) + assert "500" not in ports + + +class TestExtractEntitiesFilePaths: + def test_extracts_absolute_path(self): + es = extract_entities("Error in /home/user/project/src/main.py line 42") + assert "file_path" in es.entities + assert any("/home/user" in p for p in es.entities["file_path"]) + + def test_url_path_not_treated_as_file_path(self): + es = extract_entities("See https://example.com/foo/bar/baz for more") + # /foo/bar/baz comes from a URL, should not also appear as file_path + file_paths = es.entities.get("file_path", set()) + # The URL itself is captured; path inside URL should not be double-captured + urls = es.entities.get("url", set()) + assert any("example.com" in u for u in urls) + + +class TestExtractEntitiesClassNames: + def test_extracts_camel_case_class(self): + es = extract_entities("The ProbeAnswer class stores results") + assert "class_name" in es.entities + classes = {c.lower() for c in es.entities["class_name"]} + assert "probeanswer" in classes + + def test_extracts_multiple_classes(self): + es = extract_entities("Use DimensionScore and AggregateResult") + classes = {c.lower() for c in es.entities.get("class_name", set())} + assert "dimensionscore" in classes + assert "aggregateresult" in classes + + +class TestExtractEntitiesPackages: + def test_extracts_pip_install(self): + es = extract_entities("Run: pip install requests httpx") + assert "package" in es.entities + assert "requests" in es.entities["package"] + + def test_extracts_npm_install(self): + es = extract_entities("npm install express") + assert "package" in es.entities + assert "express" in es.entities["package"] + + +class TestExtractEntitiesHttpStatus: + def test_extracts_404_not_found(self): + es = extract_entities("Got 404 Not Found from the API") + assert "http_status" in es.entities + assert "404" in es.entities["http_status"] + + def test_extracts_500_internal(self): + es = extract_entities("Server returned 500 Internal Server Error") + assert "http_status" in es.entities + assert "500" in es.entities["http_status"] + + +# --------------------------------------------------------------------------- +# compute_coverage +# --------------------------------------------------------------------------- + +class TestComputeCoverageEmpty: + def test_empty_suffix_returns_ones(self): + suffix = EntitySet() + kept = EntitySet(entities={"exception": {"ValueError"}}) + unweighted, weighted, breakdown = compute_coverage(suffix, kept) + assert unweighted == pytest.approx(1.0) + assert weighted == pytest.approx(1.0) + assert breakdown == {} + + def test_empty_kept_with_suffix_gives_zero_coverage(self): + suffix = EntitySet(entities={"exception": {"ValueError"}}) + kept = EntitySet() + unweighted, weighted, breakdown = compute_coverage(suffix, kept) + assert unweighted == pytest.approx(0.0) + assert weighted == pytest.approx(0.0) + + +class TestComputeCoverageFull: + def test_identical_sets_give_full_coverage(self): + entities = {"exception": {"valueerror", "typeerror"}} + suffix = EntitySet(entities=entities) + kept = EntitySet(entities=entities) + unweighted, weighted, breakdown = compute_coverage(suffix, kept) + assert unweighted == pytest.approx(1.0) + assert weighted == pytest.approx(1.0) + + def test_breakdown_has_type_info(self): + suffix = EntitySet(entities={"exception": {"valueerror"}}) + kept = EntitySet(entities={"exception": {"valueerror"}}) + _, _, breakdown = compute_coverage(suffix, kept) + assert "exception" in breakdown + assert breakdown["exception"]["covered"] == 1 + assert breakdown["exception"]["total"] == 1 + assert breakdown["exception"]["coverage"] == pytest.approx(1.0) + + +class TestComputeCoveragePartial: + def test_half_covered(self): + suffix = EntitySet(entities={"exception": {"valueerror", "typeerror"}}) + kept = EntitySet(entities={"exception": {"valueerror"}}) + unweighted, _, _ = compute_coverage(suffix, kept) + assert unweighted == pytest.approx(0.5) + + def test_unmatched_type_in_kept_not_counted(self): + suffix = EntitySet(entities={"exception": {"valueerror"}}) + # kept has a different type — shouldn't count toward coverage + kept = EntitySet(entities={"port": {"8080"}}) + unweighted, _, _ = compute_coverage(suffix, kept) + assert unweighted == pytest.approx(0.0) + + def test_weighted_coverage_differs_by_type_importance(self): + # file_path has weight 1.0, function has weight 0.5 + # If we have one file_path and one function uncovered, weights differ + suffix = EntitySet(entities={ + "file_path": {"/foo/bar/baz"}, + "function": {"my_function"}, + }) + kept = EntitySet(entities={"file_path": {"/foo/bar/baz"}}) + _, weighted, _ = compute_coverage(suffix, kept) + # file_path covered, function not — weighted should be > 0 and < 1 + assert 0.0 < weighted < 1.0 diff --git a/tests/test_eval_evidence_coverage.py b/tests/test_eval_evidence_coverage.py new file mode 100644 index 0000000..29897d5 --- /dev/null +++ b/tests/test_eval_evidence_coverage.py @@ -0,0 +1,253 @@ +"""Tests for lib/eval/evidence_coverage.py — dataclasses, _dcg, compute_evidence_coverage.""" + +from __future__ import annotations + +import math +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.eval.evidence_coverage import ( + DIFFICULTY_WEIGHTS, + DimensionCoverage, + EvidenceCoverageResult, + ProbeCoverage, + _dcg, + compute_evidence_coverage, +) +from lib.eval.probes import Probe, ProbeSet + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _probe(id="p1", dimension="progress", difficulty="medium", + evidence_turns=None) -> Probe: + return Probe( + id=id, + dimension=dimension, + tier="factual", + question="Q?", + gold_answer="A", + difficulty=difficulty, + evidence_turns=evidence_turns or [], + ) + + +def _probe_set(*probes) -> ProbeSet: + return ProbeSet(probes=list(probes)) + + +# --------------------------------------------------------------------------- +# DIFFICULTY_WEIGHTS constant +# --------------------------------------------------------------------------- + +class TestDifficultyWeights: + def test_easy_is_one(self): + assert DIFFICULTY_WEIGHTS["easy"] == 1.0 + + def test_medium_is_two(self): + assert DIFFICULTY_WEIGHTS["medium"] == 2.0 + + def test_hard_is_three(self): + assert DIFFICULTY_WEIGHTS["hard"] == 3.0 + + +# --------------------------------------------------------------------------- +# ProbeCoverage dataclass +# --------------------------------------------------------------------------- + +class TestProbeCoverage: + def test_fields_stored(self): + pc = ProbeCoverage( + probe_id="p1", + dimension="progress", + difficulty="medium", + evidence_turns=[0, 1, 2], + kept_evidence=[0, 1], + dropped_evidence=[2], + coverage=2/3, + ) + assert pc.probe_id == "p1" + assert pc.dimension == "progress" + assert pc.evidence_turns == [0, 1, 2] + assert pc.kept_evidence == [0, 1] + assert pc.dropped_evidence == [2] + assert pc.coverage == pytest.approx(2/3) + + +# --------------------------------------------------------------------------- +# DimensionCoverage dataclass +# --------------------------------------------------------------------------- + +class TestDimensionCoverage: + def test_fields_stored(self): + dc = DimensionCoverage( + dimension="instruction", + weight=0.25, + mean_coverage=0.8, + probe_count=4, + coverages=[0.5, 0.75, 1.0, 0.9], + ) + assert dc.dimension == "instruction" + assert dc.mean_coverage == pytest.approx(0.8) + assert dc.probe_count == 4 + + def test_default_coverages_empty(self): + dc = DimensionCoverage(dimension="noise", weight=0.05, mean_coverage=0.0, probe_count=0) + assert dc.coverages == [] + + +# --------------------------------------------------------------------------- +# EvidenceCoverageResult dataclass +# --------------------------------------------------------------------------- + +class TestEvidenceCoverageResult: + def test_defaults(self): + ecr = EvidenceCoverageResult(method="dedup", budget=80_000) + assert ecr.composite == 0.0 + assert ecr.ndcg == 0.0 + assert ecr.dimensions == [] + assert ecr.probe_details == [] + + def test_dimension_map_property(self): + dc = DimensionCoverage(dimension="progress", weight=0.25, mean_coverage=0.7, probe_count=2) + ecr = EvidenceCoverageResult(method="dedup", budget=80_000, dimensions=[dc]) + assert "progress" in ecr.dimension_map + assert ecr.dimension_map["progress"] is dc + + def test_to_dict_has_required_keys(self): + ecr = EvidenceCoverageResult(method="dedup", budget=80_000) + d = ecr.to_dict() + for key in ("method", "budget", "composite", "ndcg", "dimensions", "probe_details"): + assert key in d + + def test_to_dict_method_budget(self): + ecr = EvidenceCoverageResult(method="eitf", budget=40_000) + d = ecr.to_dict() + assert d["method"] == "eitf" + assert d["budget"] == 40_000 + + +# --------------------------------------------------------------------------- +# _dcg +# --------------------------------------------------------------------------- + +class TestDcg: + def test_empty_returns_zero(self): + assert _dcg([]) == pytest.approx(0.0) + + def test_single_item(self): + result = _dcg([(1.0, 2.0)]) + assert result == pytest.approx(2.0 / math.log2(2)) + + def test_zero_score_returns_zero(self): + assert _dcg([(0.0, 3.0), (0.0, 1.0)]) == pytest.approx(0.0) + + def test_sorted_by_weight_desc(self): + # Higher weight item should come first (lower denominator) + hard = (0.5, 3.0) + easy = (0.5, 1.0) + result = _dcg([easy, hard]) + expected = (0.5 * 3.0) / math.log2(2) + (0.5 * 1.0) / math.log2(3) + assert result == pytest.approx(expected) + + +# --------------------------------------------------------------------------- +# compute_evidence_coverage — core paths +# --------------------------------------------------------------------------- + +class TestComputeEvidenceCoverageEmpty: + def test_no_probes_returns_zero_composite(self): + ps = _probe_set() + result = compute_evidence_coverage(ps, {0, 1}, "dedup", 80_000) + assert result.composite == pytest.approx(0.0) + + def test_no_probes_returns_ndcg_zero(self): + ps = _probe_set() + result = compute_evidence_coverage(ps, {0, 1}, "dedup", 80_000) + assert result.ndcg == pytest.approx(0.0) + + def test_no_probes_returns_no_probe_details(self): + ps = _probe_set() + result = compute_evidence_coverage(ps, set(), "dedup", 80_000) + assert result.probe_details == [] + + def test_probe_without_evidence_turns_is_skipped(self): + ps = _probe_set(_probe("p1", evidence_turns=[])) + result = compute_evidence_coverage(ps, {0, 1}, "dedup", 80_000) + assert result.probe_details == [] + + def test_method_and_budget_stored(self): + ps = _probe_set() + result = compute_evidence_coverage(ps, set(), "setcover", 60_000) + assert result.method == "setcover" + assert result.budget == 60_000 + + +class TestComputeEvidenceCoverageFullCoverage: + def test_all_evidence_kept_gives_coverage_one(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[0, 1, 2])) + result = compute_evidence_coverage(ps, {0, 1, 2}, "dedup", 80_000) + assert result.probe_details[0].coverage == pytest.approx(1.0) + + def test_all_evidence_kept_gives_ndcg_one(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[0, 1])) + result = compute_evidence_coverage(ps, {0, 1}, "dedup", 80_000) + assert result.ndcg == pytest.approx(1.0) + + def test_all_evidence_kept_gives_non_zero_composite(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[5])) + result = compute_evidence_coverage(ps, {5}, "dedup", 80_000) + assert result.composite > 0.0 + + +class TestComputeEvidenceCoverageZeroCoverage: + def test_no_evidence_kept_gives_coverage_zero(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[3, 4])) + result = compute_evidence_coverage(ps, set(), "dedup", 80_000) + assert result.probe_details[0].coverage == pytest.approx(0.0) + + def test_no_evidence_kept_ndcg_zero(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[3])) + result = compute_evidence_coverage(ps, {0, 1, 2}, "dedup", 80_000) + assert result.ndcg == pytest.approx(0.0) + + +class TestComputeEvidenceCoveragePartial: + def test_partial_coverage_value(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[0, 1, 2, 3])) + result = compute_evidence_coverage(ps, {0, 1}, "dedup", 80_000) + assert result.probe_details[0].coverage == pytest.approx(0.5) + + def test_kept_and_dropped_lists(self): + ps = _probe_set(_probe("p1", "progress", evidence_turns=[0, 1, 2])) + result = compute_evidence_coverage(ps, {0, 2}, "dedup", 80_000) + pc = result.probe_details[0] + assert set(pc.kept_evidence) == {0, 2} + assert set(pc.dropped_evidence) == {1} + + +class TestComputeEvidenceCoverageDimension: + def test_multiple_probes_mean_coverage(self): + ps = _probe_set( + _probe("p1", "progress", evidence_turns=[0, 1]), + _probe("p2", "progress", evidence_turns=[2, 3]), + ) + # Keep all for p1 (coverage=1.0), none for p2 (coverage=0.0) + result = compute_evidence_coverage(ps, {0, 1}, "dedup", 80_000) + progress = result.dimension_map["progress"] + assert progress.mean_coverage == pytest.approx(0.5) + assert progress.probe_count == 2 + + def test_ndcg_between_zero_and_one_for_partial(self): + ps = _probe_set( + _probe("p1", "progress", difficulty="medium", evidence_turns=[0]), + _probe("p2", "instruction", difficulty="hard", evidence_turns=[1]), + ) + result = compute_evidence_coverage(ps, {0}, "dedup", 80_000) + assert 0.0 <= result.ndcg <= 1.0 diff --git a/tests/test_eval_judge.py b/tests/test_eval_judge.py new file mode 100644 index 0000000..1dfbbe7 --- /dev/null +++ b/tests/test_eval_judge.py @@ -0,0 +1,237 @@ +"""Tests for lib/eval/judge.py — ProbeAnswer, JudgeResult, and scoring helpers.""" + +from __future__ import annotations + +import asyncio +import sys +from pathlib import Path +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.eval.judge import ( + ANSWER_MODELS, + JUDGE_MODEL, + JudgeResult, + ProbeAnswer, + _openrouter_generate_async, + _score_one_answer, + generate_answers, + score_answers, +) +from lib.eval.probes import Probe, ProbeSet + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_probe(id="p1", question="Q?", gold_answer="A") -> Probe: + return Probe( + id=id, + dimension="progress", + tier="factual", + question=question, + gold_answer=gold_answer, + ) + + +def _make_answer(probe_id="p1", score=-1) -> ProbeAnswer: + return ProbeAnswer( + probe_id=probe_id, + model_key="cheap", + model_label="haiku", + answer="some answer", + score=score, + ) + + +async def _run_score_one(raw_response: str) -> ProbeAnswer: + """Run _score_one_answer with a mocked _openrouter_generate_async.""" + answer = _make_answer() + probe = _make_probe() + + async def _mock_gen(*args, **kwargs): + return raw_response + + sem = asyncio.Semaphore(1) + client = MagicMock() + + with patch("lib.eval.judge._openrouter_generate_async", side_effect=_mock_gen): + await _score_one_answer(client, sem, answer, probe, "judge-model") + + return answer + + +# --------------------------------------------------------------------------- +# ProbeAnswer dataclass +# --------------------------------------------------------------------------- + +class TestProbeAnswerDataclass: + def test_default_score_is_minus_one(self): + a = ProbeAnswer(probe_id="p1", model_key="cheap", model_label="haiku", answer="x") + assert a.score == -1 + + def test_default_judge_reasoning_is_empty(self): + a = ProbeAnswer(probe_id="p1", model_key="cheap", model_label="haiku", answer="x") + assert a.judge_reasoning == "" + + def test_fields_stored(self): + a = ProbeAnswer(probe_id="myid", model_key="capable", model_label="gpt4", answer="hello") + assert a.probe_id == "myid" + assert a.model_key == "capable" + assert a.model_label == "gpt4" + assert a.answer == "hello" + + def test_custom_score_and_reasoning(self): + a = ProbeAnswer( + probe_id="p2", + model_key="cheap", + model_label="haiku", + answer="blah", + score=3, + judge_reasoning="perfect answer", + ) + assert a.score == 3 + assert a.judge_reasoning == "perfect answer" + + +# --------------------------------------------------------------------------- +# JudgeResult dataclass +# --------------------------------------------------------------------------- + +class TestJudgeResultDataclass: + def test_method_and_budget_stored(self): + jr = JudgeResult(method="dedup", budget=80_000) + assert jr.method == "dedup" + assert jr.budget == 80_000 + + def test_default_answers_is_empty_list(self): + jr = JudgeResult(method="eitf", budget=40_000) + assert jr.answers == [] + + def test_answers_mutable(self): + jr = JudgeResult(method="dedup", budget=80_000) + jr.answers.append(_make_answer()) + assert len(jr.answers) == 1 + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +class TestConstants: + def test_judge_model_is_string(self): + assert isinstance(JUDGE_MODEL, str) + assert len(JUDGE_MODEL) > 0 + + def test_answer_models_has_capable_and_cheap(self): + assert "capable" in ANSWER_MODELS + assert "cheap" in ANSWER_MODELS + + def test_answer_model_entries_have_model_key(self): + for key, cfg in ANSWER_MODELS.items(): + assert "model" in cfg + assert "label" in cfg + + +# --------------------------------------------------------------------------- +# generate_answers — no probes +# --------------------------------------------------------------------------- + +class TestGenerateAnswersEmpty: + def test_no_probes_returns_empty_list(self, monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + probe_set = ProbeSet(probes=[]) + result = generate_answers("some context", probe_set) + assert result == [] + + def test_no_probes_resets_counter(self, monkeypatch): + import lib.eval.judge as judge_mod + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + probe_set = ProbeSet(probes=[]) + judge_mod._answer_counter = 99 + generate_answers("ctx", probe_set) + assert judge_mod._answer_counter == 0 + + +# --------------------------------------------------------------------------- +# score_answers — no-API paths +# --------------------------------------------------------------------------- + +class TestScoreAnswersEmpty: + def test_no_answers_returns_same_list(self, monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + probe_set = ProbeSet(probes=[_make_probe()]) + result = score_answers([], probe_set) + assert result == [] + + +class TestScoreAnswersMissingProbe: + def test_missing_probe_sets_score_zero(self, monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + probe_set = ProbeSet(probes=[]) # no probes registered + answer = _make_answer(probe_id="ghost_probe") + score_answers([answer], probe_set) + assert answer.score == 0 + + def test_missing_probe_sets_reasoning(self, monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "test-key") + probe_set = ProbeSet(probes=[]) + answer = _make_answer(probe_id="ghost_probe") + score_answers([answer], probe_set) + assert "Probe not found" in answer.judge_reasoning + + +# --------------------------------------------------------------------------- +# _openrouter_generate_async — no API key +# --------------------------------------------------------------------------- + +class TestOpenrouterMissingKey: + def test_missing_key_raises_runtime_error(self, monkeypatch): + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + import httpx + + async def run(): + async with httpx.AsyncClient() as client: + return await _openrouter_generate_async(client, "model", "sys", "user") + + with pytest.raises(RuntimeError, match="OPENROUTER_API_KEY"): + asyncio.run(run()) + + +# --------------------------------------------------------------------------- +# _score_one_answer — JSON parsing, fence stripping, clamping +# --------------------------------------------------------------------------- + +class TestScoreOneAnswer: + def test_valid_json_sets_score(self): + answer = asyncio.run(_run_score_one('{"score": 2, "reasoning": "partial"}')) + assert answer.score == 2 + + def test_valid_json_sets_reasoning(self): + answer = asyncio.run(_run_score_one('{"score": 3, "reasoning": "complete match"}')) + assert answer.judge_reasoning == "complete match" + + def test_markdown_fence_stripped(self): + fenced = "```json\n{\"score\": 1, \"reasoning\": \"ok\"}\n```" + answer = asyncio.run(_run_score_one(fenced)) + assert answer.score == 1 + + def test_score_clamped_above_three(self): + answer = asyncio.run(_run_score_one('{"score": 9, "reasoning": "too high"}')) + assert answer.score == 3 + + def test_score_clamped_below_zero(self): + answer = asyncio.run(_run_score_one('{"score": -5, "reasoning": "too low"}')) + assert answer.score == 0 + + def test_bad_json_sets_score_zero(self): + answer = asyncio.run(_run_score_one("not valid json at all")) + assert answer.score == 0 + + def test_bad_json_sets_judge_error_reasoning(self): + answer = asyncio.run(_run_score_one("not valid json at all")) + assert "Judge error" in answer.judge_reasoning diff --git a/tests/test_eval_report.py b/tests/test_eval_report.py new file mode 100644 index 0000000..c3c9091 --- /dev/null +++ b/tests/test_eval_report.py @@ -0,0 +1,200 @@ +"""Tests for lib/eval/report.py — export_json and export_trace.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path +from unittest.mock import MagicMock + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.eval.report import export_json, export_trace +from lib.eval.aggregate import AggregateResult, DimensionScore + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_result(**overrides) -> AggregateResult: + defaults = dict( + method="dedup", + budget=80_000, + model_key="cheap", + model_label="claude-haiku", + dimensions=[], + composite=0.75, + ndcg=0.82, + speed_s=1.5, + kept_tokens=20_000, + total_tokens=100_000, + ) + defaults.update(overrides) + return AggregateResult(**defaults) + + +def _make_dim(**overrides) -> DimensionScore: + defaults = dict( + dimension="progress", + weight=1.0, + mean_score=0.8, + probe_count=5, + raw_scores=[2, 3, 2, 3, 3], + ) + defaults.update(overrides) + return DimensionScore(**defaults) + + +# --------------------------------------------------------------------------- +# export_json +# --------------------------------------------------------------------------- + +class TestExportJson: + def test_creates_file(self, tmp_path): + path = tmp_path / "results.json" + export_json([], path) + assert path.exists() + + def test_empty_results_writes_empty_array(self, tmp_path): + path = tmp_path / "results.json" + export_json([], path) + data = json.loads(path.read_text()) + assert data == [] + + def test_single_result_structure(self, tmp_path): + result = _make_result() + path = tmp_path / "results.json" + export_json([result], path) + data = json.loads(path.read_text()) + assert len(data) == 1 + entry = data[0] + assert entry["method"] == "dedup" + assert entry["budget"] == 80_000 + assert entry["model_key"] == "cheap" + assert entry["composite"] == pytest.approx(0.75) + assert entry["ndcg"] == pytest.approx(0.82) + + def test_keeps_speed_and_token_counts(self, tmp_path): + result = _make_result(speed_s=2.5, kept_tokens=15_000, total_tokens=80_000) + path = tmp_path / "results.json" + export_json([result], path) + data = json.loads(path.read_text()) + assert data[0]["speed_s"] == pytest.approx(2.5) + assert data[0]["kept_tokens"] == 15_000 + assert data[0]["total_tokens"] == 80_000 + + def test_dimension_scores_included(self, tmp_path): + dim = _make_dim(dimension="error_solution", mean_score=0.9, probe_count=3) + result = _make_result(dimensions=[dim]) + path = tmp_path / "results.json" + export_json([result], path) + data = json.loads(path.read_text()) + dims = data[0]["dimensions"] + assert "error_solution" in dims + assert dims["error_solution"]["score"] == pytest.approx(0.9) + assert dims["error_solution"]["probe_count"] == 3 + + def test_multiple_results(self, tmp_path): + results = [_make_result(method="dedup"), _make_result(method="eitf")] + path = tmp_path / "results.json" + export_json(results, path) + data = json.loads(path.read_text()) + assert len(data) == 2 + methods = {e["method"] for e in data} + assert methods == {"dedup", "eitf"} + + def test_output_is_valid_json(self, tmp_path): + path = tmp_path / "results.json" + export_json([_make_result()], path) + # Should parse without exception + json.loads(path.read_text()) + + +# --------------------------------------------------------------------------- +# export_trace +# --------------------------------------------------------------------------- + +class TestExportTrace: + def _make_probe_set(self, probes=None): + ps = MagicMock() + ps.probes = probes or [] + return ps + + def _make_answer(self, probe_id="p1", score=2, **overrides): + a = MagicMock() + a.probe_id = probe_id + a.score = score + a.model_key = overrides.get("model_key", "cheap") + a.model_label = overrides.get("model_label", "haiku") + a.answer = overrides.get("answer", "some answer") + a.judge_reasoning = overrides.get("judge_reasoning", "looks good") + return a + + def _make_probe(self, id="p1", dimension="progress", tier="factual", + difficulty="medium", question="Q?", gold_answer="A", + evidence_turns=None): + p = MagicMock() + p.id = id + p.dimension = dimension + p.tier = tier + p.difficulty = difficulty + p.question = question + p.gold_answer = gold_answer + p.evidence_turns = evidence_turns or [] + return p + + def test_creates_trace_file(self, tmp_path): + probe_set = self._make_probe_set() + result_path = export_trace("dedup", 80_000, probe_set, [], tmp_path) + assert result_path.exists() + + def test_trace_file_in_trace_dir(self, tmp_path): + probe_set = self._make_probe_set() + result_path = export_trace("eitf", 40_000, probe_set, [], tmp_path) + assert result_path.parent == tmp_path + + def test_trace_filename_contains_method_and_budget(self, tmp_path): + probe_set = self._make_probe_set() + result_path = export_trace("setcover", 60_000, probe_set, [], tmp_path) + assert "setcover" in result_path.name + assert "60000" in result_path.name + + def test_trace_has_method_and_budget(self, tmp_path): + probe_set = self._make_probe_set() + result_path = export_trace("dedup", 80_000, probe_set, [], tmp_path) + data = json.loads(result_path.read_text()) + assert data["method"] == "dedup" + assert data["budget"] == 80_000 + + def test_empty_answers_produces_empty_entries(self, tmp_path): + probe_set = self._make_probe_set() + result_path = export_trace("dedup", 80_000, probe_set, [], tmp_path) + data = json.loads(result_path.read_text()) + assert data["entries"] == [] + + def test_answer_with_matching_probe_included(self, tmp_path): + probe = self._make_probe(id="p1", question="What happened?", gold_answer="Error") + probe_set = self._make_probe_set(probes=[probe]) + answer = self._make_answer(probe_id="p1", score=3, answer="Error occurred") + result_path = export_trace("dedup", 80_000, probe_set, [answer], tmp_path) + data = json.loads(result_path.read_text()) + assert len(data["entries"]) == 1 + entry = data["entries"][0] + assert entry["probe_id"] == "p1" + assert entry["score"] == 3 + + def test_answer_without_matching_probe_skipped(self, tmp_path): + probe_set = self._make_probe_set(probes=[]) + answer = self._make_answer(probe_id="missing_probe") + result_path = export_trace("dedup", 80_000, probe_set, [answer], tmp_path) + data = json.loads(result_path.read_text()) + assert data["entries"] == [] + + def test_creates_trace_dir_if_missing(self, tmp_path): + nested = tmp_path / "a" / "b" / "traces" + probe_set = self._make_probe_set() + export_trace("dedup", 80_000, probe_set, [], nested) + assert nested.is_dir() diff --git a/tests/test_formatter.py b/tests/test_formatter.py new file mode 100644 index 0000000..6c48a05 --- /dev/null +++ b/tests/test_formatter.py @@ -0,0 +1,219 @@ +"""Tests for lib/formatter.py — write_summary_text, write_compacted_jsonl, write_scores_csv.""" + +from __future__ import annotations + +import csv +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.formatter import write_compacted_jsonl, write_scores_csv, write_summary_text +from lib.parser import Turn +from lib.selector import SelectionResult +from lib.types import ScoredTurn + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(kind: str, index: int, text: str = "content") -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +def _selection(*turns: Turn) -> SelectionResult: + sr = SelectionResult() + sr.kept_turns = list(turns) + return sr + + +def _scored(turn: Turn, score: float, tokens: int) -> ScoredTurn: + return ScoredTurn(turn=turn, score=score, tokens=tokens) + + +# --------------------------------------------------------------------------- +# write_summary_text +# --------------------------------------------------------------------------- + +class TestWriteSummaryText: + def test_creates_file(self, tmp_path): + result = _selection(_turn("user", 0, "hello")) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + assert path.exists() + + def test_empty_selection_creates_empty_file(self, tmp_path): + result = _selection() + path = tmp_path / "summary.txt" + write_summary_text(result, path) + assert path.read_text() == "" + + def test_user_turn_labeled(self, tmp_path): + result = _selection(_turn("user", 0, "what is the fix?")) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + text = path.read_text() + assert "User" in text + assert "what is the fix?" in text + + def test_system_turn_labeled_assistant(self, tmp_path): + result = _selection(_turn("system", 1, "here is the fix")) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + text = path.read_text() + assert "Assistant" in text + assert "here is the fix" in text + + def test_turn_index_in_output(self, tmp_path): + result = _selection(_turn("user", 7, "msg")) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + assert "7" in path.read_text() + + def test_multiple_turns_separated(self, tmp_path): + result = _selection( + _turn("user", 0, "question"), + _turn("system", 1, "answer"), + ) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + text = path.read_text() + assert "question" in text + assert "answer" in text + assert "---" in text + + def test_long_turn_truncated(self, tmp_path): + long_text = "x" * 5000 + result = _selection(_turn("user", 0, long_text)) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + text = path.read_text() + assert "truncated" in text + + def test_blank_turn_skipped(self, tmp_path): + t = Turn(kind="user", index=0) # no content lines + result = _selection(t) + path = tmp_path / "summary.txt" + write_summary_text(result, path) + # Empty turn should be skipped → empty file + assert path.read_text() == "" + + +# --------------------------------------------------------------------------- +# write_compacted_jsonl +# --------------------------------------------------------------------------- + +class TestWriteCompactedJsonl: + def test_creates_file(self, tmp_path): + result = _selection(_turn("user", 0)) + path = tmp_path / "out.jsonl" + write_compacted_jsonl(result, path) + assert path.exists() + + def test_empty_selection_creates_empty_file(self, tmp_path): + result = _selection() + path = tmp_path / "out.jsonl" + write_compacted_jsonl(result, path) + assert path.read_text().strip() == "" + + def test_each_line_is_valid_json(self, tmp_path): + result = _selection(_turn("user", 0, "hello"), _turn("system", 1, "world")) + path = tmp_path / "out.jsonl" + write_compacted_jsonl(result, path) + for line in path.read_text().splitlines(): + if line.strip(): + json.loads(line) # should not raise + + def test_records_written_in_turn_order(self, tmp_path): + t0 = _turn("user", 0, "first") + t1 = _turn("system", 1, "second") + result = _selection(t0, t1) + path = tmp_path / "out.jsonl" + write_compacted_jsonl(result, path) + lines = [json.loads(l) for l in path.read_text().splitlines() if l.strip()] + contents = [l.get("message", {}).get("content", "") for l in lines] + assert contents[0] == "first" + assert contents[1] == "second" + + def test_multiple_records_per_turn(self, tmp_path): + t = Turn(kind="user", index=0) + t.append({"message": {"content": "rec1"}}) + t.append({"message": {"content": "rec2"}}) + result = _selection(t) + path = tmp_path / "out.jsonl" + write_compacted_jsonl(result, path) + lines = [l for l in path.read_text().splitlines() if l.strip()] + assert len(lines) == 2 + + +# --------------------------------------------------------------------------- +# write_scores_csv +# --------------------------------------------------------------------------- + +class TestWriteScoresCsv: + def test_creates_file(self, tmp_path): + t = _turn("system", 0) + scored = [_scored(t, 0.75, 200)] + path = tmp_path / "scores.csv" + write_scores_csv(scored, {0}, path) + assert path.exists() + + def test_header_row(self, tmp_path): + path = tmp_path / "scores.csv" + write_scores_csv([], set(), path) + with open(path) as f: + reader = csv.reader(f) + header = next(reader) + assert "turn_index" in header + assert "score" in header + assert "tokens" in header + assert "kept" in header + + def test_score_written(self, tmp_path): + t = _turn("system", 5, "content") + scored = [_scored(t, 0.9876, 300)] + path = tmp_path / "scores.csv" + write_scores_csv(scored, {5}, path) + with open(path) as f: + reader = csv.DictReader(f) + rows = list(reader) + assert len(rows) == 1 + assert float(rows[0]["score"]) == pytest.approx(0.9876, abs=0.001) + + def test_kept_flag_true_when_in_kept_indices(self, tmp_path): + t = _turn("system", 3) + scored = [_scored(t, 0.5, 100)] + path = tmp_path / "scores.csv" + write_scores_csv(scored, {3}, path) + with open(path) as f: + reader = csv.DictReader(f) + rows = list(reader) + assert rows[0]["kept"] == "True" + + def test_kept_flag_false_when_not_in_kept_indices(self, tmp_path): + t = _turn("system", 4) + scored = [_scored(t, 0.3, 50)] + path = tmp_path / "scores.csv" + write_scores_csv(scored, set(), path) + with open(path) as f: + reader = csv.DictReader(f) + rows = list(reader) + assert rows[0]["kept"] == "False" + + def test_sorted_by_turn_index(self, tmp_path): + t5 = _turn("system", 5) + t2 = _turn("system", 2) + scored = [_scored(t5, 0.9, 100), _scored(t2, 0.1, 200)] + path = tmp_path / "scores.csv" + write_scores_csv(scored, set(), path) + with open(path) as f: + reader = csv.DictReader(f) + rows = list(reader) + indices = [int(r["turn_index"]) for r in rows] + assert indices == sorted(indices) diff --git a/tests/test_llama_embed.py b/tests/test_llama_embed.py new file mode 100644 index 0000000..ed52587 --- /dev/null +++ b/tests/test_llama_embed.py @@ -0,0 +1,235 @@ +"""Tests for lib/llama_embed.py — constants, _instruct, LlamaEmbedScorer (mocked).""" + +from __future__ import annotations + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.llama_embed import ( + DOC_INSTRUCTION, + MAX_DOC_CHARS, + QUERY_INSTRUCTION, + LlamaEmbedScorer, + _instruct, +) +from lib.parser import Turn +from lib.types import ScoredTurn + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(kind: str, index: int, text: str) -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +def _mock_health(): + """Return a mock for httpx.get that simulates a healthy server.""" + resp = MagicMock() + resp.raise_for_status.return_value = None + return resp + + +def _mock_embed_response(embeddings: list[list[float]]): + """Return a mock httpx.post response for /v1/embeddings.""" + resp = MagicMock() + resp.raise_for_status.return_value = None + resp.json.return_value = { + "data": [ + {"index": i, "embedding": emb} + for i, emb in enumerate(embeddings) + ] + } + return resp + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +class TestConstants: + def test_query_instruction_is_str(self): + assert isinstance(QUERY_INSTRUCTION, str) + + def test_query_instruction_nonempty(self): + assert len(QUERY_INSTRUCTION) > 0 + + def test_doc_instruction_is_str(self): + assert isinstance(DOC_INSTRUCTION, str) + + def test_doc_instruction_nonempty(self): + assert len(DOC_INSTRUCTION) > 0 + + def test_max_doc_chars_is_int(self): + assert isinstance(MAX_DOC_CHARS, int) + + def test_max_doc_chars_positive(self): + assert MAX_DOC_CHARS > 0 + + +# --------------------------------------------------------------------------- +# _instruct +# --------------------------------------------------------------------------- + +class TestInstruct: + def test_basic_format(self): + result = _instruct("do task", "hello") + assert result == "Instruct: do task\nQuery: hello" + + def test_instruction_in_result(self): + result = _instruct("my instruction", "query text") + assert "Instruct: my instruction" in result + + def test_query_in_result(self): + result = _instruct("inst", "my query") + assert "Query: my query" in result + + def test_newline_separator(self): + result = _instruct("a", "b") + lines = result.split("\n") + assert len(lines) == 2 + + def test_empty_inputs(self): + result = _instruct("", "") + assert result == "Instruct: \nQuery: " + + def test_returns_string(self): + assert isinstance(_instruct("x", "y"), str) + + +# --------------------------------------------------------------------------- +# LlamaEmbedScorer — init with mocked health check +# --------------------------------------------------------------------------- + +class TestLlamaEmbedScorerInit: + def test_url_constructed_from_base(self): + with patch("lib.llama_embed.httpx.get", return_value=_mock_health()): + scorer = LlamaEmbedScorer("http://localhost:8080") + assert scorer.url == "http://localhost:8080/v1/embeddings" + + def test_trailing_slash_stripped(self): + with patch("lib.llama_embed.httpx.get", return_value=_mock_health()): + scorer = LlamaEmbedScorer("http://localhost:8080/") + assert scorer.url == "http://localhost:8080/v1/embeddings" + + def test_health_check_called(self): + with patch("lib.llama_embed.httpx.get", return_value=_mock_health()) as mock_get: + LlamaEmbedScorer("http://localhost:8080") + mock_get.assert_called_once() + call_url = mock_get.call_args[0][0] + assert "health" in call_url + + def test_health_failure_propagates(self): + bad_resp = MagicMock() + bad_resp.raise_for_status.side_effect = Exception("Connection refused") + with patch("lib.llama_embed.httpx.get", return_value=bad_resp): + with pytest.raises(Exception, match="Connection refused"): + LlamaEmbedScorer("http://localhost:8080") + + +# --------------------------------------------------------------------------- +# LlamaEmbedScorer.score_turns — mocked _embed +# --------------------------------------------------------------------------- + +class TestLlamaEmbedScorerScoreTurns: + def _make_scorer(self) -> LlamaEmbedScorer: + with patch("lib.llama_embed.httpx.get", return_value=_mock_health()): + return LlamaEmbedScorer("http://localhost:8080") + + def test_single_turn_returns_one_result(self): + scorer = self._make_scorer() + t = _turn("system", 0, "single content") + token_counts = {0: 100} + query_emb = np.array([[1.0, 0.0, 0.0, 0.0]]) + doc_embs = np.array([[1.0, 0.0, 0.0, 0.0]], dtype=np.float32) + + call_count = [0] + def fake_embed(texts): + idx = call_count[0] + call_count[0] += 1 + return query_emb if idx == 0 else doc_embs + + with patch.object(scorer, "_embed", side_effect=fake_embed): + result = scorer.score_turns([t], "query", token_counts) + assert len(result) == 1 + + def test_returns_one_scored_turn_per_system_turn(self): + scorer = self._make_scorer() + dim = 4 + turns = [_turn("system", i, f"content{i}") for i in range(3)] + token_counts = {i: 50 for i in range(3)} + # Query emb + 3 doc embs + query_emb = np.array([[1.0, 0.0, 0.0, 0.0]]) + doc_embs = np.eye(3, 4, dtype=np.float32) + + call_count = [0] + def fake_embed(texts): + idx = call_count[0] + call_count[0] += 1 + if idx == 0: + return query_emb + return doc_embs + + with patch.object(scorer, "_embed", side_effect=fake_embed): + result = scorer.score_turns(turns, "query", token_counts) + assert len(result) == 3 + + def test_scores_are_in_zero_one_range(self): + scorer = self._make_scorer() + turns = [_turn("system", 0, "hello"), _turn("system", 1, "world")] + token_counts = {0: 30, 1: 40} + dim = 4 + query_emb = np.array([[1.0, 0.0, 0.0, 0.0]]) + doc_embs = np.array([[0.5, 0.5, 0.5, 0.5], [-0.5, 0.5, 0.5, 0.5]], dtype=np.float32) + + call_count = [0] + def fake_embed(texts): + idx = call_count[0] + call_count[0] += 1 + return query_emb if idx == 0 else doc_embs + + with patch.object(scorer, "_embed", side_effect=fake_embed): + result = scorer.score_turns(turns, "query", token_counts) + for st in result: + assert 0.0 <= st.score <= 1.0 + + def test_tokens_set_from_token_counts(self): + scorer = self._make_scorer() + t = _turn("system", 7, "text") + query_emb = np.array([[1.0, 0.0]]) + doc_embs = np.array([[1.0, 0.0]], dtype=np.float32) + + call_count = [0] + def fake_embed(texts): + idx = call_count[0] + call_count[0] += 1 + return query_emb if idx == 0 else doc_embs + + with patch.object(scorer, "_embed", side_effect=fake_embed): + result = scorer.score_turns([t], "query", {7: 999}) + assert result[0].tokens == 999 + + def test_returns_scored_turn_instances(self): + scorer = self._make_scorer() + t = _turn("system", 0, "content") + query_emb = np.array([[1.0, 0.0]]) + doc_embs = np.array([[0.0, 1.0]], dtype=np.float32) + + call_count = [0] + def fake_embed(texts): + idx = call_count[0] + call_count[0] += 1 + return query_emb if idx == 0 else doc_embs + + with patch.object(scorer, "_embed", side_effect=fake_embed): + result = scorer.score_turns([t], "query", {0: 50}) + assert all(isinstance(st, ScoredTurn) for st in result) diff --git a/tests/test_llm_compact.py b/tests/test_llm_compact.py new file mode 100644 index 0000000..870fe30 --- /dev/null +++ b/tests/test_llm_compact.py @@ -0,0 +1,129 @@ +"""Tests for lib/llm_compact.py — make_synthetic_turn and llm_compact guards.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.llm_compact import ( + COMPACT_SYSTEM_PROMPT, + MODEL, + make_synthetic_turn, +) +from lib.parser import Turn, extract_text + + +# --------------------------------------------------------------------------- +# make_synthetic_turn +# --------------------------------------------------------------------------- + +class TestMakeSyntheticTurn: + def test_returns_turn_instance(self): + t = make_synthetic_turn("hello") + assert isinstance(t, Turn) + + def test_kind_is_system(self): + t = make_synthetic_turn("some summary") + assert t.kind == "system" + + def test_default_index_zero(self): + t = make_synthetic_turn("summary text") + assert t.index == 0 + + def test_custom_index_stored(self): + t = make_synthetic_turn("summary", index=42) + assert t.index == 42 + + def test_has_one_line(self): + t = make_synthetic_turn("content") + assert len(t.lines) == 1 + + def test_line_type_is_assistant(self): + t = make_synthetic_turn("content") + assert t.lines[0]["type"] == "assistant" + + def test_message_role_is_assistant(self): + t = make_synthetic_turn("content") + assert t.lines[0]["message"]["role"] == "assistant" + + def test_summary_text_in_message_content(self): + text = "Error: FileNotFoundError at /home/user/project/main.py" + t = make_synthetic_turn(text) + assert t.lines[0]["message"]["content"] == text + + def test_extract_text_returns_summary(self): + text = "ValueError in /src/module.py on line 42" + t = make_synthetic_turn(text) + assert extract_text(t) == text + + def test_empty_summary_produces_valid_turn(self): + t = make_synthetic_turn("") + assert t.kind == "system" + assert extract_text(t) == "" + + def test_multiline_summary_preserved(self): + text = "line1\nline2\nline3" + t = make_synthetic_turn(text) + assert extract_text(t) == text + + def test_summary_with_special_chars(self): + text = "Token: abc_123-456 / path: /usr/local/lib.so" + t = make_synthetic_turn(text) + assert extract_text(t) == text + + def test_different_indices_independent(self): + t1 = make_synthetic_turn("a", index=0) + t2 = make_synthetic_turn("b", index=99) + assert t1.index == 0 + assert t2.index == 99 + assert t1.lines[0]["message"]["content"] == "a" + assert t2.lines[0]["message"]["content"] == "b" + + +# --------------------------------------------------------------------------- +# llm_compact — missing API key guard +# --------------------------------------------------------------------------- + +class TestLlmCompactApiKeyGuard: + def test_missing_api_key_raises_runtime_error(self, monkeypatch): + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + from lib.llm_compact import llm_compact + from lib.parser import Turn as T + + with pytest.raises(RuntimeError, match="OPENROUTER_API_KEY"): + llm_compact([], budget=1000) + + def test_missing_api_key_message_contains_var_name(self, monkeypatch): + monkeypatch.delenv("OPENROUTER_API_KEY", raising=False) + from lib.llm_compact import llm_compact + + with pytest.raises(RuntimeError) as exc_info: + llm_compact([], budget=500) + assert "OPENROUTER_API_KEY" in str(exc_info.value) + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +class TestConstants: + def test_compact_system_prompt_is_str(self): + assert isinstance(COMPACT_SYSTEM_PROMPT, str) + + def test_compact_system_prompt_mentions_file_paths(self): + assert "file path" in COMPACT_SYSTEM_PROMPT.lower() or "path" in COMPACT_SYSTEM_PROMPT + + def test_compact_system_prompt_nonempty(self): + assert len(COMPACT_SYSTEM_PROMPT) > 0 + + def test_model_is_str(self): + assert isinstance(MODEL, str) + + def test_model_nonempty(self): + assert len(MODEL) > 0 diff --git a/tests/test_pareto.py b/tests/test_pareto.py new file mode 100644 index 0000000..138ef2a --- /dev/null +++ b/tests/test_pareto.py @@ -0,0 +1,204 @@ +"""Tests for lib/pareto.py — METHOD_STYLES, plot_entity_coverage, plot_type_breakdown.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") # non-interactive backend for tests +import matplotlib.pyplot as plt +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.pareto import METHOD_STYLES, plot_entity_coverage, plot_type_breakdown + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_result(method: str, speed: float, coverage: float, budget: int = 10000) -> dict: + return { + "method": method, + "speed_s": speed, + "weighted_coverage": coverage, + "budget": budget, + "kept_tokens": int(budget * 0.8), + } + + +def _make_ax(): + fig, ax = plt.subplots() + return ax + + +# --------------------------------------------------------------------------- +# METHOD_STYLES +# --------------------------------------------------------------------------- + +class TestMethodStyles: + def test_expected_methods_present(self): + for method in ("dedup", "eitf", "setcover", "llama-embed", "llama-rerank", "claude-code"): + assert method in METHOD_STYLES, f"{method!r} missing from METHOD_STYLES" + + def test_each_entry_has_color(self): + for method, style in METHOD_STYLES.items(): + assert "color" in style, f"{method!r} missing color" + + def test_each_entry_has_marker(self): + for method, style in METHOD_STYLES.items(): + assert "marker" in style, f"{method!r} missing marker" + + def test_each_entry_has_label(self): + for method, style in METHOD_STYLES.items(): + assert "label" in style, f"{method!r} missing label" + + def test_color_is_hex_or_named_string(self): + for method, style in METHOD_STYLES.items(): + color = style["color"] + assert isinstance(color, str) and len(color) > 0 + + def test_labels_are_non_empty_strings(self): + for method, style in METHOD_STYLES.items(): + assert isinstance(style["label"], str) and len(style["label"]) > 0 + + +# --------------------------------------------------------------------------- +# plot_entity_coverage +# --------------------------------------------------------------------------- + +class TestPlotEntityCoverage: + def teardown_method(self): + plt.close("all") + + def test_empty_results_does_not_crash(self): + ax = _make_ax() + plot_entity_coverage(ax, []) + + def test_single_result_does_not_crash(self): + ax = _make_ax() + plot_entity_coverage(ax, [_make_result("dedup", 1.0, 0.8)]) + + def test_multiple_results_does_not_crash(self): + ax = _make_ax() + results = [ + _make_result("dedup", 1.0, 0.7), + _make_result("eitf", 2.0, 0.85), + _make_result("setcover", 3.0, 0.9), + ] + plot_entity_coverage(ax, results) + + def test_same_method_multiple_points(self): + ax = _make_ax() + results = [ + _make_result("dedup", 1.0, 0.7, budget=10000), + _make_result("dedup", 2.0, 0.8, budget=20000), + ] + plot_entity_coverage(ax, results) + + def test_unknown_method_uses_fallback(self): + ax = _make_ax() + results = [_make_result("unknown_method", 1.0, 0.5)] + plot_entity_coverage(ax, results) # should not crash + + def test_xlabel_set(self): + ax = _make_ax() + plot_entity_coverage(ax, [_make_result("eitf", 1.0, 0.8)]) + assert ax.get_xlabel() != "" + + def test_ylabel_set(self): + ax = _make_ax() + plot_entity_coverage(ax, [_make_result("eitf", 1.0, 0.8)]) + assert ax.get_ylabel() != "" + + def test_title_set(self): + ax = _make_ax() + plot_entity_coverage(ax, [_make_result("eitf", 1.0, 0.8)]) + assert ax.get_title() != "" + + def test_show_legend_false_no_crash(self): + ax = _make_ax() + plot_entity_coverage(ax, [_make_result("dedup", 1.0, 0.8)], show_legend=False) + + def test_large_budget_uses_K_label(self): + # Budget ≥ 1000 → budget_label contains 'K' + ax = _make_ax() + plot_entity_coverage(ax, [_make_result("eitf", 1.0, 0.8, budget=50000)]) + # If it doesn't crash and the annotation is applied — pass + + def test_claude_code_uses_larger_scatter_size(self): + # claude-code uses size=300, others 180 — just check no crash + ax = _make_ax() + results = [ + _make_result("claude-code", 1.0, 0.95), + _make_result("dedup", 2.0, 0.7), + ] + plot_entity_coverage(ax, results) + + +# --------------------------------------------------------------------------- +# plot_type_breakdown +# --------------------------------------------------------------------------- + +class TestPlotTypeBreakdown: + def teardown_method(self): + plt.close("all") + + def test_empty_type_coverage_does_not_crash(self): + ax = _make_ax() + results = [{"method": "dedup", "budget": 10000, "type_coverage": {}}] + plot_type_breakdown(ax, results) + + def test_missing_type_coverage_key_does_not_crash(self): + ax = _make_ax() + results = [{"method": "dedup", "budget": 10000}] + plot_type_breakdown(ax, results) + + def test_single_type_does_not_crash(self): + ax = _make_ax() + results = [{ + "method": "eitf", + "budget": 10000, + "type_coverage": { + "exception": {"coverage": 0.9, "weight": 1.0, "covered": 9, "total": 10}, + }, + }] + plot_type_breakdown(ax, results) + + def test_multiple_types_sorted(self): + ax = _make_ax() + results = [{ + "method": "eitf", + "budget": 10000, + "type_coverage": { + "url": {"coverage": 0.5, "weight": 0.5, "covered": 5, "total": 10}, + "exception": {"coverage": 0.9, "weight": 1.0, "covered": 9, "total": 10}, + "port": {"coverage": 0.7, "weight": 0.8, "covered": 7, "total": 10}, + }, + }] + plot_type_breakdown(ax, results) + + def test_largest_budget_selected(self): + ax = _make_ax() + results = [ + {"method": "dedup", "budget": 5000, "type_coverage": {}}, + {"method": "dedup", "budget": 20000, "type_coverage": { + "exception": {"coverage": 0.8, "weight": 1.0, "covered": 8, "total": 10}, + }}, + ] + # Should use the budget=20000 entry for the bar chart — no crash + plot_type_breakdown(ax, results) + + def test_title_contains_budget(self): + ax = _make_ax() + results = [{ + "method": "eitf", + "budget": 30000, + "type_coverage": { + "exception": {"coverage": 0.8, "weight": 1.0, "covered": 8, "total": 10}, + }, + }] + plot_type_breakdown(ax, results) + assert "30" in ax.get_title() # budget 30K in title diff --git a/tests/test_parser.py b/tests/test_parser.py new file mode 100644 index 0000000..153ca84 --- /dev/null +++ b/tests/test_parser.py @@ -0,0 +1,224 @@ +"""Tests for lib/parser.py — Turn, extract_text, _is_user_message.""" + +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.parser import Turn, _is_user_message, extract_text, parse_jsonl + + +# --------------------------------------------------------------------------- +# Turn dataclass +# --------------------------------------------------------------------------- + +class TestTurn: + def test_default_fields(self): + t = Turn(kind="user") + assert t.kind == "user" + assert t.lines == [] + assert t.index == 0 + + def test_custom_fields(self): + t = Turn(kind="system", lines=[{"foo": "bar"}], index=3) + assert t.kind == "system" + assert t.lines == [{"foo": "bar"}] + assert t.index == 3 + + def test_append_adds_line(self): + t = Turn(kind="user") + t.append({"type": "user", "message": {"content": "hello"}}) + assert len(t.lines) == 1 + assert t.lines[0]["type"] == "user" + + def test_append_multiple(self): + t = Turn(kind="system") + t.append({"a": 1}) + t.append({"b": 2}) + assert len(t.lines) == 2 + + +# --------------------------------------------------------------------------- +# _is_user_message +# --------------------------------------------------------------------------- + +class TestIsUserMessage: + def test_string_content_is_user(self): + record = {"type": "user", "message": {"content": "hello"}} + assert _is_user_message(record) is True + + def test_text_block_is_user(self): + record = { + "type": "user", + "message": {"content": [{"type": "text", "text": "hi"}]}, + } + assert _is_user_message(record) is True + + def test_tool_result_block_is_not_user(self): + record = { + "type": "user", + "message": {"content": [{"type": "tool_result", "content": "ok"}]}, + } + assert _is_user_message(record) is False + + def test_non_user_type_is_false(self): + record = {"type": "assistant", "message": {"content": "reply"}} + assert _is_user_message(record) is False + + def test_tool_injection_with_source_uuid_is_false(self): + record = { + "type": "user", + "sourceToolAssistantUUID": "abc-123", + "message": {"content": "injected"}, + } + assert _is_user_message(record) is False + + def test_empty_content_list(self): + record = {"type": "user", "message": {"content": []}} + # Empty list: no tool_result blocks → considered user + assert _is_user_message(record) is True + + +# --------------------------------------------------------------------------- +# extract_text +# --------------------------------------------------------------------------- + +class TestExtractTextStringContent: + def test_simple_string_content(self): + t = Turn(kind="user") + t.append({"message": {"content": "Hello world"}}) + assert extract_text(t) == "Hello world" + + def test_multiple_records_joined(self): + t = Turn(kind="user") + t.append({"message": {"content": "line one"}}) + t.append({"message": {"content": "line two"}}) + result = extract_text(t) + assert "line one" in result + assert "line two" in result + + def test_empty_turn_returns_empty(self): + t = Turn(kind="user") + assert extract_text(t) == "" + + +class TestExtractTextListContent: + def test_text_block(self): + t = Turn(kind="system") + t.append({"message": {"content": [{"type": "text", "text": "hello"}]}}) + assert extract_text(t) == "hello" + + def test_thinking_block(self): + t = Turn(kind="system") + t.append({"message": {"content": [ + {"type": "thinking", "thinking": "let me think"} + ]}}) + assert "let me think" in extract_text(t) + + def test_tool_use_block_includes_name(self): + t = Turn(kind="system") + t.append({"message": {"content": [ + {"type": "tool_use", "name": "Read", "input": {"file_path": "/tmp/foo.py"}} + ]}}) + result = extract_text(t) + assert "Read" in result + assert "file_path" in result + assert "/tmp/foo.py" in result + + def test_tool_use_input_truncated_at_500(self): + long_input = "x" * 600 + t = Turn(kind="system") + t.append({"message": {"content": [ + {"type": "tool_use", "name": "Write", "input": {"content": long_input}} + ]}}) + result = extract_text(t) + assert "..." in result + assert len(result) < 700 # shouldn't be excessively long + + def test_tool_result_string_content(self): + t = Turn(kind="user") + t.append({"message": {"content": [ + {"type": "tool_result", "content": "file contents here"} + ]}}) + assert "file contents here" in extract_text(t) + + def test_tool_result_list_content(self): + t = Turn(kind="user") + t.append({"message": {"content": [ + { + "type": "tool_result", + "content": [{"type": "text", "text": "nested text"}], + } + ]}}) + assert "nested text" in extract_text(t) + + def test_non_dict_block_skipped(self): + t = Turn(kind="system") + t.append({"message": {"content": ["a string block"]}}) + # Should not crash, returns empty or whatever + result = extract_text(t) + assert isinstance(result, str) + + def test_multiple_blocks_concatenated(self): + t = Turn(kind="system") + t.append({"message": {"content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"}, + ]}}) + result = extract_text(t) + assert "first" in result + assert "second" in result + + +# --------------------------------------------------------------------------- +# parse_jsonl +# --------------------------------------------------------------------------- + +class TestParseJsonl: + def _write_jsonl(self, tmp_path: Path, records: list[dict]) -> Path: + p = tmp_path / "conv.jsonl" + with open(p, "w") as f: + for r in records: + f.write(json.dumps(r) + "\n") + return p + + def test_empty_file_returns_empty(self, tmp_path): + p = self._write_jsonl(tmp_path, []) + turns = parse_jsonl(p) + assert turns == [] + + def test_single_user_message(self, tmp_path): + records = [ + {"type": "user", "message": {"content": "hello"}}, + ] + p = self._write_jsonl(tmp_path, records) + turns = parse_jsonl(p) + assert any(t.kind == "user" for t in turns) + + def test_user_then_assistant(self, tmp_path): + records = [ + {"type": "user", "message": {"content": "question"}}, + {"type": "assistant", "message": {"content": [{"type": "text", "text": "answer"}]}}, + ] + p = self._write_jsonl(tmp_path, records) + turns = parse_jsonl(p) + kinds = [t.kind for t in turns] + assert "user" in kinds + assert "system" in kinds + + def test_turns_indexed_sequentially(self, tmp_path): + records = [ + {"type": "user", "message": {"content": "msg1"}}, + {"type": "assistant", "message": {"content": [{"type": "text", "text": "resp1"}]}}, + {"type": "user", "message": {"content": "msg2"}}, + ] + p = self._write_jsonl(tmp_path, records) + turns = parse_jsonl(p) + indices = [t.index for t in turns] + assert indices == sorted(indices) + assert len(set(indices)) == len(indices) # all unique diff --git a/tests/test_scorer_pure_fns.py b/tests/test_scorer_pure_fns.py new file mode 100644 index 0000000..11932b1 --- /dev/null +++ b/tests/test_scorer_pure_fns.py @@ -0,0 +1,153 @@ +"""Tests for lib/scorer.py — pure helper functions (no model loading). + +Covers _format_instruct and _last_token_pool. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.scorer import _format_instruct, _last_token_pool + + +# --------------------------------------------------------------------------- +# _format_instruct +# --------------------------------------------------------------------------- + +class TestFormatInstruct: + def test_basic_format(self): + result = _format_instruct("do task", "hello world") + assert result == "Instruct: do task\nQuery: hello world" + + def test_instruction_included(self): + result = _format_instruct("find errors", "some query text") + assert "Instruct: find errors" in result + + def test_query_included(self): + result = _format_instruct("find errors", "some query text") + assert "Query: some query text" in result + + def test_newline_separates_parts(self): + result = _format_instruct("inst", "qry") + parts = result.split("\n") + assert len(parts) == 2 + assert parts[0].startswith("Instruct:") + assert parts[1].startswith("Query:") + + def test_empty_instruction(self): + result = _format_instruct("", "query") + assert result == "Instruct: \nQuery: query" + + def test_empty_query(self): + result = _format_instruct("instruction", "") + assert result == "Instruct: instruction\nQuery: " + + def test_multiline_query_preserved(self): + query = "line1\nline2" + result = _format_instruct("inst", query) + # The query section starts after "Query: " + assert result.endswith(f"\nQuery: {query}") + + def test_returns_string(self): + result = _format_instruct("a", "b") + assert isinstance(result, str) + + +# --------------------------------------------------------------------------- +# _last_token_pool — right-padding (variable-length sequences) +# --------------------------------------------------------------------------- + +class TestLastTokenPoolRightPadding: + """When attention_mask has right padding, pool at actual last non-pad position.""" + + def _make_hidden(self, batch: int, seq: int, dim: int) -> torch.Tensor: + return torch.zeros(batch, seq, dim) + + def test_single_seq_full_mask_returns_last_token(self): + hidden = torch.zeros(1, 4, 3) + hidden[0, 3, :] = torch.tensor([1.0, 2.0, 3.0]) + mask = torch.ones(1, 4, dtype=torch.long) # no padding → right-pad variant + result = _last_token_pool(hidden, mask) + assert result.shape == (1, 3) + assert torch.allclose(result[0], torch.tensor([1.0, 2.0, 3.0])) + + def test_right_padded_returns_last_real_token(self): + hidden = torch.zeros(1, 5, 4) + hidden[0, 2, :] = torch.tensor([1.0, 2.0, 3.0, 4.0]) + # seq len = 3, padded to 5 + mask = torch.tensor([[1, 1, 1, 0, 0]]) + result = _last_token_pool(hidden, mask) + assert torch.allclose(result[0], torch.tensor([1.0, 2.0, 3.0, 4.0])) + + def test_batch_with_different_lengths(self): + hidden = torch.zeros(2, 4, 3) + # Seq 0: real length 2, last real token at index 1 + hidden[0, 1, :] = torch.tensor([10.0, 20.0, 30.0]) + # Seq 1: real length 3, last real token at index 2 + hidden[1, 2, :] = torch.tensor([40.0, 50.0, 60.0]) + mask = torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]]) + result = _last_token_pool(hidden, mask) + assert result.shape == (2, 3) + assert torch.allclose(result[0], torch.tensor([10.0, 20.0, 30.0])) + assert torch.allclose(result[1], torch.tensor([40.0, 50.0, 60.0])) + + def test_output_shape_matches_batch_and_dim(self): + hidden = torch.randn(3, 6, 8) + mask = torch.ones(3, 6, dtype=torch.long) + result = _last_token_pool(hidden, mask) + assert result.shape == (3, 8) + + def test_single_token_sequence(self): + hidden = torch.zeros(1, 1, 4) + hidden[0, 0, :] = torch.tensor([7.0, 8.0, 9.0, 10.0]) + mask = torch.tensor([[1]]) + result = _last_token_pool(hidden, mask) + assert torch.allclose(result[0], torch.tensor([7.0, 8.0, 9.0, 10.0])) + + +# --------------------------------------------------------------------------- +# _last_token_pool — left-padding (Qwen3 default) +# --------------------------------------------------------------------------- + +class TestLastTokenPoolLeftPadding: + """When attention_mask has left padding (all last columns = 1), pool last index.""" + + def _left_pad_mask(self, batch: int, seq: int, real_len: int) -> torch.Tensor: + """Create a left-padded mask where real tokens are at the end.""" + mask = torch.zeros(batch, seq, dtype=torch.long) + mask[:, seq - real_len :] = 1 + return mask + + def test_left_padded_returns_last_token(self): + hidden = torch.zeros(1, 5, 3) + hidden[0, 4, :] = torch.tensor([5.0, 6.0, 7.0]) + # Left-padded: real tokens at indices 2,3,4 + mask = torch.tensor([[0, 0, 1, 1, 1]]) + result = _last_token_pool(hidden, mask) + assert torch.allclose(result[0], torch.tensor([5.0, 6.0, 7.0])) + + def test_batch_left_padded_all_return_index_minus1(self): + batch, seq, dim = 3, 6, 4 + hidden = torch.randn(batch, seq, dim) + # All sequences fully populated (mask all-ones is NOT left-padding per this fn) + # Create left-padding: pad at start, real at end + mask = torch.zeros(batch, seq, dtype=torch.long) + mask[:, 3:] = 1 # last 3 tokens are real for all seqs + result = _last_token_pool(hidden, mask) + # Left-padding detected → returns last (index -1) hidden state + assert result.shape == (batch, dim) + assert torch.allclose(result, hidden[:, -1, :]) + + def test_left_pad_output_shape(self): + hidden = torch.randn(2, 5, 16) + # Left-padded: last column all 1s → left_padding=True + mask = torch.zeros(2, 5, dtype=torch.long) + mask[:, -1] = 1 + result = _last_token_pool(hidden, mask) + assert result.shape == (2, 16) diff --git a/tests/test_selector.py b/tests/test_selector.py new file mode 100644 index 0000000..ffdd8e5 --- /dev/null +++ b/tests/test_selector.py @@ -0,0 +1,194 @@ +"""Tests for lib/selector.py — SelectionResult and select_turns.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.parser import Turn +from lib.selector import SelectionResult, select_turns +from lib.types import ScoredTurn + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(kind: str, index: int, text: str = "x") -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +def _scored(turn: Turn, score: float, tokens: int) -> ScoredTurn: + return ScoredTurn(turn=turn, score=score, tokens=tokens) + + +# --------------------------------------------------------------------------- +# SelectionResult dataclass +# --------------------------------------------------------------------------- + +class TestSelectionResult: + def test_defaults(self): + sr = SelectionResult() + assert sr.kept_turns == [] + assert sr.dropped_turns == [] + assert sr.kept_scored == [] + assert sr.user_tokens == 0 + assert sr.short_system_tokens == 0 + assert sr.scored_kept_tokens == 0 + assert sr.scored_dropped_tokens == 0 + assert sr.total_input_tokens == 0 + assert sr.budget == 0 + + +# --------------------------------------------------------------------------- +# select_turns — empty inputs +# --------------------------------------------------------------------------- + +class TestSelectTurnsEmpty: + def test_no_turns_returns_empty(self): + result = select_turns([], [], {}, budget=1000) + assert result.kept_turns == [] + assert result.total_input_tokens == 0 + + def test_no_scored_all_user_kept(self): + turns = [_turn("user", 0), _turn("user", 1)] + token_counts = {0: 100, 1: 100} + result = select_turns(turns, [], token_counts, budget=500) + assert len(result.kept_turns) == 2 + assert result.user_tokens == 200 + + +# --------------------------------------------------------------------------- +# select_turns — user turns always kept +# --------------------------------------------------------------------------- + +class TestSelectTurnsUserAlwaysKept: + def test_user_turns_kept_even_over_budget(self): + turns = [_turn("user", 0), _turn("user", 1)] + token_counts = {0: 1000, 1: 1000} + result = select_turns(turns, [], token_counts, budget=10) + user_indices = [t.index for t in result.kept_turns if t.kind == "user"] + assert set(user_indices) == {0, 1} + + def test_user_token_count_accumulated(self): + turns = [_turn("user", 0), _turn("user", 1)] + token_counts = {0: 300, 1: 400} + result = select_turns(turns, [], token_counts, budget=10_000) + assert result.user_tokens == 700 + + +# --------------------------------------------------------------------------- +# select_turns — short system turns always kept +# --------------------------------------------------------------------------- + +class TestSelectTurnsShortSystem: + def test_short_system_kept(self): + turns = [_turn("system", 0)] + token_counts = {0: 100} # <= default threshold 300 + result = select_turns(turns, [], token_counts, budget=1000, short_threshold=300) + assert any(t.index == 0 for t in result.kept_turns) + + def test_short_system_token_count(self): + turns = [_turn("system", 0), _turn("system", 1)] + token_counts = {0: 150, 1: 200} + result = select_turns(turns, [], token_counts, budget=1000, short_threshold=300) + assert result.short_system_tokens == 350 + + def test_long_system_not_in_short(self): + t = _turn("system", 0) + turns = [t] + token_counts = {0: 500} + result = select_turns(turns, [], token_counts, budget=1000, short_threshold=300) + assert result.short_system_tokens == 0 + + +# --------------------------------------------------------------------------- +# select_turns — last system turn always kept +# --------------------------------------------------------------------------- + +class TestSelectTurnsLastSystem: + def test_last_system_turn_always_kept(self): + turns = [ + _turn("user", 0), + _turn("system", 1), + ] + token_counts = {0: 10, 1: 10_000} + # Even with tiny budget, last system should be kept + result = select_turns(turns, [], token_counts, budget=11) + kept_indices = {t.index for t in result.kept_turns} + assert 1 in kept_indices + + +# --------------------------------------------------------------------------- +# select_turns — scored long system turns +# --------------------------------------------------------------------------- + +class TestSelectTurnsScoredSystem: + def test_high_score_system_kept_within_budget(self): + user_t = _turn("user", 0) + sys_t = _turn("system", 1) # high score + sys_t2 = _turn("system", 2) # low score, and last system (always kept) + + scored = [ + _scored(sys_t, score=0.9, tokens=100), + _scored(sys_t2, score=0.1, tokens=100), + ] + turns = [user_t, sys_t, sys_t2] + token_counts = {0: 10, 1: 100, 2: 100} + # Budget: user(10) + last_system(100) + sys_t(100) = 210 → set budget=220 + result = select_turns( + turns, scored, token_counts, + budget=220, + short_threshold=50, + ) + kept_indices = {t.index for t in result.kept_turns} + assert 1 in kept_indices # high-score kept within budget + + def test_over_budget_system_dropped(self): + user_t = _turn("user", 0) + sys_t = _turn("system", 1) + scored = [_scored(sys_t, score=0.9, tokens=1000)] + turns = [user_t, sys_t] + token_counts = {0: 10, 1: 1000} + result = select_turns( + turns, scored, token_counts, + budget=50, + short_threshold=50, + ) + # sys_t should be dropped (but last-system rule may save it — check dropped_turns) + # The last-system-turn rule keeps sys_t regardless of budget + # So it ends up in kept_turns + # This test verifies scored_dropped is handled if there are multiple + assert result is not None # no crash + + def test_total_input_tokens_sum(self): + turns = [ + _turn("user", 0), + _turn("system", 1), + _turn("system", 2), + ] + token_counts = {0: 100, 1: 200, 2: 300} + result = select_turns(turns, [], token_counts, budget=10_000) + assert result.total_input_tokens == 600 + + def test_budget_stored_in_result(self): + result = select_turns([], [], {}, budget=42_000) + assert result.budget == 42_000 + + def test_kept_turns_in_index_order(self): + turns = [ + _turn("user", 0), + _turn("system", 1), + _turn("user", 2), + _turn("system", 3), + ] + token_counts = {0: 10, 1: 10, 2: 10, 3: 10} + result = select_turns(turns, [], token_counts, budget=10_000) + indices = [t.index for t in result.kept_turns] + assert indices == sorted(indices) diff --git a/tests/test_setcover.py b/tests/test_setcover.py new file mode 100644 index 0000000..3604045 --- /dev/null +++ b/tests/test_setcover.py @@ -0,0 +1,113 @@ +"""Tests for lib/setcover.py — setcover_scores.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.setcover import setcover_scores +from lib.parser import Turn +from lib.types import ScoredTurn + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(kind: str, index: int, text: str) -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +# --------------------------------------------------------------------------- +# setcover_scores +# --------------------------------------------------------------------------- + +class TestSetcoverScoresEmpty: + def test_no_system_turns_returns_empty(self): + turns = [_turn("user", 0, "hello")] + result = setcover_scores(turns, [], {0: 10}) + assert result == [] + + def test_empty_all_turns_returns_empty(self): + result = setcover_scores([], [], {}) + assert result == [] + + +class TestSetcoverScoresBasic: + def test_returns_scored_turns(self): + t = _turn("system", 1, "ValueError at /path/to/module.py") + all_turns = [_turn("user", 0, "question"), t] + result = setcover_scores(all_turns, [t], {0: 10, 1: 80}) + assert len(result) == 1 + assert isinstance(result[0], ScoredTurn) + + def test_scores_normalized_to_zero_one(self): + turns = [ + _turn("user", 0, "what happened?"), + _turn("system", 1, "ValueError in /src/main.py at port :8080"), + _turn("system", 2, "ModuleNotFoundError for package requests"), + _turn("system", 3, "generic content without entities"), + ] + system = turns[1:] + token_counts = {i: 60 for i in range(4)} + result = setcover_scores(turns, system, token_counts) + for st in result: + assert 0.0 <= st.score <= 1.0 + + def test_max_score_is_one(self): + turns = [ + _turn("user", 0, "question"), + _turn("system", 1, "ValueError at /home/user/project/main.py"), + _turn("system", 2, "just plain text no entities"), + ] + system = turns[1:] + token_counts = {0: 10, 1: 80, 2: 40} + result = setcover_scores(turns, system, token_counts) + scores = [st.score for st in result] + assert max(scores) == pytest.approx(1.0) + + def test_one_result_per_system_turn(self): + turns = [_turn("system", i, f"content{i}") for i in range(5)] + result = setcover_scores(turns, turns, {i: 50 for i in range(5)}) + assert len(result) == 5 + + def test_tokens_set_from_token_counts(self): + t = _turn("system", 2, "ValueError occurred") + all_turns = [t] + result = setcover_scores(all_turns, [t], {2: 456}) + assert result[0].tokens == 456 + + def test_exclusive_entity_gets_higher_score(self): + # Turn 1 has a unique entity (only appears once in system turns) + # Turn 2 has an entity shared across many turns (less exclusive) + common = "ValueError" # appears in all turns + unique_path = "/unique/path/to/special_module.py" # only in turn 1 + + turns = [ + _turn("user", 0, f"{common}"), + _turn("system", 1, f"{common} at {unique_path}"), + _turn("system", 2, f"{common} again"), + _turn("system", 3, f"{common} once more"), + ] + system = turns[1:] + token_counts = {i: 60 for i in range(4)} + result = setcover_scores(turns, system, token_counts) + score_map = {st.turn.index: st.score for st in result} + # Turn 1 has the unique path, should score at least as high as others + assert score_map[1] >= score_map[2] + + def test_all_plain_text_scores_all_zero_or_valid(self): + # No extractable entities → all raw scores are 0 → max_score=1.0 default + turns = [ + _turn("system", 0, "plain words here aaaa"), + _turn("system", 1, "more plain words bbbbb"), + ] + result = setcover_scores(turns, turns, {0: 50, 1: 50}) + for st in result: + assert 0.0 <= st.score <= 1.0 diff --git a/tests/test_types_and_probes.py b/tests/test_types_and_probes.py new file mode 100644 index 0000000..86ab373 --- /dev/null +++ b/tests/test_types_and_probes.py @@ -0,0 +1,226 @@ +"""Tests for lib/types.py (ScoredTurn, build_query, random_scores) and +lib/eval/probes.py (Probe, ProbeSet, _format_turns_for_prompt).""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) + +from lib.parser import Turn +from lib.types import ScoredTurn, build_query, random_scores +from lib.eval.probes import Probe, ProbeSet, _format_turns_for_prompt + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _turn(text: str, kind: str = "user", index: int = 0) -> Turn: + t = Turn(kind=kind, index=index) + t.append({"message": {"content": text}}) + return t + + +# --------------------------------------------------------------------------- +# ScoredTurn dataclass +# --------------------------------------------------------------------------- + +class TestScoredTurn: + def test_fields_stored(self): + t = _turn("hello") + st = ScoredTurn(turn=t, score=0.75, tokens=120) + assert st.turn is t + assert st.score == pytest.approx(0.75) + assert st.tokens == 120 + + +# --------------------------------------------------------------------------- +# build_query +# --------------------------------------------------------------------------- + +class TestBuildQuery: + def test_single_user_turn(self): + turns = [_turn("what is the bug?")] + result = build_query(turns) + assert "what is the bug?" in result + + def test_uses_last_three_turns(self): + turns = [_turn(f"msg{i}") for i in range(5)] + result = build_query(turns) + # Only last 3 should appear + assert "msg4" in result + assert "msg3" in result + assert "msg2" in result + # Earlier messages excluded + assert "msg0" not in result + assert "msg1" not in result + + def test_fewer_than_three_returns_all(self): + turns = [_turn("first"), _turn("second")] + result = build_query(turns) + assert "first" in result + assert "second" in result + + def test_truncated_at_max_chars(self): + big_text = "x" * 5000 + turns = [_turn(big_text)] + result = build_query(turns, max_chars=500) + assert len(result) <= 500 + + def test_empty_turns_returns_empty(self): + result = build_query([]) + assert result == "" + + def test_parts_separated_by_divider(self): + turns = [_turn("alpha"), _turn("beta"), _turn("gamma")] + result = build_query(turns) + assert "---" in result + + +# --------------------------------------------------------------------------- +# random_scores +# --------------------------------------------------------------------------- + +class TestRandomScores: + def test_returns_one_per_turn(self): + turns = [_turn("a", kind="system", index=i) for i in range(4)] + token_counts = {0: 100, 1: 200, 2: 150, 3: 50} + scored = random_scores(turns, token_counts) + assert len(scored) == 4 + + def test_returns_scored_turns(self): + turns = [_turn("hello", kind="system", index=0)] + scored = random_scores(turns, {0: 100}) + assert isinstance(scored[0], ScoredTurn) + + def test_scores_between_zero_and_one(self): + turns = [_turn("x", kind="system", index=i) for i in range(10)] + token_counts = {i: 50 for i in range(10)} + scored = random_scores(turns, token_counts) + for st in scored: + assert 0.0 <= st.score <= 1.0 + + def test_tokens_from_token_counts(self): + t = _turn("hello", kind="system", index=7) + scored = random_scores([t], {7: 333}) + assert scored[0].tokens == 333 + + def test_missing_token_count_defaults_to_zero(self): + t = _turn("hello", kind="system", index=99) + scored = random_scores([t], {}) + assert scored[0].tokens == 0 + + +# --------------------------------------------------------------------------- +# Probe dataclass +# --------------------------------------------------------------------------- + +class TestProbeDataclass: + def test_fields_stored(self): + p = Probe( + id="esr_001", + dimension="error_solution", + tier="factual", + question="What caused the error?", + gold_answer="Missing import", + ) + assert p.id == "esr_001" + assert p.dimension == "error_solution" + assert p.tier == "factual" + assert p.question == "What caused the error?" + assert p.gold_answer == "Missing import" + + def test_default_evidence_turns_empty(self): + p = Probe(id="p1", dimension="progress", tier="factual", + question="Q?", gold_answer="A") + assert p.evidence_turns == [] + + def test_default_difficulty_medium(self): + p = Probe(id="p1", dimension="progress", tier="factual", + question="Q?", gold_answer="A") + assert p.difficulty == "medium" + + +# --------------------------------------------------------------------------- +# ProbeSet dataclass +# --------------------------------------------------------------------------- + +class TestProbeSet: + def test_defaults(self): + ps = ProbeSet() + assert ps.probes == [] + assert ps.conv_hash == "" + assert ps.split_ratio == pytest.approx(0.70) + assert ps.version == "1" + + def test_to_dict_structure(self): + p = Probe(id="p1", dimension="progress", tier="factual", + question="Q?", gold_answer="A") + ps = ProbeSet(probes=[p], conv_hash="abc123", split_ratio=0.7, version="1") + d = ps.to_dict() + assert d["conv_hash"] == "abc123" + assert d["split_ratio"] == pytest.approx(0.7) + assert d["version"] == "1" + assert len(d["probes"]) == 1 + assert d["probes"][0]["id"] == "p1" + + def test_from_dict_roundtrip(self): + p = Probe(id="p1", dimension="instruction", tier="comprehension", + question="Why?", gold_answer="Because", + evidence_turns=[0, 2], difficulty="hard") + ps = ProbeSet(probes=[p], conv_hash="xyz", split_ratio=0.8, version="2") + d = ps.to_dict() + restored = ProbeSet.from_dict(d) + assert restored.conv_hash == "xyz" + assert restored.split_ratio == pytest.approx(0.8) + assert len(restored.probes) == 1 + rp = restored.probes[0] + assert rp.id == "p1" + assert rp.dimension == "instruction" + assert rp.difficulty == "hard" + assert rp.evidence_turns == [0, 2] + + def test_from_dict_missing_optional_fields(self): + data = {"probes": []} + ps = ProbeSet.from_dict(data) + assert ps.probes == [] + assert ps.conv_hash == "" + assert ps.split_ratio == pytest.approx(0.70) + + +# --------------------------------------------------------------------------- +# _format_turns_for_prompt +# --------------------------------------------------------------------------- + +class TestFormatTurnsForPrompt: + def test_single_turn_included(self): + t = _turn("hello world", kind="user", index=0) + result = _format_turns_for_prompt([t]) + assert "hello world" in result + + def test_turn_header_included(self): + t = _turn("content", kind="user", index=5) + result = _format_turns_for_prompt([t]) + assert "Turn 5" in result + assert "user" in result + + def test_multiple_turns_ordered(self): + turns = [_turn(f"msg{i}", index=i) for i in range(3)] + result = _format_turns_for_prompt(turns) + # All messages should appear + for i in range(3): + assert f"msg{i}" in result + + def test_truncated_at_max_chars(self): + big_text = "y" * 200_000 + t = _turn(big_text, index=0) + result = _format_turns_for_prompt([t], max_chars=1000) + assert "truncated" in result + + def test_empty_list_returns_empty(self): + result = _format_turns_for_prompt([]) + assert result == ""