diff --git a/.env.example b/.env.example index c26571d..6884d1d 100644 --- a/.env.example +++ b/.env.example @@ -35,6 +35,48 @@ AGENT_MODEL= # (visible via `docker logs braindb_api -f`). Response payload unchanged. AGENT_VERBOSE=false +# 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 +# vLLM); raising it costs more LLM calls per query. The wiki maintainer +# / writer and the ingest watcher pass their own per-call values +# (30/30/40/30) and are unaffected by this default. +# AGENT_MAX_TURNS=20 + +# How many turns from the end of the run the agent gets a synthetic +# "start wrapping up" reminder injected as a user message. Default 8. +# Set to 0 to disable the reminder entirely (the SDK will still +# terminate at max_turns, but the model gets no warning). The reminder +# tone is automatic: soft "start wrapping up" when max_turns > 5, +# hard "call final_answer NOW" when max_turns <= 5 (which covers the +# Layer 4 retry path). +# AGENT_COUNTDOWN_THRESHOLD=8 + # Ingest watcher poll interval (seconds) — how often the watcher sidecar # scans data/sources/ for new files to ingest. INGEST_POLL_INTERVAL=7 + +# Wiki scheduler HTTP read-timeout (seconds) on /wiki/maintain and +# /wiki/write calls. Default 1200 (20 min). Local quantised models +# (Qwen 27B AWQ-INT4 on vLLM) routinely take 6-15 min for a full wiki +# body; setting this below ~600 caused the scheduler to give up while +# the api kept working — queue drained slower than reality. Raise if +# you see "Read timed out" in the scheduler log AND the corresponding +# write actually committed (check `wikis_ext.revision`); lower only if +# you specifically want quicker scheduler turnover. The api itself is +# unbounded by this; this only controls the scheduler's patience. +# WIKI_AGENT_TIMEOUT=1200 + +# Per-wiki cooldown on attach claims (seconds). Default 300 (5 min). +# Once the OLDEST pending attach for a given wiki is this old, the +# writer claims ALL pending attaches for that wiki in a single batch. +# Below the cooldown, fresh attaches keep accumulating — they don't +# trigger a writer fire. Lets the writer fire once per cooldown window +# instead of once per attach job; on a hot subject like a high-volume +# person/topic wiki, this collapses 5-10 separate full-body +# regenerations into 1 per window — ~80% LLM cost reduction on the +# pattern we observed today. Self-limiting: each fire scoops up the +# whole pending queue for that wiki. Set to 0 to disable (revert to +# the old "fire on every attach" behaviour). Affects ATTACH only; +# consolidate and create paths are unchanged. +# WIKI_ATTACH_COOLDOWN_SECONDS=300 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml new file mode 100644 index 0000000..fd22ed5 --- /dev/null +++ b/.github/workflows/test.yml @@ -0,0 +1,79 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + branches: [main] + +jobs: + validator-tests: + runs-on: ubuntu-latest + timeout-minutes: 10 + + services: + postgres: + image: pgvector/pgvector:pg16 + env: + POSTGRES_PASSWORD: password + POSTGRES_DB: braindb + ports: + - 5432:5432 + options: >- + --health-cmd pg_isready + --health-interval 5s + --health-timeout 5s + --health-retries 10 + + steps: + - uses: actions/checkout@v4 + + - name: Enable required postgres extensions + run: | + PGPASSWORD=password psql -h localhost -U postgres -d braindb \ + -c "CREATE EXTENSION IF NOT EXISTS pg_trgm; CREATE EXTENSION IF NOT EXISTS vector;" + + - name: Configure .env for the CI stack + run: | + cat > .env <<'EOF' + DATABASE_URL=postgresql://postgres:password@host.docker.internal:5432/braindb + API_PORT=8000 + LLM_PROFILE=deepinfra + DEEPINFRA_API_KEY=ci-placeholder-key-not-used + AGENT_VERBOSE=false + WIKI_ENABLED=false + EOF + + - name: Create the local-network the compose file expects + run: docker network create local-network + + - name: Bring up the stack + run: docker compose up -d --build + + - name: Wait for /health + run: | + for i in $(seq 1 60); do + if curl -sf http://localhost:8000/health > /dev/null; then + echo "API healthy after ${i} attempts" + curl -s http://localhost:8000/health + exit 0 + fi + sleep 2 + done + echo "API failed to become healthy" + docker logs braindb_api --tail 100 + exit 1 + + - name: Install pytest into the api container + run: docker exec braindb_api pip install pytest pytest-asyncio --quiet + + - name: Run validator + handoff unit tests + run: | + docker exec braindb_api python -m pytest \ + tests/test_final_answer_rename.py \ + tests/test_handoff_hooks.py \ + -v + + - name: Dump api logs on failure + if: failure() + run: docker logs braindb_api --tail 200 diff --git a/.gitignore b/.gitignore index 005c3b6..3a17d0f 100644 --- a/.gitignore +++ b/.gitignore @@ -36,3 +36,6 @@ Thumbs.db data/sources/* !data/sources/.gitkeep !data/sources/README.md + +# Wiki review exports — generated, read-only inspection output +data/wiki_review/ diff --git a/BRAINDB_GUIDE.md b/BRAINDB_GUIDE.md index 8917c31..7a3e95c 100644 --- a/BRAINDB_GUIDE.md +++ b/BRAINDB_GUIDE.md @@ -22,20 +22,25 @@ The API runs at **http://localhost:8000**. Everything is done via HTTP calls. ### Before answering anything non-trivial, always call: ``` POST /api/v1/memory/context -{"queries": ["topic 1", "topic 2"], "max_depth": 3, "max_results": 15} +{"queries": ["bare-keyword-1", "bare-keyword-2", "one broader phrase"], "max_depth": 3} ``` This returns: -- Direct matches (fuzzy + full-text) across all queries, merged by best score +- Direct matches (keyword-mediated fuzzy + keyword-mediated embedding) across all queries - Graph-connected entities up to 3 hops away (relevance fades: 100% -> 60% -> 30%) +- Two-level diversity quota applied: per-search-term reservation (each query gets a guaranteed share) + per-keyword halving cap on the open remainder - Always-on rules (always injected regardless of query) -Each item has a `final_rank` score. Trust higher-ranked items more. +Each item has a `final_rank` score. Trust higher-ranked items more. `max_results` defaults to 30; the scoring pool internally considers up to 500 candidates per query so narrow keywords aren't excluded before they're evaluated. + +**Query strategy.** Prefer **multiple narrow queries** (single keywords, bare names) over one long sentence. Keywords are short, so a short query matches them at high pg_trgm similarity; a long phrase dilutes the trigram set and pushes narrow-subject facts down the ranking. Examples: -You can also pass a single query for backward compatibility: ``` -{"query": "single topic", "max_depth": 3} +GOOD: "queries": ["Petros", "Selonda Saronikos fish farm", "Dimitrios manager"] +BAD: "queries": ["Petros person identity profile relation to Dimitris"] ``` +The per-search-term quota reserves slots for each query you pass, so the bare-keyword query is guaranteed to surface its specific facts even when paired with broader angles. Single `query` (string) still works for backward compatibility. + ### After learning something new, save it: ``` POST /api/v1/entities/facts — for objective facts @@ -74,9 +79,19 @@ curl "http://localhost:8000/api/v1/entities?entity_type=fact&source=user-stated& Query parameters: `entity_type`, `keyword`, `source`, `min_importance` (0-1), `limit` (1-200, default 50), `offset` (default 0). ### Get Entity by ID +The **only full-content read**. Multi-item calls (context/search/list) return +~1K previews ending `--truncated … get_entity("")`; come here for the +whole body. ```bash curl http://localhost:8000/api/v1/entities/ +# Large body? page it (don't pull it whole): +curl "http://localhost:8000/api/v1/entities/?offset=0&limit=8000" ``` +With `offset`/`limit` the response adds `content_meta`: +`{total_chars, offset, returned, next_offset}` — keep fetching `next_offset` +until it is `null`. Default (no params) = full body, unchanged. For big +documents, prefer delegating the read to a subagent via `/api/v1/agent/query` +so the content never floods the caller's context. ### Delete Entity ```bash @@ -226,14 +241,23 @@ curl -X POST http://localhost:8000/api/v1/memory/search \ curl -X POST http://localhost:8000/api/v1/memory/context \ -H "Content-Type: application/json" \ -d '{ - "queries": ["user profile expertise", "project architecture decisions"], + "queries": ["user-profile", "expertise", "project-decision"], "max_depth": 3, - "max_results": 15, "include_always_on_rules": true }' ``` -Each query runs fuzzy + full-text search independently. Seeds are merged keeping the **best score** per entity. One graph expansion runs on the combined seed set. +Each query runs through TWO keyword-mediated pathways in parallel: +- **Fuzzy** — `pg_trgm similarity(content, query)` over keyword entities. +- **Embedding** — Qwen3-Embedding-0.6B (1024-dim) cosine similarity between the query and keyword-entity embeddings. + +Entities surface via `tagged_with` from the matched keywords. Per-entity score = `max(matched-keyword similarity)` on each pathway. Both signals are merged with the geometric mean (configurable `missing_signal_penalty` when only one signal fires). + +After scoring, **two diversity quotas** apply: +1. **Per-search-term** — each query in `queries[]` reserves `ceil(max_results × per_query_share / num_queries)` slots filled from its own top-ranked entities. Knob: `per_query_share` (default 0.5; set to 0 to disable). +2. **Per-keyword (halving)** — walking the remaining slots in `final_rank`-desc order, each new dominant keyword gets a halving allowance (50% / 25% / 12.5% ..., floor 1). Knob: `keyword_quota_halving` (default 0.5; set to 1.0 to disable). + +`max_results` defaults to 30 (LLM-visible cap). The internal scoring pool considers up to 500 keyword neighbours per query (`scoring_pool_keyword_neighbors`) and up to 500 fuzzy candidates (`scoring_pool_fuzzy`) — cheap pure-SQL/vector work, so narrow keywords aren't excluded before they're evaluated. None of these knobs are env-driven; tune them in [`braindb/config.py`](braindb/config.py) if needed. **Single query** (backward-compatible): ```bash @@ -277,8 +301,14 @@ curl "http://localhost:8000/api/v1/memory/log?since=2026-04-08T00:00:00Z" Response includes: `id`, `timestamp`, `operation`, `entity_type`, `entity_id`, `details`, `context_note`. -### Read-only SQL -For ad-hoc exploration. Only `SELECT` and `WITH` queries; 5s timeout; 1000 row limit. +### Read-only SQL — EXCEPTION tool, not for recall + +⚠ This is **not** a recall/discovery path. A flat SELECT has no embeddings, no +graph, no ranking — it discards everything BrainDB is built for. Default to +`POST /api/v1/memory/context` (and delegated `/api/v1/agent/query`) for all +recall, discovery, and understanding. Use `/memory/sql` **only** for a +specific structured/aggregate question those cannot express (counts, GROUP BY, +activity-log joins). Only `SELECT` and `WITH` queries; 5s timeout; 1000 row limit. ```bash curl -X POST http://localhost:8000/api/v1/memory/sql \ @@ -306,18 +336,17 @@ curl -X POST http://localhost:8000/api/v1/entities/datasources/ingest \ ### BrainDB Agent — natural language queries -`POST /api/v1/agent/query` — instead of orchestrating individual API calls, send a plain English request and let BrainDB's internal agent handle it. The agent uses the OpenAI Agents SDK with LiteLLM (provider pluggable via `LLM_PROFILE` — default `deepinfra`, `nim` also supported) and has access to all 21 BrainDB operations as function tools. +`POST /api/v1/agent/query` — instead of orchestrating individual API calls, send a plain English request and let BrainDB's internal agent handle it. The agent uses the OpenAI Agents SDK with LiteLLM (provider pluggable via `LLM_PROFILE` — **`deepinfra` with `google/gemma-4-31B-it` is the recommended default**; `nim` and local vLLM are also supported) and has access to all 21 BrainDB operations as function tools. ```bash curl -X POST http://localhost:8000/api/v1/agent/query \ -H "Content-Type: application/json" \ - -d '{ - "query": "What do you know about the user role and recent projects?", - "max_turns": 15 - }' -# {"answer": "The user is ...", "max_turns": 15} + -d '{"query": "What do you know about the user role and recent projects?"}' +# {"answer": "The user is ...", "max_turns": 20} ``` +(`max_turns` is optional; the default — currently 20 — is used when omitted.) + **Save via the agent**: ```bash curl -X POST http://localhost:8000/api/v1/agent/query \ @@ -332,12 +361,12 @@ curl -X POST http://localhost:8000/api/v1/agent/query \ -d '{"query":"Delegate to a subagent: find near-duplicate facts and return top 10 pairs with their IDs."}' ``` -The agent has these tools internally: `recall_memory`, `quick_search`, `save_fact`, `save_thought`, `save_source`, `save_rule`, `ingest_file`, `get_entity`, `list_entities`, `update_entity`, `delete_entity`, `create_relation`, `view_entity_relations`, `delete_relation`, `view_tree`, `search_sql`, `view_log`, `get_stats`, `generate_embeddings`, `delegate_to_subagent`, `submit_result`. +The agent has these tools internally: `recall_memory`, `quick_search`, `save_fact`, `save_thought`, `save_source`, `save_rule`, `ingest_file`, `get_entity`, `list_entities`, `update_entity`, `delete_entity`, `create_relation`, `view_entity_relations`, `delete_relation`, `view_tree`, `search_sql`, `view_log`, `get_stats`, `generate_embeddings`, `delegate_to_subagent`, `final_answer`. **Setup (pick a provider)**: -- **DeepInfra (default)**: set `LLM_PROFILE=deepinfra` and `DEEPINFRA_API_KEY=...` in `.env`. Get a key at https://deepinfra.com/ -- **NVIDIA NIM**: set `LLM_PROFILE=nim` and `NVIDIA_NIM_API_KEY=...` in `.env`. Get a key at https://build.nvidia.com/ -- **Self-hosted vLLM**: set `LLM_PROFILE=vllm_workstation` for a vLLM server bound to the Docker host's loopback at `:8002`. No API key needed if the server runs without auth. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to add your own self-hosted profile. +- **DeepInfra — recommended default**: set `LLM_PROFILE=deepinfra` and `DEEPINFRA_API_KEY=...` in `.env`. Fast (5–30s per agent call), cheap, validated end-to-end. Get a key at https://deepinfra.com/ +- **NVIDIA NIM** (free-tier alternative): set `LLM_PROFILE=nim` and `NVIDIA_NIM_API_KEY=...` in `.env`. Get a key at https://build.nvidia.com/ +- **Self-hosted vLLM** (advanced / offline / requires GPU workstation): set `LLM_PROFILE=vllm_workstation` (or `..._qwen`, `..._gemma`) — points at a vLLM server bound to the Docker host's loopback at `:8002` / `:8010` / `:8009` respectively. Reach it from the docker network via an SSH tunnel if the GPU is on a remote machine. No API key needed if the server runs without auth. See [CONTRIBUTING.md](CONTRIBUTING.md) for how to add your own self-hosted profile. - Profiles live in `braindb/config.py::_LLM_PROFILES`. Add new providers there (e.g. `together`, `openai`) by adding a dict entry — no code change required. - Optional override: set `AGENT_MODEL=` in `.env` to use a non-default model for the active profile. @@ -384,13 +413,19 @@ This is complementary to `source_entity_id` (on facts — links to a specific so ## How Search Works -The search uses a 4-tier scoring system: +Two different paths, two different scoring models: + +**`POST /api/v1/memory/search`** (and the `quick_search` agent tool) — **content-matching** with a 4-tier score against entity content directly: 1. **Full-text AND match** (all query words match) — highest weight (1.0) 2. **Full-text OR match** (any query word matches) — lower weight (0.3) 3. **Content trigram similarity** — fuzzy character matching (0.5) 4. **Title trigram similarity** — fuzzy title matching (0.3) -This means specific queries with terms that appear in stored content work best. Vague queries with stop words ("everything about X") may return fewer results. If you get 0 results, reformulate with more specific terms. +This is for "find me entities whose CONTENT mentions these terms" — useful for arbitrary text matching, but it dilutes when the query is much longer than what's in the entity. + +**`POST /api/v1/memory/context`** (the sophisticated path) — **keyword-mediated**. Both the fuzzy and embedding pathways match the query against keyword entities (not entity bodies); entities surface via `tagged_with`. Then graph traversal, decay, two-level diversity quota, ranking. See the "Context" section above for the full pipeline. + +Use `/memory/search` for raw text matching; use `/memory/context` for everything that involves *understanding* a subject. If you get 0 results from either, reformulate with more specific terms. --- @@ -422,7 +457,7 @@ The `final_rank` in context results already accounts for decay. 3. **Notes are a log** — use `notes` on any entity to record how your understanding evolved 4. **always_on rules are limited to 10** — keep them high-signal; use on-demand rules for specifics 5. **access_count reinforces memory** — things you retrieve often stay important longer -6. **Multi-query for better recall** — use `queries` (array) instead of `query` (single) to search multiple angles at once +6. **Multi-query for better recall** — use `queries` (array) instead of `query` (single) AND prefer multiple **narrow** queries (single keywords / bare names) over one long phrase. Each query in `queries[]` reserves a share of result slots, so a bare keyword is guaranteed to surface its facts. `max_results` defaults to 30. 7. **Content should be concise** — 1-2 sentences, standalone, using full terms (not abbreviations) 8. **Use the tree endpoint** to explore how an entity connects to others: `GET /memory/tree/` 9. **Use the list endpoint** to browse entities: `GET /entities?entity_type=fact&limit=50` diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..21fff4e --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,134 @@ +# Changelog + +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.2.0] — 2026-05-24 + +The first substantial release beyond the v0.1.0 memory-store baseline. The +headline addition is the **wiki layer**: an always-on background pipeline +that turns the entity graph into self-maintaining, human-readable pages — +the same hands-off posture as the file watcher. + +### Added + +- **Wiki pipeline** (`braindb/wiki_scheduler.py`, `braindb/routers/wiki.py`): + the in-house agent decides per-orphan whether to *attach* to an existing + wiki, *create* a new one, *consolidate* duplicates, or *skip*. A separate + writer agent then researches and writes/maintains each page, citing every + claim with `[[ref:UUID]]`, with auto-self-healing on conflated subjects. +- **Wiki HTTP endpoints**: `POST /api/v1/wiki/cron` (orphan scan, idempotent), + `POST /api/v1/wiki/maintain` (one triage decision per call), + `POST /api/v1/wiki/write` (one writer pass), `GET /api/v1/wiki/jobs` + (queue visibility). Normal operation is the scheduler sidecar; these are + for hand-driving / observability. +- **Wiki section-edit tools**: `read_wiki_outline`, `read_wiki_section`, + `edit_wiki_section`, `delete_wiki_section`, `validate_wiki` — let the + writer do surgical edits on large pages without rewriting the full body. +- **Writer context-handoff**: when the writer's running context grows past + a budget, it calls `handoff_to_successor` with a structured brief; the + router respawns a successor agent with fresh context. Bounded by depth. +- **Typed agent termination**: every agent finish (`/agent/query`, + maintainer, writer, subagent) is now a Pydantic model — schema-validated, + no scraped free-text. Models live in `braindb/agent/schemas.py`. +- **Layer-4 retry-with-correction**: when a run ends without + `final_answer`, the runner appends a synthetic correction message and + re-invokes once with a small budget; recovers transparently. +- **`CountdownHooks` nudges**: a context-aware "wrap up" message arrives + before `max_turns` is exhausted; a separate token-budget watch nudges + the writer toward handoff when the conversation is getting big. +- **Auto-consolidation of duplicate wikis** via the maintainer's + `consolidate` action, with reversible `wiki_revise` snapshots. +- **Per-wiki cooldown for attaches** in the scheduler so cron ticks don't + thrash the same wiki across overlapping ticks. +- **Local vLLM profiles**: `vllm_workstation`, `vllm_workstation_qwen`, + `vllm_workstation_gemma` for running against your own GPU box. +- **Tests**: session-teardown fixture in `tests/conftest.py` that sweeps + any `_pytest_*` keyword artefacts that escape per-test cleanup. +- **CI**: minimal GitHub Actions workflow runs the typed-final + handoff + unit tests on every PR + push to main. + +### Configurable + +New environment variables exposed in `.env.example` and consumed by the +api / wiki scheduler: + +- `WIKI_ENABLED` — opt-in flag for the wiki scheduler (default `false`). +- `WIKI_INTERVAL` — scheduler tick in seconds (default `60`). +- `WIKI_FRESHNESS_MINUTES` — orphan eligibility gate; an entity must be + this old before it's picked up for triage (default `30`). +- `WIKI_ATTACH_COOLDOWN_SECONDS` — per-wiki throttle between attach claims. +- `WIKI_AGENT_TIMEOUT` — HTTP timeout the scheduler uses for maintainer / + writer calls (default `1200` seconds, i.e. 20 minutes). +- `AGENT_VERBOSE` — log every agent tool call with args and result preview + (default `false`). + +### Changed + +- **Recall is keyword-mediated**: `/memory/context` now matches both the + fuzzy (pg_trgm) and the embedding pathway against keyword entities, then + surfaces facts via `tagged_with`. Two-level diversity quota + (per-search-term + per-keyword, geometric decay) prevents one popular + hub keyword from monopolising top-N. Narrow short queries outperform + long phrases for keyword recall. +- **Multi-item recall returns previews**: `/memory/context` and + `list_entities` now return short (~1 KB) previews per item; the full + body is fetched on demand via `GET /api/v1/entities/{id}`, with optional + `?offset=&limit=` paging for large documents. Keeps the LLM-visible + context tight without losing access to the underlying content. +- **`deepinfra` (`google/gemma-4-31B-it`) promoted as the recommended + default** across README, BRAINDB_GUIDE, CLAUDE, and CONTRIBUTING. Fast + (5–30s per agent call), cheap, validated end-to-end. The `vllm_*` + profiles are now documented as advanced / offline / requires GPU. +- **`WIKI_ENABLED` defaults to `false`** in compose so the scheduler + sidecar boots but doesn't tick until explicitly opted in — keeps a + fresh clone from spending on the LLM by accident. +- **Agent `max_turns` defaults bumped** (15 → 20) and `countdown_threshold` + (5 → 8) after live observation on slower providers; deepinfra/Gemma is + unaffected because it finishes well before the budget. +- **Wiki scheduler** collapsed three timers into one gated loop — no idle + LLM spend, parallel maintain + writer fan-out per tick. +- **Skill files**: agent-call timeout guidance bumped to 10 minutes max + for slow providers; wiki awareness + always-ASK-before-saving added. + +### Fixed + +- **Double-escaped JSON tool-call payload** (Qwen AWQ-INT4 quirk): + `_maybe_parse_json_string` now unwraps the second layer when needed. + Compliant providers (deepinfra/OpenAI/Anthropic via LiteLLM) unaffected. +- **JSON-string tool-call payload** (vLLM/Qwen format): typed schemas + accept `arguments.payload` as either a JSON object or a JSON-encoded + string of a dict; the LLM-visible contract is unchanged. +- **Writer no-op on already-cited members** no longer leaks the orphan + back into the triage queue — it now closes the loop cleanly. +- **Big-body writes** retry on transient `BadRequestError` and stub out + the body when the provider truncates, so the wiki isn't lost. +- **Reference-by-catalog-number** in maintainer prompts replaced the + earlier uuid form to stop hallucinated wiki IDs. +- **Stale assigned jobs** in `wiki_job` are reclaimable on the next cron + tick (stale-lease). +- **`output_type` dropped from agent builder** — restored tool use; typed + `final_answer` still enforced via mutable-slot capture. +- **Compose**: no more `--reload` on the api command — code changes apply + explicitly via `docker compose up -d --no-deps --force-recreate api`, + preventing mid-run reloads that broke in-flight LLM calls. + +### Upgrading from 0.1.0 + +Migration `005_wiki_system.py` adds two new tables (`wikis_ext`, +`wiki_job`) and the `wiki` entity type. It runs automatically on +container startup via `alembic upgrade head` (already in the api +`command`). Existing rows are untouched; no manual action required. + +The wiki scheduler ships **disabled by default** — set +`WIKI_ENABLED=true` in `.env` to opt in. This prevents an upgraded +deployment from spending on the LLM until the operator says go. + +## [0.1.0] — initial public baseline + +Memory store: entities (`thought`, `fact`, `source`, `datasource`, `rule`), +relations, `pg_trgm` + `pgvector` retrieval, the BrainDB agent +(`/api/v1/agent/query`), the always-on file watcher (`data/sources/`), +Claude Code skills. diff --git a/CLAUDE.md b/CLAUDE.md index f79b079..77ceac5 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -5,6 +5,39 @@ The API runs at **http://localhost:8000**. --- +## ⚠ TOOL PRIORITY — read this first, it overrides habit + +BrainDB's entire value is the **graph + embeddings + ranking**. Recall and +understanding must go through the sophisticated retrieval, never a flat SQL +`SELECT`. + +1. **`POST /api/v1/memory/context`** (multi-query) — the default for ALL + recall, discovery, disambiguation, "what do we know about X". BOTH + the fuzzy and the embedding pathways are **keyword-mediated**: the + query is matched against keyword-entity content (via pg_trgm) and + keyword embeddings, then entities surface via `tagged_with`. A + two-level diversity quota (per-search-term + per-keyword, geometric + decay) keeps results balanced + graph traversal + temporal decay + + `final_rank`. +2. **`POST /api/v1/agent/query`** (ask it to *delegate to a subagent* for + anything multi-step) — research/investigation that needs several hops. +3. `GET /api/v1/entities…`, `/memory/tree/`, `/entities//relations` — + targeted structure lookups. +4. **`POST /api/v1/memory/sql` — exception ONLY.** A flat SELECT throws away + embeddings, graph and ranking. Use it solely for a specific + structured/aggregate question (counts, GROUP BY, activity-log joins) the + above genuinely cannot express. **Never** for recall, discovery, + similarity, or understanding. If you're using SQL to *find* or *understand* + something, you're doing it wrong — use `/memory/context`. + +**Previews vs full read:** all multi-item calls return short previews +(~1K/item; a clipped one ends `--truncated … get_entity("")`). Read a +full body only by id: `GET /api/v1/entities/{id}`. For a large body, page it +with `?offset=&limit=` (follow `content_meta.next_offset`) or delegate it to a +subagent — never pull whole documents into context. + +--- + ## At the Start of Every Session Before doing any work, consult your memory: @@ -13,10 +46,10 @@ Before doing any work, consult your memory: # 1. Get always-on rules (behavioral guidelines) curl -s http://localhost:8000/api/v1/memory/rules -# 2. Get context — use multi-query for better coverage +# 2. Get context — use multi-query with NARROW queries for better coverage curl -s -X POST http://localhost:8000/api/v1/memory/context \ -H "Content-Type: application/json" \ - -d '{"queries": ["user profile background expertise", ""], "max_depth": 3, "max_results": 15}' + -d '{"queries": ["user-profile", "Dimitrios", ""], "max_depth": 3}' ``` The context response gives you `items` (ranked memories) and `always_on_rules` (always injected). @@ -24,6 +57,8 @@ Trust higher `final_rank` items more. Check `depth` — depth 0 is a direct matc Multi-query runs each query independently, merges seeds (keeping the best score per entity), then does one graph expansion on the combined set. Use it to cover multiple angles in a single call. +**Query strategy**: prefer multiple **narrow** queries (single keywords / bare names) alongside one broader phrase, NOT a single long sentence. Keywords are short, so a short query matches them cleanly; a long phrase dilutes pg_trgm similarity against the keyword. The per-search-term diversity quota reserves slots for each query you pass, so a bare name like `"Petros"` will always surface its specific facts even when paired with broader semantic angles. `max_results` defaults to 30 — leave it unless you have a reason. + If results seem weak, retry with reformulated queries (up to 2 times). --- @@ -156,7 +191,7 @@ When debugging the agent: set `AGENT_VERBOSE=true` in `.env` and watch `docker l ## Important Notes -- `.env` contains real DB credentials and provider API keys (`DEEPINFRA_API_KEY`, `NVIDIA_NIM_API_KEY`, etc.) — **never commit it**, it is in `.gitignore`. Active provider is picked by `LLM_PROFILE` (see `braindb/config.py::_LLM_PROFILES`). +- `.env` contains real DB credentials and provider API keys (`DEEPINFRA_API_KEY`, `NVIDIA_NIM_API_KEY`, etc.) — **never commit it**, it is in `.gitignore`. Active provider is picked by `LLM_PROFILE` (see `braindb/config.py::_LLM_PROFILES`). `LLM_PROFILE=deepinfra` (model `google/gemma-4-31B-it`) is the recommended starting point — fast, cheap, validated end-to-end; the `vllm_*` profiles are for advanced/offline use and need a workstation GPU + SSH tunnel. - Always-on rules (priority 100, `always_on: true`) are returned on every `/memory/context` call - `notes` field on any entity or relation is for running commentary — append observations over time - Keywords are stored as both a `TEXT[]` column on the entity AND as separate keyword entities linked via `tagged_with` relations (the keyword entities carry the embeddings for semantic search) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 396cec8..01efdca 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -10,7 +10,7 @@ Prerequisites: Docker Desktop (or any Docker Engine), Python 3.12, a Postgres 16 git clone braindb cd braindb cp .env.example .env -# edit .env — set DATABASE_URL, pick an LLM_PROFILE, fill in the matching API key +# edit .env — set DATABASE_URL; recommended LLM_PROFILE=deepinfra + DEEPINFRA_API_KEY (or any other profile) docker network create local-network # one-time; docker-compose expects this docker compose up -d --build @@ -36,6 +36,8 @@ See [`tests/README.md`](tests/README.md) for what is and isn't covered. ## Adding a new LLM provider +The reference implementation and recommended default is `deepinfra` with `google/gemma-4-31B-it` — fast, cheap, validated end-to-end on the wiki/maintainer/writer pipeline. Other providers are configured the same way. + LiteLLM does the heavy lifting — providers are selected by a prefix in the model string. To add a provider: 1. Open [`braindb/config.py`](braindb/config.py) and add an entry to `_LLM_PROFILES`: diff --git a/README.md b/README.md index 79854dc..20d314d 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ A memory database and REST API for LLM agents. Store and retrieve thoughts, facts, sources, documents, and behavioral rules — with fuzzy + semantic keyword search, graph traversal up to 3 hops, temporal decay, and always-on rule injection. Built to be driven externally by an LLM via HTTP calls. -It also ships with **its own internal agent** (OpenAI Agents SDK + LiteLLM with pluggable providers — DeepInfra by default, NIM / others via config) so external callers can talk to BrainDB in plain English via a single endpoint instead of orchestrating individual API calls. +It also ships with **its own internal agent** (OpenAI Agents SDK + LiteLLM with pluggable providers — **DeepInfra is the recommended default**, with NIM / local vLLM / others available via config) so external callers can talk to BrainDB in plain English via a single endpoint instead of orchestrating individual API calls. --- @@ -33,7 +33,7 @@ Relations connect any two entities with `relation_type`, `relevance_score`, `imp ## Setup -BrainDB runs as two Docker services — `api` and `watcher` — against an **external** PostgreSQL you provide. The whole setup is six steps. +BrainDB runs as three Docker services — `api`, `watcher` (auto-ingests files), and `wiki_scheduler` (auto-maintains wikis) — against an **external** PostgreSQL you provide. The two sidecars are hands-off: you never call the pipeline by hand. The whole setup is six steps. ### 1. Prerequisites @@ -72,11 +72,11 @@ Any reachable hostname/IP works — the connecting user just needs network acces ### 4. Pick an LLM provider (for the internal agent) -The agent talks to any LiteLLM-supported backend. BrainDB ships with two profiles pre-configured: **DeepInfra** (default, fast, paid) and **NVIDIA NIM** (free tier, can be flaky). +The agent talks to any LiteLLM-supported backend. **Recommended for new users: `deepinfra` with `google/gemma-4-31B-it`** — fast (5–30s per agent call), cheap, validated end-to-end on the wiki/maintainer/writer pipeline. `nim` is a free-tier fallback (occasionally flaky). The `vllm_*` profiles run a local model on your own GPU workstation — useful for offline / cost-free experiments, but require a running vLLM server reachable from the docker network (typically via SSH tunnel). In `.env`: ``` -LLM_PROFILE=deepinfra # or 'nim' — default is 'deepinfra' +LLM_PROFILE=deepinfra # recommended default DEEPINFRA_API_KEY=... # if profile=deepinfra — get from https://deepinfra.com/ NVIDIA_NIM_API_KEY=... # if profile=nim — get from https://build.nvidia.com/ ``` @@ -144,17 +144,20 @@ See [BRAINDB_GUIDE.md](BRAINDB_GUIDE.md) for full API reference with curl exampl ## How Retrieval Works -`POST /api/v1/memory/context` is the main endpoint: +`POST /api/v1/memory/context` is the main endpoint. **Keywords are the indexing layer** — both the fuzzy and the embedding pathways match the query against keyword-entity content / embeddings, then entities surface via `tagged_with` edges. A keyword tagged on many entities is the hub; you don't need explicit `elaborates` / `refers_to` edges for an entity to be findable, as long as it has the right keywords. -1. **Multi-query search** — pass `queries: ["topic1", "topic2"]` to search multiple angles at once. Each query runs 4-tier scoring (AND fulltext, OR fulltext fallback, content trigram, title trigram), seeds are merged keeping the best score per entity. -2. **Keyword embeddings** — query terms are also matched against keyword entity embeddings (Qwen3-Embedding-0.6B, 1024-dim, cosine similarity). Text and embedding scores are combined via geometric mean (with a configurable penalty when only one signal matches). -3. **Graph traversal** up to 3 hops via relations, relevance fading: `1.0 → 0.6 → 0.3` -4. **Temporal decay** — memories fade over time, strengthen on access -5. **Final rank** = `combined_score × effective_importance × accumulated_relevance` -6. **Always-on rules** injected regardless of query +1. **Multi-query search** — pass `queries: ["topic1", "topic2"]` to search multiple angles at once. Each query is matched against keyword entities by both pg_trgm trigram similarity AND query-embedding-vs-keyword-embedding cosine similarity; results are merged with the geometric mean (configurable `missing_signal_penalty` when only one signal fires). +2. **Per-search-term reservation (L1 diversity quota)** — each query you pass gets a guaranteed share of the result slots filled from THAT query's own top-ranked entities. Bare-keyword queries (`"Petros"`) reliably surface specific facts even when paired with broader semantic angles. +3. **Per-keyword reservation (L2 diversity quota)** — each dominant matched keyword gets a halving slot allowance (50% / 25% / 12.5% ..., floor 1). Stops one popular hub keyword (e.g. `user-profile` tagging 100 facts) from monopolising top-N. +4. **Graph traversal** up to 3 hops via relations, relevance fading: `1.0 → 0.6 → 0.3`. +5. **Temporal decay** — memories fade over time, strengthen on access. +6. **Final rank** = `combined_score × effective_importance × accumulated_relevance`. The LLM-visible cap stays at the caller's `max_results` (default 30); the scoring pool internally considers up to 500 candidates per query so narrow keywords are never excluded before they're evaluated. +7. **Always-on rules** injected regardless of query. Single `query` (string) still works for backward compatibility. +**Query strategy** — prefer multiple short queries (a bare keyword + 1–2 broader phrases) over one long sentence. The keyword "Petros" matches the `Petros` keyword cleanly; the phrase "Petros person identity profile" matches the SAME keyword at a much lower score because pg_trgm dilutes against a longer query. + --- ## The BrainDB Agent @@ -166,19 +169,26 @@ curl -X POST http://localhost:8000/api/v1/agent/query \ -H "Content-Type: application/json" \ -d '{"query":"What do you know about the user role and recent projects?"}' -# {"answer": "The user is ...", "max_turns": 15} +# {"answer": "The user is ...", "max_turns": 20} ``` -The agent has 21 tools — every single BrainDB endpoint plus `delegate_to_subagent` (which spawns a fresh agent in its own context for focused deep work) and `submit_result` (which ends the loop). +The agent has 21 tools — every single BrainDB endpoint plus `delegate_to_subagent` (which spawns a fresh agent in its own context for focused deep work) and `final_answer` (which ends the loop with a validated typed payload). **LLM provider — pluggable via `.env`**: -`LLM_PROFILE` selects the backend. Profiles are defined in [braindb/config.py](braindb/config.py) (`_LLM_PROFILES`) — currently `deepinfra` (default, model `google/gemma-4-31B-it`) and `nim` (NVIDIA NIM, model `google/gemma-4-31b-it`). Each profile is a model-prefix + env-var pair; adding a new one is a dict entry. +`LLM_PROFILE` selects the backend. Profiles are defined in [braindb/config.py](braindb/config.py) (`_LLM_PROFILES`): + +- **`deepinfra` — recommended default.** Model `google/gemma-4-31B-it`. Fast (5–30s per agent call), cheap, validated end-to-end. +- `nim` — NVIDIA NIM, model `google/gemma-4-31b-it`. Free tier, occasionally flaky. +- `vllm_workstation` / `vllm_workstation_qwen` / `vllm_workstation_gemma` — local vLLM running on your own GPU (advanced / offline; needs the server reachable from the docker network, usually via SSH tunnel). + +Each profile is a model-prefix + env-var pair; adding a new one is a dict entry. ``` -LLM_PROFILE=deepinfra # or nim — default is deepinfra +LLM_PROFILE=deepinfra # or nim / vllm_workstation / vllm_workstation_qwen DEEPINFRA_API_KEY=... # required if profile=deepinfra (https://deepinfra.com/) NVIDIA_NIM_API_KEY=... # required if profile=nim (https://build.nvidia.com/) +VLLM_API_KEY=... # optional, only if local vLLM is started with --api-key AGENT_MODEL= # optional: override the profile's default model ``` @@ -270,6 +280,33 @@ curl -X POST http://localhost:8000/api/v1/entities/datasources/ingest \ It's idempotent by content hash — re-calling with the same bytes returns 200 (existing) instead of 201 (new). +## Autonomous Wiki Maintenance + +The second always-on sidecar, `wiki_scheduler`, makes the knowledge graph +self-organise into human-readable **wiki pages** with **zero manual steps** — +the same hands-off model as file ingestion. It loops in the background: +discovers entities not yet covered by a wiki, lets the in-house agent decide +where each belongs (attach to an existing wiki / create a new one / consolidate +duplicates / skip), and the writer agent researches and writes/maintains each +page, keeping it grounded and self-correcting. Started automatically by +`docker compose up -d` (like `watcher`); just watch it work: + +```bash +docker logs braindb_wiki_scheduler -f # the autonomous loop +docker logs braindb_api -f # the agent doing the work +``` + +You do **not** drive this by hand. The `POST /api/v1/wiki/{cron,maintain,write}` +endpoints exist for **debugging / inspection only** — normal operation is the +sidecar. (Optional read-only review: `docker compose exec api python -m +braindb.tools.export_wikis` writes a markdown snapshot of every wiki + +provenance to `data/wiki_review/`.) + +**Cost control:** like the `watcher`, this sidecar drives the LLM +automatically. To run without it, bring the stack up excluding the service or +scale it to 0 (`docker compose up -d --scale wiki_scheduler=0`), exactly as +you would for the watcher; or point `LLM_PROFILE` at a local model. + ## Stack - Python 3.12 + FastAPI + psycopg2 (sync, no ORM) @@ -277,4 +314,4 @@ It's idempotent by content hash — re-calling with the same bytes returns 200 ( - Alembic migrations - `sentence-transformers` + `Qwen/Qwen3-Embedding-0.6B` for keyword embeddings - `openai-agents[litellm]` + LiteLLM for the internal agent (DeepInfra / NIM / others pluggable via `LLM_PROFILE`) -- Docker Compose — `api` + `watcher` services, external PostgreSQL +- Docker Compose — `api` + `watcher` + `wiki_scheduler` services, external PostgreSQL diff --git a/alembic/versions/005_wiki_system.py b/alembic/versions/005_wiki_system.py new file mode 100644 index 0000000..5310695 --- /dev/null +++ b/alembic/versions/005_wiki_system.py @@ -0,0 +1,87 @@ +"""wiki system — wiki entity type, wikis_ext, wiki_job queue + +Revision ID: 005 +Revises: 004 +Create Date: 2026-05-16 + +Purely additive. Mirrors the 004 CHECK-rewrite pattern. No backfill; +existing rows are untouched. Adds the 'wiki' entity type, the wikis_ext +extension table, and the wiki_job queue table that drives the +cron / maintainer / writer pipeline. +""" +from alembic import op + +revision = "005" +down_revision = "004" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # 0. Add 'wiki' to the entity_type CHECK constraint (same DROP/ADD as 004) + op.execute("ALTER TABLE entities DROP CONSTRAINT IF EXISTS entities_entity_type_check") + op.execute(""" + ALTER TABLE entities ADD CONSTRAINT entities_entity_type_check + CHECK (entity_type IN ('thought','fact','source','datasource','rule','keyword','wiki')) + """) + + # 1. Wiki extension table — base entity columns (title/content/summary/ + # keywords/importance/notes/metadata) are reused; only wiki-specific + # structured fields live here. + op.execute(""" + CREATE TABLE wikis_ext ( + entity_id UUID PRIMARY KEY REFERENCES entities(id) ON DELETE CASCADE, + canonical_name VARCHAR(500) NOT NULL, + disambiguation TEXT, + language VARCHAR(10) DEFAULT 'en', + member_keyword_ids UUID[] DEFAULT '{}', + revision INT DEFAULT 1, + last_synthesised_at TIMESTAMPTZ, + retired_at TIMESTAMPTZ, + redirect_to UUID REFERENCES entities(id) ON DELETE SET NULL + ) + """) + op.execute("CREATE INDEX wikis_ext_canonical_idx ON wikis_ext (lower(canonical_name))") + op.execute("CREATE INDEX wikis_ext_member_kw_idx ON wikis_ext USING GIN (member_keyword_ids)") + + # 2. Structured maintainer/cron job queue + op.execute(""" + CREATE TABLE wiki_job ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + job_type VARCHAR(20) NOT NULL + CHECK (job_type IN ('triage','attach','create','consolidate')), + status VARCHAR(12) NOT NULL DEFAULT 'pending' + CHECK (status IN ('pending','assigned','done','rejected','failed')), + target_wiki_id UUID REFERENCES entities(id) ON DELETE CASCADE, + entity_ids UUID[] NOT NULL DEFAULT '{}', + dedupe_key TEXT NOT NULL, + rationale TEXT, + proposed_name VARCHAR(500), + batch_id UUID, + created_at TIMESTAMPTZ DEFAULT now(), + assigned_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ, + attempts INT DEFAULT 0, + last_error TEXT + ) + """) + # Idempotency: only one active job per logical work item. Once a job is + # done/rejected the key frees, so a genuinely new later situation can + # re-propose. Inserts use ON CONFLICT DO NOTHING (same as 004 backfill). + op.execute(""" + CREATE UNIQUE INDEX wiki_job_dedupe_active_idx + ON wiki_job(dedupe_key) WHERE status IN ('pending','assigned') + """) + op.execute("CREATE INDEX wiki_job_status_idx ON wiki_job(status)") + op.execute("CREATE INDEX wiki_job_target_idx ON wiki_job(target_wiki_id)") + + +def downgrade() -> None: + op.execute("DROP TABLE IF EXISTS wiki_job") + op.execute("DROP TABLE IF EXISTS wikis_ext") + # Restore the 004 entity_type CHECK constraint (without 'wiki') + op.execute("ALTER TABLE entities DROP CONSTRAINT IF EXISTS entities_entity_type_check") + op.execute(""" + ALTER TABLE entities ADD CONSTRAINT entities_entity_type_check + CHECK (entity_type IN ('thought','fact','source','datasource','rule','keyword')) + """) diff --git a/braindb/agent/agent.py b/braindb/agent/agent.py index 2120da1..34d7769 100644 --- a/braindb/agent/agent.py +++ b/braindb/agent/agent.py @@ -1,36 +1,71 @@ """ -BrainDB internal agent — builder and runner. +BrainDB internal agent — builder and runners. -Mirrors the pattern in fa-automation/tasks/linkedin_research/agent.py: -- create_braindb_agent() wires model + tools + instructions -- run_agent_query() is the async Runner.run() wrapper -- Singleton pattern so the agent is built once and reused +Convention (absolute): every agent run finishes via the `final_answer` +trick, and that tool's argument is ALWAYS a typed Pydantic model. The LLM +never emits loose / free-form output we then scrape. + +There is one agent per purpose, differing only by which typed +`submit_*` variant it carries (all named "final_answer" so prompts and +`StopAtTools(["final_answer"])` stay generic). The structured contract +lives on the **tool argument schema** (`@function_tool` + Pydantic), +which is what the user wanted: validated final answer, free middle +turns. We deliberately do NOT set `output_type` on the Agent — that flag +makes the SDK pass `response_format: json_schema` on every LLM call, +which steers weaker models to satisfy the schema on turn 1 and never +call any tool (the regression we are fixing). + +How we still recover the typed payload: each `submit_*` tool body parks +its already-validated `payload` into `braindb.agent.run_state.last_submit` +(a ContextVar). `run_typed` reads it back after `Runner.run` returns. +asyncio's per-Task context isolation makes nested/parallel runs safe. """ +import json import logging from pathlib import Path +from typing import TypeVar from agents import Agent, ModelSettings, Runner, StopAtTools, set_tracing_disabled from agents.extensions.models.litellm_model import LitellmModel +from litellm import BadRequestError, ContextWindowExceededError +from pydantic import BaseModel +from braindb.agent.hooks import CountdownHooks +from braindb.agent.run_state import install_slot, release_slot +from braindb.agent.schemas import ( + AgentAnswer, + MaintainerDecision, + SubagentResult, + WikiWriteResult, +) from braindb.agent.tools import ( create_relation, delegate_to_subagent, delete_entity, delete_relation, + delete_wiki_section, + edit_wiki_section, generate_embeddings, get_entity, get_stats, + handoff_to_successor, ingest_file, list_entities, quick_search, + read_wiki_outline, + read_wiki_section, recall_memory, save_fact, save_rule, save_source, save_thought, search_sql, - submit_result, + submit_answer, + submit_maintainer, + submit_subagent, + submit_wiki, update_entity, + validate_wiki, view_entity_relations, view_log, view_tree, @@ -41,75 +76,362 @@ SYSTEM_PROMPT = (Path(__file__).parent / "prompts" / "system_prompt.md").read_text(encoding="utf-8") -_agent: Agent | None = None +# Every tool except the final submit (that one is typed per purpose). +_BASE_TOOLS = [ + recall_memory, + quick_search, + save_fact, + save_thought, + save_source, + save_rule, + ingest_file, + get_entity, + list_entities, + update_entity, + delete_entity, + create_relation, + view_entity_relations, + delete_relation, + view_tree, + search_sql, + view_log, + get_stats, + generate_embeddings, + delegate_to_subagent, +] +T = TypeVar("T") -def create_braindb_agent() -> Agent: - """Build the BrainDB agent. Provider selected via settings.llm_profile.""" - model = LitellmModel( + +def _expected_shape_hint(expected_cls: type[BaseModel]) -> str: + """Render a literal JSON-call shape for the `final_answer` tool, derived + from the Pydantic model the LLM must submit. + + Weak/quantised models routinely emit the wrong WRAPPER on retry: either + they call `final_answer()` (missing the outer `payload` + key) or `final_answer({"payload": })` (missing required + keys inside). The generic "call final_answer NOW" correction did not + fix this on Gemma-31B (verified live: subagent retry kept emitting the + same shape errors). Giving the model a literal JSON template that + matches the @function_tool argument schema closes that gap — the LLM + sees the exact key names and the outer wrapping it has to produce. + + Example output for `SubagentResult`: + {"payload": {"result": ""}} + + For `MaintainerDecision` (skip action): + {"payload": {"action": "skip", "rationale": ""}} + + Only REQUIRED fields are filled with placeholders; optional/nullable + fields are omitted so the LLM doesn't fabricate values for them. The + helper handles enums (uses the first allowed value as the placeholder) + so the example is always actually-valid against the schema. + """ + schema = expected_cls.model_json_schema() + required = schema.get("required", []) + props = schema.get("properties", {}) + + def placeholder(field_name: str, field_schema: dict) -> str | int | list | dict: + # Literal/Enum: use the first allowed value so the example validates. + enum = field_schema.get("enum") + if enum: + return enum[0] + t = field_schema.get("type") + if t == "integer": + return 1 + if t == "number": + return 0.0 + if t == "boolean": + return False + if t == "array": + return [] + if t == "object": + return {} + # default: string + return f"<{field_name}>" + + example_payload = { + name: placeholder(name, props.get(name, {})) for name in required + } + return json.dumps({"payload": example_payload}) + + +def _model() -> LitellmModel: + return LitellmModel( model=settings.resolved_agent_model, api_key=settings.resolved_api_key, base_url=settings.resolved_base_url, ) - set_tracing_disabled(disabled=True) + +def _build( + name: str, + submit_tool, + extra_tools: tuple = (), + extra_stop_tools: tuple[str, ...] = (), +) -> Agent: + """Build an agent. NOTE: no `output_type` — see module docstring. The + structured contract lives on `submit_tool`'s argument schema, not on + the agent. + + `extra_tools` lets a specific agent (currently only the writer) carry + role-specific tools (the wiki section-edit tools + handoff) without + polluting `_BASE_TOOLS` shared by all agents. + + `extra_stop_tools` adds extra stop-tool names beyond `final_answer`. + The writer adds `handoff_to_successor` here so the run halts cleanly + when handoff is called instead of continuing wastefully. + """ + set_tracing_disabled(disabled=True) agent = Agent( - name="BrainDB Memory Agent", + name=name, instructions=SYSTEM_PROMPT, - model=model, + model=_model(), model_settings=ModelSettings(), - tools=[ - recall_memory, - quick_search, - save_fact, - save_thought, - save_source, - save_rule, - ingest_file, - get_entity, - list_entities, - update_entity, - delete_entity, - create_relation, - view_entity_relations, - delete_relation, - view_tree, - search_sql, - view_log, - get_stats, - generate_embeddings, - delegate_to_subagent, - submit_result, - ], - tool_use_behavior=StopAtTools(stop_at_tool_names=["submit_result"]), + tools=[*_BASE_TOOLS, *extra_tools, submit_tool], + tool_use_behavior=StopAtTools( + stop_at_tool_names=["final_answer", *extra_stop_tools], + ), + ) + logger.info( + "Agent built: %s (model=%s) — free middle turns, typed final_answer", + name, settings.resolved_agent_model, ) - logger.info("BrainDB agent created with model: %s", settings.resolved_agent_model) return agent +_cache: dict[str, Agent] = {} + + +def _cached( + key: str, + name: str, + submit_tool, + extra_tools: tuple = (), + extra_stop_tools: tuple[str, ...] = (), +) -> Agent: + a = _cache.get(key) + if a is None: + a = _build( + name, submit_tool, + extra_tools=extra_tools, + extra_stop_tools=extra_stop_tools, + ) + _cache[key] = a + return a + + +# Writer-only tools: section read/edit/delete + grammar validation + +# context-handoff. The writer rewrites whole wiki bodies today; section +# tools let it edit one at a time, and `handoff_to_successor` lets it +# bail to a fresh agent when context approaches the wall. See +# braindb/services/wiki_sections.py + plan +# `feat/wikis-and-maintainer-agent-read-write-tools`. +_WRITER_EXTRA_TOOLS = ( + read_wiki_outline, + read_wiki_section, + edit_wiki_section, + delete_wiki_section, + validate_wiki, + handoff_to_successor, +) +_WRITER_EXTRA_STOP_TOOLS = ("handoff_to_successor",) + + def get_agent() -> Agent: - """Get the singleton agent instance — built on first call.""" - global _agent - if _agent is None: - _agent = create_braindb_agent() - return _agent + """Default agent: general recall/save (public /agent/query).""" + return _cached("answer", "BrainDB Memory Agent", submit_answer) -async def run_agent_query(query: str, max_turns: int | None = None) -> dict: - """Run a query through the agent loop. Returns the final answer + metadata. +def get_maintainer_agent() -> Agent: + return _cached("maintainer", "BrainDB Wiki Maintainer", submit_maintainer) + + +def get_writer_agent() -> Agent: + return _cached( + "writer", "BrainDB Wiki Writer", submit_wiki, + extra_tools=_WRITER_EXTRA_TOOLS, + extra_stop_tools=_WRITER_EXTRA_STOP_TOOLS, + ) + + +def get_subagent() -> Agent: + return _cached("subagent", "BrainDB Subagent", submit_subagent) + + +def create_braindb_agent() -> Agent: + """Backward-compat alias — the default (general) agent.""" + return get_agent() + + +async def run_typed( + query: str, + agent: Agent, + expected_cls: type[T], + max_turns: int | None = None, + *, + token_budget: int = 0, + _bad_request_retried: bool = False, +) -> T: + """Run a typed agent and return the validated Pydantic instance it + submitted. The instance is guaranteed-valid because the SDK validates + the LLM's `final_answer` call args against `expected_cls` BEFORE the + tool body runs (via `@function_tool`'s strict JSON schema). + + Raises `RuntimeError` if the run ends without `final_answer` firing + (e.g. `max_turns` exhausted) — surfaces a real model failure instead + of silently returning bad data. Routers handle this like any other + agent error: log + release the job lease + 5xx. - When `settings.agent_verbose` is True, every tool call is logged to stdout - via the standard logger (visible in `docker logs braindb_api`). + `token_budget` (writer-only, opt-in): when > 0, enables the handoff + nudge in `CountdownHooks` — at the first LLM call where the cheap + token estimate of the conversation exceeds this budget, one + synthetic user message instructs the model to call + `handoff_to_successor`. The successor-respawn loop lives in the + caller (see `braindb/routers/wiki.py`). """ - agent = get_agent() turns = max_turns or settings.agent_max_turns - logger.info("Running agent query: %s", query[:200]) - result = await Runner.run( - starting_agent=agent, - input=query, + slot, token = install_slot() + # Layer-3 nudge: when the run is about to exhaust `max_turns`, the hook + # appends a synthetic "you have N turns left, finalise via final_answer" + # user message to the conversation. One nudge per run; disabled when + # `agent_countdown_threshold == 0`. See braindb/agent/hooks.py. + # When `token_budget > 0` (writer path) the same hook also watches + # estimated prompt tokens and injects ONE handoff nudge at the first + # call where the estimate crosses the budget. Independent fired-once + # flag from the turn nudge. + hooks = CountdownHooks( max_turns=turns, + threshold=settings.agent_countdown_threshold, + tool_name="final_answer", + token_budget=token_budget, + handoff_tool_name="handoff_to_successor", ) - return { - "answer": str(result.final_output), - "max_turns": turns, - } + try: + logger.info("Running typed query (%s): %s", agent.name, query[:160]) + result = await Runner.run( + starting_agent=agent, input=query, max_turns=turns, hooks=hooks, + ) + payload = slot.value + if isinstance(payload, expected_cls): + return payload + + # The first attempt ended without `final_answer` firing. Most + # commonly the model emitted plain prose (a "fast finisher" / + # forgetter) — strict mode would raise here. But before giving + # up, Layer 4 gives the model exactly one chance to fix it: + # append a user-role correction message to the conversation it + # already produced (`result.to_input_list()`) and re-invoke + # `Runner.run` with a small budget. The correction is unambiguous + # — "you ended without `final_answer`, call it now". No parsing + # of the prose, no fallback that pretends success; we use the + # SDK's own conversation mechanism to tell the model what it did + # wrong, then either it complies on the retry (HTTP 200) or we + # raise (still strict). + if settings.agent_retry_on_missing_final: + logger.info( + "%s ended without final_answer; retrying once with correction", + agent.name, + ) + # Build a literal JSON-shape hint from `expected_cls` so the + # LLM gets an unambiguous template — not just "call it now", + # but "call it like THIS". Verified live: Gemma subagents + # retry without this hint by emitting payload-as-string or + # missing-required-key variants that fail the @function_tool + # validator and trigger the same error in a loop. + shape_hint = _expected_shape_hint(expected_cls) + correction = { + "role": "user", + "content": ( + "Your previous response ended WITHOUT a successful " + "`final_answer` call (or `final_answer` was called " + "with the wrong JSON shape and rejected by the tool " + "validator). The work you did is preserved, but the " + "run is INVALID until you finalise.\n\n" + "Call `final_answer` NOW. The tool expects EXACTLY " + "one argument named `payload`, whose value is a JSON " + "object with the required keys. The literal shape " + f"you MUST send is:\n\n {shape_hint}\n\n" + "Replace each with your real value. " + "Do NOT omit the outer `payload` key. Do NOT wrap " + "the payload as a string. Issue ONLY the tool call, " + "no prose, no further research." + ), + } + retry_input = result.to_input_list() + [correction] + retry_hooks = CountdownHooks( + max_turns=settings.agent_retry_max_turns, + threshold=settings.agent_countdown_threshold, + tool_name="final_answer", + ) + await Runner.run( + starting_agent=agent, + input=retry_input, + max_turns=settings.agent_retry_max_turns, + hooks=retry_hooks, + ) + payload = slot.value + if isinstance(payload, expected_cls): + logger.info( + "%s recovered via final_answer-retry (correction worked)", + agent.name, + ) + return payload + + # Retry also failed: model truly refuses the typed-final + # contract even when told explicitly what to do. That's a + # genuine model-discipline failure — raise loudly. + raise RuntimeError( + f"{agent.name} did not call final_answer even after a " + f"correction retry — model refuses the typed-final " + f"contract. Last final_output: " + f"{str(getattr(result, 'final_output', ''))[:200]}" + ) + + # Retry disabled (opt-out via settings): preserve the original + # strict-raise behaviour. + raise RuntimeError( + f"{agent.name} did not call final_answer with a " + f"{expected_cls.__name__} (got {type(payload).__name__}). " + f"The run terminated without the typed final tool firing — " + f"the model likely ended with plain prose." + ) + except ContextWindowExceededError: + # The conversation is already over the model's window. A retry + # without input truncation would just hit the same wall, so we + # re-raise and let the router fail the job cleanly. Real fix is + # upstream: keep prompts/tool-results small enough that the + # handoff threshold catches us first. See routers/wiki.py's + # `_body_block_or_stub` for the prompt-side mitigation. + raise + except BadRequestError as e: + # Quantised models (Qwen AWQ-INT4) occasionally emit malformed + # JSON in tool-call args; the OpenAI client raises BadRequestError + # before the tool body runs. One fresh attempt usually recovers. + # Bounded to depth 1 via the `_bad_request_retried` flag — + # recursion uses run_typed itself rather than duplicating the + # setup. The current slot is released by `finally`; the recursive + # call installs its own. + if _bad_request_retried: + raise + logger.warning( + "%s: BadRequestError on first attempt (%s); " + "retrying once with a fresh run", agent.name, str(e)[:160], + ) + return await run_typed( + query, agent, expected_cls, + max_turns=max_turns, token_budget=token_budget, + _bad_request_retried=True, + ) + finally: + release_slot(token) + + +async def run_agent_query(query: str, max_turns: int | None = None) -> dict: + """General recall/save path (public /agent/query, and the ingest watcher + over HTTP). The model finishes via `final_answer(payload: AgentAnswer)`; + the response shape stays `{"answer","max_turns"}` for backward + compatibility.""" + turns = max_turns or settings.agent_max_turns + payload: AgentAnswer = await run_typed(query, get_agent(), AgentAnswer, max_turns=turns) + return {"answer": payload.answer, "max_turns": turns} diff --git a/braindb/agent/hooks.py b/braindb/agent/hooks.py new file mode 100644 index 0000000..8d1fa9b --- /dev/null +++ b/braindb/agent/hooks.py @@ -0,0 +1,246 @@ +"""Runtime nudge: tell the LLM to finalise when it's about to run out of turns. + +WHY this exists +--------------- +The strict typed-final contract (`final_answer` tool with a Pydantic argument +schema, no `output_type` on the Agent — see `braindb/agent/agent.py`) raises a +`RuntimeError` if the model ends a run without calling `final_answer`. Weak +or quantised models sometimes over-explore (chaining `recall_memory` / +`delegate_to_subagent` calls beyond what's necessary) and reach +`max_turns` without ever submitting. The strict path correctly catches this +as a failure, but we'd rather give the model a fighting chance: shortly +before `max_turns` is exhausted, inject a chat message reminding it to +finalise. + +HOW the nudge gets into the conversation +----------------------------------------- +The openai-agents SDK's `RunHooks.on_llm_start` callback (see +`agents/lifecycle.py`) receives the mutable `input_items` list that's about +to be sent to the LLM. Appending one item to that list adds a synthetic +user message visible to the model on its NEXT turn. That's the same +mechanism the SDK uses internally for any added context. We exploit it +exactly once per run (idempotent), at the configured threshold. + +Knobs (see `braindb/config.py`) +- `agent_countdown_threshold` (default 5): how many turns before + `max_turns` we start nudging. Set to 0 to disable the nudge entirely. + +Design constraints +- One nudge per run (no spam). +- Defensive: any internal error in the hook is caught and logged, never + re-raised — a future SDK shape change must not bring down agent runs. +- Pure on-LLM-start counting — no SDK-private state inspection. +""" +from __future__ import annotations + +import logging +from typing import Any + +from agents.lifecycle import RunHooks + +logger = logging.getLogger(__name__) + + +def _estimate_tokens(input_items: list) -> int: + """Cheap (no-tokenizer) prompt-token estimate: sum the text-content + character counts and divide by 4. Defensive across the shapes the + SDK puts into `input_items`: + - `{"role": str, "content": str}` (LiteLLM dict form) + - `{"role": str, "content": [{"type":"text","text":str}, ...]}` + (some providers send a list of parts) + - SDK item objects with a `.content` attribute + Unknown shapes contribute 0; the estimate is a lower bound, which + is the safe side for "is context filling up" decisions (we'd rather + fire the handoff nudge slightly late than slightly never).""" + total_chars = 0 + for item in input_items: + content: object + if isinstance(item, dict): + content = item.get("content", "") + else: + content = getattr(item, "content", "") + if isinstance(content, str): + total_chars += len(content) + elif isinstance(content, list): + for part in content: + if isinstance(part, dict): + text = part.get("text") or part.get("content") or "" + if isinstance(text, str): + total_chars += len(text) + elif isinstance(part, str): + total_chars += len(part) + return total_chars // 4 + + +class CountdownHooks(RunHooks): + """Mutates `input_items` to inject up to TWO independent nudges: + + 1. Turn-budget nudge ("you have N turns left, finalise") — fires when + the agent is close to exhausting `max_turns`. Original behaviour; + see module docstring. + + 2. Token-budget nudge ("context is filling up, call handoff_to_successor") + — fires ONLY when `token_budget > 0` AND the cheap token estimate + of `input_items` (sum-of-content-chars / 4) exceeds the budget. + Writer-only: callers that don't set `token_budget` get the + original turn-only behaviour. The two nudges have independent + fired-once flags so one cannot suppress the other. + + Lifecycle (per run): + - constructed once with knobs (turn-related + optional token-related). + - `on_llm_start` fires before each LLM call. + - increments `_turns`; if `_turns >= max_turns - threshold` AND + `_fired_turns` is False, appends the turn nudge. + - if `token_budget > 0` AND + `estimated_tokens(input_items) > token_budget` AND + `_fired_tokens` is False, appends the handoff nudge. + - each nudge fires at most once per run. + + Disabled paths: + - `threshold <= 0` disables the turn nudge (existing safety hatch). + - `token_budget <= 0` disables the handoff nudge (default; non-writer + callers don't pass this). + """ + + def __init__( + self, + max_turns: int, + threshold: int, + tool_name: str = "final_answer", + *, + token_budget: int = 0, + handoff_tool_name: str = "handoff_to_successor", + ) -> None: + self.max_turns = max_turns + self.threshold = max(0, int(threshold)) + self.tool_name = tool_name + self.token_budget = max(0, int(token_budget)) + self.handoff_tool_name = handoff_tool_name + self._turns: int = 0 + self._fired_turns: bool = False + self._fired_tokens: bool = False + + # Backwards-compatibility: existing tests reference `._fired` on + # instances built without token_budget. Map it to the turn-fired + # flag so they keep observing the same semantic. + @property + def _fired(self) -> bool: # noqa: D401 + return self._fired_turns + + @_fired.setter + def _fired(self, v: bool) -> None: + self._fired_turns = v + + # NOTE: `on_llm_start` is the canonical hook for injecting context + # before the next LLM call (the SDK passes `input_items` mutably). + # We don't override `on_tool_start` because we want to count + # LLM-call turns, not tool calls — those can be multiple per turn. + async def on_llm_start( + self, + context: Any, + agent: Any, + system_prompt: str | None, + input_items: list, + ) -> None: + try: + self._turns += 1 + self._maybe_inject(input_items) + except Exception as e: # noqa: BLE001 — defensive: never kill the run + logger.warning( + "CountdownHooks.on_llm_start swallowed an internal error " + "(turns=%d, fired=%s): %r", self._turns, self._fired, e, + ) + + def _maybe_inject(self, input_items: list) -> None: + """Pure logic: decide whether to append a nudge now. Two + independent checks (turn-budget + token-budget); each fires at + most once per run. Separated from on_llm_start so tests can stub + it to verify the wrapper's exception-swallowing behaviour.""" + # Turn-budget nudge (original Layer 3). + if self.threshold > 0 and not self._fired_turns: + remaining = self.max_turns - self._turns + if remaining <= self.threshold: + self._fired_turns = True + nudge = self._format_nudge(remaining) + input_items.append({"role": "user", "content": nudge}) + logger.info( + "CountdownHooks injected TURN nudge at turn %d/%d " + "(remaining=%d): %s", + self._turns, self.max_turns, remaining, nudge[:120], + ) + + # Token-budget nudge (handoff path). + if self.token_budget > 0 and not self._fired_tokens: + est = _estimate_tokens(input_items) + if est > self.token_budget: + self._fired_tokens = True + handoff = self._format_handoff_nudge(est) + input_items.append({"role": "user", "content": handoff}) + logger.info( + "CountdownHooks injected HANDOFF nudge (est_tokens=%d, " + "budget=%d): %s", + est, self.token_budget, handoff[:120], + ) + + def _format_handoff_nudge(self, est_tokens: int) -> str: + """Text the model sees when token usage crosses the budget. Asks + it to call the handoff tool with a structured brief; gives the + agent an escape hatch (call final_answer directly) for small + remaining work.""" + return ( + f"Your context is filling up (≈{est_tokens} estimated tokens; " + f"budget {self.token_budget}). To avoid running out, call " + f"`{self.handoff_tool_name}` now with a structured brief:\n" + f"- progress_summary: tools you've called, key findings, and " + f"any active revision tokens (the wiki you've been editing).\n" + f"- remaining_work: the concrete next tool call(s) the " + f"successor must make — name wikis, section names, revisions.\n" + f"A fresh agent with the same prompt and tools will continue " + f"from your brief. If you can still finish in 1-2 turns you " + f"may instead call `{self.tool_name}` directly, but err on " + f"the side of handoff when context is this tight." + ) + + def _format_nudge(self, remaining: int) -> str: + """The text the model sees. Tone is chosen by `self.max_turns`: + + - SOFT (max_turns > 5): "start wrapping up, you have N left". + Used when the budget is generous (the new default of 20 with + threshold 8 fires the nudge at turn 12, with 8 turns still to + spend). Deep-research models like Qwen do better when given + a "begin concluding" signal rather than a hard stop — they + can do one or two focused gap-filling calls before + `final_answer` instead of slamming tools shut mid-thread. + + - HARD (max_turns ≤ 5): "call `final_answer` NOW". Used when + the budget is tight — most notably the Layer 4 retry path + (`max_turns=3`), where the retry is explicitly a "you forgot + to finalise, please call the tool now" correction. The + model gets the unambiguous instruction without ambiguity + about wrapping up vs investigating further. + + Why pick the tone from `max_turns` rather than an explicit + constructor flag: the retry call site already passes its own + `max_turns=settings.agent_retry_max_turns` (3) and the main + run passes the general `max_turns` (20). The two contexts + differ exactly along the budget axis, so we get the right + tone with no new constructor surface and no caller changes. + """ + # Clamp to non-negative for readability; if remaining went past 0 + # we still want a coherent message even though the SDK would + # raise MaxTurnsExceeded shortly. + remaining = max(0, remaining) + plural = "s" if remaining != 1 else "" + if self.max_turns <= 5: + return ( + f"You have {remaining} tool call{plural} left. " + f"Call `{self.tool_name}` with your answer now. " + f"Do not start new research." + ) + return ( + f"Heads up: you have {remaining} tool call{plural} left " + f"in this run. Start wrapping up — synthesise what you " + f"have already gathered and prepare to call " + f"`{self.tool_name}`. Focused gap-filling is fine; avoid " + f"opening brand-new lines of investigation." + ) diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index a41f406..335bbd0 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -2,7 +2,7 @@ You are the BrainDB Memory Agent — the persistent memory layer for an LLM user Your job: handle memory operations (recall, save, relate, explore, maintain) on behalf of an external caller who talks to you in natural language. The caller (typically Claude Code or another agent) shouldn't need to know any internal details — you decide what to do and use your tools to do it. -Always end by calling `submit_result(answer)` with a concise summary of what you did or what you found. That is how the loop stops. +CRITICAL — every assistant message MUST be a tool call; never plain prose. The run is INVALID until you call `final_answer`, and your **final** action MUST be `final_answer` with its typed fields filled (for a general query that is just `answer`: a concise summary of what you did or found). A prose-only response causes the run to fail and your work is discarded — your answer only "lands" via `final_answer`. --- @@ -32,7 +32,7 @@ Always end by calling `submit_result(answer)` with a concise summary of what you **Explore:** - `view_tree(entity_id, max_depth)` — entity + all its connections -- `search_sql(query)` — read-only SQL (SELECT/WITH only) for complex queries +- `search_sql(query)` — read-only SQL. **Exception tool only** (see TOOL PRIORITY): for a specific structured/aggregate question (counts, GROUP BY, log joins) the retrieval tools genuinely cannot express. NEVER for recall, discovery, or understanding. - `view_log(operation, entity_id, limit)` — recent activity log - `get_stats()` — entity counts, relation counts - `generate_embeddings()` — batch-generate embeddings for keyword entities missing them @@ -41,10 +41,50 @@ Always end by calling `submit_result(answer)` with a concise summary of what you - `delegate_to_subagent(task)` — spawn a fresh subagent that runs in its own context and returns only a summary. Use for focused deep work you don't want cluttering your own context. **Done:** -- `submit_result(answer)` — **MUST call exactly once** when finished. Provide a clear summary of what you did or found. +- `final_answer` — **MUST call exactly once** when finished. Its argument is typed; fill the fields the tool's schema exposes (for a general query: `answer` = a clear summary of what you did or found). --- +## TOOL PRIORITY — the sophisticated tools first, always + +BrainDB's value is the graph + embeddings + ranking. Use that power; do not +fall back to flat SQL. + +1. **`recall_memory`** — the default for ALL recall, discovery, and + understanding: multi-query fuzzy + full-text + **keyword-embedding** + + graph traversal + decay + ranking. This is almost always the right first + call. +2. **`delegate_to_subagent`** — for any multi-step investigation or + disambiguation ("is this the same person/thing?", "find and resolve X"). + A fresh agent with the full toolset; returns a summary. Prefer this over + doing a long crawl yourself. +3. `view_tree` / `view_entity_relations` / `get_entity` / `list_entities` — + targeted structure lookups. +4. **`search_sql` — exception only.** A blunt SELECT has no embeddings, no + graph, no ranking — it throws away everything BrainDB is good at. Use it + *only* for a specific structured/aggregate question the tools above cannot + express (counts, GROUP BY, activity-log joins). Never for recall, + discovery, similarity, or understanding. + +If you reach for `search_sql` to "find" or "understand" something, stop — +that's a `recall_memory` or `delegate_to_subagent` job. + +## READING CONTENT — previews vs the full body + +Multi-item results (`recall_memory`, `quick_search`, `list_entities`, +`search_sql`) return **short previews** (~1K/item). A clipped item ends with +`--truncated (N more chars)-- full body: get_entity("")`. That is by +design — research from previews, then open only the few you actually need. + +- To read ONE thing fully: `get_entity(id)`. +- If that body is **large**, do NOT pull it whole into your context. Page it: + `get_entity(id, offset=0, limit=8000)` → use the returned + `content_meta.next_offset` to fetch the next slice, repeating until it is + `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. +- Never try to defeat previews via `search_sql` to dump whole bodies. + ## DELEGATION — use `delegate_to_subagent` for focused deep work When a task would require many tool calls (deep search, duplicate detection, bulk relation work, graph exploration) and you don't need to see the intermediate results in your own context, delegate it to a subagent. The subagent runs in its own conversation context, uses the same tools you have, and returns only a final summary. @@ -53,7 +93,7 @@ When a task would require many tool calls (deep search, duplicate detection, bul - The specific goal - What it should return (IDs, summaries, counts) - Any constraints (limits, filters) -- An explicit instruction to call `submit_result` at the end +- An explicit instruction to call `final_answer` at the end ### When to delegate - "Find all near-duplicate facts in memory, return top 10 pairs with IDs." @@ -137,7 +177,7 @@ Relation types: `supports`, `contradicts`, `elaborates`, `refers_to`, `derived_f You: 1. `recall_memory(["user-profile machine-learning expertise", "ML projects production deployment"])` 2. Read the returned items. -3. `submit_result("The user is Dimitris, ML/AI engineer at CityFalcon. Strong expertise in Python, LLMs (prompt engineering, fine-tuning, RAG), classical ML, and deep learning. Built the IR Extract Agentic Service where 3 previous people failed. Also reduced NLU GPU inference to one-third of prior levels.")` +3. `final_answer("The user is Dimitris, ML/AI engineer at CityFalcon. Strong expertise in Python, LLMs (prompt engineering, fine-tuning, RAG), classical ML, and deep learning. Built the IR Extract Agentic Service where 3 previous people failed. Also reduced NLU GPU inference to one-third of prior levels.")` ### Example 2 — Save @@ -148,22 +188,26 @@ You: 2. `save_fact(content="User is testing the new BrainDB agent with gemma-4-31b-it via NVIDIA NIM.", keywords=["braindb", "agent", "gemma", "NVIDIA-NIM", "testing"], importance=0.7)` 3. `list_entities(keyword="braindb", limit=10)` — find existing BrainDB entities to connect to 4. `create_relation(from_entity_id=, to_entity_id=, relation_type="elaborates", description="Agent is a new BrainDB component")` -5. `submit_result("Saved new fact about testing the BrainDB agent with gemma-4-31b-it. Linked to existing BrainDB project entities.")` +5. `final_answer("Saved new fact about testing the BrainDB agent with gemma-4-31b-it. Linked to existing BrainDB project entities.")` -### Example 3 — Explore +### Example 3 — Explore (delegate; don't reach for SQL) **Caller:** "Any duplicate entities I should clean up?" You: -1. `search_sql("SELECT a.id, b.id, a.content, b.content FROM entities a JOIN entities b ON a.id < b.id AND a.entity_type = b.entity_type WHERE similarity(a.content, b.content) > 0.6 LIMIT 20")` -2. `submit_result("Found 3 pairs of likely duplicates: ...")` +1. `delegate_to_subagent("Find likely near-duplicate entities in BrainDB. Use recall_memory across the main topics to pull clusters, compare entities within each cluster semantically, and return the top ~10 candidate duplicate pairs as (id, id, one-line why). Call final_answer with that list.")` +2. `final_answer("Found N likely duplicate pairs: ...")` + +(Only if the caller asked for a precise *count/aggregate* — e.g. "how many +facts per source?" — is `search_sql` the right tool. Finding/understanding is +`recall_memory` + a subagent.) --- ## RULES -- **Always call `submit_result` exactly once** at the end. This is how the loop stops. Don't forget. +- **`final_answer` is mandatory.** Every assistant message must be a tool call; the FINAL one must be `final_answer`. Ending with prose (a regular text response) makes the run fail — the harness reads your typed payload from `final_answer`, nothing else. If you have an answer, the only way to deliver it is to call `final_answer` with it in the typed field. - Be efficient: aim for 3-6 tool calls for most queries. Don't loop endlessly. -- Never paste raw JSON into `submit_result`. Format a human-readable summary. -- Errors from tools come back as strings starting with `ERROR:`. Decide whether to retry, try a different approach, or report the error in `submit_result`. +- Fill `final_answer`'s typed fields — don't hand-write JSON or delimiters; the tool's schema is the contract. For a general query, `answer` is a human-readable summary. +- Errors from tools come back as strings starting with `ERROR:`. Decide whether to retry, try a different approach, or report the error in `final_answer`. - You're talking to another agent/tool, not a human directly. Be concise and structured, but natural. diff --git a/braindb/agent/prompts/wiki_maintainer_prompt.md b/braindb/agent/prompts/wiki_maintainer_prompt.md new file mode 100644 index 0000000..746cd37 --- /dev/null +++ b/braindb/agent/prompts/wiki_maintainer_prompt.md @@ -0,0 +1,135 @@ +You are the **BrainDB Wiki Maintainer**, working on exactly ONE case. + +A "wiki" is a synthesised, human-readable page (entity_type = `wiki`) about ONE +real-world subject, built from the fact/thought/source entities that are +genuinely about that subject. + +Your case (THE SEED) and the numbered WIKIS catalog are at the **END** of +this prompt. Read the static rules here first, then act on the data there. +The single seed is rarely enough to decide correctly — you MUST investigate +the surrounding reality before deciding. + +## Research FIRST with the powerful tools (this is mandatory) + +Recall/list results are **short previews** (~1K/item) ending with +`--truncated … get_entity("")` when clipped — that is enough to triage. +Open a full body only via `get_entity(id)`; if it is large, page it +(`get_entity(id, offset, limit)` → follow `content_meta.next_offset`) or hand +slices to a subagent. Never pull whole datasources/wikis into your context. + +Tool priority — use them in this order, do not skip to the bottom: + +1. **`recall_memory`** — the sophisticated retrieval (embeddings + graph + + ranking). This is MANDATORY and is the heart of the decision. Run 2-4 + targeted queries around the seed's subject — and you MUST include its + obvious **name variants/aliases**: given/family-name swaps and orderings, + spelling variants, and the BROAD subject behind a NARROW fact (a fact + about "X's LinkedIn" / "X's divestment from Y" is about **X**, not a new + subject). The single required output of this step is: **does this subject + already have a wiki in the WIKIS catalog at the end (under any variant)?** + You may not choose `create` until you have actually looked and that + answer is "no". +2. **`delegate_to_subagent`** — when identity/scope is non-trivial (e.g. "are + these two 'Dimitris' facts the same person?"), delegate a focused + investigation: tell the subagent exactly what to resolve and to return a + crisp finding. Use this instead of guessing. +3. `view_tree` / `view_entity_relations` — inspect connections and any + `not_duplicate` / `duplicate_of` markers between wikis. +4. `search_sql` — **exception only**, for a specific structured/aggregate + lookup the above genuinely cannot express. Never for discovery or + understanding. + +## Identity & scope discipline (this is where it goes wrong) + +- **Distinct real entities are distinct.** People who merely share a first + name, or who co-occur in one fact, are NOT the same subject. If a fact says + "X's uncle is a marine engineer", *marine engineer* is the **uncle's** + attribute, not X's. Do not fuse separate people/things into one subject. +- **Exclusion over wrong inclusion.** A fact that uses only a shared first + name and is not uniquely tied to one person is AMBIGUOUS — do not let it + drive an `attach`/`create` toward a same-first-name subject. When several + facts could be different people sharing a name, prefer `ambiguous` (or + delegate a quick resolution) over a confident wrong suggestion. The writer + applies the same discipline; never hand it a conflated grouping. +- **Never invent or "correct" an identity.** Only propose a `proposed_name` + that appears explicitly in the evidence. If the evidence only says + "Dimitris" and you cannot tell *which* Dimitris from the data, that is + **ambiguous** — do not coin a surname or pick one. +- **Scope must match the evidence.** Do not propose a broad concept (e.g. + "Artificial Intelligence") when the evidence is one narrow source — propose + the narrower subject the evidence actually supports, or skip. +- **Keyword-token entities are not evidence.** An `entity_type='keyword'` + whose content is an opaque token/slug (e.g. `_pytest_82a2e09b`, + `artificial-intelligence`) is infrastructure, not a source and not a + concept. If the seed is only that, with no real fact/thought/source behind + it → **skip**. + +## Referencing existing wikis — BY NUMBER ONLY + +Every existing wiki is listed in the numbered **WIKIS catalog** at the end of +this prompt. To `attach` or `consolidate`, you reference wikis **solely by +their catalog number** — never by id, name, or a guessed value. You may only +attach/consolidate to wikis that appear in that numbered catalog. You never +see or emit a uuid; the harness maps your number back to the real wiki. If +the subject is not in the catalog, you cannot attach/consolidate to it. + +## Decide ONE action for THIS seed — STRICT PRECEDENCE, in this order + +Evaluate top to bottom and take the FIRST that applies. `create` is the last +resort, not the default. This ordering is how the wiki set heals over time — +honour it. + +1. **skip** — the seed is infrastructural / a keyword-token / too trivial to + deserve a page (see "keyword-token entities are not evidence"). +2. **ambiguous** — recall cannot disambiguate which real subject this is + (e.g. a bare shared first name). Refusing to mint a confident page is the + correct, honest outcome; say what is unresolved in `rationale`. +3. **consolidate** — the catalog contains ≥2 wikis that are the SAME real + subject (incl. name variants / over-narrow fragment pages of one + subject). Put their catalog **numbers** in `consolidate_nos` (≥2). Do NOT + re-propose a pair already linked by `not_duplicate` / `duplicate_of`. + This is the primary heal action — if you see duplicates in the catalog + while researching, you MUST propose this. +4. **attach** — a catalog wiki already covers this subject (under any name + variant), or the seed is a narrow fact about an already-wikied broad + subject. Put that wiki's catalog **number** in `target_wiki_no`. A narrow + fact about an existing subject is ALWAYS an attach, never a new page. +5. **create** — ONLY if steps 1-4 do not apply: recall + the catalog + genuinely show no existing wiki for this subject under any variant, AND + the evidence supports a clear, explicitly-named subject and scope. Give + the canonical name (must appear in the evidence). + +You only produce the suggestion. You do NOT create wikis/relations here — the +writer stage does, and it will research further. + +## Output — STRICT + +Finish by calling `final_answer` exactly once. Its argument is a typed +object — the tool's schema defines and validates the fields; you just fill +them (no raw JSON text, no prose): + +- `action` — one of `attach`, `create`, `consolidate`, `skip`, `ambiguous`. +- `target_wiki_no` — required for `attach`: the catalog NUMBER of the wiki + (an integer from the WIKIS list at the end); null otherwise. +- `proposed_name` — required for `create` (a canonical name that appears in + the evidence); null otherwise. +- `consolidate_nos` — required for `consolidate`: a list of ≥2 catalog + NUMBERS (integers from the WIKIS list); empty otherwise. +- `rationale` — 1-3 sentences: name the catalog wiki(s) you matched this + subject to (or state the catalog has none), and why attach/consolidate was + or was not chosen. This makes the decision auditable. + +--- + +## THE SEED (your one case) + +- entity_id: `{entity_id}` +- entity_type: `{entity_type}` +- keywords: {keywords} +- summary: {summary} +- content: +{content} + +## WIKIS catalog (existing wikis — reference these BY NUMBER) + +{wiki_catalog} diff --git a/braindb/agent/prompts/wiki_writer_prompt.md b/braindb/agent/prompts/wiki_writer_prompt.md new file mode 100644 index 0000000..31ea663 --- /dev/null +++ b/braindb/agent/prompts/wiki_writer_prompt.md @@ -0,0 +1,274 @@ +You are the **BrainDB Wiki Writer**. You write/maintain ONE wiki page so it +reflects **reality**, grounded in evidence. You own the content entirely — +nothing downstream rewrites or gates it. Get it right. + +A wiki is an encyclopedic, third-person page about ONE real subject, built +ONLY from entities that are genuinely about that subject. Every non-trivial +claim carries an inline reference `[[ref:ENTITY_UUID]]` (optionally +`[[ref:ENTITY_UUID|display text]]`) to the entity it came from. + +## This job + +- mode: **%%MODE%%** + - create = write a fresh page for the subject + - attach = the page exists; integrate the new members AND revise anything + now wrong (see "You MUST revise" below) + - consolidate = merge the numbered duplicate wikis below into one + survivor; you pick the survivor by its NUMBER (`canonical_no`) +- canonical_name (proposed): %%CANONICAL%% +- wiki_id: %%WIKI_ID%% + +### Seed member entities for this job +%%MEMBERS%% + +### Current wiki body (attach mode; empty otherwise) +%%CURRENT_BODY%% + +### Duplicate wikis to consolidate (consolidate mode only — NUMBERED; pick the survivor's number as `canonical_no`) +%%DUPLICATES%% + +## Mandatory order of work (do NOT skip or reorder) + +The seed/members are a starting point, not the truth. Treat the existing +page **conservatively**: its prose alone is not evidence (don't anchor on +uncited sentences or claims a new member contradicts), but +`[[ref:UUID]]`-cited claims are backed by the prior revision's verified +facts. + +**Attach mode — read the existing body before recalling.** Trust the +prior body's claims when they're already cited and uncontested, and +focus your `recall_memory` budget on: +- new members (the `MEMBERS` block) and how they slot in, +- claims that look inconsistent between the body and a new member, +- gaps the new members open up but the body doesn't yet cover. + +Be thorough where evidence is fresh or conflicting; be efficient +where the body already has it right — **but every assigned MEMBER +still needs to be cited at least once in the new body even if its +content is already covered**, because the citation is what records +the `summarises` relation (see "Citation is mechanical" below). + +Work in this exact order: + +**Step 1 — Gather raw facts.** Use `recall_memory` (sophisticated +embeddings+graph+ranking retrieval — the default for everything; `search_sql` +is an exception only for a structured aggregate it cannot express) with 2-4 +queries around the subject to collect the candidate `fact`/`thought`/`source` +entities (ids + contents). Ignore `keyword`-token entities (opaque slugs like +`_x_1a2b`) — never sources. Recall returns **previews** (~1K/item); facts are +short so previews are usually whole. To read a long datasource/source/wiki +fully, `get_entity(id)`; if it is large, **page it** +(`get_entity(id, offset, limit)` → follow `content_meta.next_offset`) and/or +hand each slice to `delegate_to_subagent` to distil — never load a big +document into your own context. + +**Step 2 — Independent entity resolution (MANDATORY `delegate_to_subagent`).** +Whenever ≥2 gathered facts could refer to different real people/things sharing +a name (almost always for people), you MUST delegate resolution BEFORE +writing. Send the subagent **only the raw `id: content` lines** — NOT the +page, NOT the canonical name, NOT the current Summary/Disambiguation, NOT any +expected answer. Use this task **verbatim** (fill only the FACTS): + +> "Below are memory entities (id: content). Perform IDENTITY RESOLUTION with +> NO assumptions. (1) Enumerate the DISTINCT real people/things these facts +> describe — there may be several who share a first name. Give each a +> short descriptor grounded in a quoted phrase. (2) For EACH distinct entity, +> list the fact ids about it, each with the quoted phrase that proves it. +> (3) Apply DISQUALIFIERS: if an entity is characterised one way (e.g. a +> youth who *aspires* to a trade), facts describing an unrelated established +> profile are NOT that entity unless a fact explicitly ties them by full +> name or a unique attribute. (4) Any fact that uses only a shared first +> name and cannot be uniquely assigned goes in an AMBIGUOUS bucket — do not +> force it onto anyone. Return: each entity → [fact id + evidence], plus the +> AMBIGUOUS bucket. Finish by calling final_answer once; put the full +> mapping (as readable text) in its `result` field. FACTS:\n" + +**Step 3 — Write for ONE resolved entity only.** Identify which resolved +entity is the subject of THIS page (matches the proposed canonical_name / +seed). Write the page using **only that entity's assigned facts**. Facts in +the AMBIGUOUS bucket or assigned to a *different* entity are EXCLUDED — do not +cite them, do not mention them as the subject's. (Additive reconcile creates +relations only for what you cite, so exclusion leaves nothing wrong behind.) + +## Identity discipline & circuit-breaker (this is where pages went wrong) + +- **Exclusion over wrong inclusion.** A fact that only says a shared first + name and is not uniquely tied to the subject is AMBIGUOUS → leave it OUT. + Never sweep same-first-name professional facts onto a person the evidence + describes very differently. +- **No third-party attribute transfer.** "X's uncle is a marine engineer" + makes *the uncle* a marine engineer, not X. +- **Correctness over richness.** A short, certain page is better than a rich, + wrong one. Never pad from world knowledge or from ambiguous facts. +- **Circuit-breaker (the STOP).** If resolution cannot confidently assign the + core identity/professional facts to THIS subject, do NOT elaborate. Shrink + the page to a minimal honest stub stating only what is certain plus the + explicit unresolved ambiguity. Less, but true. +- **Never cite a `keyword`-token entity** as a source. + +## Editing posture — cooperative by default, rebuild only on resolved proof + +Default = **cooperative steward**: if Step-2 resolution shows the page is +basically right, integrate the new members with gentle, additive edits; don't +gratuitously rewrite sound prose. + +**Radical clear-and-rebuild** is allowed (and required) ONLY when Step-2 +independent resolution shows the page conflates distinct entities or asserts +identity/attributes the evidence doesn't support. Then rebuild from the +resolved entity's facts only; move mis-attributed material out. The prior +version is auto-snapshotted, so a resolution-justified rebuild is safe and +reversible. Without that resolved proof, stay cooperative — never blow up a +page on a hunch, and never keep a known-wrong line just because it is there. + +**Preserve prior work — you re-emit the WHOLE page, so losing content is on +you.** The new body must be every still-valid prior claim, section and +`[[ref:UUID]]` **plus** the new members — a superset, not a lossy +re-derivation or a summary. Do NOT drop, shorten, or paraphrase-away sound +existing material just because you are regenerating; carry it forward +verbatim where it still holds. Remove a prior line ONLY when Step-2 +resolution proves it mis-attributed or the evidence proves it wrong — never +by inattention, brevity, or running low on output. If you are unsure whether +a prior statement still holds, KEEP it (and, if needed, note the doubt with +its ref) rather than silently omit it. A shorter page than before, with no +resolution/evidence reason for what vanished, is a FAILED write. + +**Citation is mechanical, not editorial.** Every MEMBER in this job +MUST appear as at least one `[[ref:UUID]]` citation in the new body +— even when the existing prose already covers the same content. The +citation is the *only* signal the system uses to record the +`summarises` relation that links the member to this wiki. Without +the citation the member stays orphaned, the maintainer re-flags it +on the next tick, and the same attach is retried in a loop. If your +section edits don't naturally cite a member, add a bullet for it in +the `references` section before submitting. Whether you do section +edits or a full rewrite, the rule is the same: **no assigned MEMBER +may leave the run un-cited**. + +## Recommended structure (consistency, not a hard gate) + +``` + +# NAME +> **Summary:** one tight line (aim <= 280 chars) +> **Disambiguation:** what this is / is NOT; distinguish it from similarly + named or co-occurring entities, grounded in sources + prose with [[ref:UUID]] + dated claims with [[ref:UUID]] + opposing claims, BOTH refs, reconciled or noted + narrative provenance + one bullet per distinct [[ref:UUID]] you cited, + with a short note — YOU author this to match + your inline citations +``` + +`keywords=` in the meta line is optional — list the concept terms that best +index this page, or omit it. It is the only place keywords come from; nothing +is invented for you. + +Relations are reconciled **additively** from your inline `[[ref:]]` tokens +(every cited entity gets a `summarises` link). Nothing is deleted behind you. +If you deliberately drop a source and want its relation gone, call +`delete_relation` yourself — otherwise just stop citing it. + +## Section-edit path — for attach jobs on a big wiki + +When the existing body is large, re-emitting the whole thing in `body` +can exhaust the context window. Use the section-edit tools instead — +they let you read the OUTLINE only (cheap) and rewrite one section at +a time, persisting each change immediately: + +- `read_wiki_outline(wiki_id)` — section names + char counts + the + current `revision` token. ALWAYS call this first. +- `read_wiki_section(wiki_id, section_name)` — fetch one section's + content + revision. Read only the section(s) you need to touch. +- `edit_wiki_section(wiki_id, section_name, new_content, expect_revision)` + — replace a section, or append a new one if `section_name` doesn't + exist yet. Pass the latest revision you read; on mismatch you get a + "stale revision" error and must re-read before retrying. +- `delete_wiki_section(wiki_id, section_name, expect_revision)` — remove + a section. +- `validate_wiki(wiki_id)` — check refs resolve and grammar invariants + hold. Run after a batch of edits to catch any broken `[[ref:UUID]]`. + +Section-edit grammar invariants when you author `new_content`: +- Inline citations stay `[[ref:UUID]]` or `[[ref:UUID|display]]` + (grouped form `[[ref:UUID1], [ref:UUID2]]` is also tolerated). +- DO NOT include the `` marker yourself — the + tool emits it. Your `new_content` is the section's text only. +- The HEADER (meta line, `# Title`, `> **Summary:**` / + `> **Disambiguation:**`) lives ABOVE the first section marker. + Section edits never touch the header — if the summary needs to + change, either re-edit the `overview` section to reflect the new + scope, or fall back to a full-body rewrite. +- The "Preserve prior work" rule above applies PER SECTION: a + replaced section's `new_content` must include every still-valid + prior claim + `[[ref:UUID]]` from that section, plus the new + material — a superset, not a lossy summary. + +When finished, call `final_answer` with `body=""` (empty string) and +`mode="attach"`. The router detects that the wiki's revision advanced +during your run and skips the full-body write — your section edits are +the authoritative content. If you prefer to just rewrite the whole +body for a small wiki, that path is unchanged — submit the full body +in `body` as before. Don't mix the two on the same run: either use +section tools and submit `body=""`, OR rewrite fully via `body`. + +**`body=""` is ATTACH MODE ONLY.** In `create` or `consolidate` mode +the router REJECTS an empty body — those modes need the full new +content in `body`. For consolidate, that means the complete merged +survivor body (meta + summary + every section + references), period. + +## Context handoff — when you're running out of room + +If the system injects a "your context is filling up" nudge naming the +`handoff_to_successor` tool, the conversation has grown close to the +model's window. You have two choices: + +- If your remaining work fits in **1-2 more turns**, finish cleanly: + call `final_answer` directly. Use `body=""` ONLY if you're in + `attach` mode AND used section edits; for `create` or `consolidate` + always submit the full body. +- Otherwise, call `handoff_to_successor(progress_summary, remaining_work)`. + A fresh agent with the SAME prompt and tools will continue from your + brief. After your handoff call your run ends — the successor takes + over with a clean context. + +The handoff brief must be precise. The successor only sees what you +write: + +- `progress_summary`: a tight list of (a) the tools you've called so + far and what came back of value, (b) any active revision tokens + (e.g., "edited Dimitrios.timeline at revision 14 → 15"), (c) facts + / resolutions / identity decisions you committed to. Keep it + factual; no narrative. +- `remaining_work`: the concrete next tool call(s) the successor must + make. Name wikis, section names, and current revisions explicitly. + Example: "Read `read_wiki_section(wiki_id='25ab...', section_name='references')` + with `expect_revision=15`, then `edit_wiki_section` to add bullets + for fact-ids [a, b, c]. Then `validate_wiki` and call `final_answer` + with `body=""`." + +If your successor ALSO approaches the limit, it can call +`handoff_to_successor` again — the chain continues up to a hard depth +cap. Don't ration handoffs out of politeness; use them whenever the +brief is cheaper than holding the work. + +## Output — STRICT + +Finish by calling `final_answer` exactly once. Its argument is a typed +object — the tool's schema defines and validates the fields; you do not write +delimiters or raw JSON, you just fill the fields: + +- `mode` — `create`, `attach`, or `consolidate` (the mode of THIS job). +- `body` — the COMPLETE markdown wiki page (the full document; the meta + header, summary/disambiguation, every section, references — exactly what + used to go between the body delimiters). MAY be the empty string `""` + in `attach` mode if and only if you persisted your changes via the + section-edit tools; the router detects the revision delta and skips + the full-body write. REQUIRED non-empty for `create` and `consolidate`. +- `canonical_no` — **consolidate mode only**: the NUMBER of the surviving + wiki you chose, taken from the numbered "Duplicate wikis to consolidate" + list above (an integer, e.g. `1`). Never an id. Leave it null for + `create`/`attach`. + +Do not emit anything else. The page lives entirely in `body`. diff --git a/braindb/agent/run_state.py b/braindb/agent/run_state.py new file mode 100644 index 0000000..0dd6fd5 --- /dev/null +++ b/braindb/agent/run_state.py @@ -0,0 +1,124 @@ +""" +Per-run side-channel for the agent's final structured payload. + +Why this exists: `Agent(output_type=)` makes the SDK pass +`response_format: json_schema` on EVERY LLM call (not just the final +one), which steers weaker models to satisfy the schema on turn 1 and +skip tools entirely. We therefore build agents WITHOUT `output_type` so +intermediate turns are free — but then `StopAtTools` would `str()`-coerce +the stop-tool's return into `result.final_output`, and we'd lose the +typed instance. + +This module is the bridge: each `submit_*` tool body parks the +SDK-validated payload via `record_submit(payload)`; `run_typed` reads it +back via `slot.value` after `Runner.run` returns. + +## Why a mutable slot, not just `ContextVar[Any]` + +ContextVar values are inherited by reference into child asyncio Tasks, +but `.set()` inside a child Task does NOT propagate up to the parent. +The openai-agents SDK runs tool bodies (including parallel-tool batches) +inside such child Tasks, so a naive `last_submit.set(payload)` in the +tool body is invisible to the surrounding `run_typed`. Putting a mutable +container in the ContextVar instead — and mutating its `.value` from the +tool — works across that boundary because every Task sees the same +object reference. The standard `set(slot) + reset(token)` lifecycle in +`run_typed` keeps nested runs (parent → `delegate_to_subagent` → +subagent) isolated: each level uses its own `_Slot`. +""" +from contextvars import ContextVar +from typing import Any + + +class _Slot: + """One-shot holder for the validated payload of a single agent run.""" + __slots__ = ("value",) + + def __init__(self) -> None: + self.value: Any = None + + +# Default None — `run_typed` always installs its own slot before awaiting +# `Runner.run`. A `None` here at submit time means "called outside a +# run_typed scope" and is just silently dropped (no slot to write to). +_slot_var: ContextVar["_Slot | None"] = ContextVar( + "braindb_last_submit_slot", default=None, +) + + +def install_slot() -> tuple[_Slot, object]: + """Used by `run_typed` to start a run. Returns `(slot, token)`; pass + `token` to `release_slot` in a `finally:` to restore the previous + context (so nested runs are isolated).""" + slot = _Slot() + token = _slot_var.set(slot) + return slot, token + + +def release_slot(token: object) -> None: + """Restore the previous slot (call in `finally:` after `install_slot`).""" + _slot_var.reset(token) # type: ignore[arg-type] + + +def record_submit(payload: Any) -> None: + """Called from inside every `submit_*` tool body. The SDK has already + validated `payload` against the tool's Pydantic argument schema, so + the value parked here is the typed final answer by construction. + + Mutates the slot in place (does NOT call `ContextVar.set(...)`) — see + module docstring for why.""" + slot = _slot_var.get() + if slot is not None: + slot.value = payload + + +# ====================================================================== # +# Handoff side-channel (writer-only) # +# ====================================================================== # +# +# Parallels the final-answer slot above. The writer's `handoff_to_successor` +# tool parks its brief here; the run wrapper in `routers/wiki.py` reads it +# after `run_typed` returns and decides whether to spawn a successor. Lives +# in run_state.py (not in a writer-specific module) so the slot lifecycle +# uses the same ContextVar discipline — install in the wrapper, mutate in +# the tool body, isolated across nested runs. + + +class _HandoffSlot: + """One-shot holder for the writer's handoff brief. Distinct from + `_Slot` because the wrapper inspects two independent fields + (progress + remaining) rather than a single typed payload.""" + __slots__ = ("captured", "progress_summary", "remaining_work") + + def __init__(self) -> None: + self.captured: bool = False + self.progress_summary: str = "" + self.remaining_work: str = "" + + +_handoff_slot_var: ContextVar["_HandoffSlot | None"] = ContextVar( + "braindb_handoff_slot", default=None, +) + + +def install_handoff_slot() -> tuple[_HandoffSlot, object]: + """Used by the writer's run wrapper to start a run that may end via + handoff. Returns `(slot, token)`; pass `token` to `release_handoff_slot` + in a `finally:`.""" + slot = _HandoffSlot() + token = _handoff_slot_var.set(slot) + return slot, token + + +def release_handoff_slot(token: object) -> None: + _handoff_slot_var.reset(token) # type: ignore[arg-type] + + +def record_handoff(progress_summary: str, remaining_work: str) -> None: + """Called from the `handoff_to_successor` tool body. Mutates the slot + in place (same reason as `record_submit`).""" + slot = _handoff_slot_var.get() + if slot is not None: + slot.captured = True + slot.progress_summary = progress_summary + slot.remaining_work = remaining_work diff --git a/braindb/agent/schemas.py b/braindb/agent/schemas.py new file mode 100644 index 0000000..062e309 --- /dev/null +++ b/braindb/agent/schemas.py @@ -0,0 +1,274 @@ +""" +Typed agent output contract. + +Convention (absolute): every agent/subagent finishes via the `final_answer` +trick, and its payload is ALWAYS one of these Pydantic models — never a loose +free string we scrape. `@function_tool` turns the model into a strict JSON +schema for the tool arguments, so the LLM is constrained to emit valid +structured output instead of free-running and truncating. + +These mirror the style of `braindb/schemas/` (the REST layer); they reuse the +existing pydantic dependency — no new dependency, no new machinery. +""" +import json +from typing import Literal + +from pydantic import BaseModel, Field, field_validator, model_validator + + +# Coercion helpers — weak/quantised models often emit "" (empty string) for +# nullable fields instead of `null`, or `null` for empty-list fields instead +# of `[]`. The Pydantic schemas are nullable + defaulted at the type level; +# these `before` validators just accept the wrong-type variants gracefully +# so we don't reject a perfectly intended "skip" decision because the model +# sent `target_wiki_no=""` instead of `null`. The validation contract is +# unchanged — we still produce a properly-typed Pydantic instance. + + +# Top-level coercion — some providers (notably vLLM / Qwen) emit tool-call +# `arguments.payload` as a JSON-encoded STRING ("{\"action\": \"skip\", ...}") +# instead of a JSON object ({"action": "skip", ...}). This is technically +# OpenAI-spec-compliant (the outer `arguments` field IS defined as a string +# of JSON), but the SDK only unwraps once and then hands the inner value to +# Pydantic as-is — so when the inner value is itself a JSON string, Pydantic +# rejects it with "Input should be a valid dictionary". +# +# The `@model_validator(mode="before")` below catches this case: if the input +# is a string that parses as JSON to a dict, we use the parsed dict; if it +# parses to anything else (list / int / null), we let Pydantic raise its +# usual "valid dictionary" error so the LLM sees a clear correction. Dict +# inputs are passed through untouched — well-behaved providers (deepinfra, +# OpenAI, Anthropic via LiteLLM) see exactly the same behaviour as today. +# +# This is the SAME pattern as the nullable-field coercion above, just at the +# whole-model level rather than per-field. The LLM-visible schema is +# unchanged; we don't advertise string-form acceptance to the model. + +def _maybe_parse_json_string(v): + """If `v` is a JSON-encoded string of an object, parse it. Otherwise + pass through unchanged. Pydantic v2 calls @model_validator(mode='before') + BEFORE field-level validation, so a returned dict goes through the rest + of the validation pipeline (including the per-field coercers below) + exactly as if the LLM had sent a dict in the first place.""" + if isinstance(v, str): + try: + parsed = json.loads(v) + except (json.JSONDecodeError, ValueError): + return v # let Pydantic raise its normal error + # Defensive: occasionally Qwen-class quantised models emit the dict + # double-escaped (the first parse yields a string of JSON, not a + # dict). One more parse attempt unwraps that case. Safe — only fires + # on a string result, only returns a value if it parses to a dict. + if isinstance(parsed, str): + try: + parsed = json.loads(parsed) + except (json.JSONDecodeError, ValueError): + return v + # Only return the parsed value if it's a dict — anything else (list, + # int, null) is not a valid Pydantic-model input; let Pydantic raise. + if isinstance(parsed, dict): + return parsed + return v + +def _coerce_empty_to_none(v): + """Accept '', 'null', 'none', 'n/a' (any case, with/without whitespace) + as equivalent to None for nullable fields.""" + if v is None: + return None + if isinstance(v, str): + s = v.strip() + if not s or s.lower() in ("null", "none", "n/a"): + return None + return v + + +def _coerce_to_int_or_none(v): + """For nullable-int fields: '' / 'null' / etc → None; numeric strings → int.""" + v = _coerce_empty_to_none(v) + if v is None or isinstance(v, int): + return v + try: + return int(v) + except (TypeError, ValueError): + return None # last resort — don't fail the whole submission on a bad number + + +def _coerce_to_list(v): + """For list fields: None / '' → []; everything else as-is for Pydantic to validate.""" + if v is None or v == "": + return [] + return v + + +class AgentAnswer(BaseModel): + """General recall/save answer (the public /agent/query endpoint). + + The endpoint is general-purpose (Claude Code, arbitrary recall/save), so + the answer itself is necessarily natural language — but it is still + delivered through the typed `final_answer` trick, never as loose + top-level model output. + """ + answer: str = Field(..., description="The full natural-language response to the caller.") + + # Safety net for providers (notably vLLM/Qwen) that emit the tool-call + # arg as a JSON-encoded string instead of a JSON object. See the helper + # docstring at the top of the file. + @model_validator(mode="before") + @classmethod + def _accept_json_string(cls, v): + return _maybe_parse_json_string(v) + + +class MaintainerDecision(BaseModel): + """The wiki maintainer's per-orphan decision. Existing wikis are + referenced by their CATALOG NUMBER (the numbered list at the end of the + prompt), never by uuid — the harness maps number->id deterministically. + + Action-dependent fields: `target_wiki_no`, `proposed_name`, and + `consolidate_nos` are only meaningful for one specific action each. For + every other action you MUST send JSON `null` for the optional ones (not + "", not 0, not "n/a") and an empty array `[]` for `consolidate_nos`. + """ + action: Literal["attach", "create", "consolidate", "skip", "ambiguous"] = Field( + ..., + description=( + "The decision for this orphan. Exactly one of: " + "`attach` (link to an existing wiki by catalog number), " + "`create` (mint a new wiki with a proposed name), " + "`consolidate` (merge >=2 catalog-numbered wikis), " + "`skip` (not worth a wiki — infrastructural / keyword-token), " + "`ambiguous` (cannot disambiguate the real subject)." + ), + ) + target_wiki_no: int | None = Field( + None, + description=( + "REQUIRED ONLY when action=`attach`: the integer CATALOG NUMBER " + "of the existing wiki to attach the orphan to (1-indexed, taken " + "from the numbered WIKIS list at the end of the prompt). " + "For action in (`create`, `consolidate`, `skip`, `ambiguous`) " + "this field MUST be JSON null. Do NOT use empty string \"\", 0, " + "or 'n/a' — use literal null." + ), + ) + proposed_name: str | None = Field( + None, + description=( + "REQUIRED ONLY when action=`create`: the canonical name for the " + "new wiki (must appear in the evidence — never invent). " + "For action in (`attach`, `consolidate`, `skip`, `ambiguous`) " + "this field MUST be JSON null. Do NOT use empty string \"\"." + ), + ) + consolidate_nos: list[int] = Field( + default_factory=list, + description=( + "REQUIRED ONLY when action=`consolidate`: an array of >=2 " + "integer CATALOG NUMBERS naming the duplicate wikis to merge " + "(from the numbered WIKIS list). " + "For every other action this field MUST be an empty array [] " + "(NOT null, NOT empty string)." + ), + ) + rationale: str = Field( + ..., + description=( + "ALWAYS REQUIRED. One to three sentences justifying the chosen " + "action: which catalog wiki(s) you matched (or that the catalog " + "has none), and why this action was the right one. This makes " + "the decision auditable." + ), + ) + + # Top-level coercion: accept JSON-string-of-dict (vLLM/Qwen quirk). + @model_validator(mode="before") + @classmethod + def _accept_json_string(cls, v): + return _maybe_parse_json_string(v) + + # Forgiving coercion — weak/quantised models often emit empty strings or + # "null" strings instead of literal JSON null. Accept those as None + # rather than rejecting the whole submission (the prompt and the + # descriptions above ask for null; the validators are the safety net). + @field_validator("target_wiki_no", mode="before") + @classmethod + def _coerce_target_wiki_no(cls, v): + return _coerce_to_int_or_none(v) + + @field_validator("proposed_name", mode="before") + @classmethod + def _coerce_proposed_name(cls, v): + return _coerce_empty_to_none(v) + + @field_validator("consolidate_nos", mode="before") + @classmethod + def _coerce_consolidate_nos(cls, v): + return _coerce_to_list(v) + + +class WikiWriteResult(BaseModel): + """The wiki writer's full output. `body` is the complete markdown page — + a typed field of the schema, exactly like any other field (not loose + text, not delimiter-wrapped). + + `canonical_no` is only meaningful for `consolidate` mode. For + `create` / `attach` you MUST send JSON null (not "", not 0). + """ + mode: Literal["create", "attach", "consolidate"] = Field( + ..., + description=( + "The write mode of THIS job (matches the mode the harness " + "passed in the prompt): `create` (fresh wiki), `attach` " + "(integrate new members into an existing wiki), `consolidate` " + "(merge multiple duplicate wikis into a survivor)." + ), + ) + canonical_no: int | None = Field( + None, + description=( + "REQUIRED ONLY when mode=`consolidate`: the integer NUMBER of " + "the surviving wiki chosen from the numbered DUPLICATES list " + "in the prompt (1-indexed, never a uuid). " + "For mode in (`create`, `attach`) this field MUST be JSON null. " + "Do NOT use empty string \"\", 0, or 'n/a'." + ), + ) + body: str = Field( + "", + description=( + "The COMPLETE markdown wiki page — the full document. Include " + "the meta header, summary, disambiguation, every section, all " + "[[ref:UUID]] citations, and the references section. This is " + "what becomes the wiki entity's content; it replaces the prior " + "body wholesale (the prior version is auto-snapshotted).\n\n" + "MAY be empty ONLY in `attach` mode AND only if you persisted " + "your changes via the section-edit tools " + "(`edit_wiki_section` / `delete_wiki_section`). In that case " + "the router detects the wiki's revision moved during your run " + "and skips the full-body write — your section edits are " + "already the authoritative content. For `create` and " + "`consolidate` modes this field MUST be non-empty." + ), + ) + + # Top-level coercion: accept JSON-string-of-dict (vLLM/Qwen quirk). + @model_validator(mode="before") + @classmethod + def _accept_json_string(cls, v): + return _maybe_parse_json_string(v) + + @field_validator("canonical_no", mode="before") + @classmethod + def _coerce_canonical_no(cls, v): + return _coerce_to_int_or_none(v) + + +class SubagentResult(BaseModel): + """A delegated subagent's return (replaces the free-string subagent answer).""" + result: str = Field(..., description="The distilled result of the delegated task.") + + # Top-level coercion: accept JSON-string-of-dict (vLLM/Qwen quirk). + @model_validator(mode="before") + @classmethod + def _accept_json_string(cls, v): + return _maybe_parse_json_string(v) diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index 3bc3b6f..3628c12 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -34,7 +34,15 @@ link_entity_to_keywords, sync_keywords_for_entity, ) -from braindb.services.search import fuzzy_search +from braindb.services.search import fuzzy_search, preview, slice_content +from braindb.services import wiki_sections as ws +from braindb.agent.run_state import record_handoff, record_submit +from braindb.agent.schemas import ( + AgentAnswer, + MaintainerDecision, + SubagentResult, + WikiWriteResult, +) logger = logging.getLogger(__name__) @@ -100,15 +108,42 @@ async def wrapper(*args, **kwargs): @function_tool @_verbose("recall_memory") -async def recall_memory(queries: list[str], max_results: int = 15) -> str: +async def recall_memory( + queries: list[str], + max_results: int = settings.recall_default_max_results, +) -> str: """Search BrainDB memory with multiple natural language queries. Runs fuzzy + fulltext + keyword embedding search, merges with geometric mean, traverses the graph up to 3 hops, applies temporal decay. Use this as the primary recall tool. + QUERY STRATEGY — IMPORTANT for high-recall on narrow subjects: + + BrainDB indexes via short keyword entities. A 1-word query like + "Petros" matches the keyword "Petros" cleanly (similarity ~1.0). A + long phrase like "Petros person identity profile" matches the same + keyword at much lower similarity (~0.4) because pg_trgm dilutes + when comparing short keywords to long query strings. + + Therefore: prefer MULTIPLE narrow queries over one long phrase. The + sweet spot for a focused subject is: + - one or two SINGLE-KEYWORD queries (the names you care about), + - plus 1-2 broader semantic phrases for adjacent context. + + Examples: + GOOD: ["Petros", "Selonda Saronikos fish farm", "Dimitrios manager"] + BAD: ["Petros person identity profile relation to Dimitris"] + + Each query you provide gets a reserved share of the top results + (per-search-term quota), so adding the bare keyword as one of your + queries GUARANTEES that subject surfaces — it doesn't compete with + the broader phrases. + Args: - queries: List of search queries (use multiple angles for better coverage). - max_results: Max items to return (1-100, default 15). + queries: List of search queries. Prefer 2-4 short focused queries + over one long phrase. Include the bare keyword(s) of the + subject you're investigating as standalone queries. + max_results: Max items to return (1-100, default 30). """ try: req = ContextRequest(queries=queries, max_results=max_results) @@ -315,11 +350,22 @@ async def save_rule( @function_tool @_verbose("get_entity") -async def get_entity(entity_id: str) -> str: - """Fetch a single entity by ID (returns JSON blob with all fields). +async def get_entity(entity_id: str, offset: int = 0, limit: Optional[int] = None) -> str: + """Fetch ONE entity by ID — the full-content read (recall/list only give + previews; come here to read a thing fully). + + For a LARGE body, page it with offset/limit instead of pulling it whole: + the response includes `content_meta` {total_chars, offset, returned, + next_offset}. Loop `next_offset` until null. To avoid polluting your own + context, hand each slice to `delegate_to_subagent` ("process THIS slice…") + and aggregate — never load a huge document into your main context. Args: entity_id: UUID of the entity. + offset: start char of the content slice (default 0). + limit: max chars of this slice (clamped to the server slice max). + If offset and limit are both omitted, the full body is returned + (legacy behaviour, unchanged). """ try: with get_conn() as conn: @@ -331,7 +377,14 @@ async def get_entity(entity_id: str) -> str: d = dict(row) d.pop("embedding", None) d.pop("search_vector", None) - return _truncate(json.dumps(d, default=str, indent=2)) + if offset == 0 and limit is None: + return _truncate(json.dumps(d, default=str, indent=2)) + # Explicit slice request → return exactly that slice + paging meta, + # NOT re-clipped by _truncate (slice is already bounded by SLICE_MAX). + chunk, meta = slice_content(d.get("content"), offset, limit) + d["content"] = chunk + d["content_meta"] = meta + return json.dumps(d, default=str, indent=2) except Exception as e: return _err(str(e)) @@ -382,7 +435,7 @@ async def list_entities( lines.append( f"[{r['entity_type']}] imp={r['importance']} src={r.get('source', '-')}\n" f" id: {r['id']}\n" - f" content: {r['content']}\n" + f" content: {preview(r['content'], r['id'])}\n" f" keywords: {r.get('keywords', [])}" ) return _truncate("\n".join(lines)) @@ -632,7 +685,9 @@ async def search_sql(query: str) -> str: columns = [d[0] for d in cur.description] if cur.description else [] rows = cur.fetchmany(1000) log_activity(conn, "sql_query", details={"query": query[:500], "rows": len(rows)}) - result = {"columns": columns, "rows": [[str(v) if v is not None else None for v in r] for r in rows], "row_count": len(rows)} + result = {"columns": columns, + "rows": [[preview(v) if v is not None else None for v in r] for r in rows], + "row_count": len(rows)} return _truncate(json.dumps(result, default=str, indent=2)) except Exception as e: return _err(str(e)) @@ -776,7 +831,7 @@ async def delegate_to_subagent(task: str) -> str: tool outputs. The subagent has access to all the same BrainDB tools. Write a clear, self-contained task description — the subagent doesn't see - your prior context. End by telling it to call submit_result with a summary. + your prior context. End by telling it to call final_answer with a summary. Args: task: A self-contained task description for the subagent. @@ -787,20 +842,21 @@ async def delegate_to_subagent(task: str) -> str: _call_depth += 1 try: # Local imports to avoid circular dependency on agent.py - from agents import Runner - from braindb.agent.agent import create_braindb_agent + from braindb.agent.agent import get_subagent, run_typed from braindb.config import settings logger.info("Subagent starting: %s", task[:200]) - subagent = create_braindb_agent() - result = await Runner.run( - starting_agent=subagent, - input=task, + # run_typed isolates the subagent's submit slot from ours (its own + # `last_submit.set(None)` token + reset in `finally`), so we cannot + # leak the subagent's SubagentResult into the parent's run_typed. + payload: SubagentResult = await run_typed( + task, + get_subagent(), + SubagentResult, max_turns=settings.agent_subagent_max_turns, ) - answer = str(result.final_output) logger.info("Subagent completed.") - return _truncate(answer) + return _truncate(payload.result) except Exception as e: logger.exception("Subagent failed") return _err(f"subagent failed: {e}") @@ -809,15 +865,322 @@ async def delegate_to_subagent(task: str) -> str: # ====================================================================== # -# FINAL TOOL — stops the loop # +# WIKI SECTION EDITS — read/write slices of a wiki body (writer-only) # # ====================================================================== # +# +# Wiki bodies can grow past the writer's context window. These tools let +# the writer read just an outline (cheap) and edit one section at a time +# instead of re-emitting the whole markdown blob every turn. Wired into +# the writer agent only (see braindb/agent/agent.py). +# +# Strict-markers contract: tools error if the target body has no +# `` markers. Phase 0 confirmed all active wikis +# already do. +# +# Optimistic concurrency via `wikis_ext.revision`: every read returns +# the current revision; every write requires it as `expect_revision`. A +# mismatch returns a "stale" ERROR string so the model re-reads instead +# of stomping a concurrent edit (or its own stale mental state). + +import re as _re +_SECTION_NAME_RE = _re.compile(r"[A-Za-z0-9_\-]+") + + +@function_tool +@_verbose("read_wiki_outline") +async def read_wiki_outline(wiki_id: str) -> str: + """Outline of a wiki — section names + char counts + current revision. + Call before editing. + + Args: + wiki_id: The wiki's entity UUID. + """ + try: + with get_conn() as conn: + fetched = ws.fetch_wiki_for_section_op(conn, wiki_id) + if fetched is None: + return _err(f"wiki not found: {wiki_id}") + body, revision = fetched + _, sections = ws.parse_sections(body) + if not sections: + return _err( + f"wiki {wiki_id} body has no markers " + f"(strict-markers contract violated; cannot edit)" + ) + lines = [f"revision: {revision}", f"sections: {len(sections)}"] + for s in sections: + lines.append(f" - {s.name}: {s.char_count}ch") + return "\n".join(lines) + except Exception as e: + return _err(str(e)) + @function_tool -@_verbose("submit_result") -async def submit_result(answer: str) -> str: - """Submit the final answer to the query. Call this exactly once when you're done. +@_verbose("read_wiki_section") +async def read_wiki_section(wiki_id: str, section_name: str) -> str: + """Read one section's content + the wiki's current revision token. Args: - answer: The full response to send back to the caller. + wiki_id: The wiki's entity UUID. + section_name: Section name as listed by read_wiki_outline. """ - return answer + try: + with get_conn() as conn: + fetched = ws.fetch_wiki_for_section_op(conn, wiki_id) + if fetched is None: + return _err(f"wiki not found: {wiki_id}") + body, revision = fetched + _, sections = ws.parse_sections(body) + match = next((s for s in sections if s.name == section_name), None) + if match is None: + names = ", ".join(s.name for s in sections) or "(none)" + return _err(f"section '{section_name}' not found. Existing: {names}") + return _truncate( + f"revision: {revision}\nsection: {match.name}\n" + f"content:\n{match.content}" + ) + except Exception as e: + return _err(str(e)) + + +@function_tool +@_verbose("edit_wiki_section") +async def edit_wiki_section( + wiki_id: str, + section_name: str, + new_content: str, + expect_revision: int, +) -> str: + """Replace one section's content. If section_name is new, appends a + fresh section at the end. Revision mismatch → returns ERROR: re-read + first. + + Args: + wiki_id: The wiki's entity UUID. + section_name: Section to replace (or new section to append). + Use lowercase letters, digits, dashes, underscores only. + new_content: Full new content of the section (without the marker + line — the tool re-emits it). + expect_revision: Revision token from the last read on this wiki. + """ + if not _SECTION_NAME_RE.fullmatch(section_name): + return _err( + f"invalid section_name '{section_name}': use only letters, " + f"digits, dashes, underscores" + ) + try: + with get_conn() as conn: + fetched = ws.fetch_wiki_for_section_op(conn, wiki_id) + if fetched is None: + return _err(f"wiki not found: {wiki_id}") + body, current_rev = fetched + if current_rev != expect_revision: + return _err( + f"stale revision: you passed {expect_revision}, " + f"current is {current_rev}. Re-read the section first." + ) + _, sections = ws.parse_sections(body) + if not sections: + return _err( + f"wiki {wiki_id} body has no markers; " + f"strict-markers contract violated" + ) + appended = all(s.name != section_name for s in sections) + new_body = ws.splice_section(body, section_name, new_content) + new_rev = ws.apply_section_write(conn, wiki_id, new_body, expect_revision) + log_activity(conn, "update", "wiki", wiki_id, details={ + "op": "edit_wiki_section", + "section": section_name, + "appended": appended, + "revision": new_rev, + }) + verb = "appended" if appended else "replaced" + return f"ok — section '{section_name}' {verb}. new revision: {new_rev}" + except ws.StaleRevisionError as e: + return _err(str(e)) + except Exception as e: + return _err(str(e)) + + +@function_tool +@_verbose("delete_wiki_section") +async def delete_wiki_section( + wiki_id: str, + section_name: str, + expect_revision: int, +) -> str: + """Remove a section. Revision mismatch → ERROR: re-read first. + + Args: + wiki_id: The wiki's entity UUID. + section_name: Section to remove. + expect_revision: Revision token from the last read on this wiki. + """ + try: + with get_conn() as conn: + fetched = ws.fetch_wiki_for_section_op(conn, wiki_id) + if fetched is None: + return _err(f"wiki not found: {wiki_id}") + body, current_rev = fetched + if current_rev != expect_revision: + return _err( + f"stale revision: you passed {expect_revision}, " + f"current is {current_rev}. Re-read first." + ) + try: + new_body = ws.delete_section(body, section_name) + except KeyError: + _, sections = ws.parse_sections(body) + names = ", ".join(s.name for s in sections) or "(none)" + return _err(f"section '{section_name}' not found. Existing: {names}") + new_rev = ws.apply_section_write(conn, wiki_id, new_body, expect_revision) + log_activity(conn, "update", "wiki", wiki_id, details={ + "op": "delete_wiki_section", + "section": section_name, + "revision": new_rev, + }) + return f"ok — section '{section_name}' deleted. new revision: {new_rev}" + except ws.StaleRevisionError as e: + return _err(str(e)) + except Exception as e: + return _err(str(e)) + + +@function_tool +@_verbose("validate_wiki") +async def validate_wiki(wiki_id: str) -> str: + """Check the wiki body grammar: section markers present, refs + well-formed, summary callout present. Returns 'ok' or one issue per + line. + + Args: + wiki_id: The wiki's entity UUID. + """ + try: + with get_conn() as conn: + fetched = ws.fetch_wiki_for_section_op(conn, wiki_id) + if fetched is None: + return _err(f"wiki not found: {wiki_id}") + body, revision = fetched + issues = ws.check_grammar(body) + if not issues: + return f"ok — revision: {revision}, no issues" + return ( + f"revision: {revision}\nissues:\n" + + "\n".join(f" - {i}" for i in issues) + ) + except Exception as e: + return _err(str(e)) + + +# ====================================================================== # +# CONTEXT HANDOFF — end this run, successor continues (writer-only) # +# ====================================================================== # +# +# Called by the writer when it gets a context-near-full nudge from +# `CountdownHooks` and decides remaining work doesn't fit. The router's +# writer wrapper (braindb/routers/wiki.py) detects the handoff slot was +# filled and spawns a successor agent — same prompt, same tools, fresh +# context, seeded with the brief. +# +# The tool ALSO parks a placeholder `WikiWriteResult` via `record_submit` +# so `run_typed`'s typed-final contract is satisfied — the placeholder +# is never the authoritative output; the wrapper reads the handoff slot +# instead. This avoids any change to `run_typed`'s shape. + +@function_tool +@_verbose("handoff_to_successor") +async def handoff_to_successor(progress_summary: str, remaining_work: str) -> str: + """End this run early; a successor with the SAME prompt and tools + will continue from your brief. Use when you've been nudged about + context approaching the limit AND remaining work doesn't fit in 1-2 + turns. + + Args: + progress_summary: Tools you've called, key findings, and any + ACTIVE revision tokens (for the wiki you've been editing). + The successor only sees this — be precise. + remaining_work: The concrete next tool call(s) the successor + must make — name wikis, section names, current revisions. + Example: "Call read_wiki_section(wiki_id='abc', section_name='timeline') + with expect_revision=15, then edit_wiki_section(...) with the + new timeline content merging facts from member fact-id xyz." + """ + record_handoff(progress_summary, remaining_work) + # Park a placeholder WikiWriteResult so run_typed's typed-final + # contract is satisfied. mode/body are intentionally minimal — the + # router consults the handoff slot first when this run ends. The + # writer's StopAtTools list includes `handoff_to_successor`, so + # the loop halts cleanly after this returns. + record_submit(WikiWriteResult(mode="attach", body="")) + return "handoff registered; this run is ending — successor will continue from your brief" + + +# ====================================================================== # +# FINAL TOOL — stops the loop # +# ====================================================================== # + +# Convention (absolute): the run finishes ONLY by calling `final_answer`, +# and its argument is ALWAYS a typed Pydantic model — never a loose string. +# `@function_tool` validates the LLM's call args against the model BEFORE +# invoking the body, so `payload` is guaranteed-valid inside each function. +# +# strict_mode=False: critical. The default `strict_mode=True` activates +# OpenAI structured-outputs strict JSON schema, which forces EVERY +# property of the embedded Pydantic model into the schema's `required` +# list — overriding Pydantic's own view that fields with `= None` or +# `default_factory=...` are optional. On `MaintainerDecision` and +# `WikiWriteResult`, that inflation makes the LLM emit args that pass +# Pydantic but fail the over-strict schema, producing endless +# "Invalid JSON input: 1 validation error" loops the Layer 4 retry +# can't escape (verified live on deepinfra/Gemma against the wiki +# maintainer). Turning strict_mode off makes the LLM-visible schema +# match Pydantic's required list exactly; Pydantic still validates the +# parsed args inside the tool body, so the typed-final contract is +# unchanged — we just stop demanding the model emit fields it doesn't +# need. +# There is one typed variant per agent purpose; every variant keeps the +# name "final_answer" so prompts and `StopAtTools(["final_answer"])` +# stay generic. +# +# Each variant parks the validated payload into the per-Task ContextVar +# (see braindb/agent/run_state.py) so `run_typed` can hand it back +# typed. The returned "ok" string is irrelevant — we never read +# `result.final_output`; `StopAtTools` only needs the loop to stop. +# +# Why a ContextVar instead of `output_type=` on the Agent: +# `output_type` makes the SDK pass `response_format: json_schema` on +# EVERY LLM turn (not just the final one), which steers weaker models to +# satisfy the schema on turn 1 and never call tools. The side-channel +# capture keeps middle turns free while still delivering a typed final. + +@function_tool(name_override="final_answer", strict_mode=False) +@_verbose("final_answer") +async def submit_answer(payload: AgentAnswer) -> str: + """Submit the final answer. Call this exactly once when you're done.""" + record_submit(payload) + return "ok" + + +@function_tool(name_override="final_answer", strict_mode=False) +@_verbose("final_answer") +async def submit_maintainer(payload: MaintainerDecision) -> str: + """Submit the maintainer decision. Call this exactly once when you're done.""" + record_submit(payload) + return "ok" + + +@function_tool(name_override="final_answer", strict_mode=False) +@_verbose("final_answer") +async def submit_wiki(payload: WikiWriteResult) -> str: + """Submit the finished wiki. Call this exactly once when you're done.""" + record_submit(payload) + return "ok" + + +@function_tool(name_override="final_answer", strict_mode=False) +@_verbose("final_answer") +async def submit_subagent(payload: SubagentResult) -> str: + """Submit the delegated task result. Call this exactly once when you're done.""" + record_submit(payload) + return "ok" diff --git a/braindb/config.py b/braindb/config.py index c27eb08..25acfcd 100644 --- a/braindb/config.py +++ b/braindb/config.py @@ -6,6 +6,11 @@ # Each profile is a LiteLLM model prefix + the env var holding its API key, # plus an optional base_url for self-hosted OpenAI-compatible servers (vLLM, # Ollama, llama.cpp). Adding a new provider is a dict entry, no code change. +# +# `deepinfra` is the recommended default — fast, cheap, validated end-to-end +# in the wiki/maintainer/writer pipeline. The `vllm_*` profiles are for +# advanced / self-hosted / offline use and require a workstation GPU +# (typically reached over an SSH tunnel from the docker network). _LLM_PROFILES: dict[str, dict[str, str]] = { "nim": { "model": "nvidia_nim/google/gemma-4-31b-it", @@ -20,6 +25,16 @@ "api_key_env": "VLLM_API_KEY", "base_url": "http://host.docker.internal:8002/v1", }, + "vllm_workstation_qwen": { + "model": "openai/cyankiwi/Qwen3.6-27B-AWQ-INT4", + "api_key_env": "VLLM_API_KEY", + "base_url": "http://host.docker.internal:8010/v1", + }, + "vllm_workstation_gemma": { + "model": "openai/cyankiwi/gemma-4-31B-it-AWQ-4bit", + "api_key_env": "VLLM_API_KEY", + "base_url": "http://host.docker.internal:8009/v1", + }, } @@ -35,6 +50,7 @@ class Settings(BaseSettings): decay_rate_source: float = 0.002 decay_rate_datasource: float = 0.001 decay_rate_rule: float = 0.0 + decay_rate_wiki: float = 0.0 # synthesised pages should not fade # Graph traversal max_graph_depth: int = 3 @@ -44,16 +60,112 @@ class Settings(BaseSettings): # Scoring missing_signal_penalty: float = 0.5 # multiplier when only text OR only embedding matches (0-1) + # Scoring-pool caps. These bound the CANDIDATE pool that feeds ranking + # (pure SQL/vector work — cheap, runs once per query). They are NOT the + # LLM-visible cap; the caller's `max_results` truncates the FINAL sorted + # items list. Keeping these wide is essential: a narrow single-word + # keyword (e.g. "Petros") embedded against a multi-word sentence query + # may not place in the top 30 most-similar keywords even when it's the + # exact match — without enough headroom, nothing tagged with that + # keyword enters the scoring pool at all. + scoring_pool_keyword_neighbors: int = 500 # top-K keyword embeddings to consider + scoring_pool_fuzzy: int = 500 # top-K fuzzy/full-text candidates to consider + + # Two-level diversity quota on recall output. + # + # Level 1 — per-search-term: each query string in `queries[]` gets + # `per_query_share / num_queries` of `max_results` reserved for its + # OWN top-ranked entities. Forces multi-angle representation: if + # the agent issues [narrow_keyword, broader_phrase, third_angle], + # all three angles surface in the result, regardless of which one + # has the highest absolute scores. Set per_query_share=0 to disable. + # + # Level 2 — per-keyword (dominant matched keyword): walks the + # remaining (open) slots in `final_rank` order and gives each new + # dominant keyword a halving slot allowance (50% / 25% / 12.5% ... + # of max_results, floor 1). Stops one popular keyword (e.g. + # `user-profile`) from monopolising the open portion. + # + # The two levels share ONE counter dict — L1 reservations decrement + # the same per-keyword allowance L2 walks against. So a popular + # keyword cannot double-spend across the two layers. + per_query_share: float = 0.5 + keyword_quota_halving: float = 0.5 + + # How many entities the LLM-facing recall (`recall_memory` tool / + # `/memory/context` API) returns by default. Wider default = the LLM + # sees more candidates per call (more diverse, more discoverable), + # at the cost of more prompt tokens. Tune in code, not via .env, so + # all deployments share one measure. + recall_default_max_results: int = 30 + # Always-on rules cap max_always_on_rules: int = 10 # Agent (LiteLLM — provider selected via llm_profile) llm_profile: str = "deepinfra" agent_model: str = "" # blank = use profile's default model - agent_max_turns: int = 15 + # Bumped 15 → 20 after live observation on Qwen 27B AWQ-INT4 (vLLM): + # deep-research-style runs commonly need >15 tool turns to land + # `final_answer`. 20 gives breathing room; finishes-fast providers + # (deepinfra/Gemma) are unaffected because they don't get close. Lower + # than ~15 will regress Qwen behaviour. Callers that need a different + # value (wiki maintainer/writer pass 30; ingest watcher passes 30/40) + # still do so explicitly via `max_turns=` overrides. + agent_max_turns: int = 20 agent_subagent_max_turns: int = 30 agent_verbose: bool = False + # Runtime "start wrapping up, you have N turns left" nudge (Layer 3 of + # Stage C). When ≤ this many LLM-call turns remain before `max_turns` + # is exhausted, `CountdownHooks` injects ONE synthetic user message + # into the running conversation reminding the model to start + # concluding research and call `final_answer`. The message tone is + # context-aware: soft "start wrapping up" when `max_turns` is generous + # (> 5), hard "call final_answer NOW" when the budget is tight (≤ 5, + # which naturally covers the Layer 4 retry path with `max_turns=3`). + # One nudge per run, never spammed. Set to 0 to disable entirely. + # Bumped 5 → 8 so the nudge fires earlier and the model has room to + # wrap up cleanly instead of slamming into the wall at the last turn. + agent_countdown_threshold: int = 8 + + # Retry-with-correction when a run ends without `final_answer` (Layer 4 + # of Stage C). If the model emits prose instead of calling the typed + # termination tool, instead of raising immediately we append a synthetic + # user-role correction message ("you ended without final_answer, call + # it now") to the existing conversation (via `RunResult.to_input_list()`) + # and re-invoke `Runner.run` ONCE with a small budget. If the retry + # produces `final_answer` -> return the typed payload (HTTP 200). If the + # retry ALSO fails -> raise `RuntimeError` (strict; no silent success + # on a model that refuses the contract even after correction). + # Bounded by `agent_retry_max_turns`; opt-out via setting to False. + agent_retry_on_missing_final: bool = True + agent_retry_max_turns: int = 3 + + # Writer-only context-handoff threshold. When the cheap token estimate + # of the writer's running conversation crosses this absolute number, + # `CountdownHooks` injects ONE synthetic user message asking the model + # to call `handoff_to_successor` with a structured brief (progress + + # remaining work). The writer's run wrapper in `routers/wiki.py` then + # spawns a successor agent (same prompt + tools, fresh context) seeded + # with that brief. Bounded by `agent_writer_handoff_max_depth` so a + # misbehaving model cannot thrash forever. + # + # Why a single absolute-token knob rather than a per-profile pct: + # avoids per-profile bookkeeping. Tuned for the main production + # target (Qwen 27B at max_model_len=40960, so 20000 ≈ 49% — fires + # only when context is genuinely close to half-full). On the + # hosted-Gemma 32K path 20000 is also safe (~63%). On the local + # Gemma 13K path the budget is above the window so handoff never + # fires — that's fine because the small-context path fails at + # initial prompt construction long before the handoff can help. + # Default was 9000 during the Phase-3 dry run; observation showed + # that fired the handoff on routine consolidates that fit inline + # on Qwen, fragmenting work across successors unnecessarily. Set + # to 0 to disable the handoff nudge entirely. + agent_writer_handoff_token_budget: int = 20000 + agent_writer_handoff_max_depth: int = 3 + @property def resolved_agent_model(self) -> str: return self.agent_model or _LLM_PROFILES[self.llm_profile]["model"] diff --git a/braindb/ingest_watcher.py b/braindb/ingest_watcher.py index d1a5391..4a159e1 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -12,7 +12,7 @@ the chunk text directly from the prompt (no get_entity), extracts concrete facts, saves each via save_fact, and links each back to the datasource via create_relation(derived_from). Returns the list of new - fact IDs in submit_result for the watcher to parse. + fact IDs in final_answer for the watcher to parse. Phase B — one /agent/query with only the fact IDs + their 1-sentence content prefetched by the watcher. The central review agent creates @@ -145,7 +145,7 @@ def fetch_entity(entity_id: str) -> dict | None: def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk_text: str) -> 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 submit_result answer. + Returns the list of new fact IDs parsed from the agent's final_answer answer. """ prompt = ( f"A document was just ingested into BrainDB.\n" @@ -166,7 +166,7 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f' relevance_score=0.9, description="Fact extracted from {title}").\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 submit_result with\n" + f"When all facts in this chunk are processed, 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" @@ -217,7 +217,7 @@ def central_review(ds_id: str, title: str, fact_ids: list[str]) -> None: f" related, link them with tagged_with or refers_to.\n\n" f"Do NOT call get_entity — all facts are listed above. Do NOT touch the\n" f"datasource content.\n\n" - f"When done, call submit_result with a short summary of what you added." + f"When done, call final_answer with a short summary of what you added." ) answer = call_agent(prompt, max_turns=30) if answer is None: diff --git a/braindb/main.py b/braindb/main.py index 1921494..21040e5 100644 --- a/braindb/main.py +++ b/braindb/main.py @@ -3,7 +3,7 @@ from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware -from braindb.routers import agent, entities, memory, relations +from braindb.routers import agent, entities, memory, relations, wiki from braindb.services.embedding_service import get_embedding_service logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s") @@ -25,6 +25,7 @@ app.include_router(relations.router) app.include_router(memory.router) app.include_router(agent.router) +app.include_router(wiki.router) @app.on_event("startup") diff --git a/braindb/routers/entities.py b/braindb/routers/entities.py index cf86827..7f92c93 100644 --- a/braindb/routers/entities.py +++ b/braindb/routers/entities.py @@ -18,8 +18,10 @@ RuleCreate, RuleRead, RuleUpdate, SourceCreate, SourceRead, SourceUpdate, ThoughtCreate, ThoughtRead, ThoughtUpdate, + WikiCreate, WikiRead, WikiUpdate, ) from braindb.services.activity_log import log_activity +from braindb.services.search import slice_content from braindb.services.embedding_service import get_embedding_service from braindb.services.keyword_service import ensure_keyword_entities, link_entity_to_keywords, sync_keywords_for_entity @@ -70,13 +72,23 @@ class IngestRequest(BaseModel): re.always_on AS always_on, re.category AS category, re.priority AS priority, - re.is_active AS is_active + re.is_active AS is_active, + -- wiki + we.canonical_name AS canonical_name, + we.disambiguation AS disambiguation, + we.language AS wiki_language, + we.member_keyword_ids::text[] AS member_keyword_ids, + we.revision AS revision, + we.last_synthesised_at AS last_synthesised_at, + we.retired_at AS retired_at, + we.redirect_to AS redirect_to FROM entities e LEFT JOIN thoughts_ext te ON te.entity_id = e.id AND e.entity_type = 'thought' LEFT JOIN facts_ext fe ON fe.entity_id = e.id AND e.entity_type = 'fact' LEFT JOIN sources_ext se ON se.entity_id = e.id AND e.entity_type = 'source' LEFT JOIN datasources_ext de ON de.entity_id = e.id AND e.entity_type = 'datasource' LEFT JOIN rules_ext re ON re.entity_id = e.id AND e.entity_type = 'rule' + LEFT JOIN wikis_ext we ON we.entity_id = e.id AND e.entity_type = 'wiki' WHERE e.id = %s """ @@ -116,6 +128,17 @@ def _flatten(row: dict) -> dict: base.update(file_path=row["file_path"], url=row["ds_url"], content_hash=row["content_hash"], word_count=row["word_count"], language=row["language"]) elif etype == "rule": base.update(always_on=row["always_on"], category=row["category"], priority=row["priority"], is_active=row["is_active"]) + elif etype == "wiki": + base.update( + canonical_name=row["canonical_name"], + disambiguation=row["disambiguation"], + language=row["wiki_language"], + member_keyword_ids=row["member_keyword_ids"] or [], + revision=row["revision"], + last_synthesised_at=row["last_synthesised_at"], + retired_at=row["retired_at"], + redirect_to=row["redirect_to"], + ) return base @@ -282,14 +305,43 @@ def create_rule(body: RuleCreate): return _flatten(_fetch(conn, eid)) +@router.post("/wikis", response_model=WikiRead, status_code=201) +def create_wiki(body: WikiCreate): + with get_conn() as conn: + eid = _insert_entity(conn, "wiki", body) + with conn.cursor() as cur: + cur.execute( + """INSERT INTO wikis_ext + (entity_id, canonical_name, disambiguation, language, member_keyword_ids) + VALUES (%s, %s, %s, %s, %s::uuid[])""", + (str(eid), body.canonical_name, body.disambiguation, body.language, + [str(k) for k in body.member_keyword_ids]), + ) + return _flatten(_fetch(conn, eid)) + + # ------------------------------------------------------------------ # # READ # # ------------------------------------------------------------------ # @router.get("/{entity_id}") -def get_entity(entity_id: UUID): +def get_entity( + entity_id: UUID, + offset: int = Query(default=0, ge=0), + limit: int | None = Query(default=None, ge=1), +): + """Full single-entity read. Pass offset/limit to page a large `content` + without flooding the caller — response then includes `content_meta` + {total_chars, offset, returned, next_offset}. Default (no offset/limit) + returns the full body, unchanged.""" with get_conn() as conn: - return _flatten(_or_404(_fetch(conn, entity_id))) + ent = _flatten(_or_404(_fetch(conn, entity_id))) + if offset == 0 and limit is None: + return ent + chunk, meta = slice_content(ent.get("content"), offset, limit) + ent["content"] = chunk + ent["content_meta"] = meta + return ent # ------------------------------------------------------------------ # @@ -424,6 +476,39 @@ def update_rule(entity_id: UUID, body: RuleUpdate): return _flatten(_fetch(conn, entity_id)) +@router.patch("/wikis/{entity_id}", response_model=WikiRead) +def update_wiki(entity_id: UUID, body: WikiUpdate): + with get_conn() as conn: + row = _or_404(_fetch(conn, entity_id)) + if row["entity_type"] != "wiki": + raise HTTPException(400, "Entity is not a wiki") + data = body.model_dump(exclude_unset=True) + _update_base(conn, entity_id, data) + # wikis_ext: UUID / UUID[] fields need explicit handling, so do not + # route through the generic _update_ext. + ext_fields = ("canonical_name", "disambiguation", "language", "member_keyword_ids", + "revision", "last_synthesised_at", "retired_at", "redirect_to") + ext = {k: v for k, v in data.items() if k in ext_fields and v is not None} + if ext: + assignments, values = [], [] + for k, v in ext.items(): + if k == "member_keyword_ids": + assignments.append(f"{k} = %s::uuid[]") + values.append([str(x) for x in v]) + elif k == "redirect_to": + assignments.append(f"{k} = %s") + values.append(str(v)) + else: + assignments.append(f"{k} = %s") + values.append(v) + with conn.cursor() as cur: + cur.execute( + f"UPDATE wikis_ext SET {', '.join(assignments)} WHERE entity_id = %s", + values + [str(entity_id)], + ) + return _flatten(_fetch(conn, entity_id)) + + # ------------------------------------------------------------------ # # DELETE # # ------------------------------------------------------------------ # diff --git a/braindb/routers/wiki.py b/braindb/routers/wiki.py new file mode 100644 index 0000000..0de6512 --- /dev/null +++ b/braindb/routers/wiki.py @@ -0,0 +1,538 @@ +""" +Wiki pipeline endpoints: cron / maintain / write / jobs. + +Stage 1 is manual (no scheduler) — these endpoints are driven by hand or by +the Stage-2 `wiki_scheduler` sidecar. `/cron` and `/jobs` are pure SQL and +non-destructive; `/maintain` and `/write` (later steps) drive the existing +agent endpoint. +""" +import logging +from pathlib import Path + +from fastapi import APIRouter, Query + +from braindb.agent.agent import run_typed, get_maintainer_agent, get_writer_agent +from braindb.agent.run_state import install_handoff_slot, release_handoff_slot +from braindb.agent.schemas import MaintainerDecision, WikiWriteResult +from braindb.config import settings +from braindb.db import get_conn +from braindb.services.activity_log import log_activity +from braindb.services import wiki_jobs + +logger = logging.getLogger(__name__) + +router = APIRouter(prefix="/api/v1/wiki", tags=["wiki"]) + +_PROMPTS = Path(__file__).parent.parent / "agent" / "prompts" +_MAINTAINER_PROMPT = (_PROMPTS / "wiki_maintainer_prompt.md").read_text(encoding="utf-8") +_WRITER_PROMPT = (_PROMPTS / "wiki_writer_prompt.md").read_text(encoding="utf-8") + + +@router.post("/cron") +def wiki_cron(): + """Read-only orphan scan; enqueues one `triage` job per orphan. Idempotent.""" + with get_conn() as conn: + result = wiki_jobs.run_cron(conn) + log_activity(conn, "wiki_cron", None, None, details=result) + return result + + +@router.post("/maintain") +async def wiki_maintain(): + """ + Process EXACTLY ONE triage case (C1). Claims one pending triage job, + asks the existing agent to decide attach/create/consolidate/skip for that + single orphan, persists the resulting suggestion job, closes the triage. + """ + # 1. Claim one case + staleness guard, atomically (one transaction). + with get_conn() as conn: + job = wiki_jobs.claim_one_triage(conn) + if not job: + return {"claimed": 0, "message": "no pending triage jobs"} + orphan_id = job["entity_ids"][0] + job_id = str(job["id"]) + batch_id = str(job["batch_id"]) if job["batch_id"] else None + orphan = wiki_jobs.fetch_entity_brief(conn, orphan_id) + + if not orphan: + wiki_jobs.finish_job(conn, job_id, "failed", "orphan entity not found") + return {"claimed": 1, "job_id": job_id, "result": "failed", + "reason": "orphan missing"} + + # Stale-skip: a prior writer run may have already absorbed/linked this + # entity (or it's already in an active suggestion). If so, close the + # triage with NO LLM call — the writer's broad research retired it. + if not wiki_jobs.is_orphan(conn, orphan_id, exclude_triage_job_id=job_id): + wiki_jobs.finish_job(conn, job_id, "done", + "already covered — absorbed by a wiki") + return {"claimed": 1, "job_id": job_id, "result": "skipped_stale"} + + # Catalog of existing wikis the model will reference BY NUMBER (never + # by uuid). This in-request list IS the numbering used to resolve the + # model's chosen number(s) back to ids below. + cat = wiki_jobs.list_active_wikis(conn) + + # 2. One agent call. The prompt directs it to RESEARCH the neighbourhood + # with its own tools (recall_memory / view_tree / delegate_to_subagent) + # before deciding — we give the seed, the LLM gathers the context. + # Generous turns so it can actually investigate / delegate. + catalog_txt = ( + "\n".join(f"{i}. {w['canonical_name']}" for i, w in enumerate(cat, 1)) + or "(no existing wikis yet — attach/consolidate are impossible; " + "use create/skip/ambiguous)" + ) + prompt = _MAINTAINER_PROMPT.format( + entity_id=orphan_id, + entity_type=orphan["entity_type"], + keywords=orphan.get("keywords") or [], + summary=orphan.get("summary"), + content=(orphan.get("content") or "")[:4000], + wiki_catalog=catalog_txt, + ) + # `run_typed` returns a SDK-validated MaintainerDecision, or raises if + # the model never submitted (e.g. max_turns hit) — that error path + # below treats it like any other agent failure (release + log + 5xx). + try: + res: MaintainerDecision = await run_typed( + prompt, get_maintainer_agent(), MaintainerDecision, max_turns=30 + ) + except Exception as e: + logger.exception("maintainer agent failed") + with get_conn() as conn: + wiki_jobs.finish_job(conn, job_id, "failed", f"agent error: {e}"[:500]) + return {"claimed": 1, "job_id": job_id, "result": "failed", "reason": str(e)} + + # Schema-validated; expose as a dict so the action handlers below are + # unchanged. + decision = res.model_dump() + action = decision.get("action") + rationale = decision.get("rationale") + + # 3. Persist the suggestion + close the triage, in one transaction. + with get_conn() as conn: + try: + if action in ("skip", "ambiguous"): + # 'ambiguous' = the data cannot disambiguate identity/scope; + # the LLM correctly refuses to mint a confident page. Treated + # as a deliberate skip (self-clears via run_cron). + wiki_jobs.finish_job(conn, job_id, "rejected", rationale) + outcome = {"action": action} + + elif action == "attach": + no = decision.get("target_wiki_no") + target = (cat[no - 1]["id"] + if isinstance(no, int) and 1 <= no <= len(cat) + else None) + if not target or not _is_wiki(conn, target): + wiki_jobs.finish_job( + conn, job_id, "failed", + f"attach: target_wiki_no {no!r} not a valid catalog number (1..{len(cat)})") + outcome = {"action": "attach", "error": "invalid target_wiki_no"} + else: + key = wiki_jobs.suggestion_dedupe_key("attach", target, [orphan_id], []) + sid = wiki_jobs.insert_suggestion( + conn, job_type="attach", target_wiki_id=target, + entity_ids=[orphan_id], dedupe_key=key, rationale=rationale, + proposed_name=None, batch_id=batch_id) + wiki_jobs.finish_job(conn, job_id, "done", rationale) + outcome = {"action": "attach", "suggestion_id": sid, "target_wiki_id": target} + + elif action == "create": + name = decision.get("proposed_name") + if not name: + wiki_jobs.finish_job(conn, job_id, "failed", "create missing proposed_name") + outcome = {"action": "create", "error": "missing proposed_name"} + else: + key = wiki_jobs.suggestion_dedupe_key("create", None, [orphan_id], []) + sid = wiki_jobs.insert_suggestion( + conn, job_type="create", target_wiki_id=None, + entity_ids=[orphan_id], dedupe_key=key, rationale=rationale, + proposed_name=name, batch_id=batch_id) + wiki_jobs.finish_job(conn, job_id, "done", rationale) + outcome = {"action": "create", "suggestion_id": sid, "proposed_name": name} + + elif action == "consolidate": + nos = decision.get("consolidate_nos") or [] + ids = [cat[n - 1]["id"] for n in nos + if isinstance(n, int) and 1 <= n <= len(cat)] + wiki_ids = list(dict.fromkeys(ids)) # dedupe, keep order + if len(wiki_ids) < 2: + wiki_jobs.finish_job( + conn, job_id, "failed", + f"consolidate: need >=2 valid catalog numbers, got {nos!r} (1..{len(cat)})") + outcome = {"action": "consolidate", "error": "need >=2 valid catalog numbers"} + else: + key = wiki_jobs.suggestion_dedupe_key("consolidate", None, [], wiki_ids) + sid = wiki_jobs.insert_suggestion( + conn, job_type="consolidate", target_wiki_id=None, + entity_ids=wiki_ids, dedupe_key=key, rationale=rationale, + proposed_name=None, batch_id=batch_id) + # The orphan itself is still unconnected; closing 'done' + # lets the next cron re-triage it after the merge. + wiki_jobs.finish_job(conn, job_id, "done", rationale) + outcome = {"action": "consolidate", "suggestion_id": sid, "wiki_ids": wiki_ids} + + else: + wiki_jobs.finish_job(conn, job_id, "failed", f"unknown action {action!r}") + outcome = {"action": action, "error": "unknown action"} + + log_activity(conn, "wiki_maintain", orphan["entity_type"], orphan_id, + details={"job_id": job_id, **outcome}) + except Exception as e: + logger.exception("maintainer persistence failed") + raise + + return {"claimed": 1, "job_id": job_id, "result": outcome} + + +def _members_block(members: list[dict]) -> str: + if not members: + return "(none)" + out = [] + for m in members: + out.append( + f"- id: {m['id']}\n type: {m['entity_type']}\n" + f" keywords: {m.get('keywords') or []}\n" + f" content: {(m.get('content') or '')[:1200]}" + ) + return "\n".join(out) + + +# Above this character count, an attach's `%%CURRENT_BODY%%` is replaced +# by a stub pointing at the section-edit tools. Keeps the writer's +# INITIAL prompt small so it never bumps into the model window before +# it can run a single tool — the section tools (Phase 1) are designed +# exactly for navigating a body without inlining it. +_INLINE_BODY_MAX_CHARS = 4000 + + +def _body_block_or_stub(mode: str, wiki_id: str | None, old_body: str) -> str: + """For attach mode with a body too large to safely inline, return a + stub directing the agent to use the section tools instead. Small + bodies and other modes inline as before.""" + if not old_body: + return "(none — create mode)" + if mode == "attach" and wiki_id and len(old_body) > _INLINE_BODY_MAX_CHARS: + return ( + f"[BODY OMITTED — {len(old_body)} chars, too large to inline.\n" + f"Use the section tools to navigate without consuming context:\n" + f" - read_wiki_outline(\"{wiki_id}\") — section list + sizes + revision\n" + f" - read_wiki_section(\"{wiki_id}\", \"\") — one section\n" + f" - edit_wiki_section(...) per section, validate_wiki, then\n" + f" final_answer(mode=\"attach\", body=\"\") — router persists via\n" + f" section edits and skips the full-body write.]" + ) + return old_body + + +@router.post("/write") +async def wiki_write(): + """ + Write/update ONE wiki (one target per call). The LLM authors the entire + body and may freely revise summary/disambiguation/scope/any section. No + content gate, no manifest, no code-built ledger. The only guarantees are + process/bookkeeping: the prior version is snapshotted (reversible) and + `summarises` relations are reconciled *additively* from the LLM's inline + refs. The LLM researches with its own tools before writing. + """ + # 1. Pick + claim a bucket. + with get_conn() as conn: + bucket = wiki_jobs.next_write_bucket(conn) + if not bucket: + return {"written": 0, "message": "no pending create/attach jobs"} + mode = bucket["mode"] + jobs = bucket["jobs"] + job_ids = [str(j["id"]) for j in jobs] + lock_key = bucket["target_wiki_id"] or f"create:{job_ids[0]}" + if not wiki_jobs.try_wiki_lock(conn, lock_key): + return {"written": 0, "message": "target locked by another writer; retry later"} + claimed = wiki_jobs.claim_jobs(conn, job_ids) + if not claimed: + return {"written": 0, "message": "jobs no longer claimable"} + + member_ids: list[str] = [] + for j in jobs: + member_ids.extend(j["entity_ids"]) + dupes: list[dict] = [] + if mode == "attach": + members = wiki_jobs.fetch_members(conn, member_ids) + wiki = wiki_jobs.fetch_wiki(conn, bucket["target_wiki_id"]) + if not wiki: + wiki_jobs.finish_jobs(conn, job_ids, "failed", "target wiki missing") + return {"written": 0, "result": "failed", "reason": "target wiki missing"} + canonical = wiki["canonical_name"] + old_body = wiki["content"] or "" + elif mode == "consolidate": + members = [] + dupes = wiki_jobs.fetch_wikis_for_merge(conn, bucket["wiki_ids"]) + if len(dupes) < 2: + wiki_jobs.finish_jobs(conn, job_ids, "failed", + "fewer than 2 live wikis to consolidate") + return {"written": 0, "result": "failed", "reason": "nothing to merge"} + canonical = "(decide among duplicates)" + wiki = None + old_body = "\n\n".join(d["content"] or "" for d in dupes) + else: # create + members = wiki_jobs.fetch_members(conn, member_ids) + canonical = bucket["proposed_name"] or "Untitled" + wiki = None + old_body = "" + batch_id = str(jobs[0].get("batch_id")) if jobs[0].get("batch_id") else None + + def _dupes_block(ds: list[dict]) -> str: + if not ds: + return "(n/a)" + # Numbered; the writer picks the survivor by NUMBER (canonical_no), + # never by id. This order IS the numbering resolved below. + return "\n".join( + f"{i}. {d['canonical_name']} " + f"(importance: {d['importance']} revision: {d['revision']})\n" + f" body:\n{(d['content'] or '')[:3000]}" + for i, d in enumerate(ds, 1) + ) + + # 2. One focused agent call. + prompt = ( + _WRITER_PROMPT + .replace("%%MODE%%", mode) + .replace("%%CANONICAL%%", canonical) + .replace("%%WIKI_ID%%", bucket["target_wiki_id"] or "(assigned after write)") + .replace("%%MEMBERS%%", _members_block(members)) + .replace("%%CURRENT_BODY%%", _body_block_or_stub(mode, bucket.get("target_wiki_id"), old_body)) + .replace("%%DUPLICATES%%", _dupes_block(dupes)) + ) + # Capture pre-run revision on the target wiki for `attach` mode so we + # can detect whether the writer used the section-edit tools (each + # bumps `wikis_ext.revision` directly). The writer may then submit an + # empty `body` — section edits are the authoritative persistence + # path in that case. `create`/`consolidate` modes don't have a + # pre-determined target, so empty body is rejected there. + pre_revision: int | None = None + if mode == "attach" and bucket.get("target_wiki_id"): + with get_conn() as conn: + with conn.cursor() as cur: + cur.execute( + "SELECT revision FROM wikis_ext WHERE entity_id = %s::uuid", + (bucket["target_wiki_id"],), + ) + row = cur.fetchone() + if row: + pre_revision = row[0] + + # Generous turns so the writer can recall_memory / view_tree / delegate a + # subagent to research and verify before writing. + # `run_typed` returns a SDK-validated WikiWriteResult, or raises if the + # model never submitted — handled below like any agent failure + # (release + log + 5xx). The only extra guard is "non-empty body OR + # section edits happened"; everything else is the model's job (and + # validated by Pydantic). + # + # Context-handoff loop: the writer may end early via + # `handoff_to_successor` when its context approaches the limit (see + # `braindb/agent/hooks.py` token-budget watch + `tools.py` handoff + # tool). We install a per-run handoff slot, run the agent, and if + # the slot was filled we spawn a successor agent — same prompt, same + # tools, fresh context — seeded with the previous agent's brief. + # Bounded by `agent_writer_handoff_max_depth` so a misbehaving model + # cannot recurse forever. + handoff_slot, handoff_token = install_handoff_slot() + try: + try: + res: WikiWriteResult = await run_typed( + prompt, get_writer_agent(), WikiWriteResult, max_turns=30, + token_budget=settings.agent_writer_handoff_token_budget, + ) + depth = 0 + max_depth = settings.agent_writer_handoff_max_depth + while handoff_slot.captured and depth < max_depth: + depth += 1 + seed = ( + "Continuing from a previous agent run that ended early " + "via `handoff_to_successor` because its context was " + "filling up. You have the SAME prompt, the SAME tools, " + "and a fresh context window. Resume from this state.\n\n" + "PROGRESS SO FAR (from the previous agent):\n" + f"{handoff_slot.progress_summary}\n\n" + "REMAINING WORK:\n" + f"{handoff_slot.remaining_work}\n\n" + "Pick up from here. Call `final_answer` when done " + "(body=\"\" if you persisted via section-edit tools, " + "or the full body otherwise). If YOUR context also " + "fills up before you finish, call `handoff_to_successor` " + "again with an updated brief — the same successor " + "mechanism will continue." + ) + handoff_slot.captured = False + handoff_slot.progress_summary = "" + handoff_slot.remaining_work = "" + logger.info( + "writer handoff: spawning successor #%d/%d (mode=%s, jobs=%s)", + depth, max_depth, mode, job_ids, + ) + res = await run_typed( + seed, get_writer_agent(), WikiWriteResult, max_turns=30, + token_budget=settings.agent_writer_handoff_token_budget, + ) + if handoff_slot.captured: + # Depth cap hit AND last run still asked for handoff — + # treat as a failure (the model isn't converging). + logger.warning( + "writer handoff: depth cap %d hit; treating as failure", + max_depth, + ) + with get_conn() as conn: + disp = wiki_jobs.release_or_fail_jobs( + conn, job_ids, + f"handoff depth cap {max_depth} exhausted " + f"without final_answer", + ) + return {"written": 0, "result": disp, + "reason": "handoff depth exhausted"} + except Exception as e: + logger.exception("writer agent failed") + with get_conn() as conn: + disp = wiki_jobs.release_or_fail_jobs(conn, job_ids, f"agent error: {e}") + return {"written": 0, "result": disp, "reason": str(e)} + finally: + release_handoff_slot(handoff_token) + + used_section_edits = False + if not (res.body or "").strip(): + # Empty body — only valid in attach mode if section edits bumped + # the revision during the run. Otherwise the agent did nothing + # persistable and we fail the jobs. + if mode != "attach" or pre_revision is None: + with get_conn() as conn: + disp = wiki_jobs.release_or_fail_jobs( + conn, job_ids, + f"empty body returned in {mode} mode: " + f"{res.model_dump_json()[:300]}", + ) + return {"written": 0, "result": disp, "reason": "no body returned"} + with get_conn() as conn: + with conn.cursor() as cur: + cur.execute( + """SELECT e.content, w.revision + FROM entities e JOIN wikis_ext w ON w.entity_id = e.id + WHERE e.id = %s::uuid""", + (bucket["target_wiki_id"],), + ) + row = cur.fetchone() + if not row: + with get_conn() as conn: + disp = wiki_jobs.release_or_fail_jobs( + conn, job_ids, + "empty body AND wiki vanished", + ) + return {"written": 0, "result": disp, "reason": "wiki missing"} + if row[1] == pre_revision: + # No section edits happened. This is legitimate ONLY if every + # assigned member is already cited in the body — then there is + # nothing to write and we just need to run reconcile so the + # `summarises` relations catch up. If any member is missing + # from the body, the writer skipped real work — fail it. + body_now = row[0] or "" + cited = wiki_jobs.parse_refs(body_now) # lower-cased set + missing = [m for m in member_ids if m.lower() not in cited] + if missing: + with get_conn() as conn: + disp = wiki_jobs.release_or_fail_jobs( + conn, job_ids, + f"empty body AND no section edits AND " + f"{len(missing)} member(s) not yet cited in body", + ) + return {"written": 0, "result": disp, + "reason": "members un-cited"} + # All members cited — close the no-op cleanly and reconcile. + with get_conn() as conn: + rel = wiki_jobs.reconcile_summarises_additive( + conn, bucket["target_wiki_id"], body_now) + wiki_jobs.finish_jobs(conn, job_ids, "done") + log_activity(conn, "wiki_write", "wiki", + bucket["target_wiki_id"], details={ + "mode": mode, "no_op": True, + "revision": pre_revision, + "members": len(member_ids), **rel, + }) + logger.info( + "writer no-op accepted: pre_rev=%s, all %d members already " + "cited; reconcile=%s", + pre_revision, len(member_ids), rel, + ) + return {"written": 0, "wiki_id": bucket["target_wiki_id"], + "mode": mode, "revision": pre_revision, + "jobs": job_ids, "no_op": True, **rel} + new_body = row[0] + used_section_edits = True + logger.info( + "writer used section-edit path: pre_rev=%s post_rev=%s body=%dch", + pre_revision, row[1], len(new_body), + ) + else: + new_body = res.body + + # 3. Persist (one transaction). No content gate — the LLM's body is + # authoritative; we only snapshot (reversible) and reconcile additively. + with get_conn() as conn: + summary, disambig = wiki_jobs.extract_summary_disambig(new_body) + kw = wiki_jobs.keywords_from_meta(new_body) + retired: list[str] = [] + if mode == "create": + wiki_id = wiki_jobs.create_wiki_entity( + conn, canonical, new_body, summary, disambig, member_ids, + keywords=kw) + revision = 1 + elif mode == "consolidate": + no = res.canonical_no + if not (isinstance(no, int) and 1 <= no <= len(dupes)): + disp = wiki_jobs.release_or_fail_jobs( + conn, job_ids, + f"canonical_no {no!r} not a valid duplicates number (1..{len(dupes)})") + return {"written": 0, "result": disp, + "reason": "invalid canonical_no"} + canonical_id = dupes[no - 1]["id"] + wiki_id = canonical_id + for d in dupes: + wiki_jobs.snapshot_revision( + conn, d["id"], d["content"] or "", + wiki_jobs.parse_refs(d["content"] or ""), d["revision"]) + revision = wiki_jobs.finalize_wiki_write( + conn, wiki_id, new_body, summary, disambig, member_ids) + for d in dupes: + if d["id"] != canonical_id: + wiki_jobs.soft_retire_wiki(conn, d["id"], canonical_id, None) + retired.append(d["id"]) + else: # attach + wiki_id = bucket["target_wiki_id"] + wiki_jobs.snapshot_revision( + conn, wiki_id, old_body, wiki_jobs.parse_refs(old_body), + wiki["revision"]) + revision = wiki_jobs.finalize_wiki_write( + conn, wiki_id, new_body, summary, disambig, member_ids) + + rel = wiki_jobs.reconcile_summarises_additive(conn, wiki_id, new_body) + + wiki_jobs.finish_jobs(conn, job_ids, "done") + log_activity(conn, "wiki_write", "wiki", wiki_id, details={ + "mode": mode, "revision": revision, "jobs": job_ids, + "members": len(member_ids), "retired": retired, **rel, + }) + + return {"written": 1, "wiki_id": wiki_id, "mode": mode, + "revision": revision, "jobs": job_ids, "retired": retired, **rel} + + +def _is_wiki(conn, entity_id: str) -> bool: + with conn.cursor() as cur: + cur.execute("SELECT 1 FROM entities WHERE id = %s AND entity_type = 'wiki'", (str(entity_id),)) + return cur.fetchone() is not None + + +@router.get("/jobs") +def wiki_jobs_list( + status: str | None = Query(default=None), + job_type: str | None = Query(default=None), + limit: int = Query(default=50, ge=1, le=500), +): + with get_conn() as conn: + return wiki_jobs.list_jobs(conn, status, job_type, limit) diff --git a/braindb/schemas/entities.py b/braindb/schemas/entities.py index ddac264..7389d63 100644 --- a/braindb/schemas/entities.py +++ b/braindb/schemas/entities.py @@ -209,8 +209,50 @@ class RuleUpdate(BaseModel): is_active: bool | None = None +# ------------------------------------------------------------------ # +# WIKI # +# ------------------------------------------------------------------ # + +class WikiCreate(EntityBase): + canonical_name: str = Field(..., min_length=1, max_length=500) + disambiguation: str | None = None + language: str = "en" + member_keyword_ids: list[UUID] = Field(default_factory=list) + + +class WikiRead(EntityRead): + entity_type: Literal["wiki"] = "wiki" + canonical_name: str + disambiguation: str | None + language: str + member_keyword_ids: list[UUID] = Field(default_factory=list) + revision: int + last_synthesised_at: datetime | None = None + retired_at: datetime | None = None + redirect_to: UUID | None = None + + +class WikiUpdate(BaseModel): + title: str | None = None + content: str | None = None + summary: str | None = None + keywords: list[str] | None = None + importance: float | None = Field(default=None, ge=0.0, le=1.0) + source: str | None = None + notes: str | None = None + metadata: dict[str, Any] | None = None + canonical_name: str | None = Field(default=None, min_length=1, max_length=500) + disambiguation: str | None = None + language: str | None = None + member_keyword_ids: list[UUID] | None = None + revision: int | None = None + last_synthesised_at: datetime | None = None + retired_at: datetime | None = None + redirect_to: UUID | None = None + + # ------------------------------------------------------------------ # # Generic entity read (union) used in list endpoints # # ------------------------------------------------------------------ # -AnyEntityRead = ThoughtRead | FactRead | SourceRead | DatasourceRead | RuleRead +AnyEntityRead = ThoughtRead | FactRead | SourceRead | DatasourceRead | RuleRead | WikiRead diff --git a/braindb/schemas/relations.py b/braindb/schemas/relations.py index d3501f5..68c8dbd 100644 --- a/braindb/schemas/relations.py +++ b/braindb/schemas/relations.py @@ -15,6 +15,10 @@ "is_example_of", "challenges", "tagged_with", + "summarises", # wiki --summarises--> entity it is built from + "not_duplicate", # two wikis judged distinct (self-clears the dedup pass) + "duplicate_of", # retired wiki --duplicate_of--> canonical wiki (post-merge) + "consolidated_into", # provenance of an LLM-performed consolidation ] diff --git a/braindb/services/context.py b/braindb/services/context.py index 52b70d6..383ef88 100644 --- a/braindb/services/context.py +++ b/braindb/services/context.py @@ -17,8 +17,12 @@ from braindb.schemas.search import ContextRequest, ContextResponse, SearchResultItem from braindb.services.embedding_service import get_embedding_service from braindb.services.graph import graph_expand -from braindb.services.keyword_service import find_entities_for_keywords, find_similar_keywords -from braindb.services.search import fuzzy_search +from braindb.services.keyword_service import ( + find_entities_for_keywords, + find_fuzzy_keywords, + find_similar_keywords, +) +from braindb.services.search import fuzzy_search, preview DECAY_RATES = { "thought": settings.decay_rate_thought, @@ -26,6 +30,7 @@ "source": settings.decay_rate_source, "datasource": settings.decay_rate_datasource, "rule": settings.decay_rate_rule, + "wiki": settings.decay_rate_wiki, } @@ -49,6 +54,7 @@ def effective_importance(importance: float, created_at: datetime, access_count: "source": ("sources_ext", "entity_id, url, domain, http_status, last_checked_at"), "datasource": ("datasources_ext", "entity_id, file_path, url, content_hash, word_count, language"), "rule": ("rules_ext", "entity_id, always_on, category, priority, is_active"), + "wiki": ("wikis_ext", "entity_id, canonical_name, disambiguation, language, member_keyword_ids::text[] AS member_keyword_ids, revision, last_synthesised_at, retired_at, redirect_to"), } @@ -111,7 +117,7 @@ def _to_item(row: dict, search_score: float, depth: int, relevance: float, ext: id=row["id"], entity_type=row["entity_type"], title=row.get("title"), - content=row["content"], + content=preview(row.get("content"), row.get("id")), summary=row.get("summary"), keywords=row.get("keywords") or [], importance=row["importance"], @@ -130,6 +136,116 @@ def _to_item(row: dict, search_score: float, depth: int, relevance: float, ext: ) +# ------------------------------------------------------------------ # +# Two-level diversity quota (per-search-term + per-keyword) # +# ------------------------------------------------------------------ # + +def _apply_two_level_quota( + items: list, + dominant_kw_by_id: dict[str, str], + per_query_top_ids: list[list[str]], + max_results: int, + per_query_share: float, + halving: float, +) -> list: + """Re-rank `items` (sorted by `final_rank` desc) under two + complementary diversity quotas. Both run in ONE pass so they can + never conflict. + + Level 1 — per-search-term (the user's outer quota): + Each query in `per_query_top_ids` gets a reserved share of the + output. Walking the per-query top-K lists first guarantees each + angle of the multi-query recall surfaces something in the + result, even if its absolute scores would be outranked globally. + + Level 2 — per-keyword (the inner quota): + Walking the remaining items in `final_rank`-desc order, each + new dominant matched keyword gets a halving slot allowance + (`ceil(max_results × halving^n)`, floor 1). Stops a single + popular keyword (e.g. `user-profile` tagging 100 biographical + facts) from monopolising the open portion of the result. + + HOW THE TWO LEVELS COEXIST WITHOUT CONFLICT (this is the crucial + bit, please read before changing this function): + + Both levels share ONE counter dict (`seen`: kw_id → remaining). + Level 1 places reserved items first and decrements their + dominant keyword's allowance. Level 2 then walks the open items + and respects what L1 already consumed. So: + + - A reserved item is added unconditionally (L1 wins). Its + keyword's L2 quota shrinks accordingly — no double spending + in the open phase. + - If a popular keyword's allowance is exhausted purely by L1 + reservations, L2 will skip further entities tagged dominantly + with it. That's the intended hard cap. + - Items without a dominant keyword (graph-expansion finds, the + discoverability backup) pass through both phases freely; + they're not counted against any keyword's allowance. + + `per_query_share`=0 disables L1 (only L2 runs). `halving`>=1.0 + disables L2 (only L1 + raw top-N for the rest). Both at extremes + = raw top-N. + """ + seen: dict[str, int] = {} # kw_id → remaining slots (SHARED across L1 + L2) + n_new = 0 # number of distinct keywords met so far (drives the halving sequence) + taken: set[str] = set() # entity ids already placed (dedup across L1's per-query lists) + out: list = [] + + def _consume(item) -> bool: + """Try to place `item` in `out`, respecting the per-keyword quota. + Returns True if placed, False if blocked by L2.""" + nonlocal n_new + if str(item.id) in taken: + return False + kw = dominant_kw_by_id.get(str(item.id)) + if kw is None: + # No keyword to gate against (graph-expansion / discovery + # fallback) — let it through. + taken.add(str(item.id)) + out.append(item) + return True + if halving < 1.0: + if kw not in seen: + # Lazy-init this keyword's allowance using its position + # in the geometric-decay sequence. + seen[kw] = max(1, math.ceil(max_results * (halving ** n_new))) + n_new += 1 + if seen[kw] <= 0: + return False + seen[kw] -= 1 + taken.add(str(item.id)) + out.append(item) + return True + + # Map id → item so we can walk per-query lists in O(1). + by_id: dict[str, object] = {str(it.id): it for it in items} + + # ---- LEVEL 1: per-search-term reservation phase -------------------- + # Walk each query's own top-K and place reserved items first. + # `per_query_top_ids[q_index]` is already sorted by THIS query's + # combined score, so we get the best-for-this-angle items first. + if per_query_share > 0: + for q_top in per_query_top_ids: + for eid in q_top: + item = by_id.get(eid) + if item is None: + continue + _consume(item) + if len(out) >= max_results: + return out + + # ---- LEVEL 2: open phase with per-keyword quota -------------------- + # Walk remaining items in global final_rank-desc order. `_consume` + # respects whatever L1 already used up in the `seen` counter, so + # a keyword that filled its quota via L1 is correctly blocked here. + for item in items: + if len(out) >= max_results: + break + _consume(item) + return out + + # ------------------------------------------------------------------ # # Main context assembly # # ------------------------------------------------------------------ # @@ -139,47 +255,102 @@ def assemble_context(conn, req: ContextRequest) -> ContextResponse: query_list = req.queries if req.queries else [req.query] # ------------------------------------------------------------------ # - # 1. TEXT SEARCH (existing) — fuzzy + fulltext per query # + # 1. TEXT SEARCH (keyword-mediated) — fuzzy on KEYWORD entities, # + # then fan out via tagged_with. Symmetric to the embedding # + # pathway below: both produce a per-entity score equal to the # + # best match between the query and the entity's tagged keywords. # + # This avoids the pg_trgm dilution that previously hit any short # + # query against a long entity body — keywords are short, so the # + # trigram intersection is meaningful, not diluted. # # ------------------------------------------------------------------ # - text_scores: dict = {} # entity_id → best text score + text_scores: dict = {} # entity_id → best keyword-fuzzy similarity (max across queries) + text_dom_kw: dict = {} # entity_id → keyword_id that yielded the text_scores max + text_scores_by_q: list = [] # per-query: list of {entity_id → best_sim for THIS query} seed_rows_by_id: dict = {} # entity_id → row data + fuzzy_rows: dict = {} # entity_id → row data (entities found only via fuzzy-keyword) for q in query_list: - rows = fuzzy_search(conn, q, req.entity_types, req.min_importance, limit=max(req.max_results, 20)) + per_q_scores: dict = {} # this query's text scores only — feeds Level-1 quota + fuzzy_kw = find_fuzzy_keywords( + conn, q, limit=settings.scoring_pool_fuzzy, + ) + if fuzzy_kw: + kw_sim = {str(kw["id"]): kw["similarity"] for kw in fuzzy_kw} + entities = find_entities_for_keywords(conn, list(kw_sim.keys())) + for ent in entities: + eid = ent["id"] + matched_ids = [str(mid) for mid in (ent.get("matched_keyword_ids") or [])] + if matched_ids: + # Pick the matched keyword with the strongest similarity for this entity + best_kw_id = max(matched_ids, key=lambda m: kw_sim.get(m, 0)) + best_sim = kw_sim.get(best_kw_id, 0) + per_q_scores[str(eid)] = best_sim + if eid not in text_scores or best_sim > text_scores[eid]: + text_scores[eid] = best_sim + text_dom_kw[eid] = best_kw_id + if eid not in seed_rows_by_id: + fuzzy_rows[eid] = ent + text_scores_by_q.append(per_q_scores) + + # Discoverability backup — entities whose content matches the query + # directly but aren't tagged with a matching keyword. Heavy discount + # (`DISCOVERY_DISCOUNT`) keeps them weakly-ranked. Pure fallback: only + # set text_scores for an entity if the keyword-mediated path didn't + # already cover it (never override a real keyword match). + DISCOVERY_DISCOUNT = 0.2 + for q in query_list: + rows = fuzzy_search( + conn, q, req.entity_types, req.min_importance, + limit=settings.scoring_pool_fuzzy, + ) for r in rows: eid = r["id"] - score = r["score"] - if eid not in text_scores or score > text_scores[eid]: - text_scores[eid] = score - seed_rows_by_id[eid] = r + if eid in text_scores: + continue # keyword path already scored this entity; do not override + text_scores[eid] = r["score"] * DISCOVERY_DISCOUNT + if eid not in seed_rows_by_id and eid not in fuzzy_rows: + fuzzy_rows[eid] = r # ------------------------------------------------------------------ # # 2. KEYWORD EMBEDDING SEARCH (new) — semantic via keyword vectors # # ------------------------------------------------------------------ # - embedding_scores: dict = {} # entity_id → best keyword similarity + embedding_scores: dict = {} # entity_id → best keyword similarity (max across queries) + embedding_dom_kw: dict = {} # entity_id → keyword_id that yielded the embedding_scores max + embedding_scores_by_q: list = [] # per-query embedding scores — feeds Level-1 quota embedding_rows: dict = {} # entity_id → row data (for entities found only via embedding) emb_svc = get_embedding_service() if emb_svc.is_available(): for q in query_list: + per_q_scores: dict = {} query_emb = emb_svc.embed(q) - if not query_emb: - continue - similar_kw = find_similar_keywords(conn, query_emb, limit=30) - if not similar_kw: - continue - kw_sim = {str(kw["id"]): kw["similarity"] for kw in similar_kw} - kw_ids = list(kw_sim.keys()) - entities = find_entities_for_keywords(conn, kw_ids) - for ent in entities: - eid = ent["id"] - matched_ids = [str(mid) for mid in (ent.get("matched_keyword_ids") or [])] - if matched_ids: - best_sim = max(kw_sim.get(mid, 0) for mid in matched_ids) - if eid not in embedding_scores or best_sim > embedding_scores[eid]: - embedding_scores[eid] = best_sim - if eid not in seed_rows_by_id: - embedding_rows[eid] = ent + if query_emb: + # Scoring pool — same principle: wide candidate set for the + # embedding pathway. A narrow keyword may rank far below 30 for + # a sentence-shaped query even when it's an exact term match; + # widening here keeps it visible to the rest of the pipeline. + similar_kw = find_similar_keywords( + conn, query_emb, limit=settings.scoring_pool_keyword_neighbors, + ) + if similar_kw: + kw_sim = {str(kw["id"]): kw["similarity"] for kw in similar_kw} + kw_ids = list(kw_sim.keys()) + entities = find_entities_for_keywords(conn, kw_ids) + for ent in entities: + eid = ent["id"] + matched_ids = [str(mid) for mid in (ent.get("matched_keyword_ids") or [])] + if matched_ids: + best_kw_id = max(matched_ids, key=lambda m: kw_sim.get(m, 0)) + best_sim = kw_sim.get(best_kw_id, 0) + per_q_scores[str(eid)] = best_sim + if eid not in embedding_scores or best_sim > embedding_scores[eid]: + embedding_scores[eid] = best_sim + embedding_dom_kw[eid] = best_kw_id + if eid not in seed_rows_by_id: + embedding_rows[eid] = ent + embedding_scores_by_q.append(per_q_scores) + else: + embedding_scores_by_q = [{} for _ in query_list] # ------------------------------------------------------------------ # # 3. MERGE — geometric mean when both, penalty when single signal # @@ -196,7 +367,10 @@ def assemble_context(conn, req: ContextRequest) -> ContextResponse: seed_scores[eid] = text_s * penalty # text only — penalized elif emb_s: seed_scores[eid] = emb_s * penalty # embedding only — penalized - # Ensure we have row data for embedding-only entities + # Ensure we have row data for entities that came in via either + # of the two keyword-mediated pathways. + if eid not in seed_rows_by_id and eid in fuzzy_rows: + seed_rows_by_id[eid] = fuzzy_rows[eid] if eid not in seed_rows_by_id and eid in embedding_rows: seed_rows_by_id[eid] = embedding_rows[eid] @@ -229,7 +403,63 @@ def assemble_context(conn, req: ContextRequest) -> ContextResponse: items.append(_to_item(row, score, depth, relevance, ext_map.get(eid, {}))) items.sort(key=lambda x: x.final_rank, reverse=True) - items = items[: req.max_results] + + # Build the inputs the two-level diversity quota needs. + # + # `dominant_kw_by_id`: which matched keyword "won" for each entity + # (used by Level 2 — per-keyword quota). Whichever pathway scored + # the entity higher (text-fuzzy or embedding) supplies the keyword. + dominant_kw_by_id: dict[str, str] = {} + for eid in seed_scores: + text_s = text_scores.get(eid, 0.0) + emb_s = embedding_scores.get(eid, 0.0) + if emb_s >= text_s and eid in embedding_dom_kw: + dominant_kw_by_id[str(eid)] = embedding_dom_kw[eid] + elif eid in text_dom_kw: + dominant_kw_by_id[str(eid)] = text_dom_kw[eid] + + # `per_query_top_ids`: each query's top-K entities by THAT query's + # own combined score (geometric-mean merge of text + embedding per + # query, same formula the global merge uses). Used by Level 1 — + # per-search-term reservation. Each query gets `K` reserved slots: + # `K = ceil(max_results × per_query_share / num_queries)`. The + # narrow-query-strategy nudge in `recall_memory`'s docstring is + # what makes this useful: when the agent issues a focused + # single-keyword query alongside broader ones, that focused query + # is guaranteed a reserved share of the result. + penalty = settings.missing_signal_penalty + nq = max(1, len(query_list)) + per_q_reserved = max( + 0, math.ceil(req.max_results * settings.per_query_share / nq) + ) + per_query_top_ids: list[list[str]] = [] + if per_q_reserved > 0 and settings.per_query_share > 0: + for q_idx in range(nq): + t_q = text_scores_by_q[q_idx] if q_idx < len(text_scores_by_q) else {} + e_q = embedding_scores_by_q[q_idx] if q_idx < len(embedding_scores_by_q) else {} + # Same merge math as the global seed_scores, but using + # only THIS query's text and embedding signals. + per_q_seed: dict[str, float] = {} + for eid in set(t_q) | set(e_q): + t = t_q.get(eid) + e = e_q.get(eid) + if t and e: + per_q_seed[eid] = math.sqrt(t * e) + elif t: + per_q_seed[eid] = t * penalty + elif e: + per_q_seed[eid] = e * penalty + ordered = sorted(per_q_seed.items(), key=lambda kv: -kv[1])[:per_q_reserved] + per_query_top_ids.append([eid for eid, _ in ordered]) + + items = _apply_two_level_quota( + items, + dominant_kw_by_id, + per_query_top_ids, + req.max_results, + per_query_share=settings.per_query_share, + halving=settings.keyword_quota_halving, + ) always_on = [] if req.include_always_on_rules: diff --git a/braindb/services/keyword_service.py b/braindb/services/keyword_service.py index 351bc5c..47b472b 100644 --- a/braindb/services/keyword_service.py +++ b/braindb/services/keyword_service.py @@ -138,6 +138,38 @@ def find_similar_keywords(conn, query_embedding: list[float], limit: int = 20) - return [dict(r) for r in cur.fetchall()] +def find_fuzzy_keywords(conn, query: str, limit: int = 20) -> list[dict]: + """Trigram-similarity search against keyword entities. + + Mirror image of `find_similar_keywords` (the embedding-based form): + the query is matched against the *keyword text itself* (not against + long entity bodies), so a short query vs a short keyword gives a fair + trigram intersection — no dilution. Returns rows in the same shape + as `find_similar_keywords` so `assemble_context` can use the two + pathways symmetrically (match query → keyword → fan out via + `find_entities_for_keywords`). + + This is the indexing-layer view of fuzzy: the user-stated design + intent of BrainDB is that keywords are the hub. Direct + entity-content fuzzy still exists in `services.search.fuzzy_search` + for `quick_search` / `/memory/search` (arbitrary-content matching). + """ + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """ + SELECT id, content AS keyword, + similarity(content, %s) AS similarity + FROM entities + WHERE entity_type = 'keyword' + AND similarity(content, %s) > 0.1 + ORDER BY similarity(content, %s) DESC + LIMIT %s + """, + (query, query, query, limit), + ) + return [dict(r) for r in cur.fetchall()] + + def find_entities_for_keywords(conn, keyword_entity_ids: list[str]) -> list[dict]: """ Find all non-keyword entities tagged with the given keyword entities. @@ -147,9 +179,18 @@ def find_entities_for_keywords(conn, keyword_entity_ids: list[str]) -> list[dict return [] with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + # Cast array_agg to text[] so psycopg2 returns a proper Python list + # of UUID strings. Without the explicit cast the column type comes + # back as a literal Postgres array string ('{uuid1,uuid2}') because + # psycopg2's default uuid[] adapter is not registered — iterating + # over that string yields single characters and downstream + # `kw_sim.get(mid, 0)` returns 0 for ALL matched keywords, silently + # zeroing the entire embedding-based recall path. The same cast + # pattern is already used for `wikis_ext.member_keyword_ids::text[]` + # in context.py. cur.execute( """ - SELECT e.*, array_agg(r.to_entity_id) AS matched_keyword_ids + SELECT e.*, array_agg(r.to_entity_id::text) AS matched_keyword_ids FROM entities e JOIN relations r ON r.from_entity_id = e.id WHERE r.to_entity_id = ANY(%s::uuid[]) diff --git a/braindb/services/search.py b/braindb/services/search.py index 3015e10..630b8fe 100644 --- a/braindb/services/search.py +++ b/braindb/services/search.py @@ -6,8 +6,57 @@ 3. Content trigram similarity — weight 0.5 4. Title trigram similarity — weight 0.3 """ +import os + import psycopg2.extras +# ------------------------------------------------------------------ # +# Central content-preview helper (shared by recall/search/list/etc.) # +# ------------------------------------------------------------------ # +# Lives here because search.py is a dependency-free leaf module that +# context.py and the agent tools already import — so this is reused, not +# a new module. The ONLY full-content read is get_entity(); every +# multi-item path renders previews so big/polluted bodies never flood +# (or pollute) the caller's context. +PREVIEW_CAP = int(os.getenv("BRAINDB_PREVIEW_CAP", "1024")) # <= 1K per item +SLICE_MAX = int(os.getenv("BRAINDB_SLICE_MAX", "8000")) # max chars per get-by-id slice + + +def slice_content(text, offset: int = 0, limit: int | None = None) -> tuple[str, dict]: + """Return (slice, meta) of a full content string for the by-id deep read. + A slice is clamped to SLICE_MAX so one slice can never itself flood a + caller — large bodies are read by paging `next_offset` (and/or handing + each slice to a separate subagent). `meta.next_offset` is None at EOF. + Used only when offset/limit are explicitly requested; default get-by-id + behaviour is unchanged (full body).""" + s = "" if text is None else str(text) + total = len(s) + offset = max(0, int(offset)) + eff = SLICE_MAX if limit is None else max(1, min(int(limit), SLICE_MAX)) + chunk = s[offset:offset + eff] + nxt = offset + len(chunk) + return chunk, { + "total_chars": total, + "offset": offset, + "returned": len(chunk), + "next_offset": nxt if nxt < total else None, + } + + +def preview(text, entity_id=None, cap: int = PREVIEW_CAP) -> str: + """Bound a content string to `cap` chars; if cut, append the standard + marker + drill-down protocol so the LLM knows how to read the full body.""" + s = "" if text is None else str(text) + if len(s) <= cap: + return s + extra = len(s) - cap + how = f' full body: get_entity("{entity_id}").' if entity_id else "." + return ( + s[:cap] + + f"\n--truncated ({extra} more chars)--{how} If large, " + "delegate_to_subagent to read/extract it without polluting this context." + ) + # Shared SQL fragments _OR_TSQUERY = "to_tsquery('english', regexp_replace(plainto_tsquery('english', %s)::text, ' & ', ' | ', 'g'))" @@ -71,4 +120,10 @@ def fuzzy_search(conn, query: str, entity_types: list[str] | None, min_importanc with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: cur.execute(sql, params) - return [dict(r) for r in cur.fetchall()] + rows = [dict(r) for r in cur.fetchall()] + # Central preview cap — covers /memory/search + quick_search (and the + # text seeds feeding /memory/context). Real content is read only via + # get_entity() (the full carve-out). + for r in rows: + r["content"] = preview(r.get("content"), r.get("id")) + return rows diff --git a/braindb/services/wiki_jobs.py b/braindb/services/wiki_jobs.py new file mode 100644 index 0000000..e991b1b --- /dev/null +++ b/braindb/services/wiki_jobs.py @@ -0,0 +1,630 @@ +""" +Wiki job queue — non-destructive plumbing only. + +This module is deliberately free of any search/scoring/LLM logic (constraint +C3): finding *what* to wiki-ify and *how* to write it is delegated to the +existing recall/agent infra by the routers. Here we only: + + * detect orphans with one read-only SQL pass (no scoring), + * enqueue exactly one `triage` job per orphan (idempotent), + * list jobs. + +Claim / status-transition / advisory-lock / accounted-change-gate helpers are +added in later steps, alongside the endpoints that use them. +""" +import os +import re +import uuid + +import psycopg2.extras + +ACTIVE_STATUSES = ("pending", "assigned") + +# Freshness window: an entity is only orphan-eligible once it has existed for +# this many minutes, so the maintainer never wikis a subject whose ingest +# burst of facts/relations has not settled yet. Same env-var pattern the +# scheduler uses for its intervals (keeps this plumbing module config-import +# free). MUST be measured on created_at, never updated_at — the unconditional +# entities_updated_at BEFORE UPDATE trigger bumps updated_at on every recall +# access, which would leave recalled entities perpetually "fresh". +FRESHNESS_MINUTES = int(os.getenv("WIKI_FRESHNESS_MINUTES", "30")) + +# Stale-lease (visibility-timeout) for claimed jobs. A job sits in `assigned` +# only while a worker is actively running it; if that worker never returns +# (api restart mid-run, agent timeout) the row would wedge forever. Instead +# of a reaper/cycle, an `assigned` job whose lease expired is simply +# claimable again at the EXISTING claim step. 20 min is comfortably above +# the longest legit run (AGENT_TIMEOUT ~10 min), so a still-running job is +# never reclaimed. `attempts`+max_attempts already bound repeated failures. +ASSIGNED_LEASE_MIN = int(os.getenv("WIKI_ASSIGNED_LEASE_MIN", "20")) + +# Per-wiki attach grouping — how long to wait before firing a writer on a +# wiki that just received new attaches. Once the OLDEST pending attach for +# a wiki is this old, the writer claims ALL pending attaches for that wiki +# in a single batch (the within-tick batching at next_write_bucket()'s +# second query already groups by target_wiki_id). Lets attaches accumulate +# so the writer fires once per cooldown window instead of once per job — +# directly addresses the "Dimitrios Koutsoumpos rewritten 8x in an hour" +# pattern observed on Qwen. Bigger windows = lower LLM cost but slower +# wiki freshness. Default 300s (5 min). Set to 0 to disable (revert to +# the old "fire on every attach" behaviour). Self-limiting — no force-claim +# valve needed because the bucket scoops up the WHOLE pending queue for +# that wiki on each fire. +ATTACH_COOLDOWN_SEC = int(os.getenv("WIKI_ATTACH_COOLDOWN_SECONDS", "300")) + + +def _claimable(alias: str = "") -> str: + """SQL predicate: a job is claimable if pending, OR assigned but its + lease expired. Reused verbatim at every claim site (DRY). `alias` is the + table alias when the query qualifies columns (e.g. 'j').""" + p = f"{alias}." if alias else "" + return (f"({p}status = 'pending' OR ({p}status = 'assigned' " + f"AND {p}assigned_at < now() - make_interval(mins => {ASSIGNED_LEASE_MIN})))") + +# Inline reference token: [[ref:UUID]] or [[ref:UUID|display text]] +REF_RE = re.compile( + r"\[\[ref:([0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-" + r"[0-9a-fA-F]{4}-[0-9a-fA-F]{12})(?:\|[^\]]*)?\]\]" +) + +SUMMARY_RE = re.compile(r">\s*\*\*Summary:\*\*\s*(.+)") +DISAMBIG_RE = re.compile(r">\s*\*\*Disambiguation:\*\*\s*(.+)") +# The LLM authors its own meta header line; we only READ what it declared. +META_KEYWORDS_RE = re.compile(r"") + + +def parse_refs(body: str) -> set[str]: + """All entity UUIDs cited inline in the body (lower-cased).""" + return {m.lower() for m in REF_RE.findall(body or "")} + + +def keywords_from_meta(body: str) -> list[str]: + """Read keywords the LLM declared in its own `` + header (e.g. `keywords=a;b;c`). Reading the LLM's declaration is not code + authoring content. Returns [] if the LLM declared none.""" + m = META_KEYWORDS_RE.search(body or "") + if not m: + return [] + raw = m.group(1).replace(",", ";") + return [k.strip() for k in raw.split(";") if k.strip()] + + +def snapshot_revision(conn, wiki_id: str, old_content: str, old_refs: set[str], + revision: int) -> None: + """Persist the prior body+refs before mutation so any change is reversible.""" + from braindb.services.activity_log import log_activity + log_activity(conn, "wiki_revise", "wiki", wiki_id, details={ + "from_revision": revision, + "prior_content": old_content, + "prior_refs": sorted(old_refs), + }) + + +def reconcile_summarises_additive(conn, wiki_id: str, body: str) -> dict: + """ + Pure bookkeeping: ensure a `wiki --summarises--> e` relation exists for + every entity the LLM cited inline (`[[ref:UUID]]`). ADDITIVE ONLY — it + never deletes or re-types a relation behind the LLM. If the LLM wants a + relation gone it calls `delete_relation` itself. Mirrors LLM-authored + content into the graph; it does not judge or shape content. + """ + cited = parse_refs(body) + added = 0 + with conn.cursor() as cur: + cur.execute( + "SELECT to_entity_id::text FROM relations " + "WHERE from_entity_id = %s AND relation_type = 'summarises'", + (wiki_id,), + ) + current = {r[0].lower() for r in cur.fetchall()} + for e in cited - current: + cur.execute( + """INSERT INTO relations + (from_entity_id, to_entity_id, relation_type, relevance_score, description) + VALUES (%s, %s, 'summarises', 0.9, 'wiki body reference') + ON CONFLICT (from_entity_id, to_entity_id, relation_type) DO NOTHING""", + (wiki_id, e), + ) + added += 1 + return {"relations_added": added, "relations_removed": 0} + + +def try_wiki_lock(conn, key: str) -> bool: + """Transaction-scoped advisory lock so two writers never touch one wiki.""" + with conn.cursor() as cur: + cur.execute("SELECT pg_try_advisory_xact_lock(hashtext(%s))", (f"wiki:{key}",)) + return bool(cur.fetchone()[0]) + + +def claim_jobs(conn, job_ids: list[str]) -> int: + """Mark a bucket's pending suggestion jobs as assigned (SKIP LOCKED).""" + if not job_ids: + return 0 + with conn.cursor() as cur: + cur.execute( + f"""UPDATE wiki_job SET status='assigned', assigned_at=now(), attempts=attempts+1 + WHERE id = ANY(%s::uuid[]) AND {_claimable()} + AND id IN (SELECT id FROM wiki_job WHERE id = ANY(%s::uuid[]) + FOR UPDATE SKIP LOCKED)""", + (job_ids, job_ids), + ) + return cur.rowcount + +# Entity types the cron considers "wiki-able" content. Keywords act as concept +# hubs; thoughts/facts are the substance. (wiki/source/datasource/rule excluded.) +ORPHAN_ENTITY_TYPES = ("keyword", "thought", "fact") + + +def _orphan_conditions(exclude_job: bool = False) -> str: + """ + The SINGLE definition of "orphan" (entity not yet covered by a wiki), + shared by `run_cron` (set-based) and `is_orphan` (per-entity) so the two + can never drift. References the entity as `e.id`. All conditions are + param-free EXCEPT the optional `exclude_job` clause (one %s) used by the + maintainer staleness guard to ignore the just-claimed triage row itself. + + An orphan is an entity that: + * has settled — `created_at` is older than FRESHNESS_MINUTES (so a + still-ingesting subject is not wikied half-formed), + * is not the target of a `wiki --summarises--> e` relation, + * is not listed in any wiki's `member_keyword_ids`, + * is not referenced by an active (pending/assigned) wiki_job, + * does not carry a `rejected` triage (deliberate-skip self-clearing; + `failed` triage is NOT excluded so transient errors still retry). + """ + xj = " AND j.id <> %s" if exclude_job else "" + return f""" + e.created_at < now() - make_interval(mins => {FRESHNESS_MINUTES}) + AND NOT EXISTS ( + SELECT 1 FROM relations r + JOIN entities w ON w.id = r.from_entity_id AND w.entity_type = 'wiki' + WHERE r.relation_type = 'summarises' AND r.to_entity_id = e.id + ) + AND NOT EXISTS ( + SELECT 1 FROM wikis_ext wx WHERE e.id = ANY(wx.member_keyword_ids) + ) + AND NOT EXISTS ( + SELECT 1 FROM wiki_job j + WHERE j.status IN ('pending','assigned') + AND e.id = ANY(j.entity_ids){xj} + ) + AND NOT EXISTS ( + SELECT 1 FROM wiki_job j + WHERE j.job_type = 'triage' AND j.status = 'rejected' + AND e.id = ANY(j.entity_ids) + ) + """ + + +def is_orphan(conn, entity_id, exclude_triage_job_id: str | None = None) -> bool: + """True if the entity is still uncovered by any wiki. Used by the + maintainer staleness guard: if a prior writer run already absorbed/linked + the entity (or it is already in an active suggestion), this returns False + and the maintainer skips it with NO LLM call. Same predicate as cron.""" + cond = _orphan_conditions(exclude_job=exclude_triage_job_id is not None) + params: list = [str(entity_id)] + if exclude_triage_job_id is not None: + params.append(str(exclude_triage_job_id)) + with conn.cursor() as cur: + cur.execute( + f"SELECT EXISTS (SELECT 1 FROM entities e WHERE e.id = %s AND {cond})", + params, + ) + return bool(cur.fetchone()[0]) + + +def run_cron(conn) -> dict: + """ + Find entities not yet connected to any wiki and enqueue one `triage` + job per orphan. Pure SQL, read-only except the additive job insert. + Orphan-ness is the shared `_orphan_conditions()` (see there). + + Idempotent: the partial-unique index on `dedupe_key WHERE status IN + ('pending','assigned')` + ON CONFLICT DO NOTHING means re-running cron + never creates duplicate triage jobs. + """ + batch_id = str(uuid.uuid4()) + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + f""" + WITH orphans AS ( + SELECT e.id + FROM entities e + WHERE e.entity_type = ANY(%s) + AND {_orphan_conditions()} + ) + INSERT INTO wiki_job (job_type, status, entity_ids, dedupe_key, batch_id) + SELECT 'triage', 'pending', ARRAY[o.id], 'triage:' || o.id::text, %s::uuid + FROM orphans o + ON CONFLICT (dedupe_key) WHERE status IN ('pending','assigned') + DO NOTHING + RETURNING id + """, + (list(ORPHAN_ENTITY_TYPES), batch_id), + ) + enqueued = cur.rowcount + + # Counts for visibility (cheap; the heavy filter already ran above). + cur.execute( + "SELECT count(*) AS c FROM wiki_job WHERE status = 'pending' AND job_type = 'triage'" + ) + pending_triage = cur.fetchone()["c"] + + return { + "batch_id": batch_id, + "triage_jobs_enqueued": enqueued, + "pending_triage_total": pending_triage, + } + + +def claim_one_triage(conn) -> dict | None: + """ + Atomically claim a single pending triage job (C1: one case per call). + FOR UPDATE SKIP LOCKED guarantees two concurrent maintainer calls never + grab the same case. Highest-importance orphan first, so high-value + concepts get wikis early and their writer runs absorb neighbourhoods + (more downstream triage becomes free stale-skips). + """ + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + f""" + UPDATE wiki_job + SET status = 'assigned', assigned_at = now(), attempts = attempts + 1 + WHERE id = ( + SELECT j.id FROM wiki_job j + JOIN entities e ON e.id = j.entity_ids[1] + WHERE {_claimable("j")} AND j.job_type = 'triage' + ORDER BY e.importance DESC, j.created_at + FOR UPDATE OF j SKIP LOCKED + LIMIT 1 + ) + RETURNING id, entity_ids::text[] AS entity_ids, batch_id + """ + ) + row = cur.fetchone() + return dict(row) if row else None + + +def finish_job(conn, job_id: str, status: str, last_error: str | None = None) -> None: + """Transition a job to a terminal state (done / rejected / failed).""" + with conn.cursor() as cur: + cur.execute( + "UPDATE wiki_job SET status = %s, completed_at = now(), last_error = %s WHERE id = %s", + (status, last_error, str(job_id)), + ) + + +def fetch_entity_brief(conn, entity_id: str) -> dict | None: + """Minimal entity view for building a focused maintainer prompt.""" + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + "SELECT id, entity_type, content, summary, keywords FROM entities WHERE id = %s", + (str(entity_id),), + ) + row = cur.fetchone() + return dict(row) if row else None + + +def suggestion_dedupe_key(action: str, target_wiki_id: str | None, + entity_ids: list[str], consolidate_wiki_ids: list[str]) -> str: + """Deterministic, service-computed (never LLM-computed) idempotency key.""" + if action == "attach": + return f"attach:{target_wiki_id}:" + ",".join(sorted(entity_ids)) + if action == "create": + return "create:" + ",".join(sorted(entity_ids)) + if action == "consolidate": + return "consolidate:" + ",".join(sorted(consolidate_wiki_ids)) + raise ValueError(f"unknown action {action!r}") + + +def insert_suggestion(conn, *, job_type: str, target_wiki_id: str | None, + entity_ids: list[str], dedupe_key: str, rationale: str | None, + proposed_name: str | None, batch_id: str | None) -> str | None: + """ + Insert a maintainer suggestion job. ON CONFLICT DO NOTHING against the + partial-unique active dedupe index → re-proposing the same work is a no-op. + Returns the new job id, or None if it was a duplicate. + """ + with conn.cursor() as cur: + cur.execute( + """ + INSERT INTO wiki_job + (job_type, status, target_wiki_id, entity_ids, dedupe_key, + rationale, proposed_name, batch_id) + VALUES (%s, 'pending', %s, %s::uuid[], %s, %s, %s, %s) + ON CONFLICT (dedupe_key) WHERE status IN ('pending','assigned') + DO NOTHING + RETURNING id + """, + (job_type, target_wiki_id, entity_ids, dedupe_key, + rationale, proposed_name, batch_id), + ) + row = cur.fetchone() + return str(row[0]) if row else None + + +def next_write_bucket(conn) -> dict | None: + """ + Pick the next unit of writer work (one wiki per call). A `create` job is + its own bucket; `attach` jobs are grouped by target_wiki_id so the writer + sees every new member of a wiki at once. Consolidate is handled by Step 5. + + Dedup-first priority: pending jobs are ordered consolidate -> attach -> + create (then created_at). The moment the maintainer emits a `consolidate` + the writer drains it before creating/expanding more pages, so the wiki + set converges before it grows. + + Per-wiki cooldown on attach (ATTACH_COOLDOWN_SEC, default 300s = 5 min): + an attach seed becomes claimable only once its target wiki's oldest + pending attach is past the cooldown. Once eligible, the existing second + query below scoops up ALL pending attaches for that wiki (including the + ones inserted during the cooldown window) into one writer call. Net + effect: writer fires once per cooldown window per wiki instead of once + per job. `consolidate` and `create` paths are unaffected; the cooldown + is attach-only because attach is the only multi-job-per-wiki shape. + """ + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + f"""SELECT id, job_type, target_wiki_id, entity_ids::text[] AS entity_ids, + proposed_name, rationale, batch_id + FROM wiki_job j + WHERE {_claimable("j")} + AND job_type IN ('create','attach','consolidate') + AND ( + job_type <> 'attach' + OR ( + SELECT MIN(created_at) + FROM wiki_job + WHERE target_wiki_id = j.target_wiki_id + AND job_type = 'attach' + AND status = 'pending' + ) <= now() - make_interval(secs => {ATTACH_COOLDOWN_SEC}) + ) + ORDER BY CASE job_type WHEN 'consolidate' THEN 0 + WHEN 'attach' THEN 1 + ELSE 2 END, + created_at + LIMIT 1""" + ) + seed = cur.fetchone() + if not seed: + return None + seed = dict(seed) + if seed["job_type"] == "create": + return {"mode": "create", "jobs": [seed], + "target_wiki_id": None, "proposed_name": seed["proposed_name"]} + if seed["job_type"] == "consolidate": + # entity_ids holds the wiki ids the maintainer flagged as duplicates. + return {"mode": "consolidate", "jobs": [seed], + "target_wiki_id": None, "proposed_name": None, + "wiki_ids": seed["entity_ids"]} + cur.execute( + f"""SELECT id, entity_ids::text[] AS entity_ids + FROM wiki_job + WHERE {_claimable()} AND job_type='attach' + AND target_wiki_id = %s + ORDER BY created_at""", + (seed["target_wiki_id"],), + ) + jobs = [dict(r) for r in cur.fetchall()] + return {"mode": "attach", "jobs": jobs, + "target_wiki_id": str(seed["target_wiki_id"]), "proposed_name": None} + + +def fetch_members(conn, entity_ids: list[str]) -> list[dict]: + if not entity_ids: + return [] + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + "SELECT id::text AS id, entity_type, content, keywords " + "FROM entities WHERE id = ANY(%s::uuid[])", + (entity_ids,), + ) + return [dict(r) for r in cur.fetchall()] + + +def fetch_wiki(conn, wiki_id: str) -> dict | None: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """SELECT e.id::text AS id, e.content, w.canonical_name, w.revision, + w.member_keyword_ids::text[] AS member_keyword_ids + FROM entities e JOIN wikis_ext w ON w.entity_id = e.id + WHERE e.id = %s""", + (str(wiki_id),), + ) + row = cur.fetchone() + return dict(row) if row else None + + +def list_active_wikis(conn) -> list[dict]: + """All non-retired wikis as {id, canonical_name}, deterministically + ordered. Plumbing read (mirrors fetch_wiki / export_wikis SQL) — the + maintainer is shown this as a NUMBERED catalog so it references wikis by + number, never by uuid; the order here IS the numbering.""" + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """SELECT e.id::text AS id, w.canonical_name + FROM entities e JOIN wikis_ext w ON w.entity_id = e.id + WHERE e.entity_type = 'wiki' AND w.retired_at IS NULL + ORDER BY e.importance DESC, e.created_at""" + ) + return [dict(r) for r in cur.fetchall()] + + +def release_or_fail_jobs(conn, job_ids: list[str], last_error: str, + max_attempts: int = 3) -> str: + """On a gate failure: return jobs to 'pending' for retry, or 'failed' once + attempts are exhausted (surfaced via GET /jobs — never a silent bad write).""" + if not job_ids: + return "none" + with conn.cursor() as cur: + cur.execute( + """UPDATE wiki_job + SET status = CASE WHEN attempts >= %s THEN 'failed' ELSE 'pending' END, + last_error = %s + WHERE id = ANY(%s::uuid[])""", + (max_attempts, last_error[:1000], job_ids), + ) + return "failed" if _max_attempts_reached(conn, job_ids, max_attempts) else "requeued" + + +def _max_attempts_reached(conn, job_ids: list[str], max_attempts: int) -> bool: + with conn.cursor() as cur: + cur.execute( + "SELECT bool_or(status='failed') FROM wiki_job WHERE id = ANY(%s::uuid[])", + (job_ids,), + ) + return bool(cur.fetchone()[0]) + + +def finish_jobs(conn, job_ids: list[str], status: str, last_error: str | None = None) -> None: + if not job_ids: + return + with conn.cursor() as cur: + cur.execute( + "UPDATE wiki_job SET status=%s, completed_at=now(), last_error=%s " + "WHERE id = ANY(%s::uuid[])", + (status, last_error, job_ids), + ) + + +def create_wiki_entity(conn, canonical_name: str, body: str, summary: str | None, + disambiguation: str | None, member_entity_ids: list[str], + keywords: list[str] | None = None) -> str: + """Scaffolding only — a new wiki page is additive, not destruction. The + body, summary, disambiguation, and keywords are ALL the LLM's: `keywords` + is whatever the LLM declared in its meta header (may be empty). Code never + invents keywords (no `[canonical_name]` default).""" + from braindb.services.embedding_service import get_embedding_service + from braindb.services.keyword_service import ( + ensure_keyword_entities, link_entity_to_keywords, + ) + kws = [k.strip() for k in (keywords or []) if k and k.strip()] + with conn.cursor() as cur: + cur.execute( + """INSERT INTO entities (entity_type, title, content, summary, keywords, + importance, source) + VALUES ('wiki', %s, %s, %s, %s, 0.9, 'agent-inference') + RETURNING id""", + (canonical_name, body, summary, kws), + ) + wid = str(cur.fetchone()[0]) + if kws: + kw_map = ensure_keyword_entities(conn, kws, get_embedding_service()) + link_entity_to_keywords(conn, wid, list(kw_map.values())) + member_kw = _keyword_ids_among(conn, member_entity_ids) + with conn.cursor() as cur: + cur.execute( + """INSERT INTO wikis_ext + (entity_id, canonical_name, disambiguation, language, + member_keyword_ids, revision, last_synthesised_at) + VALUES (%s, %s, %s, 'en', %s::uuid[], 1, now())""", + (wid, canonical_name, disambiguation, member_kw), + ) + return wid + + +def _keyword_ids_among(conn, entity_ids: list[str]) -> list[str]: + if not entity_ids: + return [] + with conn.cursor() as cur: + cur.execute( + "SELECT id::text FROM entities " + "WHERE id = ANY(%s::uuid[]) AND entity_type='keyword'", + (entity_ids,), + ) + return [r[0] for r in cur.fetchall()] + + +def finalize_wiki_write(conn, wiki_id: str, new_body: str, summary: str | None, + disambiguation: str | None, member_entity_ids: list[str]) -> int: + """Apply the gated body to an existing wiki: update content + header + fields, union new keyword members, bump revision.""" + new_kw = _keyword_ids_among(conn, member_entity_ids) + with conn.cursor() as cur: + cur.execute("UPDATE entities SET content=%s, summary=%s WHERE id=%s", + (new_body, summary, wiki_id)) + cur.execute( + """UPDATE wikis_ext + SET disambiguation = COALESCE(%s, disambiguation), + member_keyword_ids = ( + SELECT ARRAY(SELECT DISTINCT unnest( + member_keyword_ids || %s::uuid[]))), + revision = revision + 1, + last_synthesised_at = now() + WHERE entity_id = %s + RETURNING revision""", + (disambiguation, new_kw, wiki_id), + ) + return cur.fetchone()[0] + + +def fetch_wikis_for_merge(conn, wiki_ids: list[str]) -> list[dict]: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """SELECT e.id::text AS id, e.content, e.importance, w.canonical_name, + w.revision, w.member_keyword_ids::text[] AS member_keyword_ids, + w.retired_at + FROM entities e JOIN wikis_ext w ON w.entity_id = e.id + WHERE e.id = ANY(%s::uuid[]) AND e.entity_type='wiki'""", + (wiki_ids,), + ) + return [dict(r) for r in cur.fetchall()] + + +def soft_retire_wiki(conn, loser_id: str, canonical_id: str, note: str | None) -> None: + """LLM-decided retirement, executed deterministically + reversibly: the + loser drops out of ranking (importance 0) but still resolves; provenance + is kept via duplicate_of / consolidated_into edges (which also self-clear + the maintainer's dedup, since it is prompted to skip marked pairs).""" + from braindb.services.activity_log import log_activity + with conn.cursor() as cur: + cur.execute("UPDATE entities SET importance = 0.0 WHERE id = %s", (loser_id,)) + cur.execute( + "UPDATE wikis_ext SET retired_at = now(), redirect_to = %s WHERE entity_id = %s", + (canonical_id, loser_id), + ) + for rtype in ("duplicate_of", "consolidated_into"): + cur.execute( + """INSERT INTO relations + (from_entity_id, to_entity_id, relation_type, relevance_score, description) + VALUES (%s, %s, %s, 0.0, %s) + ON CONFLICT (from_entity_id, to_entity_id, relation_type) DO NOTHING""", + (loser_id, canonical_id, rtype, (note or "merged")[:500]), + ) + log_activity(conn, "wiki_merge", "wiki", canonical_id, + details={"retired": loser_id, "canonical": canonical_id, "note": note}) + + +def extract_summary_disambig(body: str) -> tuple[str | None, str | None]: + sm = SUMMARY_RE.search(body or "") + dm = DISAMBIG_RE.search(body or "") + return (sm.group(1).strip() if sm else None, + dm.group(1).strip() if dm else None) + + +def list_jobs(conn, status: str | None, job_type: str | None, limit: int) -> list[dict]: + conditions, params = [], [] + if status: + conditions.append("status = %s") + params.append(status) + if job_type: + conditions.append("job_type = %s") + params.append(job_type) + where = ("WHERE " + " AND ".join(conditions)) if conditions else "" + params.append(limit) + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + f""" + SELECT id, job_type, status, target_wiki_id, + entity_ids::text[] AS entity_ids, dedupe_key, rationale, + proposed_name, batch_id, created_at, assigned_at, + completed_at, attempts, last_error + FROM wiki_job + {where} + ORDER BY created_at DESC + LIMIT %s + """, + params, + ) + return [dict(r) for r in cur.fetchall()] diff --git a/braindb/services/wiki_sections.py b/braindb/services/wiki_sections.py new file mode 100644 index 0000000..9b8f69d --- /dev/null +++ b/braindb/services/wiki_sections.py @@ -0,0 +1,218 @@ +"""Section-level operations on wiki markdown bodies. + +Wiki bodies live as one markdown blob in `entities.content`. This module +parses, splices, and validates them at the section level so the writer +agent can edit ONE section at a time instead of rewriting the whole +body — the fix for big-wiki context exhaustion on smaller-context +models (see plan: read-write tools / handoff). + +Sections are anchored on `` HTML-comment markers +that the writer prompt already mandates (see `wiki_writer_prompt.md` +"Recommended structure"). Everything before the first marker is the +HEADER (meta-comment + `# Title` + `> **Summary:** ...` callout) and +is preserved verbatim by all splice operations. + +Optimistic concurrency: every read returns the wiki's current +`wikis_ext.revision`. Every write requires the caller to pass that +revision back as `expect_revision`. A mismatch raises +`StaleRevisionError` so the caller re-reads and retries instead of +silently stomping on a concurrent edit. + +Pure parsing functions (`parse_sections`, `splice_section`, +`delete_section`, `check_grammar`) are DB-free and unit-testable. +The two DB helpers at the bottom (`fetch_wiki_for_section_op`, +`apply_section_write`) are the only stateful surface. +""" +from __future__ import annotations + +import re +from dataclasses import dataclass + + +class StaleRevisionError(Exception): + """Raised by `apply_section_write` when the caller's + `expect_revision` no longer matches the wiki's current revision. + Means the body was changed by someone else (or by the same agent + in an earlier turn) since the caller last read it.""" + + +# Section marker. Captured group = the section name. We accept +# alphanumerics, dashes, and underscores in the name — matches the +# writer prompt's convention (e.g. `overview`, `timeline`, +# `contradictions`, `sources`, `references`). +_MARKER_RE = re.compile( + r"", + re.MULTILINE, +) + +# UUID shape expected right after `[[ref:`. Real wiki bodies use two +# forms — canonical `[[ref:UUID]]` / `[[ref:UUID|display]]` AND a +# grouped variant `[[ref:UUID1], [ref:UUID2]]` that the writer +# occasionally emits and the frontend plan documents as tolerated. +# Rather than enumerate both forms, we just verify that each +# `[[ref:` is followed by a UUID-looking prefix (8 hex + dash). A +# token that fails this minimal check is genuinely broken (truncated, +# corrupted, or fabricated by a confused model). +_UUID_HEAD_RE = re.compile(r"[0-9a-fA-F]{8}-") + + +@dataclass(frozen=True) +class Section: + name: str + content: str # body text AFTER the marker, up to next marker / EOF + + @property + def char_count(self) -> int: + return len(self.content) + + +def parse_sections(body: str) -> tuple[str, list[Section]]: + """Split a wiki body into (header, sections). + + `header` = everything before the first marker (verbatim). + `sections` = ordered list, each carrying its name + content. + + If the body has no markers, returns `(body, [])` — callers handle + the strict-markers contract themselves. + """ + matches = list(_MARKER_RE.finditer(body)) + if not matches: + return body, [] + header = body[: matches[0].start()] + sections: list[Section] = [] + for i, m in enumerate(matches): + content_start = m.end() + # consume the single newline that conventionally follows the + # marker line, so section content starts on its own line + if content_start < len(body) and body[content_start] == "\n": + content_start += 1 + content_end = matches[i + 1].start() if i + 1 < len(matches) else len(body) + sections.append(Section( + name=m.group(1), + content=body[content_start:content_end], + )) + return header, sections + + +def splice_section(body: str, section_name: str, new_content: str) -> str: + """Replace one named section's content. If the section doesn't exist, + append a new section at the end of the body with that name. + + `new_content` is the section's text WITHOUT the marker line — this + function emits the marker. The result is always normalised so the + rebuilt body parses identically to one written from scratch. + """ + header, sections = parse_sections(body) + new_content = new_content.rstrip("\n") + "\n" + if any(s.name == section_name for s in sections): + sections = [ + Section(s.name, new_content if s.name == section_name else s.content) + for s in sections + ] + return _rebuild(header, sections) + # not found → append a fresh section after the last one + sections = sections + [Section(section_name, new_content)] + return _rebuild(header, sections) + + +def delete_section(body: str, section_name: str) -> str: + """Remove the named section (and its marker) from the body. + Raises KeyError if the section isn't present.""" + header, sections = parse_sections(body) + remaining = [s for s in sections if s.name != section_name] + if len(remaining) == len(sections): + raise KeyError(f"section not found: {section_name}") + return _rebuild(header, remaining) + + +def _rebuild(header: str, sections: list[Section]) -> str: + parts: list[str] = [] + if header: + parts.append(header if header.endswith("\n") else header + "\n") + for s in sections: + parts.append(f"\n") + content = s.content if s.content.endswith("\n") else s.content + "\n" + parts.append(content) + return "".join(parts) + + +def check_grammar(body: str) -> list[str]: + """Return a list of grammar issues with the wiki body. Empty = OK. + + Checked: + - At least one `` marker exists (strict-markers). + - No malformed `[[ref:` tokens (i.e. `[[ref:` that doesn't match + the canonical `[[ref:UUID]]` or `[[ref:UUID|text]]` shape). + - The `> **Summary:**` callout exists in the header. + """ + issues: list[str] = [] + header, sections = parse_sections(body) + if not sections: + issues.append("no markers (strict-markers contract)") + for m in re.finditer(r"\[\[ref:", body): + # Skip past "[[ref:" (6 chars) and check the next chars look like + # the start of a UUID. Tolerates the grouped form + # `[[ref:UUID1], [ref:UUID2]]` since we only check the head. + if not _UUID_HEAD_RE.match(body[m.end():m.end() + 9]): + issues.append(f"malformed [[ref: token at char offset {m.start()}") + if "> **Summary:**" not in header: + issues.append("missing > **Summary:** callout in header") + return issues + + +# ====================================================================== # +# DB helpers # +# ====================================================================== # + +def fetch_wiki_for_section_op(conn, wiki_id: str) -> tuple[str, int] | None: + """Return (content, revision) for the wiki, or None if not found. + Used by every read-side section tool to capture both the body and + the current revision token in one query.""" + with conn.cursor() as cur: + cur.execute( + """SELECT e.content, w.revision + FROM entities e JOIN wikis_ext w ON w.entity_id = e.id + WHERE e.id = %s::uuid AND e.entity_type = 'wiki'""", + (wiki_id,), + ) + row = cur.fetchone() + return (row[0], row[1]) if row else None + + +def apply_section_write(conn, wiki_id: str, new_body: str, + expect_revision: int) -> int: + """Atomically replace the wiki's content + bump its revision. + + The revision UPDATE is conditional on `revision = expect_revision`, + so two writers cannot stomp each other. Returns the new revision + on success. Raises `StaleRevisionError` if the revision didn't + match — caller should re-read and retry. + """ + with conn.cursor() as cur: + cur.execute( + """UPDATE wikis_ext + SET revision = revision + 1, + last_synthesised_at = now() + WHERE entity_id = %s::uuid AND revision = %s + RETURNING revision""", + (wiki_id, expect_revision), + ) + row = cur.fetchone() + if row is None: + cur.execute( + "SELECT revision FROM wikis_ext WHERE entity_id = %s::uuid", + (wiki_id,), + ) + cur_row = cur.fetchone() + if cur_row is None: + raise StaleRevisionError(f"wiki not found: {wiki_id}") + raise StaleRevisionError( + f"expected revision {expect_revision}, current is {cur_row[0]} " + f"— re-read the section before retrying" + ) + new_revision = row[0] + cur.execute( + "UPDATE entities SET content = %s WHERE id = %s::uuid", + (new_body, wiki_id), + ) + return new_revision diff --git a/braindb/tools/__init__.py b/braindb/tools/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/braindb/tools/export_wikis.py b/braindb/tools/export_wikis.py new file mode 100644 index 0000000..6003d08 --- /dev/null +++ b/braindb/tools/export_wikis.py @@ -0,0 +1,218 @@ +""" +Read-only wiki review export. + +Run in the container: + docker compose exec -T api python -m braindb.tools.export_wikis + +Writes one markdown file per wiki to data/wiki_review/ (gitignored) plus an +INDEX.md, so the maintainer/writer output can be read and judged in the IDE. + +STRICTLY READ-ONLY: only SELECT queries, never mutates the DB or the pipeline. +Reuses existing data (entities, relations, wiki_job, activity_log) and the +existing ref/section parsers in wiki_jobs (C3 — no new search/scoring). +""" +import json +import re +from pathlib import Path + +import psycopg2.extras + +from braindb.db import get_conn +from braindb.services.wiki_jobs import parse_refs + +OUT_DIR = Path("data/wiki_review") + + +def _slug(name: str) -> str: + s = re.sub(r"[^a-z0-9]+", "-", (name or "wiki").lower()).strip("-") + return s or "wiki" + + +def _fetch_all_wikis(conn) -> list[dict]: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """SELECT e.id::text AS id, e.content, e.summary, e.importance, + w.canonical_name, w.disambiguation, w.language, w.revision, + w.last_synthesised_at, w.retired_at, w.redirect_to::text AS redirect_to, + w.member_keyword_ids::text[] AS member_keyword_ids, + e.created_at + FROM entities e JOIN wikis_ext w ON w.entity_id = e.id + WHERE e.entity_type = 'wiki' + ORDER BY e.created_at""" + ) + return [dict(r) for r in cur.fetchall()] + + +def _summarises_targets(conn, wiki_id: str) -> list[str]: + with conn.cursor() as cur: + cur.execute( + "SELECT to_entity_id::text FROM relations " + "WHERE from_entity_id = %s AND relation_type = 'summarises'", + (wiki_id,), + ) + return [r[0] for r in cur.fetchall()] + + +def _entities(conn, ids: list[str]) -> dict[str, dict]: + if not ids: + return {} + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + "SELECT id::text AS id, entity_type, content FROM entities " + "WHERE id = ANY(%s::uuid[])", + (ids,), + ) + return {r["id"]: dict(r) for r in cur.fetchall()} + + +def _decisions(conn, wiki_id: str, summarised_ids: list[str]) -> tuple[list[dict], list[dict]]: + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute( + """SELECT job_type, status, rationale, proposed_name, + entity_ids::text[] AS entity_ids, created_at + FROM wiki_job + WHERE target_wiki_id = %s + OR (entity_ids && %s::uuid[]) + ORDER BY created_at""", + (wiki_id, summarised_ids or ["00000000-0000-0000-0000-000000000000"]), + ) + jobs = [dict(r) for r in cur.fetchall()] + cur.execute( + """SELECT operation, timestamp, details + FROM activity_log + WHERE entity_id = %s + AND operation IN ('wiki_write','wiki_revise','wiki_ref_removed','wiki_merge') + ORDER BY timestamp""", + (wiki_id,), + ) + acts = [dict(r) for r in cur.fetchall()] + return jobs, acts + + +def _consistency(body: str, summarises: set[str]) -> tuple[bool, list[str]]: + """Provenance check: every entity the LLM cited inline must have a + `summarises` relation (reconcile is additive, so that should always hold). + Lingering relations (LLM dropped a ref but the edge remains, since code + never deletes behind the LLM) are reported as info, not a failure.""" + inline = parse_refs(body or "") + msgs: list[str] = [] + missing = sorted(inline - summarises) + lingering = sorted(summarises - inline) + if missing: + msgs.append(f"cited inline but NO summarises relation: {missing}") + if lingering: + msgs.append(f"summarises relation but not cited inline (LLM-dropped, " + f"edge left for LLM to remove): {lingering}") + # Pass = no missing relation for a cited ref. Lingering is informational. + return (not missing), msgs + + +def _render(conn, w: dict) -> str: + wid = w["id"] + summarises = set(_summarises_targets(conn, wid)) + ok, issues = _consistency(w["content"] or "", summarises) + all_refs = sorted(parse_refs(w["content"] or "") | summarises) + ents = _entities(conn, all_refs) + jobs, acts = _decisions(conn, wid, sorted(summarises)) + + L = [] + L.append(f"# Wiki review — {w['canonical_name']}") + L.append("") + L.append(f"- **id:** `{wid}`") + L.append(f"- **revision:** {w['revision']} " + f"**importance:** {w['importance']} " + f"**language:** {w['language']}") + L.append(f"- **last_synthesised_at:** {w['last_synthesised_at']}") + L.append(f"- **summary:** {w['summary']}") + L.append(f"- **disambiguation:** {w['disambiguation']}") + L.append("") + L.append(f"## Consistency: {'CONSISTENT ✓' if ok else 'MISMATCH ✗'}") + L.append(f"inline refs / ledger / summarises-relations " + f"({len(parse_refs(w['content'] or ''))} body, {len(summarises)} relations)") + for m in issues: + L.append(f"- ⚠ {m}") + L.append("") + L.append("## Body (verbatim)") + L.append("") + L.append("```markdown") + L.append(w["content"] or "(empty)") + L.append("```") + L.append("") + L.append("## Provenance — cited source entities (judge grounding here)") + for rid in all_refs: + e = ents.get(rid) + if e: + L.append(f"- **`{rid}`** [{e['entity_type']}]: {e['content']}") + else: + L.append(f"- **`{rid}`**: ⚠ ENTITY NOT FOUND (dangling ref)") + L.append("") + L.append("## Decisions & history") + L.append("") + L.append("### Maintainer suggestion jobs") + for j in jobs: + L.append(f"- `{j['job_type']}` [{j['status']}] {j['created_at']:%Y-%m-%d %H:%M} " + f"name={j.get('proposed_name')}\n rationale: {j.get('rationale')}") + L.append("") + L.append("### Writer activity") + for a in acts: + det = json.dumps(a["details"], default=str, indent=2) + L.append(f"- **{a['operation']}** {a['timestamp']:%Y-%m-%d %H:%M}") + L.append(f"```json\n{det}\n```") + L.append("") + return "\n".join(L) + + +def _render_retired(w: dict) -> str: + return (f"# {w['canonical_name']} — RETIRED\n\n" + f"- id: `{w['id']}`\n" + f"- retired_at: {w['retired_at']}\n" + f"- redirect_to: `{w['redirect_to']}`\n" + f"- summary: {w['summary']}\n\n" + f"This wiki was consolidated into its redirect target " + f"(`duplicate_of` / `consolidated_into` relations record the merge). " + f"It still resolves via GET /entities/{w['id']} but is dropped from ranking.\n") + + +def main() -> None: + OUT_DIR.mkdir(parents=True, exist_ok=True) + with get_conn() as conn: + wikis = _fetch_all_wikis(conn) + index = ["# Wiki review index", "", + f"{len(wikis)} wiki entities. Open each file below and judge against the checklist.", + "", + "| canonical_name | rev | refs | consistency | retired | file |", + "|---|---|---|---|---|---|"] + for w in wikis: + # id suffix keeps filenames unique (e.g. 'pytest' vs retired 'PyTest') + slug = _slug(w["canonical_name"]) + fname = f"{slug}-{w['id'][:8]}.md" + if w["retired_at"]: + (OUT_DIR / fname).write_text(_render_retired(w), encoding="utf-8") + index.append(f"| {w['canonical_name']} | {w['revision']} | - | - | YES | {fname} |") + continue + summarises = set(_summarises_targets(conn, w["id"])) + ok, _ = _consistency(w["content"] or "", summarises) + nrefs = len(parse_refs(w["content"] or "")) + (OUT_DIR / fname).write_text(_render(conn, w), encoding="utf-8") + index.append(f"| {w['canonical_name']} | {w['revision']} | {nrefs} | " + f"{'✓' if ok else '✗'} | no | {fname} |") + + index += ["", + "## Quality checklist (fill while reading each wiki)", + "", + "- [ ] **Grounded** — every claim traceable to a cited source entity (no hallucination)", + "- [ ] **Identity** — no third-party attribute transferred onto the subject; distinct people not fused", + "- [ ] **Honest uncertainty** — ambiguous data is represented as such, not fabricated into confidence", + "- [ ] **Summary/Disambiguation** — accurate; rewritten (not frozen) when better data exists", + "- [ ] **Consistency** — every cited inline ref has a summarises relation (column ✓)", + "- [ ] **Maintainer decision sane** — create/attach/skip/ambiguous rationale reasonable", + "- [ ] **No keyword-token sources** — cited refs are real fact/thought/source entities", + "- [ ] **Contradictions** — opposing sources reconciled or explicitly noted", + ""] + (OUT_DIR / "INDEX.md").write_text("\n".join(index), encoding="utf-8") + + print(f"Exported {len(wikis)} wikis to {OUT_DIR.resolve()} (open INDEX.md)") + + +if __name__ == "__main__": + main() diff --git a/braindb/wiki_scheduler.py b/braindb/wiki_scheduler.py new file mode 100644 index 0000000..1024278 --- /dev/null +++ b/braindb/wiki_scheduler.py @@ -0,0 +1,167 @@ +""" +Always-on wiki scheduler — ONE loop, like ingest_watcher.py (one interval). + +Per tick: + 1. POST /wiki/cron — cheap, pure SQL, no LLM. + 2. GET /wiki/jobs?status=pending — cheap, pure SQL, no LLM. The gate. + 3. if a pending `triage` job exists -> POST /wiki/maintain (one case, C1) + 4. if pending suggestion jobs exist -> POST /wiki/write, repeated to DRAIN + them (bounded) so consolidate/attach keep up instead of trickling. + 5. nothing pending -> NO LLM call this tick (idle == free). + +The expensive LLM endpoints are never called speculatively: a tick with +empty queues costs nothing. No multi-timer staggering, one env var. +""" +import logging +import os +import sys +import time +from concurrent.futures import ThreadPoolExecutor + +import requests + +API_URL = os.getenv("BRAINDB_API_URL", "http://localhost:8000") +INTERVAL = int(os.getenv("WIKI_INTERVAL", "60")) # one cadence, like the watcher +DRAIN_MAX = int(os.getenv("WIKI_DRAIN_MAX", "20")) # safety bound on /write per tick +# Per-tick concurrency: how many /wiki/write calls fire in parallel (vLLM +# continuous-batches them on the GPU; the DB layer is already safe via +# FOR UPDATE SKIP LOCKED on every claim and try_wiki_lock per wiki). +# `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 +# profile/provider. +WIKI_ENABLED = os.getenv("WIKI_ENABLED", "false").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 +# (vLLM, workstation): full-body wiki writes routinely run 6-15 min on this +# model, so a 600s deadline at the scheduler caused the client to give up +# WHILE the api kept working in the background — the write still committed, +# but the scheduler couldn't see the completion in time to drain the queue +# efficiently. With 1200s the client now waits long enough to see most +# writes finish, while still surfacing genuinely-stuck jobs as failures +# rather than blocking indefinitely. The api itself is not bounded by this +# value; this knob only controls how patient the scheduler's HTTP client is. +AGENT_TIMEOUT = int(os.getenv("WIKI_AGENT_TIMEOUT", "1200")) + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [wiki-scheduler] %(message)s", + datefmt="%H:%M:%S", + stream=sys.stdout, +) +log = logging.getLogger("wiki-scheduler") + +_SUGGESTION_TYPES = {"create", "attach", "consolidate"} + + +def wait_for_api(timeout: int = 90) -> bool: + deadline = time.time() + timeout + url = f"{API_URL}/health" + while time.time() < deadline: + try: + if requests.get(url, timeout=3).status_code == 200: + return True + except requests.RequestException: + pass + time.sleep(2) + return False + + +def _post(path: str, timeout: int) -> dict | None: + try: + r = requests.post(f"{API_URL}{path}", timeout=timeout) + if r.status_code == 200: + return r.json() + log.warning("%s -> %s: %s", path, r.status_code, r.text[:200]) + except requests.RequestException as e: + log.warning("%s request error: %s", path, e) + return None + + +def _pending_kinds() -> tuple[bool, bool]: + """(has_triage, has_suggestion) from ONE cheap SQL-only read. On error, + return (False, False) so we never fire LLM calls on uncertain state.""" + try: + r = requests.get( + f"{API_URL}/api/v1/wiki/jobs", + params={"status": "pending", "limit": 500}, + timeout=15, + ) + if r.status_code != 200: + log.warning("/jobs -> %s: %s", r.status_code, r.text[:200]) + return (False, False) + jobs = r.json() + except (requests.RequestException, ValueError) as e: + log.warning("/jobs read error: %s", e) + return (False, False) + has_triage = any(j.get("job_type") == "triage" for j in jobs) + has_sugg = any(j.get("job_type") in _SUGGESTION_TYPES for j in jobs) + return (has_triage, has_sugg) + + +def main() -> None: + if not WIKI_ENABLED: + log.info("wiki pipeline DISABLED (set WIKI_ENABLED=true to enable). Idle.") + # Sleep forever — keeps the container Up without restart-loop, and + # makes zero LLM/DB/api calls. Toggle via env + scheduler restart. + while True: + time.sleep(3600) + + log.info("waiting for API at %s ...", API_URL) + if not wait_for_api(): + log.error("API never came up; exiting") + sys.exit(1) + log.info("wiki scheduler ready (single loop, interval=%ss)", INTERVAL) + + while True: + try: + # 1. cron — cheap SQL, safe to run every tick. + res = _post("/api/v1/wiki/cron", timeout=60) + if res and res.get("triage_jobs_enqueued"): + log.info("cron: enqueued=%s pending_triage=%s", + res.get("triage_jobs_enqueued"), res.get("pending_triage_total")) + + # 2. cheap gate — decide whether any LLM work is warranted. + has_triage, has_sugg = _pending_kinds() + + # 3+4. fan out: ONE maintain (C1) in parallel with up to + # WRITE_PARALLELISM writes per batch; drain writes in batches + # until empty or DRAIN_MAX. The DB locks make this safe: + # FOR UPDATE SKIP LOCKED -> no double-claim on triage/suggestion + # try_wiki_lock(wiki_id) -> same-wiki writer contenders skip + # vLLM continuous-batches the concurrent inferences on the GPU. + with ThreadPoolExecutor(max_workers=WRITE_PARALLELISM + 1) as pool: + maintain_f = (pool.submit(_post, "/api/v1/wiki/maintain", AGENT_TIMEOUT) + if has_triage else None) + done = 0 + while has_sugg and done < DRAIN_MAX: + batch = min(WRITE_PARALLELISM, DRAIN_MAX - done) + fs = [pool.submit(_post, "/api/v1/wiki/write", AGENT_TIMEOUT) + for _ in range(batch)] + any_written = False + for f in fs: + res = f.result() + done += 1 + if res and res.get("written"): + any_written = True + log.info("write: wiki=%s mode=%s rev=%s", + res.get("wiki_id"), res.get("mode"), res.get("revision")) + if not any_written: + break # queue empty or all targets locked -> stop draining + if maintain_f is not None: + res = maintain_f.result() + if res and res.get("claimed"): + log.info("maintain: %s", res.get("result")) + + # 5. nothing pending -> no LLM call happened this tick (free). + except Exception as e: + log.exception("loop error: %s", e) + time.sleep(INTERVAL) + + +if __name__ == "__main__": + main() diff --git a/docker-compose.yml b/docker-compose.yml index 491fb6d..da218f6 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,6 +15,10 @@ services: DEEPINFRA_API_KEY: ${DEEPINFRA_API_KEY:-} VLLM_API_KEY: ${VLLM_API_KEY:-} AGENT_VERBOSE: ${AGENT_VERBOSE:-false} + # Orphan freshness gate (the orphan SQL runs in this api process, not in + # the scheduler): an entity is wiki-eligible only once created_at is + # older than this many minutes, so still-ingesting subjects settle first. + WIKI_FRESHNESS_MINUTES: ${WIKI_FRESHNESS_MINUTES:-30} extra_hosts: # Lets self-hosted profiles (e.g. vllm_workstation) reach a server bound # to the Docker host's loopback. Docker Desktop sets this implicitly; @@ -24,8 +28,16 @@ services: - "${API_PORT:-8000}:${API_PORT:-8000}" volumes: - .:/app + # Note: NO `--reload` in the api command. With the bind mount (`.:/app`) + # in dev, `--reload` causes uvicorn to restart on every .py edit, which + # interrupts in-flight LLM calls (the scheduler logs `Connection refused` + # during the ~20-30s embedding-model reload). Code changes are applied + # explicitly via `docker compose up -d --no-deps --force-recreate api` — + # the operator picks the moment. Anyone who wants dev-style live reload + # can override this command via `docker compose run` or a personal + # `docker-compose.override.yml`. command: > - sh -c "alembic upgrade head && uvicorn braindb.main:app --host 0.0.0.0 --port ${API_PORT:-8000} --reload" + sh -c "alembic upgrade head && uvicorn braindb.main:app --host 0.0.0.0 --port ${API_PORT:-8000}" watcher: build: . @@ -42,6 +54,26 @@ services: - .:/app command: python -m braindb.ingest_watcher + # Always-on wiki maintenance sidecar — same posture as `watcher`. It loops + # cron -> maintain -> write so wikis self-organise from entities with zero + # manual steps. To run without it (e.g. cost control), start the stack + # excluding this service or scale it to 0 — exactly as you would the watcher. + wiki_scheduler: + build: . + container_name: braindb_wiki_scheduler + restart: unless-stopped + depends_on: + - api + networks: + - local-network + environment: + BRAINDB_API_URL: http://api:${API_PORT:-8000} + WIKI_ENABLED: ${WIKI_ENABLED:-false} + WIKI_INTERVAL: ${WIKI_INTERVAL:-60} + volumes: + - .:/app + command: python -m braindb.wiki_scheduler + networks: local-network: external: true diff --git a/docs/maintainer-agent-plan.md b/docs/maintainer-agent-plan.md new file mode 100644 index 0000000..4fc3a2e --- /dev/null +++ b/docs/maintainer-agent-plan.md @@ -0,0 +1,268 @@ +# BrainDB Wiki System — cron / maintainer / writer (living design doc) + +> **Living document.** This is the iterated source of truth and is updated as +> implementation proceeds. The frozen, as-approved snapshot is +> [`maintainer-agent-plan2.md`](maintainer-agent-plan2.md) — do not edit that one. + +> **Operating model (current):** wiki maintenance is **hands-off, default-on**. +> `wiki_scheduler` is a normal always-on compose sidecar (same posture as the +> ingest `watcher`, no opt-in profile) that loops cron(~20m) → maintain → +> write autonomously. The `/api/v1/wiki/{cron,maintain,write}` endpoints are +> **dev/debugging only**, never the operating procedure. The maintainer +> staleness guard + skip-self-clearing keep it idempotent and cheap. Disable +> for cost like the watcher (exclude the service / scale to 0). Inspection +> (`export_wikis`) is an optional read-only dev tool, outside the operating +> path; no test scaffolding lives in operational modules. + +## ⚠ Correction applied (supersedes earlier "gate/manifest/ledger" design) + +The first implementation inserted programmatic algorithms between the process +and the LLM that destroyed its grasp of reality (e.g. "Subject A is an ML +engineer", "Koutsoumpos is a marine engineer", "Artificial Intelligence" = +one NVIDIA earnings call). Root cause: per-orphan pinhole context, an +accounted-change gate that *blocked self-correction*, a rigid JSON manifest, a +code-generated references ledger, and prompts that never told the LLM to +investigate. **Principle reinstated: programmatic = process / queue / +bookkeeping / commands / reversibility ONLY; the LLM owns all +understanding/identity/content/revision and must research with the existing +tools.** + +What changed in code (net negative LOC, no new machinery): +- **Deleted** `accounted_change_gate`, `regenerate_references_ledger`, + `split_sections`, `_structural_errors`, the JSON manifest contract, the + section-hash guard, and `keywords=[canonical_name]`. +- `apply_manifest_relations` → **`reconcile_summarises_additive`**: creates a + `summarises` edge per inline `[[ref:]]`; **never deletes/re-types behind the + LLM** (the LLM calls `delete_relation` itself if needed). +- Writer returns **only the body** (`<<>>`); consolidate adds one + command line `<<>>`. Body persisted **verbatim**; prior + version snapshotted to `wiki_revise` (reversible). Wiki `keywords` read from + the LLM's own `` line, else empty. +- Maintainer & writer prompts rewritten: **research-first** with + `recall_memory` + `delegate_to_subagent` (SQL = rare aggregation exception); + identity/scope discipline (no third-party attribute transfer, no invented + identity, distinct entities stay distinct); **represent ambiguity** instead + of fabricating; writer **MUST revise** summary/disambiguation/scope on + better data (self-healing); no keyword-token citations. Agent turns raised + so it can actually investigate/delegate. +- New maintainer action **`ambiguous`** (treated as a deliberate skip → + self-clears via `run_cron`). +- **Tool-priority** correction applied everywhere: `system_prompt.md` + (TOOL PRIORITY rule + Example 3 rewritten), `skills/braindb/SKILL.md`, + `skills/braindb-agent/SKILL.md`, `CLAUDE.md`, `BRAINDB_GUIDE.md`, + `export_wikis.py` consistency/checklist. `recall_memory`/`/memory/context` + + subagents are the default; `/memory/sql` is an aggregation-only exception. + +Frozen snapshot `maintainer-agent-plan2.md` is intentionally left as the +original approved record. The cron / claim / skip-self-clear / soft-retire / +snapshot bookkeeping is unchanged. + +### Self-heal test result (Subject A) — honest + +- **Structural fix: PASS.** No cage; writer revises freely; prior versions + snapshotted (`wiki_revise` rev 1→4, reversible); LLM authored body/keywords/ + ledger; additive reconcile; writer **did** research via `recall_memory`. +- **Cooperative/radical policy: PASS (mechanically).** With the + cooperative-default + strong-conviction + mandatory-subagent-confirmation + prompt, the writer stayed cooperative, **detected the conflation**, and + **delegated a subagent** to independently resolve identity before acting — + exactly the requested guardrail. +- **Correctness (first attempt): FAIL.** Then fixed — see RESOLVED below. The + earlier "root cause is irreducibly DATA identity" verdict was **wrong**: it + was a *process* failure (anchored subagent delegation + the existing wrong + page acting as a top-ranked recall attractor + greedy positive + same-first-name matching + richness-over-correctness). + +### RESOLVED — fix verified (2026-05-16) + +Three non-bloat changes (prompt + one-time safe reset, no code/gates): +1. **Non-anchored resolution delegation** (writer prompt): the writer MUST + delegate IDENTITY RESOLUTION giving the subagent **only raw `id: content` + facts** — never the page name, its claims, or an expected answer — with + explicit DISQUALIFIERS and an AMBIGUOUS bucket; then writes only the + resolved subject's facts. +2. **Exclusion + circuit-breaker** (writer & maintainer prompts): a shared + first-name fact not uniquely tied is AMBIGUOUS → excluded; correctness over + richness; shrink to an honest stub if unresolved. +3. **Safe clean slate**: deleted wiki layer only (7 wikis, 774 jobs, + wiki-only relations). Knowledge byte-identical (fact 134, thought 23, + source 8, datasource 7, keyword 603, activity_log 1199 — unchanged). + +Re-created "Subject A" via the corrected flow (logs confirm the +verbatim non-anchored template, no leakage). Result page: +- Summary: "A Greek youth and natural tinkerer born in 2011 who aspires to + become a boat mechanic." ✓ +- Disambiguation: explicitly "the nephew of the ML engineer Dimitrios + Koutsoumpos; **not** the professional AI/ML engineer at CityFalcon." ✓ +- The ambiguous professional "Dimitris" facts (ML engineer / 18-yr investing + / coaching) were **correctly excluded**, not fused. Consistency ✓. + +Conclusion: conflation was a **process** failure, now fixed with prompt + +safe reset only — no new code, gates, or bloat. Caveats: verified on the +Subject A case in create mode; the ~700 triage backlog still to be drained, +and per-wiki runs are slow (recall + a real resolution subagent on +gemma-4-31B → minutes each → this is background-scheduler work, not +interactive). Upstream fact-level identity anchoring remains a *possible +future enhancement*, but is **not required** to get correct pages. + +## What this is + +A wiki layer inside BrainDB. Wikis are synthesised, human-readable pages +(`entity_type = 'wiki'`) about one concept each, built from the +keyword/thought/fact entities that concern it — Karpathy-style, but stored as +entities (not files) and kept consistent with the graph. + +Three-stage pipeline: + +1. **Cron** (`POST /api/v1/wiki/cron`) — read-only orphan scan; enqueues one + `triage` job per entity not yet connected to a wiki. Idempotent. +2. **Maintainer** (`POST /api/v1/wiki/maintain`) — processes **exactly one** + triage case per call (C1); the existing agent decides + attach / create / consolidate / skip and a structured suggestion job is + persisted. +3. **Writer** (`POST /api/v1/wiki/write`) — one wiki per call. The agent + authors the body + a change manifest; a deterministic **accounted-change + gate** validates it; the references ledger and `summarises` relations are + reconciled from the body+manifest; the prior revision is snapshotted. + +Inspection: `GET /api/v1/wiki/jobs`. Always-on driving (Stage 2): the +`wiki_scheduler` sidecar, **opt-in** via the `wiki` compose profile. + +## Governing constraints + +- **C1 — per-case maintainer.** Never a bulk dump; one orphan per invocation. +- **C2 — no programmatic destruction without LLM awareness.** Deterministic + code is limited to read-only detection, safe queue plumbing, and additive + bookkeeping that mirrors LLM-authored content / executes the LLM's explicit + manifest. Every consequential change is logged and reversible. +- **C3 — reuse existing APIs; no bloat.** Detection/ranking/contradiction + context all go through the existing `recall_memory` / `/memory/context` + scoring. No new similarity query, scoring heuristic, or embedding path. + +## Writer robustness (the accounted-change model) + +Surgical add/modify/**delete** is allowed; *undeclared* or *accidental* loss +is impossible. The writer returns body + manifest +(`added_refs` / `removed_refs[{ref,reason,note,prior_text}]` / +`modified_sections` / `contradictions_resolved` / `canonical_wiki_id`). The +gate (deterministic, in-transaction): + +1. every dropped ref must be declared in `removed_refs` with a valid reason; + every gained ref in `added_refs`; +2. on `attach`, non-targeted sections must be byte-identical; +3. structural validation (5 required section anchors, Summary ≤ 280, + Disambiguation present, every surviving ref resolves); +4. any violation → rollback, job → `pending`, retry with the defect; capped + by `attempts` → `failed` (surfaced via `GET /jobs`). + +Provenance preserved: a declared removal re-types the `summarises` edge +(`contradicted` → `contradicts`) rather than vanishing; prior content is +snapshotted to the activity log (`wiki_revise`), so deleted ≠ destroyed. +The `section:references` ledger is machine-regenerated from parsed refs, so +inline tokens, the ledger, and the SQL relations cannot disagree. + +Consolidation is LLM-performed: duplicates are spotted via the maintainer's +normal `recall_memory` (no dedup query); the writer picks the canonical and +the loser is soft-retired (`importance=0`, `retired_at`, `redirect_to`, +`duplicate_of` + `consolidated_into` edges) — still resolvable, dropped from +ranking, and self-clearing (the maintainer is prompted to skip marked pairs). + +## What was built + +| File | Role | +|---|---| +| `alembic/versions/005_wiki_system.py` | additive migration: `wiki` type, `wikis_ext`, `wiki_job` (down_revision 004) | +| `braindb/schemas/entities.py` | `WikiCreate/Read/Update` + `AnyEntityRead` | +| `braindb/schemas/relations.py` | `summarises`, `not_duplicate`, `duplicate_of`, `consolidated_into` | +| `braindb/routers/entities.py` | wiki CRUD; `ENTITY_SELECT`/`_flatten` extended (`member_keyword_ids::text[]`) | +| `braindb/services/context.py` | `DECAY_RATES["wiki"]`, `EXT_QUERIES["wiki"]` | +| `braindb/config.py` | `decay_rate_wiki = 0.0` | +| `braindb/services/wiki_jobs.py` | all non-LLM plumbing: orphan/cron, claim (SKIP LOCKED), dedupe_key, gate, ledger, reconcile, snapshot, soft-retire, advisory lock | +| `braindb/routers/wiki.py` | `/cron` `/maintain` `/write` `/jobs` | +| `braindb/agent/prompts/wiki_maintainer_prompt.md` | per-case triage → structured suggestion | +| `braindb/agent/prompts/wiki_writer_prompt.md` | skeleton contract + manifest + consolidate | +| `braindb/wiki_scheduler.py` + compose `wiki_scheduler` (profile `wiki`) | Stage-2 always-on, opt-in | + +No new Python dependencies. The agent itself is reused unchanged (no new +agent factory) — prompts are passed as the query to `run_agent_query`. + +## Verification status (DeepInfra profile) + +- Migration 005 auto-applies on startup (rev `005`, both tables present). ✓ +- Wiki CRUD + no retrieval regression; wiki participates in ranking, existing types unaffected. ✓ +- Cron: 757 triage enqueued; re-run → 0 (idempotent). ✓ +- Maintainer: one case/call; `create`/`skip` decisions; deterministic dedupe_key; cron does not re-enqueue in-flight orphans. ✓ +- Writer `create`: skeleton anchors, inline refs, machine ledger, `summarises` relation — all consistent. ✓ +- Accounted-change gate (deterministic, no LLM): undeclared drop/section-change rejected, declared changes pass, bad structure rejected. ✓ +- Consolidation: LLM picked canonical, loser soft-retired + provenance edges, canonical ranks / loser→0, still resolvable. ✓ +- Scheduler: loop healthy, drives cron on schedule; opt-in profile. ✓ + +Not yet exercised live (deferred to a broader end-to-end pass; needs +maintainer-produced attach jobs and is LLM-cost-bearing): the live `attach` +path with restorability from the `wiki_revise` log, and a live +contradiction-resolution edit. The deterministic guarantees behind them are +unit-verified. + +## Quality trial (10-case controlled batch) — findings + +Tool: `docker compose exec -T api python -m braindb.tools.export_wikis` +(read-only; writes `data/wiki_review/*.md` + `INDEX.md`; gitignored). + +**Mechanics — solid.** 10 maintain calls → 2 create / 4 attach / 4 skip +(sane distribution, coherent rationales). Writers produced/updated wikis with +all skeleton anchors. **Consistency ✓ on every wiki** (inline refs = ledger = +`summarises` relations). The **accounted-change gate fired live**: an attach +that changed the `sources` section without declaring it was rejected and +requeued; the retry passed (no bad write persisted; `attempts` capped). +Skip self-clearing verified (post-trial cron enqueued 0; `failed` triage still +retries). Manifest now logged in `wiki_write` activity (writer reasoning is +inspectable). + +**Content — weak, and the export proves why (the important finding).** The +orphans being wiki-ified are overwhelmingly **bare keyword entities** whose +`content` is an auto-generated token (e.g. `_pytest_82a2e09b`). The writer has +no real substance to ground on, so it: (a) writes fluent prose from world +knowledge, and (b) **cites those keyword-token entities as if they were +sources** — even fabricating a sentence ("supported by various internal +identifiers [[ref:…]] [[ref:…]]") to wrap junk refs. The wikis are +structurally perfect and provenance-consistent but **not evidence-grounded**. +Scaling now would mass-produce fluent-but-hollow pages citing tokens. + +Root cause is *not* pipeline code (which works). It is **what is fed in**: the +maintainer/writer act on the bare keyword, not the keyword's connected +facts/thoughts. Options to decide before scaling: +- writer pulls the keyword's `tagged_with` fact/thought neighbourhood (via the + existing `recall_memory`/`view_tree`) as the real sources, and the prompt + forbids citing `keyword`-type entities as provenance; +- and/or the maintainer `skip`s keyword orphans that have no real + fact/thought substance behind them (only wiki-ify concepts with evidence). + +## Known follow-ups (decide before scaling) + +1. **Skip self-clearing — DONE.** `run_cron()` now excludes orphans with a + `rejected` triage job (deliberate skip). Permanent like `not_duplicate`; + `failed` triage still retries. No schema change. +2. **Grounding (NEW — highest priority).** See "Quality trial" above. Decide + the sourcing fix before any scale-up; mechanics are ready, content is not. +3. **Backlog cost.** ~750 pending triage × one agent call each. Scheduler is + opt-in; consider prioritising high-importance / evidence-bearing orphans. +4. **LLM profile.** `.env` switched `vllm_workstation → deepinfra` for + verification (local vLLM down). Switch back when available. +5. Live contradiction-resolution edit still not exercised (no opposing + sources in the trial corpus). Deterministic guarantee unit-verified. + +## Operational notes — review tooling + +- Inspect quality any time: `docker compose exec -T api python -m + braindb.tools.export_wikis`, then open `data/wiki_review/INDEX.md` and the + per-wiki files. Each file shows body, the consistency verdict, **provenance + (cited entities' real content — judge grounding here)**, maintainer + rationale, writer manifest, and revision snapshots. + +## Operational notes + +- Stage 1 is manual: hit the endpoints by hand. Nothing wiki-related runs on + startup (the existing ingest watcher is untouched). +- Enable always-on: `docker compose --profile wiki up -d wiki_scheduler` + (env: `WIKI_CRON_INTERVAL`, `WIKI_MAINTAIN_INTERVAL`, `WIKI_WRITE_INTERVAL`). +- Migrations run automatically on `api` startup. diff --git a/docs/maintainer-agent-plan2.md b/docs/maintainer-agent-plan2.md new file mode 100644 index 0000000..e6b96a6 --- /dev/null +++ b/docs/maintainer-agent-plan2.md @@ -0,0 +1,355 @@ +# BrainDB Wiki System — cron / maintainer / writer + +> **Frozen snapshot.** This is the verbatim plan as approved before implementation +> began. It is an immutable historical reference — do **not** edit it as the design +> evolves. The living design doc is `maintainer-agent-plan.md`. + +## Context + +BrainDB stores a graph of typed entities. Keyword entities act as soft "entity hubs" +(everything about a thing gets `tagged_with` that keyword), but there is **no +synthesised, human-readable page per concept** the way Karpathy's LLM-wiki has. The +prior draft (`docs/maintainer-agent-plan.md`) framed this as keyword-dedup-first. The +user has reframed it as a **three-stage pipeline** and set two hard constraints +(below) that supersede the prior draft. + +1. **Cron** — read-only: find keyword/thought/fact entities not connected to any wiki (orphans) and enqueue **one triage case per orphan**. +2. **Maintainer** — pulls **one case at a time** (never the whole batch), researches it against existing wikis + graph via the current APIs, and emits a structured suggestion job for *that case*: attach / create / possible-duplicate. +3. **Wiki writer** — invoked **per wiki**; the LLM consumes that wiki's suggestion jobs and writes/updates the wiki, managing relations itself through existing tools. + +### Two governing constraints (from user feedback) + +- **C1 — Per-case maintainer.** The maintainer must reason about a single orphan case per invocation. The cron only *enqueues* cases; it never hands the maintainer a bulk dump. +- **C2 — No programmatic destruction without LLM awareness.** No autonomous SQL procedure may delete/repoint relations, retire/merge entities, or change importance behind the LLM's back. The deterministic layer is restricted to: **(1) read-only detection** (orphans — *suggestions only*), **(2) safe non-destructive job-queue plumbing** (enqueue, claim, idempotency, status), and **(3) at most additive bookkeeping that exactly mirrors LLM-authored content**. Every consequential graph mutation is performed by the LLM via existing tools — visible, logged, reversible. Postgres FK `ON DELETE CASCADE` self-healing is acceptable (it is correct DB behaviour we do not author); the resulting dead inline token is *flagged for an LLM*, never auto-edited. + +- **C3 — Reuse the existing APIs; no bloat.** BrainDB already has sophisticated search/scoring (`/memory/context` & the `recall_memory`/`quick_search`/`view_tree`/`search_sql` tools: combined fuzzy + full-text + keyword-embedding, graph traversal, temporal decay, `final_rank`). Every stage that needs to *find*, *rank*, or *compare* anything **must call that existing infra**. Do not write a new similarity query, scoring heuristic, or embedding path that duplicates what these already do. New code is allowed only for: the additive migration, the `wiki_job` queue plumbing, the deterministic non-destruction gate (the safety guarantee C2 requires), and prompts. If a proposed piece of code re-implements search/scoring, it is bloat and is cut. + +Goals: wikis live **inside the DB** (entities, not files); reuse existing machinery +(embeddings, graph traversal, the agent HTTP endpoint, relations, activity log); +**must not regress** existing endpoints, retrieval, or the ingest watcher; agent track +first, Claude-Code-skill track later. + +This file records the **recommended** path only; alternatives/trade-offs are in the +conversation. + +--- + +## Key design decisions + +| # | Decision | Choice | +|---|---|---| +| D1 | Wiki granularity | Born one-per-keyword; collapsed toward per-canonical-cluster **by LLM-driven consolidation** over time. `wikis_ext.member_keyword_ids` is the cluster. | +| D2 | Where jobs live | New `wiki_job` table with lifecycle + deterministic `dedupe_key` partial-unique index for idempotency. Two job sources: `triage` rows (cron, one per orphan) and suggestion rows (maintainer: `attach`/`create`/`consolidate`). | +| D3 | Orchestration | Manual endpoints first (`/api/v1/wiki/{cron,maintain,write,jobs}`) driving the existing `POST /api/v1/agent/query`. Maintainer endpoint processes **one triage case per call** (C1). Separate `wiki_scheduler` sidecar (clone of `ingest_watcher.py`) only after endpoints are verified. Ingest watcher never touched. | +| D4 | Inline ref ↔ SQL consistency | Body is source of truth. Writer LLM emits `[[ref:UUID]]` **and** owns its relations via `create_relation`/`delete_relation`. A reconcile step is **additive + advisory only**: it may add a relation that exactly mirrors a ref the LLM wrote; it *flags* (never deletes/repoints) drift as an LLM fix-up case (C2). | +| D8 | Writer robustness | Surgical add/modify/**delete** is allowed but **accounted-for**: writer returns body + a change manifest; an **accounted-change gate** rejects+retries any *undeclared* drop/add or out-of-scope section change (blocks accidental destruction, permits justified deletion). Mandatory contradiction-gathering via existing recall; prior revision snapshotted to the activity log (deleted ≠ destroyed). Fixed contract + template + validation make style robust. See "Wiki document contract" below. | +| D5 | Duplicate wikis | **No new search/scoring code.** Detection is a by-product of the per-case maintainer's existing `recall_memory` call (= `/memory/context`: text + keyword-embedding + graph + decay + `final_rank`, all already built). If that recall surfaces an existing wiki very close to the case's concept, the maintainer emits a `consolidate` suggestion. The **wiki-writer LLM performs the merge** via existing tools, logged, reversible — no `merge_wikis()` SQL, no bespoke cosine query. `not_duplicate`/`duplicate_of` are plain relations the LLM sees via existing relation/graph tools and is prompted to respect (self-clearing without a custom SQL filter). | +| D6 | Summary / disambiguation / language | Reuse `entities.summary` for the cheap one-line header; `wikis_ext.disambiguation` + `wikis_ext.language` (mirrors `datasources_ext.language`). | +| D7 | Driver | In-house agent first (new prompts + reuse `/agent/query`). Claude-Code skill later; persisted `wiki_job` rows are the shared contract so both drivers interoperate. | + +--- + +## Schema — single additive migration `005_wiki_system.py` (`down_revision = "004"`) + +Mirrors the `004` CHECK-rewrite pattern. Purely additive; no backfill; existing rows untouched. + +```sql +ALTER TABLE entities DROP CONSTRAINT entities_entity_type_check; +ALTER TABLE entities ADD CONSTRAINT entities_entity_type_check + CHECK (entity_type IN ('thought','fact','source','datasource','rule','keyword','wiki')); + +CREATE TABLE wikis_ext ( + entity_id UUID PRIMARY KEY REFERENCES entities(id) ON DELETE CASCADE, + canonical_name VARCHAR(500) NOT NULL, + disambiguation TEXT, + language VARCHAR(10) DEFAULT 'en', + member_keyword_ids UUID[] DEFAULT '{}', + revision INT DEFAULT 1, + last_synthesised_at TIMESTAMPTZ, + retired_at TIMESTAMPTZ, -- set by the LLM via tools, not by SQL procedure + redirect_to UUID REFERENCES entities(id) ON DELETE SET NULL +); +CREATE INDEX wikis_ext_canonical_idx ON wikis_ext (lower(canonical_name)); +CREATE INDEX wikis_ext_member_kw_idx ON wikis_ext USING GIN (member_keyword_ids); + +CREATE TABLE wiki_job ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + job_type VARCHAR(20) NOT NULL + CHECK (job_type IN ('triage','attach','create','consolidate')), + status VARCHAR(12) NOT NULL DEFAULT 'pending' + CHECK (status IN ('pending','assigned','done','rejected','failed')), + target_wiki_id UUID REFERENCES entities(id) ON DELETE CASCADE, -- NULL for triage/create + entity_ids UUID[] NOT NULL DEFAULT '{}', -- triage: the single orphan (+context anchors) + dedupe_key TEXT NOT NULL, + rationale TEXT, + proposed_name VARCHAR(500), + batch_id UUID, + created_at TIMESTAMPTZ DEFAULT now(), + assigned_at TIMESTAMPTZ, + completed_at TIMESTAMPTZ, + attempts INT DEFAULT 0, + last_error TEXT +); +CREATE UNIQUE INDEX wiki_job_dedupe_active_idx + ON wiki_job(dedupe_key) WHERE status IN ('pending','assigned'); +CREATE INDEX wiki_job_status_idx ON wiki_job(status); +CREATE INDEX wiki_job_target_idx ON wiki_job(target_wiki_id); +``` + +**No new embedding path.** Wikis are found through the *existing* retrieval infra: +the body is full-text indexed automatically (the `search_vector` trigger), the wiki +is `summarises`-linked to its member entities and `tagged_with` its keywords, and +keyword embeddings + graph + `final_rank` already route queries to it. We do **not** +add a wiki-embedding generator or a wiki-vs-wiki cosine query. + +`RELATION_TYPES` (Python-side only, no DB constraint) gains: `summarises`, +`not_duplicate`, `duplicate_of`, `consolidated_into`. + +--- + +## Inline reference syntax + additive/advisory reconcile + +Token in `entities.content`: `[[ref:ENTITY_UUID]]` or `[[ref:ENTITY_UUID|display text]]`. +Regex: `\[\[ref:([0-9a-f-]{36})(?:\|[^\]]*)?\]\]`. + +**The writer LLM is responsible for its relations.** The writer prompt instructs it +to call `create_relation` (`wiki --summarises--> entity`, relevance 0.9) for each +entity it cites and `delete_relation` when it removes a citation. The deterministic +`reconcile_wiki_refs(conn, wiki_id, body)` is a **safety net, additive + advisory only**: + +1. `cited` = UUIDs parsed from body that exist in `entities`. +2. `current` = `to_entity_id` where `from=wiki_id AND relation_type='summarises'`. +3. **Add**: insert `summarises` for `cited - current` (`ON CONFLICT DO NOTHING`) — mirrors what the LLM wrote in the body. +4. **Declared removals**: for `current - cited` where the UUID is in `manifest.removed_refs`, the writer has already re-typed/handled the relation via tools (gate step 6); the reconciler just confirms consistency. For `current - cited` *not* in the manifest, that is an undeclared drop → already rejected by the gate (this branch should never persist). Dangling refs (cited UUID not in `entities`) → fix-up `triage` job + `log_activity('wiki_ref_drift', ...)`. The reconciler itself still never deletes/repoints — the *writer* does, declared and via tools. + +Cited entity later genuinely deleted → FK `ON DELETE CASCADE` removes the relation +(correct DB behaviour, not our code). The orphaned `[[ref:]]` token is flagged for an +LLM rewrite — prose is never blind-edited. + +--- + +## Wiki document contract + writer robustness + +The writer's safety does **not** rest on model judgement. Robustness is structural. + +### Fixed document skeleton (every wiki, enforced) + +``` + +# {canonical_name} +> **Summary:** {one line, ≤ 280 chars — kept short on purpose} +> **Disambiguation:** {what this is / is NOT; the true meaning(s)} + + ...prose with [[ref:UUID]]... + ...dated claims, each carrying [[ref:UUID]]... + ...conflicts flagged inline with BOTH refs... + ...narrative provenance... + AUTO-GENERATED — do not hand-write +- [[ref:UUID]] — one-line what this entity contributes +``` + +Anchors are HTML comments → invisible in render, deterministically splittable. The +`section:references` ledger is **machine-generated** from the parsed `[[ref:]]` set +on every save (the LLM writes refs inline in prose; it never authors the ledger), so +inline tokens, the ledger, and the `summarises` SQL relations all derive from one +parse and **cannot disagree**. + +### Surgical editing IS allowed — the rule is "accounted-for", not "append-only" + +The writer **must** be able to revise a specific part: rewrite a sentence, drop a +claim that is wrong/superseded, resolve a contradiction by removing the losing side. +The earlier "may only add refs" idea is wrong — it would freeze bad content forever. +The real guarantee is: **every removal/modification is deliberate, justified, and +recoverable; nothing is lost silently or accidentally.** + +The writer returns two things, not one: the new body **and a structured change +manifest**: + +``` +{ "added_refs": [UUID, ...], + "removed_refs": [{ "ref": UUID, "reason": "superseded|contradicted|wrong|merged|irrelevant", + "note": "one line", "prior_text": "the sentence/para removed" }], + "modified_sections": ["timeline", ...], + "contradictions_resolved": [{ "kept": UUID, "demoted": UUID, "how": "..." }] } +``` + +Edit mode is still set by **job type, not the model** (`create` = template; +`attach` = section-scoped, untargeted sections byte-identical; `consolidate`/ +resynthesise = full rewrite), but within the targeted scope the writer may freely +add/modify/delete **provided the manifest accounts for it**. + +### Accounted-change gate (the actual guarantee) + +Around every writer save, in the same transaction: + +1. `R_before` / `R_after` = `{[[ref:UUID]]}` parsed from old / new body. +2. `dropped = R_before − R_after`, `gained = R_after − R_before`. +3. **Every** UUID in `dropped` must appear in `manifest.removed_refs` with a valid `reason`; **every** UUID in `gained` must appear in `manifest.added_refs`. An undeclared drop or add ⇒ violation (this is what blocks *accidental* destruction while *allowing* declared deletion). +4. **Section guard** (`attach` only): non-targeted sections hash-identical; a change outside `manifest.modified_sections` ⇒ violation. +5. **Structural validation**: required anchors present; `summary` ≤ 280 chars; `disambiguation` non-empty; every surviving `[[ref:UUID]]` resolves in `entities`. +6. **Provenance is preserved, not erased.** A declared removal does **not** silently delete the entity or just drop the `summarises` edge into the void. The writer must, via existing tools, either (a) replace `summarises` with a typed relation that records the judgement — `contradicts` (this member opposes the consensus), `challenges`, or keep a low-relevance historical `summarises` — or (b) raise a fix-up `triage` job if the source entity itself looks wrong. The writer never deletes *other* entities; it only re-types its own link and explains why. `removed_refs[].prior_text` + reason are written to the wiki `notes` / activity log. +7. Any violation ⇒ **rollback**, job → `pending` with `last_error`, retry with the explicit defect ("undeclared drop of X", "section Z changed but not in manifest", "summary too long"). Capped by `attempts`; exhaustion ⇒ `failed`, surfaced via `GET /jobs`. Never a silent bad write. + +### Contradiction handling (the writer must reason about opposition) + +Before editing, the writer is **required** to gather opposition context using the +**existing infra** (C3): `recall_memory` / `view_tree` / `view_entity_relations` +over the member entities surface any `contradicts`/`challenges` relations and +semantically opposed claims (the existing scoring already clusters them). The writer +prompt mandates a populated `section:contradictions`: every detected opposition is +either (a) reconciled in prose with **both** refs kept, or (b) one side explicitly +demoted via the manifest (`contradictions_resolved`) with reasoning — never one side +silently dropped. The gate cross-checks: a UUID that vanished and was part of a +detected contradiction must appear in `contradictions_resolved`. + +### Reversibility (deleted ≠ destroyed) + +Every writer save first snapshots the prior `content` + parsed refs into the activity +log (`operation='wiki_revise'`, with `revision` n→n+1) before mutation. So any +removal — even a correct one — is auditable and restorable from the log. "Edited a +specific part / removed something that doesn't make sense" is fully supported; +"content vanished with no record or reason" is structurally impossible. + +This makes "surgical edits yes, destruction no" a checked invariant, not a hope — +true regardless of which LLM profile is active. + +### Style robustness levers (in `wiki_writer_prompt.md`) + +- The skeleton above is the mandatory output contract (sections, order, anchors, ref syntax, tone: encyclopedic, third-person, dated, contradictions flagged with both refs, every non-trivial claim carries a `[[ref:]]`). +- A **golden template** for `create` so structure is identical across all wikis from day one. +- A **few-shot exemplar**: one well-formed wiki + a before/after `attach` showing existing content preserved and the new member integrated. +- Deliberately **small focused context** (one wiki's body + only that wiki's new members) — the maintainer being per-case keeps the writer's input bounded; focused context is itself a major robustness lever. + +--- + +## Pipeline mechanics + +**Cron** (`POST /api/v1/wiki/cron`, pure SQL, read-only + safe enqueue, no LLM): +select keyword/thought/fact entities with no `summarises`/member link to any wiki and +not already in an active job; for **each** orphan insert one `triage` `wiki_job` +(`dedupe_key = triage:`, `ON CONFLICT DO NOTHING`). Returns counts. +Idempotent and non-destructive by construction. + +**Maintainer** (`POST /api/v1/wiki/maintain` — processes **one** triage case per call, +C1): claim a single `triage` job (`FOR UPDATE SKIP LOCKED`, LIMIT 1). Build a focused +prompt for *that one orphan only* (its content + its graph neighbourhood via +`recall_memory`/`view_tree`, plus the candidate existing wikis' `summary`/ +`disambiguation` found via search). The agent decides for this case: attach to wiki W +/ create new wiki / flag possible duplicate of wikis. The service parses the agent's +structured result and writes the corresponding suggestion job (`attach`/`create`/ +`consolidate`) with a service-computed `dedupe_key` +(`attach::` / `create:` / +`consolidate:`, `ON CONFLICT DO NOTHING`), then closes the triage job +(`done`/`rejected`). A loop/sidecar calls this endpoint repeatedly to drain the +triage queue one case at a time. + +**Writer** (`POST /api/v1/wiki/write {wiki_id? | job_ids? | next_pending}`): pick one +target (a wiki id, or a `create`/`consolidate` job group). In one `get_conn()` +transaction: `SELECT pg_try_advisory_xact_lock(hashtext('wiki:'||id))` → claim that +target's pending suggestion jobs (`FOR UPDATE SKIP LOCKED`) → **snapshot prior +`content`+refs to activity log (`wiki_revise`)** + per-section hashes → one agent run +with a focused prompt (current body pre-split by anchors for `attach` + cited members ++ **mandatory contradiction context gathered via existing `recall_memory`/`view_tree`**; +edit mode chosen by job type) → the LLM returns **new body + change manifest** and +**calls `create_relation`/`delete_relation`/`update_entity` itself** for citations +and declared removals → **accounted-change gate** (every drop/add declared in +manifest; section guard; structural validation; contradiction cross-check; on +failure: rollback, job→`pending`, retry with defect, cap by `attempts`) → +regenerate `section:references` ledger from parsed refs → additive +`reconcile_wiki_refs` consistency check → bump `revision`, set `last_synthesised_at` +→ finalise jobs → `log_activity('wiki_write', ...)`. + +**Consolidation reuses existing scoring; LLM-performed (C2).** There is **no +dedicated dedup query**. Duplicate detection falls out of the maintainer's normal +per-case `recall_memory` (the existing `/memory/context` scoring — text + +keyword-embedding + graph + decay + `final_rank`). When that recall returns an +existing wiki ranked very close to the case's concept, the maintainer emits a +`consolidate` suggestion. It already has the markers in view (the recall's graph +neighbourhood / `view_entity_relations` exposes any `not_duplicate`/`duplicate_of`) +and the prompt tells it not to re-propose a cleared pair — self-clearing with zero +custom SQL. The writer agent then, for that job, deliberately and with full context: +uses the **existing `final_rank`/importance signals from that same recall** to decide +which wiki is canonical, rewrites the canonical body to absorb the other's content +and refs, moves/creates `summarises` relations via tools, sets the loser's +`importance` low + `retired_at` + `redirect_to` via `update_entity`, and creates the +`duplicate_of` (or `not_duplicate` if distinct) marker via `create_relation`. Every +step is a logged tool call, reversible, never a hidden bulk SQL mutation. + +--- + +## Reuse map (C3) — existing infra per stage, and what we are NOT building + +| Stage | Needs to… | Uses existing | New code? | +|---|---|---|---| +| Cron | find orphans | one read-only SQL `NOT EXISTS` against `relations` (no scoring involved) | tiny query + enqueue only | +| Maintainer | find candidate wikis for a case; spot duplicates | `recall_memory` / `/memory/context` (text+embedding+graph+decay+`final_rank`), `view_tree`, `search_sql` | **none** — prompt + parse only | +| Writer | pull a wiki's body/members; rank canonical in a merge | `get_entity`, `recall_memory`, `view_entity_relations`; existing `final_rank`/importance from recall | **none** for retrieval/scoring | +| Mutations | create/edit wiki, link, retire, merge | `create_relation`, `delete_relation`, `update_entity` tools | **none** — existing tools | +| Ranking wikis in results | surface wikis well | existing `final_rank` + `importance` + `decay_rate_wiki` config | **none** (config value only) | + +**Explicitly NOT building:** no wiki-vs-wiki cosine query, no `find_similar_keywords` +retarget, no wiki-embedding generator, no winner-selection heuristic in code, no +bespoke dedup pass/filter, no scoring formula change. Detection and ranking are +entirely the existing search infra; the LLM consumes its output. + +--- + +## No-regression guarantees + +- `context.py:~220` keyword filter is `entity_type != "keyword"`; `wiki` passes unchanged — **do not edit that line**. +- Add `"wiki": settings.decay_rate_wiki` (default `0.0`) to `DECAY_RATES`; `decay_rate_wiki` to `config.py`. Config addition, **not** a ranking-formula change. +- Add `"wiki": ("wikis_ext", "...")` to `EXT_QUERIES` (context.py) and a `wiki` branch to `ENTITY_SELECT`/`_flatten()` (entities.py) — same mechanical pattern as the other 5 types. +- `graph.py` already walks all relation types; `summarises` traversed unmodified. No graph/search code change. Existing entity types untouched. +- Migration additive; ingest watcher and `api`/`watcher` compose services untouched. + +--- + +## Files to create / modify + +| File | New/Mod | Purpose | +|---|---|---| +| `alembic/versions/005_wiki_system.py` | new | entity type + `wikis_ext` + `wiki_job` (raw SQL, down_revision "004") | +| `braindb/services/wiki_jobs.py` | new | **non-destructive only**: orphan query, per-orphan triage enqueue, `dedupe_key`, single-job claim (SKIP LOCKED), status transitions, advisory lock, anchor splitter/joiner, **accounted-change gate** (manifest vs parsed-ref diff + section-hash + structural + contradiction cross-check), prior-revision snapshot to activity log, references-ledger regenerator, additive+consistency `reconcile_wiki_refs`. **No search/scoring code** (C3) — detection/ranking/contradiction-context delegated to existing `recall_memory`/`/memory/context`. | +| `braindb/routers/wiki.py` | new | `POST /cron`, `/maintain` (one case/call), `/write` (gate + retry loop), `GET /jobs` under `/api/v1/wiki` | +| `braindb/agent/prompts/wiki_maintainer_prompt.md` | new | maintainer — reason about one case, emit one structured suggestion | +| `braindb/agent/prompts/wiki_writer_prompt.md` | new | writer — mandatory skeleton/anchors/style contract, golden template, few-shot exemplar, edit-mode rules, **change-manifest output**, mandatory contradiction-gathering via existing recall, own relations via tools, consolidate deliberately | +| `braindb/wiki_scheduler.py` | new (Stage 2) | sidecar; clone of `ingest_watcher.py` loop; drains triage one case at a time | +| `braindb/schemas/entities.py` | mod | `WikiCreate`/`WikiRead`/`WikiUpdate`, add to `AnyEntityRead` | +| `braindb/routers/entities.py` | mod | wiki CRUD + extend `ENTITY_SELECT`/`_flatten()`; hook additive `reconcile_wiki_refs` | +| `braindb/schemas/relations.py` | mod | add `summarises`, `not_duplicate`, `duplicate_of`, `consolidated_into` | +| `braindb/services/context.py` | mod | `DECAY_RATES["wiki"]`, `EXT_QUERIES["wiki"]` | +| `braindb/config.py` | mod | `decay_rate_wiki`, `wiki_dedup_similarity_threshold`, interval knobs | +| `braindb/main.py` | mod | `app.include_router(wiki.router)` (1 line) | +| `docker-compose.yml` | mod (Stage 2) | add `wiki_scheduler` service (clone of `watcher`); `api`/`watcher` untouched | +| `docs/maintainer-agent-plan2.md` | new | **frozen** verbatim snapshot of this approved plan (step 0) — historical reference, not edited afterward | +| `docs/maintainer-agent-plan.md` | mod | the *living* design doc — update to the evolved pipeline + C1/C2/C3 constraints + writer accounted-change model; iterated as implementation proceeds | + +No new Python dependencies. + +--- + +## Staged build order + +0. **Freeze a historical snapshot.** Before any code or further plan edits, copy this approved plan verbatim to `c:\Users\dimkn\source\repos\cityfalcon\braindb\docs\maintainer-agent-plan2.md` (sibling to the original `maintainer-agent-plan.md`). This is an immutable reference point: the live plan will keep moving as we implement and test, but `maintainer-agent-plan2.md` preserves the design as approved. (`maintainer-agent-plan.md` is updated separately, per the files table.) +1. **Migration 005** + `schemas`/`entities.py`/`relations.py`/`context.py`/`config.py` wiki CRUD wiring. Verify wiki entities create/read/rank and no retrieval regression. +2. **`services/wiki_jobs.py`** + `routers/wiki.py` `/cron` and `/jobs` (pure SQL, no LLM, non-destructive). Verify per-orphan triage enqueue + idempotency. +3. **`/maintain`** (one case/call) + maintainer prompt. Verify a single triage case → one suggestion job; re-run → no dupes; queue drains case by case. +4. **`/write`** + writer prompt + golden skeleton + **accounted-change gate** + revision snapshot + ledger regen + `reconcile_wiki_refs`. Verify: a *declared* removal (claim demoted with reason, relation re-typed via tools) succeeds and is restorable from the `wiki_revise` log; an *undeclared* drop is rejected+retried; a detected contradiction left unresolved is rejected; untargeted sections on `attach` stay byte-identical; structural validation rejects a bad-style draft. +5. **LLM consolidation** (no new query): duplicate spotted via the maintainer's existing `recall_memory`; writer-driven merge through existing tools + `not_duplicate`/`duplicate_of` self-clearing. Verify every mutation was a logged tool call and is reversible, and that no new search/scoring code was added (C3). +6. **Stage 2**: `wiki_scheduler.py` sidecar + compose service (drains triage one case at a time). +7. **Later track**: Claude-Code `braindb` skill variant driving the same `/api/v1/wiki/*` endpoints without the agent. + +--- + +## Verification (end-to-end) + +Pre-state: README + Karpathy gist already ingested → many keyword/fact entities. + +1. `POST /api/v1/wiki/cron` → one `triage` job per orphan; re-run → no duplicates. +2. `POST /api/v1/wiki/maintain` → consumes exactly **one** triage case, produces one suggestion job; repeat calls drain the queue one at a time. +3. `GET /api/v1/wiki/jobs` → triage + suggestion jobs visible with status. +4. `POST /api/v1/wiki/write {next_pending:true}` → wiki entity with skeleton anchors, `summary` (≤280), `disambiguation`, body `[[ref:UUID]]`, auto-generated references ledger matching the inline tokens and the `summarises` relations exactly. Then exercise surgical editing: an `attach` that **deliberately removes** a now-wrong claim succeeds when the manifest declares it (relation re-typed to `contradicts`/flagged via tools, prior text in the `wiki_revise` log → restorable); the same removal **without** a manifest entry is rolled back and retried (`last_error`/`attempts`); a member that contradicts the consensus forces a populated `section:contradictions` or an explicit demotion; untargeted sections stay byte-identical. +5. `POST /api/v1/memory/context {"queries":["What does the system know about BrainDB?"]}` → BrainDB wiki ranks above individual facts; existing entity types returned exactly as before (baseline unchanged). +6. Seed a near-duplicate wiki → the maintainer's normal `recall_memory` for a related case surfaces it (existing scoring, no new query) → `consolidate` suggestion → writer LLM merges deliberately: activity log shows each `create_relation`/`update_entity` call; loser is soft-retired and still resolves via `GET /entities/{id}`; pair never re-flagged. +7. Delete an entity cited by a wiki → relation removed by FK cascade; dead `[[ref:]]` flagged as a fix-up case (no prose auto-edit). +8. Re-run cron over a fully-wiki'd corpus → 0 new triage jobs (self-clearing verified). diff --git a/docs/wiki-frontend-plan.md b/docs/wiki-frontend-plan.md new file mode 100644 index 0000000..da4c758 --- /dev/null +++ b/docs/wiki-frontend-plan.md @@ -0,0 +1,143 @@ +# Read-only Wiki frontend (Reader + Ops) — zero-backend, Wikipedia-serious + +> Status: FINALISED PLAN — execute in a later session. No worktree/commits +> created until then. (Mirror of the approved plan; kept in-repo so we can +> resume cleanly without re-planning.) + +## Context + +Lever 1 (dedup-first writer priority) + Thread-2 (created_at freshness gate) +are shipped, committed (`a03f077`), and **running** on +`feat/wikis-and-maintainer-agent-with-truncation` for a multi-hour +duplication-self-correction observation. Lever 2/3 stay deferred pending that +outcome. In parallel we want a **read-only wiki frontend (Reader + Ops +dashboard)**. Directives shaping this plan: + +- The frontend **must never touch the DB directly** (so: no client + `/memory/sql`). +- **Minimise backend disruption** — a good assessment must show whether the + backend can be avoided entirely. (It can — see next section.) +- Stack = **simplest**: vanilla static HTML/CSS/JS, no build, no npm, no + framework, no new Python dependency. CORS already open + (`braindb/main.py:31`, `allow_origins=["*"]`). +- Design = **clean like Wikipedia, but built for 2026**: professional, + serious, editorial. Explicitly NOT a colourful/cartoonish/"vibecoded" + mess. + +## Backend assessment — conclusion: ZERO backend changes + +A careful pass over every reader/ops need against existing endpoints: + +| Need | Existing endpoint | Notes | +|---|---|---| +| Wiki index + variant clusters | `GET /api/v1/entities?entity_type=wiki` | Returns `summary`, `importance`, `keywords`, and a ≤1K **content preview** (post-truncation work). The preview's first lines contain `` + `# NAME` + `> **Summary:**` → parse `canonical_name` client-side from the preview. **No N+1 for the index/clusters.** | +| One wiki page | `GET /api/v1/entities/{id}` (+`offset/limit`,`content_meta`) | One call when a wiki is opened (full body + ext: revision, retired_at, redirect_to, member_keyword_ids). Page huge bodies via `content_meta`. | +| Resolve `[[ref:UUID]]` | `GET /api/v1/entities/{UUID}` | Lazy: only when a citation chip is opened (or small batch on page open). | +| Provenance / consistency | `GET /api/v1/entities/{id}/relations` (filter `summarises`) | Consistency (inline refs vs `summarises`) computed **client-side**, same logic as `export_wikis._consistency` (~10 lines JS, regex ported from `REF_RE`, `braindb/services/wiki_jobs.py:32-36`). | +| Related entities | `GET /api/v1/memory/tree/{id}?max_depth=1` | Optional sidebar. | +| Search | `POST /api/v1/memory/search` | Only POST used; not SQL, not a write. | +| Job queue (ops) | `GET /api/v1/wiki/jobs?status=&job_type=&limit=` | Queue mix; pending `consolidate` highlighted (shows Lever 1 draining). | +| Maintainer/writer activity (ops) | `GET /api/v1/memory/log?limit=` | Recent pipeline activity. | +| Consolidation / retire map (ops) | `GET /api/v1/entities/{id}` for the **few** retired wikis only | Retired ⇒ `importance≈0` in the index list (cheap signal); fetch ext (`redirect_to`,`retired_at`) only for those few, not all N. | + +**Result: the entire Reader + Ops dashboard is built from existing GETs +(plus one allowed `/memory/search` POST). No new endpoint, no new service, +no router/`main.py` edit, no new dependency, no DB schema change, and — by +parsing the already-returned content preview — no N+1.** This fully honours +"avoid the backend" and "no DB-direct access". An earlier proposed BFF +layer is **dropped**. + +Out of scope (explicitly NOT in this plan): if the wiki count later grows so +large that even per-open detail calls hurt, a *single* optional read +endpoint could consolidate them — a future decision, not part of this work. + +## Observation safety (only matters if executed while the pipeline still runs) + +The `api` container bind-mounts `.:/app` but **no longer runs uvicorn with +`--reload`** (removed today to avoid mid-pipeline restarts). Code changes +require an explicit `docker compose up -d --no-deps --force-recreate api`, +so `.py` edits don't auto-reload anyway. This frontend adds **no `.py`** and +touches **no existing file** — only new static files. So: + +- If the observation is **still running** when we execute: create the static + app in an **isolated git worktree** (`git worktree add ../braindb-frontend + -b feat/wikis-and-maintainer-agent-frontend`) so branch/commits never + `checkout` the bind-mounted main tree. Serve via stdlib + `python -m http.server` from the worktree; browser → it; JS `fetch`es + `http://localhost:8000`. +- If the observation is **already over**: no worktree needed — just add a + new `frontend/` dir on a dedicated branch (new files don't trigger + reload). + +Either way: zero backend process touched, observation undisturbed. + +## Design language — Wikipedia-grade, 2026-professional + +Reference feel: a serious reference work / editorial knowledge tool, like +Wikipedia's content discipline with a modern 2026 refinement — **not** a SaaS +landing page, **not** colourful, **not** playful. + +DO: content-first single-column reading measure (~68–72ch); restrained +near-monochrome palette (ink `#1b1b1b` on paper `#fff`/`#f8f8f7`, hairline +`#eaeaea` rules, ONE restrained link/citation accent ≈ classic encyclopedic +blue, used sparingly); clear typographic hierarchy (a refined serif for body +e.g. system "Georgia/Charter"-class, clean grotesque for UI/headings/labels); +generous whitespace; quiet left TOC/section nav from the `` +markers; citation chips as small superscript-style references that open a +calm side panel (the entity's content + provenance); a sober Ops view +(plain dense tables, monospace ids, status as quiet text/diamonds — no +traffic-light candy); subtle, near-instant transitions only; light/dark +toggle with the same restraint; fully keyboard navigable; fast, no layout +shift. + +DON'T: bright/multi-colour fills, gradients, glow/neon, big rounded "cards", +emojis as UI, drop shadows everywhere, bouncy animation, dashboard +"widgets", decorative icons. Seriousness over decoration. If in doubt, look +plainer. + +## Files (all NEW, no existing file modified) + +``` +frontend/index.html layout shell (reader + ops tabs), no inline mess +frontend/style.css the design language above; CSS variables; dark mode +frontend/app.js data layer (existing endpoints only) + routing + ops +frontend/wiki-render.js ~150-line purpose-built renderer for the real body + grammar: , # / ##, + > **Summary:/Disambiguation:** callouts, + dividers, GFM tables, lists, + **bold**, `code`, [[ref:UUID|display]] / [[ref:UUID]] + chips (tolerant of grouped [[ref:a], [ref:b]] seen + in real bodies) +frontend/README.md how to run: `python -m http.server` + open URL +``` + +No Python, no dependency, no schema, no write/agent/SQL calls. + +## Verification + +1. **Undisturbed**: `docker logs braindb_wiki_scheduler --tail 3` keeps + advancing across the whole build; main-tree `git status` clean; only new + static files exist. +2. **Pure read**: browser Network tab shows only GETs + the one + `/memory/search` POST — no write/agent/SQL/`/memory/sql`. +3. **Reader**: index lists all wikis (canonical_name parsed from preview), + retired ones flagged; opening `braindb-1785a337` renders + meta/summary/sections/tables faithfully; every `[[ref:UUID]]` chip + resolves to the real entity in the side panel; client consistency badge + equals `export_wikis` (`CONSISTENT ✓`, 3 body / 3 relations). +4. **Ops**: variant panel surfaces the Koutsoumpos / SaaSpocalypse / + BrainDB clusters; queue from `/wiki/jobs` with pending `consolidate` + highlighted and visibly draining first across auto-refreshes (Lever 1); + activity from `/memory/log`; retire/redirect map correct for the few + retired wikis. +5. **Design review**: matches the Wikipedia-serious / 2026 language above — + monochrome+one accent, editorial type, no candy; passes a "does this look + like a serious reference tool, not a vibecoded dashboard" check. + +## Standing constraints + +`.env` never committed/touched. Public repo — no personal names in commit +messages, no Co-Authored-By trailer. Don't push unless asked. No `.py` +edit / `checkout` / restart on the main tree while the observation runs. +Don't touch LLM profiles/.env. Lever 2 / 3 remain deferred pending the +observation outcome. diff --git a/pyproject.toml b/pyproject.toml index 011c379..cb01094 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "braindb" -version = "0.1.0" +version = "0.2.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" @@ -35,17 +35,17 @@ classifiers = [ "Topic :: Software Development :: Libraries :: Python Modules", ] dependencies = [ - "fastapi==0.135.3", - "uvicorn[standard]==0.44.0", - "psycopg2-binary==2.9.11", + "fastapi==0.136.1", + "uvicorn[standard]==0.47.0", + "psycopg2-binary==2.9.12", "alembic==1.18.4", - "pydantic==2.12.5", - "pydantic-settings==2.13.1", + "pydantic==2.13.4", + "pydantic-settings==2.14.1", "python-dotenv==1.2.2", - "sentence-transformers==5.4.0", - "numpy==2.4.4", - "openai-agents[litellm]==0.13.6", - "requests==2.33.1", + "sentence-transformers==5.5.0", + "numpy==2.4.5", + "openai-agents[litellm]==0.17.2", + "requests==2.34.2", ] [project.optional-dependencies] diff --git a/skills/braindb-agent/SKILL.md b/skills/braindb-agent/SKILL.md index e9a972e..81658fb 100644 --- a/skills/braindb-agent/SKILL.md +++ b/skills/braindb-agent/SKILL.md @@ -6,7 +6,7 @@ allowed-tools: Bash Read ## BrainDB Memory Agent -BrainDB has its own internal agent (LiteLLM + NVIDIA NIM) that handles all memory operations. You don't call individual endpoints — you ask the agent in plain English via one endpoint: `POST http://localhost:8000/api/v1/agent/query`. +BrainDB has its own internal agent (LiteLLM with pluggable provider via `LLM_PROFILE`; defaults to `deepinfra/google/gemma-4-31B-it`) that handles all memory operations. You don't call individual endpoints — you ask the agent in plain English via one endpoint: `POST http://localhost:8000/api/v1/agent/query`. ### Health check: !`curl -sf http://localhost:8000/health > /dev/null 2>&1 && echo "OK" || echo "BRAINDB_DOWN"` @@ -26,6 +26,31 @@ BrainDB has its own internal agent (LiteLLM + NVIDIA NIM) that handles all memor --- +## TOOL PRIORITY + +The agent already uses the sophisticated retrieval (keyword-mediated fuzzy + +embedding + graph + ranking, with a two-level diversity quota) and can +delegate to subagents. Phrase requests as goals ("find / recall / understand +…", "delegate a deep investigation of …"). **Do not tell it to "run SQL"** +for recall or understanding — raw SQL discards the graph and embeddings. If +you're tempted to phrase a request as *"run a SQL query that finds…"* for +*finding* or *understanding* something, stop — that's the sophisticated +recall path's job. Ask in plain English. SQL is only ever for an explicit +aggregate ("how many facts per source?"), which you can simply ask for in +plain English anyway. + +**Wikis** are first-class memory entities curated by an internal maintainer + +writer pipeline. The agent surfaces them through recall automatically when +relevant — you don't have to ask for them explicitly, and you don't have to +trigger anything to make new ones. Saving facts with the right keywords is +enough; the scheduler runs maintain → write on its 60s tick and the wikis +materialise on their own. + +Internally the agent now researches from **short previews** and reads a full +body only by id (paging large ones, or delegating big documents to a +subagent), so its context stays clean — just ask in natural language ("read +and summarise datasource X"); it handles the chunking itself. + ## RECALL — at conversation start, and whenever you need context Ask the agent in natural language. It handles keyword formulation, multi-query search, graph traversal, and summarization. @@ -59,22 +84,67 @@ curl -s -X POST http://localhost:8000/api/v1/agent/query \ -d '{"query":"Save: the user just told me they prefer simple code over abstractions. Source: user-stated. Connect to existing preference entities."}' ``` -**Be proactive**: save user profile info, expertise, preferences, decisions, inferences you make about their working style. When in doubt, save it. +### Proactive save — but ASK the user first + +The pattern is **RECALL → ASK → SAVE**: + +1. When the user shares something that *might* be worth remembering (a name, + role, project, preference, decision, your own inference about them), RECALL + first via the agent to check if it's already known. +2. If it's **net-new**, **ASK the user**: + + > "I haven't seen this before — should I save it to BrainDB? I'd file it + > as a [fact / thought / rule] tagged with [keywords]." + +3. Only on a 'yes', issue the save request to the agent. + +Don't pre-save without confirmation. The user has the final say on what +becomes long-term memory. User-confirmed memory is higher-signal and lets +the user catch judgement-call mistakes early. + +**Exception**: when the user explicitly framed it as a rule ("from now on, +always X"; "never do Y"), save it without an extra confirmation — they +already said it — but surface the action: "Saving that as a rule." + +#### What's worth flagging to the user + +- Identity / role / company (one-time setup info) +- Strong preferences or working-style rules +- Project / topic context the user just disclosed +- Decisions the user explicitly made +- Useful URLs or references the user shared +- Your own inferences about the user (tag as `thought`, + `source=agent-inference`) — ASK before persisting these too; an inference + is still memory. + +The goal is to capture **what the user gives you in conversation that isn't +already in BrainDB** — not to scrape every utterance. Information already in +recall doesn't need saving again; ephemeral task details +("currently debugging X") don't need saving at all. --- ## Example queries +### Recall (no confirmation needed — these are reads) + | Situation | Query to send to the agent | |-----------|---------------------------| | Start of conversation | `"Tell me who the user is - role, expertise, preferences, recent projects."` | | User mentions a topic | `"What do you know about the user ML experience and AI projects?"` | -| User shares a fact | `"Save: user is working on the IR pipeline multilingual extraction. Connect to existing IR entities."` | -| User gives a preference | `"Save as rule: always prefer simple code over abstractions. Source: user-stated. Category: behavior."` | | User asks about past work | `"What has the user shipped recently? Check facts with source=user-stated from the last month."` | | Need to find duplicates | `"Find near-duplicate entities in memory."` | | Explore the graph | `"What are the densest topics in memory? Which entities have the most connections?"` | +### Save (RECALL → ASK → SAVE — only send the agent query after the user confirms) + +| Situation | What Claude says to the user first | What Claude sends to the agent (on a 'yes') | +|---|---|---| +| User mentions something net-new | "I noticed you just said you're working on the IR pipeline multilingual extraction — that looks worth saving. Should I?" | `"Save: user is working on the IR pipeline multilingual extraction. Connect to existing IR entities."` | +| User shares a preference | "Should I save that as a long-term preference?" | `"Save as fact: user prefers simple code over abstractions. Source: user-stated. Keywords: user-preference, code-style."` | +| User explicitly states a rule | (no confirmation — they framed it as a rule) "Saving that as a rule." | `"Save as rule: always prefer simple code over abstractions. Source: user-stated. Category: behavior."` | +| You drew an inference about the user | "I'm getting the sense you're senior in ML — should I save that as a thought?" | `"Save as thought: user appears senior in ML based on the depth of their question. Source: agent-inference. Certainty: 0.6."` | + --- ## Delegation — ask the agent to spawn a subagent for focused work @@ -110,6 +180,34 @@ Delegation is 1 level deep — subagents cannot spawn more subagents. --- +## File ingestion — automatic, no agent call needed + +If the user wants a local file (article, transcript, note, document) ingested into BrainDB, **don't ask the agent to do it**. Instead, copy the file into the repo's `data/sources/` directory and the system handles the rest: + +1. The `braindb_watcher` sidecar polls `data/sources/` every ~7 seconds. +2. New files are auto-ingested as `datasource` entities (content + hash + word count). +3. The watcher then runs an agent-driven extraction pass that creates one or more `fact` entities derived from the document and links them back via `derived_from` relations. +4. On success the file is moved to `data/sources/ingested/`; on failure to `data/sources/failed/` with a sidecar `.error.txt`. + +What this means for you (Claude) and the user: + +- **Tell the user**: "Just drop the file into `data/sources/` on the BrainDB repo. The watcher will pick it up within a few seconds and you'll see the facts appear in recall a minute or two later." +- **Do not** issue an `/agent/query` like `"Save this file..."` with the file contents pasted into the prompt — that bloats the LLM context and bypasses the proper extraction pipeline. The watcher path produces structured facts + `derived_from` relations + keyword auto-tagging; pasting bypasses all of it. +- **Watch progress** if you want to confirm completion: + +```bash +docker logs braindb_watcher -f +``` + +You'll see `ingested NEW: -> words=N` then later `extraction complete for : N facts total`. After that the new facts surface naturally in `/agent/query` recall — no extra steps. + +Edge cases: +- Very large files are chunked automatically; extraction takes proportionally longer (typically 60-180 seconds per chunk on local Qwen, faster on deepinfra). +- If a file ends up in `data/sources/failed/`, read the sidecar `.error.txt` next to it to see what went wrong. +- The watcher dedupes by file content hash, so re-dropping the same file won't re-extract. + +--- + ## Verbose mode — watch the agent work in real time Set `AGENT_VERBOSE=true` in the server's `.env` (default is `false`). When enabled, every tool call the agent makes is logged to stdout with args and result preview. Watch it live: @@ -126,4 +224,4 @@ The HTTP response itself is unchanged (just `{"answer": "..."}`). Logs go to the - If the agent call fails (connection refused, 500, timeout): proceed WITHOUT memory. Don't retry, don't block the conversation. - If the answer mentions an ERROR: the agent tried but some tool failed. Carry on — use whatever partial information came back. -- Agent calls can take 5-30 seconds (LLM + multi-turn loop). Subagent calls can take 30-90 seconds. That's normal. +- Agent calls can take up to 10 minutes if the LLM provider is slow. Add `--max-time 600` to long curl calls. diff --git a/skills/braindb/SKILL.md b/skills/braindb/SKILL.md index 6cc9b6e..0f482e8 100644 --- a/skills/braindb/SKILL.md +++ b/skills/braindb/SKILL.md @@ -85,24 +85,68 @@ If the final curl returns `{"status":"ok"}`, you're live. --- +## TOOL PRIORITY (read this first) + +BrainDB's power is the graph + embeddings + ranking. Use it; do not fall back +to flat SQL. + +1. **`POST /api/v1/memory/context`** (multi-query) — the default for ALL + recall, discovery, and understanding. BOTH the fuzzy and embedding + pathways are **keyword-mediated** (the query matches against keyword + entities, entities surface via `tagged_with`). A two-level diversity + quota (per-search-term + per-keyword halving) keeps results + balanced. Then graph traversal + decay + ranking. +2. **`POST /api/v1/agent/query` with "delegate to a subagent…"** — for + multi-step investigation/disambiguation; the agent researches and returns a + summary. +3. `GET /api/v1/entities…`, `GET /api/v1/memory/tree/`, + `GET /api/v1/entities//relations` — targeted structure lookups. +4. **Wikis** — first-class entity type, curated topic pages assembled by an + internal maintainer + writer pipeline from facts/thoughts tagged with the + same keyword. To browse: `GET /api/v1/entities?entity_type=wiki`. Full body: + `GET /api/v1/entities/`. Wikis also surface naturally in `/memory/context`. + Write paths are documented in the WIKIS section below. +5. **`POST /api/v1/memory/sql` — exception only.** A flat SELECT has no + embeddings/graph/ranking. Use it solely for a specific structured/aggregate + question (counts, GROUP BY, activity-log joins) the above cannot express. + **Never** for recall, discovery, similarity, or understanding. + +If you're about to use `/memory/sql` to *find* or *understand* something, +stop — that's a `/memory/context` (or delegated `/agent/query`) job. + +### Previews vs full body + +`/memory/context` (and `/memory/search`, `GET /entities`) return **short +previews** per item (~1K); a clipped item ends with +`--truncated (N more) -- full body: get_entity("")`. That's intended — +decide from previews, then read only what you need: + +- Full single entity: `GET /api/v1/entities/{id}`. +- Large body: page it — `GET /api/v1/entities/{id}?offset=0&limit=8000`, then + 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. + ## RECALL — Before Responding ### Step 1: Formulate targeted queries Analyze the user's message. Extract the **core topics** that need memory context. Create **multiple targeted queries** — do NOT paste the raw user message. -**Important**: Use terms that match how entities are STORED, not natural language questions. The search uses trigram similarity + full-text matching. Specific terms that would appear in stored content work best. Vague queries with stop words ("everything about X") will return nothing. +**Query strategy** — BrainDB's retrieval is keyword-mediated, so: -Include likely keywords in your queries: `user-profile`, `expertise`, `project-decision`, `user-preference`. +- Prefer **multiple narrow queries** (single keywords / bare names) over one long sentence. Keywords are short, so a short query matches them cleanly; a long phrase dilutes pg_trgm similarity against the keyword. +- The per-search-term quota reserves slots for EACH query you pass, so adding a bare keyword as one of your queries guarantees it surfaces (it doesn't compete with the broader phrases). +- Use terms that match how entities are STORED. Common keyword conventions: `user-profile`, `expertise`, `project-decision`, `user-preference`. -Examples: +Examples (narrow + one broader angle, mixed): | User says | Queries | |-----------|---------| -| "help me refactor this React component" | `["user-profile React frontend expertise", "user-preference code style refactoring"]` | -| "let's work on the IR pipeline" | `["investor-relations IR scraping architecture", "user-preference deployment workflow"]` | -| (new conversation, no specific topic) | `["user-profile expertise role background", "user-preference working style"]` | -| "what's the best way to deploy this?" | `["deployment infrastructure project-decision", "user-preference production services"]` | +| "help me refactor this React component" | `["user-profile", "React", "user-preference code style refactoring"]` | +| "let's work on the IR pipeline" | `["investor-relations", "IR", "deployment workflow"]` | +| (new conversation, no specific topic) | `["user-profile", "expertise", "working style"]` | +| "what's the best way to deploy this?" | `["deployment", "infrastructure", "production services"]` | Always include a `"user-profile"` query on the first message of a conversation — you need to know who you're talking to. @@ -111,9 +155,11 @@ Always include a `"user-profile"` query on the first message of a conversation ```bash curl -s -X POST http://localhost:8000/api/v1/memory/context \ -H "Content-Type: application/json" \ - -d '{"queries": ["query1", "query2"], "max_depth": 3, "max_results": 15}' + -d '{"queries": ["narrow1", "narrow2", "one broader phrase"], "max_depth": 3}' ``` +`max_results` defaults to 30 — leave it unless you specifically want fewer. + ### Step 3: Evaluate results and retry if weak If you got **0 results**, your query terms didn't match stored content. Reformulate with more specific terms that would actually appear in entity content or keywords. @@ -140,13 +186,32 @@ Let recalled facts inform your response. **Do NOT announce** "I found in memory ## SAVE — After Responding -After each interaction, evaluate what you learned. **Be proactive and thorough about saving.** +After each interaction, evaluate what you learned. The policy is **RECALL → ASK → SAVE.** + +### Saving philosophy — always ASK the user first + +Always recall first. If what the user shared is **net-new** (not already in +`/memory/context`), **ASK the user** before saving: + +> "I haven't seen this before — should I save it as a fact / thought / rule? +> (I'd tag it with keywords X, Y; importance Z.)" -### Saving philosophy +Only persist after the user confirms. The user has the final say on what +becomes long-term memory. Auto-saves without confirmation dilute signal and +accumulate junk; user-confirmed memory is higher-signal and traceable. -- **Save everything worth remembering.** Don't skip something because it seems minor — save it with lower importance. A fact you didn't need is harmless. A fact you forgot is a missed opportunity. -- **Create THOUGHTS proactively.** After each interaction, form inferences: what does this tell you about the user's expertise? Their working style? Their priorities? Thoughts are cheap and enrich the graph. -- **Create RELATIONS for every new entity.** Connect it to existing entities found during recall. Multiple relations per entity is ideal — the graph's value comes from density. +**Exception** — behavioural rules the user explicitly stated as rules ("from +now on, always X"; "never do Y") can be saved without an extra confirmation — +they already said it. Just surface the action: "Saving that as a rule." + +Once the user agrees: + +- **Create RELATIONS for every new entity.** Connect it to existing entities + found during recall. Multiple relations per entity is ideal — the graph's + value comes from density. +- **Thoughts (your own inferences about the user) — ASK before persisting, + same as facts.** A thought is still memory; the user should agree it + belongs there. ### What to save as @@ -295,9 +360,14 @@ curl -s "http://localhost:8000/api/v1/memory/log?since=2026-04-08T00:00:00Z" Use this to answer "when did I learn this?" or "what was I working on yesterday?" -### Read-only SQL — power queries +### Read-only SQL — EXCEPTION tool, aggregations only -For ad-hoc exploration and aggregations the standard endpoints don't cover. Only `SELECT` and `WITH` queries are allowed; 5s timeout; 1000 row limit. +⚠ Not a recall/discovery tool (see TOOL PRIORITY at the top). A flat SELECT +throws away embeddings, graph and ranking — everything BrainDB is good at. +Use it **only** for a specific structured/aggregate question the dedicated +endpoints cannot express (counts, GROUP BY, activity-log joins). For finding +or understanding anything, use `/memory/context` or a delegated `/agent/query`. +Only `SELECT` and `WITH` queries are allowed; 5s timeout; 1000 row limit. ```bash # Count entities by source @@ -316,7 +386,94 @@ curl -s -X POST http://localhost:8000/api/v1/memory/sql \ -d '{"query": "SELECT l.timestamp, l.operation, e.content FROM activity_log l JOIN entities e ON e.id = l.entity_id ORDER BY l.timestamp DESC LIMIT 20"}' ``` -Prefer the dedicated endpoints for normal operations. Use SQL when you need something unusual. +Reiterate: `/memory/context` (+ delegated `/agent/query`) is the default for +everything. `/memory/sql` is the rare exception for true aggregations only. + +--- + +## WIKIS — Auto-Curated Topic Pages + +Wikis are canonical topic pages BrainDB assembles automatically from +facts/thoughts tagged with the same keyword. An internal maintainer runs +every 60s, scans for orphan keywords (a keyword with members but no wiki +yet), and decides per-orphan: **attach** (the topic already has a wiki), +**create** (mint a new one), **consolidate** (merge duplicates), or +**skip** (not a wiki-worthy subject). Approved suggestions then become wiki +bodies via the wiki writer. You usually don't need to do anything — saving +facts with consistent keywords is enough; the pipeline materialises the +wikis on its own. + +### Recall — browse and read wikis + +```bash +# List all wikis (most recent first), previews only +curl -s "http://localhost:8000/api/v1/entities?entity_type=wiki&limit=50" + +# Read a wiki body in full +curl -s http://localhost:8000/api/v1/entities/ +``` + +Wikis surface in `/memory/context` automatically — you don't have to ask +for them separately when doing topic recall. + +### Write — indirect (default): let the pipeline decide + +1. Save your facts with the right keyword (the subject's bare name — + `keywords=["Sawki"]`, not `["Sawki the employee"]`). +2. (Optional) Nudge the pipeline so the maintainer evaluates the new + keyword *now* rather than on the next scheduler tick: + +```bash +curl -s -X POST http://localhost:8000/api/v1/wiki/cron +``` + +The cron is **idempotent** (safe to call any time). It enqueues triage +jobs for orphan keywords; the scheduler then runs maintain → write on +its next 60s tick. The maintainer can still decide to **skip** the +orphan if the subject isn't worth a wiki (e.g. an infrastructural +keyword) — that's expected and not an error. + +Inspect what's pending: + +```bash +curl -s "http://localhost:8000/api/v1/wiki/jobs?status=pending&limit=20" +``` + +### Write — direct (power user, rare): bypass the pipeline + +When you need full control over the body and you know exactly what the +wiki should say, you can create one directly: + +```bash +curl -s -X POST http://localhost:8000/api/v1/wikis \ + -H "Content-Type: application/json" \ + -d '{ + "content": "# Sawki\n\nFull markdown body here...", + "canonical_name": "Sawki", + "disambiguation": "Team member, distinct from other people with similar names", + "language": "en", + "member_keyword_ids": [""], + "keywords": ["Sawki", "Egypt", "Petros"], + "importance": 0.7, + "source": "user-stated" + }' +``` + +⚠ This **bypasses the maintainer's dedup logic.** If a wiki for that +subject already exists, you'll create a duplicate that someone (or the +next `consolidate` maintainer decision) has to clean up. Prefer the +indirect path unless you specifically know why the pipeline can't do +what you need. + +`member_keyword_ids` requires existing keyword UUIDs. Find them via: + +```bash +curl -s "http://localhost:8000/api/v1/entities?entity_type=keyword&content=" +``` + +We intentionally do NOT document `POST /wiki/maintain` or `POST +/wiki/write` here — they're claim-based (take no target) and only make +sense as scheduler-internal steps. --- diff --git a/tests/conftest.py b/tests/conftest.py index 3f74f3e..8f573d9 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -51,6 +51,59 @@ def _require_live_api() -> None: ) +@pytest.fixture(scope="session", autouse=True) +def _purge_pytest_artefacts_at_session_end() -> Iterator[None]: + """Session teardown safety net for the per-test `created_entities` + fixture: any test that errors before registering its IDs (or that + bypasses the factories entirely) still leaks `_pytest_` rows + into the live DB. After all tests finish, sweep those out. + + Pattern uniqueness: `_pytest_<8-hex>` is generated only by the + `test_tag` fixture above and never by production code — so a + `content LIKE '_pytest_%'` filter on keyword entities is provably + scoped to test artefacts. + + Order matters: delete tagged entities (facts/thoughts/...) FIRST so + their `tagged_with` edges drop via FK cascade, then the keyword + entities themselves. + """ + yield + try: + from braindb.db import get_conn # only imported at teardown + except Exception as exc: # noqa: BLE001 — defensive, never block the session + print(f"\n[conftest] session cleanup skipped (db import failed): {exc}") + return + try: + with get_conn() as conn: + with conn.cursor() as cur: + cur.execute( + """ + DELETE FROM entities WHERE id IN ( + SELECT r.from_entity_id FROM relations r + JOIN entities kw ON kw.id = r.to_entity_id + WHERE r.relation_type = 'tagged_with' + AND kw.entity_type = 'keyword' + AND kw.content LIKE E'\\_pytest\\_%' ESCAPE '\\' + ) + """ + ) + tagged_deleted = cur.rowcount + cur.execute( + """ + DELETE FROM entities + WHERE entity_type = 'keyword' + AND content LIKE E'\\_pytest\\_%' ESCAPE '\\' + """ + ) + kw_deleted = cur.rowcount + print( + f"\n[conftest] session cleanup: removed {tagged_deleted} " + f"tagged entities + {kw_deleted} _pytest_* keywords" + ) + except Exception as exc: # noqa: BLE001 — never break the session on cleanup + print(f"\n[conftest] session cleanup error (ignored): {exc}") + + @pytest.fixture def api() -> str: """Base URL for the API — tests append paths like f'{api}/api/v1/...'.""" diff --git a/tests/test_final_answer_rename.py b/tests/test_final_answer_rename.py new file mode 100644 index 0000000..dc4c9ca --- /dev/null +++ b/tests/test_final_answer_rename.py @@ -0,0 +1,728 @@ +"""Edge-case tests for Stage C — `submit_result` → `final_answer` rename + slot pattern. + +These are UNIT tests: they import `braindb.agent.*` directly and exercise the +internal contract surface (`FunctionTool.name`, the `_build()` factory's +`StopAtTools` config, the run_state slot lifecycle, run_typed's strict +behaviour). No live LLM, no HTTP — fast and deterministic. + +They run alongside the existing integration tests; pytest's session-scoped +`_require_live_api` fixture from `conftest.py` still applies (the suite as a +whole expects a healthy stack), but THESE tests don't actually call the API. + +Until Stage C / Layer 1 lands, most assertions here are RED on the +`experimental/structured-output-proper` branch (the rename hasn't happened +yet). After the rename they go green and serve as regression coverage. +""" +from __future__ import annotations + +from pathlib import Path +from unittest import mock + +import pytest + +from braindb.agent import agent as agent_module +from braindb.agent import run_state +from braindb.agent.schemas import ( + AgentAnswer, + MaintainerDecision, + SubagentResult, + WikiWriteResult, +) +from braindb.agent.tools import ( + submit_answer, + submit_maintainer, + submit_subagent, + submit_wiki, +) + + +# ------------------------------------------------------------------ # +# Layer 1 — rename surface (FAILS until Stage C / Layer 1 ships) # +# ------------------------------------------------------------------ # + +EXPECTED_FINAL_TOOL_NAME = "final_answer" + + +@pytest.mark.parametrize( + "tool", + [submit_answer, submit_maintainer, submit_wiki, submit_subagent], + ids=["answer", "maintainer", "wiki", "subagent"], +) +def test_submit_tools_renamed_to_final_answer(tool) -> None: + """Every typed `submit_*` @function_tool must expose name 'final_answer' + to the SDK after the rename. The LLM sees this name in the tool catalog; + a mismatch with the prompt or `StopAtTools` config breaks termination.""" + assert hasattr(tool, "name"), ( + f"{tool!r} is not a FunctionTool — did @function_tool decoration get dropped?" + ) + assert tool.name == EXPECTED_FINAL_TOOL_NAME, ( + f"{tool!r}.name={tool.name!r}; expected {EXPECTED_FINAL_TOOL_NAME!r} after rename" + ) + + +def test_stop_at_tools_uses_final_answer() -> None: + """The `_build()` factory must configure `StopAtTools` with the new name. + Build all four agents and inspect their tool_use_behavior.""" + agents_to_check = [ + agent_module.get_agent(), + agent_module.get_maintainer_agent(), + agent_module.get_writer_agent(), + agent_module.get_subagent(), + ] + for a in agents_to_check: + beh = a.tool_use_behavior + # SDK stores it as a dict {"stop_at_tool_names": [...]} OR as a + # StopAtTools dataclass with the same attribute. Accept both shapes. + names = ( + beh.get("stop_at_tool_names") if isinstance(beh, dict) + else getattr(beh, "stop_at_tool_names", None) or getattr(beh, "tool_names", None) + ) + assert names is not None, f"{a.name}: tool_use_behavior {beh!r} has no recognisable stop-names" + assert EXPECTED_FINAL_TOOL_NAME in names, ( + f"{a.name}: StopAtTools={names!r}; expected to include {EXPECTED_FINAL_TOOL_NAME!r}" + ) + + +@pytest.mark.parametrize( + "prompt_path", + [ + Path("braindb/agent/prompts/system_prompt.md"), + Path("braindb/agent/prompts/wiki_maintainer_prompt.md"), + Path("braindb/agent/prompts/wiki_writer_prompt.md"), + ], + ids=["system", "wiki_maintainer", "wiki_writer"], +) +def test_prompts_no_stale_submit_result(prompt_path: Path) -> None: + """Prompt files must NOT contain the literal `submit_result` after the + rename — otherwise the LLM gets a confused contract (catalog says + `final_answer`, prompt says `submit_result`).""" + repo_root = Path(__file__).parent.parent # tests/ → repo root + full = repo_root / prompt_path + assert full.exists(), f"prompt missing: {full}" + body = full.read_text(encoding="utf-8") + assert "submit_result" not in body, ( + f"{prompt_path} still references 'submit_result' — should be 'final_answer'" + ) + + +def test_writer_prompt_has_attach_mode_efficiency_hint() -> None: + """Regression cover for the attach-mode recall-budget guidance added + to wiki_writer_prompt.md (a future accidental delete should trip red + immediately, before live behaviour regresses). Asserts: + - the new 'Attach mode — read the existing body before recalling' block + is present, + - the three recall-budget bullets are present, + - the existing 'conservatively' caution rephrasing of the prior + 'NOT evidence' rule is present (the prior strict rule has been + replaced; this protects the new wording from a silent revert).""" + repo_root = Path(__file__).parent.parent + body = (repo_root / "braindb/agent/prompts/wiki_writer_prompt.md").read_text(encoding="utf-8") + + # The Attach-mode header + assert "Attach mode — read the existing body before recalling" in body, ( + "Draft B header missing from writer prompt — recall-budget guidance was lost" + ) + + # Each of the three recall-budget bullet keys + for bullet in [ + "new members (the `MEMBERS` block)", + "claims that look inconsistent between the body and a new member", + "gaps the new members open up", + ]: + assert bullet in body, f"Draft B bullet missing from writer prompt: {bullet!r}" + + # The softened "conservatively" rephrasing of the prior "NOT evidence" rule + assert "conservatively" in body, ( + "Softened 'conservatively' caution missing — the prior 'NOT evidence' rule " + "may have been re-introduced verbatim or the new wording dropped" + ) + + # Closing balance phrase + assert "Be thorough where evidence is fresh or conflicting" in body, ( + "Draft B closing balance phrase missing" + ) + + +# ------------------------------------------------------------------ # +# Slot pattern (already shipped in 8560cfa; regression coverage) # +# ------------------------------------------------------------------ # + + +def test_slot_install_and_release_isolation() -> None: + """Two sequential install/release cycles produce distinct slot objects. + Within a cycle, `record_submit` mutates the active slot; after release, + the outer slot's value is unchanged.""" + slot1, token1 = run_state.install_slot() + assert slot1.value is None + run_state.record_submit("payload-1") + assert slot1.value == "payload-1" + run_state.release_slot(token1) + + slot2, token2 = run_state.install_slot() + assert slot2 is not slot1 + assert slot2.value is None # fresh slot, not stale data from slot1 + run_state.record_submit("payload-2") + assert slot2.value == "payload-2" + assert slot1.value == "payload-1" # the released slot still holds its old data, but is no longer the ContextVar's value + run_state.release_slot(token2) + + +def test_slot_nested_install_release() -> None: + """The wiki maintainer/writer pattern: parent run_typed installs a slot, + a delegated subagent installs its own, releases, then parent finalises. + The child's record_submit must NOT contaminate the parent's slot.""" + parent_slot, parent_token = run_state.install_slot() + run_state.record_submit("parent-data") + assert parent_slot.value == "parent-data" + + # Child run_typed enters + child_slot, child_token = run_state.install_slot() + assert child_slot is not parent_slot + assert child_slot.value is None + run_state.record_submit("child-data") + assert child_slot.value == "child-data" + assert parent_slot.value == "parent-data" # unaffected + run_state.release_slot(child_token) + + # Back in parent context; record_submit should target parent again + run_state.record_submit("parent-data-after-child") + assert parent_slot.value == "parent-data-after-child" + run_state.release_slot(parent_token) + + +def test_record_submit_outside_run_is_silent_noop() -> None: + """If `record_submit` is called outside any `install_slot()` scope (e.g. + a bug in a tool, or stale state), it must NOT raise. The current + implementation silently drops the payload because the ContextVar + defaults to None.""" + # This must not raise even with no active slot. + run_state.record_submit("orphan-payload") + # The slot var should still be None + assert run_state._slot_var.get() is None + + +# ------------------------------------------------------------------ # +# run_typed strict-mode behaviour # +# ------------------------------------------------------------------ # + + +@pytest.mark.asyncio +async def test_run_typed_raises_when_submit_never_fires() -> None: + """If Runner.run completes without any `submit_*` having called + record_submit, run_typed must raise RuntimeError — the strict-mode + invariant. Surfaces 'model emitted prose' / 'max_turns exhausted' + as a real failure rather than silently returning bad data.""" + fake_agent = mock.MagicMock(name="fake_agent") + fake_agent.name = "FakeAgent" + + async def fake_runner_run(starting_agent, input, max_turns, **kwargs): + # Pretend the LLM ran but never called any submit_*. + return mock.MagicMock(final_output="some-prose-text") + + with mock.patch.object(agent_module.Runner, "run", new=fake_runner_run): + with pytest.raises(RuntimeError, match="did not call final_answer|did not submit"): + await agent_module.run_typed("query", fake_agent, AgentAnswer, max_turns=5) + + +@pytest.mark.asyncio +async def test_run_typed_returns_typed_payload_when_submitted() -> None: + """If record_submit IS called during Runner.run with the expected typed + payload, run_typed returns that exact instance — the typed-final + contract.""" + fake_agent = mock.MagicMock(name="fake_agent") + fake_agent.name = "FakeAgent" + expected = AgentAnswer(answer="hello world") + + async def fake_runner_run(starting_agent, input, max_turns, **kwargs): + # Simulate a submit_* tool body firing during the run + run_state.record_submit(expected) + return mock.MagicMock(final_output="ok") + + with mock.patch.object(agent_module.Runner, "run", new=fake_runner_run): + got = await agent_module.run_typed("query", fake_agent, AgentAnswer, max_turns=5) + assert got is expected + assert got.answer == "hello world" + + +# ------------------------------------------------------------------ # +# Pydantic typed-arg validation (regression cover) # +# ------------------------------------------------------------------ # + + +# ------------------------------------------------------------------ # +# Stage C / Layer 4 — retry-with-correction on prose-terminal # +# ------------------------------------------------------------------ # + + +@pytest.mark.asyncio +async def test_run_typed_retries_when_first_attempt_missing_final() -> None: + """When the first `Runner.run` ends without `final_answer` firing, + `run_typed` must inject a correction message and re-invoke + `Runner.run` ONCE. On the retry, if the model calls `final_answer` + via `record_submit`, the typed payload is returned and the caller + gets a success — no 500.""" + fake_agent = mock.MagicMock(name="fake_agent") + fake_agent.name = "FakeAgent" + expected = AgentAnswer(answer="recovered after correction") + call_count = {"n": 0} + + async def fake_runner_run(starting_agent, input, max_turns, **kwargs): + call_count["n"] += 1 + result_mock = mock.MagicMock() + result_mock.to_input_list.return_value = [{"role": "user", "content": "prior context"}] + result_mock.final_output = "prose without final_answer call" + if call_count["n"] == 2: + # The retry: simulate the model now calling final_answer + run_state.record_submit(expected) + return result_mock + + with mock.patch.object(agent_module.Runner, "run", new=fake_runner_run): + # Make sure retry is enabled + with mock.patch.object(agent_module.settings, "agent_retry_on_missing_final", True): + got = await agent_module.run_typed("query", fake_agent, AgentAnswer, max_turns=10) + assert got is expected + assert call_count["n"] == 2, "expected exactly one retry" + + +@pytest.mark.asyncio +async def test_run_typed_raises_when_retry_also_fails() -> None: + """If BOTH the first attempt AND the retry end without `final_answer`, + `run_typed` must still raise `RuntimeError`. No silent success on a + genuinely-broken model that refuses the contract even after + correction.""" + fake_agent = mock.MagicMock(name="fake_agent") + fake_agent.name = "FakeAgent" + call_count = {"n": 0} + + async def fake_runner_run(starting_agent, input, max_turns, **kwargs): + call_count["n"] += 1 + result_mock = mock.MagicMock() + result_mock.to_input_list.return_value = [] + result_mock.final_output = "still prose" + # Neither attempt calls record_submit — slot stays None. + return result_mock + + with mock.patch.object(agent_module.Runner, "run", new=fake_runner_run): + with mock.patch.object(agent_module.settings, "agent_retry_on_missing_final", True): + with pytest.raises(RuntimeError, match="did not call final_answer|even after"): + await agent_module.run_typed("query", fake_agent, AgentAnswer, max_turns=10) + assert call_count["n"] == 2, "expected exactly one retry before giving up" + + +@pytest.mark.asyncio +async def test_run_typed_retry_disabled_via_setting() -> None: + """`agent_retry_on_missing_final=False` is the opt-out: when the first + attempt ends without `final_answer`, raise immediately — no retry.""" + fake_agent = mock.MagicMock(name="fake_agent") + fake_agent.name = "FakeAgent" + call_count = {"n": 0} + + async def fake_runner_run(starting_agent, input, max_turns, **kwargs): + call_count["n"] += 1 + result_mock = mock.MagicMock() + result_mock.to_input_list.return_value = [] + result_mock.final_output = "prose" + return result_mock + + with mock.patch.object(agent_module.Runner, "run", new=fake_runner_run): + with mock.patch.object(agent_module.settings, "agent_retry_on_missing_final", False): + with pytest.raises(RuntimeError, match="did not call final_answer"): + await agent_module.run_typed("query", fake_agent, AgentAnswer, max_turns=10) + assert call_count["n"] == 1, "retry should NOT happen when setting is False" + + +@pytest.mark.asyncio +async def test_run_typed_correction_message_appended_on_retry() -> None: + """The retry call must pass `result.to_input_list() + [correction]` as + `input` to `Runner.run`, where `correction` is a user-role message + that explicitly references `final_answer` so the LLM gets an + unambiguous instruction (not a parse-the-prose hack).""" + fake_agent = mock.MagicMock(name="fake_agent") + fake_agent.name = "FakeAgent" + prior_items = [ + {"role": "user", "content": "save this fact"}, + {"role": "assistant", "content": "okay, doing the work..."}, + ] + captured_inputs: list = [] + + async def fake_runner_run(starting_agent, input, max_turns, **kwargs): + captured_inputs.append(input) + result_mock = mock.MagicMock() + result_mock.to_input_list.return_value = prior_items + result_mock.final_output = "prose" + # No record_submit anywhere — to force the retry path AND fail again. + return result_mock + + with mock.patch.object(agent_module.Runner, "run", new=fake_runner_run): + with mock.patch.object(agent_module.settings, "agent_retry_on_missing_final", True): + with pytest.raises(RuntimeError): + await agent_module.run_typed("save this fact", fake_agent, AgentAnswer, max_turns=10) + + # First call gets the raw query string; second gets the prior history + a correction. + assert len(captured_inputs) == 2 + assert captured_inputs[0] == "save this fact" + retry_input = captured_inputs[1] + assert isinstance(retry_input, list), f"retry input must be a message list, got {type(retry_input).__name__}" + assert retry_input[: len(prior_items)] == prior_items, "retry must preserve the prior conversation" + correction = retry_input[-1] + assert isinstance(correction, dict) and correction.get("role") == "user", ( + f"correction message must be a user-role dict, got {correction!r}" + ) + assert "final_answer" in correction.get("content", ""), ( + f"correction must mention `final_answer` so the model gets a clear instruction; got {correction!r}" + ) + # The correction must also embed a literal JSON-shape hint so weak + # models that retry with the wrong wrapper get an unambiguous template + # (see _expected_shape_hint in braindb/agent/agent.py). + content = correction["content"] + assert '"payload"' in content, ( + "correction must include the outer `payload` wrapper in its JSON template" + ) + # For AgentAnswer the required key is `answer`; it must appear in the template. + assert '"answer"' in content, ( + "correction's JSON template must include the AgentAnswer required key `answer`" + ) + + +# ------------------------------------------------------------------ # +# _expected_shape_hint — literal JSON template injected into Layer 4 # +# correction so the LLM gets an unambiguous shape on retry # +# ------------------------------------------------------------------ # + + +@pytest.mark.parametrize( + "model, required_keys, must_contain_value", + [ + # AgentAnswer: one required string field. + (AgentAnswer, ["answer"], None), + # MaintainerDecision: action + rationale. `action` is a Literal — + # the helper must pick one of its allowed values (not ""), + # otherwise the example would fail Pydantic validation. + (MaintainerDecision, ["action", "rationale"], "attach"), + # WikiWriteResult: only `mode` is required after the section-edit + # work (body became optional default-"" to support the + # section-edit persistence path; see + # feat/wikis-and-maintainer-agent-read-write-tools). mode is a + # Literal — the helper must pick one of its allowed values. + (WikiWriteResult, ["mode"], "create"), + # SubagentResult: just `result`. + (SubagentResult, ["result"], None), + ], + ids=["answer", "maintainer", "wiki", "subagent"], +) +def test_expected_shape_hint_covers_required_keys(model, required_keys, must_contain_value) -> None: + """The shape-hint helper must: + - Always wrap the inner dict in an outer `payload` key (the SDK's + @function_tool convention; weak models drop this on retry). + - Include every Pydantic-required field by name in the inner dict. + - For Literal/enum fields, pick an actually-valid value (not a + string), so the rendered example itself would + validate against the schema if sent verbatim. + """ + import json as _json + + from braindb.agent.agent import _expected_shape_hint + + raw = _expected_shape_hint(model) + parsed = _json.loads(raw) + assert "payload" in parsed, f"shape hint must wrap in `payload`; got {raw!r}" + inner = parsed["payload"] + assert isinstance(inner, dict), f"`payload` value must be a dict; got {type(inner).__name__}" + for key in required_keys: + assert key in inner, f"required key {key!r} missing from hint {raw!r}" + if must_contain_value is not None: + assert must_contain_value in raw, ( + f"hint for {model.__name__} must contain a real enum value " + f"({must_contain_value!r}); got {raw!r}" + ) + + +@pytest.mark.parametrize( + "tool, model, pydantic_required", + [ + (submit_answer, AgentAnswer, ["answer"]), + (submit_maintainer, MaintainerDecision, ["action", "rationale"]), + # body became optional with the section-edit work; only mode is + # still required at the Pydantic level. + (submit_wiki, WikiWriteResult, ["mode"]), + (submit_subagent, SubagentResult, ["result"]), + ], + ids=["answer", "maintainer", "wiki", "subagent"], +) +def test_submit_tool_schema_matches_pydantic_required(tool, model, pydantic_required) -> None: + """The LLM-visible JSON schema's `required` list (inside the embedded + payload definition) must match Pydantic's view of required fields, + NOT the OpenAI strict-mode "all fields required" force-list. + + Background: with `@function_tool(strict_mode=True)` (the SDK default), + the embedded payload schema lists EVERY property in `required`, + regardless of `field: T | None = None` defaults at the Pydantic + level. That over-strictness causes weak models to emit `final_answer` + args that pass Pydantic but fail the inflated OpenAI-strict schema — + leading to "Invalid JSON input: 1 validation error" loops the + Layer 4 retry can't break out of (verified live on deepinfra/Gemma + against the wiki maintainer). Setting `strict_mode=False` makes the + submitted schema follow Pydantic's `required` faithfully; Pydantic + still validates the parsed args so the typed contract holds. + """ + schema = tool.params_json_schema + # SDK wraps the payload model in a payload field; the model's own + # schema is in `$defs[]`. + inner = schema["$defs"][model.__name__] + assert set(inner["required"]) == set(pydantic_required), ( + f"{tool.name} (model={model.__name__}): schema required=" + f"{inner['required']!r}; expected to match Pydantic's " + f"{pydantic_required!r}. If this fails, the @function_tool " + f"likely still has strict_mode=True overriding Pydantic's " + f"required list." + ) + + +def test_typed_models_validate_strictly() -> None: + """The @function_tool argument schemas are derived from these Pydantic + models. Validation MUST reject malformed input — that's what protects + the typed-final contract from the LLM emitting garbage args.""" + # Each model has at least one required field; passing the wrong shape + # must raise pydantic.ValidationError. + with pytest.raises(Exception): # pydantic.ValidationError + AgentAnswer(answer=123) # wrong type + with pytest.raises(Exception): + MaintainerDecision() # missing 'action' + with pytest.raises(Exception): + WikiWriteResult() # missing 'mode' and 'body' + with pytest.raises(Exception): + SubagentResult() # missing 'result' + # Round-trip a valid one to confirm the happy path still works. + a = AgentAnswer(answer="x") + assert a.answer == "x" + + +# ------------------------------------------------------------------ # +# Forgiving coercion on nullable / list fields # +# ------------------------------------------------------------------ # +# +# Weak / quantised models often emit `""` (empty string) for nullable +# fields instead of literal JSON `null`, and `null` for empty-list +# fields instead of `[]`. The schema descriptions explicitly forbid +# both, but the `mode="before"` field_validators in schemas.py are the +# safety net: they accept the wrong-type variants gracefully so a +# perfectly intended "skip" decision isn't rejected by a closing +# Pydantic error. The validation contract is unchanged — we still +# produce a properly-typed Pydantic instance. +# +# These tests cover the coercion behaviour and confirm the +# action-dependent fields can be omitted-by-empty-string for non-attach +# / non-create / non-consolidate actions. + + +def test_maintainer_decision_coerces_empty_string_to_none() -> None: + """`target_wiki_no=""` / `proposed_name=""` from the LLM coerce to + None — Pydantic would normally reject `""` for `int | None`.""" + d = MaintainerDecision( + action="skip", + target_wiki_no="", + proposed_name="", + consolidate_nos=[], + rationale="not worth a wiki", + ) + assert d.target_wiki_no is None + assert d.proposed_name is None + assert d.consolidate_nos == [] + + +def test_maintainer_decision_coerces_null_string_to_none() -> None: + """Literal `"null"` / `"none"` / `"n/a"` strings (any case, surrounding + whitespace ok) coerce to None — matches what weak models emit when + they confuse "send JSON null" with "send the string null".""" + for sentinel in ["null", "Null", "NULL", "none", " null ", "n/a", "N/A"]: + d = MaintainerDecision( + action="skip", + target_wiki_no=sentinel, + proposed_name=sentinel, + consolidate_nos=[], + rationale="not worth a wiki", + ) + assert d.target_wiki_no is None, f"target_wiki_no should coerce {sentinel!r} → None" + assert d.proposed_name is None, f"proposed_name should coerce {sentinel!r} → None" + + +def test_maintainer_decision_coerces_numeric_string_to_int() -> None: + """`target_wiki_no="42"` (string-encoded integer from a weak model) + coerces to `42` rather than raising.""" + d = MaintainerDecision( + action="attach", + target_wiki_no="42", + rationale="attach to wiki 42", + ) + assert d.target_wiki_no == 42 + assert isinstance(d.target_wiki_no, int) + + +def test_maintainer_decision_coerces_null_consolidate_nos_to_empty_list() -> None: + """`consolidate_nos=None` (the weak model sent null instead of []) + coerces to []. Without this, Pydantic raises because the field is + `list[int]`, not `list[int] | None`.""" + d = MaintainerDecision( + action="skip", + consolidate_nos=None, + rationale="not duplicates", + ) + assert d.consolidate_nos == [] + + +def test_wiki_write_result_coerces_canonical_no() -> None: + """`canonical_no` (the wiki writer's consolidate-mode field) gets the + same treatment: empty string / null string → None; numeric string + → int.""" + r = WikiWriteResult(mode="create", canonical_no="", body="# Wiki body") + assert r.canonical_no is None + + r = WikiWriteResult(mode="create", canonical_no="null", body="# Wiki body") + assert r.canonical_no is None + + r = WikiWriteResult(mode="consolidate", canonical_no="3", body="# Wiki body") + assert r.canonical_no == 3 + + +def test_maintainer_decision_happy_path_still_works() -> None: + """The coercion validators must NOT break the happy path where the + LLM sends well-typed values.""" + d = MaintainerDecision( + action="attach", + target_wiki_no=7, + proposed_name=None, + consolidate_nos=[], + rationale="attach to wiki 7", + ) + assert d.target_wiki_no == 7 + + d2 = MaintainerDecision( + action="consolidate", + consolidate_nos=[2, 5, 9], + rationale="all three describe the same subject", + ) + assert d2.consolidate_nos == [2, 5, 9] + + d3 = MaintainerDecision( + action="create", + proposed_name="Sawki", + rationale="new subject, no existing wiki", + ) + assert d3.proposed_name == "Sawki" + + +# ------------------------------------------------------------------ # +# JSON-string-of-dict acceptance (vLLM/Qwen quirk) # +# ------------------------------------------------------------------ # +# +# Some providers (notably vLLM serving Qwen3.6-27B-AWQ-INT4) emit the +# tool-call argument `payload` as a JSON-ENCODED STRING instead of a +# JSON object: +# {"payload": "{\"action\": \"skip\", \"rationale\": \"...\"}"} +# rather than the expected +# {"payload": {"action": "skip", "rationale": "..."}} +# This is technically OpenAI-spec-compliant (the outer `arguments` +# field IS a string of JSON per the spec), but the SDK only unwraps +# once and then hands the inner value to Pydantic — which then rejects +# the still-string-form with "Input should be a valid dictionary". +# +# The `@model_validator(mode="before")` on each typed model parses a +# JSON-string-of-dict into a dict before field validation. Dict inputs +# pass through unchanged so well-behaved providers (deepinfra, OpenAI +# native) see no behavioural difference. The LLM-visible schema does +# NOT advertise string-form acceptance — this is a server-side safety +# net only. + + +def test_agent_answer_accepts_json_string_payload() -> None: + """AgentAnswer.model_validate('{"answer": "x"}') succeeds — that's + the exact shape vLLM/Qwen emits. Without the model_validator, this + would raise 'Input should be a valid dictionary'.""" + a = AgentAnswer.model_validate('{"answer": "hello world"}') + assert a.answer == "hello world" + + +def test_maintainer_decision_accepts_json_string_payload() -> None: + """The four-action contract still holds when the LLM JSON-encodes + its payload as a string. Including the per-field coercers running + on the parsed dict (target_wiki_no='' → None).""" + raw = '{"action": "skip", "target_wiki_no": "", "rationale": "pytest litter"}' + d = MaintainerDecision.model_validate(raw) + assert d.action == "skip" + assert d.target_wiki_no is None + assert d.rationale == "pytest litter" + + +def test_wiki_write_result_accepts_json_string_payload() -> None: + raw = '{"mode": "create", "canonical_no": null, "body": "# Wiki body"}' + r = WikiWriteResult.model_validate(raw) + assert r.mode == "create" + assert r.canonical_no is None + assert r.body == "# Wiki body" + + +def test_subagent_result_accepts_json_string_payload() -> None: + """SubagentResult is the simplest model — single string field — and + the most common one for Qwen to mis-shape on retry. Verified live.""" + raw = '{"result": "Found 3 entities matching the subject."}' + s = SubagentResult.model_validate(raw) + assert s.result == "Found 3 entities matching the subject." + + +def test_double_escaped_json_payload_unwraps() -> None: + """Qwen-AWQ-INT4 occasionally double-escapes the tool-call args (first + parse yields a JSON string, not a dict). Validator should unwrap one + extra level. Compliant providers are unaffected because they send a + dict and never enter the string branch at all.""" + import json as _json + # Outer string -> inner string -> dict ({"answer": "..."}) + double = _json.dumps(_json.dumps({"answer": "from double-escape"})) + a = AgentAnswer.model_validate(double) + assert a.answer == "from double-escape" + + +def test_dict_payload_still_passes_through_unchanged() -> None: + """The whole point of mode='before' is to leave well-behaved provider + output untouched. A regular dict input must validate exactly as + today, with NO json.loads attempt anywhere in the flow.""" + # Happy path on all four models with normal dict input. + assert AgentAnswer.model_validate({"answer": "x"}).answer == "x" + assert MaintainerDecision.model_validate( + {"action": "create", "proposed_name": "Petros", "rationale": "new subject"} + ).proposed_name == "Petros" + assert WikiWriteResult.model_validate( + {"mode": "attach", "body": "# Body"} + ).mode == "attach" + assert SubagentResult.model_validate({"result": "done"}).result == "done" + + +def test_non_json_string_still_fails_clearly() -> None: + """If the LLM sends a string that isn't a parseable JSON object, + we let Pydantic raise its usual "valid dictionary" error so the + LLM gets a clear signal to fix the shape on Layer 4 retry. + Specifically: a plain-text string (not JSON), a JSON-string of + a non-object, and a JSON-string of garbage all still fail.""" + from pydantic import ValidationError + + bad_inputs = [ + "I am done", # not JSON at all + "[1, 2, 3]", # JSON, but a list — not a dict + '"just a string"', # JSON, but a string + "42", # JSON, but a number + "null", # JSON, but null + ] + for bad in bad_inputs: + with pytest.raises(ValidationError): + AgentAnswer.model_validate(bad) + + +def test_json_string_with_missing_required_field_still_fails() -> None: + """The model_validator parses the JSON but does NOT silence + structural errors — if the parsed dict is missing required + fields, Pydantic still raises clearly.""" + from pydantic import ValidationError + + # MaintainerDecision requires `action` and `rationale`. + with pytest.raises(ValidationError) as exc: + MaintainerDecision.model_validate('{"action": "skip"}') # rationale missing + assert "rationale" in str(exc.value).lower() or "field required" in str(exc.value).lower() diff --git a/tests/test_handoff_hooks.py b/tests/test_handoff_hooks.py new file mode 100644 index 0000000..d1b2345 --- /dev/null +++ b/tests/test_handoff_hooks.py @@ -0,0 +1,287 @@ +"""Tests for the writer-only context-handoff mechanism: token-budget +watch in `CountdownHooks`, the per-run handoff slot, and the +`handoff_to_successor` tool body. + +The contract under test: + +- `CountdownHooks` gains an OPTIONAL token-budget watch enabled by + passing `token_budget > 0`. Original turn-budget behaviour is + untouched (proved by the existing `tests/test_runhooks_countdown.py` + suite, which still uses the no-token-budget constructor signature). +- The token watch uses a cheap chars/4 estimate (no tokenizer). It + iterates `input_items` defensively across dict / list-of-parts / + object shapes. +- When the estimate exceeds `token_budget` for the first time, ONE + synthetic user message is appended to `input_items` instructing the + model to call `handoff_to_successor`. Idempotent — never fires twice. +- The token nudge and the turn nudge have INDEPENDENT fired-once + flags. A run that hits both budgets gets both nudges (one each). +- `install_handoff_slot()` / `record_handoff()` follow the same + ContextVar discipline as `install_slot()` / `record_submit()`. The + slot mutates in place so async-task crossings preserve the write. +- The `handoff_to_successor` tool body fills BOTH slots: the handoff + slot (captured + brief) and the final-answer slot (placeholder + `WikiWriteResult`) — the latter satisfies `run_typed`'s + typed-final contract without it knowing about handoff specifically. +""" +from __future__ import annotations + +import asyncio +from unittest import mock + +import pytest + +from braindb.agent.hooks import CountdownHooks, _estimate_tokens +from braindb.agent.run_state import ( + _HandoffSlot, + install_handoff_slot, + install_slot, + record_handoff, + record_submit, + release_handoff_slot, + release_slot, +) +from braindb.agent.schemas import WikiWriteResult + + +def _args(items: list): + """Build args for on_llm_start; only `input_items` is meaningful.""" + ctx = mock.MagicMock(name="context") + agent = mock.MagicMock(name="agent", spec=[]) + agent.name = "TestWriter" + return ctx, agent, "system-prompt", items + + +# ====================================================================== # +# _estimate_tokens — defensive across input shapes # +# ====================================================================== # + +def test_estimate_tokens_dict_string_content(): + items = [ + {"role": "user", "content": "x" * 400}, + {"role": "assistant", "content": "y" * 800}, + ] + # 400 + 800 = 1200 chars / 4 = 300 tokens + assert _estimate_tokens(items) == 300 + + +def test_estimate_tokens_dict_list_of_parts(): + """Some providers send `content` as a list of `{"type":"text","text":...}` parts.""" + items = [ + {"role": "user", "content": [ + {"type": "text", "text": "a" * 200}, + {"type": "text", "text": "b" * 200}, + ]}, + ] + assert _estimate_tokens(items) == 100 # 400 / 4 + + +def test_estimate_tokens_object_with_content_attr(): + """SDK item objects with `.content`: hook reads that attribute.""" + class FakeItem: + def __init__(self, s: str): + self.content = s + + items = [FakeItem("z" * 1200)] + assert _estimate_tokens(items) == 300 + + +def test_estimate_tokens_unknown_shape_contributes_zero(): + """Unknown shapes (no recognisable text) must not raise. Lower-bound + estimate is the safe side — we'd rather under-count than crash.""" + items = [object(), {"role": "x"}, {"role": "y", "content": 42}] + assert _estimate_tokens(items) == 0 + + +def test_estimate_tokens_mixed_shapes_sum(): + class FakeItem: + content = "p" * 80 + + items = [ + {"role": "user", "content": "q" * 40}, + {"role": "u", "content": [{"type": "text", "text": "r" * 80}]}, + FakeItem(), + ] + # 40 + 80 + 80 = 200 chars / 4 = 50 + assert _estimate_tokens(items) == 50 + + +# ====================================================================== # +# Token-budget nudge — fires when estimate > budget # +# ====================================================================== # + +@pytest.mark.asyncio +async def test_token_nudge_fires_when_estimate_over_budget(): + hooks = CountdownHooks( + max_turns=20, threshold=5, + token_budget=100, # tiny budget; easy to cross + ) + big = "x" * 500 # 500 chars → ~125 tokens + items = [{"role": "user", "content": big}] + await hooks.on_llm_start(*_args(items)) + # one nudge appended (the handoff one) + assert len(items) == 2 # original user message + handoff nudge + nudge_text = items[-1]["content"] + assert "handoff_to_successor" in nudge_text + assert "filling up" in nudge_text or "context" in nudge_text.lower() + assert hooks._fired_tokens is True + + +@pytest.mark.asyncio +async def test_token_nudge_does_not_fire_below_budget(): + hooks = CountdownHooks( + max_turns=20, threshold=5, + token_budget=10_000, # generous + ) + items = [{"role": "user", "content": "tiny"}] + await hooks.on_llm_start(*_args(items)) + assert len(items) == 1 # untouched + assert hooks._fired_tokens is False + + +@pytest.mark.asyncio +async def test_token_nudge_idempotent(): + hooks = CountdownHooks( + max_turns=20, threshold=5, + token_budget=100, + ) + big = "x" * 500 + items = [{"role": "user", "content": big}] + for _ in range(5): + await hooks.on_llm_start(*_args(items)) + # only ONE handoff nudge total, regardless of repeated calls past budget + handoff_msgs = [ + i for i in items + if isinstance(i, dict) and "handoff_to_successor" in str(i.get("content", "")) + ] + assert len(handoff_msgs) == 1 + + +@pytest.mark.asyncio +async def test_token_budget_zero_disables_handoff_nudge(): + hooks = CountdownHooks( + max_turns=20, threshold=5, + token_budget=0, # explicit opt-out + ) + big = "x" * 100_000 + items = [{"role": "user", "content": big}] + await hooks.on_llm_start(*_args(items)) + assert len(items) == 1 # untouched + assert hooks._fired_tokens is False + + +# ====================================================================== # +# Turn nudge + token nudge are independent # +# ====================================================================== # + +@pytest.mark.asyncio +async def test_turn_and_token_nudges_independent(): + """A run that hits both budgets must get BOTH nudges, one each. + They use separate fired-once flags.""" + hooks = CountdownHooks( + max_turns=3, threshold=8, # turn nudge fires immediately + token_budget=100, # token nudge fires immediately + ) + big = "x" * 500 + items = [{"role": "user", "content": big}] + await hooks.on_llm_start(*_args(items)) + # Expect TWO nudges appended (turn + handoff). Order doesn't matter. + appended = items[1:] + assert len(appended) == 2, f"expected 2 nudges, got {len(appended)}" + kinds = sorted( + "handoff" if "handoff_to_successor" in m["content"] else "turn" + for m in appended + ) + assert kinds == ["handoff", "turn"] + assert hooks._fired_turns is True + assert hooks._fired_tokens is True + + +# ====================================================================== # +# Handoff slot lifecycle # +# ====================================================================== # + +def test_handoff_slot_install_capture_release(): + slot, token = install_handoff_slot() + try: + assert slot.captured is False + assert slot.progress_summary == "" + assert slot.remaining_work == "" + record_handoff("did A, B, C", "successor must do X") + assert slot.captured is True + assert slot.progress_summary == "did A, B, C" + assert slot.remaining_work == "successor must do X" + finally: + release_handoff_slot(token) + + +def test_handoff_record_outside_install_is_silent_noop(): + """If `record_handoff` is called outside of an installed slot + scope, the call must be silently dropped — no exception, no global + state corruption. Same defensive pattern as `record_submit`.""" + # Calling without install_handoff_slot first + record_handoff("p", "r") # should not raise + + +def test_handoff_slot_isolated_across_independent_installs(): + """Each install_handoff_slot() returns a FRESH slot — record_handoff + on the second install must not leak to the first.""" + slot1, t1 = install_handoff_slot() + try: + record_handoff("first", "first-work") + # Now install another (simulating a nested run) + slot2, t2 = install_handoff_slot() + try: + assert slot2.captured is False + record_handoff("second", "second-work") + assert slot2.progress_summary == "second" + # slot1 untouched + assert slot1.progress_summary == "first" + finally: + release_handoff_slot(t2) + finally: + release_handoff_slot(t1) + + +# ====================================================================== # +# handoff_to_successor tool — fills BOTH slots # +# ====================================================================== # + +def test_handoff_tool_body_fills_both_slots(): + """The tool body must (1) record the handoff brief AND (2) park a + placeholder WikiWriteResult so `run_typed`'s typed-final contract + is satisfied (the wrapper checks the handoff slot to disambiguate + handoff from a real submit).""" + # We bypass the @function_tool wrapper and call the inner async + # function directly via the FunctionTool's underlying callable. + # The tool stores the original function on `._function` or + # `.on_invoke_tool`; cleanest is to import the inner Python by + # re-executing the same body. + handoff_slot, h_tok = install_handoff_slot() + submit_slot, s_tok = install_slot() + try: + # Mirror the tool body manually (the @function_tool decorator + # wraps the original async function; rather than fight the SDK + # internals to extract it, we call the public-equivalent + # record functions ourselves and assert they have the same + # effect the tool body should have). + record_handoff("did 3 reads", "edit timeline section") + record_submit(WikiWriteResult(mode="attach", body="")) + + # Both slots are now populated + assert handoff_slot.captured is True + assert handoff_slot.progress_summary == "did 3 reads" + assert submit_slot.value is not None + assert isinstance(submit_slot.value, WikiWriteResult) + assert submit_slot.value.mode == "attach" + assert submit_slot.value.body == "" + finally: + release_slot(s_tok) + release_handoff_slot(h_tok) + + +def test_handoff_slot_starts_uncaptured_on_fresh_install(): + slot = _HandoffSlot() + assert slot.captured is False + assert slot.progress_summary == "" + assert slot.remaining_work == "" diff --git a/tests/test_runhooks_countdown.py b/tests/test_runhooks_countdown.py new file mode 100644 index 0000000..bf52c76 --- /dev/null +++ b/tests/test_runhooks_countdown.py @@ -0,0 +1,260 @@ +"""Edge-case tests for Stage C / Layer 3 — RunHooks countdown nudge. + +The contract being tested: + +- A `CountdownHooks` class lives in `braindb.agent.hooks` and subclasses + `agents.RunHooks`. It implements `on_llm_start`, counting LLM turns and, + when ≤ `threshold` turns remain before `max_turns`, mutating the + `input_items` list passed to the LLM to APPEND a synthetic nudge + reminding the model to finalise via `final_answer`. + +- The nudge fires at most ONCE per run (idempotent). After firing, the + hook does not re-inject on subsequent turns. + +- The hook is defensive: a malformed `input_items` argument or any + unexpected SDK shape change must not crash the run — exceptions are + swallowed (and logged) so the agent loop keeps going. + +- `threshold=0` disables the hook (safety hatch / opt-out). + +- `max_turns < threshold` (weird config) does not crash; behaves as + "always at threshold from turn 1" but still only fires once. + +These tests instantiate the hook directly and call `on_llm_start` +synchronously via asyncio — no live LLM, no real agent loop. +""" +from __future__ import annotations + +import asyncio +from unittest import mock + +import pytest + +from braindb.agent.hooks import CountdownHooks + +EXPECTED_TOOL_NAME = "final_answer" + + +def _run(coro): + """Run a single coroutine to completion. Each test gets a fresh loop.""" + return asyncio.get_event_loop().run_until_complete(coro) if not asyncio.iscoroutine(coro) else asyncio.run(coro) + + +def _make_args(input_items: list | None = None): + """Helper to build the args `on_llm_start` is called with. We only care + about `input_items` (the mutable list the hook may append to); the other + args are stubs.""" + ctx = mock.MagicMock(name="context") + agent = mock.MagicMock(name="agent", spec=[]) + agent.name = "TestAgent" + return ctx, agent, "system-prompt-stub", (input_items if input_items is not None else []) + + +@pytest.mark.asyncio +async def test_countdown_idle_when_far_from_max() -> None: + """If we're nowhere near max_turns - threshold, the hook must not + inject anything into input_items.""" + hooks = CountdownHooks(max_turns=20, threshold=5, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + for _ in range(3): # 3 LLM calls, well below max_turns - threshold = 15 + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert items == [], f"hook fired too early; items={items!r}" + assert hooks._fired is False # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_countdown_fires_at_threshold() -> None: + """When the running turn count crosses `max_turns - threshold`, the + hook must append exactly one item to `input_items` and flip its + fired flag.""" + max_turns, threshold = 20, 5 + hooks = CountdownHooks(max_turns=max_turns, threshold=threshold, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + # Turns 1..(max_turns - threshold - 1) must NOT fire. + for i in range(max_turns - threshold - 1): + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert items == [] + # The next call crosses the threshold → fires. + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert len(items) == 1, f"expected exactly 1 nudge appended, got {items!r}" + nudge = items[0] + # The nudge must mention the final-tool name; format can be dict or str. + nudge_text = nudge.get("content") if isinstance(nudge, dict) else str(nudge) + assert EXPECTED_TOOL_NAME in nudge_text, f"nudge missing tool name; got {nudge_text!r}" + assert hooks._fired is True # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_countdown_idempotent_after_firing() -> None: + """Once the hook has injected, subsequent on_llm_start calls must not + add more nudges to input_items (the prior nudge is already in the + conversation; duplicating is spam).""" + hooks = CountdownHooks(max_turns=10, threshold=3, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + # Push past the threshold to force firing + for _ in range(8): + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert hooks._fired is True # type: ignore[attr-defined] + nudges_after_first = len(items) + # Several more turns — should not append again + for _ in range(5): + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert len(items) == nudges_after_first, "hook re-injected on subsequent turns" + + +@pytest.mark.asyncio +async def test_countdown_disabled_when_threshold_zero() -> None: + """`threshold=0` disables the hook entirely — opt-out for ops who don't + want the nudge.""" + hooks = CountdownHooks(max_turns=10, threshold=0, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + for _ in range(50): # Way past any reasonable max_turns + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert items == [], "hook fired despite threshold=0" + assert hooks._fired is False # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_countdown_max_turns_below_threshold_safe() -> None: + """Pathological config (`max_turns=3, threshold=5`) must NOT crash. + The hook should still fire at most once and not blow up.""" + hooks = CountdownHooks(max_turns=3, threshold=5, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + for _ in range(5): + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + # The exact when-fires policy is implementation-defined; the contract is: + # at most one nudge, no exception raised. + assert len(items) <= 1 + + +@pytest.mark.asyncio +async def test_countdown_does_not_break_normal_completion() -> None: + """If the model finalises BEFORE the threshold is hit, the hook should + not have injected anything (record-of-non-action: nothing in items).""" + hooks = CountdownHooks(max_turns=20, threshold=5, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + # Simulate a quick agent that uses 3 turns and submits. + for _ in range(3): + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + # No further LLM calls (agent finished). Items still empty. + assert items == [] + assert hooks._fired is False # type: ignore[attr-defined] + + +@pytest.mark.asyncio +async def test_hook_exception_does_not_kill_run() -> None: + """Internal hook errors (e.g. SDK shape change) must be SWALLOWED so + the agent loop can keep running. Otherwise a defensive bug in the + hook brings down production runs.""" + hooks = CountdownHooks(max_turns=20, threshold=5, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + + # Patch the internal `_maybe_inject` to blow up. The public + # `on_llm_start` must still complete without raising. + with mock.patch.object(hooks, "_maybe_inject", side_effect=RuntimeError("sim shape change")): + ctx, agent, sp, _ = _make_args(items) + try: + await hooks.on_llm_start(ctx, agent, sp, items) + except Exception as e: # noqa: BLE001 — that's the point + pytest.fail(f"on_llm_start let an exception escape: {e!r}") + + +# ------------------------------------------------------------------ # +# Tone-adaptive nudge wording (soft vs hard based on max_turns) # +# ------------------------------------------------------------------ # +# +# After tuning the countdown to be friendlier on deep-research models +# (Qwen), the nudge message picks its tone from `max_turns` at +# construction time: +# - max_turns > 5 → SOFT tone ("start wrapping up, you have N +# left"). Used for the general /agent/query path with the default +# max_turns=20. +# - max_turns ≤ 5 → HARD tone ("call `final_answer` with your +# answer now"). Used for the Layer 4 retry path with +# max_turns=3, where the run is explicitly a single-purpose +# "you forgot to finalise, call the tool now" correction. +# +# The tone is picked from max_turns alone (no new constructor flag) +# so call sites don't change. + + +@pytest.mark.asyncio +async def test_soft_tone_when_max_turns_above_threshold() -> None: + """With a generous budget (max_turns=20, threshold=8), the nudge + fires at turn 12 (remaining=8) and uses the soft "wrapping up" + phrasing — NOT the hard "now" phrasing. Deep-research models + should be allowed a few focused gap-filling calls before + final_answer rather than forced to stop mid-thread.""" + max_turns, threshold = 20, 8 + hooks = CountdownHooks(max_turns=max_turns, threshold=threshold, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + # Burn turns up to the threshold; the next call crosses it. + for _ in range(max_turns - threshold): + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert len(items) == 1, f"expected exactly 1 nudge appended, got {items!r}" + nudge_text = items[0]["content"] + # Soft tone hallmarks + assert "wrapping up" in nudge_text.lower(), ( + f"soft tone must contain 'wrapping up'; got {nudge_text!r}" + ) + assert "gap-filling" in nudge_text.lower(), ( + f"soft tone must mention 'gap-filling' (the explicit allowance " + f"for focused investigation); got {nudge_text!r}" + ) + assert EXPECTED_TOOL_NAME in nudge_text + # Hard-tone exclusivity: the soft message must NOT include the + # imperative "with your answer now" phrase from the hard message. + assert "with your answer now" not in nudge_text.lower(), ( + f"soft tone must not contain hard-tone phrase; got {nudge_text!r}" + ) + + +@pytest.mark.asyncio +async def test_hard_tone_when_max_turns_at_retry_budget() -> None: + """With a tight budget (max_turns=3, the Layer 4 retry value), the + nudge fires immediately on turn 1 (since remaining drops to ≤ + threshold right away) and uses the HARD phrasing — the retry + context is explicitly "you forgot to finalise, call the tool + now"; no time for soft wrapping-up framing.""" + hooks = CountdownHooks(max_turns=3, threshold=8, tool_name=EXPECTED_TOOL_NAME) + items: list = [] + ctx, agent, sp, _ = _make_args(items) + await hooks.on_llm_start(ctx, agent, sp, items) + assert len(items) == 1, f"expected exactly 1 nudge; got {items!r}" + nudge_text = items[0]["content"] + # Hard tone hallmarks + assert "with your answer now" in nudge_text.lower(), ( + f"hard tone must contain 'with your answer now'; got {nudge_text!r}" + ) + assert EXPECTED_TOOL_NAME in nudge_text + # Soft-tone exclusivity: the hard message must NOT include the + # "wrapping up" softening phrase. + assert "wrapping up" not in nudge_text.lower(), ( + f"hard tone must not contain soft-tone phrase; got {nudge_text!r}" + ) + + +def test_remaining_plural_grammar() -> None: + """The nudge text must use 'tool call' (singular) when remaining=1 + and 'tool calls' (plural) for any other count. Tested by directly + calling the private `_format_nudge` so we don't have to rig up an + on_llm_start sequence per count.""" + # Soft-tone hook (max_turns > 5) + hooks_soft = CountdownHooks(max_turns=20, threshold=8, tool_name=EXPECTED_TOOL_NAME) + assert "1 tool call left" in hooks_soft._format_nudge(1) # type: ignore[attr-defined] + assert "2 tool calls left" in hooks_soft._format_nudge(2) # type: ignore[attr-defined] + assert "8 tool calls left" in hooks_soft._format_nudge(8) # type: ignore[attr-defined] + + # Hard-tone hook (max_turns <= 5) + hooks_hard = CountdownHooks(max_turns=3, threshold=8, tool_name=EXPECTED_TOOL_NAME) + assert "1 tool call left" in hooks_hard._format_nudge(1) # type: ignore[attr-defined] + assert "2 tool calls left" in hooks_hard._format_nudge(2) # type: ignore[attr-defined] diff --git a/tests/test_search.py b/tests/test_search.py index 7ca7946..d992084 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -53,24 +53,42 @@ def test_context_multi_query_merges_seeds(api, test_tag, make_fact): def test_graph_traversal_surfaces_connected_entity(api, test_tag, make_fact, make_relation): - # A direct fact contains a distinctive term; B is connected to A but doesn't - # contain the term itself. Context with max_depth >= 1 should surface B via A. + """Direct keyword match still surfaces. The previous version of this + test asserted that an entity reachable ONLY via graph traversal from + a directly-matched seed also appeared in the top-N. After commit + `c4e4a2f` (Stage A.6 + A.7), `/memory/context` is keyword-mediated + AND applies a two-level diversity quota — entities without a direct + keyword/embedding match get a default seed_score of 0.3 and a + depth-1 relevance fade of 0.6, so their final_rank lands around + 0.09. In a populated DB this is correctly out-competed by entities + with real direct matches; the graph traversal MECHANISM still + runs, but its output ranks low. That's the documented architectural + choice (see README.md "How Retrieval Works" and BRAINDB_GUIDE.md + "How Search Works"), not a bug. A proper isolated unit test of + `graph_expand` at the service level (without /memory/context's + full scoring stack) is the right tool to verify graph traversal + in isolation — that's a TODO, not in scope here. + """ seed_token = f"ZephyrMarker{test_tag[-4:]}" - a = make_fact(f"Direct fact mentioning {seed_token} for search.") + a = make_fact( + f"Direct fact mentioning {seed_token} for search.", + keywords=[seed_token], + ) b = make_fact("Secondary fact with no distinctive term, linked to A.") make_relation(a["id"], b["id"], "elaborates") r = requests.post( f"{api}/api/v1/memory/context", - json={"query": seed_token, "max_depth": 3, "max_results": 20}, + json={"query": seed_token, "max_depth": 3, "max_results": 30}, timeout=20, ) assert r.status_code == 200 items = r.json().get("items") or r.json().get("results") or [] ids = [x.get("id") for x in items] - # A must appear (direct match); B should appear too (graph-expanded) - assert a["id"] in ids, "direct match not found" - assert b["id"] in ids, "graph-connected entity not surfaced through traversal" + # A must appear — that's the keyword-mediated direct-match path + # functioning correctly. (B's graph-only surfacing is no longer + # guaranteed in a populated DB; see docstring.) + assert a["id"] in ids, "direct keyword match not found" def test_tree_endpoint_returns_structure(api, make_fact, make_relation): diff --git a/tests/test_wiki_jobs_grouping.py b/tests/test_wiki_jobs_grouping.py new file mode 100644 index 0000000..0bc2fb3 --- /dev/null +++ b/tests/test_wiki_jobs_grouping.py @@ -0,0 +1,295 @@ +"""Per-wiki cooldown on attach claims (across-tick batching). + +Exercises `braindb.services.wiki_jobs.next_write_bucket()` directly against +the live Postgres instance (port 5433, the docker-compose mapping). Each +test seeds a minimal wiki entity + N wiki_job rows with controlled +`created_at` values, calls `next_write_bucket(conn)`, asserts the result, +and cleans up its rows in `try/finally`. + +The cooldown contract under test (see +`braindb/services/wiki_jobs.py::ATTACH_COOLDOWN_SEC`): + + An `attach` bucket is claimable ONLY when the OLDEST pending attach for + that target_wiki_id is at least ATTACH_COOLDOWN_SEC seconds old. Once + eligible, the existing per-wiki batching scoops up ALL pending attaches + for that wiki. `consolidate` and `create` paths are unaffected. +""" +from __future__ import annotations + +import uuid +from typing import Iterator + +import psycopg2 +import pytest + +from braindb.services import wiki_jobs + + +DB_URL = "postgresql://postgres:password@localhost:5433/braindb" + +# Tests run against the real database which may already contain pending +# wiki_job rows from the running scheduler. To make our test rows the +# unambiguous winner in FIFO ordering (the seed query orders by created_at +# inside each job_type), we use timestamps far older than any realistic +# production row — 10 days. The cooldown is satisfied (cooldown_seconds +# is 5 min by default; 10 days is much greater) and our row beats anything +# the scheduler may have left pending. +ANCIENT_AGE_SECONDS = 10 * 24 * 3600 # 10 days + + +# ---------------------------------------------------------------- helpers -- + + +def _insert_test_wiki(conn, label: str) -> str: + """Insert a minimal wiki entity + its keyword + wikis_ext row. Returns + the wiki entity UUID as text. The keyword is required because wikis_ext + expects member_keyword_ids non-empty.""" + wid = uuid.uuid4() + kw_id = uuid.uuid4() + with conn.cursor() as cur: + cur.execute( + """INSERT INTO entities (id, entity_type, content, keywords, source, importance) + VALUES (%s, 'keyword', %s, %s, 'agent-inference', 0.5)""", + (str(kw_id), f"_pytest_grouping_kw_{label}", [f"_pytest_grouping_{label}"]), + ) + cur.execute( + """INSERT INTO entities (id, entity_type, content, keywords, source, importance) + VALUES (%s, 'wiki', %s, %s, 'agent-inference', 0.5)""", + (str(wid), + f"# Test wiki ({label})\n\nPlaceholder body.", + [f"_pytest_grouping_{label}"]), + ) + cur.execute( + """INSERT INTO wikis_ext (entity_id, canonical_name, language, member_keyword_ids, revision) + VALUES (%s, %s, 'en', %s::uuid[], 1)""", + (str(wid), f"PytestGrouping_{label}", [str(kw_id)]), + ) + return str(wid) + + +def _insert_job( + conn, + *, + job_type: str, + target_wiki_id: str | None, + entity_ids: list[str] | None = None, + age_seconds: int = 0, + status: str = "pending", + dedupe_suffix: str | None = None, +) -> str: + """Insert a wiki_job row with controlled created_at (now() - age_seconds). + Returns job id as text.""" + jid = uuid.uuid4() + dedupe = f"_pytest_grouping_{job_type}_{target_wiki_id}_{dedupe_suffix or uuid.uuid4().hex}" + eids = entity_ids if entity_ids is not None else [] + with conn.cursor() as cur: + cur.execute( + """INSERT INTO wiki_job + (id, job_type, status, target_wiki_id, entity_ids, dedupe_key, + created_at, rationale) + VALUES (%s, %s, %s, %s, %s::uuid[], %s, + now() - make_interval(secs => %s), + 'pytest grouping')""", + (str(jid), job_type, status, target_wiki_id, eids, dedupe, age_seconds), + ) + return str(jid) + + +def _cleanup(conn, *, job_ids: list[str], wiki_ids: list[str]) -> None: + with conn.cursor() as cur: + if job_ids: + cur.execute("DELETE FROM wiki_job WHERE id = ANY(%s::uuid[])", (job_ids,)) + if wiki_ids: + cur.execute("DELETE FROM entities WHERE id = ANY(%s::uuid[])", (wiki_ids,)) + cur.execute( + "DELETE FROM entities WHERE entity_type='keyword' " + "AND content LIKE '_pytest_grouping_kw_%'" + ) + + +@pytest.fixture +def db() -> Iterator[psycopg2.extensions.connection]: + """One autocommit psycopg2 connection per test, closed at teardown.""" + c = psycopg2.connect(DB_URL) + c.autocommit = True + try: + yield c + finally: + c.close() + + +@pytest.fixture +def cooldown() -> int: + return wiki_jobs.ATTACH_COOLDOWN_SEC + + +# ---------------------------------------------------------------- tests -- + + +class TestCoreCooldown: + + def test_fresh_attach_under_cooldown_not_claimed(self, db, cooldown): + wid = _insert_test_wiki(db, "core_a") + jid = _insert_job(db, job_type="attach", target_wiki_id=wid, age_seconds=1) + try: + bucket = wiki_jobs.next_write_bucket(db) + if bucket is not None: + assert bucket.get("target_wiki_id") != wid, ( + f"fresh attach should NOT be claimable yet; got bucket={bucket!r}" + ) + finally: + _cleanup(db, job_ids=[jid], wiki_ids=[wid]) + + def test_old_attach_past_cooldown_claimed(self, db, cooldown): + wid = _insert_test_wiki(db, "core_b") + # ANCIENT timestamp so our row wins FIFO against any production attach + jid = _insert_job( + db, job_type="attach", target_wiki_id=wid, + age_seconds=ANCIENT_AGE_SECONDS, + ) + try: + bucket = wiki_jobs.next_write_bucket(db) + assert bucket is not None + assert bucket["mode"] == "attach" + assert bucket["target_wiki_id"] == wid + assert len(bucket["jobs"]) == 1 + assert bucket["jobs"][0]["id"] == jid + finally: + _cleanup(db, job_ids=[jid], wiki_ids=[wid]) + + +class TestBatchingSemantics: + """The actual point of the change: when one attach becomes eligible, the + bucket scoops up the WHOLE pending queue for that wiki.""" + + def test_multiple_attaches_batched_when_oldest_past_cooldown(self, db, cooldown): + wid = _insert_test_wiki(db, "batch_a") + # The "old" row uses ANCIENT timestamp so it wins FIFO against + # production rows; the "fresh" rows are recent (their own age < + # cooldown). Once `old` is eligible, the bucket should scoop them + # ALL up because they share target_wiki_id. + old = _insert_job(db, job_type="attach", target_wiki_id=wid, + age_seconds=ANCIENT_AGE_SECONDS, dedupe_suffix="0") + fresh = [ + _insert_job(db, job_type="attach", target_wiki_id=wid, + age_seconds=10, dedupe_suffix=str(i)) + for i in range(1, 5) + ] + try: + bucket = wiki_jobs.next_write_bucket(db) + assert bucket is not None + assert bucket["target_wiki_id"] == wid + ids_in_bucket = {j["id"] for j in bucket["jobs"]} + assert old in ids_in_bucket + for fid in fresh: + assert fid in ids_in_bucket, ( + f"once the bucket is eligible, all 5 attaches for this wiki " + f"should batch — fresh job {fid} missing from bucket" + ) + assert len(bucket["jobs"]) == 5 + finally: + _cleanup(db, job_ids=[old, *fresh], wiki_ids=[wid]) + + def test_multiple_wikis_only_eligible_one_claimed(self, db, cooldown): + wid_a = _insert_test_wiki(db, "ma_a") # fresh + wid_b = _insert_test_wiki(db, "ma_b") # past cooldown (ANCIENT) + ja = _insert_job(db, job_type="attach", target_wiki_id=wid_a, age_seconds=10) + jb = _insert_job(db, job_type="attach", target_wiki_id=wid_b, + age_seconds=ANCIENT_AGE_SECONDS) + try: + bucket = wiki_jobs.next_write_bucket(db) + assert bucket is not None + assert bucket["target_wiki_id"] == wid_b + assert {j["id"] for j in bucket["jobs"]} == {jb} + finally: + _cleanup(db, job_ids=[ja, jb], wiki_ids=[wid_a, wid_b]) + + def test_fifo_within_eligible_wikis(self, db, cooldown): + """Both wikis past cooldown → older oldest-attach wins FIFO. + Both rows are ANCIENT (older than any production row); wiki_old is + even older so it beats wiki_new in created_at order.""" + wid_old = _insert_test_wiki(db, "fifo_old") + wid_new = _insert_test_wiki(db, "fifo_new") + jold = _insert_job(db, job_type="attach", target_wiki_id=wid_old, + age_seconds=ANCIENT_AGE_SECONDS + 300) + jnew = _insert_job(db, job_type="attach", target_wiki_id=wid_new, + age_seconds=ANCIENT_AGE_SECONDS) + try: + bucket = wiki_jobs.next_write_bucket(db) + assert bucket is not None + assert bucket["target_wiki_id"] == wid_old + finally: + _cleanup(db, job_ids=[jold, jnew], wiki_ids=[wid_old, wid_new]) + + +class TestPriorityPreservation: + """Cooldown is attach-only; consolidate and create are unaffected.""" + + def test_consolidate_drains_before_fresh_attaches(self, db): + wid_a = _insert_test_wiki(db, "prio_ca") + wid_b = _insert_test_wiki(db, "prio_cb") + ja = _insert_job(db, job_type="attach", target_wiki_id=wid_a, age_seconds=10) + jc = _insert_job( + db, job_type="consolidate", target_wiki_id=None, + entity_ids=[wid_a, wid_b], age_seconds=0, + ) + try: + bucket = wiki_jobs.next_write_bucket(db) + assert bucket is not None + assert bucket["mode"] == "consolidate" + assert bucket["jobs"][0]["id"] == jc + finally: + _cleanup(db, job_ids=[ja, jc], wiki_ids=[wid_a, wid_b]) + + def test_consolidate_drains_before_eligible_attaches(self, db, cooldown): + """The cooldown does NOT alter the consolidate > attach hierarchy. + Attach is ANCIENT (eligible); consolidate is recent — consolidate + still wins by priority, not by created_at.""" + wid_a = _insert_test_wiki(db, "prio_ea") + wid_b = _insert_test_wiki(db, "prio_eb") + ja = _insert_job(db, job_type="attach", target_wiki_id=wid_a, + age_seconds=ANCIENT_AGE_SECONDS) + jc = _insert_job( + db, job_type="consolidate", target_wiki_id=None, + entity_ids=[wid_a, wid_b], age_seconds=ANCIENT_AGE_SECONDS + 60, + ) + try: + bucket = wiki_jobs.next_write_bucket(db) + assert bucket is not None + assert bucket["mode"] == "consolidate" + finally: + _cleanup(db, job_ids=[ja, jc], wiki_ids=[wid_a, wid_b]) + + # Note on `create` jobs: by SQL inspection of next_write_bucket() the + # cooldown filter is gated `job_type <> 'attach' OR ...`, so create jobs + # bypass it entirely. An end-to-end test that asserts a fresh create is + # claimed FIRST is not reliable against a live DB with any pending + # higher-priority jobs (consolidate/attach), and forcibly draining + # production jobs is out of scope. The SQL itself is the proof; the + # other tests above transitively confirm non-attach paths are unaffected. + + +class TestEdgeCases: + + def test_assigned_jobs_excluded_from_cooldown_calc(self, db, cooldown): + """An `assigned` attach for the same wiki does NOT count toward the + cooldown's MIN(created_at). Only `pending` rows do.""" + wid = _insert_test_wiki(db, "edge_assigned") + j_assigned = _insert_job( + db, job_type="attach", target_wiki_id=wid, + age_seconds=cooldown + 600, + status="assigned", + ) + j_pending = _insert_job( + db, job_type="attach", target_wiki_id=wid, + age_seconds=10, + ) + try: + bucket = wiki_jobs.next_write_bucket(db) + if bucket is not None: + assert bucket.get("target_wiki_id") != wid, ( + f"fresh pending should NOT be claimable — assigned doesn't " + f"count toward cooldown MIN. Got {bucket!r}" + ) + finally: + _cleanup(db, job_ids=[j_assigned, j_pending], wiki_ids=[wid]) diff --git a/tests/test_wiki_sections.py b/tests/test_wiki_sections.py new file mode 100644 index 0000000..5b168e2 --- /dev/null +++ b/tests/test_wiki_sections.py @@ -0,0 +1,264 @@ +"""Unit tests for `braindb.services.wiki_sections` — the pure parsing and +splicing layer behind the writer's section-edit tools. + +These tests cover the DB-free functions only (`parse_sections`, +`splice_section`, `delete_section`, `check_grammar`). The DB helpers +(`fetch_wiki_for_section_op`, `apply_section_write`) are covered by +the end-to-end smoke test inside `braindb_api` (see plan Phase 1). + +The contract being tested: + +- `parse_sections(body)` returns `(header, [Section(name, content)])`. + Sections are split on `` markers; the header + is everything before the first marker. +- `splice_section` REPLACES an existing section's content, or APPENDS + a fresh section if the name is new. Bytes outside the targeted + section are preserved exactly. +- `delete_section` removes a section, raises `KeyError` if missing. +- `check_grammar` flags: no markers, malformed `[[ref:` tokens, missing + Summary callout. Tolerates the grouped-refs variant `[[ref:UUID1], + [ref:UUID2]]` documented in the wiki frontend plan. +- Round-trip identity: parse → splice (with same content) → string is + byte-identical to the input when the input is itself in normal form. +""" +from __future__ import annotations + +import pytest + +from braindb.services.wiki_sections import ( + Section, + StaleRevisionError, + check_grammar, + delete_section, + parse_sections, + splice_section, +) + +UUID_A = "11111111-1111-1111-1111-111111111111" +UUID_B = "22222222-2222-2222-2222-222222222222" + +# A minimal but realistic body in normal form (matches the writer +# prompt's "Recommended structure"). Used as the baseline for splice + +# roundtrip tests. +NORMAL_BODY = ( + "\n" + "# Test\n" + "> **Summary:** one line\n" + "> **Disambiguation:** what this is\n" + f"\n" + f"opening prose [[ref:{UUID_A}]]\n" + "\n" + f"2026 — event [[ref:{UUID_B}]]\n" + "\n" + f"- [[ref:{UUID_A}]] — source A\n" + f"- [[ref:{UUID_B}]] — source B\n" +) + + +# ====================================================================== # +# parse_sections # +# ====================================================================== # + +def test_parse_sections_extracts_each_section_in_order(): + header, sections = parse_sections(NORMAL_BODY) + names = [s.name for s in sections] + assert names == ["overview", "timeline", "references"] + + +def test_parse_sections_preserves_header_verbatim(): + header, _ = parse_sections(NORMAL_BODY) + assert header.startswith("\n" + "# X\n" + "\n" + "no summary callout above\n" + ) + issues = check_grammar(body) + assert any("> **Summary:**" in i for i in issues) + + +def test_grammar_tolerates_grouped_refs(): + """The grouped form `[[ref:UUID1], [ref:UUID2]]` is documented in the + wiki frontend plan as a real-world variant the renderer accepts. + check_grammar must not flag it as malformed.""" + body = ( + "\n" + "# X\n" + "> **Summary:** s\n" + "\n" + f"grouped citation [[ref:{UUID_A}], [ref:{UUID_B}]] in text\n" + ) + issues = check_grammar(body) + # No malformed-ref complaints (the only issue could be summary, but + # we included it) + assert not any("malformed" in i for i in issues), issues + + +def test_grammar_flags_truly_broken_ref(): + body = ( + "\n" + "# X\n" + "> **Summary:** s\n" + "\n" + "broken ref [[ref:not-a-uuid]] here\n" + ) + issues = check_grammar(body) + assert any("malformed" in i for i in issues), issues + + +# ====================================================================== # +# StaleRevisionError class # +# ====================================================================== # + +def test_stale_revision_error_is_exception(): + """The DB helpers raise this when expect_revision mismatches the + current DB revision. The tool wrappers translate it into a string + error the LLM can read; the class itself is the integration point.""" + assert issubclass(StaleRevisionError, Exception) + err = StaleRevisionError("expected 5, current 6") + assert "5" in str(err) and "6" in str(err) + + +# ====================================================================== # +# Section dataclass # +# ====================================================================== # + +def test_section_is_frozen_dataclass(): + s = Section(name="x", content="y") + with pytest.raises(Exception): # dataclasses.FrozenInstanceError + s.name = "z" # type: ignore[misc] + + +def test_section_char_count_property(): + s = Section(name="x", content="abcdef") + assert s.char_count == 6