From 0ca9464a676c70d13f4f538db6e1ea5d8490f625 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Wed, 3 Jun 2026 23:46:22 +0100 Subject: [PATCH 01/17] feat(bench): add BEAM benchmark foundation (compose + submodule + isolation) Foundation only -- no harness code yet, no benchmark runs. Lays the scaffolding for v0.5.0's BEAM target so the next commit can implement the adapter / runner / judge / score with the safety + trust posture locked in. What ships: - benchmarks/beam/upstream/ as git submodule pinned to mohammadtavakoli78/BEAM@3e12035 (ICLR 2026, arXiv:2510.27246). Our harness will call their eval module unmodified and import the judge prompt (src/prompts.py::unified_llm_judge_base_prompt) at runtime -- zero prompt content in our codebase. - docker-compose.bench.yml: standalone bench stack with strict isolation. Separate Postgres container on host port 5433, separate database (braindb_bench), separate volume, separate BrainDB API on :8001, separate watcher data dir (./data_bench/sources/). The personal braindb stack is never touched. - .env.bench.example + benchmarks/beam/README.md document the bench config knobs, trust model, isolation guarantees, and Qwen-as-judge caveat (Qwen not directly comparable to mem0's GPT-4o or Hindsight's Gemini-judged BEAM scores; cross-judge calibration planned). - Granular .gitignore: personal A/B content stays local-only, benchmarks/beam/ is tracked, generated artefacts ignored. --- .env.bench.example | 41 ++++++++++++ .gitignore | 19 ++++++ .gitmodules | 3 + benchmarks/beam/README.md | 110 +++++++++++++++++++++++++++++++ benchmarks/beam/upstream | 1 + docker-compose.bench.yml | 133 ++++++++++++++++++++++++++++++++++++++ 6 files changed, 307 insertions(+) create mode 100644 .env.bench.example create mode 100644 .gitmodules create mode 100644 benchmarks/beam/README.md create mode 160000 benchmarks/beam/upstream create mode 100644 docker-compose.bench.yml diff --git a/.env.bench.example b/.env.bench.example new file mode 100644 index 0000000..e718f49 --- /dev/null +++ b/.env.bench.example @@ -0,0 +1,41 @@ +# Copy this to .env.bench and fill in the values you need. +# +# Bench-only environment. Used by docker-compose.bench.yml + the bench +# runner under benchmarks/. The personal stack reads .env (separate file) +# and is unaffected by anything here. +# +# DO NOT commit .env.bench — it's gitignored. Only this template is tracked. + +# --- LLM provider for the bench BrainDB API --- +# Bench defaults to local Qwen via vLLM on the workstation (per plan). +# Override here if you want a different provider for bench runs. +LLM_PROFILE_BENCH=vllm_workstation_qwen + +# Optional model override (e.g. to try a smaller Qwen variant for bench). +AGENT_MODEL= + +# --- Provider keys (only fill in what your LLM_PROFILE_BENCH needs) --- +NVIDIA_NIM_API_KEY= +DEEPINFRA_API_KEY= + +# Self-hosted vLLM: leave empty unless your vLLM was started with --api-key. +# The resolver supplies the literal "EMPTY" placeholder if blank. +VLLM_API_KEY= + +# HuggingFace token — for the bench dataset cache + the BEAM dataset download +HF_TOKEN= + +# --- Bench-specific tuning knobs (defaults match docker-compose.bench.yml) --- +# How fresh new entities must be before the wiki maintainer considers them +# eligible (prod default: 30 min). Bench: 1 min so warmup stays short. +WIKI_FRESHNESS_MINUTES_BENCH=1 + +# Wiki scheduler tick (prod default: 60s). Bench: 5s so per-conversation +# warmup is dominated by LLM time, not scheduler wait. +WIKI_INTERVAL_BENCH=5 + +# Ingest watcher poll interval (prod default: 7s). Bench: 3s for snappy smoke. +INGEST_POLL_INTERVAL_BENCH=3 + +# Agent verbose logging (every tool call printed). Bench default: true. +AGENT_VERBOSE_BENCH=true diff --git a/.gitignore b/.gitignore index 3a17d0f..b10de4a 100644 --- a/.gitignore +++ b/.gitignore @@ -39,3 +39,22 @@ data/sources/* # Wiki review exports — generated, read-only inspection output data/wiki_review/ + +# --- Benchmarks --- +# Personal A/B bench content (5-Q personal recall + run traces) — local only, +# tied to the user's own BrainDB content, not for the public repo. +benchmarks/README.md +benchmarks/questions.json +benchmarks/run_agent.sh +benchmarks/runs/ + +# Generated artefacts inside any public benchmark subdir (e.g. benchmarks/beam/) +# — full LLM traces, dataset caches, per-question JSONLs. Only results/*.md +# gets committed. +benchmarks/*/runs/ +benchmarks/*/dataset_cache/ +benchmarks/*/answers/ + +# Bench compose: separate host data dir for the bench watcher, separate env file +data_bench/ +.env.bench diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..9669c51 --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "benchmarks/beam/upstream"] + path = benchmarks/beam/upstream + url = https://github.com/mohammadtavakoli78/BEAM.git diff --git a/benchmarks/beam/README.md b/benchmarks/beam/README.md new file mode 100644 index 0000000..71f7c04 --- /dev/null +++ b/benchmarks/beam/README.md @@ -0,0 +1,110 @@ +# BrainDB on BEAM (Beyond a Million Tokens) + +Public benchmark harness running BrainDB against the **BEAM** memory benchmark +([arXiv:2510.27246](https://arxiv.org/abs/2510.27246), ICLR 2026, Tavakoli et al. +of U Alberta + UMass Amherst). + +**Status**: Step 0 scaffolding (this file, the bench compose, the upstream +submodule). Phase 1 harness code (`adapter.py`, `bench.py`, `judge.py`, +`score.py`, `warmup.py`) is being implemented in the same commit series. + +--- + +## Trust model + +The integrity bar is: **the parts that score us are theirs, used unmodified**. +Nothing in our code paraphrases, copies, or reinterprets the benchmark. + +| Component | Where it lives | Why this matters | +|---|---|---| +| BEAM dataset | `Mohammadta/BEAM` on HuggingFace; pinned by sha256 in our harness | Same SHA → same dataset; anyone can verify | +| BEAM eval script | Git submodule `upstream/` pinned to a specific commit SHA of `mohammadtavakoli78/BEAM` | We call `python -m src.evaluation.run_evaluation` unmodified | +| Judge prompt | `from src.prompts import unified_llm_judge_base_prompt` — loaded at runtime from the submodule | Zero prompt content in our codebase; no risk of paraphrase or typo | +| Adapter (dataset → BrainDB ingest) | `adapter.py` (ours, ~150 LOC) | Transparent: read the code | +| Bench runner (per-conversation reset + ingest + warmup + answer) | `bench.py` (ours, ~120 LOC) | Same | +| Judge runner (calls Qwen with the upstream prompt) | `judge.py` (ours, ~80 LOC) | Anyone can re-judge our `answers.json` with any model | +| Eval wrapper (invokes upstream eval) | `score.py` (ours, ~50 LOC) | Just a thin caller | + +**The "no copy" rule** is satisfied at the strongest level: the judge prompt +text never leaves the upstream submodule. We import it as a Python string +and pass it to our chosen judge model. Hindsight and mem0 both do the same +(they wire it into their own judge runners); the upstream eval code itself +already imports it from `src/prompts.py::unified_llm_judge_base_prompt`. + +--- + +## Isolation from your personal BrainDB + +The bench runs in `docker-compose.bench.yml` — a completely separate stack +from your personal `docker-compose.yml`. Layered isolation: + +- Separate Docker project namespace: `name: braindb_bench` +- **Separate Postgres container** (`braindb_bench_postgres`) on host port 5433 +- Separate Postgres database: `braindb_bench` (never `braindb`) +- Separate Postgres data volume: `braindb_bench_pgdata` +- **Separate BrainDB API on port 8001** (personal stays on 8000) +- **Separate host data directory**: `./data_bench/sources/` — the bench + watcher polls this; personal watcher continues polling `./data/sources/` + and never sees bench files +- Hard-coded safety assertion in `bench.py`: the active `DATABASE_URL` + MUST contain `braindb_bench` literally, or the runner refuses to start +- Explicit invocation: bench requires + `docker compose -f docker-compose.bench.yml --env-file .env.bench up`; + a plain `docker compose up` runs the personal stack as normal + +Before the first benchmark run, snapshot your personal Postgres volume as a +paranoia tarball: + +```bash +docker run --rm -v braindb_pgdata:/data -v "$(pwd)":/backup alpine \ + tar czf /backup/braindb_personal_backup_$(date +%Y%m%d).tar.gz /data +``` + +One-command restore if anything ever goes wrong. + +--- + +## Caveats on the published number + +We use **Qwen 3.6-27B** (local, via the workstation vLLM tunnel) as the LLM +judge. Published BEAM scores used GPT-4o (mem0) or Gemini-2.5-flash-lite +(Hindsight). **Our Qwen-judged number is NOT directly comparable to those +published numbers** — judge models systematically differ. + +The mitigations: + +1. We publish `answers.json` (our raw answers, judge-independent) so anyone + with another model's API access can re-judge our work and verify. +2. We run a 30-question stratified Claude Sonnet calibration sample (~$15–30 + Anthropic credit) so the delta between Qwen and Claude scoring is + published alongside the headline. +3. BEAM's ceiling is ~73% (Hindsight @ 1M, Gemini judge) / ~64% (mem0 @ 1M, + GPT-4o judge) — lower than LongMemEval's ~95%. Less headroom → small judge + biases shift the number more. This makes the calibration sample + **mandatory for credibility**, not optional. + +Bench-mode config tunes the cadence of the wiki maintenance pipeline (5s +tick instead of 60s; writer concurrency 5 instead of 1) so per-conversation +warmup completes in minutes, not 30+ min. **The pipeline itself is +identical to production** — same extraction prompts, same wiki maintainer, +same writer. Only the throughput knobs differ, and they are listed +explicitly in `docker-compose.bench.yml` so reviewers see them. + +--- + +## Citations + +If you reference BrainDB's BEAM numbers, please also cite the original BEAM +paper — it's their benchmark, we just ran on it: + +``` +@inproceedings{tavakoli2026beam, + title={Beyond a Million Tokens: Benchmarking and Enhancing Long-Term Memory in LLMs}, + author={Tavakoli, Mohammad and Salemi, Alireza and Ye, Carrie and Abdalla, Mohamed and Zamani, Hamed and Mitchell, J. Ross}, + booktitle={ICLR 2026}, + year={2026} +} +``` + +BEAM dataset and code are CC BY-SA 4.0 / MIT respectively; see +`upstream/LICENSE`. diff --git a/benchmarks/beam/upstream b/benchmarks/beam/upstream new file mode 160000 index 0000000..3e12035 --- /dev/null +++ b/benchmarks/beam/upstream @@ -0,0 +1 @@ +Subproject commit 3e12035532eb85768f1a7cd779832b650c4b2ef9 diff --git a/docker-compose.bench.yml b/docker-compose.bench.yml new file mode 100644 index 0000000..b5931f0 --- /dev/null +++ b/docker-compose.bench.yml @@ -0,0 +1,133 @@ +# Standalone bench stack — NOT an overlay on docker-compose.yml. +# +# Runs alongside the personal BrainDB stack with strict isolation at every +# layer: separate Docker project namespace, separate Postgres container, +# separate volume, separate ports, separate host data dir. The personal +# `braindb` database is never touched by anything here. +# +# Run with: +# docker compose -f docker-compose.bench.yml --env-file .env.bench up +# +# Plain `docker compose up` runs the personal stack as normal — the bench +# stack cannot fire by accident. +# +# Safety properties (also enforced by code in benchmarks/beam/bench.py): +# - Postgres database is named `braindb_bench` (never `braindb`) +# - Postgres runs in its OWN container on host port 5433 (not 5432) +# - The bench BrainDB API is on host port 8001 (not 8000) +# - The bench watcher mounts ./data_bench/ as /app/data, so it never +# sees files in ./data/ (the personal watcher's directory) +# - The bench runner asserts the active DATABASE_URL contains +# `braindb_bench` before any destructive op; refuses otherwise + +name: braindb_bench + +services: + postgres_bench: + # pgvector + pg_trgm baked in; matches what alembic migrations expect. + image: pgvector/pgvector:pg17 + container_name: braindb_bench_postgres + restart: unless-stopped + environment: + POSTGRES_DB: braindb_bench + POSTGRES_USER: braindb_bench + # NOT a secret — bench-local only, never reachable from outside the + # host. The hard-coded safety check in bench.py looks for the + # literal "braindb_bench" substring in DATABASE_URL. + POSTGRES_PASSWORD: bench_local_only + ports: + - "5433:5432" + volumes: + - braindb_bench_pgdata:/var/lib/postgresql/data + networks: + - bench-network + healthcheck: + test: ["CMD-SHELL", "pg_isready -U braindb_bench -d braindb_bench"] + interval: 5s + timeout: 3s + retries: 10 + + api_bench: + build: . + container_name: braindb_bench_api + restart: unless-stopped + depends_on: + postgres_bench: + condition: service_healthy + networks: + - bench-network + environment: + # MUST contain "braindb_bench" or the bench runner refuses to start. + DATABASE_URL: postgresql://braindb_bench:bench_local_only@postgres_bench:5432/braindb_bench + API_PORT: 8001 + HF_TOKEN: ${HF_TOKEN:-} + # Bench defaults to local Qwen via the workstation tunnel — matches the + # plan's Qwen-as-judge story. Override via .env.bench if needed. + LLM_PROFILE: ${LLM_PROFILE_BENCH:-vllm_workstation_qwen} + AGENT_MODEL: ${AGENT_MODEL:-} + NVIDIA_NIM_API_KEY: ${NVIDIA_NIM_API_KEY:-} + DEEPINFRA_API_KEY: ${DEEPINFRA_API_KEY:-} + VLLM_API_KEY: ${VLLM_API_KEY:-EMPTY} + # Verbose by default during bench — every tool call is logged so the + # smoke tests can be audited. + AGENT_VERBOSE: ${AGENT_VERBOSE_BENCH:-true} + # Bench-specific tuning: wiki maintenance fires faster than prod's + # 30-min freshness window so per-conversation warmup stays in minutes. + WIKI_FRESHNESS_MINUTES: ${WIKI_FRESHNESS_MINUTES_BENCH:-1} + extra_hosts: + # Lets self-hosted profiles (vllm_workstation_qwen) reach a server + # bound to the Docker host's loopback or the SSH tunnel. + - "host.docker.internal:host-gateway" + ports: + - "8001:8001" + volumes: + # Bind the repo (for code reload) but OVERRIDE /app/data with the + # bench-only directory. The personal watcher polls ./data/sources/; + # the bench watcher polls ./data_bench/sources/. They cannot see + # each other's files. + - .:/app + - ./data_bench:/app/data + command: > + sh -c "alembic upgrade head && uvicorn braindb.main:app --host 0.0.0.0 --port 8001" + + watcher_bench: + build: . + container_name: braindb_bench_watcher + restart: unless-stopped + depends_on: + - api_bench + networks: + - bench-network + environment: + BRAINDB_API_URL: http://api_bench:8001 + # Faster poll than prod's 7s, so the smoke test feels responsive. + INGEST_POLL_INTERVAL: ${INGEST_POLL_INTERVAL_BENCH:-3} + volumes: + - .:/app + - ./data_bench:/app/data + command: python -m braindb.ingest_watcher + + wiki_scheduler_bench: + build: . + container_name: braindb_bench_wiki_scheduler + restart: unless-stopped + depends_on: + - api_bench + networks: + - bench-network + environment: + BRAINDB_API_URL: http://api_bench:8001 + # Wiki pipeline runs identical to prod — only the cadence changes. + WIKI_ENABLED: 'true' + WIKI_INTERVAL: ${WIKI_INTERVAL_BENCH:-5} + volumes: + - .:/app + command: python -m braindb.wiki_scheduler + +volumes: + braindb_bench_pgdata: + name: braindb_bench_pgdata + +networks: + bench-network: + name: braindb_bench_network From afd32e21121c0eed2a10778fe456ddc3343d68d2 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Wed, 3 Jun 2026 23:48:46 +0100 Subject: [PATCH 02/17] fix(bench): move bench Postgres to host port 5434 (5433 already taken) The personal Postgres in this environment (postgres_container, timescale/timescaledb:latest-pg16) already binds host port 5433. Bench Postgres claimed the same port and would have refused to start. Move bench Postgres to 5434. Verified port 5434 (and 8001 for the bench API) free with netstat before commit. README updated to match. No other isolation properties change. --- benchmarks/beam/README.md | 4 +++- docker-compose.bench.yml | 8 ++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/benchmarks/beam/README.md b/benchmarks/beam/README.md index 71f7c04..a40ea13 100644 --- a/benchmarks/beam/README.md +++ b/benchmarks/beam/README.md @@ -39,7 +39,9 @@ The bench runs in `docker-compose.bench.yml` — a completely separate stack from your personal `docker-compose.yml`. Layered isolation: - Separate Docker project namespace: `name: braindb_bench` -- **Separate Postgres container** (`braindb_bench_postgres`) on host port 5433 +- **Separate Postgres container** (`braindb_bench_postgres`) on host port 5434 + (host port 5433 is already used by the personal `postgres_container` in + this environment; bench takes 5434 to avoid a collision) - Separate Postgres database: `braindb_bench` (never `braindb`) - Separate Postgres data volume: `braindb_bench_pgdata` - **Separate BrainDB API on port 8001** (personal stays on 8000) diff --git a/docker-compose.bench.yml b/docker-compose.bench.yml index b5931f0..8f2dec7 100644 --- a/docker-compose.bench.yml +++ b/docker-compose.bench.yml @@ -13,7 +13,9 @@ # # Safety properties (also enforced by code in benchmarks/beam/bench.py): # - Postgres database is named `braindb_bench` (never `braindb`) -# - Postgres runs in its OWN container on host port 5433 (not 5432) +# - Postgres runs in its OWN container on host port 5434 (the user's +# personal Postgres is already on 5433, so bench takes 5434 to avoid +# a host-port collision) # - The bench BrainDB API is on host port 8001 (not 8000) # - The bench watcher mounts ./data_bench/ as /app/data, so it never # sees files in ./data/ (the personal watcher's directory) @@ -36,7 +38,9 @@ services: # literal "braindb_bench" substring in DATABASE_URL. POSTGRES_PASSWORD: bench_local_only ports: - - "5433:5432" + # Personal Postgres is already on 5433 in this environment; bench + # takes 5434 to avoid a host-port collision. Verified before commit. + - "5434:5432" volumes: - braindb_bench_pgdata:/var/lib/postgresql/data networks: From a6677554376f9ff5d6be06e2a4a8a2ad025c59aa Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 01:54:44 +0100 Subject: [PATCH 03/17] feat(bench): Phase 1 BEAM harness (adapter, runner, judge, score, warmup) Adapter renders BEAM rows to markdown for the existing watcher pipeline. Runner orchestrates per-conversation reset -> ingest -> warmup -> answer -> record, with a hard-coded DATABASE_URL safety assertion (must contain literal "braindb_bench" or refuses to proceed). Judge imports upstream BEAM's unified_llm_judge_base_prompt VERBATIM from the git submodule (zero prompt content in our codebase), calls Qwen via OpenAI-compat, replicates upstream's per-nugget 0/0.5/1.0 scoring. Score reporter aggregates + writes results.md with the Qwen-judge caveat. First smoke run found + fixed two harness bugs (both folded into this commit): - Watcher is non-recursive AND ingest endpoint hardcodes the file path prefix "data/sources/ingested/..." (braindb/ingest_watcher.py:278). Writing adapter output to a beam/ subdir fails ingest with a 404. Adapter now writes top-level data_bench/sources/; bench-conv filenames (e.g. beam_100k_conv_001.md) are unique enough at top level. - Warmup with settle_seconds=30 was too tight for slow chunk processing: declared "settled" between datasource creation and first extracted fact (a 3-5 min gap on big documents). Bumped default to 180s. First smoke also measured BrainDB extraction throughput on workstation Qwen at ~4.5 min/chunk (1M conversation, chunk 1: 5 facts + ~20 relations + per-edge LLM scoring). Implies BEAM-1M ~77h/conv (infeasible in the 2-3 day budget) and BEAM-100K ~9h/conv (at the edge). 100K --limit 1 smoke running now to confirm one full end-to-end pipeline. Status: harness is complete and verified to ingest end-to-end. The extraction throughput is a separate v0.4.0-cost / v0.5.0-evaluation trade-off to surface and decide on next session (e.g. ablate per-edge scoring at bench time, or pivot to LoCoMo with smaller conversations). --- benchmarks/__init__.py | 0 benchmarks/beam/__init__.py | 9 + benchmarks/beam/adapter.py | 212 ++++++++++++++++ benchmarks/beam/bench.py | 423 +++++++++++++++++++++++++++++++ benchmarks/beam/config.py | 87 +++++++ benchmarks/beam/judge.py | 382 ++++++++++++++++++++++++++++ benchmarks/beam/requirements.txt | 13 + benchmarks/beam/score.py | 148 +++++++++++ benchmarks/beam/warmup.py | 184 ++++++++++++++ 9 files changed, 1458 insertions(+) create mode 100644 benchmarks/__init__.py create mode 100644 benchmarks/beam/__init__.py create mode 100644 benchmarks/beam/adapter.py create mode 100644 benchmarks/beam/bench.py create mode 100644 benchmarks/beam/config.py create mode 100644 benchmarks/beam/judge.py create mode 100644 benchmarks/beam/requirements.txt create mode 100644 benchmarks/beam/score.py create mode 100644 benchmarks/beam/warmup.py diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/benchmarks/beam/__init__.py b/benchmarks/beam/__init__.py new file mode 100644 index 0000000..55d3445 --- /dev/null +++ b/benchmarks/beam/__init__.py @@ -0,0 +1,9 @@ +"""BrainDB on BEAM benchmark harness. + +Standards-compliant: upstream BEAM (dataset + eval script + judge prompt) +is invoked verbatim via the git submodule at ``benchmarks/beam/upstream/``. +Our code (this package) is the adapter only: dataset -> BrainDB ingest, +question -> BrainDB /agent/query, our answers -> upstream eval. + +See ``benchmarks/beam/README.md`` for the full trust model. +""" diff --git a/benchmarks/beam/adapter.py b/benchmarks/beam/adapter.py new file mode 100644 index 0000000..afea073 --- /dev/null +++ b/benchmarks/beam/adapter.py @@ -0,0 +1,212 @@ +"""BEAM dataset -> BrainDB ingest adapter. + +For each conversation in the BEAM dataset (Mohammadta/BEAM on HuggingFace), +render the full chat history as a single markdown file under +``data_bench/sources/beam/.md``. The bench watcher picks +it up and runs the same production extraction pipeline that ingests any +other document (same chunker, same fact extraction, same wiki maintainer). + +The probing-questions JSON (the actual benchmark questions + rubrics) is +NOT ingested — those are loaded separately by the bench runner and +asked via ``/api/v1/agent/query`` after the warmup barrier passes. + +Also exposes a small helper API the runner imports: + + iter_conversations(split: str) -> Iterator[Conversation] + write_conversation_md(conv, output_dir: Path) -> Path + load_probing_questions(conv) -> dict[str, list[dict]] + +CLI: + + python -m benchmarks.beam.adapter --split 1M --output data_bench/sources/beam/ + python -m benchmarks.beam.adapter --split 1M --index 0 # dry-run, prints to stdout +""" +from __future__ import annotations + +import argparse +import ast +import os +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Iterator + +REPO_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_CACHE = REPO_ROOT / "benchmarks" / "beam" / "dataset_cache" +DEFAULT_OUTPUT = REPO_ROOT / "data_bench" / "sources" / "beam" + + +@dataclass +class Conversation: + """One BEAM row, surfaced as a small typed wrapper.""" + split: str # "100K" / "500K" / "1M" + index: int # row index within the split + conversation_id: str # BEAM's own id (e.g. "1") + raw: dict # the underlying HF row + + @property + def slug(self) -> str: + """Stable filename component, e.g. 'beam_1m_conv_001'.""" + return f"beam_{self.split.lower()}_conv_{int(self.conversation_id):03d}" + + +def _hf_env() -> None: + """Point HuggingFace at the gitignored bench cache.""" + cache = str(DEFAULT_CACHE) + os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + os.environ.setdefault("HF_DATASETS_CACHE", cache) + os.environ.setdefault("HF_HOME", cache) + + +def iter_conversations(split: str = "1M") -> Iterator[Conversation]: + """Yield each conversation in the requested BEAM split.""" + _hf_env() + from datasets import load_dataset # type: ignore + ds = load_dataset("Mohammadta/BEAM") + if split not in ds: + raise ValueError(f"unknown split {split!r}; available: {list(ds.keys())}") + for i, row in enumerate(ds[split]): + yield Conversation( + split=split, + index=i, + conversation_id=str(row["conversation_id"]), + raw=row, + ) + + +def load_probing_questions(conv: Conversation) -> dict[str, list[dict]]: + """Parse ``probing_questions`` field (Python-dict-as-string) into a + plain dict keyed by category (abstention, multi_session_reasoning, etc.). + """ + pq = conv.raw["probing_questions"] + if isinstance(pq, dict): + return pq + return ast.literal_eval(pq) + + +def _user_profile_summary(profile: dict) -> str: + """Extract a short human-readable summary from BEAM's user_profile dict. + The dict contains keys like 'user_info' with multi-line strings; we keep + the first ~10 lines so the agent has persona context without the markdown + file becoming dominated by profile prose. + """ + info = profile.get("user_info", "") if isinstance(profile, dict) else str(profile) + lines = [ln.rstrip() for ln in info.split("\n") if ln.strip()][:12] + return "\n".join(lines) + + +def _render_message(msg: dict) -> str: + """Render one chat message as a markdown turn.""" + role = msg.get("role", "unknown").lower() + content = (msg.get("content") or "").rstrip() + # BEAM messages have 'time_anchor' per message; we surface it on user + # turns only, to anchor temporal-reasoning questions without doubling + # the prefix on every assistant reply. + ts = msg.get("time_anchor") + if role == "user": + header = f"**User** ({ts}):" if ts else "**User:**" + elif role == "assistant": + header = "**Assistant:**" + else: + header = f"**{role.title()}:**" + return f"{header}\n\n{content}" + + +def render_conversation_md(conv: Conversation) -> str: + """Render the entire conversation as a single markdown document. + + Format: a short frontmatter-style header (id, category, subtopics, + profile snippet), then one section per batch (with the batch's time + anchor), then all messages in that batch as alternating User/Assistant + blocks. This is the SAME shape ingested by BrainDB's watcher in + production — no benchmark-only extraction prompts, no special handling. + """ + raw = conv.raw + seed = raw.get("conversation_seed") or {} + profile = raw.get("user_profile") or {} + chat_batches = raw.get("chat") or [] + + parts: list[str] = [] + parts.append(f"# BEAM conversation {conv.slug}") + parts.append("") + parts.append(f"- conversation_id: {conv.conversation_id}") + parts.append(f"- split: {conv.split}") + if seed.get("category"): + parts.append(f"- category: {seed['category']}") + subtopics = seed.get("subtopics") or [] + if subtopics: + parts.append(f"- subtopics: {', '.join(map(str, subtopics))}") + profile_text = _user_profile_summary(profile) + if profile_text: + parts.append("") + parts.append("## User profile") + parts.append("") + parts.append(profile_text) + + for batch_idx, batch in enumerate(chat_batches, start=1): + if not batch: + continue + ts = batch[0].get("time_anchor") if isinstance(batch[0], dict) else None + parts.append("") + parts.append(f"## Batch {batch_idx}" + (f" — {ts}" if ts else "")) + parts.append("") + for msg in batch: + if not isinstance(msg, dict) or not (msg.get("content") or "").strip(): + continue + parts.append(_render_message(msg)) + parts.append("") + + # Ensure file ends with a single trailing newline + text = "\n".join(parts).rstrip() + "\n" + # Normalise stray Windows-style line endings just in case the dataset + # carries any (won't change semantics; keeps the watcher's chunker happy) + text = re.sub(r"\r\n?", "\n", text) + return text + + +def write_conversation_md(conv: Conversation, output_dir: Path = DEFAULT_OUTPUT) -> Path: + """Render and write one conversation; return the file path.""" + output_dir.mkdir(parents=True, exist_ok=True) + out_path = output_dir / f"{conv.slug}.md" + text = render_conversation_md(conv) + out_path.write_text(text, encoding="utf-8") + return out_path + + +def _cli() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--split", default="1M", choices=["100K", "500K", "1M"]) + p.add_argument("--output", default=str(DEFAULT_OUTPUT), + help=f"output directory (default: {DEFAULT_OUTPUT})") + p.add_argument("--index", type=int, default=None, + help="if set, dry-run: render only this row index to stdout, do not write") + p.add_argument("--limit", type=int, default=None, + help="if set, write only the first N conversations from the split") + args = p.parse_args() + + out_dir = Path(args.output) + + if args.index is not None: + # Dry run: print one rendered conversation to stdout. + convs = list(iter_conversations(args.split)) + if args.index >= len(convs): + print(f"index {args.index} out of range for split {args.split} (size {len(convs)})") + return 1 + text = render_conversation_md(convs[args.index]) + print(text) + return 0 + + count = 0 + for conv in iter_conversations(args.split): + if args.limit is not None and count >= args.limit: + break + path = write_conversation_md(conv, out_dir) + size_kb = path.stat().st_size / 1024 + print(f"wrote {path} ({size_kb:.1f} KB)") + count += 1 + print(f"\nwrote {count} conversation(s) to {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(_cli()) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py new file mode 100644 index 0000000..b972938 --- /dev/null +++ b/benchmarks/beam/bench.py @@ -0,0 +1,423 @@ +"""BrainDB on BEAM bench runner. + +Per-conversation lifecycle (zero interleaving within a conversation): + + Phase A — reset bench DB + write conversation .md to data_bench/sources/beam/ + Phase B — warmup wait (extraction + wiki pipeline drains) + Phase C — answer all of this conversation's probing questions via /agent/query + Phase D — record answers, probing_questions, warmup_stats, meta to runs//conv_/ + +Strict safety: every destructive op (DB reset, file write to data_bench/) +is gated on the literal substring ``braindb_bench`` appearing in the active +``DATABASE_URL`` — never touches the personal braindb database. + +CLI: + + python -m benchmarks.beam.bench --split 1M --limit 3 # smoke (3 convs) + python -m benchmarks.beam.bench --split 1M # full (35 convs) + python -m benchmarks.beam.bench --split 1M --limit 3 --fail-fast +""" +from __future__ import annotations + +import argparse +import datetime as dt +import json +import os +import subprocess +import sys +import time +from pathlib import Path +from typing import Iterable + +import psycopg2 +import requests + +from benchmarks.beam.adapter import ( + Conversation, + iter_conversations, + load_probing_questions, + write_conversation_md, +) +from benchmarks.beam.config import ( + BENCH_API_BASE, + BENCH_DATABASE_URL, + DATA_BENCH_SOURCES, + RUNS_DIR, + assert_bench_database_url, +) +from benchmarks.beam.warmup import wait_for_warmup + +# Tables that get truncated per conversation. alembic_version is NOT in this +# list — we want migrations to stay applied. Everything else is wiped. +_TRUNCATE_SQL = """ +TRUNCATE TABLE + entities, + relations, + wikis_ext, + wiki_job, + facts_ext, + thoughts_ext, + sources_ext, + datasources_ext, + rules_ext, + activity_log +RESTART IDENTITY CASCADE; +""" + +# Per-question agent timeout. BrainDB's /agent/query can take a while on +# local Qwen — multi-turn agent loops + extraction can be 30s to several +# minutes per question. 10 minutes hard ceiling per question. +QUESTION_TIMEOUT_SECONDS = 600 + + +# ----------------------------- helpers --------------------------------------- + +def _now_iso() -> str: + return dt.datetime.now(dt.timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _short_ts() -> str: + return dt.datetime.now(dt.timezone.utc).strftime("%Y%m%d_%H%M%S") + + +def _git_sha() -> str: + try: + out = subprocess.check_output( + ["git", "rev-parse", "--short", "HEAD"], + stderr=subprocess.DEVNULL, + cwd=Path(__file__).resolve().parents[2], + ) + return out.decode().strip() + except Exception: + return "unknown" + + +def _git_dirty() -> bool: + try: + out = subprocess.check_output( + ["git", "status", "--porcelain"], + stderr=subprocess.DEVNULL, + cwd=Path(__file__).resolve().parents[2], + ) + return bool(out.strip()) + except Exception: + return False + + +# ----------------------------- preflight ------------------------------------- + +def check_bench_api_healthy(base: str = BENCH_API_BASE, timeout: float = 5) -> None: + try: + r = requests.get(f"{base}/health", timeout=timeout) + r.raise_for_status() + body = r.json() + except Exception as e: + raise RuntimeError( + f"bench API not reachable at {base} ({e}). " + f"Did you run: docker compose -f docker-compose.bench.yml up -d ?" + ) + if body.get("status") != "ok": + raise RuntimeError(f"bench API unhealthy: {body}") + + +# ----------------------------- DB reset -------------------------------------- + +def reset_bench_db() -> None: + """TRUNCATE all data tables in the bench DB. Schema (alembic_version) + is preserved so the API stays connected and migrations stay applied. + """ + assert_bench_database_url() + conn = psycopg2.connect(BENCH_DATABASE_URL) + conn.autocommit = True + try: + with conn.cursor() as cur: + cur.execute(_TRUNCATE_SQL) + finally: + conn.close() + + +def _clear_sources_dir() -> None: + """Remove any stale .md files (and their .error.txt sidecars) left in + the bench watcher's source dir and its ingested/ + failed/ siblings. + Without this, files from previous conversations would either be + re-ingested by the next warmup OR poison the warmup barrier by leaving + pending content in the watch dir. + """ + DATA_BENCH_SOURCES.mkdir(parents=True, exist_ok=True) + for f in DATA_BENCH_SOURCES.glob("*.md"): + try: + f.unlink() + except OSError: + pass + # The watcher creates ingested/ + failed/ as siblings of WATCH_DIR + # (see ingest_watcher.py lines 39-40). Clean both so cumulative runs + # don't pile up on disk and so a previous failed file doesn't trick + # the next run into thinking ingest already happened. + for sibling in ("ingested", "failed"): + d = DATA_BENCH_SOURCES / sibling + if d.exists(): + for f in d.iterdir(): + if f.is_file(): + try: + f.unlink() + except OSError: + pass + + +# ----------------------------- ask agent ------------------------------------- + +def answer_one_question( + question_text: str, + base: str = BENCH_API_BASE, + timeout: float = QUESTION_TIMEOUT_SECONDS, +) -> tuple[str, dict]: + """POST /agent/query with the question; return (answer_text, raw_payload). + + Returns the answer text and the full payload (so we can capture tool + counts, latency, etc.) for the per-question record. Errors get caught + in the caller and recorded as `[ERROR: ...]` so the run continues. + """ + started = time.monotonic() + r = requests.post( + f"{base}/api/v1/agent/query", + json={"query": question_text}, + timeout=timeout, + ) + elapsed = time.monotonic() - started + r.raise_for_status() + payload = r.json() + payload["_elapsed_seconds"] = round(elapsed, 2) + # BrainDB's /agent/query response shape: {"answer": "..."} at minimum. + return payload.get("answer", ""), payload + + +# ----------------------------- per-conv lifecycle ---------------------------- + +def run_one_conversation( + conv: Conversation, + run_dir: Path, + *, + warmup_timeout: float = 1800, + warmup_settle_seconds: float = 180, + question_timeout: float = QUESTION_TIMEOUT_SECONDS, +) -> dict: + """Full A-B-C-D lifecycle for one conversation. Returns a small stats dict.""" + conv_dir = run_dir / f"conv_{int(conv.conversation_id):03d}" + conv_dir.mkdir(parents=True, exist_ok=True) + + started = time.monotonic() + + # ---- Phase A: reset + ingest ---- + print(f"[A] reset bench DB + write conversation markdown ...", flush=True) + reset_bench_db() + _clear_sources_dir() + md_path = write_conversation_md(conv, DATA_BENCH_SOURCES) + md_kb = md_path.stat().st_size / 1024 + print(f"[A done] {md_path.name} ({md_kb:.0f} KB) dropped at {_now_iso()}", flush=True) + + # ---- Phase B: warmup wait ---- + print(f"[B] warmup wait ...", flush=True) + warmup_stats = wait_for_warmup( + settle_seconds=warmup_settle_seconds, + timeout_seconds=warmup_timeout, + verbose=True, + ) + print(f"[B done] {warmup_stats}", flush=True) + + # ---- Phase C: answer all questions ---- + probing = load_probing_questions(conv) + print(f"[C] answering {sum(len(v) for v in probing.values())} questions " + f"across {len(probing)} categories ...", flush=True) + answers_by_category: dict[str, list[dict]] = {} + raw_payloads_by_category: dict[str, list[dict]] = {} + question_errors = 0 + + for category, questions in probing.items(): + category_answers: list[dict] = [] + category_payloads: list[dict] = [] + for i, q in enumerate(questions, start=1): + question_text = q["question"] + try: + ans, payload = answer_one_question( + question_text, timeout=question_timeout + ) + print( + f" [{category} {i}/{len(questions)}] OK " + f"({payload.get('_elapsed_seconds', '?')}s, " + f"{len(ans)} chars)", + flush=True, + ) + except Exception as e: + ans = f"[ERROR: {type(e).__name__}: {e}]" + payload = {"_error": str(e)} + question_errors += 1 + print(f" [{category} {i}/{len(questions)}] ERROR: {e}", flush=True) + category_answers.append({"question": question_text, "llm_response": ans}) + category_payloads.append({"question": question_text, "payload": payload}) + answers_by_category[category] = category_answers + raw_payloads_by_category[category] = category_payloads + + # ---- Phase D: record artefacts ---- + print(f"[D] recording artefacts to {conv_dir} ...", flush=True) + conv_dir.joinpath("answers.json").write_text( + json.dumps(answers_by_category, indent=2, ensure_ascii=False), encoding="utf-8" + ) + conv_dir.joinpath("probing_questions.json").write_text( + json.dumps(probing, indent=2, ensure_ascii=False), encoding="utf-8" + ) + conv_dir.joinpath("warmup_stats.json").write_text( + json.dumps(warmup_stats, indent=2), encoding="utf-8" + ) + conv_dir.joinpath("raw_payloads.json").write_text( + json.dumps(raw_payloads_by_category, indent=2, ensure_ascii=False), + encoding="utf-8", + ) + conv_meta = { + "conversation_id": conv.conversation_id, + "split": conv.split, + "slug": conv.slug, + "category": (conv.raw.get("conversation_seed") or {}).get("category"), + "ingested_md_size_kb": int(md_kb), + "question_errors": question_errors, + "wall_clock_s": round(time.monotonic() - started, 1), + "started_at": _now_iso(), + } + conv_dir.joinpath("conversation_meta.json").write_text( + json.dumps(conv_meta, indent=2), encoding="utf-8" + ) + + print(f"[done] conv {conv.slug} in {conv_meta['wall_clock_s']:.0f}s, " + f"errors={question_errors}\n", flush=True) + return conv_meta + + +# ----------------------------- run-level ------------------------------------- + +def _generate_run_id(split: str, limit: int | None) -> str: + sha = _git_sha() + dirty = "+dirty" if _git_dirty() else "" + n = f"_n{limit}" if limit is not None else "" + return f"{_short_ts()}_{split}{n}_{sha}{dirty}" + + +def _save_run_config(run_dir: Path, args: argparse.Namespace) -> None: + cfg = { + "split": args.split, + "limit": args.limit, + "fail_fast": args.fail_fast, + "warmup_timeout": args.warmup_timeout, + "warmup_settle_seconds": args.warmup_settle_seconds, + "question_timeout": args.question_timeout, + "started_at": _now_iso(), + "git_sha": _git_sha(), + "git_dirty": _git_dirty(), + "bench_api_base": BENCH_API_BASE, + "python_version": sys.version, + } + run_dir.joinpath("run_config.json").write_text( + json.dumps(cfg, indent=2), encoding="utf-8" + ) + + +def _conversations_to_run(args: argparse.Namespace) -> Iterable[Conversation]: + convs = list(iter_conversations(args.split)) + if args.limit is not None: + convs = convs[: args.limit] + if args.only_ids: + wanted = set(s.strip() for s in args.only_ids.split(",")) + convs = [c for c in convs if c.conversation_id in wanted] + return convs + + +# ----------------------------- CLI ------------------------------------------- + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + p.add_argument("--split", default="1M", choices=["100K", "500K", "1M"]) + p.add_argument("--limit", type=int, default=None, + help="run only the first N conversations from the split") + p.add_argument("--only-ids", default=None, + help="comma-separated list of conversation_ids to run (overrides --limit ordering)") + p.add_argument("--fail-fast", action="store_true", + help="abort the whole run if any single conversation raises") + p.add_argument("--warmup-timeout", type=float, default=1800, + help="seconds before warmup gives up on convergence") + p.add_argument("--warmup-settle-seconds", type=float, default=180, + help="seconds of quiet on entities before declaring warmup clear " + "(default 180; big enough to span the gap between datasource " + "creation and the first extracted fact for slow chunk processing)") + p.add_argument("--question-timeout", type=float, default=QUESTION_TIMEOUT_SECONDS, + help="per-question HTTP timeout on /agent/query") + return p.parse_args() + + +def main() -> int: + args = _parse_args() + assert_bench_database_url() + check_bench_api_healthy() + + convs = list(_conversations_to_run(args)) + if not convs: + print("no conversations selected; exiting") + return 1 + + run_id = _generate_run_id(args.split, args.limit if not args.only_ids else None) + run_dir = RUNS_DIR / run_id + run_dir.mkdir(parents=True, exist_ok=True) + _save_run_config(run_dir, args) + + print(f"=== BEAM bench run: {run_id} ===") + print(f"split={args.split} conversations={len(convs)} run_dir={run_dir}") + print(f"bench API: {BENCH_API_BASE}") + print(f"git SHA: {_git_sha()}{' (dirty)' if _git_dirty() else ''}") + print() + + metas: list[dict] = [] + run_started = time.monotonic() + for i, conv in enumerate(convs, start=1): + print(f"--- Conversation {i}/{len(convs)} ({conv.slug}, " + f"category={(conv.raw.get('conversation_seed') or {}).get('category')}) ---") + try: + meta = run_one_conversation( + conv, + run_dir, + warmup_timeout=args.warmup_timeout, + warmup_settle_seconds=args.warmup_settle_seconds, + question_timeout=args.question_timeout, + ) + metas.append(meta) + except Exception as e: + print(f"!!! conversation {conv.slug} failed: {type(e).__name__}: {e}", + file=sys.stderr, flush=True) + metas.append({ + "conversation_id": conv.conversation_id, + "slug": conv.slug, + "fatal_error": f"{type(e).__name__}: {e}", + }) + if args.fail_fast: + _write_run_summary(run_dir, metas, time.monotonic() - run_started) + return 1 + + _write_run_summary(run_dir, metas, time.monotonic() - run_started) + print(f"\n=== run complete: {run_id} ===") + return 0 + + +def _write_run_summary(run_dir: Path, metas: list[dict], wall_clock_s: float) -> None: + summary = { + "finished_at": _now_iso(), + "wall_clock_s": round(wall_clock_s, 1), + "conversations_attempted": len(metas), + "conversations_succeeded": sum(1 for m in metas if "fatal_error" not in m), + "total_question_errors": sum(int(m.get("question_errors", 0)) for m in metas), + "per_conv_meta": metas, + } + run_dir.joinpath("run_summary.json").write_text( + json.dumps(summary, indent=2), encoding="utf-8" + ) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/beam/config.py b/benchmarks/beam/config.py new file mode 100644 index 0000000..81ef41e --- /dev/null +++ b/benchmarks/beam/config.py @@ -0,0 +1,87 @@ +"""Bench-only configuration constants. + +Reads from env vars with safe defaults that match ``docker-compose.bench.yml``. +The hard-coded safety sentinel (``braindb_bench``) is the load-bearing piece +here: every destructive op gates on it. +""" +from __future__ import annotations + +import os +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +# Top-level sources dir, NOT a beam/ subdir. The watcher is non-recursive +# (line 318 of braindb/ingest_watcher.py uses WATCH_DIR.iterdir()) and the +# ingest endpoint hardcodes "data/sources/ingested/..." as the file path it +# expects (line 278), so a subdirectory triggers a 404 on ingest. Bench +# files at top-level are still uniquely named (e.g. beam_1m_conv_001.md) +# so they never collide with anything else in this dedicated dir. +DATA_BENCH_SOURCES = REPO_ROOT / "data_bench" / "sources" +ANSWERS_DIR = REPO_ROOT / "benchmarks" / "beam" / "answers" +RUNS_DIR = REPO_ROOT / "benchmarks" / "beam" / "runs" + +# ---- Bench DB ---------------------------------------------------------------- +# Hosted by docker-compose.bench.yml's `postgres_bench` service. From the host +# the bench Postgres is reachable on port 5434 (5433 is the personal Postgres +# in this environment). Inside the bench Docker network it's `postgres_bench:5432`. +BENCH_DATABASE_URL = os.getenv( + "BENCH_DATABASE_URL", + "postgresql://braindb_bench:bench_local_only@localhost:5434/braindb_bench", +) + +# The literal substring every bench DATABASE_URL must contain. If a destructive +# op is about to fire and the active URL does NOT contain this string, the +# code refuses to proceed. This is the load-bearing safety check. +BENCH_DB_SENTINEL = "braindb_bench" + +# ---- Bench BrainDB API ------------------------------------------------------ +BENCH_API_BASE = os.getenv("BENCH_API_BASE", "http://localhost:8001") + +# ---- Judge LLM (OpenAI-compatible endpoint) --------------------------------- +# Bench defaults to the workstation Qwen via vLLM on the SSH tunnel. The judge +# is fully outside BrainDB; the bench API is independent of which judge runs. +QWEN_BASE_URL = os.getenv("QWEN_BASE_URL", "http://localhost:8010/v1") +QWEN_MODEL = os.getenv("QWEN_MODEL", "") # empty -> let vLLM's served model resolve +QWEN_API_KEY = os.getenv("QWEN_API_KEY", "EMPTY") # vLLM defaults to no auth + + +def _mask_password(url: str) -> str: + """Return a printable form of the URL with the password redacted.""" + if "@" not in url or "://" not in url: + return url + scheme_and_creds, host_path = url.split("@", 1) + if ":" not in scheme_and_creds.split("//", 1)[1]: + return url + head, creds = scheme_and_creds.split("//", 1) + user, _ = creds.split(":", 1) + return f"{head}//{user}:***@{host_path}" + + +def assert_bench_database_url(url: str = BENCH_DATABASE_URL) -> None: + """REFUSE to proceed unless URL contains the bench sentinel literally. + + This is intentionally a literal substring check, NOT a hostname/database + parse, so that misconfigured env vars (typos, copy-paste mishaps) fail + closed rather than falling through to the personal DB. + """ + if BENCH_DB_SENTINEL not in url: + raise RuntimeError( + f"REFUSING to proceed: connection string must contain " + f"{BENCH_DB_SENTINEL!r} as a safety check against operating on " + f"the personal braindb database. Got: {_mask_password(url)}" + ) + + +__all__ = [ + "REPO_ROOT", + "DATA_BENCH_SOURCES", + "ANSWERS_DIR", + "RUNS_DIR", + "BENCH_DATABASE_URL", + "BENCH_DB_SENTINEL", + "BENCH_API_BASE", + "QWEN_BASE_URL", + "QWEN_MODEL", + "QWEN_API_KEY", + "assert_bench_database_url", +] diff --git a/benchmarks/beam/judge.py b/benchmarks/beam/judge.py new file mode 100644 index 0000000..7e08117 --- /dev/null +++ b/benchmarks/beam/judge.py @@ -0,0 +1,382 @@ +"""BEAM judge runner. + +For each (question, llm_response) pair our bench produced, call the +configured LLM judge with **the exact upstream BEAM judge prompt**, then +parse the per-nugget score and aggregate per question / per category. + +The judge prompt itself is imported VERBATIM from the upstream submodule: + + from prompts import unified_llm_judge_base_prompt + +No prompt text lives in our codebase. Anyone wanting to verify can diff +``benchmarks/beam/upstream/src/prompts.py`` against the same SHA on +``github.com/mohammadtavakoli78/BEAM`` — same prompt, same logic. + +What we DO NOT do here (intentionally, to keep deps light): +* invoke upstream's ``compute_metrics.py``, which depends on spacy, + sentence-transformers, scipy, nltk, rouge_score etc. (~50 heavy deps). + BEAM's BLEU/ROUGE/cosine-similarity helpers are useful for some + question types but the headline score is the LLM judge's nugget + average, and that's what we replicate exactly here. + +If you want to invoke upstream's evaluator directly (full standards +compliance with all metrics), install ``benchmarks/beam/upstream/requirements.txt`` +into a separate venv and run:: + + python -m src.evaluation.run_evaluation \ + --probing_questions \ + --answers_directory \ + --output + +Our ``answers.json`` files are already in the exact format their evaluator +expects. + +CLI:: + + python -m benchmarks.beam.judge --run-dir benchmarks/beam/runs/ + python -m benchmarks.beam.judge --run-dir --only-conv 1 +""" +from __future__ import annotations + +import argparse +import json +import re +import sys +import time +from pathlib import Path + +import requests + +from benchmarks.beam.config import ( + QWEN_API_KEY, + QWEN_BASE_URL, + QWEN_MODEL, + REPO_ROOT, +) + +# ---- Import the upstream judge prompt VERBATIM ------------------------------- + +_UPSTREAM_SRC = REPO_ROOT / "benchmarks" / "beam" / "upstream" / "src" +sys.path.insert(0, str(_UPSTREAM_SRC)) +try: + from prompts import unified_llm_judge_base_prompt # type: ignore # noqa: E402 +finally: + # Don't pollute sys.path for any sibling imports + try: + sys.path.remove(str(_UPSTREAM_SRC)) + except ValueError: + pass + + +# Sanity check at import time: the upstream prompt must contain the three +# placeholders we replace. If upstream changes them, we want to fail loudly, +# not silently send a wrong prompt to the judge. +for _ph in ("", "", ""): + if _ph not in unified_llm_judge_base_prompt: + raise RuntimeError( + f"upstream judge prompt missing expected placeholder {_ph!r}; " + "submodule may have drifted from the SHA this code was written against." + ) + + +# ---- LLM call (OpenAI-compatible chat completions) -------------------------- + +def _resolve_qwen_model() -> str: + """If QWEN_MODEL is unset, ask vLLM what it's serving and use that.""" + if QWEN_MODEL: + return QWEN_MODEL + try: + r = requests.get(f"{QWEN_BASE_URL}/models", timeout=10, + headers={"Authorization": f"Bearer {QWEN_API_KEY}"}) + r.raise_for_status() + data = r.json().get("data", []) + if data: + return data[0]["id"] + except Exception: + pass + return "qwen" # final fallback — vLLM usually accepts any string + + +def _call_judge(prompt_text: str, model_id: str, timeout: float = 180) -> str: + payload = { + "model": model_id, + "messages": [{"role": "user", "content": prompt_text}], + "temperature": 0, + } + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {QWEN_API_KEY}", + } + r = requests.post( + f"{QWEN_BASE_URL}/chat/completions", + json=payload, + headers=headers, + timeout=timeout, + ) + r.raise_for_status() + body = r.json() + return body["choices"][0]["message"]["content"] + + +# ---- Verdict parsing -------------------------------------------------------- + +_SCORE_RE = re.compile(r'"score"\s*:\s*([0-9.]+)') + + +def _parse_judge_response(response_text: str) -> dict: + """Return ``{"score": float in {0.0, 0.5, 1.0}, "raw": str, ...}``. + + Upstream uses ``json_repair.repair_json`` to be robust to malformed JSON. + We try the obvious json.loads first, then a regex fallback. If neither + yields a parseable score, we record score=0.0 + the raw text so the + reviewer can spot-check. Upstream also clamps to {0, 0.5, 1.0}. + """ + text = (response_text or "").strip() + # Strip code-fence wrappers if present + if text.startswith("```"): + text = re.sub(r"^```[a-zA-Z]*\s*", "", text) + text = re.sub(r"\s*```\s*$", "", text) + + parsed = None + try: + parsed = json.loads(text) + except Exception: + m = _SCORE_RE.search(text) + if m: + parsed = {"score": float(m.group(1))} + + if not isinstance(parsed, dict) or "score" not in parsed: + return {"score": 0.0, "parse_error": True, "raw": text[:500]} + + try: + score = float(parsed["score"]) + except (TypeError, ValueError): + return {"score": 0.0, "parse_error": True, "raw": text[:500]} + + # Clamp to allowed values per BEAM rubric scheme (0 / 0.5 / 1.0) + if score <= 0.25: + score = 0.0 + elif score <= 0.75: + score = 0.5 + else: + score = 1.0 + + parsed["score"] = score + parsed["raw"] = text[:500] + return parsed + + +# ---- Per-conversation scoring ----------------------------------------------- + +def _judge_one_question( + question: str, + llm_response: str, + rubric: list[str], + model_id: str, + timeout: float, +) -> dict: + nugget_results = [] + cumulative = 0.0 + parse_errors = 0 + for nugget in rubric: + prompt_text = ( + unified_llm_judge_base_prompt + .replace("", question) + .replace("", nugget) + .replace("", llm_response) + ) + try: + raw = _call_judge(prompt_text, model_id, timeout=timeout) + verdict = _parse_judge_response(raw) + except Exception as e: + verdict = {"score": 0.0, "error": f"{type(e).__name__}: {e}"} + if verdict.get("parse_error"): + parse_errors += 1 + nugget_results.append({"nugget": nugget, "verdict": verdict}) + cumulative += float(verdict.get("score", 0.0)) + avg = cumulative / len(rubric) if rubric else 0.0 + return { + "nugget_count": len(rubric), + "nugget_average_score": round(avg, 4), + "parse_errors": parse_errors, + "nuggets": nugget_results, + } + + +def judge_conversation(conv_dir: Path, model_id: str, timeout: float = 180) -> dict: + """Score every (answer, rubric) pair in one conv subdirectory. + + Reads ``answers.json`` + ``probing_questions.json``. Writes + ``verdicts.json`` (per-question detail) + ``category_scores.json`` + (per-category aggregates) into the same conv subdir. Returns the + aggregate dict. + """ + answers = json.loads((conv_dir / "answers.json").read_text(encoding="utf-8")) + probing = json.loads((conv_dir / "probing_questions.json").read_text(encoding="utf-8")) + + started = time.monotonic() + by_category: dict[str, list[dict]] = {} + category_scores: dict[str, dict] = {} + + for category, questions in answers.items(): + rubrics = [item.get("rubric", []) for item in probing.get(category, [])] + per_q = [] + for i, ans_item in enumerate(questions): + question = ans_item["question"] + response = ans_item["llm_response"] + rubric = rubrics[i] if i < len(rubrics) else [] + if not rubric: + per_q.append({ + "question": question, + "skipped": True, + "reason": "no rubric for this index", + }) + continue + qstart = time.monotonic() + result = _judge_one_question(question, response, rubric, model_id, timeout) + result["question"] = question + result["llm_response_preview"] = (response or "")[:300] + result["elapsed_s"] = round(time.monotonic() - qstart, 2) + per_q.append(result) + print( + f" [{conv_dir.name} {category} {i+1}/{len(questions)}] " + f"score={result['nugget_average_score']:.2f} " + f"nuggets={result['nugget_count']} " + f"parse_err={result['parse_errors']} " + f"({result['elapsed_s']}s)", + flush=True, + ) + by_category[category] = per_q + non_skipped = [q for q in per_q if "skipped" not in q] + cat_avg = ( + sum(q["nugget_average_score"] for q in non_skipped) / len(non_skipped) + if non_skipped else 0.0 + ) + category_scores[category] = { + "question_count": len(non_skipped), + "average_score": round(cat_avg, 4), + } + + overall_questions = [ + q for cat in by_category.values() for q in cat if "skipped" not in q + ] + overall_avg = ( + sum(q["nugget_average_score"] for q in overall_questions) / len(overall_questions) + if overall_questions else 0.0 + ) + + aggregate = { + "overall_average_score": round(overall_avg, 4), + "total_questions_judged": len(overall_questions), + "categories": category_scores, + "model_id": model_id, + "judge_endpoint": QWEN_BASE_URL, + "wall_clock_s": round(time.monotonic() - started, 1), + } + + (conv_dir / "verdicts.json").write_text( + json.dumps(by_category, indent=2, ensure_ascii=False), encoding="utf-8" + ) + (conv_dir / "category_scores.json").write_text( + json.dumps(aggregate, indent=2), encoding="utf-8" + ) + return aggregate + + +# ---- CLI -------------------------------------------------------------------- + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--run-dir", required=True, + help="path to a run dir under benchmarks/beam/runs/") + p.add_argument("--only-conv", default=None, + help="comma-separated conversation_ids to judge (default: all)") + p.add_argument("--timeout", type=float, default=180, + help="per-judge-call HTTP timeout") + p.add_argument("--model", default=None, + help="override judge model id (otherwise auto-resolved from vLLM /v1/models)") + return p.parse_args() + + +def main() -> int: + args = _parse_args() + run_dir = Path(args.run_dir).resolve() + if not run_dir.exists(): + print(f"run dir not found: {run_dir}") + return 1 + + model_id = args.model or _resolve_qwen_model() + print(f"judge: {QWEN_BASE_URL} model={model_id}") + + only_ids = None + if args.only_conv: + only_ids = {s.strip() for s in args.only_conv.split(",")} + + conv_dirs = sorted(p for p in run_dir.iterdir() if p.is_dir() and p.name.startswith("conv_")) + if only_ids: + conv_dirs = [p for p in conv_dirs if p.name.split("_")[1].lstrip("0") in only_ids] + + if not conv_dirs: + print(f"no conversation subdirs found under {run_dir}") + return 1 + + aggregates: list[dict] = [] + for conv_dir in conv_dirs: + print(f"\n--- {conv_dir.name} ---") + try: + agg = judge_conversation(conv_dir, model_id, timeout=args.timeout) + aggregates.append({"conv": conv_dir.name, **agg}) + except Exception as e: + print(f"!!! failed: {e}") + aggregates.append({"conv": conv_dir.name, "fatal_error": str(e)}) + + summary = _aggregate_run(aggregates) + (run_dir / "judge_summary.json").write_text( + json.dumps(summary, indent=2), encoding="utf-8" + ) + print(f"\nwrote {run_dir / 'judge_summary.json'}") + print(f"overall: {summary['overall_average_score']:.4f} " + f"({summary['total_questions_judged']} questions)") + return 0 + + +def _aggregate_run(per_conv: list[dict]) -> dict: + """Aggregate per-conv aggregates into a run-level summary.""" + cat_sums: dict[str, dict] = {} + total_qs = 0 + overall_num = 0.0 + + for agg in per_conv: + if "fatal_error" in agg: + continue + total_qs += agg.get("total_questions_judged", 0) + overall_num += agg.get("overall_average_score", 0.0) * agg.get("total_questions_judged", 0) + for cat, cat_agg in agg.get("categories", {}).items(): + entry = cat_sums.setdefault(cat, {"question_count": 0, "weighted_score": 0.0}) + qc = cat_agg.get("question_count", 0) + entry["question_count"] += qc + entry["weighted_score"] += cat_agg.get("average_score", 0.0) * qc + + cat_scores = { + cat: { + "question_count": entry["question_count"], + "average_score": ( + round(entry["weighted_score"] / entry["question_count"], 4) + if entry["question_count"] else 0.0 + ), + } + for cat, entry in cat_sums.items() + } + + return { + "overall_average_score": round(overall_num / total_qs, 4) if total_qs else 0.0, + "total_questions_judged": total_qs, + "conversations_judged": sum(1 for a in per_conv if "fatal_error" not in a), + "conversations_attempted": len(per_conv), + "categories": cat_scores, + "per_conversation": per_conv, + } + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/beam/requirements.txt b/benchmarks/beam/requirements.txt new file mode 100644 index 0000000..41bbbb9 --- /dev/null +++ b/benchmarks/beam/requirements.txt @@ -0,0 +1,13 @@ +# Bench-only Python deps (host environment, NOT the BrainDB container). +# Install with: pip install -r benchmarks/beam/requirements.txt + +# HuggingFace dataset loader (BEAM dataset hosted at Mohammadta/BEAM) +datasets>=3.0,<4.0 +huggingface_hub>=0.25 + +# HTTP client for talking to the bench BrainDB API (:8001) and the +# Qwen OpenAI-compatible endpoint +requests>=2.32 + +# Psycopg2 binary for the bench-DB safety check + reset SQL +psycopg2-binary>=2.9 diff --git a/benchmarks/beam/score.py b/benchmarks/beam/score.py new file mode 100644 index 0000000..8f2b00f --- /dev/null +++ b/benchmarks/beam/score.py @@ -0,0 +1,148 @@ +"""BEAM run reporter. + +Reads ``judge_summary.json`` from a run dir and writes a human-readable +``results.md`` in the same dir. Aggregate scores are already computed by +``judge.py`` (using the upstream BEAM judge prompt verbatim); this script +just formats them for review. + +Optionally also invokes upstream's full ``run_evaluation`` if the user has +installed the upstream deps in their venv — this gives BLEU/ROUGE/cosine +metrics alongside the LLM-judge score. Off by default (heavy install). + +CLI:: + + python -m benchmarks.beam.score --run-dir benchmarks/beam/runs/ + python -m benchmarks.beam.score --run-dir --invoke-upstream +""" +from __future__ import annotations + +import argparse +import json +import subprocess +from pathlib import Path + + +def _fmt_pct(score: float) -> str: + return f"{score * 100:.2f}%" + + +def _render_markdown(summary: dict, run_config: dict | None) -> str: + lines: list[str] = [] + lines.append(f"# BEAM run — {_fmt_pct(summary['overall_average_score'])}") + lines.append("") + if run_config: + lines.append("## Run config") + lines.append("") + for k in ("split", "limit", "started_at", "git_sha", "git_dirty"): + if k in run_config: + lines.append(f"- **{k}**: `{run_config[k]}`") + lines.append("") + lines.append("## Headline") + lines.append("") + lines.append(f"- **Overall score** (LLM-judge nugget average, all questions): " + f"**{_fmt_pct(summary['overall_average_score'])}**") + lines.append(f"- **Questions judged**: {summary['total_questions_judged']}") + lines.append(f"- **Conversations**: {summary['conversations_judged']}" + f" / {summary['conversations_attempted']}") + lines.append("") + lines.append("> **Caveat**: This score uses Qwen 3.6-27B as the LLM judge " + "via the workstation vLLM endpoint. Published BEAM scores from " + "Hindsight (73.9% @ 1M) and mem0 (64.1% @ 1M) used Gemini-2.5-" + "flash-lite and GPT-4o respectively. Our score is NOT directly " + "comparable to those numbers until a cross-judge calibration " + "sample is published (see plan: Phase 3).") + lines.append("") + lines.append("## Per-category breakdown") + lines.append("") + lines.append("| Category | Questions | Average score |") + lines.append("|---|---|---|") + for cat, info in sorted(summary["categories"].items()): + lines.append(f"| {cat} | {info['question_count']} | " + f"{_fmt_pct(info['average_score'])} |") + lines.append("") + lines.append("## Per-conversation detail") + lines.append("") + lines.append("| Conversation | Questions | Score |") + lines.append("|---|---|---|") + for conv in summary["per_conversation"]: + if "fatal_error" in conv: + lines.append(f"| {conv['conv']} | — | **FATAL: {conv['fatal_error']}** |") + else: + lines.append( + f"| {conv['conv']} | {conv.get('total_questions_judged', 0)} | " + f"{_fmt_pct(conv.get('overall_average_score', 0.0))} |" + ) + lines.append("") + return "\n".join(lines).rstrip() + "\n" + + +def _invoke_upstream(run_dir: Path) -> int: + """Optionally invoke upstream's run_evaluation across all per-conv pairs. + Requires upstream/requirements.txt installed in the active environment. + """ + print("WARNING: invoking upstream run_evaluation requires " + "benchmarks/beam/upstream/requirements.txt installed. " + "This adds spacy + sentence-transformers + scipy etc. (~50 deps).") + upstream_src = run_dir.resolve().parents[2] / "upstream" / "src" + conv_dirs = sorted(p for p in run_dir.iterdir() if p.is_dir() and p.name.startswith("conv_")) + for conv_dir in conv_dirs: + ans = conv_dir / "answers.json" + pq = conv_dir / "probing_questions.json" + out = conv_dir / "upstream_eval_output.json" + if not ans.exists() or not pq.exists(): + print(f"skip {conv_dir.name}: missing answers.json or probing_questions.json") + continue + print(f"running upstream eval on {conv_dir.name} ...") + env_cmd = [ + "python", "-c", + "import sys; sys.path.insert(0, r'" + + str(upstream_src).replace("\\", "\\\\") + + "'); " + "from src.evaluation.run_evaluation import run_evaluation; " + "from src.llm import gpt_llm; " + "run_evaluation(" + f"probing_questions_address=r'{pq}', " + f"answers_directory=r'{ans}', " + f"output_address=r'{out}', " + "model=gpt_llm)" + ] + result = subprocess.run(env_cmd, capture_output=True, text=True) + if result.returncode != 0: + print(f" failed: {result.stderr[:400]}") + else: + print(f" -> {out}") + return 0 + + +def main() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--run-dir", required=True) + p.add_argument("--invoke-upstream", action="store_true", + help="also invoke upstream's run_evaluation (requires upstream deps)") + args = p.parse_args() + + run_dir = Path(args.run_dir).resolve() + summary_path = run_dir / "judge_summary.json" + if not summary_path.exists(): + print(f"no judge_summary.json in {run_dir}. Run `python -m benchmarks.beam.judge --run-dir {run_dir}` first.") + return 1 + + summary = json.loads(summary_path.read_text(encoding="utf-8")) + run_config_path = run_dir / "run_config.json" + run_config = json.loads(run_config_path.read_text(encoding="utf-8")) if run_config_path.exists() else None + + md = _render_markdown(summary, run_config) + results_path = run_dir / "results.md" + results_path.write_text(md, encoding="utf-8") + print(f"wrote {results_path}") + print() + print(md) + + if args.invoke_upstream: + _invoke_upstream(run_dir) + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/beam/warmup.py b/benchmarks/beam/warmup.py new file mode 100644 index 0000000..9e4b63d --- /dev/null +++ b/benchmarks/beam/warmup.py @@ -0,0 +1,184 @@ +"""Warmup barrier — wait for BrainDB's async pipeline to settle. + +After a conversation .md file is dropped in ``data_bench/sources/beam/``, +BEFORE the runner asks any of that conversation's questions, this module +blocks until BrainDB has finished thinking: + +* the watcher has consumed the file (sources/ is empty) +* no new entities have been created in the last ``settle_seconds`` (extraction done) +* the wiki_job queue is fully drained (status NOT IN pending/assigned) + +These three conditions must hold for ``consecutive_clear_required`` +consecutive polls before the barrier returns, to avoid flapping on a +single quiet poll that happened to fall between events. + +Asking questions before the barrier returns would evaluate a half-formed +memory state and produce a number that does not represent BrainDB's true +performance. This is the standard two-phase eval pattern. + +CLI usage: + + python -m benchmarks.beam.warmup + python -m benchmarks.beam.warmup --timeout 600 --settle-seconds 60 +""" +from __future__ import annotations + +import argparse +import time +from pathlib import Path + +import psycopg2 + +from benchmarks.beam.config import ( + BENCH_DATABASE_URL, + DATA_BENCH_SOURCES, + assert_bench_database_url, +) + + +def _query_one(conn, sql: str): + with conn.cursor() as cur: + cur.execute(sql) + return cur.fetchone() + + +def _entity_count(conn) -> int: + row = _query_one(conn, "SELECT COUNT(*) FROM entities") + return int(row[0]) if row else 0 + + +def _seconds_since_last_entity(conn) -> float | None: + row = _query_one( + conn, + "SELECT EXTRACT(EPOCH FROM (NOW() - MAX(created_at))) FROM entities", + ) + return float(row[0]) if row and row[0] is not None else None + + +def _pending_wiki_jobs(conn) -> int: + row = _query_one( + conn, + "SELECT COUNT(*) FROM wiki_job WHERE status IN ('pending','assigned')", + ) + return int(row[0]) if row else 0 + + +def _unprocessed_files() -> list[Path]: + if not DATA_BENCH_SOURCES.exists(): + return [] + return sorted(p for p in DATA_BENCH_SOURCES.glob("*.md") if p.is_file()) + + +def wait_for_warmup( + *, + settle_seconds: float = 180.0, + consecutive_clear_required: int = 2, + poll_interval: float = 5.0, + timeout_seconds: float = 1800.0, + log_interval: float = 10.0, + verbose: bool = True, +) -> dict: + """Block until ingest + wiki pipeline have drained for this conversation. + + Returns a small stats dict (wall_clock_s, entities, relations, wikis). + Raises TimeoutError if convergence does not happen within timeout. + """ + assert_bench_database_url() + + conn = psycopg2.connect(BENCH_DATABASE_URL) + conn.autocommit = True + + start = time.monotonic() + deadline = start + timeout_seconds + consecutive_clear = 0 + last_log = -log_interval # force first iteration to log + + try: + while time.monotonic() < deadline: + files_remaining = len(_unprocessed_files()) + entity_age = _seconds_since_last_entity(conn) + pending = _pending_wiki_jobs(conn) + entity_count = _entity_count(conn) + + if entity_age is None: + # No entities yet. If files are still in sources/, watcher + # hasn't started; if not, it may have just finished — give + # it a tiny grace period before declaring "no work to do". + clear = files_remaining == 0 and (time.monotonic() - start) > 10 + else: + clear = ( + files_remaining == 0 + and entity_age >= settle_seconds + and pending == 0 + ) + + elapsed = time.monotonic() - start + if verbose and (elapsed - last_log >= log_interval): + age_str = f"{entity_age:.0f}s" if entity_age is not None else "n/a" + print( + f"[warmup t={elapsed:6.0f}s] files_left={files_remaining} " + f"entities={entity_count} last_entity_age={age_str} " + f"pending_wiki={pending} clear={clear} " + f"({consecutive_clear}/{consecutive_clear_required})", + flush=True, + ) + last_log = elapsed + + if clear: + consecutive_clear += 1 + if consecutive_clear >= consecutive_clear_required: + return _final_stats(conn, start) + else: + consecutive_clear = 0 + + time.sleep(poll_interval) + + raise TimeoutError( + f"warmup did not converge within {timeout_seconds:.0f}s " + f"(files_left={len(_unprocessed_files())}, " + f"last_entity_age={_seconds_since_last_entity(conn)}, " + f"pending_wiki={_pending_wiki_jobs(conn)})" + ) + finally: + conn.close() + + +def _final_stats(conn, start: float) -> dict: + return { + "wall_clock_s": round(time.monotonic() - start, 1), + "entities": _entity_count(conn), + "relations": int(_query_one(conn, "SELECT COUNT(*) FROM relations")[0]), + "wikis": int( + _query_one(conn, "SELECT COUNT(*) FROM entities WHERE entity_type='wiki'")[0] + ), + "wiki_jobs_done": int( + _query_one(conn, "SELECT COUNT(*) FROM wiki_job WHERE status='done'")[0] + ), + } + + +def _cli() -> int: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--timeout", type=float, default=1800) + p.add_argument("--settle-seconds", type=float, default=180, + help="seconds of no INSERT on entities before declaring extraction settled " + "(default 180; large enough to span the gap between datasource creation " + "and first extracted fact for big documents)") + p.add_argument("--poll-interval", type=float, default=5) + p.add_argument("--consecutive-clear", type=int, default=2) + p.add_argument("--quiet", action="store_true") + args = p.parse_args() + + stats = wait_for_warmup( + settle_seconds=args.settle_seconds, + consecutive_clear_required=args.consecutive_clear, + poll_interval=args.poll_interval, + timeout_seconds=args.timeout, + verbose=not args.quiet, + ) + print(f"\nwarmup complete: {stats}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(_cli()) From 9e06aa5fbb0f3fcb98ae40a023dcc50fabe2629b Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 19:09:12 +0100 Subject: [PATCH 04/17] fix(bench): per-conversation DBs + warmup no longer blocks on wiki queue MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Recovery from the 10h timeout that lost the first BEAM-100K run. Two bench-side problems, both purely on bench code (no BrainDB changes): 1. Warmup blocked on `pending_wiki == 0`. The wiki maintainer queues candidates faster than the writer drains them on big documents (grew from 24 pending at t=3min to 1701 at t=10h on the lost run). The queue never converged. Phase C never ran. 2. Per-conversation isolation via TRUNCATE meant each conversation's memory state was destroyed when moving to the next. No way to inspect what happened in conv_007 after the fact. Fix: - warmup.py: wait_for_warmup gains block_on_wiki_queue (default False) and database_url (default falls back to BENCH_DATABASE_URL). The new default drops `pending == 0` from the `clear` condition. Wikis continue async in the background while Phase C answers questions; recall queries on existing entities don't need an empty wiki queue. CLI gains --wait-for-wikis as the opt-in for the strict variant. - bench.py: replaces reset_bench_db() with three helpers: create_conv_db(conv_id) -> CREATE DATABASE braindb_conv_NNN restart_api_with_db(db_url) -> docker compose up -d --force-recreate api_bench with BENCH_DATABASE_URL overridden wait_for_api_healthy() -> poll /health until ready Each conversation now gets its own Postgres database that persists for inspection after the run. The api_bench container restarts with a new DATABASE_URL between conversations (~30-60s embedding reload per restart × N convs is negligible vs ingest wall-clock). - config.py: safety sentinel becomes a tuple ("braindb_bench", "braindb_conv_"). assert_bench_database_url accepts either pattern; rejects bare "braindb" (personal) as before. Adds BENCH_ADMIN_DATABASE_URL for connecting to Postgres's maintenance DB (CREATE / DROP). - docker-compose.bench.yml: * api_bench's DATABASE_URL becomes ${BENCH_DATABASE_URL:-...default} so the runner can override per conversation * adds new bench_runner service (profile: "runner") that has the docker CLI installed + docker socket mounted, so it can recreate api_bench between conversations. Profile means it doesn't auto-start with `docker compose up`; explicit: docker compose -f docker-compose.bench.yml run --rm bench_runner \ python -m benchmarks.beam.bench --split 100K --limit 1 - benchmarks/beam/runner.Dockerfile (new): small image for the bench_runner service. python:3.12-slim + docker.io + datasets + huggingface_hub + psycopg2-binary + requests. Repo code itself comes via bind mount. BrainDB code (braindb/) is completely untouched. Central review will still fail on large documents (it puts all facts in a single agent query, which exceeds the 10K char limit). That's a BrainDB design choice, accepted as-is for this session. Watcher logs a warning, extraction succeeds, facts land, the bench proceeds. The BrainDB central-review redesign is a separate focused session — to be validated by re-running the bench and measuring the delta vs this number. --- benchmarks/beam/bench.py | 190 +++++++++++++++++++++++------- benchmarks/beam/config.py | 32 +++-- benchmarks/beam/runner.Dockerfile | 37 ++++++ benchmarks/beam/warmup.py | 30 ++++- docker-compose.bench.yml | 60 +++++++++- 5 files changed, 294 insertions(+), 55 deletions(-) create mode 100644 benchmarks/beam/runner.Dockerfile diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index b972938..9951a88 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -2,20 +2,30 @@ Per-conversation lifecycle (zero interleaving within a conversation): - Phase A — reset bench DB + write conversation .md to data_bench/sources/beam/ - Phase B — warmup wait (extraction + wiki pipeline drains) - Phase C — answer all of this conversation's probing questions via /agent/query - Phase D — record answers, probing_questions, warmup_stats, meta to runs//conv_/ - -Strict safety: every destructive op (DB reset, file write to data_bench/) -is gated on the literal substring ``braindb_bench`` appearing in the active -``DATABASE_URL`` — never touches the personal braindb database. + Phase A — create new Postgres database `braindb_conv_NNN`, restart + api_bench with DATABASE_URL pointing at it, wait for healthy, + write conversation .md to data_bench/sources/ + Phase B — warmup wait (extraction settles; wikis run async in background) + Phase C — answer this conversation's 20 probing questions via /agent/query + Phase D — record answers + probing_questions + warmup_stats + meta to + runs//conv_/ + +After all conversations: 20 (or 35) databases sit in postgres_bench, each +fully self-contained and inspectable (`docker exec braindb_bench_postgres +psql -U braindb_bench -d braindb_conv_007 -c '\\dt'`). + +Strict safety: every destructive op (CREATE DATABASE, DROP DATABASE, etc.) +is gated on the literal substring `braindb_bench` OR `braindb_conv_` +appearing in the URL. The personal `braindb` database is never touched. CLI: - python -m benchmarks.beam.bench --split 1M --limit 3 # smoke (3 convs) - python -m benchmarks.beam.bench --split 1M # full (35 convs) + python -m benchmarks.beam.bench --split 100K --limit 1 # smoke (1 conv) + python -m benchmarks.beam.bench --split 100K # full 100K (20 convs) python -m benchmarks.beam.bench --split 1M --limit 3 --fail-fast + +Intended to run from inside the `bench_runner` container (which has the +docker CLI installed + docker socket mounted so it can recreate api_bench). """ from __future__ import annotations @@ -47,22 +57,22 @@ ) from benchmarks.beam.warmup import wait_for_warmup -# Tables that get truncated per conversation. alembic_version is NOT in this -# list — we want migrations to stay applied. Everything else is wiped. -_TRUNCATE_SQL = """ -TRUNCATE TABLE - entities, - relations, - wikis_ext, - wiki_job, - facts_ext, - thoughts_ext, - sources_ext, - datasources_ext, - rules_ext, - activity_log -RESTART IDENTITY CASCADE; -""" +# Admin URL + DB-URL base for per-conversation database creation. The +# bench_runner service in docker-compose.bench.yml supplies these via env; +# fall back to sensible defaults that match the bench network. +_ADMIN_DB_URL = os.getenv( + "BENCH_ADMIN_DATABASE_URL", + "postgresql://braindb_bench:bench_local_only@postgres_bench:5432/postgres", +) +_DB_BASE_URL = os.getenv( + "BENCH_DB_BASE_URL", + "postgresql://braindb_bench:bench_local_only@postgres_bench:5432", +) + +# docker-compose file path inside the bench_runner container. The compose +# file is mounted via the repo bind mount; `COMPOSE_FILE` env var hints +# the same path so `docker compose ...` picks it up automatically. +_COMPOSE_FILE = os.getenv("COMPOSE_FILE", "/app/docker-compose.bench.yml") # Per-question agent timeout. BrainDB's /agent/query can take a while on # local Qwen — multi-turn agent loops + extraction can be 30s to several @@ -120,20 +130,98 @@ def check_bench_api_healthy(base: str = BENCH_API_BASE, timeout: float = 5) -> N raise RuntimeError(f"bench API unhealthy: {body}") -# ----------------------------- DB reset -------------------------------------- +# ----------------------------- per-conv DB orchestration --------------------- + +def _conv_db_name(conv_id: int) -> str: + """Return the database name for a given conversation_id (zero-padded).""" + return f"braindb_conv_{int(conv_id):03d}" + -def reset_bench_db() -> None: - """TRUNCATE all data tables in the bench DB. Schema (alembic_version) - is preserved so the API stays connected and migrations stay applied. +def _conv_db_url(conv_id: int) -> str: + """Full DATABASE_URL for the conversation's database.""" + return f"{_DB_BASE_URL}/{_conv_db_name(conv_id)}" + + +def create_conv_db(conv_id: int) -> str: + """Create a fresh Postgres database `braindb_conv_NNN`. + + If a database with the same name exists from a prior run, drop + recreate + it (the per-conv DB is supposed to be a clean slate). Returns the full + connection URL for the new database. + + Safety: hits postgres_bench via the admin connection (`postgres` maintenance + DB) — never the user's personal `braindb`. """ - assert_bench_database_url() - conn = psycopg2.connect(BENCH_DATABASE_URL) - conn.autocommit = True + db_name = _conv_db_name(conv_id) + db_url = _conv_db_url(conv_id) + assert_bench_database_url(db_url) + admin = psycopg2.connect(_ADMIN_DB_URL) + admin.autocommit = True try: - with conn.cursor() as cur: - cur.execute(_TRUNCATE_SQL) + with admin.cursor() as cur: + # Terminate any leftover backend on this DB before drop (otherwise + # DROP fails with "is being accessed by other users"). + cur.execute( + "SELECT pg_terminate_backend(pid) FROM pg_stat_activity " + "WHERE datname = %s AND pid <> pg_backend_pid()", + (db_name,), + ) + cur.execute(f'DROP DATABASE IF EXISTS "{db_name}"') + cur.execute(f'CREATE DATABASE "{db_name}" OWNER braindb_bench') finally: - conn.close() + admin.close() + return db_url + + +def restart_api_with_db(database_url: str) -> None: + """Force-recreate the api_bench container with BENCH_DATABASE_URL set + to the given URL. The container's startup runs alembic migrations on + the new (empty) DB, then starts uvicorn. + + Uses `docker compose up -d --force-recreate api_bench` from inside the + bench_runner container (which has the docker CLI + socket mounted). + """ + assert_bench_database_url(database_url) + env = os.environ.copy() + env["BENCH_DATABASE_URL"] = database_url + result = subprocess.run( + [ + "docker", "compose", + "-f", _COMPOSE_FILE, + "up", "-d", "--force-recreate", "--no-deps", "api_bench", + ], + env=env, + capture_output=True, + text=True, + ) + if result.returncode != 0: + raise RuntimeError( + f"docker compose up -d --force-recreate api_bench failed:\n" + f" stdout: {result.stdout[:500]}\n stderr: {result.stderr[:500]}" + ) + + +def wait_for_api_healthy(timeout: float = 180, poll: float = 2) -> None: + """Block until the bench api responds with status=ok. + + The api needs to (1) connect to the new DB, (2) run alembic migrations, + (3) load the embedding model into memory. On first start this is ~30-60s; + on subsequent restarts (model already cached) it can still be ~20-30s. + """ + start = time.monotonic() + last_err = None + while time.monotonic() - start < timeout: + try: + r = requests.get(f"{BENCH_API_BASE}/health", timeout=3) + if r.status_code == 200 and r.json().get("status") == "ok": + return + last_err = f"HTTP {r.status_code}" + except Exception as e: + last_err = str(e) + time.sleep(poll) + raise TimeoutError( + f"bench api did not become healthy within {timeout:.0f}s (last: {last_err})" + ) def _clear_sources_dir() -> None: @@ -200,20 +288,30 @@ def run_one_conversation( warmup_timeout: float = 1800, warmup_settle_seconds: float = 180, question_timeout: float = QUESTION_TIMEOUT_SECONDS, + block_on_wiki_queue: bool = False, ) -> dict: """Full A-B-C-D lifecycle for one conversation. Returns a small stats dict.""" conv_dir = run_dir / f"conv_{int(conv.conversation_id):03d}" conv_dir.mkdir(parents=True, exist_ok=True) started = time.monotonic() - - # ---- Phase A: reset + ingest ---- - print(f"[A] reset bench DB + write conversation markdown ...", flush=True) - reset_bench_db() + conv_db_name = _conv_db_name(int(conv.conversation_id)) + conv_db_url = _conv_db_url(int(conv.conversation_id)) + + # ---- Phase A: fresh DB + restart api + ingest ---- + print(f"[A] create DB '{conv_db_name}' + restart api + write conversation ...", + flush=True) + create_conv_db(int(conv.conversation_id)) + restart_api_with_db(conv_db_url) + wait_for_api_healthy() _clear_sources_dir() md_path = write_conversation_md(conv, DATA_BENCH_SOURCES) md_kb = md_path.stat().st_size / 1024 - print(f"[A done] {md_path.name} ({md_kb:.0f} KB) dropped at {_now_iso()}", flush=True) + print( + f"[A done] api restarted on {conv_db_name}; " + f"{md_path.name} ({md_kb:.0f} KB) dropped at {_now_iso()}", + flush=True, + ) # ---- Phase B: warmup wait ---- print(f"[B] warmup wait ...", flush=True) @@ -221,6 +319,8 @@ def run_one_conversation( settle_seconds=warmup_settle_seconds, timeout_seconds=warmup_timeout, verbose=True, + block_on_wiki_queue=block_on_wiki_queue, + database_url=conv_db_url, ) print(f"[B done] {warmup_stats}", flush=True) @@ -281,6 +381,7 @@ def run_one_conversation( "question_errors": question_errors, "wall_clock_s": round(time.monotonic() - started, 1), "started_at": _now_iso(), + "bench_db_name": conv_db_name, } conv_dir.joinpath("conversation_meta.json").write_text( json.dumps(conv_meta, indent=2), encoding="utf-8" @@ -348,6 +449,12 @@ def _parse_args() -> argparse.Namespace: help="seconds of quiet on entities before declaring warmup clear " "(default 180; big enough to span the gap between datasource " "creation and the first extracted fact for slow chunk processing)") + p.add_argument("--wait-for-wikis", action="store_true", + help="Strict warmup mode: also wait for the wiki_job queue to be " + "fully drained before answering. Off by default because the " + "wiki writer can be slower than the maintainer queues for " + "large documents, so the queue may never converge. Wikis " + "continue async in the background regardless.") p.add_argument("--question-timeout", type=float, default=QUESTION_TIMEOUT_SECONDS, help="per-question HTTP timeout on /agent/query") return p.parse_args() @@ -386,6 +493,7 @@ def main() -> int: warmup_timeout=args.warmup_timeout, warmup_settle_seconds=args.warmup_settle_seconds, question_timeout=args.question_timeout, + block_on_wiki_queue=args.wait_for_wikis, ) metas.append(meta) except Exception as e: diff --git a/benchmarks/beam/config.py b/benchmarks/beam/config.py index 81ef41e..91764de 100644 --- a/benchmarks/beam/config.py +++ b/benchmarks/beam/config.py @@ -29,10 +29,21 @@ "postgresql://braindb_bench:bench_local_only@localhost:5434/braindb_bench", ) -# The literal substring every bench DATABASE_URL must contain. If a destructive -# op is about to fire and the active URL does NOT contain this string, the -# code refuses to proceed. This is the load-bearing safety check. -BENCH_DB_SENTINEL = "braindb_bench" +# Admin URL — connects to Postgres's built-in `postgres` maintenance DB so we +# can CREATE / DROP per-conversation databases. Same credentials as bench DB. +BENCH_ADMIN_DATABASE_URL = os.getenv( + "BENCH_ADMIN_DATABASE_URL", + "postgresql://braindb_bench:bench_local_only@localhost:5434/postgres", +) + +# Allowed DB-name patterns. Every URL passed to destructive ops MUST contain one +# of these substrings; otherwise we refuse (protection against the personal +# `braindb` database). Two patterns: +# "braindb_bench" — the shared bench DB (legacy / fallback) +# "braindb_conv_" — per-conversation databases (braindb_conv_001, _002, ...) +BENCH_DB_SENTINELS = ("braindb_bench", "braindb_conv_") +# Kept for backwards compatibility (some tests / scripts may import this name). +BENCH_DB_SENTINEL = BENCH_DB_SENTINELS[0] # ---- Bench BrainDB API ------------------------------------------------------ BENCH_API_BASE = os.getenv("BENCH_API_BASE", "http://localhost:8001") @@ -58,16 +69,19 @@ def _mask_password(url: str) -> str: def assert_bench_database_url(url: str = BENCH_DATABASE_URL) -> None: - """REFUSE to proceed unless URL contains the bench sentinel literally. + """REFUSE to proceed unless URL contains one of the bench sentinels. This is intentionally a literal substring check, NOT a hostname/database parse, so that misconfigured env vars (typos, copy-paste mishaps) fail closed rather than falling through to the personal DB. + + Allowed patterns: ``braindb_bench`` (shared bench DB) or ``braindb_conv_`` + (per-conversation databases like ``braindb_conv_001``). """ - if BENCH_DB_SENTINEL not in url: + if not any(s in url for s in BENCH_DB_SENTINELS): raise RuntimeError( - f"REFUSING to proceed: connection string must contain " - f"{BENCH_DB_SENTINEL!r} as a safety check against operating on " + f"REFUSING to proceed: connection string must contain one of " + f"{BENCH_DB_SENTINELS!r} as a safety check against operating on " f"the personal braindb database. Got: {_mask_password(url)}" ) @@ -78,7 +92,9 @@ def assert_bench_database_url(url: str = BENCH_DATABASE_URL) -> None: "ANSWERS_DIR", "RUNS_DIR", "BENCH_DATABASE_URL", + "BENCH_ADMIN_DATABASE_URL", "BENCH_DB_SENTINEL", + "BENCH_DB_SENTINELS", "BENCH_API_BASE", "QWEN_BASE_URL", "QWEN_MODEL", diff --git a/benchmarks/beam/runner.Dockerfile b/benchmarks/beam/runner.Dockerfile new file mode 100644 index 0000000..b1cf2ea --- /dev/null +++ b/benchmarks/beam/runner.Dockerfile @@ -0,0 +1,37 @@ +# benchmarks/beam/runner.Dockerfile +# +# Small orchestrator image for the bench. Lives separately from BrainDB's +# main image because: +# - it needs the docker CLI (to recreate api_bench between conversations) +# - it only needs bench-side Python deps (no FastAPI, no embeddings, no +# pgvector — just enough to drive the bench) +# +# Built only when running the bench: +# docker compose -f docker-compose.bench.yml --profile runner build bench_runner +# Run with: +# docker compose -f docker-compose.bench.yml run --rm bench_runner \ +# python -m benchmarks.beam.bench --split 100K --limit 1 + +FROM python:3.12-slim + +# Docker CLI lets the runner do `docker compose up -d --force-recreate api_bench` +# between conversations to swap DATABASE_URL. ca-certificates + curl are +# for TLS + simple health probes. +RUN apt-get update \ + && apt-get install -y --no-install-recommends docker.io curl ca-certificates \ + && rm -rf /var/lib/apt/lists/* + +# Bench Python deps. Repo code itself comes in via the `.:/app` bind mount; +# we just need the runtime libraries here. +RUN pip install --no-cache-dir \ + datasets \ + huggingface_hub \ + requests \ + psycopg2-binary + +WORKDIR /app +ENV PYTHONPATH=/app + +# Default: sit idle so `docker compose run --rm bench_runner ` can +# spawn ephemeral instances with whatever command we want. +CMD ["sleep", "infinity"] diff --git a/benchmarks/beam/warmup.py b/benchmarks/beam/warmup.py index 9e4b63d..d2d15ca 100644 --- a/benchmarks/beam/warmup.py +++ b/benchmarks/beam/warmup.py @@ -77,15 +77,31 @@ def wait_for_warmup( timeout_seconds: float = 1800.0, log_interval: float = 10.0, verbose: bool = True, + block_on_wiki_queue: bool = False, + database_url: str | None = None, ) -> dict: - """Block until ingest + wiki pipeline have drained for this conversation. + """Block until ingest is done (and optionally wiki pipeline drained). + + By default (``block_on_wiki_queue=False``) the barrier waits only for + extraction to settle (no new entity INSERTs for ``settle_seconds`` AND + no files remaining in ``data_bench/sources/``). Wiki processing runs + asynchronously in the background while Phase C asks questions; that is + correct behaviour because recall queries on existing entities don't + need the wiki queue to be empty. + + Set ``block_on_wiki_queue=True`` (CLI ``--wait-for-wikis``) for the old + strict behaviour that also waits for ``wiki_job`` to be fully drained. + That convention was a bench-side mistake that turned the wiki layer + into a single-point-of-failure for large documents where the maintainer + queues candidates faster than the writer drains them. Returns a small stats dict (wall_clock_s, entities, relations, wikis). Raises TimeoutError if convergence does not happen within timeout. """ - assert_bench_database_url() + url = database_url or BENCH_DATABASE_URL + assert_bench_database_url(url) - conn = psycopg2.connect(BENCH_DATABASE_URL) + conn = psycopg2.connect(url) conn.autocommit = True start = time.monotonic() @@ -109,7 +125,7 @@ def wait_for_warmup( clear = ( files_remaining == 0 and entity_age >= settle_seconds - and pending == 0 + and (not block_on_wiki_queue or pending == 0) ) elapsed = time.monotonic() - start @@ -167,6 +183,11 @@ def _cli() -> int: p.add_argument("--poll-interval", type=float, default=5) p.add_argument("--consecutive-clear", type=int, default=2) p.add_argument("--quiet", action="store_true") + p.add_argument("--wait-for-wikis", action="store_true", + help="Strict mode: also require wiki_job queue empty before " + "declaring warmup done. Off by default because the wiki " + "writer can be slower than the maintainer queues, so the " + "queue may never converge for large documents.") args = p.parse_args() stats = wait_for_warmup( @@ -175,6 +196,7 @@ def _cli() -> int: poll_interval=args.poll_interval, timeout_seconds=args.timeout, verbose=not args.quiet, + block_on_wiki_queue=args.wait_for_wikis, ) print(f"\nwarmup complete: {stats}") return 0 diff --git a/docker-compose.bench.yml b/docker-compose.bench.yml index 8f2dec7..f4e2e66 100644 --- a/docker-compose.bench.yml +++ b/docker-compose.bench.yml @@ -61,8 +61,10 @@ services: networks: - bench-network environment: - # MUST contain "braindb_bench" or the bench runner refuses to start. - DATABASE_URL: postgresql://braindb_bench:bench_local_only@postgres_bench:5432/braindb_bench + # MUST contain "braindb_bench" OR "braindb_conv_" or the bench runner + # refuses to start. The bench runner overrides this via the BENCH_DATABASE_URL + # env var when it creates per-conversation databases (braindb_conv_001, ...). + DATABASE_URL: ${BENCH_DATABASE_URL:-postgresql://braindb_bench:bench_local_only@postgres_bench:5432/braindb_bench} API_PORT: 8001 HF_TOKEN: ${HF_TOKEN:-} # Bench defaults to local Qwen via the workstation tunnel — matches the @@ -128,6 +130,60 @@ services: - .:/app command: python -m braindb.wiki_scheduler + # The bench_runner container is the orchestrator: it creates per-conversation + # databases, restarts api_bench with each new DATABASE_URL, drops the + # conversation .md into the watcher's dir, waits for warmup, asks the + # probing questions, records answers. Uses the same image as api_bench (so + # it has Python + BrainDB deps) plus the docker CLI installed at startup + # and the docker socket mounted so it can recreate api_bench. + # + # Profile "runner" means it does NOT auto-start with `docker compose up`; + # invoke it explicitly via: + # docker compose -f docker-compose.bench.yml run --rm bench_runner \ + # python -m benchmarks.beam.bench --split 100K --limit 1 + # bench_runner — orchestrator for the BEAM bench. Creates per-conversation + # Postgres databases, restarts api_bench with each new DATABASE_URL, + # drops the .md into the watcher's dir, waits for warmup, asks the probing + # questions. See benchmarks/beam/runner.Dockerfile. + # + # Profile "runner" means it does NOT auto-start with `docker compose up`. + # Invoke it explicitly: + # docker compose -f docker-compose.bench.yml run --rm bench_runner \ + # python -m benchmarks.beam.bench --split 100K --limit 1 + bench_runner: + build: + context: . + dockerfile: benchmarks/beam/runner.Dockerfile + container_name: braindb_bench_runner + profiles: ["runner"] + depends_on: + postgres_bench: + condition: service_healthy + networks: + - bench-network + environment: + # Bench-side runner reaches the api via Docker network DNS, no host port. + BENCH_API_BASE: http://api_bench:8001 + # Per-conversation DB creation uses admin creds against postgres_bench's + # built-in `postgres` maintenance DB. + BENCH_ADMIN_DATABASE_URL: postgresql://braindb_bench:bench_local_only@postgres_bench:5432/postgres + # bench.py mints per-conv URLs by appending `/braindb_conv_NNN` to this. + BENCH_DB_BASE_URL: postgresql://braindb_bench:bench_local_only@postgres_bench:5432 + # Workstation Qwen via Docker host gateway. + QWEN_BASE_URL: ${QWEN_BASE_URL_BENCH:-http://host.docker.internal:8010/v1} + QWEN_API_KEY: ${VLLM_API_KEY:-EMPTY} + HF_TOKEN: ${HF_TOKEN:-} + # Docker CLI invocations need to know which compose file + project. + COMPOSE_FILE: /app/docker-compose.bench.yml + COMPOSE_PROJECT_NAME: braindb_bench + extra_hosts: + - "host.docker.internal:host-gateway" + volumes: + - .:/app + - ./data_bench:/app/data + # Docker socket lets the runner recreate api_bench between conversations. + - /var/run/docker.sock:/var/run/docker.sock + volumes: braindb_bench_pgdata: name: braindb_bench_pgdata From cc3b0680103f859877c2112d7ce18459e4d977c9 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 21:42:20 +0100 Subject: [PATCH 05/17] fix(bench): docker Python SDK for api_bench recreate + correct CLI install Two follow-up refinements on top of 9e06aa5, surfaced when actually running the bench on the workstation: 1. restart_api_with_db now uses the Docker Python SDK instead of subprocess(docker compose up -d --force-recreate). The subprocess path broke because `docker compose -f docker-compose.bench.yml` resolves the relative bind-mount paths (`.:/app`, `./data_bench:/app/data`) against the bench_runner container's filesystem, not the host's, leaving the recreated api_bench with empty mounts and alembic failing on missing script_location. The SDK path inspects the existing container's already- host-resolved HostConfig.Binds and preserves them on recreate. 2. runner.Dockerfile: the apt `docker.io` package on python:3.12-slim only ships `docker-init`, not the `docker` CLI itself (FileNotFoundError on `docker` invocation). Switched to downloading the official static CLI binary from download.docker.com plus the compose plugin from the docker/compose GitHub release. Also adds the `docker` Python SDK to the pip install list (used by the new restart_api_with_db). BrainDB code (braindb/) remains completely untouched. Pure bench-harness follow-up. --- benchmarks/beam/bench.py | 97 ++++++++++++++++++++++++------- benchmarks/beam/runner.Dockerfile | 34 ++++++++--- 2 files changed, 101 insertions(+), 30 deletions(-) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index 9951a88..7cdbf89 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -174,32 +174,87 @@ def create_conv_db(conv_id: int) -> str: def restart_api_with_db(database_url: str) -> None: - """Force-recreate the api_bench container with BENCH_DATABASE_URL set - to the given URL. The container's startup runs alembic migrations on - the new (empty) DB, then starts uvicorn. - - Uses `docker compose up -d --force-recreate api_bench` from inside the - bench_runner container (which has the docker CLI + socket mounted). + """Recreate the api_bench container with DATABASE_URL set to ``database_url``. + + We use the Docker Python SDK rather than ``docker compose`` because bench.py + runs inside the bench_runner container. ``docker compose -f .../docker-compose.bench.yml + up`` would resolve the compose file's relative paths (e.g. ``.:/app``) against + the runner's view of the filesystem, not the host's. The recreated api_bench + would then bind-mount the runner's /app (which is the workstation host dir, + but only because of OUR bind mount — fragile and wrong from the daemon's POV). + + The SDK path captures the existing api_bench container's config (image, + binds, networks, ports, etc.) and recreates it with that config plus the + updated ``DATABASE_URL`` env var. """ assert_bench_database_url(database_url) - env = os.environ.copy() - env["BENCH_DATABASE_URL"] = database_url - result = subprocess.run( - [ - "docker", "compose", - "-f", _COMPOSE_FILE, - "up", "-d", "--force-recreate", "--no-deps", "api_bench", - ], - env=env, - capture_output=True, - text=True, - ) - if result.returncode != 0: + import docker as docker_sdk # imported lazily so the module loads on hosts + # that don't have the SDK installed + client = docker_sdk.from_env() + try: + existing = client.containers.get("braindb_bench_api") + except docker_sdk.errors.NotFound: # type: ignore[attr-defined] raise RuntimeError( - f"docker compose up -d --force-recreate api_bench failed:\n" - f" stdout: {result.stdout[:500]}\n stderr: {result.stderr[:500]}" + "braindb_bench_api container not found — bring the compose stack " + "up first: `docker compose -f docker-compose.bench.yml up -d`" ) + cfg = existing.attrs + image = cfg["Config"]["Image"] + host_cfg = cfg["HostConfig"] + net_settings = cfg["NetworkSettings"] + + # Preserve existing volumes (host paths the daemon already resolved when + # the container was first created via `docker compose up`). + binds = list(host_cfg.get("Binds") or []) + + # Networks: rejoin all networks the container was on. Use the first as the + # primary `network=` arg, then attach the rest after start. + networks = list((net_settings.get("Networks") or {}).keys()) + + # Ports: preserve the existing host port mapping. + port_bindings = {} + for port_proto, mappings in (host_cfg.get("PortBindings") or {}).items(): + if mappings: + port_bindings[port_proto] = mappings[0]["HostPort"] + + # Env: replace DATABASE_URL (or add if missing). + env_list = list(cfg["Config"].get("Env") or []) + new_env = [e for e in env_list if not e.startswith("DATABASE_URL=")] + new_env.append(f"DATABASE_URL={database_url}") + + # Stop + remove existing. + existing.stop(timeout=30) + existing.remove() + + # Recreate with same shape + new env. + container = client.containers.run( + image, + name="braindb_bench_api", + detach=True, + environment=new_env, + volumes=binds, + ports=port_bindings, + network=networks[0] if networks else None, + extra_hosts={ + h.split(":", 1)[0]: h.split(":", 1)[1] + for h in (host_cfg.get("ExtraHosts") or []) + } or None, + restart_policy=( + {"Name": host_cfg["RestartPolicy"]["Name"]} + if host_cfg.get("RestartPolicy") and host_cfg["RestartPolicy"].get("Name") + else None + ), + command=cfg["Config"].get("Cmd"), + working_dir=cfg["Config"].get("WorkingDir") or None, + ) + # Attach to additional networks (run() only takes one). + for net_name in networks[1:]: + try: + client.networks.get(net_name).connect(container) + except docker_sdk.errors.NotFound: # type: ignore[attr-defined] + pass + def wait_for_api_healthy(timeout: float = 180, poll: float = 2) -> None: """Block until the bench api responds with status=ok. diff --git a/benchmarks/beam/runner.Dockerfile b/benchmarks/beam/runner.Dockerfile index b1cf2ea..01e1d7b 100644 --- a/benchmarks/beam/runner.Dockerfile +++ b/benchmarks/beam/runner.Dockerfile @@ -2,7 +2,8 @@ # # Small orchestrator image for the bench. Lives separately from BrainDB's # main image because: -# - it needs the docker CLI (to recreate api_bench between conversations) +# - it needs the docker CLI + compose plugin (to recreate api_bench +# between conversations) # - it only needs bench-side Python deps (no FastAPI, no embeddings, no # pgvector — just enough to drive the bench) # @@ -14,20 +15,35 @@ FROM python:3.12-slim -# Docker CLI lets the runner do `docker compose up -d --force-recreate api_bench` -# between conversations to swap DATABASE_URL. ca-certificates + curl are -# for TLS + simple health probes. -RUN apt-get update \ - && apt-get install -y --no-install-recommends docker.io curl ca-certificates \ - && rm -rf /var/lib/apt/lists/* +# Install docker CLI + compose plugin directly from upstream static builds. +# The apt `docker.io` package on python:3.12-slim only ships docker-init, +# not the `docker` CLI itself, so we grab the official static binaries. +ARG DOCKER_VERSION=27.3.1 +ARG COMPOSE_VERSION=v2.30.3 +RUN set -eux \ + && apt-get update \ + && apt-get install -y --no-install-recommends curl ca-certificates \ + && curl -fsSL "https://download.docker.com/linux/static/stable/x86_64/docker-${DOCKER_VERSION}.tgz" \ + | tar -xzC /usr/local/bin --strip-components=1 docker/docker \ + && chmod +x /usr/local/bin/docker \ + && mkdir -p /usr/local/lib/docker/cli-plugins \ + && curl -fsSL "https://github.com/docker/compose/releases/download/${COMPOSE_VERSION}/docker-compose-linux-x86_64" \ + -o /usr/local/lib/docker/cli-plugins/docker-compose \ + && chmod +x /usr/local/lib/docker/cli-plugins/docker-compose \ + && rm -rf /var/lib/apt/lists/* \ + && docker --version \ + && docker compose version # Bench Python deps. Repo code itself comes in via the `.:/app` bind mount; -# we just need the runtime libraries here. +# we just need the runtime libraries here. `docker` is the Python SDK used +# by bench.py to recreate api_bench between conversations (avoids the +# compose-from-inside-container relative-path resolution issue). RUN pip install --no-cache-dir \ datasets \ huggingface_hub \ requests \ - psycopg2-binary + psycopg2-binary \ + docker WORKDIR /app ENV PYTHONPATH=/app From c185ee08422f9790e0296789346bf867d43ede25 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 21:46:51 +0100 Subject: [PATCH 06/17] feat(bench): per-chunk relation creation + 1200-word chunks + 40K endpoint cap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Five surgical changes that together eliminate central_review's HTTP 422 on large documents and roughly halve the per-conversation wall-clock. All in one revertible commit; central_review function body kept dormant for easy restore. 1. routers/agent.py — AgentQueryRequest.query.max_length 10000 -> 40000. Endpoint-scoped (HTTP /api/v1/agent/query only); internal callers (run_typed from wiki maintainer/writer/subagent) are unaffected. Qwen 27B runs at max_model_len=40960; 40K input chars ~= 10K tokens leaves ~70% of the window free for system prompt + tool defs + ~35-turn tool-call iteration + final output. No schemas, migrations, or API contract change. 2. ingest_watcher.py — CHUNK_WORDS 600 -> 1200. Halves chunk count on big documents (130 chunks -> 65 chunks on a BEAM-100K conversation), so halves total /agent/query round-trips. Per-chunk prompt grows from ~4 KB to ~7-8 KB, still 5x under the new 40K cap. CHUNK_OVERLAP stays at 75. 3. ingest_watcher.py — extract_facts_from_chunk prompt gains step (c): "AFTER all facts in this chunk are saved, look back at the facts you just created and add cross-fact relations WITHIN this chunk where genuinely meaningful." Relation types restricted to supports / contradicts / elaborates / similar_to / is_example_of. Explicit "do NOT create relations to facts from OTHER chunks" because the agent only sees this chunk's facts. Each per-chunk agent now does the same kinds of work central_review used to do, scoped to its chunk. 4. ingest_watcher.py — central_review call removed from enrich_datasource. The function body is intentionally KEPT in place as dormant code with a pointer comment explaining how to re-enable. Restoring the old behaviour = uncommenting one line. No cross-chunk semantic relations and no whole- document holistic thought get created — both already failed today on big documents (HTTP 422), so this is not a regression for big-doc behaviour; small-doc behaviour gains chunk-scoped relations instead of doc-scoped ones, with the same kinds of writes (same Pydantic tool schemas, same entity types, same downstream consumers). 5. ingest_watcher.py — max_turns 40 -> 35 on the per-chunk call. Math: ~10 facts/chunk * 2 calls (save_fact + derived_from) + ~0-5 cross-fact relations + 1 final_answer = ~21-26 turns, with ~9 turns of margin. config.py and .env / .env.bench audited: nothing of ours is env-driven. tests/test_split_chunks.py references CHUNK_WORDS by symbol (not literal 600), so the bump is test-clean. pytest tests/test_split_chunks.py: 10/10 green. pytest tests/test_ingest.py + test_ingest_watcher_no_dictation.py: 6/6 green. No test references central_review. Rollback: `git revert ` puts everything back atomically. Worst-case edge tweaks (drop max_length to 30K, drop CHUNK_WORDS to 1000, bump max_turns to 45) are each single-number edits if the smoke run surfaces an issue. --- braindb/ingest_watcher.py | 30 ++++++++++++++++++++++++------ braindb/routers/agent.py | 9 ++++++++- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/braindb/ingest_watcher.py b/braindb/ingest_watcher.py index cb3fea2..8257537 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -42,7 +42,11 @@ ALLOWED_EXTS = {".md", ".txt", ".json", ".yaml", ".yml", ".csv", ".log", ".html", ".xml"} SKIP_NAMES = {"README.md", ".gitkeep"} -CHUNK_WORDS = 600 # target chunk size (~4-5k token context per extraction call — NIM-friendly) +CHUNK_WORDS = 1200 # target chunk size (~7-8k char prompt, ~1.5K tokens of chunk text per call). + # Bumped 600 -> 1200 once /agent/query's max_length cap was lifted to 40000; + # halves the chunk count on big documents (so halves total round-trips). + # The per-chunk agent now also creates cross-fact relations within the chunk + # (see extract_facts_from_chunk), so there is no separate central_review pass. CHUNK_OVERLAP = 75 # words of overlap between adjacent chunks — catches facts that span a boundary INGEST_TIMEOUT = 60 @@ -155,7 +159,7 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f"Below is chunk {idx}/{total} of the document between markers.\n" f"Extract the concrete, standalone FACTS from this chunk — specific claims,\n" f"numbers, events, named decisions. Ignore filler, opinion, and generic\n" - f"statements. Aim for quality over quantity: typically 3-8 facts per chunk.\n\n" + f"statements. Aim for quality over quantity: typically 5-10 facts per chunk.\n\n" f"For EACH fact:\n" f' a) Call save_fact(content="",\n' f' source="document", keywords=[2-4 precise tags],\n' @@ -166,14 +170,23 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f'to_entity_id="{ds_id}", relation_type="derived_from",\n' f' description="Fact extracted from {title}"). Judge\n' f" relevance_score and importance_score per the tool's docstring.\n\n" + f"AFTER all facts in this chunk are saved, look back at the facts you just\n" + f"created and add cross-fact relations WITHIN this chunk where genuinely\n" + f"meaningful (not forced):\n" + f" c) For pairs that semantically connect, call create_relation(from_entity_id=,\n" + f' to_entity_id=, relation_type=, description=\"\").\n" + f" Skip if no pair clearly relates. Quality over quantity — typically 0-5\n" + f" such relations per chunk.\n\n" f"Do NOT call get_entity. Do NOT call update_entity on the datasource.\n" - f"Do NOT touch the datasource content — it is read-only.\n\n" - f"When all facts in this chunk are processed, call final_answer with\n" + f"Do NOT touch the datasource content — it is read-only.\n" + f"Do NOT create relations to facts from OTHER chunks — you only see this one.\n\n" + f"When all facts AND cross-fact relations are saved, call final_answer with\n" f"exactly this format so the watcher can parse it:\n" f' "Saved N facts from chunk {idx}/{total}: , , ..."\n\n' f"\n{chunk_text}\n" ) - answer = call_agent(prompt, max_turns=40) + answer = call_agent(prompt, max_turns=35) if not answer: return [] fact_ids = UUID_RE.findall(answer) @@ -256,7 +269,12 @@ def enrich_datasource(ds: dict, file_path: Path) -> None: log.warning("chunk %d/%d produced no facts", i, len(chunks)) log.info("extraction complete for %s: %d facts total", ds_id[:8], len(all_fact_ids)) - central_review(ds_id, title, all_fact_ids) + # central_review(ds_id, title, all_fact_ids) + # ^ kept dormant; per-chunk handles cross-fact relations now (see + # extract_facts_from_chunk step c). Re-enable this line to restore the + # whole-document cross-fact synthesis pass — note it fails on big docs + # because the fact list builds a >10K-char prompt (see central_review + # docstring for the prompt-construction issue). def process_file(path: Path) -> None: diff --git a/braindb/routers/agent.py b/braindb/routers/agent.py index 3337bf2..be355ee 100644 --- a/braindb/routers/agent.py +++ b/braindb/routers/agent.py @@ -20,7 +20,14 @@ class AgentQueryRequest(BaseModel): - query: str = Field(..., min_length=1, max_length=10000) + # Bumped 10000 -> 40000 to accommodate the ingest watcher's per-chunk + # extraction prompts at the larger CHUNK_WORDS=1200 size (chunk text + # ~6 KB + boilerplate + cross-fact-relation instructions ~7.5 KB total, + # well under the new cap). Sized so that even at 40K input chars + # (~10K tokens) the LLM has plenty of headroom: Qwen 27B runs at + # max_model_len=40960, so input + system prompt + tool defs + + # ~35-turn tool-call iteration leaves ~30% of the window for output. + query: str = Field(..., min_length=1, max_length=40000) max_turns: int | None = Field(default=None, ge=1, le=60) From d9cca46c35cae9cf99f1d0af9ffac1fb7dd0782e Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 22:46:16 +0100 Subject: [PATCH 07/17] fix(bench): restore api_bench network alias on SDK recreate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit restart_api_with_db() was using docker SDK containers.run(network=name) to attach the recreated api_bench to bench-network, but run()'s network kwarg does not accept aliases — so the recreated container lost the compose-set service-name alias `api_bench`. The bench_runner uses `BENCH_API_BASE: http://api_bench:8001` for DNS, and the alias loss made the first conversation fail with "Failed to resolve 'api_bench'" inside the 180s healthcheck loop. Fix: switch to containers.create() + networks.get(name).connect(container, aliases=...) + container.start(). Capture aliases from each original network and always include "api_bench" as a fallback (docker inspect sometimes returns null aliases even when DNS works, so always-add is safer). Also disconnect from the default bridge create() picks up by default. Caught by the monitor on the smoke run for the BrainDB big-chunks fix. No BrainDB code touched. --- benchmarks/beam/bench.py | 37 +++++++++++++++++++++++++++---------- 1 file changed, 27 insertions(+), 10 deletions(-) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index 7cdbf89..6a632fe 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -208,9 +208,17 @@ def restart_api_with_db(database_url: str) -> None: # the container was first created via `docker compose up`). binds = list(host_cfg.get("Binds") or []) - # Networks: rejoin all networks the container was on. Use the first as the - # primary `network=` arg, then attach the rest after start. - networks = list((net_settings.get("Networks") or {}).keys()) + # Networks: capture name + aliases from each. Compose sets the SERVICE + # NAME (`api_bench`) as a network alias so other services can DNS-resolve + # it; the SDK doesn't auto-restore that alias on container recreate, so + # we must replay it explicitly. We always include "api_bench" as a + # fallback alias even if the existing container's aliases came back + # empty (which docker inspect does sometimes report). + networks_to_connect = {} + for net_name, net_cfg in (net_settings.get("Networks") or {}).items(): + aliases = set(net_cfg.get("Aliases") or []) + aliases.add("api_bench") + networks_to_connect[net_name] = sorted(aliases) # Ports: preserve the existing host port mapping. port_bindings = {} @@ -227,15 +235,18 @@ def restart_api_with_db(database_url: str) -> None: existing.stop(timeout=30) existing.remove() - # Recreate with same shape + new env. - container = client.containers.run( + # Recreate with same shape + new env. Use create()+start() instead of + # run() so we can connect networks with aliases BEFORE the container + # boots (run(network=...) does not accept aliases). Without the alias + # replay, the recreated container loses its `api_bench` DNS name and + # bench_runner (which talks to it via `http://api_bench:8001`) fails + # to resolve it. + container = client.containers.create( image, name="braindb_bench_api", - detach=True, environment=new_env, volumes=binds, ports=port_bindings, - network=networks[0] if networks else None, extra_hosts={ h.split(":", 1)[0]: h.split(":", 1)[1] for h in (host_cfg.get("ExtraHosts") or []) @@ -248,12 +259,18 @@ def restart_api_with_db(database_url: str) -> None: command=cfg["Config"].get("Cmd"), working_dir=cfg["Config"].get("WorkingDir") or None, ) - # Attach to additional networks (run() only takes one). - for net_name in networks[1:]: + # Disconnect from the default bridge that `create` joined the container to, + # then attach the captured user-defined networks with their aliases. + try: + client.networks.get("bridge").disconnect(container, force=True) + except docker_sdk.errors.APIError: # type: ignore[attr-defined] + pass # not on the bridge — fine + for net_name, aliases in networks_to_connect.items(): try: - client.networks.get(net_name).connect(container) + client.networks.get(net_name).connect(container, aliases=aliases) except docker_sdk.errors.NotFound: # type: ignore[attr-defined] pass + container.start() def wait_for_api_healthy(timeout: float = 180, poll: float = 2) -> None: From 8827b2962a20a8ab4b70611e13ddbceda5b40f08 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:15:46 +0100 Subject: [PATCH 08/17] fix(bench): warmup tracks relations age too + bumps settle to 300s MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Caught by the v3 smoke run on the new big-chunks code path: warmup converged after only chunk 1 of 57 had completed, because: - The new per-chunk extract_facts_from_chunk prompt asks the agent to (a) save_fact + create_relation(derived_from) for each fact, (b) then create_relation between facts in the same chunk, and (c) the agent routinely calls delegate_to_subagent + recall_memory before final_answer. - Phase (a) bursts into the entities table for the first ~2 minutes of each chunk. Phases (b) and (c) write to the relations table (or do pure reasoning) for the next 2-4 minutes — the entities table is quiet during this whole tail. - The old _seconds_since_last_entity probe interpreted that quiet entities-table tail as "extraction is done", so warmup converged at t=281s while chunk 1's agent was still working and chunks 2..N had not yet started. Fix: - Switch warmup to _seconds_since_last_activity, which is the seconds since the most recent INSERT in EITHER entities OR relations (GREATEST in Postgres ignores NULLs, so empty relations falls through to entities cleanly). - Bump settle_seconds default from 180s -> 300s. The new probe is refreshed by every create_relation call, so 5 minutes of no DB writes truly means the agent loop has finished, not just paused in its relation-creation tail. - Update the TimeoutError message + CLI help to reflect the new probe. No BrainDB code touched. Bench-only fix on feat/benchmark-big-chunks. --- benchmarks/beam/warmup.py | 51 +++++++++++++++++++++++++++------------ 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/benchmarks/beam/warmup.py b/benchmarks/beam/warmup.py index d2d15ca..7aceecb 100644 --- a/benchmarks/beam/warmup.py +++ b/benchmarks/beam/warmup.py @@ -47,10 +47,28 @@ def _entity_count(conn) -> int: return int(row[0]) if row else 0 -def _seconds_since_last_entity(conn) -> float | None: +def _seconds_since_last_activity(conn) -> float | None: + """Seconds since the last INSERT in EITHER the entities OR the relations + table. Tracking both is necessary because the per-chunk extraction agent + (after the big-chunks change) does its save_fact() bursts up front, then + spends 2-4 minutes calling create_relation() for cross-fact edges and + optionally delegate_to_subagent() / recall_memory(). During that tail + the entities table is quiet but the agent IS still working. If warmup + only watched entities, it would falsely declare "extraction done" while + chunk 1's agent is still doing its relation work and chunks 2..N have + not started. + + GREATEST in Postgres ignores NULL arguments, so an empty relations + table (early in the run) falls back gracefully to the entities max. + Both NULL (truly empty DB) → returns NULL → caller treats as "no data + yet". + """ row = _query_one( conn, - "SELECT EXTRACT(EPOCH FROM (NOW() - MAX(created_at))) FROM entities", + """SELECT EXTRACT(EPOCH FROM (NOW() - GREATEST( + (SELECT MAX(created_at) FROM entities), + (SELECT MAX(created_at) FROM relations) + )))""", ) return float(row[0]) if row and row[0] is not None else None @@ -71,7 +89,7 @@ def _unprocessed_files() -> list[Path]: def wait_for_warmup( *, - settle_seconds: float = 180.0, + settle_seconds: float = 300.0, consecutive_clear_required: int = 2, poll_interval: float = 5.0, timeout_seconds: float = 1800.0, @@ -112,25 +130,26 @@ def wait_for_warmup( try: while time.monotonic() < deadline: files_remaining = len(_unprocessed_files()) - entity_age = _seconds_since_last_entity(conn) + activity_age = _seconds_since_last_activity(conn) pending = _pending_wiki_jobs(conn) entity_count = _entity_count(conn) - if entity_age is None: - # No entities yet. If files are still in sources/, watcher - # hasn't started; if not, it may have just finished — give - # it a tiny grace period before declaring "no work to do". + if activity_age is None: + # No entities or relations yet. If files are still in + # sources/, watcher hasn't started; if not, it may have + # just finished — give it a tiny grace period before + # declaring "no work to do". clear = files_remaining == 0 and (time.monotonic() - start) > 10 else: clear = ( files_remaining == 0 - and entity_age >= settle_seconds + and activity_age >= settle_seconds and (not block_on_wiki_queue or pending == 0) ) elapsed = time.monotonic() - start if verbose and (elapsed - last_log >= log_interval): - age_str = f"{entity_age:.0f}s" if entity_age is not None else "n/a" + age_str = f"{activity_age:.0f}s" if activity_age is not None else "n/a" print( f"[warmup t={elapsed:6.0f}s] files_left={files_remaining} " f"entities={entity_count} last_entity_age={age_str} " @@ -152,7 +171,7 @@ def wait_for_warmup( raise TimeoutError( f"warmup did not converge within {timeout_seconds:.0f}s " f"(files_left={len(_unprocessed_files())}, " - f"last_entity_age={_seconds_since_last_entity(conn)}, " + f"last_activity_age={_seconds_since_last_activity(conn)}, " f"pending_wiki={_pending_wiki_jobs(conn)})" ) finally: @@ -176,10 +195,12 @@ def _final_stats(conn, start: float) -> dict: def _cli() -> int: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--timeout", type=float, default=1800) - p.add_argument("--settle-seconds", type=float, default=180, - help="seconds of no INSERT on entities before declaring extraction settled " - "(default 180; large enough to span the gap between datasource creation " - "and first extracted fact for big documents)") + p.add_argument("--settle-seconds", type=float, default=300, + help="seconds of no INSERT on entities OR relations before declaring " + "extraction settled (default 300). Per-chunk extraction agents do " + "save_fact bursts up front, then 2-4 minutes of create_relation + " + "subagent / recall_memory work; tracking both tables prevents the " + "barrier from falsely converging during a chunk's relation tail.") p.add_argument("--poll-interval", type=float, default=5) p.add_argument("--consecutive-clear", type=int, default=2) p.add_argument("--quiet", action="store_true") From 15196c4d4564f9c022359b3ac20a633dbc7e3011 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Thu, 4 Jun 2026 23:41:15 +0100 Subject: [PATCH 09/17] fix(bench): bench-side warmup_settle_seconds default also 180 -> 300 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous warmup.py fix bumped its OWN default 180 -> 300, but bench.py had its own warmup_settle_seconds parameter default at 180, which is the value actually passed to wait_for_warmup() — so the warmup.py default never took effect. Caught on the v4 smoke run: warmup converged at activity_age=183s while chunks 4..57 were still pending; Phase C kicked off against a 37-fact DB (~1/4 of the document). Now both layers default to 300s and the CLI flag matches. No BrainDB code touched. Two-line bench-only fix on feat/benchmark-big-chunks. --- benchmarks/beam/bench.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index 6a632fe..504ad10 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -358,7 +358,7 @@ def run_one_conversation( run_dir: Path, *, warmup_timeout: float = 1800, - warmup_settle_seconds: float = 180, + warmup_settle_seconds: float = 300, question_timeout: float = QUESTION_TIMEOUT_SECONDS, block_on_wiki_queue: bool = False, ) -> dict: @@ -517,10 +517,11 @@ def _parse_args() -> argparse.Namespace: help="abort the whole run if any single conversation raises") p.add_argument("--warmup-timeout", type=float, default=1800, help="seconds before warmup gives up on convergence") - p.add_argument("--warmup-settle-seconds", type=float, default=180, - help="seconds of quiet on entities before declaring warmup clear " - "(default 180; big enough to span the gap between datasource " - "creation and the first extracted fact for slow chunk processing)") + p.add_argument("--warmup-settle-seconds", type=float, default=300, + help="seconds of quiet on entities AND relations before declaring " + "warmup clear (default 300). Per-chunk agents have ~60-90s quiet " + "stretches doing subagent / recall_memory work between save_fact " + "and create_relation bursts; 300s comfortably covers those.") p.add_argument("--wait-for-wikis", action="store_true", help="Strict warmup mode: also wait for the wiki_job queue to be " "fully drained before answering. Off by default because the " From 587acd9f5066c936972d50a278cefb1a36854721 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Fri, 5 Jun 2026 07:24:43 +0100 Subject: [PATCH 10/17] fix(bench): chunk + question retries, max_turns 40, generous timeouts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Six surgical edits across three files to make the full BEAM run robust enough that user can leave it for 5 days without supervision. braindb/ingest_watcher.py: - AGENT_TIMEOUT 1200s -> 1800s (30 min single-attempt cap; chunks average ~6 min, the cross-fact-relation tail can stretch past 10 min on dense chunks). - max_turns 35 -> 40 on the per-chunk call (was deliberately lowered in c185ee0 thinking 35 was enough; v4 smoke caught 2/57 chunks hitting "Max turns (3) exceeded" — the agent's Layer-4 retry ran out after primary budget exhausted at 35). - Per-chunk retry loop: up to 3 attempts with 5s/10s backoff. Retries ONLY when call_agent returns None (HTTP error) OR the parse returns zero fact UUIDs (agent never reached final_answer). If the first attempt's parse DID return UUIDs, do not retry — those facts are already committed via tool calls and a retry would dupe them. Catches the transient-failure class without the duplication risk. benchmarks/beam/bench.py: - answer_one_question: 3-attempt retry around the POST with 5s/10s backoff. Retries on ConnectionError / Timeout / HTTP 5xx. Does NOT retry on HTTP 4xx (request-shape errors will repeat). After three failures the last exception bubbles to the caller which records the question as `[ERROR: ...]` and the run continues to the next question. - wait_for_api_healthy default timeout 180s -> 300s (api restart needs alembic + uvicorn + embedding model load; 5 min has comfortable slack over the observed ~30-60s). - run_one_conversation warmup_settle_seconds default 300 -> 600 to match the new warmup.py default (otherwise this layer overrides). benchmarks/beam/warmup.py: - settle_seconds default 300 -> 600. With the new per-chunk cross-fact-relation prompt, in-chunk quiet stretches can reach 3 min (subagent / recall_memory work). 10 min gives generous slack while still catching genuinely stuck convergence (the warmup_timeout=1800s cap fires if extraction itself stalls). Pytest tests/test_split_chunks.py + tests/test_ingest.py + tests/test_ingest_watcher_no_dictation.py: 16/16 green. This is the last set of changes before the full --split 100K run. After this commit, do not touch the run autonomously. --- benchmarks/beam/bench.py | 64 +++++++++++++++++++++++++++------------ benchmarks/beam/warmup.py | 10 +++--- braindb/ingest_watcher.py | 31 ++++++++++++++----- 3 files changed, 74 insertions(+), 31 deletions(-) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index 504ad10..699d6c9 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -273,12 +273,14 @@ def restart_api_with_db(database_url: str) -> None: container.start() -def wait_for_api_healthy(timeout: float = 180, poll: float = 2) -> None: +def wait_for_api_healthy(timeout: float = 300, poll: float = 2) -> None: """Block until the bench api responds with status=ok. The api needs to (1) connect to the new DB, (2) run alembic migrations, (3) load the embedding model into memory. On first start this is ~30-60s; on subsequent restarts (model already cached) it can still be ~20-30s. + Bumped 180s -> 300s for generous slack on cold-cache restarts and to + absorb an occasional slow alembic migration without tripping. """ start = time.monotonic() last_err = None @@ -333,22 +335,44 @@ def answer_one_question( ) -> tuple[str, dict]: """POST /agent/query with the question; return (answer_text, raw_payload). - Returns the answer text and the full payload (so we can capture tool - counts, latency, etc.) for the per-question record. Errors get caught - in the caller and recorded as `[ERROR: ...]` so the run continues. + Three-attempt retry with 5s/10s backoff. Retries on connection errors, + timeouts, and HTTP 5xx — these are the transient classes (API recycling, + LLM backend blip, vLLM batching pause). Does NOT retry on HTTP 4xx — + those are request-shape errors and will repeat. After 3 failed attempts + the last exception is raised so the caller records `[ERROR: ...]` and + the run continues to the next question. """ - started = time.monotonic() - r = requests.post( - f"{base}/api/v1/agent/query", - json={"query": question_text}, - timeout=timeout, - ) - elapsed = time.monotonic() - started - r.raise_for_status() - payload = r.json() - payload["_elapsed_seconds"] = round(elapsed, 2) - # BrainDB's /agent/query response shape: {"answer": "..."} at minimum. - return payload.get("answer", ""), payload + last_exc: Exception | None = None + for attempt in range(1, 4): + started = time.monotonic() + try: + r = requests.post( + f"{base}/api/v1/agent/query", + json={"query": question_text}, + timeout=timeout, + ) + except (requests.ConnectionError, requests.Timeout) as e: + last_exc = e + else: + if 500 <= r.status_code < 600: + last_exc = requests.HTTPError( + f"HTTP {r.status_code} on attempt {attempt}/3: {r.text[:200]}", + response=r, + ) + else: + r.raise_for_status() # 4xx -> raise immediately, caller records + elapsed = time.monotonic() - started + payload = r.json() + payload["_elapsed_seconds"] = round(elapsed, 2) + if attempt > 1: + payload["_attempts"] = attempt + # BrainDB's /agent/query response shape: {"answer": "..."} at minimum. + return payload.get("answer", ""), payload + # Got a transient error — back off and retry unless this was the last try + if attempt < 3: + time.sleep(5 * attempt) # 5s, 10s + assert last_exc is not None + raise last_exc # ----------------------------- per-conv lifecycle ---------------------------- @@ -358,7 +382,7 @@ def run_one_conversation( run_dir: Path, *, warmup_timeout: float = 1800, - warmup_settle_seconds: float = 300, + warmup_settle_seconds: float = 600, question_timeout: float = QUESTION_TIMEOUT_SECONDS, block_on_wiki_queue: bool = False, ) -> dict: @@ -517,11 +541,11 @@ def _parse_args() -> argparse.Namespace: help="abort the whole run if any single conversation raises") p.add_argument("--warmup-timeout", type=float, default=1800, help="seconds before warmup gives up on convergence") - p.add_argument("--warmup-settle-seconds", type=float, default=300, + p.add_argument("--warmup-settle-seconds", type=float, default=600, help="seconds of quiet on entities AND relations before declaring " - "warmup clear (default 300). Per-chunk agents have ~60-90s quiet " + "warmup clear (default 600). Per-chunk agents have 2-4 min quiet " "stretches doing subagent / recall_memory work between save_fact " - "and create_relation bursts; 300s comfortably covers those.") + "and create_relation bursts; 10 min gives generous slack.") p.add_argument("--wait-for-wikis", action="store_true", help="Strict warmup mode: also wait for the wiki_job queue to be " "fully drained before answering. Off by default because the " diff --git a/benchmarks/beam/warmup.py b/benchmarks/beam/warmup.py index 7aceecb..76ce97e 100644 --- a/benchmarks/beam/warmup.py +++ b/benchmarks/beam/warmup.py @@ -89,7 +89,7 @@ def _unprocessed_files() -> list[Path]: def wait_for_warmup( *, - settle_seconds: float = 300.0, + settle_seconds: float = 600.0, consecutive_clear_required: int = 2, poll_interval: float = 5.0, timeout_seconds: float = 1800.0, @@ -195,12 +195,14 @@ def _final_stats(conn, start: float) -> dict: def _cli() -> int: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument("--timeout", type=float, default=1800) - p.add_argument("--settle-seconds", type=float, default=300, + p.add_argument("--settle-seconds", type=float, default=600, help="seconds of no INSERT on entities OR relations before declaring " - "extraction settled (default 300). Per-chunk extraction agents do " + "extraction settled (default 600). Per-chunk extraction agents do " "save_fact bursts up front, then 2-4 minutes of create_relation + " "subagent / recall_memory work; tracking both tables prevents the " - "barrier from falsely converging during a chunk's relation tail.") + "barrier from falsely converging during a chunk's relation tail. " + "10 min gives comfortable slack over the worst observed quiet " + "stretch (~3 min) plus chunk-transition gaps.") p.add_argument("--poll-interval", type=float, default=5) p.add_argument("--consecutive-clear", type=int, default=2) p.add_argument("--quiet", action="store_true") diff --git a/braindb/ingest_watcher.py b/braindb/ingest_watcher.py index 8257537..89ba659 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -50,7 +50,7 @@ CHUNK_OVERLAP = 75 # words of overlap between adjacent chunks — catches facts that span a boundary INGEST_TIMEOUT = 60 -AGENT_TIMEOUT = int(os.getenv("AGENT_TIMEOUT", "1200")) # env-overridable; matches wiki_scheduler pattern. Generous default for slow / deep-exploring LLM runs. +AGENT_TIMEOUT = int(os.getenv("AGENT_TIMEOUT", "1800")) # env-overridable; matches wiki_scheduler pattern. 30-min single-attempt cap leaves plenty of headroom for big-chunk runs where the cross-fact-relation tail extends past the ~6-min average. UUID_RE = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") logging.basicConfig( @@ -186,12 +186,29 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f' "Saved N facts from chunk {idx}/{total}: , , ..."\n\n' f"\n{chunk_text}\n" ) - answer = call_agent(prompt, max_turns=35) - if not answer: - return [] - fact_ids = UUID_RE.findall(answer) - # Filter out the datasource id if the model happened to echo it - return [fid for fid in fact_ids if fid != ds_id] + # 3-attempt retry around the per-chunk call. Retry ONLY when answer is None + # (HTTP error, timeout, connection failure) OR the parse produced zero fact + # UUIDs (agent never reached final_answer). If the first attempt DID save + # facts (parsed UUIDs in the answer string), do not retry — the facts are + # already committed via tool calls, and a retry would dupe them. + for attempt in range(1, 4): + answer = call_agent(prompt, max_turns=40) + if answer: + fact_ids = UUID_RE.findall(answer) + real_ids = [fid for fid in fact_ids if fid != ds_id] + if real_ids: + if attempt > 1: + log.info("chunk %d/%d succeeded on attempt %d/3", idx, total, attempt) + return real_ids + if attempt < 3: + backoff = 5 * attempt # 5s, then 10s + log.warning( + "chunk %d/%d attempt %d/3 produced no facts; retrying in %ds", + idx, total, attempt, backoff, + ) + time.sleep(backoff) + log.warning("chunk %d/%d FAILED after 3 attempts", idx, total) + return [] def central_review(ds_id: str, title: str, fact_ids: list[str]) -> None: From 8c1117491f5033c3d44f111dda7ad4cf91a61908 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Fri, 5 Jun 2026 09:01:22 +0100 Subject: [PATCH 11/17] fix(bench): warmup_timeout 1800s -> 43200s (12h) for big-chunk extraction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Caught on the morning run: every conversation timed out at 1800s warmup because 100K-conv extraction takes ~5.7h (57 chunks * ~6 min each at CHUNK_WORDS=1200) — there is no way for warmup to converge within 30 minutes when extraction is actively progressing. When the bench gave up, api_bench was recreated for the next conv mid-extraction, polluting both the old and new conv databases with cross-talk facts. Fix: bump warmup_timeout default from 1800 -> 43200 (12h). This is 2x the projected 100K extraction time and gives generous slack for retry overhead and edge cases. settle_seconds=600s (set in the prior commit) still catches genuine stalls inside the 12h window, so the full timeout only burns when something is truly stuck. Both the warmup.py default + CLI default and the bench.py run_one_conversation default + CLI default are updated together to keep them in lockstep (otherwise the bench layer overrides). For future 500K / 1M splits this knob will need to scale up further (500K convs are ~30h, 1M convs are ~60h). 100K is the first run target; later splits will bump the CLI flag explicitly. No BrainDB code touched. Two-number bench-only fix. --- benchmarks/beam/bench.py | 10 +++++++--- benchmarks/beam/warmup.py | 7 +++++-- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index 699d6c9..e53a27e 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -381,7 +381,7 @@ def run_one_conversation( conv: Conversation, run_dir: Path, *, - warmup_timeout: float = 1800, + warmup_timeout: float = 43200, warmup_settle_seconds: float = 600, question_timeout: float = QUESTION_TIMEOUT_SECONDS, block_on_wiki_queue: bool = False, @@ -539,8 +539,12 @@ def _parse_args() -> argparse.Namespace: help="comma-separated list of conversation_ids to run (overrides --limit ordering)") p.add_argument("--fail-fast", action="store_true", help="abort the whole run if any single conversation raises") - p.add_argument("--warmup-timeout", type=float, default=1800, - help="seconds before warmup gives up on convergence") + p.add_argument("--warmup-timeout", type=float, default=43200, + help="seconds before warmup gives up on convergence (default 43200 = 12h). " + "Must exceed total per-conversation extraction time (a 100K conv takes " + "~5.7h at 1200-word chunks). settle_seconds=600s catches genuine stalls " + "inside this window so the full 12h is only burned in the " + "extraction-keeps-progressing case.") p.add_argument("--warmup-settle-seconds", type=float, default=600, help="seconds of quiet on entities AND relations before declaring " "warmup clear (default 600). Per-chunk agents have 2-4 min quiet " diff --git a/benchmarks/beam/warmup.py b/benchmarks/beam/warmup.py index 76ce97e..40ff89e 100644 --- a/benchmarks/beam/warmup.py +++ b/benchmarks/beam/warmup.py @@ -92,7 +92,7 @@ def wait_for_warmup( settle_seconds: float = 600.0, consecutive_clear_required: int = 2, poll_interval: float = 5.0, - timeout_seconds: float = 1800.0, + timeout_seconds: float = 43200.0, log_interval: float = 10.0, verbose: bool = True, block_on_wiki_queue: bool = False, @@ -194,7 +194,10 @@ def _final_stats(conn, start: float) -> dict: def _cli() -> int: p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) - p.add_argument("--timeout", type=float, default=1800) + p.add_argument("--timeout", type=float, default=43200, + help="seconds before warmup gives up on convergence (default 43200 = 12h, " + "covers a ~6h 100K-extraction with 2x slack; settle_seconds catches " + "genuine stalls inside this window).") p.add_argument("--settle-seconds", type=float, default=600, help="seconds of no INSERT on entities OR relations before declaring " "extraction settled (default 600). Per-chunk extraction agents do " From 931ba6cc924481ee4b815b99d32c574f26762612 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Fri, 5 Jun 2026 10:00:13 +0100 Subject: [PATCH 12/17] =?UTF-8?q?fix(bench):=20extraction=20prompt=20?= =?UTF-8?q?=E2=80=94=20fact-count=20range=205-10=20=E2=86=92=201-2=20to=20?= =?UTF-8?q?30=20by=20density?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Observed during the in-flight 100K run: agent consistently hits the ~5-10 facts/chunk band the prompt suggests, treating it as both target and implicit cap. For 1200-word (~4-5 page) chunks this looks light relative to realistic conversation density. One-line phrase swap inside extract_facts_from_chunk's prompt template. Replaces: "typically 5-10 facts per chunk" with: "typically from 1-2 to 30 depending on density and length of the chunk" Everything else in the prompt is preserved: same "concrete, standalone FACTS — specific claims, numbers, events, named decisions" framing, same "Ignore filler, opinion, and generic statements", same "quality over quantity" anchor, same a/b/c steps, same Do-NOT rules, same final_answer format. No bump to max_turns (currently 40). The agent will likely shift distribution upward but rarely hit 30; if "Max turns exceeded" failures spike we can bump max_turns in a follow-up. Tests: pytest tests/test_split_chunks.py + test_ingest.py + test_ingest_watcher_no_dictation.py = 16/16 green. The no-dictation test specifically checks the prompt doesn't pre-dictate certainty / importance / relevance_score values; this change doesn't touch those. Easy revert: git revert this SHA, push, workstation pulls, restart. --- braindb/ingest_watcher.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/braindb/ingest_watcher.py b/braindb/ingest_watcher.py index 89ba659..10e4c4b 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -159,7 +159,7 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f"Below is chunk {idx}/{total} of the document between markers.\n" f"Extract the concrete, standalone FACTS from this chunk — specific claims,\n" f"numbers, events, named decisions. Ignore filler, opinion, and generic\n" - f"statements. Aim for quality over quantity: typically 5-10 facts per chunk.\n\n" + f"statements. Aim for quality over quantity: typically from 1-2 to 30 depending on density and length of the chunk.\n\n" f"For EACH fact:\n" f' a) Call save_fact(content="",\n' f' source="document", keywords=[2-4 precise tags],\n' From 6483d496b8179a920c9f55c58e6bb06964ae3ab2 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 7 Jun 2026 01:47:34 +0100 Subject: [PATCH 13/17] feat(bench): byte ranges in fact notes so agent stops chunk-guessing Symptoms: each Phase C question takes 30-60+ min (42 get_entity + 4 recursive delegate_to_subagent calls per question). One subagent crashed with ContextWindowExceededError after 891s. Root cause: facts had notes "Extracted from chunk X/Y" but get_entity only takes byte offsets. Agent guessed chunk -> offset (36/57 x 526K = 333K) and paged randomly. Fix: - split_chunks_with_offsets returns (text, byte_start, byte_end) per chunk; original whitespace preserved. split_chunks is now a 1-line back-compat wrapper - single source of truth, all 10 existing tests pass unchanged. - extract_facts_from_chunk threads byte_start/byte_end into the prompt so save_fact records notes="Extracted from chunk X/Y (bytes A-B)". - enrich_datasource uses the offset-aware chunker. - system_prompt.md + skills/braindb/SKILL.md each get ONE descriptive sentence explaining the new notes format. No behavioural tightening. 6 new unit tests for split_chunks_with_offsets cover: offsets recover first/last word via text[byte_start:byte_end], byte_end <= len(text), monotonic byte_starts, back-compat with split_chunks. Co-Authored-By: Claude Opus 4.7 --- braindb/agent/prompts/system_prompt.md | 1 + braindb/ingest_watcher.py | 69 ++++++++++++++----- skills/braindb/SKILL.md | 1 + tests/test_split_chunks.py | 95 +++++++++++++++++++++++++- 4 files changed, 146 insertions(+), 20 deletions(-) diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index 7d289d3..1bcb891 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -89,6 +89,7 @@ design — research from previews, then open only the few you actually need. `null`. For anything sizable, hand each slice to `delegate_to_subagent` ("process THIS slice and return only the distilled result") and aggregate — your main context must stay small. +- Fact `notes` may include a byte range like `(bytes 245760-252960)` — pass directly as `get_entity(id, offset=245760, limit=7200)` to read the exact source slice the fact was extracted from. - Never try to defeat previews via `search_sql` to dump whole bodies. ## DELEGATION — use `delegate_to_subagent` for focused deep work diff --git a/braindb/ingest_watcher.py b/braindb/ingest_watcher.py index 10e4c4b..ba2d9b0 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -113,27 +113,48 @@ def call_agent(prompt: str, max_turns: int = 8) -> str | None: return r.json().get("answer") or "" -def split_chunks(text: str, chunk_words: int = CHUNK_WORDS, overlap: int = CHUNK_OVERLAP) -> list[str]: - """Split text into word-bounded chunks with configurable overlap. - - Always splits at whitespace, never mid-word. Each chunk (after the first) - starts `overlap` words before the previous chunk's tail, so facts that - straddle a boundary are still visible in at least one chunk. +def split_chunks_with_offsets( + text: str, chunk_words: int = CHUNK_WORDS, overlap: int = CHUNK_OVERLAP, +) -> list[tuple[str, int, int]]: + """Like :func:`split_chunks` but each entry is + ``(chunk_text, byte_start, byte_end)`` where ``byte_start``/``byte_end`` + are positions in the ORIGINAL text (so ``text[byte_start:byte_end]`` + recovers the original whitespace, not just the joined words). + + Used at extraction time so the recall-side agent can locate the exact + source slice for a fact via ``get_entity(id, offset=byte_start, + limit=byte_end - byte_start)`` instead of guessing byte offsets from + a word-based chunk number. """ - words = text.split() - if not words: + word_positions = [(m.group(), m.start(), m.end()) for m in re.finditer(r"\S+", text)] + if not word_positions: return [] if overlap >= chunk_words: overlap = 0 # nonsensical config, fall back to no overlap step = chunk_words - overlap - chunks: list[str] = [] + result: list[tuple[str, int, int]] = [] i = 0 - while i < len(words): - chunks.append(" ".join(words[i:i + chunk_words])) - if i + chunk_words >= len(words): + while i < len(word_positions): + end = min(i + chunk_words, len(word_positions)) + slice_ = word_positions[i:end] + result.append(( + " ".join(w[0] for w in slice_), + slice_[0][1], + slice_[-1][2], + )) + if end == len(word_positions): break i += step - return chunks + return result + + +def split_chunks(text: str, chunk_words: int = CHUNK_WORDS, overlap: int = CHUNK_OVERLAP) -> list[str]: + """List of chunk strings. Thin back-compat wrapper around + :func:`split_chunks_with_offsets` — keeps the existing API stable for + callers that don't need byte offsets, and ensures the chunking logic + has a single source of truth. + """ + return [t for (t, _, _) in split_chunks_with_offsets(text, chunk_words, overlap)] def fetch_entity(entity_id: str) -> dict | None: @@ -146,10 +167,19 @@ def fetch_entity(entity_id: str) -> dict | None: return None -def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk_text: str) -> list[str]: +def extract_facts_from_chunk( + ds_id: str, title: str, idx: int, total: int, chunk_text: str, + byte_start: int, byte_end: int, +) -> list[str]: """Ask one agent call to extract facts from a chunk, save each via save_fact, and link each back to the datasource via create_relation(derived_from). Returns the list of new fact IDs parsed from the agent's final_answer answer. + + ``byte_start``/``byte_end`` are the chunk's offset/length in the original + datasource text; they are baked into each saved fact's ``notes`` field so + a downstream recall agent can read the exact source slice via + ``get_entity(id, offset=byte_start, limit=byte_end-byte_start)`` instead + of guessing the location from the chunk number. """ prompt = ( f"A document was just ingested into BrainDB.\n" @@ -163,7 +193,8 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f"For EACH fact:\n" f' a) Call save_fact(content="",\n' f' source="document", keywords=[2-4 precise tags],\n' - f' notes="Extracted from {title} chunk {idx}/{total}"). Judge\n' + f' notes="Extracted from {title} chunk {idx}/{total} ' + f'(bytes {byte_start}-{byte_end})"). Judge\n' f" certainty and importance per the tool's docstring guidance.\n" f" Record the returned fact id.\n" f" b) Call create_relation(from_entity_id=, " @@ -269,16 +300,18 @@ def enrich_datasource(ds: dict, file_path: Path) -> None: log.warning("enrichment read failed for %s: %s", ds_id[:8], e) return - chunks = split_chunks(text) + chunks = split_chunks_with_offsets(text) if not chunks: log.info("enrichment skipped for %s: empty content", ds_id[:8]) return log.info("extraction started for %s: %d chunks", ds_id[:8], len(chunks)) all_fact_ids: list[str] = [] - for i, chunk in enumerate(chunks, start=1): + for i, (chunk, b_start, b_end) in enumerate(chunks, start=1): log.info("extracting facts chunk %d/%d for %s", i, len(chunks), ds_id[:8]) - fact_ids = extract_facts_from_chunk(ds_id, title, i, len(chunks), chunk) + fact_ids = extract_facts_from_chunk( + ds_id, title, i, len(chunks), chunk, b_start, b_end, + ) if fact_ids: log.info("chunk %d/%d saved %d facts", i, len(chunks), len(fact_ids)) all_fact_ids.extend(fact_ids) diff --git a/skills/braindb/SKILL.md b/skills/braindb/SKILL.md index 8f8e8e9..5a01cea 100644 --- a/skills/braindb/SKILL.md +++ b/skills/braindb/SKILL.md @@ -139,6 +139,7 @@ decide from previews, then read only what you need: follow `content_meta.next_offset` until it is `null`. For big documents, prefer `POST /api/v1/agent/query` with "delegate to a subagent to read and distil entity " so the heavy content never enters this conversation. +- Fact `notes` may include a byte range like `(bytes 245760-252960)` — pass directly as `GET /api/v1/entities/{id}?offset=245760&limit=7200` to read the exact source slice the fact was extracted from. ## RECALL — Before Responding diff --git a/tests/test_split_chunks.py b/tests/test_split_chunks.py index c6da0c5..9f136e4 100644 --- a/tests/test_split_chunks.py +++ b/tests/test_split_chunks.py @@ -2,9 +2,17 @@ Pure-function tests for the ingest watcher's chunk splitter. Covers: default chunk size, custom size, overlap, word boundaries, edge cases -(empty, single word, exact chunk-size boundary). +(empty, single word, exact chunk-size boundary). Also tests +``split_chunks_with_offsets`` — the offset-aware variant whose byte ranges +get stored in fact notes so the recall agent can locate the exact source +slice via ``get_entity(offset, limit)`` instead of guessing from chunk number. """ -from braindb.ingest_watcher import split_chunks, CHUNK_WORDS, CHUNK_OVERLAP +from braindb.ingest_watcher import ( + CHUNK_OVERLAP, + CHUNK_WORDS, + split_chunks, + split_chunks_with_offsets, +) def test_empty_text(): @@ -95,3 +103,86 @@ def test_words_are_preserved_exactly(): for word in chunk.split(): seen.add(word) assert seen == {"one", "two", "three", "four", "five", "six", "seven", "eight", "nine", "ten"} + + +# ---------- split_chunks_with_offsets ---------------------------------------- + + +def test_with_offsets_empty(): + assert split_chunks_with_offsets("") == [] + assert split_chunks_with_offsets(" \n\t ") == [] + + +def test_with_offsets_single_chunk_basic(): + text = "hello world there" + out = split_chunks_with_offsets(text, chunk_words=10, overlap=0) + assert len(out) == 1 + chunk_text, b_start, b_end = out[0] + assert chunk_text == "hello world there" + assert b_start == 0 + assert b_end == len(text) + + +def test_with_offsets_recover_first_and_last_word(): + """For every chunk, text[b_start:b_end] must START with the first chunk + word and END with the last chunk word. This is the core guarantee the + recall agent will rely on when paging by byte offset.""" + # Use varied whitespace (multiple spaces, newlines, tabs) to make sure + # offsets track the original text, not the joined-words form. + text = ( + "alpha beta\nGamma\tdelta epsilon " + "zeta eta theta iota kappa " + "lambda mu nu xi omicron pi rho " + "sigma tau upsilon phi chi psi omega" + ) + out = split_chunks_with_offsets(text, chunk_words=5, overlap=1) + assert len(out) >= 4 + for chunk_text, b_start, b_end in out: + slice_ = text[b_start:b_end] + first_word = chunk_text.split()[0] + last_word = chunk_text.split()[-1] + assert slice_.lstrip().startswith(first_word), ( + f"slice {slice_!r} does not start with {first_word!r}" + ) + assert slice_.rstrip().endswith(last_word), ( + f"slice {slice_!r} does not end with {last_word!r}" + ) + + +def test_with_offsets_byte_end_in_bounds(): + """No byte_end should ever exceed the length of the original text.""" + text = " ".join(f"w{i}" for i in range(200)) + out = split_chunks_with_offsets(text, chunk_words=50, overlap=10) + n = len(text) + for chunk_text, b_start, b_end in out: + assert 0 <= b_start <= b_end <= n + # And byte_end - byte_start is non-trivial + assert b_end > b_start + + +def test_with_offsets_starts_advance_monotonically(): + """Each successive chunk's byte_start must be > the previous chunk's + byte_start (we may overlap into the previous chunk's bytes, but we always + advance forward by at least one word).""" + text = " ".join(f"w{i}" for i in range(500)) + out = split_chunks_with_offsets(text, chunk_words=50, overlap=10) + starts = [b_start for (_, b_start, _) in out] + assert starts == sorted(starts) + assert len(set(starts)) == len(starts), "duplicate byte_starts found" + + +def test_split_chunks_back_compat_with_offsets(): + """split_chunks must return exactly the chunk-text components of + split_chunks_with_offsets for the same inputs — proving the back-compat + wrapper has zero behavioural drift from the offset-aware function. + """ + cases = [ + "", + "hello", + " ".join(f"w{i}" for i in range(CHUNK_WORDS + 5)), + " ".join(f"w{i}" for i in range(CHUNK_WORDS * 3 + 17)), + ] + for text in cases: + a = split_chunks(text) + b = [t for (t, _, _) in split_chunks_with_offsets(text)] + assert a == b, f"divergence on text len={len(text)}" From 0f438201f6e1ffd362b7de82d5adac002003e5e6 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 7 Jun 2026 20:14:31 +0100 Subject: [PATCH 14/17] feat(bench): wiki-wait window between Phase B and Phase C MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds optional --wiki-wait-seconds CLI flag. When > 0, the bench sleeps between [B done] and [C], printing [Bw] / [Bw done] markers. During the sleep the wiki writer + maintainer have vLLM essentially to themselves (no extraction, no question agent competing), so wikis develop more before recall queries hit during Phase C. Recommended value for BEAM 100K: 28800s (8h) — wiki writer at full vLLM speed drains ~120 writes/hour, so 8h captures most of the typical 1600-job queue. Longer waits hit diminishing returns because the maintainer keeps queuing new triage jobs at roughly the rate the writer drains them. Default 0 (no wait) — opt-in. Co-Authored-By: Claude Opus 4.7 --- benchmarks/beam/bench.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/benchmarks/beam/bench.py b/benchmarks/beam/bench.py index e53a27e..71ba92c 100644 --- a/benchmarks/beam/bench.py +++ b/benchmarks/beam/bench.py @@ -385,6 +385,7 @@ def run_one_conversation( warmup_settle_seconds: float = 600, question_timeout: float = QUESTION_TIMEOUT_SECONDS, block_on_wiki_queue: bool = False, + wiki_wait_seconds: float = 0, ) -> dict: """Full A-B-C-D lifecycle for one conversation. Returns a small stats dict.""" conv_dir = run_dir / f"conv_{int(conv.conversation_id):03d}" @@ -420,6 +421,18 @@ def run_one_conversation( ) print(f"[B done] {warmup_stats}", flush=True) + # ---- Phase Bw (optional): wait for wiki writer to drain its queue ---- + # Between [B done] (extraction settled) and [C] (questions start), give + # the wiki writer/maintainer dedicated vLLM time so wikis develop before + # recall queries hit. Skipped when wiki_wait_seconds == 0 (default). + if wiki_wait_seconds and wiki_wait_seconds > 0: + print( + f"[Bw] waiting {wiki_wait_seconds:.0f}s for wiki queue to develop ...", + flush=True, + ) + time.sleep(wiki_wait_seconds) + print(f"[Bw done]", flush=True) + # ---- Phase C: answer all questions ---- probing = load_probing_questions(conv) print(f"[C] answering {sum(len(v) for v in probing.values())} questions " @@ -558,6 +571,14 @@ def _parse_args() -> argparse.Namespace: "continue async in the background regardless.") p.add_argument("--question-timeout", type=float, default=QUESTION_TIMEOUT_SECONDS, help="per-question HTTP timeout on /agent/query") + p.add_argument("--wiki-wait-seconds", type=float, default=0, + help="seconds to wait after [B done] before starting Phase C, " + "so the wiki writer can drain its queue against an " + "otherwise-idle vLLM (default 0 = no wait). Recommended " + "value 28800 (8 hours) for BEAM 100K convs — captures most " + "of the queue-drain benefit; longer waits hit diminishing " + "returns because the maintainer keeps queuing new triage " + "jobs at roughly the rate the writer drains them.") return p.parse_args() @@ -595,6 +616,7 @@ def main() -> int: warmup_settle_seconds=args.warmup_settle_seconds, question_timeout=args.question_timeout, block_on_wiki_queue=args.wait_for_wikis, + wiki_wait_seconds=args.wiki_wait_seconds, ) metas.append(meta) except Exception as e: From 2c9a54b927f212fa488e1667b4dca7b611bd1400 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Tue, 9 Jun 2026 09:21:33 +0100 Subject: [PATCH 15/17] =?UTF-8?q?chore(repo):=20public=20defaults=20?= =?UTF-8?q?=E2=80=94=20deepinfra=20LLM=20profile=20+=20wiki=20ON?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two coordinated default flips so a fresh clone runs end-to-end with just a hosted-provider API key, no workstation tunnel needed. 1) LLM profile defaults flipped from workstation-Qwen to deepinfra in three places that together define the bench's runtime LLM: - docker-compose.bench.yml LLM_PROFILE_BENCH default - docker-compose.bench.yml QWEN_BASE_URL/QWEN_MODEL pass-through (empty default; Python config supplies the real default) - benchmarks/beam/config.py judge defaults (QWEN_BASE_URL/MODEL/API_KEY) QWEN_* env-var names kept for backwards compat with existing .env.bench files; the values can point at any OpenAI-compatible endpoint. Self-hosted Qwen runs supported via the documented overrides in .env.bench.example. 2) WIKI_ENABLED defaults flipped from false to true in both braindb/wiki_scheduler.py and docker-compose.yml. The wiki pipeline (maintainer + writer) is part of what makes BrainDB a knowledge layer vs. a fact store, so a fresh stack starts it automatically. Idle ticks are pure SQL (no LLM cost). Set WIKI_ENABLED=false to opt out. CHANGELOG and .env.example updated to surface both changes. Co-Authored-By: Claude Opus 4.7 --- .env.bench.example | 18 +++++++++++++++--- .env.example | 8 ++++++++ CHANGELOG.md | 18 ++++++++++++++++++ benchmarks/beam/config.py | 17 ++++++++++++----- braindb/wiki_scheduler.py | 9 +++++---- docker-compose.bench.yml | 12 ++++++++---- docker-compose.yml | 2 +- 7 files changed, 67 insertions(+), 17 deletions(-) diff --git a/.env.bench.example b/.env.bench.example index e718f49..f837743 100644 --- a/.env.bench.example +++ b/.env.bench.example @@ -7,9 +7,10 @@ # DO NOT commit .env.bench — it's gitignored. Only this template is tracked. # --- LLM provider for the bench BrainDB API --- -# Bench defaults to local Qwen via vLLM on the workstation (per plan). -# Override here if you want a different provider for bench runs. -LLM_PROFILE_BENCH=vllm_workstation_qwen +# Public default: deepinfra (hosted, cheap, validated). For self-hosted Qwen +# runs, comment out the first line and uncomment the workstation line below. +LLM_PROFILE_BENCH=deepinfra +# LLM_PROFILE_BENCH=vllm_workstation_qwen # Optional model override (e.g. to try a smaller Qwen variant for bench). AGENT_MODEL= @@ -39,3 +40,14 @@ INGEST_POLL_INTERVAL_BENCH=3 # Agent verbose logging (every tool call printed). Bench default: true. AGENT_VERBOSE_BENCH=true + +# --- Judge LLM (scores answers against the BEAM rubric) --- +# Independent of LLM_PROFILE_BENCH. The judge speaks OpenAI-compatible REST, +# so any provider works. Names are QWEN_* for backwards compatibility; the +# values can point anywhere (deepinfra, OpenAI, self-hosted vLLM, ...). +# +# Defaults (in benchmarks/beam/config.py) target deepinfra. +# For self-hosted Qwen via vLLM, uncomment and adjust: +# QWEN_BASE_URL_BENCH=http://host.docker.internal:8010/v1 +# QWEN_MODEL_BENCH=cyankiwi/Qwen3.6-27B-AWQ-INT4 +# QWEN_API_KEY=EMPTY diff --git a/.env.example b/.env.example index 6884d1d..b091ede 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,14 @@ AGENT_MODEL= # (visible via `docker logs braindb_api -f`). Response payload unchanged. AGENT_VERBOSE=false +# Wiki pipeline (maintainer + writer) — ON by default. The wiki layer is +# what turns BrainDB's fact graph into navigable summary pages. Each tick +# (every WIKI_INTERVAL seconds) the scheduler runs cheap SQL bookkeeping +# and only fires LLM-heavy /wiki/maintain or /wiki/write when work is +# pending — idle ticks cost zero. Set to false if you want a pure +# fact-store stack with no LLM-heavy background work. +# WIKI_ENABLED=true + # Agent turn budget — how many tool-call turns the general /agent/query # is allowed before the SDK forces termination. Default 20. Lowering # this below ~15 degrades deep-research models (notably local Qwen via diff --git a/CHANGELOG.md b/CHANGELOG.md index 0a149fe..de19487 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,24 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Changed + +- **`WIKI_ENABLED` now defaults to `true`** in both `braindb/wiki_scheduler.py` + and `docker-compose.yml`. A fresh `docker compose up` starts the wiki + maintainer + writer automatically. The pipeline is part of what makes + BrainDB a knowledge layer; idle ticks cost zero LLM (cheap SQL only), + and the previous opt-in default was a friction point in onboarding. + Set `WIKI_ENABLED=false` in `.env` to opt out (e.g. pure fact-store + deployments). +- **Bench repo defaults flipped from workstation Qwen to deepinfra.** + `docker-compose.bench.yml`'s `LLM_PROFILE_BENCH` default is now + `deepinfra`; the judge endpoint defaults in `benchmarks/beam/config.py` + also point at deepinfra. Self-hosted Qwen runs are still supported — + set `LLM_PROFILE_BENCH=vllm_workstation_qwen` plus the judge overrides + in `.env.bench` (see `.env.bench.example`). + ## [0.4.0] — 2026-06-03 Headline: a focused pass on recall quality and the `view_tree` tool. The diff --git a/benchmarks/beam/config.py b/benchmarks/beam/config.py index 91764de..9b3b712 100644 --- a/benchmarks/beam/config.py +++ b/benchmarks/beam/config.py @@ -49,11 +49,18 @@ BENCH_API_BASE = os.getenv("BENCH_API_BASE", "http://localhost:8001") # ---- Judge LLM (OpenAI-compatible endpoint) --------------------------------- -# Bench defaults to the workstation Qwen via vLLM on the SSH tunnel. The judge -# is fully outside BrainDB; the bench API is independent of which judge runs. -QWEN_BASE_URL = os.getenv("QWEN_BASE_URL", "http://localhost:8010/v1") -QWEN_MODEL = os.getenv("QWEN_MODEL", "") # empty -> let vLLM's served model resolve -QWEN_API_KEY = os.getenv("QWEN_API_KEY", "EMPTY") # vLLM defaults to no auth +# Defaults target deepinfra so a fresh clone runs end-to-end with just an API key. +# For self-hosted Qwen via vLLM, set in .env.bench: +# QWEN_BASE_URL_BENCH=http://host.docker.internal:8010/v1 +# QWEN_MODEL_BENCH=cyankiwi/Qwen3.6-27B-AWQ-INT4 +# QWEN_API_KEY=EMPTY +# Env-var names are kept as QWEN_* for backwards compat with existing .env.bench +# files; the values can point at any OpenAI-compatible endpoint. `or` is used +# (not the second arg to getenv) so an empty string from compose pass-through +# also falls through to the default. +QWEN_BASE_URL = os.getenv("QWEN_BASE_URL") or "https://api.deepinfra.com/v1/openai" +QWEN_MODEL = os.getenv("QWEN_MODEL") or "google/gemma-4-31B-it" +QWEN_API_KEY = os.getenv("QWEN_API_KEY") or "" def _mask_password(url: str) -> str: diff --git a/braindb/wiki_scheduler.py b/braindb/wiki_scheduler.py index 1024278..c5d2ae1 100644 --- a/braindb/wiki_scheduler.py +++ b/braindb/wiki_scheduler.py @@ -29,11 +29,12 @@ # `maintain` runs concurrently alongside writers (1 maintain in flight, C1). WRITE_PARALLELISM = int(os.getenv("WIKI_WRITE_PARALLELISM", "3")) -# Master on/off for the whole wiki pipeline. Default OFF so bringing the -# stack up never auto-starts token-heavy work. Opt in explicitly with -# WIKI_ENABLED=true (or 1/yes/on). Model-agnostic; orthogonal to any LLM +# Master on/off for the whole wiki pipeline. Default ON — the wiki pipeline +# is part of what makes BrainDB a knowledge layer (vs. just a fact store), +# so a fresh stack starts the maintainer + writer automatically. Opt out +# explicitly with WIKI_ENABLED=false. Model-agnostic; orthogonal to any LLM # profile/provider. -WIKI_ENABLED = os.getenv("WIKI_ENABLED", "false").lower() in ("1", "true", "yes", "on") +WIKI_ENABLED = os.getenv("WIKI_ENABLED", "true").lower() in ("1", "true", "yes", "on") # HTTP read-timeout (seconds) the scheduler waits on a single /wiki/maintain # or /wiki/write call before its requests client gives up and moves on. # Bumped 600 → 1200 (10 → 20 min) after live observation on Qwen 27B AWQ-INT4 diff --git a/docker-compose.bench.yml b/docker-compose.bench.yml index f4e2e66..e8eed9e 100644 --- a/docker-compose.bench.yml +++ b/docker-compose.bench.yml @@ -69,7 +69,7 @@ services: HF_TOKEN: ${HF_TOKEN:-} # Bench defaults to local Qwen via the workstation tunnel — matches the # plan's Qwen-as-judge story. Override via .env.bench if needed. - LLM_PROFILE: ${LLM_PROFILE_BENCH:-vllm_workstation_qwen} + LLM_PROFILE: ${LLM_PROFILE_BENCH:-deepinfra} AGENT_MODEL: ${AGENT_MODEL:-} NVIDIA_NIM_API_KEY: ${NVIDIA_NIM_API_KEY:-} DEEPINFRA_API_KEY: ${DEEPINFRA_API_KEY:-} @@ -169,9 +169,13 @@ services: BENCH_ADMIN_DATABASE_URL: postgresql://braindb_bench:bench_local_only@postgres_bench:5432/postgres # bench.py mints per-conv URLs by appending `/braindb_conv_NNN` to this. BENCH_DB_BASE_URL: postgresql://braindb_bench:bench_local_only@postgres_bench:5432 - # Workstation Qwen via Docker host gateway. - QWEN_BASE_URL: ${QWEN_BASE_URL_BENCH:-http://host.docker.internal:8010/v1} - QWEN_API_KEY: ${VLLM_API_KEY:-EMPTY} + # Judge LLM (OpenAI-compatible endpoint). Empty default — Python config + # in benchmarks/beam/config.py supplies the real default (deepinfra). + # Set QWEN_BASE_URL_BENCH / QWEN_MODEL_BENCH in .env.bench to point at + # a self-hosted vLLM endpoint (e.g. workstation Qwen on port 8010). + QWEN_BASE_URL: ${QWEN_BASE_URL_BENCH:-} + QWEN_MODEL: ${QWEN_MODEL_BENCH:-} + QWEN_API_KEY: ${QWEN_API_KEY:-${VLLM_API_KEY:-}} HF_TOKEN: ${HF_TOKEN:-} # Docker CLI invocations need to know which compose file + project. COMPOSE_FILE: /app/docker-compose.bench.yml diff --git a/docker-compose.yml b/docker-compose.yml index da218f6..04e2f43 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -68,7 +68,7 @@ services: - local-network environment: BRAINDB_API_URL: http://api:${API_PORT:-8000} - WIKI_ENABLED: ${WIKI_ENABLED:-false} + WIKI_ENABLED: ${WIKI_ENABLED:-true} WIKI_INTERVAL: ${WIKI_INTERVAL:-60} volumes: - .:/app From af88fc694c28dd33878da48d45dca18e1a29813f Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Tue, 9 Jun 2026 09:57:26 +0100 Subject: [PATCH 16/17] docs(bench): align benchmarks/beam/README with v0.5.0 shipped state The README still described the harness as "Step 0 scaffolding" and hard-coded Qwen as the judge. Both are stale after the v0.5.0 work: the harness is shipped (adapter / runner / judge / score / warmup all present), and the judge defaults to deepinfra with self-hosted Qwen as the documented override. Co-Authored-By: Claude Opus 4.7 --- benchmarks/beam/README.md | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/benchmarks/beam/README.md b/benchmarks/beam/README.md index a40ea13..31dd1c9 100644 --- a/benchmarks/beam/README.md +++ b/benchmarks/beam/README.md @@ -4,9 +4,10 @@ Public benchmark harness running BrainDB against the **BEAM** memory benchmark ([arXiv:2510.27246](https://arxiv.org/abs/2510.27246), ICLR 2026, Tavakoli et al. of U Alberta + UMass Amherst). -**Status**: Step 0 scaffolding (this file, the bench compose, the upstream -submodule). Phase 1 harness code (`adapter.py`, `bench.py`, `judge.py`, -`score.py`, `warmup.py`) is being implemented in the same commit series. +**Status**: harness shipped. Dataset adapter, per-conversation runner, +warmup barrier, Qwen-as-judge, scoring, and the isolated bench stack +(`docker-compose.bench.yml`) are all in this directory. See +`adapter.py`, `bench.py`, `judge.py`, `score.py`, `warmup.py`. --- @@ -22,7 +23,7 @@ Nothing in our code paraphrases, copies, or reinterprets the benchmark. | Judge prompt | `from src.prompts import unified_llm_judge_base_prompt` — loaded at runtime from the submodule | Zero prompt content in our codebase; no risk of paraphrase or typo | | Adapter (dataset → BrainDB ingest) | `adapter.py` (ours, ~150 LOC) | Transparent: read the code | | Bench runner (per-conversation reset + ingest + warmup + answer) | `bench.py` (ours, ~120 LOC) | Same | -| Judge runner (calls Qwen with the upstream prompt) | `judge.py` (ours, ~80 LOC) | Anyone can re-judge our `answers.json` with any model | +| Judge runner (calls the configured LLM with the upstream prompt) | `judge.py` (ours, ~80 LOC) | Anyone can re-judge our `answers.json` with any model | | Eval wrapper (invokes upstream eval) | `score.py` (ours, ~50 LOC) | Just a thin caller | **The "no copy" rule** is satisfied at the strongest level: the judge prompt @@ -68,10 +69,13 @@ One-command restore if anything ever goes wrong. ## Caveats on the published number -We use **Qwen 3.6-27B** (local, via the workstation vLLM tunnel) as the LLM -judge. Published BEAM scores used GPT-4o (mem0) or Gemini-2.5-flash-lite -(Hindsight). **Our Qwen-judged number is NOT directly comparable to those -published numbers** — judge models systematically differ. +The judge is an OpenAI-compatible LLM endpoint. The repo default is +**deepinfra** (`google/gemma-4-31B-it`); self-hosted Qwen 3.6-27B via a +local vLLM endpoint is supported via `.env.bench` overrides +(see `.env.bench.example`). Published BEAM scores used GPT-4o (mem0) +or Gemini-2.5-flash-lite (Hindsight); **our judge will systematically +differ from theirs**, so the headline is not directly comparable to those +published numbers without a calibration sample. The mitigations: From a366b5da94ea1e6f9e3e6b57db7161cd152c8dcd Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Tue, 9 Jun 2026 11:35:23 +0100 Subject: [PATCH 17/17] chore(release): v0.5.0 Co-Authored-By: Claude Opus 4.7 --- CHANGELOG.md | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index de19487..03c0d46 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,7 +5,7 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [Unreleased] +## [0.5.0] — 2026-06-09 ### Changed diff --git a/pyproject.toml b/pyproject.toml index 6972484..2ecc691 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "braindb" -version = "0.4.0" +version = "0.5.0" description = "Persistent memory for LLM agents — thoughts, facts, sources, and behavioral rules with fuzzy + semantic search, graph traversal, and an internal agent." readme = "README.md" license = "Apache-2.0"