diff --git a/.env.bench.example b/.env.bench.example new file mode 100644 index 0000000..f837743 --- /dev/null +++ b/.env.bench.example @@ -0,0 +1,53 @@ +# 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 --- +# 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= + +# --- 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 + +# --- 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/.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/CHANGELOG.md b/CHANGELOG.md index 0a149fe..03c0d46 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). +## [0.5.0] — 2026-06-09 + +### 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/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/benchmarks/beam/README.md b/benchmarks/beam/README.md new file mode 100644 index 0000000..31dd1c9 --- /dev/null +++ b/benchmarks/beam/README.md @@ -0,0 +1,116 @@ +# 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**: 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`. + +--- + +## 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 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 +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 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) +- **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 + +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: + +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/__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..71ba92c --- /dev/null +++ b/benchmarks/beam/bench.py @@ -0,0 +1,654 @@ +"""BrainDB on BEAM bench runner. + +Per-conversation lifecycle (zero interleaving within a conversation): + + 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 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 + +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 + +# 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 +# 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}") + + +# ----------------------------- 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 _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`. + """ + 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 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: + admin.close() + return db_url + + +def restart_api_with_db(database_url: str) -> None: + """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) + 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( + "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: 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 = {} + 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. 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", + environment=new_env, + volumes=binds, + ports=port_bindings, + 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, + ) + # 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, aliases=aliases) + except docker_sdk.errors.NotFound: # type: ignore[attr-defined] + pass + container.start() + + +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 + 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: + """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). + + 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. + """ + 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 ---------------------------- + +def run_one_conversation( + conv: Conversation, + run_dir: Path, + *, + warmup_timeout: float = 43200, + 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}" + conv_dir.mkdir(parents=True, exist_ok=True) + + started = time.monotonic() + 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] 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) + warmup_stats = wait_for_warmup( + 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) + + # ---- 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 " + 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(), + "bench_db_name": conv_db_name, + } + 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=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 " + "stretches doing subagent / recall_memory work between save_fact " + "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 " + "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") + 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() + + +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, + block_on_wiki_queue=args.wait_for_wikis, + wiki_wait_seconds=args.wiki_wait_seconds, + ) + 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..9b3b712 --- /dev/null +++ b/benchmarks/beam/config.py @@ -0,0 +1,110 @@ +"""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", +) + +# 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") + +# ---- Judge LLM (OpenAI-compatible endpoint) --------------------------------- +# 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: + """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 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 not any(s in url for s in BENCH_DB_SENTINELS): + raise RuntimeError( + 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)}" + ) + + +__all__ = [ + "REPO_ROOT", + "DATA_BENCH_SOURCES", + "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", + "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/runner.Dockerfile b/benchmarks/beam/runner.Dockerfile new file mode 100644 index 0000000..01e1d7b --- /dev/null +++ b/benchmarks/beam/runner.Dockerfile @@ -0,0 +1,53 @@ +# benchmarks/beam/runner.Dockerfile +# +# Small orchestrator image for the bench. Lives separately from BrainDB's +# main image because: +# - 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) +# +# 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 + +# 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. `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 \ + docker + +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/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/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/benchmarks/beam/warmup.py b/benchmarks/beam/warmup.py new file mode 100644 index 0000000..40ff89e --- /dev/null +++ b/benchmarks/beam/warmup.py @@ -0,0 +1,232 @@ +"""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_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() - 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 + + +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 = 600.0, + consecutive_clear_required: int = 2, + poll_interval: float = 5.0, + timeout_seconds: float = 43200.0, + log_interval: float = 10.0, + verbose: bool = True, + block_on_wiki_queue: bool = False, + database_url: str | None = None, +) -> dict: + """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. + """ + url = database_url or BENCH_DATABASE_URL + assert_bench_database_url(url) + + conn = psycopg2.connect(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()) + activity_age = _seconds_since_last_activity(conn) + pending = _pending_wiki_jobs(conn) + entity_count = _entity_count(conn) + + 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 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"{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} " + 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_activity_age={_seconds_since_last_activity(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=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 " + "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. " + "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") + 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( + settle_seconds=args.settle_seconds, + consecutive_clear_required=args.consecutive_clear, + 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 + + +if __name__ == "__main__": + raise SystemExit(_cli()) 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 cb3fea2..ba2d9b0 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -42,11 +42,15 @@ 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 -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( @@ -109,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. +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). - 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. + 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: @@ -142,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" @@ -155,30 +189,57 @@ 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 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' - 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=, " 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) - 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: @@ -239,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) @@ -256,7 +319,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) 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 new file mode 100644 index 0000000..e8eed9e --- /dev/null +++ b/docker-compose.bench.yml @@ -0,0 +1,197 @@ +# 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 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) +# - 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: + # 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: + - 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 "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 + # plan's Qwen-as-judge story. Override via .env.bench if needed. + LLM_PROFILE: ${LLM_PROFILE_BENCH:-deepinfra} + 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 + + # 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 + # 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 + 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 + +networks: + bench-network: + name: braindb_bench_network 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 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" 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)}"