From 2db15130965ba1bc4da868c9b5eae4ebffbf641e Mon Sep 17 00:00:00 2001 From: DaevMithran Date: Thu, 16 Apr 2026 11:51:48 +0530 Subject: [PATCH 1/6] feat: Add AMA benchmark --- .gitignore | 3 + benchmarks/ama-bench/README.md | 144 ++++++++++++++++++ benchmarks/ama-bench/configs/default.json | 16 ++ benchmarks/ama-bench/configs/sweep.json | 7 + benchmarks/ama-bench/docker/.env.example | 8 + benchmarks/ama-bench/docker/Dockerfile | 22 +++ benchmarks/ama-bench/docker/bridge.Dockerfile | 14 ++ .../ama-bench/docker/docker-compose.yml | 44 ++++++ benchmarks/ama-bench/package.json | 16 ++ benchmarks/ama-bench/scripts/run.sh | 62 ++++++++ benchmarks/ama-bench/scripts/setup.sh | 40 +++++ benchmarks/ama-bench/scripts/sweep.sh | 131 ++++++++++++++++ benchmarks/ama-bench/src/server.ts | 129 ++++++++++++++++ benchmarks/ama-bench/tsconfig.json | 12 ++ pnpm-lock.yaml | 24 +++ pnpm-workspace.yaml | 3 +- 16 files changed, 674 insertions(+), 1 deletion(-) create mode 100644 benchmarks/ama-bench/README.md create mode 100644 benchmarks/ama-bench/configs/default.json create mode 100644 benchmarks/ama-bench/configs/sweep.json create mode 100644 benchmarks/ama-bench/docker/.env.example create mode 100644 benchmarks/ama-bench/docker/Dockerfile create mode 100644 benchmarks/ama-bench/docker/bridge.Dockerfile create mode 100644 benchmarks/ama-bench/docker/docker-compose.yml create mode 100644 benchmarks/ama-bench/package.json create mode 100755 benchmarks/ama-bench/scripts/run.sh create mode 100755 benchmarks/ama-bench/scripts/setup.sh create mode 100755 benchmarks/ama-bench/scripts/sweep.sh create mode 100644 benchmarks/ama-bench/src/server.ts create mode 100644 benchmarks/ama-bench/tsconfig.json diff --git a/.gitignore b/.gitignore index 57dcfc8..6bc3537 100644 --- a/.gitignore +++ b/.gitignore @@ -75,3 +75,6 @@ my-app/ # ROFL config (user-generated) rofl.yaml + +# Benchmark results +benchmarks/**/results/ diff --git a/benchmarks/ama-bench/README.md b/benchmarks/ama-bench/README.md new file mode 100644 index 0000000..148aaef --- /dev/null +++ b/benchmarks/ama-bench/README.md @@ -0,0 +1,144 @@ +# AMA-Bench Benchmark for @ekai/mindmap + +Benchmarks the `@ekai/mindmap` package against [AMA-Bench](https://github.com/ekailabs/AMA-Bench), an evaluation framework for Associative Memory Ability in AI agents. + +## Architecture + +``` +AMA-Bench (Python) Bridge Server (TypeScript) +┌──────────────────────┐ ┌──────────────────────────-┐ +│ run.py │ │ server.ts │ +│ └─ ContextoMethod │ HTTP │ └─ @ekai/mindmap │ +│ │ │──────────────▶ ├─ mindmap.add() │ +│ │ construct │ /construct │ │ (embed+cluster) │ +│ │ retrieve │ /retrieve │ └─ mindmap.search() │ +│ ▼ │ │ (beam search) │ +│ LLM generates answer │ │ │ +│ Judge scores answer │ │ reads configs/default.json│ +└──────────────────────┘ └──────────────────────────-┘ +``` + +The bridge server owns all mindmap configuration (`configs/default.json`) and the embedding API key (`API_KEY` env var). The Python method is a thin HTTP client that sends trajectory data and questions. + +## Quick Start (Docker Compose) + +No local Bun or Python needed. + +```bash +cd docker +cp .env.example .env # set API_KEY for embeddings +docker compose up --build +``` + +Override defaults: + +```bash +SUBSET=mcq LLM_CONFIG=claude-sonnet.yaml docker compose up --build +``` + +### CI + +```yaml +- name: Run AMA-Bench + working-directory: benchmarks/ama-bench/docker + env: + API_KEY: ${{ secrets.API_KEY }} + run: docker compose up --build --abort-on-container-exit +``` + +## Local Setup + +### Prerequisites + +- [Bun](https://bun.sh) >= 1.0 +- Python >= 3.9 +- pnpm (for workspace install) +- `huggingface-cli` (for dataset download) + +### Install + +```bash +bash scripts/setup.sh +``` + +### Configure + +1. Edit `configs/default.json` — provider, embed model, mindmap and search params +2. Set your embedding API key: + ```bash + export API_KEY=your-openrouter-or-openai-key + ``` + +### Run benchmark + +```bash +bash scripts/run.sh +``` + +Override AMA-Bench options: + +```bash +LLM_CONFIG=/path/to/llm.yaml SUBSET=openend bash scripts/run.sh +``` + +### Run parameter sweep + +Grid-search over mindmap/search params: + +```bash +bash scripts/sweep.sh +``` + +Edit `configs/sweep.json` to change the parameter ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. + +## Configuration + +All mindmap parameters are in `configs/default.json`: + +### Tree construction (`mindmap`) + +| Parameter | Default | Description | +|---|---|---| +| `similarityThreshold` | 0.5 | Min cosine similarity to cluster items together | +| `maxDepth` | 4 | Max tree nesting depth | +| `maxChildren` | 10 | Max direct children per node | +| `rebuildInterval` | 50 | Items added before full tree rebuild | + +### Retrieval (`search`) + +| Parameter | Default | Description | +|---|---|---| +| `maxResults` | 10 | Max items returned | +| `maxTokens` | 4000 | Token budget cap for results | +| `beamWidth` | 3 | Branches explored per tree level | +| `minScore` | 0.0 | Min cosine similarity to include a result | + +## Bridge API + +| Endpoint | Method | Description | +|---|---|---| +| `/health` | GET | Health check, returns `{ status, activeEpisodes }` | +| `/construct` | POST | Add trajectory items to a mindmap instance | +| `/retrieve` | POST | Search mindmap for relevant context | +| `/reset` | POST | Clear a mindmap instance for an episode | + +## File Structure + +``` +benchmarks/ama-bench/ +├── src/server.ts # Bridge server wrapping @ekai/mindmap +├── package.json +├── configs/ +│ ├── default.json # Mindmap + search parameters +│ └── sweep.json # Parameter sweep ranges +├── scripts/ +│ ├── setup.sh # One-time setup +│ ├── run.sh # Run benchmark +│ └── sweep.sh # Run parameter sweep +├── docker/ +│ ├── docker-compose.yml +│ ├── Dockerfile # AMA-Bench runner +│ ├── bridge.Dockerfile # Bridge server +│ └── .env.example +└── results/ # Benchmark outputs (gitignored) +``` diff --git a/benchmarks/ama-bench/configs/default.json b/benchmarks/ama-bench/configs/default.json new file mode 100644 index 0000000..e631169 --- /dev/null +++ b/benchmarks/ama-bench/configs/default.json @@ -0,0 +1,16 @@ +{ + "provider": "openrouter", + "embedModel": "openai/text-embedding-3-small", + "mindmap": { + "similarityThreshold": 0.5, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 4000, + "beamWidth": 3, + "minScore": 0.0 + } +} diff --git a/benchmarks/ama-bench/configs/sweep.json b/benchmarks/ama-bench/configs/sweep.json new file mode 100644 index 0000000..1ac4ff8 --- /dev/null +++ b/benchmarks/ama-bench/configs/sweep.json @@ -0,0 +1,7 @@ +{ + "similarityThreshold": [0.3, 0.5, 0.65, 0.8], + "maxDepth": [3, 4, 6], + "beamWidth": [2, 3, 5], + "minScore": [0.0, 0.1, 0.3], + "maxResults": [5, 10, 20] +} diff --git a/benchmarks/ama-bench/docker/.env.example b/benchmarks/ama-bench/docker/.env.example new file mode 100644 index 0000000..ab945b1 --- /dev/null +++ b/benchmarks/ama-bench/docker/.env.example @@ -0,0 +1,8 @@ +# Embedding API key (OpenRouter or OpenAI) +API_KEY=your-api-key + +# AMA-Bench options (these reference files in AMA-Bench/configs/) +# Available LLM configs: gpt-5.2.yaml, qwen3-32B.yaml +LLM_CONFIG=gpt-5.2.yaml +JUDGE_CONFIG=llm_judge.yaml +SUBSET=openend diff --git a/benchmarks/ama-bench/docker/Dockerfile b/benchmarks/ama-bench/docker/Dockerfile new file mode 100644 index 0000000..dd1242f --- /dev/null +++ b/benchmarks/ama-bench/docker/Dockerfile @@ -0,0 +1,22 @@ +FROM python:3.11-slim + +RUN apt-get update && apt-get install -y --no-install-recommends \ + git curl && \ + rm -rf /var/lib/apt/lists/* + +WORKDIR /app + +# Clone AMA-Bench fork (contexto_method.py and configs already in repo) +ARG AMA_BENCH_REPO=https://github.com/ekailabs/AMA-Bench.git +RUN git clone ${AMA_BENCH_REPO} /app/AMA-Bench + +# Install Python dependencies +RUN pip install --no-cache-dir -r /app/AMA-Bench/requirements.txt + +# Download dataset +RUN pip install --no-cache-dir huggingface_hub && \ + huggingface-cli download AMA-bench/AMA-bench --repo-type dataset --local-dir /app/AMA-Bench/dataset + +WORKDIR /app/AMA-Bench + +ENTRYPOINT ["python", "src/run.py"] diff --git a/benchmarks/ama-bench/docker/bridge.Dockerfile b/benchmarks/ama-bench/docker/bridge.Dockerfile new file mode 100644 index 0000000..aea0c9d --- /dev/null +++ b/benchmarks/ama-bench/docker/bridge.Dockerfile @@ -0,0 +1,14 @@ +FROM oven/bun:1-slim + +WORKDIR /app + +COPY package.json ./ +COPY src/ ./src/ + +# Install @ekai/mindmap from npm (swap workspace ref) +RUN sed -i 's/"workspace:\*"/"latest"/' package.json && \ + bun install + +EXPOSE 3456 + +CMD ["bun", "src/server.ts"] diff --git a/benchmarks/ama-bench/docker/docker-compose.yml b/benchmarks/ama-bench/docker/docker-compose.yml new file mode 100644 index 0000000..c879954 --- /dev/null +++ b/benchmarks/ama-bench/docker/docker-compose.yml @@ -0,0 +1,44 @@ +services: + bridge: + build: + context: .. + dockerfile: docker/bridge.Dockerfile + ports: + - "3456:3456" + environment: + - BRIDGE_PORT=3456 + - API_KEY=${API_KEY} + healthcheck: + test: ["CMD", "curl", "-sf", "http://localhost:3456/health"] + interval: 5s + timeout: 3s + retries: 10 + + runner: + build: + context: . + dockerfile: Dockerfile + depends_on: + bridge: + condition: service_healthy + environment: + - CONTEXTO_BRIDGE_URL=http://bridge:3456 + command: + - --llm-server + - api + - --llm-config + - configs/${LLM_CONFIG:-gpt-5.2.yaml} + - --subset + - ${SUBSET:-openend} + - --method + - contexto + - --method-config + - configs/contexto.yaml + - --test-dir + - dataset/test + - --judge-config + - configs/${JUDGE_CONFIG:-llm_judge.yaml} + - --evaluate + - "True" + volumes: + - ../results:/app/AMA-Bench/results diff --git a/benchmarks/ama-bench/package.json b/benchmarks/ama-bench/package.json new file mode 100644 index 0000000..064598e --- /dev/null +++ b/benchmarks/ama-bench/package.json @@ -0,0 +1,16 @@ +{ + "name": "@ekai/ama-bench-bridge", + "version": "0.1.0", + "private": true, + "description": "HTTP bridge between AMA-Bench Python framework and @ekai/mindmap", + "type": "module", + "scripts": { + "start": "bun src/server.ts" + }, + "dependencies": { + "@ekai/mindmap": "workspace:*" + }, + "devDependencies": { + "@types/bun": "latest" + } +} diff --git a/benchmarks/ama-bench/scripts/run.sh b/benchmarks/ama-bench/scripts/run.sh new file mode 100755 index 0000000..929f646 --- /dev/null +++ b/benchmarks/ama-bench/scripts/run.sh @@ -0,0 +1,62 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BENCH_DIR="$(dirname "$SCRIPT_DIR")" +AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" +BRIDGE_PORT="${BRIDGE_PORT:-3456}" + +# Parse arguments (pass through to run.py) +LLM_CONFIG="${LLM_CONFIG:-$AMA_BENCH/configs/gpt-4o.yaml}" +SUBSET="${SUBSET:-openend}" +JUDGE_CONFIG="${JUDGE_CONFIG:-$AMA_BENCH/configs/llm_judge.yaml}" +METHOD_CONFIG="${METHOD_CONFIG:-$AMA_BENCH/configs/contexto.yaml}" +EXTRA_ARGS="${@}" + +echo "=== Running AMA-Bench with contexto method ===" +echo "AMA-Bench dir: $AMA_BENCH" + +# Start bridge server in background +echo "[1/3] Starting contexto bridge server on port $BRIDGE_PORT..." +cd "$BENCH_DIR" +BRIDGE_PORT=$BRIDGE_PORT bun src/server.ts & +BRIDGE_PID=$! + +# Ensure cleanup on exit +cleanup() { + echo "" + echo "Shutting down bridge server (PID: $BRIDGE_PID)..." + kill $BRIDGE_PID 2>/dev/null || true + wait $BRIDGE_PID 2>/dev/null || true +} +trap cleanup EXIT + +# Wait for bridge to be ready +echo "Waiting for bridge server..." +for i in $(seq 1 30); do + if curl -s "http://localhost:$BRIDGE_PORT/health" > /dev/null 2>&1; then + echo "Bridge server ready." + break + fi + if [ "$i" -eq 30 ]; then + echo "ERROR: Bridge server failed to start within 30 seconds." + exit 1 + fi + sleep 1 +done + +# Run AMA-Bench +echo "[2/3] Running AMA-Bench evaluation..." +cd "$AMA_BENCH" +python src/run.py \ + --llm-server api \ + --llm-config "$LLM_CONFIG" \ + --subset "$SUBSET" \ + --method contexto \ + --method-config "$METHOD_CONFIG" \ + --test-dir dataset/test \ + --judge-config "$JUDGE_CONFIG" \ + --evaluate True \ + $EXTRA_ARGS + +echo "[3/3] Benchmark complete. Results saved in $AMA_BENCH/results/" diff --git a/benchmarks/ama-bench/scripts/setup.sh b/benchmarks/ama-bench/scripts/setup.sh new file mode 100755 index 0000000..0dbf12d --- /dev/null +++ b/benchmarks/ama-bench/scripts/setup.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BENCH_DIR="$(dirname "$SCRIPT_DIR")" +AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" + +echo "=== AMA-Bench Setup for Contexto ===" +echo "AMA-Bench: $AMA_BENCH" + +# 1. Clone AMA-Bench if not present +if [ ! -d "$AMA_BENCH" ]; then + echo "[1/4] Cloning AMA-Bench..." + git clone https://github.com/ekailabs/AMA-Bench.git "$AMA_BENCH" +else + echo "[1/4] AMA-Bench already exists at $AMA_BENCH" +fi + +# 2. Install Python dependencies +echo "[2/4] Installing Python dependencies..." +pip install -r "$AMA_BENCH/requirements.txt" + +# 3. Download dataset +if [ ! -d "$AMA_BENCH/dataset" ]; then + echo "[3/4] Downloading AMA-Bench dataset..." + huggingface-cli download AMA-bench/AMA-bench --repo-type dataset --local-dir "$AMA_BENCH/dataset" +else + echo "[3/4] Dataset already downloaded." +fi + +# 4. Install bridge dependencies +echo "[4/4] Installing bridge dependencies..." +cd "$BENCH_DIR" && pnpm install + +echo "" +echo "=== Setup complete ===" +echo "Next steps:" +echo " 1. Edit configs/default.json (provider, embedModel, mindmap params)" +echo " 2. Export API_KEY=your-api-key" +echo " 3. Run: bash scripts/run.sh" diff --git a/benchmarks/ama-bench/scripts/sweep.sh b/benchmarks/ama-bench/scripts/sweep.sh new file mode 100755 index 0000000..8f4da81 --- /dev/null +++ b/benchmarks/ama-bench/scripts/sweep.sh @@ -0,0 +1,131 @@ +#!/usr/bin/env bash +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BENCH_DIR="$(dirname "$SCRIPT_DIR")" +AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" +BRIDGE_PORT="${BRIDGE_PORT:-3456}" +SWEEP_CONFIG="${SWEEP_CONFIG:-$BENCH_DIR/configs/sweep.json}" +DEFAULT_CONFIG="$BENCH_DIR/configs/default.json" +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +RESULTS_DIR="$BENCH_DIR/results/sweep_$TIMESTAMP" + +LLM_CONFIG="${LLM_CONFIG:-$AMA_BENCH/configs/gpt-4o.yaml}" +JUDGE_CONFIG="${JUDGE_CONFIG:-$AMA_BENCH/configs/llm_judge.yaml}" + +echo "=== Contexto Parameter Sweep ===" +echo "Sweep config: $SWEEP_CONFIG" +echo "Results dir: $RESULTS_DIR" +mkdir -p "$RESULTS_DIR" + +# Save original default.json to restore later +cp "$DEFAULT_CONFIG" "$RESULTS_DIR/_default.json.bak" + +# Read param arrays from sweep config +read -ra SIM_THRESHOLDS <<< "$(jq -r '.similarityThreshold | join(" ")' "$SWEEP_CONFIG")" +read -ra MAX_DEPTHS <<< "$(jq -r '.maxDepth | join(" ")' "$SWEEP_CONFIG")" +read -ra BEAM_WIDTHS <<< "$(jq -r '.beamWidth | join(" ")' "$SWEEP_CONFIG")" +read -ra MIN_SCORES <<< "$(jq -r '.minScore | join(" ")' "$SWEEP_CONFIG")" +read -ra MAX_RESULTS <<< "$(jq -r '.maxResults | join(" ")' "$SWEEP_CONFIG")" + +TOTAL=$(( ${#SIM_THRESHOLDS[@]} * ${#MAX_DEPTHS[@]} * ${#BEAM_WIDTHS[@]} * ${#MIN_SCORES[@]} * ${#MAX_RESULTS[@]} )) +echo "Total configs: $TOTAL" + +# Ensure cleanup on exit +BRIDGE_PID="" +cleanup() { + echo "" + if [ -n "$BRIDGE_PID" ]; then + echo "Shutting down bridge server..." + kill $BRIDGE_PID 2>/dev/null || true + wait $BRIDGE_PID 2>/dev/null || true + fi + # Restore original config + cp "$RESULTS_DIR/_default.json.bak" "$DEFAULT_CONFIG" +} +trap cleanup EXIT + +# Run sweep +echo "[1/2] Running parameter sweep..." +COUNT=0 +SUMMARY="$RESULTS_DIR/sweep_summary.csv" +echo "similarityThreshold,maxDepth,beamWidth,minScore,maxResults,accuracy" > "$SUMMARY" + +for st in "${SIM_THRESHOLDS[@]}"; do +for md in "${MAX_DEPTHS[@]}"; do +for bw in "${BEAM_WIDTHS[@]}"; do +for ms in "${MIN_SCORES[@]}"; do +for mr in "${MAX_RESULTS[@]}"; do + COUNT=$((COUNT + 1)) + echo "" + echo "[$COUNT/$TOTAL] st=$st md=$md bw=$bw ms=$ms mr=$mr" + + # Write config for this combo + cat > "$DEFAULT_CONFIG" </dev/null || true + wait $BRIDGE_PID 2>/dev/null || true + fi + cd "$BENCH_DIR" + BRIDGE_PORT=$BRIDGE_PORT bun src/server.ts & + BRIDGE_PID=$! + + for i in $(seq 1 30); do + if curl -s "http://localhost:$BRIDGE_PORT/health" > /dev/null 2>&1; then + break + fi + [ "$i" -eq 30 ] && echo "ERROR: Bridge timeout" && exit 1 + sleep 1 + done + + # Run benchmark + cd "$AMA_BENCH" + OUTPUT=$(python src/run.py \ + --llm-server api \ + --llm-config "$LLM_CONFIG" \ + --subset openend \ + --method contexto \ + --method-config configs/contexto.yaml \ + --test-dir dataset/test \ + --judge-config "$JUDGE_CONFIG" \ + --evaluate True 2>&1) || true + + # Parse accuracy from output + ACCURACY=$(echo "$OUTPUT" | grep -i "overall" | grep -oE '[0-9]+\.[0-9]+' | head -1 || echo "0.0") + echo " -> Accuracy: $ACCURACY" + echo "$st,$md,$bw,$ms,$mr,$ACCURACY" >> "$SUMMARY" + +done +done +done +done +done + +# Print ranked results +echo "" +echo "============================================================" +echo "SWEEP RESULTS (ranked by accuracy)" +echo "============================================================" +sort -t, -k6 -rn "$SUMMARY" | head -11 + +echo "" +echo "[2/2] Sweep complete. Results: $SUMMARY" diff --git a/benchmarks/ama-bench/src/server.ts b/benchmarks/ama-bench/src/server.ts new file mode 100644 index 0000000..d321ec9 --- /dev/null +++ b/benchmarks/ama-bench/src/server.ts @@ -0,0 +1,129 @@ +import { + createMindmap, + memoryStorage, + type Mindmap, + type MindmapConfig, + type SearchOptions, +} from '@ekai/mindmap'; +import { readFileSync } from 'fs'; +import { resolve, dirname } from 'path'; +import { fileURLToPath } from 'url'; + +// --- Config --- + +const __dirname = dirname(fileURLToPath(import.meta.url)); +const configPath = resolve(__dirname, '../configs/default.json'); +const config = JSON.parse(readFileSync(configPath, 'utf-8')); + +const provider = config.provider ?? 'openrouter'; +const embedModel = config.embedModel ?? 'openai/text-embedding-3-small'; +const apiKey = process.env.API_KEY ?? ''; +const mindmapConfig: Partial = config.mindmap ?? {}; +const searchDefaults: SearchOptions = config.search ?? {}; + +if (!apiKey) { + console.warn('[bridge] WARNING: API_KEY env var not set. Embedding calls will fail.'); +} + +// --- Types --- + +interface ConstructRequest { + episodeId: string; + items: Array<{ id: string; role: string; content: string; metadata?: Record }>; +} + +interface RetrieveRequest { + episodeId: string; + question: string; + searchOptions?: SearchOptions; +} + +interface ResetRequest { + episodeId: string; +} + +// --- State --- + +const mindmaps = new Map(); + +// --- Handlers --- + +async function handleConstruct(body: ConstructRequest) { + const { episodeId, items } = body; + + mindmaps.delete(episodeId); + + const mindmap = createMindmap({ + provider, + apiKey, + embedModel, + storage: memoryStorage(), + config: mindmapConfig, + }); + + await mindmap.add(items); + mindmaps.set(episodeId, mindmap); + + const state = await mindmap.getState(); + return { success: true, totalItems: state.stats.totalItems }; +} + +async function handleRetrieve(body: RetrieveRequest) { + const { episodeId, question, searchOptions } = body; + + const mindmap = mindmaps.get(episodeId); + if (!mindmap) { + throw new Error(`No mindmap found for episode ${episodeId}. Call /construct first.`); + } + + const opts = { ...searchDefaults, ...searchOptions }; + const result = await mindmap.search(question, opts); + const context = result.items.map((si) => si.item.content).join('\n\n'); + + return { context, totalCandidates: result.totalCandidates }; +} + +function handleReset(body: ResetRequest) { + mindmaps.delete(body.episodeId); + return { success: true }; +} + +// --- Server --- + +const PORT = parseInt(process.env.BRIDGE_PORT ?? '3456', 10); + +const server = Bun.serve({ + port: PORT, + async fetch(req) { + const url = new URL(req.url); + const path = url.pathname; + + if (req.method === 'GET' && path === '/health') { + return Response.json({ status: 'ok', activeEpisodes: mindmaps.size }); + } + + if (req.method !== 'POST') { + return Response.json({ error: 'Method not allowed' }, { status: 405 }); + } + + try { + const body = await req.json(); + + if (path === '/construct') { + return Response.json(await handleConstruct(body as ConstructRequest)); + } else if (path === '/retrieve') { + return Response.json(await handleRetrieve(body as RetrieveRequest)); + } else if (path === '/reset') { + return Response.json(handleReset(body as ResetRequest)); + } + + return Response.json({ error: 'Not found' }, { status: 404 }); + } catch (err) { + const message = err instanceof Error ? err.message : 'Unknown error'; + console.error(`[bridge] Error on ${path}:`, message); + return Response.json({ error: message }, { status: 500 }); + } + }, +}); + +console.log(`[bridge] Mindmap bridge server listening on http://localhost:${server.port}`); diff --git a/benchmarks/ama-bench/tsconfig.json b/benchmarks/ama-bench/tsconfig.json new file mode 100644 index 0000000..28579a2 --- /dev/null +++ b/benchmarks/ama-bench/tsconfig.json @@ -0,0 +1,12 @@ +{ + "compilerOptions": { + "target": "ESNext", + "module": "ESNext", + "moduleResolution": "bundler", + "types": ["bun"], + "strict": true, + "skipLibCheck": true, + "noEmit": true + }, + "include": ["src"] +} diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 63a8c86..fb87cad 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -39,6 +39,16 @@ importers: specifier: ^25.0.3 version: 25.0.3(typescript@5.9.3) + benchmarks/ama-bench: + dependencies: + '@ekai/mindmap': + specifier: workspace:* + version: link:../../packages/mindmap + devDependencies: + '@types/bun': + specifier: latest + version: 1.3.12 + packages/contexto: dependencies: openclaw: @@ -1779,6 +1789,9 @@ packages: '@types/body-parser@1.19.6': resolution: {integrity: sha512-HLFeCYgz89uk22N5Qg3dvGvsv46B8GLvKKo1zKG4NybA8U2DiEO3w9lqGg29t/tfLRJpJ6iQxnVw4OnB7MoM9g==} + '@types/bun@1.3.12': + resolution: {integrity: sha512-DBv81elK+/VSwXHDlnH3Qduw+KxkTIWi7TXkAeh24zpi5l0B2kUg9Ga3tb4nJaPcOFswflgi/yAvMVBPrxMB+A==} + '@types/chai@5.2.3': resolution: {integrity: sha512-Mw558oeA9fFbv65/y4mHtXDs9bPnFMZAL/jxdPFUpOHHIXX91mcgEHbS5Lahr+pwZFR8A7GQleRWeI6cGFC2UA==} @@ -2420,6 +2433,9 @@ packages: buffer@5.7.1: resolution: {integrity: sha512-EHcyIPBQ4BSGlvjB16k5KgAJ27CIsHY/2JBmCRReo48y9rQ3MaUzWX3KVlBa4U7MyX02HdVj0K7C3WaB3ju7FQ==} + bun-types@1.3.12: + resolution: {integrity: sha512-HqOLj5PoFajAQciOMRiIZGNoKxDJSr6qigAttOX40vJuSp6DN/CxWp9s3C1Xwm4oH7ybueITwiaOcWXoYVoRkA==} + bytes@3.1.2: resolution: {integrity: sha512-/Nf7TyzTx6S3yRJObOAV7956r8cr2+Oj8AC5dt8wSP3BQAoeX58NoHyCU8P8zGkNXStjTSi6fzO6F0pBdcYbEg==} engines: {node: '>= 0.8'} @@ -7946,6 +7962,10 @@ snapshots: '@types/connect': 3.4.38 '@types/node': 20.19.37 + '@types/bun@1.3.12': + dependencies: + bun-types: 1.3.12 + '@types/chai@5.2.3': dependencies: '@types/deep-eql': 4.0.2 @@ -8654,6 +8674,10 @@ snapshots: base64-js: 1.5.1 ieee754: 1.2.1 + bun-types@1.3.12: + dependencies: + '@types/node': 20.19.37 + bytes@3.1.2: {} cac@6.7.14: {} diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index eccc335..607f971 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -1,2 +1,3 @@ packages: - - 'packages/**' \ No newline at end of file + - 'packages/**' + - 'benchmarks/ama-bench' \ No newline at end of file From ae2cf728191e99cd89f231bb6f3dc2288e742eb6 Mon Sep 17 00:00:00 2001 From: DaevMithran Date: Fri, 17 Apr 2026 12:39:12 +0530 Subject: [PATCH 2/6] Update Readme --- benchmarks/ama-bench/README.md | 201 +++++++++++++++++++------ benchmarks/ama-bench/docker/Dockerfile | 4 +- benchmarks/ama-bench/scripts/run.sh | 19 ++- benchmarks/ama-bench/scripts/setup.sh | 16 +- 4 files changed, 178 insertions(+), 62 deletions(-) diff --git a/benchmarks/ama-bench/README.md b/benchmarks/ama-bench/README.md index 148aaef..e7a22de 100644 --- a/benchmarks/ama-bench/README.md +++ b/benchmarks/ama-bench/README.md @@ -6,7 +6,7 @@ Benchmarks the `@ekai/mindmap` package against [AMA-Bench](https://github.com/ek ``` AMA-Bench (Python) Bridge Server (TypeScript) -┌──────────────────────┐ ┌──────────────────────────-┐ +┌──────────────────────┐ ┌───────────────────────────┐ │ run.py │ │ server.ts │ │ └─ ContextoMethod │ HTTP │ └─ @ekai/mindmap │ │ │ │──────────────▶ ├─ mindmap.add() │ @@ -15,85 +15,175 @@ AMA-Bench (Python) Bridge Server (TypeScript) │ ▼ │ │ (beam search) │ │ LLM generates answer │ │ │ │ Judge scores answer │ │ reads configs/default.json│ -└──────────────────────┘ └──────────────────────────-┘ +└──────────────────────┘ └───────────────────────────┘ + +ekailabs/AMA-Bench repo ekailabs/contexto repo + src/method/contexto_method.py benchmarks/ama-bench/src/server.ts + configs/contexto.yaml benchmarks/ama-bench/configs/default.json ``` -The bridge server owns all mindmap configuration (`configs/default.json`) and the embedding API key (`API_KEY` env var). The Python method is a thin HTTP client that sends trajectory data and questions. +Two repos: +- **[ekailabs/AMA-Bench](https://github.com/ekailabs/AMA-Bench)** — Python benchmark framework + `contexto` method (thin HTTP client) +- **[ekailabs/contexto](https://github.com/ekailabs/contexto)** — Bridge server wrapping `@ekai/mindmap` + all config -## Quick Start (Docker Compose) +## Prerequisites -No local Bun or Python needed. +- [Bun](https://bun.sh) >= 1.0 +- Python >= 3.9 +- pnpm +- `huggingface-cli` (`pip install huggingface_hub`) +- An API key for [OpenRouter](https://openrouter.ai) or OpenAI (used for embeddings + LLM) -```bash -cd docker -cp .env.example .env # set API_KEY for embeddings -docker compose up --build -``` +## Running Locally -Override defaults: +### 1. Clone both repos ```bash -SUBSET=mcq LLM_CONFIG=claude-sonnet.yaml docker compose up --build +git clone https://github.com/ekailabs/contexto.git +git clone https://github.com/ekailabs/AMA-Bench.git ``` -### CI +They should be siblings: -```yaml -- name: Run AMA-Bench - working-directory: benchmarks/ama-bench/docker - env: - API_KEY: ${{ secrets.API_KEY }} - run: docker compose up --build --abort-on-container-exit +``` +parent/ +├── contexto/ +└── AMA-Bench/ ``` -## Local Setup +### 2. Install dependencies -### Prerequisites +```bash +# Install contexto workspace (includes the bridge) +cd contexto +pnpm install -- [Bun](https://bun.sh) >= 1.0 -- Python >= 3.9 -- pnpm (for workspace install) -- `huggingface-cli` (for dataset download) +# Install AMA-Bench Python deps +cd ../AMA-Bench +pip install -r requirements.txt -### Install +# Download the dataset +huggingface-cli download AMA-bench/AMA-bench --repo-type dataset --local-dir dataset +``` + +Or run the setup script which does all of the above: ```bash +cd contexto/benchmarks/ama-bench bash scripts/setup.sh ``` -### Configure +### 3. Configure the bridge + +Create `contexto/benchmarks/ama-bench/.env`: + +```bash +# Embedding API key (used by the bridge for mindmap embeddings) +API_KEY=your-openrouter-or-openai-key +``` + +Tune mindmap parameters in `contexto/benchmarks/ama-bench/configs/default.json`: + +```json +{ + "provider": "openrouter", + "embedModel": "openai/text-embedding-3-small", + "mindmap": { + "similarityThreshold": 0.5, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 4000, + "beamWidth": 3, + "minScore": 0.0 + } +} +``` + +### 4. Configure AMA-Bench LLM -1. Edit `configs/default.json` — provider, embed model, mindmap and search params -2. Set your embedding API key: - ```bash - export API_KEY=your-openrouter-or-openai-key - ``` +AMA-Bench needs an LLM config for answer generation and a judge config for scoring. Create these in `AMA-Bench/configs/`: -### Run benchmark +```yaml +# AMA-Bench/configs/openrouter.yaml +provider: "openai" +api_key: "your-openrouter-key" +model: "openai/gpt-4o" +base_url: "https://openrouter.ai/api/v1" +max_tokens: 16000 +temperature: 0.0 +``` + +```yaml +# AMA-Bench/configs/llm_judge_openrouter.yaml +provider: "openai" +api_key: "your-openrouter-key" +model: "openai/gpt-4o" +base_url: "https://openrouter.ai/api/v1" +max_tokens: 16000 +temperature: 0.0 +``` + +### 5. Run ```bash +cd contexto/benchmarks/ama-bench bash scripts/run.sh ``` -Override AMA-Bench options: +This will: +1. Start the bridge server (reads `configs/default.json` + `API_KEY` from `.env`) +2. Run AMA-Bench with the `contexto` method (208 episodes, ~35s each) +3. Evaluate answers with the LLM judge +4. Save results to `AMA-Bench/results/` +5. Shut down the bridge + +Override defaults: ```bash -LLM_CONFIG=/path/to/llm.yaml SUBSET=openend bash scripts/run.sh +LLM_CONFIG=../../../AMA-Bench/configs/openrouter.yaml \ +JUDGE_CONFIG=../../../AMA-Bench/configs/llm_judge_openrouter.yaml \ +SUBSET=openend \ +bash scripts/run.sh ``` -### Run parameter sweep +### 6. Parameter sweep (optional) -Grid-search over mindmap/search params: +Grid-search over mindmap/search params to find the optimal config: ```bash bash scripts/sweep.sh ``` -Edit `configs/sweep.json` to change the parameter ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. +Edit `configs/sweep.json` to change ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. + +## Running with Docker Compose + +No local Bun or Python needed. Everything runs in containers. +TODO: fix errors on pip installation -## Configuration +```bash +cd contexto/benchmarks/ama-bench/docker +cp .env.example .env # set API_KEY +docker compose up --build +``` -All mindmap parameters are in `configs/default.json`: +The `bridge` container starts the server, the `runner` container clones AMA-Bench, downloads the dataset, and runs the benchmark. + +### CI + +```yaml +- name: Run AMA-Bench + working-directory: benchmarks/ama-bench/docker + env: + API_KEY: ${{ secrets.API_KEY }} + run: docker compose up --build --abort-on-container-exit +``` + +## Configuration Reference ### Tree construction (`mindmap`) @@ -125,20 +215,31 @@ All mindmap parameters are in `configs/default.json`: ## File Structure ``` -benchmarks/ama-bench/ -├── src/server.ts # Bridge server wrapping @ekai/mindmap +contexto/benchmarks/ama-bench/ # Bridge + config + scripts +├── src/server.ts # Bridge server wrapping @ekai/mindmap ├── package.json +├── tsconfig.json +├── .env # API_KEY (not committed) ├── configs/ -│ ├── default.json # Mindmap + search parameters -│ └── sweep.json # Parameter sweep ranges +│ ├── default.json # Mindmap + search parameters +│ └── sweep.json # Parameter sweep ranges ├── scripts/ -│ ├── setup.sh # One-time setup -│ ├── run.sh # Run benchmark -│ └── sweep.sh # Run parameter sweep +│ ├── setup.sh # One-time setup +│ ├── run.sh # Run benchmark +│ └── sweep.sh # Run parameter sweep ├── docker/ │ ├── docker-compose.yml -│ ├── Dockerfile # AMA-Bench runner -│ ├── bridge.Dockerfile # Bridge server +│ ├── Dockerfile # AMA-Bench runner +│ ├── bridge.Dockerfile # Bridge server │ └── .env.example -└── results/ # Benchmark outputs (gitignored) +└── results/ # Benchmark outputs (gitignored) + +AMA-Bench/ # Fork of AMA-Bench +├── src/method/contexto_method.py # Python method adapter (thin HTTP client) +├── configs/ +│ ├── contexto.yaml # Method config (bridge_url only) +│ ├── openrouter.yaml # LLM config for answer generation +│ └── llm_judge_openrouter.yaml # LLM config for judge scoring +├── dataset/ # Downloaded via huggingface-cli +└── results/ # Benchmark outputs ``` diff --git a/benchmarks/ama-bench/docker/Dockerfile b/benchmarks/ama-bench/docker/Dockerfile index dd1242f..6229dff 100644 --- a/benchmarks/ama-bench/docker/Dockerfile +++ b/benchmarks/ama-bench/docker/Dockerfile @@ -1,7 +1,7 @@ -FROM python:3.11-slim +FROM --platform=linux/amd64 python:3.11-slim RUN apt-get update && apt-get install -y --no-install-recommends \ - git curl && \ + git curl make build-essential && \ rm -rf /var/lib/apt/lists/* WORKDIR /app diff --git a/benchmarks/ama-bench/scripts/run.sh b/benchmarks/ama-bench/scripts/run.sh index 929f646..3e1dc9a 100755 --- a/benchmarks/ama-bench/scripts/run.sh +++ b/benchmarks/ama-bench/scripts/run.sh @@ -6,10 +6,23 @@ BENCH_DIR="$(dirname "$SCRIPT_DIR")" AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" BRIDGE_PORT="${BRIDGE_PORT:-3456}" +# Load .env if present +if [ -f "$BENCH_DIR/.env" ]; then + set -a + source "$BENCH_DIR/.env" + set +a +fi + +# Validate required vars +if [ -z "${API_KEY:-}" ]; then + echo "ERROR: API_KEY not set. Export it or add to .env" + exit 1 +fi + # Parse arguments (pass through to run.py) -LLM_CONFIG="${LLM_CONFIG:-$AMA_BENCH/configs/gpt-4o.yaml}" +LLM_CONFIG="${LLM_CONFIG:-$AMA_BENCH/configs/openrouter.yaml}" SUBSET="${SUBSET:-openend}" -JUDGE_CONFIG="${JUDGE_CONFIG:-$AMA_BENCH/configs/llm_judge.yaml}" +JUDGE_CONFIG="${JUDGE_CONFIG:-$AMA_BENCH/configs/llm_judge_openrouter.yaml}" METHOD_CONFIG="${METHOD_CONFIG:-$AMA_BENCH/configs/contexto.yaml}" EXTRA_ARGS="${@}" @@ -19,7 +32,7 @@ echo "AMA-Bench dir: $AMA_BENCH" # Start bridge server in background echo "[1/3] Starting contexto bridge server on port $BRIDGE_PORT..." cd "$BENCH_DIR" -BRIDGE_PORT=$BRIDGE_PORT bun src/server.ts & +API_KEY=$API_KEY BRIDGE_PORT=$BRIDGE_PORT bun src/server.ts & BRIDGE_PID=$! # Ensure cleanup on exit diff --git a/benchmarks/ama-bench/scripts/setup.sh b/benchmarks/ama-bench/scripts/setup.sh index 0dbf12d..93e0750 100755 --- a/benchmarks/ama-bench/scripts/setup.sh +++ b/benchmarks/ama-bench/scripts/setup.sh @@ -8,17 +8,19 @@ AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" echo "=== AMA-Bench Setup for Contexto ===" echo "AMA-Bench: $AMA_BENCH" -# 1. Clone AMA-Bench if not present -if [ ! -d "$AMA_BENCH" ]; then - echo "[1/4] Cloning AMA-Bench..." - git clone https://github.com/ekailabs/AMA-Bench.git "$AMA_BENCH" -else - echo "[1/4] AMA-Bench already exists at $AMA_BENCH" +# 1. Clone AMA-Bench (delete if exists and clone latest) +if [ -d "$AMA_BENCH" ]; then + echo "[1/4] Removing existing AMA-Bench..." + rm -rf "$AMA_BENCH" fi +echo "[1/4] Cloning AMA-Bench..." +git clone https://github.com/ekailabs/AMA-Bench.git "$AMA_BENCH" # 2. Install Python dependencies echo "[2/4] Installing Python dependencies..." -pip install -r "$AMA_BENCH/requirements.txt" +cd $AMA_BENCH +pip install -r requirements.txt +cd .. # 3. Download dataset if [ ! -d "$AMA_BENCH/dataset" ]; then From 67477fd183df467ebb58a4ff4f8b0eaf9040e2c1 Mon Sep 17 00:00:00 2001 From: DaevMithran Date: Fri, 17 Apr 2026 14:37:05 +0530 Subject: [PATCH 3/6] Refactor --- benchmarks/ama-bench/{docker => }/.env.example | 0 benchmarks/ama-bench/README.md | 3 +-- benchmarks/ama-bench/scripts/run.sh | 2 +- benchmarks/ama-bench/scripts/setup.sh | 12 ++++++------ 4 files changed, 8 insertions(+), 9 deletions(-) rename benchmarks/ama-bench/{docker => }/.env.example (100%) diff --git a/benchmarks/ama-bench/docker/.env.example b/benchmarks/ama-bench/.env.example similarity index 100% rename from benchmarks/ama-bench/docker/.env.example rename to benchmarks/ama-bench/.env.example diff --git a/benchmarks/ama-bench/README.md b/benchmarks/ama-bench/README.md index e7a22de..9075f5b 100644 --- a/benchmarks/ama-bench/README.md +++ b/benchmarks/ama-bench/README.md @@ -160,10 +160,9 @@ bash scripts/sweep.sh Edit `configs/sweep.json` to change ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. -## Running with Docker Compose +## Running with Docker Compose [WIP] No local Bun or Python needed. Everything runs in containers. -TODO: fix errors on pip installation ```bash cd contexto/benchmarks/ama-bench/docker diff --git a/benchmarks/ama-bench/scripts/run.sh b/benchmarks/ama-bench/scripts/run.sh index 3e1dc9a..195014d 100755 --- a/benchmarks/ama-bench/scripts/run.sh +++ b/benchmarks/ama-bench/scripts/run.sh @@ -24,7 +24,7 @@ LLM_CONFIG="${LLM_CONFIG:-$AMA_BENCH/configs/openrouter.yaml}" SUBSET="${SUBSET:-openend}" JUDGE_CONFIG="${JUDGE_CONFIG:-$AMA_BENCH/configs/llm_judge_openrouter.yaml}" METHOD_CONFIG="${METHOD_CONFIG:-$AMA_BENCH/configs/contexto.yaml}" -EXTRA_ARGS="${@}" +EXTRA_ARGS="${*:-}" echo "=== Running AMA-Bench with contexto method ===" echo "AMA-Bench dir: $AMA_BENCH" diff --git a/benchmarks/ama-bench/scripts/setup.sh b/benchmarks/ama-bench/scripts/setup.sh index 93e0750..5e5f6b7 100755 --- a/benchmarks/ama-bench/scripts/setup.sh +++ b/benchmarks/ama-bench/scripts/setup.sh @@ -8,13 +8,13 @@ AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" echo "=== AMA-Bench Setup for Contexto ===" echo "AMA-Bench: $AMA_BENCH" -# 1. Clone AMA-Bench (delete if exists and clone latest) -if [ -d "$AMA_BENCH" ]; then - echo "[1/4] Removing existing AMA-Bench..." - rm -rf "$AMA_BENCH" +# 1. Clone AMA-Bench (if not already present) +if [ ! -d "$AMA_BENCH" ]; then + echo "[1/4] Cloning AMA-Bench..." + git clone https://github.com/ekailabs/AMA-Bench.git "$AMA_BENCH" +else + echo "[1/4] AMA-Bench already exists, skipping clone." fi -echo "[1/4] Cloning AMA-Bench..." -git clone https://github.com/ekailabs/AMA-Bench.git "$AMA_BENCH" # 2. Install Python dependencies echo "[2/4] Installing Python dependencies..." From 74a5266b4ffad9a4eb973bba03e6f4bb237a474e Mon Sep 17 00:00:00 2001 From: DaevMithran Date: Fri, 24 Apr 2026 14:19:14 +0530 Subject: [PATCH 4/6] feat: Create episodic summary for benchmarks --- benchmarks/ama-bench/README.md | 214 ++++++++++-------- benchmarks/ama-bench/configs/default.json | 15 +- benchmarks/ama-bench/configs/hybrid.json | 28 +++ benchmarks/ama-bench/configs/local.json | 28 +++ benchmarks/ama-bench/configs/openai.json | 27 +++ benchmarks/ama-bench/docker/Dockerfile | 22 -- benchmarks/ama-bench/docker/bridge.Dockerfile | 14 -- .../ama-bench/docker/docker-compose.yml | 44 ---- benchmarks/ama-bench/scripts/run.sh | 118 +++++++++- benchmarks/ama-bench/scripts/setup.sh | 18 +- benchmarks/ama-bench/src/episodic/summary.ts | 196 ++++++++++++++++ benchmarks/ama-bench/src/episodic/types.ts | 39 ++++ .../ama-bench/src/episodic/validation.ts | 83 +++++++ benchmarks/ama-bench/src/server.ts | 138 ++++++++++- packages/contexto/package.json | 2 +- 15 files changed, 778 insertions(+), 208 deletions(-) create mode 100644 benchmarks/ama-bench/configs/hybrid.json create mode 100644 benchmarks/ama-bench/configs/local.json create mode 100644 benchmarks/ama-bench/configs/openai.json delete mode 100644 benchmarks/ama-bench/docker/Dockerfile delete mode 100644 benchmarks/ama-bench/docker/bridge.Dockerfile delete mode 100644 benchmarks/ama-bench/docker/docker-compose.yml create mode 100644 benchmarks/ama-bench/src/episodic/summary.ts create mode 100644 benchmarks/ama-bench/src/episodic/types.ts create mode 100644 benchmarks/ama-bench/src/episodic/validation.ts diff --git a/benchmarks/ama-bench/README.md b/benchmarks/ama-bench/README.md index 9075f5b..9c65f0a 100644 --- a/benchmarks/ama-bench/README.md +++ b/benchmarks/ama-bench/README.md @@ -28,11 +28,16 @@ Two repos: ## Prerequisites +Always required: - [Bun](https://bun.sh) >= 1.0 - Python >= 3.9 - pnpm - `huggingface-cli` (`pip install huggingface_hub`) -- An API key for [OpenRouter](https://openrouter.ai) or OpenAI (used for embeddings + LLM) + +Mode-specific (see [Modes](#modes) below): +- **`MODE=local`** — Linux box with 4+ GPUs (≈80GB VRAM total) for VLLM, plus [Ollama](https://ollama.com) for embeddings. No API keys needed. +- **`MODE=openai`** — `OPENAI_API_KEY` only. No GPUs, no Ollama. +- **`MODE=hybrid`** — Ollama for embeddings + `OPENAI_API_KEY` for LLM/judge/summarizer. Mac-friendly, no GPUs. ## Running Locally @@ -66,122 +71,126 @@ pip install -r requirements.txt huggingface-cli download AMA-bench/AMA-bench --repo-type dataset --local-dir dataset ``` -Or run the setup script which does all of the above: +Or run the setup script which does all of the above (plus pulls the Ollama embedding model): ```bash cd contexto/benchmarks/ama-bench bash scripts/setup.sh ``` -### 3. Configure the bridge +### 3. Pick a mode + +`scripts/run.sh` exposes a single `MODE` switch that wires the bridge config + AMA-Bench LLM/judge consistently. Each mode maps to one preset in `configs/`. + +| `MODE` | Embed | Episodic summarizer | Answer-gen + Judge | Needs | +|---|---|---|---|---| +| `local` (default) | Ollama `qwen3-embedding:4b` | VLLM `Qwen/Qwen3-32B` | VLLM `Qwen/Qwen3-32B` | 4+ GPU box, Ollama, VLLM | +| `openai` | OpenAI `text-embedding-3-small` | OpenAI `gpt-4.1-mini` | OpenAI (`gpt-5.2` / `gpt-5.4` judge) | `OPENAI_API_KEY` | +| `hybrid` | Ollama `qwen3-embedding:4b` | OpenAI `gpt-4.1-mini` | OpenAI | Ollama + `OPENAI_API_KEY` | + +The preset files are `configs/{local,openai,hybrid}.json`. Edit them to tune `mindmap.*` / `search.*` or swap models — the bridge picks them up via the `BENCH_CONFIG` env var that `run.sh` sets. -Create `contexto/benchmarks/ama-bench/.env`: +### 3a. (Ollama-using modes) install + start Ollama + +Skip if `MODE=openai`. ```bash -# Embedding API key (used by the bridge for mindmap embeddings) -API_KEY=your-openrouter-or-openai-key +brew install ollama +ollama serve & # keep running in the background +ollama pull qwen3-embedding:4b # ~2.5GB, one-time ``` -Tune mindmap parameters in `contexto/benchmarks/ama-bench/configs/default.json`: - -```json -{ - "provider": "openrouter", - "embedModel": "openai/text-embedding-3-small", - "mindmap": { - "similarityThreshold": 0.5, - "maxDepth": 4, - "maxChildren": 10, - "rebuildInterval": 50 - }, - "search": { - "maxResults": 10, - "maxTokens": 4000, - "beamWidth": 3, - "minScore": 0.0 - } -} -``` +Ollama serves embeddings on `http://localhost:11434` with zero rate limits. -### 4. Configure AMA-Bench LLM +> Note: `@ekai/mindmap`'s package default is still `text-embedding-3-small`. The Ollama embedding is a **benchmark-only override** baked into the preset. -AMA-Bench needs an LLM config for answer generation and a judge config for scoring. Create these in `AMA-Bench/configs/`: +### 3b. (`MODE=local` only) prep VLLM -```yaml -# AMA-Bench/configs/openrouter.yaml -provider: "openai" -api_key: "your-openrouter-key" -model: "openai/gpt-4o" -base_url: "https://openrouter.ai/api/v1" -max_tokens: 16000 -temperature: 0.0 -``` +Skip for `openai` / `hybrid`. + +- Linux box with 4+ GPUs (≈80GB VRAM total — Qwen3-32B at `tensor_parallel_size=4`) +- `pip install vllm` +- Edit GPU IDs in `AMA-Bench/configs/qwen3-32B.yaml` (`vllm_launch.gpus`) -```yaml -# AMA-Bench/configs/llm_judge_openrouter.yaml -provider: "openai" -api_key: "your-openrouter-key" -model: "openai/gpt-4o" -base_url: "https://openrouter.ai/api/v1" -max_tokens: 16000 -temperature: 0.0 +`run.sh` auto-invokes `AMA-Bench/scripts/launch_vllm_32B.sh` (idempotent — skips if already up on the configured port). Answer-gen, judge, and the episodic summarizer all share the one VLLM endpoint. + +If you don't have GPUs locally, common alternatives: Lambda Labs, RunPod, Modal, Replicate, AWS EC2 (g6e.12xlarge), Bedrock (Qwen3-32B dense). Point `episodic.baseUrl` and the AMA-Bench yaml at the remote endpoint. + +### 4. (Cloud-using modes) set your OpenAI key + +Skip if `MODE=local`. + +`.env`: ``` +OPENAI_API_KEY=sk-... +``` + +Both `AMA-Bench/configs/gpt-5.2.yaml` and `llm_judge_api.yaml` are wired to read from `OPENAI_API_KEY` — no need to hard-code keys in yaml. ### 5. Run ```bash cd contexto/benchmarks/ama-bench -bash scripts/run.sh +bash scripts/run.sh # MODE=local (default) + +MODE=openai bash scripts/run.sh # cloud-only +MODE=hybrid bash scripts/run.sh # local embed + cloud LLM ``` -This will: -1. Start the bridge server (reads `configs/default.json` + `API_KEY` from `.env`) -2. Run AMA-Bench with the `contexto` method (208 episodes, ~35s each) -3. Evaluate answers with the LLM judge -4. Save results to `AMA-Bench/results/` -5. Shut down the bridge +`run.sh` will: +1. Validate the dependencies for the chosen mode (Ollama running? OpenAI key present?) +2. Start the bridge with the matching `BENCH_CONFIG` +3. Launch VLLM if `LLM_SERVER=vllm` +4. Run AMA-Bench (208 episodes for `openend`) +5. Save results to `AMA-Bench/results/` and shut down the bridge + +#### Per-knob overrides -Override defaults: +Any individual env var overrides the mode preset, so you can mix: ```bash -LLM_CONFIG=../../../AMA-Bench/configs/openrouter.yaml \ -JUDGE_CONFIG=../../../AMA-Bench/configs/llm_judge_openrouter.yaml \ -SUBSET=openend \ -bash scripts/run.sh +# Local VLLM answer-gen + OpenAI judge for comparison +MODE=local JUDGE_SERVER=api JUDGE_CONFIG=$AMA_BENCH/configs/llm_judge_api.yaml \ + bash scripts/run.sh + +# Hybrid mode but with a custom bridge config +MODE=hybrid BENCH_CONFIG=./configs/my-experiment.json bash scripts/run.sh + +# Disable episodic layer for ablation (point at any preset and edit it, +# or copy a preset and set "episodic": { "enabled": false }) +BENCH_CONFIG=./configs/no-episodic.json bash scripts/run.sh ``` -### 6. Parameter sweep (optional) +Available env vars: `MODE`, `BENCH_CONFIG`, `LLM_SERVER`, `JUDGE_SERVER`, `LLM_CONFIG`, `JUDGE_CONFIG`, `SUBSET`, `BRIDGE_PORT`. -Grid-search over mindmap/search params to find the optimal config: +### 6. Episodic summary layer -```bash -bash scripts/sweep.sh -``` +The bridge runs every turn through an LLM **before** `mindmap.add()`, producing a structured summary (`status`, `summary`, `key_findings`, `evidence_refs`, `open_questions`, `confidence`) that mirrors the production `ekailabs-api-server` behavior. Only the summary (+ formatted key findings) is embedded; raw turn text is kept in `metadata.turn.rawContent` for reference. -Edit `configs/sweep.json` to change ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. +Each preset already wires `episodic` correctly: -## Running with Docker Compose [WIP] +- `local.json` — Qwen3-32B on local VLLM, `jsonMode: false`, `noThink: true` +- `openai.json` / `hybrid.json` — `gpt-4.1-mini` on OpenAI, `jsonMode: true` -No local Bun or Python needed. Everything runs in containers. +Knobs: +- `jsonMode` — set `true` for OpenAI-style `response_format: json_object` (works on OpenAI; on VLLM requires `--guided-decoding-backend outlines`). Set `false` and rely on the system prompt's "JSON only" instruction otherwise. +- `noThink: true` — Qwen3 is a hybrid reasoning model; summarization is factual extraction, so we append `/no_think` to skip the thinking phase (saves tokens + latency). Leave off for non-Qwen3 models. +- `apiKey` — optional when `baseUrl` is localhost (VLLM accepts any value); otherwise reads `episodic.apiKey` → `API_KEY` → `OPENAI_API_KEY`. -```bash -cd contexto/benchmarks/ama-bench/docker -cp .env.example .env # set API_KEY -docker compose up --build -``` +**Toggle off for ablation:** set `"episodic": { "enabled": false }` in your bench config; raw turn content is embedded directly. + +**Cost / volume (`MODE=openai` or `hybrid`):** one `gpt-4.1-mini` call per turn. Full `openend` subset (208 episodes × ~30–100 turns) ≈ 10k–20k summarization calls — pricing varies, check OpenAI's current rate. `MODE=local` is free apart from compute. -The `bridge` container starts the server, the `runner` container clones AMA-Bench, downloads the dataset, and runs the benchmark. +### 7. Parameter sweep (optional) -### CI +Grid-search over mindmap/search params to find the optimal config: -```yaml -- name: Run AMA-Bench - working-directory: benchmarks/ama-bench/docker - env: - API_KEY: ${{ secrets.API_KEY }} - run: docker compose up --build --abort-on-container-exit +```bash +bash scripts/sweep.sh ``` +Edit `configs/sweep.json` to change ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. + ## Configuration Reference ### Tree construction (`mindmap`) @@ -214,31 +223,36 @@ The `bridge` container starts the server, the `runner` container clones AMA-Benc ## File Structure ``` -contexto/benchmarks/ama-bench/ # Bridge + config + scripts -├── src/server.ts # Bridge server wrapping @ekai/mindmap +contexto/benchmarks/ama-bench/ # Bridge + config + scripts +├── src/ +│ ├── server.ts # Bridge server wrapping @ekai/mindmap +│ └── episodic/ # Production-parity summary layer +│ ├── summary.ts +│ ├── types.ts +│ └── validation.ts ├── package.json ├── tsconfig.json -├── .env # API_KEY (not committed) +├── .env # OPENAI_API_KEY (not committed) ├── configs/ -│ ├── default.json # Mindmap + search parameters -│ └── sweep.json # Parameter sweep ranges +│ ├── default.json # Loaded if BENCH_CONFIG is unset +│ ├── local.json # MODE=local preset +│ ├── openai.json # MODE=openai preset +│ ├── hybrid.json # MODE=hybrid preset +│ └── sweep.json # Parameter sweep ranges ├── scripts/ -│ ├── setup.sh # One-time setup -│ ├── run.sh # Run benchmark -│ └── sweep.sh # Run parameter sweep -├── docker/ -│ ├── docker-compose.yml -│ ├── Dockerfile # AMA-Bench runner -│ ├── bridge.Dockerfile # Bridge server -│ └── .env.example -└── results/ # Benchmark outputs (gitignored) - -AMA-Bench/ # Fork of AMA-Bench -├── src/method/contexto_method.py # Python method adapter (thin HTTP client) +│ ├── setup.sh # One-time setup +│ ├── run.sh # Run benchmark (MODE-aware) +│ └── sweep.sh # Run parameter sweep +└── results/ # Benchmark outputs (gitignored) + +AMA-Bench/ # Sibling repo +├── src/method/contexto_method.py # Python method adapter (thin HTTP client) ├── configs/ -│ ├── contexto.yaml # Method config (bridge_url only) -│ ├── openrouter.yaml # LLM config for answer generation -│ └── llm_judge_openrouter.yaml # LLM config for judge scoring -├── dataset/ # Downloaded via huggingface-cli -└── results/ # Benchmark outputs +│ ├── contexto.yaml # Method config (bridge_url only) +│ ├── qwen3-32B.yaml # VLLM answer-gen +│ ├── llm_judge.yaml # VLLM judge +│ ├── gpt-5.2.yaml # OpenAI answer-gen +│ └── llm_judge_api.yaml # OpenAI judge +├── dataset/ # Downloaded via huggingface-cli +└── results/ # Benchmark outputs ``` diff --git a/benchmarks/ama-bench/configs/default.json b/benchmarks/ama-bench/configs/default.json index e631169..c190e54 100644 --- a/benchmarks/ama-bench/configs/default.json +++ b/benchmarks/ama-bench/configs/default.json @@ -1,6 +1,17 @@ { - "provider": "openrouter", - "embedModel": "openai/text-embedding-3-small", + "embed": { + "type": "ollama", + "baseUrl": "http://localhost:11434", + "model": "qwen3-embedding:4b" + }, + "episodic": { + "enabled": true, + "baseUrl": "http://localhost:8002/v1/chat/completions", + "model": "Qwen/Qwen3-32B", + "temperature": 0.2, + "jsonMode": false, + "noThink": true + }, "mindmap": { "similarityThreshold": 0.5, "maxDepth": 4, diff --git a/benchmarks/ama-bench/configs/hybrid.json b/benchmarks/ama-bench/configs/hybrid.json new file mode 100644 index 0000000..e1be8f4 --- /dev/null +++ b/benchmarks/ama-bench/configs/hybrid.json @@ -0,0 +1,28 @@ +{ + "_comment": "MODE=hybrid — local Ollama embeddings (zero rate limits) + OpenAI for episodic/answer-gen/judge. No GPUs required. Requires Ollama + OPENAI_API_KEY in .env.", + "embed": { + "type": "ollama", + "baseUrl": "http://localhost:11434", + "model": "qwen3-embedding:4b" + }, + "episodic": { + "enabled": true, + "baseUrl": "https://api.openai.com/v1/chat/completions", + "model": "gpt-4.1-mini", + "temperature": 0.2, + "jsonMode": true, + "noThink": false + }, + "mindmap": { + "similarityThreshold": 0.5, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 4000, + "beamWidth": 3, + "minScore": 0.0 + } +} diff --git a/benchmarks/ama-bench/configs/local.json b/benchmarks/ama-bench/configs/local.json new file mode 100644 index 0000000..58cd40e --- /dev/null +++ b/benchmarks/ama-bench/configs/local.json @@ -0,0 +1,28 @@ +{ + "_comment": "MODE=local — fully self-hosted. Embed via Ollama, episodic + answer-gen + judge via local Qwen3-32B VLLM. Requires GPU box with 4+ GPUs (~80GB VRAM) and Ollama running locally. No API keys needed.", + "embed": { + "type": "ollama", + "baseUrl": "http://localhost:11434", + "model": "qwen3-embedding:4b" + }, + "episodic": { + "enabled": true, + "baseUrl": "http://localhost:8002/v1/chat/completions", + "model": "Qwen/Qwen3-32B", + "temperature": 0.2, + "jsonMode": false, + "noThink": true + }, + "mindmap": { + "similarityThreshold": 0.5, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 4000, + "beamWidth": 3, + "minScore": 0.0 + } +} diff --git a/benchmarks/ama-bench/configs/openai.json b/benchmarks/ama-bench/configs/openai.json new file mode 100644 index 0000000..d4a4ff7 --- /dev/null +++ b/benchmarks/ama-bench/configs/openai.json @@ -0,0 +1,27 @@ +{ + "_comment": "MODE=openai — fully cloud. Embed + episodic + answer-gen + judge all hit OpenAI. Slowest + most expensive but no GPUs required. Requires OPENAI_API_KEY in .env.", + "embed": { + "type": "openai", + "model": "text-embedding-3-small" + }, + "episodic": { + "enabled": true, + "baseUrl": "https://api.openai.com/v1/chat/completions", + "model": "gpt-4.1-mini", + "temperature": 0.2, + "jsonMode": true, + "noThink": false + }, + "mindmap": { + "similarityThreshold": 0.5, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 4000, + "beamWidth": 3, + "minScore": 0.0 + } +} diff --git a/benchmarks/ama-bench/docker/Dockerfile b/benchmarks/ama-bench/docker/Dockerfile deleted file mode 100644 index 6229dff..0000000 --- a/benchmarks/ama-bench/docker/Dockerfile +++ /dev/null @@ -1,22 +0,0 @@ -FROM --platform=linux/amd64 python:3.11-slim - -RUN apt-get update && apt-get install -y --no-install-recommends \ - git curl make build-essential && \ - rm -rf /var/lib/apt/lists/* - -WORKDIR /app - -# Clone AMA-Bench fork (contexto_method.py and configs already in repo) -ARG AMA_BENCH_REPO=https://github.com/ekailabs/AMA-Bench.git -RUN git clone ${AMA_BENCH_REPO} /app/AMA-Bench - -# Install Python dependencies -RUN pip install --no-cache-dir -r /app/AMA-Bench/requirements.txt - -# Download dataset -RUN pip install --no-cache-dir huggingface_hub && \ - huggingface-cli download AMA-bench/AMA-bench --repo-type dataset --local-dir /app/AMA-Bench/dataset - -WORKDIR /app/AMA-Bench - -ENTRYPOINT ["python", "src/run.py"] diff --git a/benchmarks/ama-bench/docker/bridge.Dockerfile b/benchmarks/ama-bench/docker/bridge.Dockerfile deleted file mode 100644 index aea0c9d..0000000 --- a/benchmarks/ama-bench/docker/bridge.Dockerfile +++ /dev/null @@ -1,14 +0,0 @@ -FROM oven/bun:1-slim - -WORKDIR /app - -COPY package.json ./ -COPY src/ ./src/ - -# Install @ekai/mindmap from npm (swap workspace ref) -RUN sed -i 's/"workspace:\*"/"latest"/' package.json && \ - bun install - -EXPOSE 3456 - -CMD ["bun", "src/server.ts"] diff --git a/benchmarks/ama-bench/docker/docker-compose.yml b/benchmarks/ama-bench/docker/docker-compose.yml deleted file mode 100644 index c879954..0000000 --- a/benchmarks/ama-bench/docker/docker-compose.yml +++ /dev/null @@ -1,44 +0,0 @@ -services: - bridge: - build: - context: .. - dockerfile: docker/bridge.Dockerfile - ports: - - "3456:3456" - environment: - - BRIDGE_PORT=3456 - - API_KEY=${API_KEY} - healthcheck: - test: ["CMD", "curl", "-sf", "http://localhost:3456/health"] - interval: 5s - timeout: 3s - retries: 10 - - runner: - build: - context: . - dockerfile: Dockerfile - depends_on: - bridge: - condition: service_healthy - environment: - - CONTEXTO_BRIDGE_URL=http://bridge:3456 - command: - - --llm-server - - api - - --llm-config - - configs/${LLM_CONFIG:-gpt-5.2.yaml} - - --subset - - ${SUBSET:-openend} - - --method - - contexto - - --method-config - - configs/contexto.yaml - - --test-dir - - dataset/test - - --judge-config - - configs/${JUDGE_CONFIG:-llm_judge.yaml} - - --evaluate - - "True" - volumes: - - ../results:/app/AMA-Bench/results diff --git a/benchmarks/ama-bench/scripts/run.sh b/benchmarks/ama-bench/scripts/run.sh index 195014d..7a261ea 100755 --- a/benchmarks/ama-bench/scripts/run.sh +++ b/benchmarks/ama-bench/scripts/run.sh @@ -6,33 +6,121 @@ BENCH_DIR="$(dirname "$SCRIPT_DIR")" AMA_BENCH="$(cd "$BENCH_DIR/../../.." && pwd)/AMA-Bench" BRIDGE_PORT="${BRIDGE_PORT:-3456}" -# Load .env if present +# Load .env if present (optional — judge/LLM api keys live in AMA-Bench configs) if [ -f "$BENCH_DIR/.env" ]; then set -a source "$BENCH_DIR/.env" set +a fi -# Validate required vars -if [ -z "${API_KEY:-}" ]; then - echo "ERROR: API_KEY not set. Export it or add to .env" - exit 1 -fi +# MODE picks coherent defaults across the bridge config + AMA-Bench LLM/judge. +# local → Ollama embed + VLLM (Qwen3-32B) for episodic + answer-gen + judge. No API keys, requires GPUs. +# openai → OpenAI for everything (embed + episodic + answer-gen + judge). No GPUs, costs $$. +# hybrid → Ollama embed (no rate limits) + OpenAI for episodic + answer-gen + judge. No GPUs. +# Any individual env var below can override the MODE preset. +MODE="${MODE:-local}" + +case "$MODE" in + local) + DEFAULT_BENCH_CONFIG="$BENCH_DIR/configs/local.json" + DEFAULT_LLM_SERVER="vllm" + DEFAULT_JUDGE_SERVER="vllm" + DEFAULT_LLM_CONFIG="$AMA_BENCH/configs/qwen3-32B.yaml" + DEFAULT_JUDGE_CONFIG="$AMA_BENCH/configs/llm_judge.yaml" + ;; + openai) + DEFAULT_BENCH_CONFIG="$BENCH_DIR/configs/openai.json" + DEFAULT_LLM_SERVER="api" + DEFAULT_JUDGE_SERVER="api" + DEFAULT_LLM_CONFIG="$AMA_BENCH/configs/gpt-5.2.yaml" + DEFAULT_JUDGE_CONFIG="$AMA_BENCH/configs/llm_judge_api.yaml" + ;; + hybrid) + DEFAULT_BENCH_CONFIG="$BENCH_DIR/configs/hybrid.json" + DEFAULT_LLM_SERVER="api" + DEFAULT_JUDGE_SERVER="api" + DEFAULT_LLM_CONFIG="$AMA_BENCH/configs/gpt-5.2.yaml" + DEFAULT_JUDGE_CONFIG="$AMA_BENCH/configs/llm_judge_api.yaml" + ;; + *) + echo "ERROR: unknown MODE='$MODE'. Valid: local | openai | hybrid" + exit 1 + ;; +esac -# Parse arguments (pass through to run.py) -LLM_CONFIG="${LLM_CONFIG:-$AMA_BENCH/configs/openrouter.yaml}" +# Per-var overrides (caller can mix and match) +BENCH_CONFIG="${BENCH_CONFIG:-$DEFAULT_BENCH_CONFIG}" +LLM_SERVER="${LLM_SERVER:-$DEFAULT_LLM_SERVER}" +JUDGE_SERVER="${JUDGE_SERVER:-$DEFAULT_JUDGE_SERVER}" +LLM_CONFIG="${LLM_CONFIG:-$DEFAULT_LLM_CONFIG}" +JUDGE_CONFIG="${JUDGE_CONFIG:-$DEFAULT_JUDGE_CONFIG}" SUBSET="${SUBSET:-openend}" -JUDGE_CONFIG="${JUDGE_CONFIG:-$AMA_BENCH/configs/llm_judge_openrouter.yaml}" METHOD_CONFIG="${METHOD_CONFIG:-$AMA_BENCH/configs/contexto.yaml}" EXTRA_ARGS="${*:-}" +echo "MODE=$MODE bridge=$BENCH_CONFIG llm=$LLM_SERVER:$LLM_CONFIG judge=$JUDGE_SERVER:$JUDGE_CONFIG" + echo "=== Running AMA-Bench with contexto method ===" echo "AMA-Bench dir: $AMA_BENCH" +# Pre-flight: detect configured embed.type from chosen bench config +EMBED_TYPE="$(python3 -c "import json; print(json.load(open('$BENCH_CONFIG')).get('embed',{}).get('type','ollama'))" 2>/dev/null || echo "ollama")" + +if [ "$EMBED_TYPE" = "ollama" ]; then + OLLAMA_URL="${OLLAMA_URL:-http://localhost:11434}" + EMBED_MODEL_HINT="qwen3-embedding" + if ! curl -sf "$OLLAMA_URL/api/tags" >/dev/null; then + echo "ERROR: Ollama not reachable at $OLLAMA_URL" + echo " Install: brew install ollama" + echo " Start: ollama serve &" + exit 1 + fi + if ! curl -sf "$OLLAMA_URL/api/tags" | grep -q "$EMBED_MODEL_HINT"; then + echo "ERROR: embedding model not found in Ollama" + echo " Pull: ollama pull qwen3-embedding:0.6b" + exit 1 + fi + echo "Ollama OK ($OLLAMA_URL) — $EMBED_MODEL_HINT present" +else + # Cloud provider — bridge will read apiKey from config or API_KEY env var + if [ -z "${API_KEY:-}" ]; then + # not fatal — embed.apiKey in configs/default.json may be set directly + echo "NOTE: embed.type=$EMBED_TYPE — bridge will use embed.apiKey from configs/default.json or API_KEY env var" + else + echo "Embed: $EMBED_TYPE (using API_KEY from environment)" + fi +fi + +# Episodic summary layer: only require an API key when the baseUrl is NOT localhost +EPISODIC_ENABLED="$(python3 -c "import json; print(str(json.load(open('$BENCH_CONFIG')).get('episodic',{}).get('enabled', True)).lower())" 2>/dev/null || echo "true")" +EPISODIC_BASEURL="$(python3 -c "import json; print(json.load(open('$BENCH_CONFIG')).get('episodic',{}).get('baseUrl',''))" 2>/dev/null || echo "")" +if [ "$EPISODIC_ENABLED" = "true" ]; then + case "$EPISODIC_BASEURL" in + *localhost*|*127.0.0.1*|*0.0.0.0*) EPISODIC_LOCAL=1 ;; + *) EPISODIC_LOCAL=0 ;; + esac + if [ "$EPISODIC_LOCAL" = "0" ] && [ -z "${API_KEY:-}" ] && [ -z "${OPENAI_API_KEY:-}" ]; then + echo "ERROR: episodic baseUrl=$EPISODIC_BASEURL needs API_KEY or OPENAI_API_KEY" + echo " Set one in .env, switch to MODE=local, or disable the layer (episodic.enabled=false)." + exit 1 + fi + echo "Episodic: enabled (baseUrl=$EPISODIC_BASEURL)" +else + echo "Episodic: disabled (raw turn content will be embedded directly)" +fi + +# OpenAI key required when LLM_SERVER or JUDGE_SERVER hits the API +if [ "$LLM_SERVER" = "api" ] || [ "$JUDGE_SERVER" = "api" ]; then + if [ -z "${OPENAI_API_KEY:-}" ] && [ -z "${API_KEY:-}" ]; then + echo "ERROR: LLM_SERVER/JUDGE_SERVER=api requires OPENAI_API_KEY in .env" + exit 1 + fi +fi + # Start bridge server in background echo "[1/3] Starting contexto bridge server on port $BRIDGE_PORT..." cd "$BENCH_DIR" -API_KEY=$API_KEY BRIDGE_PORT=$BRIDGE_PORT bun src/server.ts & +BRIDGE_PORT=$BRIDGE_PORT BENCH_CONFIG="$BENCH_CONFIG" API_KEY="${API_KEY:-}" OPENAI_API_KEY="${OPENAI_API_KEY:-}" bun src/server.ts & BRIDGE_PID=$! # Ensure cleanup on exit @@ -58,17 +146,25 @@ for i in $(seq 1 30); do sleep 1 done +# If using VLLM for answer generation, launch it via AMA-Bench's helper +if [ "$LLM_SERVER" = "vllm" ]; then + echo "[1b/3] Launching VLLM answer-gen server (config: $LLM_CONFIG)..." + cd "$AMA_BENCH" + bash scripts/launch_vllm_32B.sh "$LLM_CONFIG" +fi + # Run AMA-Bench echo "[2/3] Running AMA-Bench evaluation..." cd "$AMA_BENCH" python src/run.py \ - --llm-server api \ + --llm-server "$LLM_SERVER" \ --llm-config "$LLM_CONFIG" \ --subset "$SUBSET" \ --method contexto \ --method-config "$METHOD_CONFIG" \ --test-dir dataset/test \ --judge-config "$JUDGE_CONFIG" \ + --judge-server "$JUDGE_SERVER" \ --evaluate True \ $EXTRA_ARGS diff --git a/benchmarks/ama-bench/scripts/setup.sh b/benchmarks/ama-bench/scripts/setup.sh index 5e5f6b7..d737c77 100755 --- a/benchmarks/ama-bench/scripts/setup.sh +++ b/benchmarks/ama-bench/scripts/setup.sh @@ -31,12 +31,24 @@ else fi # 4. Install bridge dependencies -echo "[4/4] Installing bridge dependencies..." +echo "[4/5] Installing bridge dependencies..." cd "$BENCH_DIR" && pnpm install +# 5. Pull Ollama embedding model (optional — won't hard-fail) +echo "[5/5] Checking Ollama..." +if command -v ollama >/dev/null 2>&1; then + echo "Pulling qwen3-embedding:0.6b (idempotent)..." + ollama pull qwen3-embedding:0.6b || echo "WARN: ollama pull failed — run 'ollama serve &' first, then retry." +else + echo "Ollama not installed. To enable local embeddings:" + echo " brew install ollama" + echo " ollama serve &" + echo " ollama pull qwen3-embedding:0.6b" +fi + echo "" echo "=== Setup complete ===" echo "Next steps:" -echo " 1. Edit configs/default.json (provider, embedModel, mindmap params)" -echo " 2. Export API_KEY=your-api-key" +echo " 1. Ensure 'ollama serve' is running and qwen3-embedding:0.6b is pulled" +echo " 2. (Optional) tune mindmap params in configs/default.json" echo " 3. Run: bash scripts/run.sh" diff --git a/benchmarks/ama-bench/src/episodic/summary.ts b/benchmarks/ama-bench/src/episodic/summary.ts new file mode 100644 index 0000000..1046865 --- /dev/null +++ b/benchmarks/ama-bench/src/episodic/summary.ts @@ -0,0 +1,196 @@ +// Ported from ekailabs-api-server/src/services/api/summary.service.ts +// +// Adapted for the benchmark: +// - Plain factory (no singleton, no env-driven init) +// - Takes baseUrl/model/apiKey from config instead of OpenRouter hardcoded +// - summarizeTurn(item, episodeId, turnIndex) replaces summarizeEpisode +// - episode metadata shape replaced with simpler `turn` shape +// +// SYSTEM_PROMPT is kept byte-identical to production so benchmark results +// reflect the same LLM behavior that ships in ekailabs-api-server. + +import type { EpisodeSummary, TurnData } from './types.js'; +import { validateSummary, toDegradedSummary } from './validation.js'; + +const SYSTEM_PROMPT = `You are a concise summarizer. Given a conversation episode (user question + assistant answer + tool outputs), produce a JSON object with exactly these fields: + +{ + "status": "complete" | "partial" | "blocked", + "summary": "", + "key_findings": ["", "", ...], + "evidence_refs": [{"type": "", "value": ""}], + "open_questions": [""], + "confidence": <0.0 to 1.0> +} + +Rules: +- Set status to "complete" if the episode fully resolved the user's request, "partial" if only partly, "blocked" if unable to proceed. +- summary should be 1-3 sentences capturing the essence. +- key_findings should have at least one entry. +- evidence_refs should reference relevant tools, files, or episodes mentioned. +- Respond ONLY with valid JSON, no markdown fences, no extra text.`; + +export interface SummarizerConfig { + baseUrl: string; // e.g. https://api.openai.com/v1/chat/completions OR http://localhost:8002/v1/chat/completions + model: string; // e.g. gpt-4o-mini, Qwen/Qwen3-32B + apiKey?: string; // optional for local VLLM (any value accepted) + temperature?: number; // default 0.2 + jsonMode?: boolean; // default true — disable for VLLM backends without guided decoding + noThink?: boolean; // default false — if true, appends '/no_think' to user message (Qwen3 hybrid-reasoning models only) +} + +export interface ConversationItemInput { + id: string; + role: string; + content: string; + metadata?: Record; +} + +export interface ConversationItemOutput { + id: string; + role: string; + content: string; + embedding: number[]; + timestamp: string; + metadata: Record; +} + +export interface Summarizer { + summarizeTurn( + item: ConversationItemInput, + episodeId: string, + turnIndex: number, + ): Promise; +} + +export function createSummarizer(config: SummarizerConfig): Summarizer { + const temperature = config.temperature ?? 0.2; + const jsonMode = config.jsonMode ?? true; + const noThink = config.noThink ?? false; + const apiKey = config.apiKey ?? 'EMPTY'; // VLLM accepts any token; OpenAI requires a real key + console.log( + `[episodic] initialized model=${config.model} baseUrl=${config.baseUrl} jsonMode=${jsonMode} noThink=${noThink}`, + ); + + async function callLLM(content: string): Promise { + // For Qwen3 hybrid-reasoning models, append /no_think to skip the thinking phase + const userContent = noThink ? `${content}\n\n/no_think` : content; + const body: Record = { + model: config.model, + messages: [ + { role: 'system', content: SYSTEM_PROMPT }, + { role: 'user', content: userContent }, + ], + temperature, + }; + if (jsonMode) body.response_format = { type: 'json_object' }; + + const response = await fetch(config.baseUrl, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify(body), + }); + + if (!response.ok) { + const body = await response.text().catch(() => '(no body)'); + throw new Error(`summarizer HTTP ${response.status}: ${body.slice(0, 200)}`); + } + + const data = (await response.json()) as { + choices?: Array<{ message?: { content?: string } }>; + }; + const text = data?.choices?.[0]?.message?.content; + if (!text) throw new Error('Empty LLM response'); + + // Strip Qwen3 reasoning blocks if present, then markdown code fences + const stripped = text + .replace(/[\s\S]*?<\/think>/g, '') + .trim() + .replace(/^```(?:json)?\s*/i, '') + .replace(/\s*```$/, '') + .trim(); + return JSON.parse(stripped); + } + + return { + async summarizeTurn( + item: ConversationItemInput, + episodeId: string, + turnIndex: number, + ): Promise { + const traceRef = crypto.randomUUID(); + const turn: TurnData = { + episodeId, + turnIndex, + role: item.role, + rawContent: item.content, + }; + + let summary: EpisodeSummary; + try { + const raw = await callLLM(item.content); + const { valid, failures } = validateSummary(raw); + + if (valid) { + const r = raw as { + summary: string; + key_findings: string[]; + status: EpisodeSummary['metadata']['status']; + evidence_refs: EpisodeSummary['metadata']['evidence_refs']; + open_questions?: string[]; + confidence?: number; + }; + summary = { + id: crypto.randomUUID(), + summary: r.summary, + key_findings: r.key_findings, + metadata: { + status: r.status, + evidence_refs: r.evidence_refs, + open_questions: r.open_questions, + confidence: r.confidence, + trace_ref: traceRef, + turn, + }, + timestamp: new Date().toISOString(), + }; + } else { + summary = toDegradedSummary(raw, failures, traceRef, turn); + } + } catch (err) { + console.error( + '[episodic] LLM call failed:', + err instanceof Error ? err.message : err, + ); + summary = toDegradedSummary(null, ['llm_call_failed'], traceRef, turn); + } + + // Build ConversationItem for the mindmap — identical shape to production: + // content = summary + "\nKey findings:\n- ..." + const contentParts = [summary.summary]; + if (summary.key_findings.length > 0) { + contentParts.push( + `\nKey findings:\n${summary.key_findings.map((f) => `- ${f}`).join('\n')}`, + ); + } + + return { + id: summary.id, + role: 'assistant', + content: contentParts.join('\n'), + embedding: [], + timestamp: summary.timestamp, + metadata: { + source: 'summary', + ...summary.metadata, + // Preserve original item id + any pre-existing metadata (e.g. turnIndex from Python) + original_id: item.id, + original_metadata: item.metadata ?? {}, + }, + }; + }, + }; +} diff --git a/benchmarks/ama-bench/src/episodic/types.ts b/benchmarks/ama-bench/src/episodic/types.ts new file mode 100644 index 0000000..2f8321f --- /dev/null +++ b/benchmarks/ama-bench/src/episodic/types.ts @@ -0,0 +1,39 @@ +// Ported from ekailabs-api-server/src/types/summary.ts +// Simplified: production's episode.{userMessage,assistantMessages,toolMessages} +// is replaced with a flat `turn` shape appropriate for a single trajectory turn. + +export type EvidenceRefType = 'episode_ref' | 'tool_ref' | 'file_ref' | 'trace_ref'; + +export interface EvidenceRef { + type: EvidenceRefType; + value: string; +} + +export interface TurnData { + episodeId: string; + turnIndex: number; + role: string; + rawContent: string; +} + +export interface EpisodeSummaryMetadata { + status: 'complete' | 'partial' | 'blocked'; + evidence_refs: EvidenceRef[]; + open_questions?: string[]; + confidence?: number; + trace_ref: string; + turn: TurnData; +} + +export interface EpisodeSummary { + id: string; + summary: string; + key_findings: string[]; + metadata: EpisodeSummaryMetadata; + timestamp: string; +} + +export interface ValidationResult { + valid: boolean; + failures: string[]; +} diff --git a/benchmarks/ama-bench/src/episodic/validation.ts b/benchmarks/ama-bench/src/episodic/validation.ts new file mode 100644 index 0000000..a9e068f --- /dev/null +++ b/benchmarks/ama-bench/src/episodic/validation.ts @@ -0,0 +1,83 @@ +// Ported from ekailabs-api-server/src/helpers/summary-validation.ts +// Semantics are identical; only the `episodeData` parameter type has changed +// (now the simpler `TurnData` shape from types.ts). + +import type { EpisodeSummary, TurnData, ValidationResult } from './types.js'; + +const VALID_STATUSES = new Set(['complete', 'partial', 'blocked']); + +export function validateSummary(raw: unknown): ValidationResult { + const failures: string[] = []; + const obj = raw as Record | null; + + if (!obj || typeof obj !== 'object') { + return { valid: false, failures: ['root (not an object)'] }; + } + + // status + if (!obj.status || !VALID_STATUSES.has(obj.status as string)) { + failures.push('status'); + } + + // summary + if (typeof obj.summary !== 'string' || (obj.summary as string).trim() === '') { + failures.push('summary'); + } + + // key_findings + if ( + !Array.isArray(obj.key_findings) || + obj.key_findings.length === 0 || + !obj.key_findings.some((f: unknown) => typeof f === 'string' && f.trim() !== '') + ) { + failures.push('key_findings'); + } + + // evidence_refs + if (!Array.isArray(obj.evidence_refs)) { + failures.push('evidence_refs'); + } + + // trace_ref is injected after the LLM call, not validated here + + return { valid: failures.length === 0, failures }; +} + +export function toDegradedSummary( + raw: unknown, + failures: string[], + traceRef: string, + turn: TurnData, +): EpisodeSummary { + const obj = (raw && typeof raw === 'object' ? raw : {}) as Record; + + const summary = + typeof obj.summary === 'string' && (obj.summary as string).trim() ? (obj.summary as string) : ''; + const key_findings = Array.isArray(obj.key_findings) + ? (obj.key_findings as unknown[]).filter((f): f is string => typeof f === 'string') + : []; + const evidence_refs = Array.isArray(obj.evidence_refs) + ? (obj.evidence_refs as EpisodeSummary['metadata']['evidence_refs']) + : []; + + console.warn( + `[episodic] degraded summary — failed fields: ${failures.join(', ')}, traceRef: ${traceRef}`, + ); + + return { + id: crypto.randomUUID(), + summary, + key_findings, + metadata: { + status: 'partial', + evidence_refs, + open_questions: Array.isArray(obj.open_questions) + ? (obj.open_questions as string[]) + : undefined, + confidence: typeof obj.confidence === 'number' ? (obj.confidence as number) : undefined, + trace_ref: traceRef, + turn, + }, + timestamp: new Date().toISOString(), + }; +} diff --git a/benchmarks/ama-bench/src/server.ts b/benchmarks/ama-bench/src/server.ts index d321ec9..f0bef78 100644 --- a/benchmarks/ama-bench/src/server.ts +++ b/benchmarks/ama-bench/src/server.ts @@ -1,6 +1,9 @@ import { + createEmbedFn, createMindmap, memoryStorage, + type EmbedFn, + type EmbedProvider, type Mindmap, type MindmapConfig, type SearchOptions, @@ -8,21 +11,117 @@ import { import { readFileSync } from 'fs'; import { resolve, dirname } from 'path'; import { fileURLToPath } from 'url'; +import { createSummarizer, type Summarizer } from './episodic/summary.js'; // --- Config --- const __dirname = dirname(fileURLToPath(import.meta.url)); -const configPath = resolve(__dirname, '../configs/default.json'); +// BENCH_CONFIG env var lets run.sh point us at a preset (local.json / openai.json / hybrid.json) +// without mutating default.json. +const configPath = process.env.BENCH_CONFIG + ? resolve(process.env.BENCH_CONFIG) + : resolve(__dirname, '../configs/default.json'); +console.log(`[bridge] Loading config from ${configPath}`); const config = JSON.parse(readFileSync(configPath, 'utf-8')); -const provider = config.provider ?? 'openrouter'; -const embedModel = config.embedModel ?? 'openai/text-embedding-3-small'; -const apiKey = process.env.API_KEY ?? ''; +interface EmbedConfig { + type?: string; // 'ollama' | 'openai' | 'openrouter' | 'gemini' + baseUrl?: string; // ollama only + model?: string; // model name (provider-specific) + apiKey?: string; // cloud providers (falls back to API_KEY env var) +} + +const embedCfg: EmbedConfig = config.embed ?? {}; +const embedType = embedCfg.type ?? 'ollama'; const mindmapConfig: Partial = config.mindmap ?? {}; const searchDefaults: SearchOptions = config.search ?? {}; -if (!apiKey) { - console.warn('[bridge] WARNING: API_KEY env var not set. Embedding calls will fail.'); +// --- Embed functions --- + +function makeOllamaEmbedFn({ baseUrl, model }: { baseUrl: string; model: string }): EmbedFn { + return async (text: string): Promise => { + const resp = await fetch(`${baseUrl}/api/embeddings`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ model, prompt: text }), + }); + if (!resp.ok) { + const body = await resp.text().catch(() => '(no body)'); + throw new Error(`ollama embed failed: ${resp.status} ${body}`); + } + const json = (await resp.json()) as { embedding: number[] }; + return json.embedding; + }; +} + +function buildEmbedFn(): EmbedFn { + if (embedType === 'ollama') { + const baseUrl = embedCfg.baseUrl ?? 'http://localhost:11434'; + const model = embedCfg.model ?? 'qwen3-embedding:0.6b'; + console.log(`[bridge] Embed: ollama ${baseUrl} model=${model}`); + return makeOllamaEmbedFn({ baseUrl, model }); + } + + // Cloud providers — delegate to @ekai/mindmap's built-in embed client + if (embedType === 'openai' || embedType === 'openrouter' || embedType === 'gemini') { + const apiKey = embedCfg.apiKey ?? process.env.API_KEY ?? ''; + if (!apiKey) { + throw new Error( + `embed.type='${embedType}' requires an apiKey (set embed.apiKey in configs/default.json or API_KEY env var).`, + ); + } + console.log(`[bridge] Embed: ${embedType} model=${embedCfg.model ?? '(provider default)'}`); + return createEmbedFn({ + provider: embedType as EmbedProvider, + apiKey, + model: embedCfg.model, + }); + } + + throw new Error( + `Unsupported embed.type '${embedType}'. Supported: 'ollama' | 'openai' | 'openrouter' | 'gemini'.`, + ); +} + +const embedFn = buildEmbedFn(); + +// --- Episodic summary layer --- + +interface EpisodicConfig { + enabled?: boolean; + baseUrl?: string; + model?: string; + apiKey?: string; + temperature?: number; + jsonMode?: boolean; + noThink?: boolean; +} + +const episodicCfg: EpisodicConfig = config.episodic ?? {}; +const episodicEnabled = episodicCfg.enabled !== false; // default on + +let summarizer: Summarizer | null = null; +if (episodicEnabled) { + const baseUrl = episodicCfg.baseUrl ?? 'https://api.openai.com/v1/chat/completions'; + // Local endpoints (VLLM/Ollama/etc.) don't need a real API key + const isLocal = /^https?:\/\/(localhost|127\.0\.0\.1|0\.0\.0\.0|\[::1?\])/i.test(baseUrl); + const apiKey = + episodicCfg.apiKey ?? process.env.API_KEY ?? process.env.OPENAI_API_KEY ?? ''; + if (!apiKey && !isLocal) { + throw new Error( + 'episodic.enabled=true but no API key found (set episodic.apiKey in configs/default.json, or API_KEY / OPENAI_API_KEY env var).', + ); + } + summarizer = createSummarizer({ + baseUrl, + model: episodicCfg.model ?? 'gpt-4o-mini', + apiKey: apiKey || undefined, + temperature: episodicCfg.temperature ?? 0.2, + jsonMode: episodicCfg.jsonMode, + noThink: episodicCfg.noThink, + }); +} else { + console.log('[episodic] disabled — raw turn content will be embedded directly'); } // --- Types --- @@ -53,19 +152,36 @@ async function handleConstruct(body: ConstructRequest) { mindmaps.delete(episodeId); + // When the episodic layer is enabled, each raw turn is sent to an LLM that + // produces a structured summary (production's SummaryService behavior). + // Only the summary content is embedded/stored; raw turn text is preserved + // in metadata.turn.rawContent for debugging and ablation. + const toStore = summarizer + ? await Promise.all( + items.map((item, idx) => summarizer!.summarizeTurn(item, episodeId, idx)), + ) + : items.map((item) => ({ + id: item.id, + role: item.role, + content: item.content, + metadata: item.metadata, + })); + const mindmap = createMindmap({ - provider, - apiKey, - embedModel, + embedFn, storage: memoryStorage(), config: mindmapConfig, }); - await mindmap.add(items); + await mindmap.add(toStore); mindmaps.set(episodeId, mindmap); const state = await mindmap.getState(); - return { success: true, totalItems: state.stats.totalItems }; + return { + success: true, + totalItems: state.stats.totalItems, + summarized: summarizer !== null, + }; } async function handleRetrieve(body: RetrieveRequest) { diff --git a/packages/contexto/package.json b/packages/contexto/package.json index 42c9ffc..ead7750 100644 --- a/packages/contexto/package.json +++ b/packages/contexto/package.json @@ -1,6 +1,6 @@ { "name": "@ekai/contexto", - "version": "0.1.12", + "version": "0.1.13", "description": "Context Engine for Long-running OpenClaw agents", "type": "module", "license": "Apache-2.0", From e0743239f884708311831a8a7d74b4974aa71b28 Mon Sep 17 00:00:00 2001 From: DaevMithran Date: Mon, 27 Apr 2026 18:46:34 +0530 Subject: [PATCH 5/6] Update local mode & benchmark readme --- benchmarks/ama-bench/README.md | 88 ++++++++++++++-- benchmarks/ama-bench/configs/bedrock.json | 29 +++++ benchmarks/ama-bench/configs/hybrid.json | 2 +- benchmarks/ama-bench/configs/local.json | 2 +- benchmarks/ama-bench/scripts/run.sh | 23 +++- benchmarks/ama-bench/src/episodic/summary.ts | 105 ++++++++++++++++--- benchmarks/ama-bench/src/server.ts | 56 +++++++++- 7 files changed, 272 insertions(+), 33 deletions(-) create mode 100644 benchmarks/ama-bench/configs/bedrock.json diff --git a/benchmarks/ama-bench/README.md b/benchmarks/ama-bench/README.md index 9c65f0a..0b2b216 100644 --- a/benchmarks/ama-bench/README.md +++ b/benchmarks/ama-bench/README.md @@ -34,10 +34,11 @@ Always required: - pnpm - `huggingface-cli` (`pip install huggingface_hub`) -Mode-specific (see [Modes](#modes) below): +Mode-specific (see [Modes](#3-pick-a-mode) below): - **`MODE=local`** — Linux box with 4+ GPUs (≈80GB VRAM total) for VLLM, plus [Ollama](https://ollama.com) for embeddings. No API keys needed. - **`MODE=openai`** — `OPENAI_API_KEY` only. No GPUs, no Ollama. - **`MODE=hybrid`** — Ollama for embeddings + `OPENAI_API_KEY` for LLM/judge/summarizer. Mac-friendly, no GPUs. +- **`MODE=bedrock`** — Ollama for embeddings + `BEDROCK_API_KEY` for everything else (Qwen3-32B for answer-gen + judge at 16k ctx, Qwen3-VL-235B-A22B (`/no_think`) for the summarizer at ~256k ctx — instruct-tuned for clean schema adherence). No GPUs, single key. ## Running Locally @@ -84,11 +85,12 @@ bash scripts/setup.sh | `MODE` | Embed | Episodic summarizer | Answer-gen + Judge | Needs | |---|---|---|---|---| -| `local` (default) | Ollama `qwen3-embedding:4b` | VLLM `Qwen/Qwen3-32B` | VLLM `Qwen/Qwen3-32B` | 4+ GPU box, Ollama, VLLM | +| `bedrock` (default) | Ollama `qwen3-embedding:0.6b` | Bedrock Qwen3-VL-235B-A22B (`/no_think`) (~256k ctx) | Bedrock `qwen.qwen3-32b-v1:0` (16k ctx) | Ollama + `BEDROCK_API_KEY` | +| `local` | Ollama `qwen3-embedding:0.6b` | VLLM `Qwen/Qwen3-32B` | VLLM `Qwen/Qwen3-32B` | 4+ GPU box, Ollama, VLLM | | `openai` | OpenAI `text-embedding-3-small` | OpenAI `gpt-4.1-mini` | OpenAI (`gpt-5.2` / `gpt-5.4` judge) | `OPENAI_API_KEY` | -| `hybrid` | Ollama `qwen3-embedding:4b` | OpenAI `gpt-4.1-mini` | OpenAI | Ollama + `OPENAI_API_KEY` | +| `hybrid` | Ollama `qwen3-embedding:0.6b` | OpenAI `gpt-4.1-mini` | OpenAI | Ollama + `OPENAI_API_KEY` | -The preset files are `configs/{local,openai,hybrid}.json`. Edit them to tune `mindmap.*` / `search.*` or swap models — the bridge picks them up via the `BENCH_CONFIG` env var that `run.sh` sets. +The preset files are `configs/{local,openai,hybrid,bedrock}.json`. Edit them to tune `mindmap.*` / `search.*` or swap models — the bridge picks them up via the `BENCH_CONFIG` env var that `run.sh` sets. ### 3a. (Ollama-using modes) install + start Ollama @@ -116,25 +118,33 @@ Skip for `openai` / `hybrid`. If you don't have GPUs locally, common alternatives: Lambda Labs, RunPod, Modal, Replicate, AWS EC2 (g6e.12xlarge), Bedrock (Qwen3-32B dense). Point `episodic.baseUrl` and the AMA-Bench yaml at the remote endpoint. -### 4. (Cloud-using modes) set your OpenAI key +### 4. (Cloud-using modes) set your API key Skip if `MODE=local`. `.env`: ``` +# MODE=openai / hybrid OPENAI_API_KEY=sk-... + +# MODE=bedrock (single key — summarizer uses Bedrock Nova 2 Lite 1M ctx, +# answer-gen + judge use Bedrock Qwen3-32B; prefer long-term keys) +BEDROCK_API_KEY=bedrock-api-key-... ``` -Both `AMA-Bench/configs/gpt-5.2.yaml` and `llm_judge_api.yaml` are wired to read from `OPENAI_API_KEY` — no need to hard-code keys in yaml. +`MODE=bedrock`: `run.sh` aliases `BEDROCK_API_KEY` → both `API_KEY` (bridge summarizer) and `OPENAI_API_KEY` (Python answer-gen + judge). Verify model ids in `configs/bedrock.json` (`amazon.nova-2-lite-v1:0` for episodic) and `AMA-Bench/configs/bedrock.yaml` (`qwen.qwen3-32b-v1:0` for answer-gen/judge) match what's available in your region. + +All cloud yaml configs (`gpt-5.2.yaml`, `llm_judge_api.yaml`, `bedrock.yaml`, `llm_judge_bedrock.yaml`) read the API key from env — no need to hard-code anywhere. ### 5. Run ```bash cd contexto/benchmarks/ama-bench -bash scripts/run.sh # MODE=local (default) +bash scripts/run.sh # MODE=bedrock (default) — local embed + Bedrock Qwen -MODE=openai bash scripts/run.sh # cloud-only -MODE=hybrid bash scripts/run.sh # local embed + cloud LLM +MODE=local bash scripts/run.sh # fully self-hosted (needs 4+ GPUs) +MODE=openai bash scripts/run.sh # cloud-only +MODE=hybrid bash scripts/run.sh # local embed + cloud LLM ``` `run.sh` will: @@ -171,6 +181,7 @@ Each preset already wires `episodic` correctly: - `local.json` — Qwen3-32B on local VLLM, `jsonMode: false`, `noThink: true` - `openai.json` / `hybrid.json` — `gpt-4.1-mini` on OpenAI, `jsonMode: true` +- `bedrock.json` — Qwen3-VL-235B-A22B (`/no_think`) (instruct, ~256k ctx) for the summarizer; `qwen.qwen3-32b-v1:0` (16k ctx) for answer-gen + judge. Same OpenAI-compat endpoint for both, single key. Investigated and ruled out: Amazon Nova (OpenAI-compat path doesn't translate content shape properly), Anthropic Claude (not exposed on OpenAI-compat, would need LiteLLM proxy), Qwen3-Coder variants (coder-tuned → emits empty `key_findings` on trivial turns and treats `/no_think` as user content). The 235B instruct model is only ~47% more expensive than coder-next and cleanly follows the production system prompt. Knobs: - `jsonMode` — set `true` for OpenAI-style `response_format: json_object` (works on OpenAI; on VLLM requires `--guided-decoding-backend outlines`). Set `false` and rely on the system prompt's "JSON only" instruction otherwise. @@ -191,6 +202,60 @@ bash scripts/sweep.sh Edit `configs/sweep.json` to change ranges. Results are saved to `results/sweep_/sweep_summary.csv`, ranked by accuracy. +## Results + +Latest full-benchmark run on the `openend` subset (208 episodes, 2496 questions) using `MODE=bedrock`: + +- **Embed**: Ollama `qwen3-embedding:0.6b` +- **Episodic summarizer**: Bedrock `qwen.qwen3-vl-235b-a22b` (`/no_think`, ~256k ctx) +- **Answer-gen**: Bedrock `qwen.qwen3-32b-v1:0` +- **Judge**: Bedrock `qwen.qwen3-32b-v1:0` + +### Overall + +| Metric | Value | +|---|---| +| Total questions | 2496 | +| Average score | 0.2568 | +| Accuracy | **0.2568** | + +### By task type + +| Task | Accuracy | Questions | +|---|---:|---:| +| 2048 | 0.1528 | 72 | +| alfworld | 0.1750 | 360 | +| babaisai | 0.4444 | 72 | +| candy_crush | 0.5694 | 72 | +| crafter | 0.3194 | 72 | +| gaia_level1 | 0.3250 | 120 | +| gaia_level2 | 0.2944 | 180 | +| gaia_level3 | 0.5000 | 60 | +| minihack | 0.3472 | 72 | +| spider2 | 0.1454 | 612 | +| swebench | 0.2500 | 432 | +| webarena | 0.3414 | 372 | + +### By domain + +| Domain | Accuracy | Questions | +|---|---:|---:| +| EMBODIED_AI | 0.1750 | 360 | +| Game | 0.3667 | 360 | +| OPENWORLD_QA | 0.3389 | 360 | +| SOFTWARE | 0.2500 | 432 | +| TEXT2SQL | 0.1454 | 612 | +| WEB | 0.3414 | 372 | + +### By QA type + +| Type | Accuracy | Questions | +|---|---:|---:| +| A | 0.2622 | 839 | +| B | 0.3003 | 596 | +| C | 0.2581 | 647 | +| D | 0.1812 | 414 | + ## Configuration Reference ### Tree construction (`mindmap`) @@ -238,6 +303,7 @@ contexto/benchmarks/ama-bench/ # Bridge + config + scripts │ ├── local.json # MODE=local preset │ ├── openai.json # MODE=openai preset │ ├── hybrid.json # MODE=hybrid preset +│ ├── bedrock.json # MODE=bedrock preset │ └── sweep.json # Parameter sweep ranges ├── scripts/ │ ├── setup.sh # One-time setup @@ -252,7 +318,9 @@ AMA-Bench/ # Sibling repo │ ├── qwen3-32B.yaml # VLLM answer-gen │ ├── llm_judge.yaml # VLLM judge │ ├── gpt-5.2.yaml # OpenAI answer-gen -│ └── llm_judge_api.yaml # OpenAI judge +│ ├── llm_judge_api.yaml # OpenAI judge +│ ├── bedrock.yaml # Bedrock answer-gen +│ └── llm_judge_bedrock.yaml # Bedrock judge ├── dataset/ # Downloaded via huggingface-cli └── results/ # Benchmark outputs ``` diff --git a/benchmarks/ama-bench/configs/bedrock.json b/benchmarks/ama-bench/configs/bedrock.json new file mode 100644 index 0000000..2a92d32 --- /dev/null +++ b/benchmarks/ama-bench/configs/bedrock.json @@ -0,0 +1,29 @@ +{ + "_comment": "MODE=bedrock — Ollama embeddings + AWS Bedrock for everything. Answer-gen + judge on Qwen3-32B (16k ctx); summarizer on Qwen3-VL-235B-A22B (~256k ctx) with noThink=true to disable hybrid reasoning. Summarization + key-finding extraction is a factual task, not a reasoning one — reasoning models (gpt-oss) waste tokens thinking and leak blocks into the response. VL variant is what Bedrock exposes; vision capability is unused but doesn't hurt text workflows. Single BEDROCK_API_KEY.", + "embed": { + "type": "ollama", + "baseUrl": "http://localhost:11434", + "model": "qwen3-embedding:0.6b" + }, + "episodic": { + "enabled": true, + "baseUrl": "https://bedrock-runtime.us-east-1.amazonaws.com/openai/v1/chat/completions", + "model": "qwen.qwen3-vl-235b-a22b", + "temperature": 0.2, + "jsonMode": false, + "noThink": true, + "maxInputChars": 600000 + }, + "mindmap": { + "similarityThreshold": 0.5, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 4000, + "beamWidth": 3, + "minScore": 0.0 + } +} diff --git a/benchmarks/ama-bench/configs/hybrid.json b/benchmarks/ama-bench/configs/hybrid.json index e1be8f4..5005ce0 100644 --- a/benchmarks/ama-bench/configs/hybrid.json +++ b/benchmarks/ama-bench/configs/hybrid.json @@ -3,7 +3,7 @@ "embed": { "type": "ollama", "baseUrl": "http://localhost:11434", - "model": "qwen3-embedding:4b" + "model": "qwen3-embedding:0.6b" }, "episodic": { "enabled": true, diff --git a/benchmarks/ama-bench/configs/local.json b/benchmarks/ama-bench/configs/local.json index 58cd40e..b0e5de1 100644 --- a/benchmarks/ama-bench/configs/local.json +++ b/benchmarks/ama-bench/configs/local.json @@ -3,7 +3,7 @@ "embed": { "type": "ollama", "baseUrl": "http://localhost:11434", - "model": "qwen3-embedding:4b" + "model": "qwen3-embedding:0.6b" }, "episodic": { "enabled": true, diff --git a/benchmarks/ama-bench/scripts/run.sh b/benchmarks/ama-bench/scripts/run.sh index 7a261ea..139cf3e 100755 --- a/benchmarks/ama-bench/scripts/run.sh +++ b/benchmarks/ama-bench/scripts/run.sh @@ -17,8 +17,9 @@ fi # local → Ollama embed + VLLM (Qwen3-32B) for episodic + answer-gen + judge. No API keys, requires GPUs. # openai → OpenAI for everything (embed + episodic + answer-gen + judge). No GPUs, costs $$. # hybrid → Ollama embed (no rate limits) + OpenAI for episodic + answer-gen + judge. No GPUs. +# bedrock → Ollama embed + AWS Bedrock (Qwen3-32B) for episodic + answer-gen + judge. No GPUs, uses BEDROCK_API_KEY. # Any individual env var below can override the MODE preset. -MODE="${MODE:-local}" +MODE="${MODE:-bedrock}" case "$MODE" in local) @@ -42,8 +43,24 @@ case "$MODE" in DEFAULT_LLM_CONFIG="$AMA_BENCH/configs/gpt-5.2.yaml" DEFAULT_JUDGE_CONFIG="$AMA_BENCH/configs/llm_judge_api.yaml" ;; + bedrock) + DEFAULT_BENCH_CONFIG="$BENCH_DIR/configs/bedrock.json" + DEFAULT_LLM_SERVER="api" + DEFAULT_JUDGE_SERVER="api" + DEFAULT_LLM_CONFIG="$AMA_BENCH/configs/bedrock.yaml" + DEFAULT_JUDGE_CONFIG="$AMA_BENCH/configs/llm_judge_bedrock.yaml" + if [ -z "${BEDROCK_API_KEY:-}" ]; then + echo "ERROR: MODE=bedrock requires BEDROCK_API_KEY in .env." + echo " Get one from AWS Console → Bedrock → API keys (prefer long-term for full runs)." + exit 1 + fi + # Bridge reads API_KEY; Python model_client reads OPENAI_API_KEY. Both paths hit + # Bedrock's OpenAI-compatible endpoint, so alias the single Bedrock bearer token to both. + export API_KEY="$BEDROCK_API_KEY" + export OPENAI_API_KEY="$BEDROCK_API_KEY" + ;; *) - echo "ERROR: unknown MODE='$MODE'. Valid: local | openai | hybrid" + echo "ERROR: unknown MODE='$MODE'. Valid: local | openai | hybrid | bedrock" exit 1 ;; esac @@ -120,7 +137,7 @@ fi # Start bridge server in background echo "[1/3] Starting contexto bridge server on port $BRIDGE_PORT..." cd "$BENCH_DIR" -BRIDGE_PORT=$BRIDGE_PORT BENCH_CONFIG="$BENCH_CONFIG" API_KEY="${API_KEY:-}" OPENAI_API_KEY="${OPENAI_API_KEY:-}" bun src/server.ts & +BRIDGE_PORT=$BRIDGE_PORT BENCH_CONFIG="$BENCH_CONFIG" API_KEY="${API_KEY:-}" OPENAI_API_KEY="${OPENAI_API_KEY:-}" EPISODIC_QUIET="${EPISODIC_QUIET:-1}" EPISODIC_DEBUG="${EPISODIC_DEBUG:-0}" EPISODIC_CONCURRENCY="${EPISODIC_CONCURRENCY:-8}" bun src/server.ts & BRIDGE_PID=$! # Ensure cleanup on exit diff --git a/benchmarks/ama-bench/src/episodic/summary.ts b/benchmarks/ama-bench/src/episodic/summary.ts index 1046865..b8d7d69 100644 --- a/benchmarks/ama-bench/src/episodic/summary.ts +++ b/benchmarks/ama-bench/src/episodic/summary.ts @@ -37,6 +37,7 @@ export interface SummarizerConfig { temperature?: number; // default 0.2 jsonMode?: boolean; // default true — disable for VLLM backends without guided decoding noThink?: boolean; // default false — if true, appends '/no_think' to user message (Qwen3 hybrid-reasoning models only) + maxInputChars?: number; // default 10_000_000 (effectively off) — only set lower for small-context models (e.g. 40000 for 32k-ctx models like Bedrock Qwen3-32B) } export interface ConversationItemInput { @@ -67,14 +68,26 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { const temperature = config.temperature ?? 0.2; const jsonMode = config.jsonMode ?? true; const noThink = config.noThink ?? false; + const maxInputChars = config.maxInputChars ?? 10_000_000; const apiKey = config.apiKey ?? 'EMPTY'; // VLLM accepts any token; OpenAI requires a real key console.log( - `[episodic] initialized model=${config.model} baseUrl=${config.baseUrl} jsonMode=${jsonMode} noThink=${noThink}`, + `[episodic] initialized model=${config.model} baseUrl=${config.baseUrl} jsonMode=${jsonMode} noThink=${noThink} maxInputChars=${maxInputChars}`, ); async function callLLM(content: string): Promise { + // Guard against oversized turns (tool dumps, doc attachments) blowing past the + // model's context window. Truncate middle-out so prefix + suffix of the content + // survive — often more informative than a simple head-truncate. + let trimmed = content; + if (content.length > maxInputChars) { + const half = Math.floor(maxInputChars / 2); + trimmed = + content.slice(0, half) + + `\n\n[... ${content.length - maxInputChars} chars truncated for summarizer context window ...]\n\n` + + content.slice(content.length - half); + } // For Qwen3 hybrid-reasoning models, append /no_think to skip the thinking phase - const userContent = noThink ? `${content}\n\n/no_think` : content; + const userContent = noThink ? `${trimmed}\n\n/no_think` : trimmed; const body: Record = { model: config.model, messages: [ @@ -85,18 +98,34 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { }; if (jsonMode) body.response_format = { type: 'json_object' }; - const response = await fetch(config.baseUrl, { - method: 'POST', - headers: { - 'Content-Type': 'application/json', - Authorization: `Bearer ${apiKey}`, - }, - body: JSON.stringify(body), - }); - - if (!response.ok) { - const body = await response.text().catch(() => '(no body)'); - throw new Error(`summarizer HTTP ${response.status}: ${body.slice(0, 200)}`); + // Retry transient errors (429 rate-limit, 5xx capacity/server errors) with + // exponential backoff. Bedrock's 235B-class models return 500s under bursty + // parallel load — a short retry usually succeeds. + const MAX_ATTEMPTS = 4; + const BASE_DELAY_MS = 1000; + let response: Response | null = null; + let lastError: string = ''; + for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt++) { + response = await fetch(config.baseUrl, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${apiKey}`, + }, + body: JSON.stringify(body), + }); + if (response.ok) break; + const isTransient = response.status === 429 || response.status >= 500; + const bodyText = await response.text().catch(() => '(no body)'); + lastError = `HTTP ${response.status}: ${bodyText.slice(0, 200)}`; + if (!isTransient || attempt === MAX_ATTEMPTS - 1) { + throw new Error(`summarizer ${lastError}`); + } + const delay = BASE_DELAY_MS * Math.pow(2, attempt) + Math.random() * 500; + await new Promise((r) => setTimeout(r, delay)); + } + if (!response || !response.ok) { + throw new Error(`summarizer ${lastError || 'unreachable'}`); } const data = (await response.json()) as { @@ -105,14 +134,55 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { const text = data?.choices?.[0]?.message?.content; if (!text) throw new Error('Empty LLM response'); - // Strip Qwen3 reasoning blocks if present, then markdown code fences + // Strip provider-specific reasoning wrappers before parsing: + // ... — Qwen3 hybrid-reasoning output + // ... — gpt-oss reasoning output + // <|channel|>...<|end|> — gpt-oss harmony-format leakage + // Then strip any markdown code fences the model wraps the JSON in. const stripped = text .replace(/[\s\S]*?<\/think>/g, '') + .replace(/[\s\S]*?<\/reasoning>/g, '') + .replace(/<\|channel\|>[\s\S]*?<\|end\|>/g, '') .trim() .replace(/^```(?:json)?\s*/i, '') .replace(/\s*```$/, '') .trim(); - return JSON.parse(stripped); + try { + return tolerantJsonParse(stripped); + } catch (err) { + // One-time diagnostic: dump the raw model response shape so we can see what the + // provider actually returns (e.g. harmony channel markers from gpt-oss, or + // prose-wrapped JSON). Keeps the rest of the run quiet via EPISODIC_DEBUG guard. + if (process.env.EPISODIC_DEBUG === '1') { + console.error( + `[episodic] parse fail — raw content (first 600 chars):\n${text.slice(0, 600)}\n---`, + ); + } + throw err; + } + } + + /** + * Extract + parse the first JSON object from a model response, tolerating two + * common Qwen-without-guided-decoding failure modes: + * 1. Leading/trailing prose ("Here's the JSON: {...}. Let me know...") + * 2. Invalid string escapes (\p, \m, \x — Qwen sometimes literalizes backslashes) + */ + function tolerantJsonParse(text: string): unknown { + const start = text.indexOf('{'); + const end = text.lastIndexOf('}'); + if (start === -1 || end === -1 || end <= start) { + throw new Error(`No JSON object found in response (head=${text.slice(0, 80)})`); + } + const body = text.slice(start, end + 1); + try { + return JSON.parse(body); + } catch { + // Escape unknown backslash sequences so the parser doesn't choke on \p, \m, etc. + // Leaves valid JSON escapes (\" \\ \/ \b \f \n \r \t \uXXXX) untouched. + const repaired = body.replace(/\\([^"\\/bfnrtu])/g, '\\\\$1'); + return JSON.parse(repaired); + } } return { @@ -158,6 +228,9 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { timestamp: new Date().toISOString(), }; } else { + // Log raw response structure to diagnose which fields the model dropped + const preview = JSON.stringify(raw).slice(0, 400); + console.warn(`[episodic] validation failed (${failures.join(',')}), raw: ${preview}`); summary = toDegradedSummary(raw, failures, traceRef, turn); } } catch (err) { diff --git a/benchmarks/ama-bench/src/server.ts b/benchmarks/ama-bench/src/server.ts index f0bef78..e2ac469 100644 --- a/benchmarks/ama-bench/src/server.ts +++ b/benchmarks/ama-bench/src/server.ts @@ -13,6 +13,26 @@ import { resolve, dirname } from 'path'; import { fileURLToPath } from 'url'; import { createSummarizer, type Summarizer } from './episodic/summary.js'; +// --- Benchmark-only log suppression --- +// The production-parity validator (episodic/validation.ts) logs a warning for every +// turn that fails schema validation (e.g. empty key_findings on blocked turns — see +// qwen3-coder-next behavior). In a 208-episode run that's hundreds of noisy lines. +// Set EPISODIC_QUIET=1 to drop those specific warnings without touching the ported +// production logic. All other logs (errors, bridge info) pass through unchanged. +if (process.env.EPISODIC_QUIET === '1') { + const origWarn = console.warn; + console.warn = (...args: unknown[]) => { + const first = typeof args[0] === 'string' ? args[0] : ''; + if ( + first.startsWith('[episodic] degraded summary') || + first.startsWith('[episodic] validation failed') + ) { + return; + } + origWarn.apply(console, args); + }; +} + // --- Config --- const __dirname = dirname(fileURLToPath(import.meta.url)); @@ -95,6 +115,7 @@ interface EpisodicConfig { temperature?: number; jsonMode?: boolean; noThink?: boolean; + maxInputChars?: number; } const episodicCfg: EpisodicConfig = config.episodic ?? {}; @@ -119,6 +140,7 @@ if (episodicEnabled) { temperature: episodicCfg.temperature ?? 0.2, jsonMode: episodicCfg.jsonMode, noThink: episodicCfg.noThink, + maxInputChars: episodicCfg.maxInputChars, }); } else { console.log('[episodic] disabled — raw turn content will be embedded directly'); @@ -145,6 +167,36 @@ interface ResetRequest { const mindmaps = new Map(); +// --- Helpers --- + +/** + * Run `tasks` with at most `limit` in flight. Preserves input order in the output. + * Used to throttle summarizer fan-out — firing 50+ Promise.all requests at Bedrock's + * large models triggers 500s/429s from capacity limits. + */ +async function mapWithConcurrency( + items: T[], + limit: number, + fn: (item: T, idx: number) => Promise, +): Promise { + const results: U[] = new Array(items.length); + let next = 0; + const workers = new Array(Math.min(limit, items.length)).fill(0).map(async () => { + while (true) { + const i = next++; + if (i >= items.length) return; + results[i] = await fn(items[i], i); + } + }); + await Promise.all(workers); + return results; +} + +const SUMMARIZER_CONCURRENCY = parseInt( + process.env.EPISODIC_CONCURRENCY ?? '8', + 10, +); + // --- Handlers --- async function handleConstruct(body: ConstructRequest) { @@ -157,8 +209,8 @@ async function handleConstruct(body: ConstructRequest) { // Only the summary content is embedded/stored; raw turn text is preserved // in metadata.turn.rawContent for debugging and ablation. const toStore = summarizer - ? await Promise.all( - items.map((item, idx) => summarizer!.summarizeTurn(item, episodeId, idx)), + ? await mapWithConcurrency(items, SUMMARIZER_CONCURRENCY, (item, idx) => + summarizer!.summarizeTurn(item, episodeId, idx), ) : items.map((item) => ({ id: item.id, From f512eaff4a6c220bc483b7033934fdd9816f76f3 Mon Sep 17 00:00:00 2001 From: DaevMithran Date: Tue, 28 Apr 2026 19:16:43 +0530 Subject: [PATCH 6/6] Use 4b embedding --- benchmarks/ama-bench/README.md | 63 ++++++++++--------- benchmarks/ama-bench/configs/bedrock-oai.json | 30 +++++++++ benchmarks/ama-bench/configs/bedrock.json | 13 ++-- benchmarks/ama-bench/configs/hybrid.json | 2 +- benchmarks/ama-bench/configs/local.json | 6 +- benchmarks/ama-bench/scripts/run.sh | 29 ++++++++- benchmarks/ama-bench/scripts/setup.sh | 8 +-- benchmarks/ama-bench/src/episodic/summary.ts | 8 ++- benchmarks/ama-bench/src/server.ts | 29 ++++++++- 9 files changed, 138 insertions(+), 50 deletions(-) create mode 100644 benchmarks/ama-bench/configs/bedrock-oai.json diff --git a/benchmarks/ama-bench/README.md b/benchmarks/ama-bench/README.md index 0b2b216..e8a6632 100644 --- a/benchmarks/ama-bench/README.md +++ b/benchmarks/ama-bench/README.md @@ -85,12 +85,13 @@ bash scripts/setup.sh | `MODE` | Embed | Episodic summarizer | Answer-gen + Judge | Needs | |---|---|---|---|---| -| `bedrock` (default) | Ollama `qwen3-embedding:0.6b` | Bedrock Qwen3-VL-235B-A22B (`/no_think`) (~256k ctx) | Bedrock `qwen.qwen3-32b-v1:0` (16k ctx) | Ollama + `BEDROCK_API_KEY` | -| `local` | Ollama `qwen3-embedding:0.6b` | VLLM `Qwen/Qwen3-32B` | VLLM `Qwen/Qwen3-32B` | 4+ GPU box, Ollama, VLLM | +| `bedrock` (default) | Ollama `qwen3-embedding:4b` | Bedrock Qwen3-VL-235B-A22B (`/no_think`) (~256k ctx) | Bedrock `qwen.qwen3-32b-v1:0` (16k ctx) | Ollama + `BEDROCK_API_KEY` | +| `local` | Ollama `qwen3-embedding:4b` | VLLM `Qwen/Qwen3-32B` | VLLM `Qwen/Qwen3-32B` | 4+ GPU box, Ollama, VLLM | | `openai` | OpenAI `text-embedding-3-small` | OpenAI `gpt-4.1-mini` | OpenAI (`gpt-5.2` / `gpt-5.4` judge) | `OPENAI_API_KEY` | -| `hybrid` | Ollama `qwen3-embedding:0.6b` | OpenAI `gpt-4.1-mini` | OpenAI | Ollama + `OPENAI_API_KEY` | +| `hybrid` | Ollama `qwen3-embedding:4b` | OpenAI `gpt-4.1-mini` | OpenAI | Ollama + `OPENAI_API_KEY` | +| `bedrock-oai` | Ollama `qwen3-embedding:4b` | OpenRouter `openai/gpt-4.1-mini` (1M ctx) | Bedrock `qwen.qwen3-32b-v1:0` (16k ctx) | Ollama + `BEDROCK_API_KEY` + `OPENROUTER_API_KEY` | -The preset files are `configs/{local,openai,hybrid,bedrock}.json`. Edit them to tune `mindmap.*` / `search.*` or swap models — the bridge picks them up via the `BENCH_CONFIG` env var that `run.sh` sets. +The preset files are `configs/{local,openai,hybrid,bedrock,bedrock-oai}.json`. Edit them to tune `mindmap.*` / `search.*` or swap models — the bridge picks them up via the `BENCH_CONFIG` env var that `run.sh` sets. ### 3a. (Ollama-using modes) install + start Ollama @@ -206,55 +207,59 @@ Edit `configs/sweep.json` to change ranges. Results are saved to `results/sweep_ Latest full-benchmark run on the `openend` subset (208 episodes, 2496 questions) using `MODE=bedrock`: -- **Embed**: Ollama `qwen3-embedding:0.6b` -- **Episodic summarizer**: Bedrock `qwen.qwen3-vl-235b-a22b` (`/no_think`, ~256k ctx) +- **Embed**: Ollama `qwen3-embedding:4b` +- **Episodic summarizer**: Bedrock `qwen.qwen3-vl-235b-a22b` (`/no_think`, ~256k ctx, `maxInputChars=600000`) - **Answer-gen**: Bedrock `qwen.qwen3-32b-v1:0` - **Judge**: Bedrock `qwen.qwen3-32b-v1:0` +- **Mindmap**: `similarityThreshold=0.45` +- **Search**: `maxResults=30`, `maxTokens=16000`, `beamWidth=3` + +Up from a baseline of 0.2568 (using `qwen3-embedding:0.6b`, `similarityThreshold=0.5`, `maxResults=10`, `maxTokens=4000`) — net **+1.7 points overall**, with the largest task-level gains on `2048` (+18.1) and `minihack` (+6.9), and the largest QA-type gain on Type B (+4.9 across 596 questions). ### Overall | Metric | Value | |---|---| | Total questions | 2496 | -| Average score | 0.2568 | -| Accuracy | **0.2568** | +| Average score | 0.2740 | +| Accuracy | **0.2740** | ### By task type | Task | Accuracy | Questions | |---|---:|---:| -| 2048 | 0.1528 | 72 | -| alfworld | 0.1750 | 360 | -| babaisai | 0.4444 | 72 | -| candy_crush | 0.5694 | 72 | -| crafter | 0.3194 | 72 | -| gaia_level1 | 0.3250 | 120 | -| gaia_level2 | 0.2944 | 180 | -| gaia_level3 | 0.5000 | 60 | -| minihack | 0.3472 | 72 | -| spider2 | 0.1454 | 612 | -| swebench | 0.2500 | 432 | -| webarena | 0.3414 | 372 | +| 2048 | 0.3333 | 72 | +| alfworld | 0.1806 | 360 | +| babaisai | 0.3194 | 72 | +| candy_crush | 0.4861 | 72 | +| crafter | 0.3472 | 72 | +| gaia_level1 | 0.3417 | 120 | +| gaia_level2 | 0.3222 | 180 | +| gaia_level3 | 0.3833 | 60 | +| minihack | 0.4167 | 72 | +| spider2 | 0.1748 | 612 | +| swebench | 0.2708 | 432 | +| webarena | 0.3656 | 372 | ### By domain | Domain | Accuracy | Questions | |---|---:|---:| -| EMBODIED_AI | 0.1750 | 360 | -| Game | 0.3667 | 360 | +| EMBODIED_AI | 0.1806 | 360 | +| Game | 0.3806 | 360 | | OPENWORLD_QA | 0.3389 | 360 | -| SOFTWARE | 0.2500 | 432 | -| TEXT2SQL | 0.1454 | 612 | -| WEB | 0.3414 | 372 | +| SOFTWARE | 0.2708 | 432 | +| TEXT2SQL | 0.1748 | 612 | +| WEB | 0.3656 | 372 | ### By QA type | Type | Accuracy | Questions | |---|---:|---:| -| A | 0.2622 | 839 | -| B | 0.3003 | 596 | -| C | 0.2581 | 647 | -| D | 0.1812 | 414 | +| A | 0.2646 | 839 | +| B | 0.3490 | 596 | +| C | 0.2736 | 647 | +| D | 0.1860 | 414 | ## Configuration Reference diff --git a/benchmarks/ama-bench/configs/bedrock-oai.json b/benchmarks/ama-bench/configs/bedrock-oai.json new file mode 100644 index 0000000..753a2e3 --- /dev/null +++ b/benchmarks/ama-bench/configs/bedrock-oai.json @@ -0,0 +1,30 @@ +{ + "_comment": "MODE=bedrock-oai — same as bedrock for embed + answer-gen + judge, but episodic summarization runs on gpt-4.1-mini via OpenRouter (1M context). Avoids the Bedrock Qwen3-VL 256k-token ceiling that causes silent 400s and timeouts on long tool-output turns. Requires BEDROCK_API_KEY + OPENROUTER_API_KEY in .env.", + "embed": { + "type": "ollama", + "baseUrl": "http://localhost:11434", + "model": "qwen3-embedding:4b" + }, + "episodic": { + "enabled": true, + "baseUrl": "https://openrouter.ai/api/v1/chat/completions", + "model": "openai/gpt-4.1-mini", + "temperature": 0.2, + "jsonMode": true, + "noThink": false, + "maxInputChars": 3000000, + "maxOutputTokens": 4096 + }, + "mindmap": { + "similarityThreshold": 0.45, + "maxDepth": 4, + "maxChildren": 10, + "rebuildInterval": 50 + }, + "search": { + "maxResults": 10, + "maxTokens": 16000, + "beamWidth": 3, + "minScore": 0.0 + } +} diff --git a/benchmarks/ama-bench/configs/bedrock.json b/benchmarks/ama-bench/configs/bedrock.json index 2a92d32..044fefc 100644 --- a/benchmarks/ama-bench/configs/bedrock.json +++ b/benchmarks/ama-bench/configs/bedrock.json @@ -1,9 +1,9 @@ { - "_comment": "MODE=bedrock — Ollama embeddings + AWS Bedrock for everything. Answer-gen + judge on Qwen3-32B (16k ctx); summarizer on Qwen3-VL-235B-A22B (~256k ctx) with noThink=true to disable hybrid reasoning. Summarization + key-finding extraction is a factual task, not a reasoning one — reasoning models (gpt-oss) waste tokens thinking and leak blocks into the response. VL variant is what Bedrock exposes; vision capability is unused but doesn't hurt text workflows. Single BEDROCK_API_KEY.", + "_comment": "MODE=bedrock — Ollama embeddings + AWS Bedrock for everything. Answer-gen + judge on Qwen3-32B (16k ctx); summarizer on Qwen3-VL-235B-A22B (~256k ctx) with noThink=true to disable hybrid reasoning. maxInputChars=600000 — Bedrock returns silent 400s on the rare turn that exceeds the 256k token ceiling (≤3 / 208 episodes in practice), but tighter truncation hurts summary quality on long episodes far more than the 400s cost. Llama 4 / Nova would be cleaner (1M ctx) but aren't exposed on Bedrock's OpenAI-compat /chat/completions surface as of 2026-04. Single BEDROCK_API_KEY.", "embed": { "type": "ollama", "baseUrl": "http://localhost:11434", - "model": "qwen3-embedding:0.6b" + "model": "qwen3-embedding:4b" }, "episodic": { "enabled": true, @@ -12,17 +12,18 @@ "temperature": 0.2, "jsonMode": false, "noThink": true, - "maxInputChars": 600000 + "maxInputChars": 600000, + "maxOutputTokens": 4096 }, "mindmap": { - "similarityThreshold": 0.5, + "similarityThreshold": 0.45, "maxDepth": 4, "maxChildren": 10, "rebuildInterval": 50 }, "search": { - "maxResults": 10, - "maxTokens": 4000, + "maxResults": 30, + "maxTokens": 16000, "beamWidth": 3, "minScore": 0.0 } diff --git a/benchmarks/ama-bench/configs/hybrid.json b/benchmarks/ama-bench/configs/hybrid.json index 5005ce0..e1be8f4 100644 --- a/benchmarks/ama-bench/configs/hybrid.json +++ b/benchmarks/ama-bench/configs/hybrid.json @@ -3,7 +3,7 @@ "embed": { "type": "ollama", "baseUrl": "http://localhost:11434", - "model": "qwen3-embedding:0.6b" + "model": "qwen3-embedding:4b" }, "episodic": { "enabled": true, diff --git a/benchmarks/ama-bench/configs/local.json b/benchmarks/ama-bench/configs/local.json index b0e5de1..2846f7c 100644 --- a/benchmarks/ama-bench/configs/local.json +++ b/benchmarks/ama-bench/configs/local.json @@ -3,7 +3,7 @@ "embed": { "type": "ollama", "baseUrl": "http://localhost:11434", - "model": "qwen3-embedding:0.6b" + "model": "qwen3-embedding:4b" }, "episodic": { "enabled": true, @@ -14,14 +14,14 @@ "noThink": true }, "mindmap": { - "similarityThreshold": 0.5, + "similarityThreshold": 0.45, "maxDepth": 4, "maxChildren": 10, "rebuildInterval": 50 }, "search": { "maxResults": 10, - "maxTokens": 4000, + "maxTokens": 16000, "beamWidth": 3, "minScore": 0.0 } diff --git a/benchmarks/ama-bench/scripts/run.sh b/benchmarks/ama-bench/scripts/run.sh index 139cf3e..cd11076 100755 --- a/benchmarks/ama-bench/scripts/run.sh +++ b/benchmarks/ama-bench/scripts/run.sh @@ -18,6 +18,7 @@ fi # openai → OpenAI for everything (embed + episodic + answer-gen + judge). No GPUs, costs $$. # hybrid → Ollama embed (no rate limits) + OpenAI for episodic + answer-gen + judge. No GPUs. # bedrock → Ollama embed + AWS Bedrock (Qwen3-32B) for episodic + answer-gen + judge. No GPUs, uses BEDROCK_API_KEY. +# bedrock-oai → bedrock for embed + answer-gen + judge, but episodic summarizer on gpt-4.1-mini via OpenRouter (1M ctx). Needs BEDROCK_API_KEY + OPENROUTER_API_KEY (or OPENAI_API_KEY for direct-OpenAI users). # Any individual env var below can override the MODE preset. MODE="${MODE:-bedrock}" @@ -59,8 +60,30 @@ case "$MODE" in export API_KEY="$BEDROCK_API_KEY" export OPENAI_API_KEY="$BEDROCK_API_KEY" ;; + bedrock-oai) + DEFAULT_BENCH_CONFIG="$BENCH_DIR/configs/bedrock-oai.json" + DEFAULT_LLM_SERVER="api" + DEFAULT_JUDGE_SERVER="api" + DEFAULT_LLM_CONFIG="$AMA_BENCH/configs/bedrock.yaml" + DEFAULT_JUDGE_CONFIG="$AMA_BENCH/configs/llm_judge_bedrock.yaml" + if [ -z "${BEDROCK_API_KEY:-}" ]; then + echo "ERROR: MODE=bedrock-oai requires BEDROCK_API_KEY in .env (answer-gen + judge → Bedrock)." + exit 1 + fi + # Summarizer key — prefer OPENROUTER_API_KEY, fall back to OPENAI_API_KEY for direct-OpenAI users. + SUMMARIZER_KEY="${OPENROUTER_API_KEY:-${OPENAI_API_KEY:-}}" + if [ -z "$SUMMARIZER_KEY" ]; then + echo "ERROR: MODE=bedrock-oai requires OPENROUTER_API_KEY (or OPENAI_API_KEY) in .env for the episodic summarizer." + exit 1 + fi + # Capture the summarizer key BEFORE we alias OPENAI_API_KEY to the Bedrock key for the + # Python answer-gen + judge pipeline. + export EPISODIC_API_KEY="$SUMMARIZER_KEY" + export API_KEY="$BEDROCK_API_KEY" + export OPENAI_API_KEY="$BEDROCK_API_KEY" + ;; *) - echo "ERROR: unknown MODE='$MODE'. Valid: local | openai | hybrid | bedrock" + echo "ERROR: unknown MODE='$MODE'. Valid: local | openai | hybrid | bedrock | bedrock-oai" exit 1 ;; esac @@ -94,7 +117,7 @@ if [ "$EMBED_TYPE" = "ollama" ]; then fi if ! curl -sf "$OLLAMA_URL/api/tags" | grep -q "$EMBED_MODEL_HINT"; then echo "ERROR: embedding model not found in Ollama" - echo " Pull: ollama pull qwen3-embedding:0.6b" + echo " Pull: ollama pull qwen3-embedding:4b" exit 1 fi echo "Ollama OK ($OLLAMA_URL) — $EMBED_MODEL_HINT present" @@ -137,7 +160,7 @@ fi # Start bridge server in background echo "[1/3] Starting contexto bridge server on port $BRIDGE_PORT..." cd "$BENCH_DIR" -BRIDGE_PORT=$BRIDGE_PORT BENCH_CONFIG="$BENCH_CONFIG" API_KEY="${API_KEY:-}" OPENAI_API_KEY="${OPENAI_API_KEY:-}" EPISODIC_QUIET="${EPISODIC_QUIET:-1}" EPISODIC_DEBUG="${EPISODIC_DEBUG:-0}" EPISODIC_CONCURRENCY="${EPISODIC_CONCURRENCY:-8}" bun src/server.ts & +BRIDGE_PORT=$BRIDGE_PORT BENCH_CONFIG="$BENCH_CONFIG" API_KEY="${API_KEY:-}" OPENAI_API_KEY="${OPENAI_API_KEY:-}" BEDROCK_API_KEY="${BEDROCK_API_KEY:-}" EPISODIC_API_KEY="${EPISODIC_API_KEY:-}" EPISODIC_QUIET="${EPISODIC_QUIET:-1}" EPISODIC_DEBUG="${EPISODIC_DEBUG:-0}" EPISODIC_CONCURRENCY="${EPISODIC_CONCURRENCY:-8}" bun src/server.ts & BRIDGE_PID=$! # Ensure cleanup on exit diff --git a/benchmarks/ama-bench/scripts/setup.sh b/benchmarks/ama-bench/scripts/setup.sh index d737c77..6c926a3 100755 --- a/benchmarks/ama-bench/scripts/setup.sh +++ b/benchmarks/ama-bench/scripts/setup.sh @@ -37,18 +37,18 @@ cd "$BENCH_DIR" && pnpm install # 5. Pull Ollama embedding model (optional — won't hard-fail) echo "[5/5] Checking Ollama..." if command -v ollama >/dev/null 2>&1; then - echo "Pulling qwen3-embedding:0.6b (idempotent)..." - ollama pull qwen3-embedding:0.6b || echo "WARN: ollama pull failed — run 'ollama serve &' first, then retry." + echo "Pulling qwen3-embedding:4b (idempotent)..." + ollama pull qwen3-embedding:4b || echo "WARN: ollama pull failed — run 'ollama serve &' first, then retry." else echo "Ollama not installed. To enable local embeddings:" echo " brew install ollama" echo " ollama serve &" - echo " ollama pull qwen3-embedding:0.6b" + echo " ollama pull qwen3-embedding:4b" fi echo "" echo "=== Setup complete ===" echo "Next steps:" -echo " 1. Ensure 'ollama serve' is running and qwen3-embedding:0.6b is pulled" +echo " 1. Ensure 'ollama serve' is running and qwen3-embedding:4b is pulled" echo " 2. (Optional) tune mindmap params in configs/default.json" echo " 3. Run: bash scripts/run.sh" diff --git a/benchmarks/ama-bench/src/episodic/summary.ts b/benchmarks/ama-bench/src/episodic/summary.ts index b8d7d69..f3afd91 100644 --- a/benchmarks/ama-bench/src/episodic/summary.ts +++ b/benchmarks/ama-bench/src/episodic/summary.ts @@ -38,6 +38,7 @@ export interface SummarizerConfig { jsonMode?: boolean; // default true — disable for VLLM backends without guided decoding noThink?: boolean; // default false — if true, appends '/no_think' to user message (Qwen3 hybrid-reasoning models only) maxInputChars?: number; // default 10_000_000 (effectively off) — only set lower for small-context models (e.g. 40000 for 32k-ctx models like Bedrock Qwen3-32B) + maxOutputTokens?: number; // default 4096 — sent as max_tokens. Without this OpenRouter/Bedrock often truncate mid-JSON. } export interface ConversationItemInput { @@ -69,9 +70,10 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { const jsonMode = config.jsonMode ?? true; const noThink = config.noThink ?? false; const maxInputChars = config.maxInputChars ?? 10_000_000; + const maxOutputTokens = config.maxOutputTokens ?? 4096; const apiKey = config.apiKey ?? 'EMPTY'; // VLLM accepts any token; OpenAI requires a real key console.log( - `[episodic] initialized model=${config.model} baseUrl=${config.baseUrl} jsonMode=${jsonMode} noThink=${noThink} maxInputChars=${maxInputChars}`, + `[episodic] initialized model=${config.model} baseUrl=${config.baseUrl} jsonMode=${jsonMode} noThink=${noThink} maxInputChars=${maxInputChars} maxOutputTokens=${maxOutputTokens}`, ); async function callLLM(content: string): Promise { @@ -95,6 +97,7 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { { role: 'user', content: userContent }, ], temperature, + max_tokens: maxOutputTokens, }; if (jsonMode) body.response_format = { type: 'json_object' }; @@ -122,6 +125,9 @@ export function createSummarizer(config: SummarizerConfig): Summarizer { throw new Error(`summarizer ${lastError}`); } const delay = BASE_DELAY_MS * Math.pow(2, attempt) + Math.random() * 500; + console.warn( + `[episodic] retry ${attempt + 1}/${MAX_ATTEMPTS - 1} — ${lastError} — backing off ${Math.round(delay)}ms`, + ); await new Promise((r) => setTimeout(r, delay)); } if (!response || !response.ok) { diff --git a/benchmarks/ama-bench/src/server.ts b/benchmarks/ama-bench/src/server.ts index e2ac469..277bc89 100644 --- a/benchmarks/ama-bench/src/server.ts +++ b/benchmarks/ama-bench/src/server.ts @@ -77,7 +77,7 @@ function makeOllamaEmbedFn({ baseUrl, model }: { baseUrl: string; model: string function buildEmbedFn(): EmbedFn { if (embedType === 'ollama') { const baseUrl = embedCfg.baseUrl ?? 'http://localhost:11434'; - const model = embedCfg.model ?? 'qwen3-embedding:0.6b'; + const model = embedCfg.model ?? 'qwen3-embedding:4b'; console.log(`[bridge] Embed: ollama ${baseUrl} model=${model}`); return makeOllamaEmbedFn({ baseUrl, model }); } @@ -116,6 +116,7 @@ interface EpisodicConfig { jsonMode?: boolean; noThink?: boolean; maxInputChars?: number; + maxOutputTokens?: number; } const episodicCfg: EpisodicConfig = config.episodic ?? {}; @@ -126,11 +127,22 @@ if (episodicEnabled) { const baseUrl = episodicCfg.baseUrl ?? 'https://api.openai.com/v1/chat/completions'; // Local endpoints (VLLM/Ollama/etc.) don't need a real API key const isLocal = /^https?:\/\/(localhost|127\.0\.0\.1|0\.0\.0\.0|\[::1?\])/i.test(baseUrl); + // EPISODIC_API_KEY is preferred when set — used by mixed modes (e.g. bedrock-oai) + // where API_KEY/OPENAI_API_KEY are aliased to a different provider's key for the + // answer-gen + judge pipeline, but the summarizer needs a separate provider's key. + // BEDROCK_API_KEY is recognized directly so the bridge works without run.sh's aliasing. + // NOTE: || (not ??) — run.sh forwards unset vars as empty strings, which ?? would + // treat as a valid value and short-circuit on. We want empty-string to fall through. const apiKey = - episodicCfg.apiKey ?? process.env.API_KEY ?? process.env.OPENAI_API_KEY ?? ''; + episodicCfg.apiKey || + process.env.EPISODIC_API_KEY || + process.env.BEDROCK_API_KEY || + process.env.API_KEY || + process.env.OPENAI_API_KEY || + ''; if (!apiKey && !isLocal) { throw new Error( - 'episodic.enabled=true but no API key found (set episodic.apiKey in configs/default.json, or API_KEY / OPENAI_API_KEY env var).', + 'episodic.enabled=true but no API key found (set episodic.apiKey in the config, or one of: EPISODIC_API_KEY / BEDROCK_API_KEY / API_KEY / OPENAI_API_KEY env vars).', ); } summarizer = createSummarizer({ @@ -141,6 +153,7 @@ if (episodicEnabled) { jsonMode: episodicCfg.jsonMode, noThink: episodicCfg.noThink, maxInputChars: episodicCfg.maxInputChars, + maxOutputTokens: episodicCfg.maxOutputTokens, }); } else { console.log('[episodic] disabled — raw turn content will be embedded directly'); @@ -203,11 +216,13 @@ async function handleConstruct(body: ConstructRequest) { const { episodeId, items } = body; mindmaps.delete(episodeId); + const t0 = performance.now(); // When the episodic layer is enabled, each raw turn is sent to an LLM that // produces a structured summary (production's SummaryService behavior). // Only the summary content is embedded/stored; raw turn text is preserved // in metadata.turn.rawContent for debugging and ablation. + const tSummStart = performance.now(); const toStore = summarizer ? await mapWithConcurrency(items, SUMMARIZER_CONCURRENCY, (item, idx) => summarizer!.summarizeTurn(item, episodeId, idx), @@ -218,6 +233,7 @@ async function handleConstruct(body: ConstructRequest) { content: item.content, metadata: item.metadata, })); + const summarizeMs = performance.now() - tSummStart; const mindmap = createMindmap({ embedFn, @@ -225,10 +241,17 @@ async function handleConstruct(body: ConstructRequest) { config: mindmapConfig, }); + const tBuildStart = performance.now(); await mindmap.add(toStore); + const buildMs = performance.now() - tBuildStart; + mindmaps.set(episodeId, mindmap); const state = await mindmap.getState(); + const totalMs = performance.now() - t0; + console.log( + `[bridge] /construct ep=${episodeId} turns=${items.length} summarize=${summarizeMs.toFixed(0)}ms build=${buildMs.toFixed(0)}ms total=${totalMs.toFixed(0)}ms`, + ); return { success: true, totalItems: state.stats.totalItems,