From 395024f0bc5b5e990158be4752183d4259135c62 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 31 May 2026 16:11:57 +0100 Subject: [PATCH 01/10] chore(prompts): elevate /memory/context + /memory/tree, demote search_sql MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - recall_memory docstring: prepend "Primary recall tool" framing. - view_tree docstring: add explicit "when to use" line, named as the entity-driven default; positioned ahead of search_sql for any "what's around this entity" question. - search_sql docstring: tightened to one assertive line — aggregates only, never for recall/discovery/understanding/neighbourhood. - system_prompt.md TOOL PRIORITY block: promote view_tree from a buried bullet to a named slot (#2), making the entity-driven vs query-driven split explicit. SQL stays as #5 exception. - skills/braindb-agent/SKILL.md: prose paragraph rewritten as a numbered priority list matching the other skill's shape. - skills/braindb/SKILL.md: split tree out of the "structure lookups" bullet into its own #2 slot; SQL bullet explicitly forbids "around this entity" questions (those are tree's job). - BRAINDB_GUIDE.md: added a top-level ⚠ TOOL PRIORITY block so the reference doc isn't the weakest spot in the guidance chain. Net token impact on agent prompts: ~+20 tokens per call (one extra phrase in three tool docstrings). Skill markdowns and the guide are not loaded into prompts. No behaviour change in any code path; this is pure messaging. --- BRAINDB_GUIDE.md | 20 ++++++++++++ braindb/agent/prompts/system_prompt.md | 31 +++++++++++-------- braindb/agent/tools.py | 13 +++++--- skills/braindb-agent/SKILL.md | 36 ++++++++++++++-------- skills/braindb/SKILL.md | 42 ++++++++++++++++---------- 5 files changed, 97 insertions(+), 45 deletions(-) diff --git a/BRAINDB_GUIDE.md b/BRAINDB_GUIDE.md index 7a3e95c..11d0a65 100644 --- a/BRAINDB_GUIDE.md +++ b/BRAINDB_GUIDE.md @@ -5,6 +5,26 @@ The API runs at **http://localhost:8000**. Everything is done via HTTP calls. --- +## ⚠ TOOL PRIORITY (read this first) + +BrainDB's value is the graph + embeddings + ranking. Use that power; do not +fall back to flat SQL. + +1. **`POST /api/v1/memory/context`** — default for **query-driven** recall, + discovery, understanding ("what do we know about X?"). Keyword-mediated + fuzzy + embeddings + graph + ranking. +2. **`GET /api/v1/memory/tree/?max_depth=N`** — default for **entity-driven** + neighbourhood exploration ("what's around entity Y?"). Returns the chain + with relation types and edge scores in one call. +3. **`POST /api/v1/agent/query`** ("delegate to a subagent") — for multi-step + investigation / disambiguation. +4. `GET /api/v1/entities…` and `/entities//relations` — direct lookups. +5. **`POST /api/v1/memory/sql` ⚠ exception only** — aggregates (counts, + GROUP BY, activity-log joins). NEVER for recall, discovery, similarity, + understanding, or "what's around this entity" — those are the tools above. + +--- + ## Entity Types | Type | What to store | diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index 335bbd0..87437e3 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -50,24 +50,29 @@ CRITICAL — every assistant message MUST be a tool call; never plain prose. The 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 +1. **`recall_memory`** — default for ALL **query-driven** recall, discovery, + and understanding: multi-query fuzzy + full-text + **keyword-embedding** + + graph traversal + decay + ranking. Use first when you don't yet know which + specific entity you want — "what do we know about X." +2. **`view_tree(entity_id, max_depth=N)`** — default for **entity-driven** + neighbourhood exploration. When you already have an entity ID (from a + previous `recall_memory` result or known beforehand) and want to see what + surrounds it 1-N hops out, with relation types and edge scores. Beats + `search_sql` for any "what's around this entity" question — SQL can't show + the chain in one call. +3. **`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. +4. `view_entity_relations` / `get_entity` / `list_entities` — direct lookups + (single-hop relations, full body of one entity, listing by filter). +5. **`search_sql` ⚠ exception only — aggregates only** (counts, GROUP BY, + activity-log joins, schema inspection). A blunt SELECT has no embeddings, + no graph, no ranking — it throws away everything BrainDB is good at. + 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. +that's a `recall_memory` or `view_tree` or `delegate_to_subagent` job. ## READING CONTENT — previews vs the full body diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index 3628c12..87d177b 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -112,10 +112,10 @@ async def recall_memory( queries: list[str], max_results: int = settings.recall_default_max_results, ) -> str: - """Search BrainDB memory with multiple natural language queries. + """⭐ Primary recall tool — use FIRST for ANY "what do we know about X" question. + 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: @@ -634,7 +634,9 @@ async def delete_relation(relation_id: str) -> str: @function_tool @_verbose("view_tree") async def view_tree(entity_id: str, max_depth: int = 2) -> str: - """View the graph of connections around an entity (incoming + outgoing). + """⭐ Use when you already have an entity ID and want its 1-N hop neighbourhood + with relation types and edge scores. Preferred over search_sql for any + "what's around this entity" question — SQL can't show the chain in one call. Args: entity_id: UUID of the root entity. @@ -668,7 +670,10 @@ async def view_tree(entity_id: str, max_depth: int = 2) -> str: @function_tool @_verbose("search_sql") async def search_sql(query: str) -> str: - """Run a read-only SQL query (SELECT/WITH only). Use for complex exploration. + """⚠ Aggregates ONLY (counts, GROUP BY, joins for stats). NEVER for recall / + discovery / understanding — that's recall_memory. NEVER for "what's around + this entity" — that's view_tree. If you're using SQL to find or understand + something, stop and pick the right tool. Args: query: SQL query — must start with SELECT or WITH. diff --git a/skills/braindb-agent/SKILL.md b/skills/braindb-agent/SKILL.md index 81658fb..1f42658 100644 --- a/skills/braindb-agent/SKILL.md +++ b/skills/braindb-agent/SKILL.md @@ -26,18 +26,30 @@ BrainDB has its own internal agent (LiteLLM with pluggable provider via `LLM_PRO --- -## 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. +## TOOL PRIORITY (read this first) + +The agent has a clear order of tools it should reach for. When you phrase a +request, lean into the sophisticated tools — don't ask it to "run SQL" for +anything to do with recall or understanding. + +1. **Query-driven recall** — *"what do we know about X?"* → the agent calls + `/memory/context` (keyword-mediated fuzzy + embedding + graph + ranking, + with diversity quotas). The default for ALL discovery and understanding. +2. **Entity-driven neighbourhood** — *"what's connected to entity Y?"* (once + you have an ID from a previous recall) → the agent calls `/memory/tree/`. + Returns the multi-hop neighbourhood with relation types and edge scores in + one call. Use this instead of SQL for "around this entity" questions. +3. **Multi-step investigation** — *"investigate / disambiguate / resolve X"* + → the agent delegates to a subagent. Keeps the main context clean. +4. **Direct lookups** — `view_entity_relations`, `get_entity`, `list_entities` + for narrow questions. +5. **`search_sql` ⚠ exception only** — for explicit aggregates (counts, + GROUP BY, activity-log joins). Never for finding / understanding / + "what's related to" — those are jobs for the tools above. + +If you're tempted to phrase a request as *"run a SQL query that finds…"* for +*finding* or *understanding* something, stop — that's the recall or tree +path's job. Ask in plain English. **Wikis** are first-class memory entities curated by an internal maintainer + writer pipeline. The agent surfaces them through recall automatically when diff --git a/skills/braindb/SKILL.md b/skills/braindb/SKILL.md index 0f482e8..ecce359 100644 --- a/skills/braindb/SKILL.md +++ b/skills/braindb/SKILL.md @@ -91,28 +91,38 @@ 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 + **query-driven** recall, discovery, and understanding ("what do we know + about X?"). 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. **`GET /api/v1/memory/tree/?max_depth=N`** — the default for + **entity-driven** neighbourhood exploration ("what's around entity Y?" + when you already have an ID from a previous recall). Returns the multi-hop + neighbourhood in ONE call, with relation types and edge scores. Beats + `/memory/sql` for any "what's connected to X" question — SQL can't show + the chain in a single call. +3. **`POST /api/v1/agent/query` with "delegate to a subagent…"** — for + multi-step investigation/disambiguation; the agent researches and returns + a summary. +4. `GET /api/v1/entities…`, `GET /api/v1/entities//relations` — direct + lookups (list-by-filter, single-hop relations). +5. **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. +6. **`POST /api/v1/memory/sql` ⚠ exception only — aggregates 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. **Never** for "what's around this entity" — that's + `/memory/tree`. If you're about to use `/memory/sql` to *find* or *understand* something, -stop — that's a `/memory/context` (or delegated `/agent/query`) job. +stop — that's a `/memory/context` or `/memory/tree` (or delegated +`/agent/query`) job. ### Previews vs full body From c91c3be3d35a09a41d91caa1cbee0cabdf20d986 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 31 May 2026 16:30:46 +0100 Subject: [PATCH 02/10] fix(scoring): propagate seed similarity through graph; soften depth multiplier MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related fixes to the recall ranking, both confirmed against live data where every depth-1 graph result was pinning at the same literal 0.27 regardless of which seed it came from. 1. Propagate seed similarity through graph hops. The graph CTE now carries a `seed_origin_id` column from each seed down through every recursive row. In context.py, the score component for a graph-discovered (non-seed) entity is now the score of its origin seed, inherited via that column, instead of the literal `0.3` fallback that made every graph-only entity rank identically. Before this fix a perfect-match seed (sim=1.0) and a weak-match seed (sim=0.3) produced the same depth-1 neighbour rank. Worse: a weak seed (sim<0.27) was outranked by its own neighbours because the fallback floor was higher than the seed's real score. Both gone now. 2. Soften the depth multiplier. The hardcoded depth step in the recursive CTE goes from [1.0 / 0.6 / 0.3] to [1.0 / 0.8 / 0.6]. Deeper hops still decay but no longer collapse — depth-2 and depth-3 items can now reach final_rank values that exceed the min_relevance threshold and surface in results, instead of vanishing as they do today. Net effect (for a seed with similarity 1.0): depth 0: 1.00 -> 1.00 (unchanged) depth 1: 0.27 -> 0.80 depth 2: 0.12 -> 0.51 depth 3: 0.05 -> 0.31 For a seed with similarity 0.5: depth 0: 0.50 -> 0.50 (unchanged) depth 1: 0.27 -> 0.40 (now correctly lower than seed) depth 2: 0.12 -> 0.26 depth 3: 0.05 -> 0.15 No new tables, no migration, no config flags, no module reorganisation. Just two surgical edits: one extra column in the CTE, one Python lookup swap, two constants nudged up. --- braindb/services/context.py | 9 ++++++++- braindb/services/graph.py | 18 ++++++++++++------ 2 files changed, 20 insertions(+), 7 deletions(-) diff --git a/braindb/services/context.py b/braindb/services/context.py index 383ef88..f9abd0f 100644 --- a/braindb/services/context.py +++ b/braindb/services/context.py @@ -397,7 +397,14 @@ def assemble_context(conn, req: ContextRequest) -> ContextResponse: items = [] for row in all_rows: eid = row["id"] - score = seed_scores.get(eid, 0.3) + # Score = the entity's own similarity if it was a seed; otherwise inherit + # the score of the seed it descended from (carried by `seed_origin_id` + # through the graph CTE). This propagates the real similarity signal + # through depth-1+ hops instead of resetting to a literal fallback. + score = seed_scores.get(eid) + if score is None: + origin = row.get("seed_origin_id") + score = seed_scores.get(str(origin), 1.0) if origin else 1.0 depth = row.get("min_depth", 0) relevance = row.get("relevance", 1.0) items.append(_to_item(row, score, depth, relevance, ext_map.get(eid, {}))) diff --git a/braindb/services/graph.py b/braindb/services/graph.py index ecc4346..63dc591 100644 --- a/braindb/services/graph.py +++ b/braindb/services/graph.py @@ -16,7 +16,11 @@ NULL::UUID AS via_relation_id, NULL::TEXT AS via_relation_type, NULL::TEXT AS via_description, - NULL::TEXT AS via_notes + NULL::TEXT AS via_notes, + -- The seed each row descended from. For seeds, it's themselves; + -- for graph-discovered rows, it propagates through the recursion + -- so context.py can inherit the seed's similarity score. + e.id AS seed_origin_id FROM entities e WHERE e.id = ANY(%s::uuid[]) @@ -32,15 +36,16 @@ * r.relevance_score * CASE t.depth + 1 WHEN 1 THEN 1.0 - WHEN 2 THEN 0.6 - ELSE 0.3 + WHEN 2 THEN 0.8 + ELSE 0.6 END )::FLOAT, t.visited || target.id, r.id, r.relation_type, r.description, - r.notes + r.notes, + t.seed_origin_id FROM traversal t JOIN relations r ON ( r.from_entity_id = t.id @@ -56,7 +61,7 @@ AND ( t.accumulated_relevance * r.relevance_score - * CASE t.depth + 1 WHEN 1 THEN 1.0 WHEN 2 THEN 0.6 ELSE 0.3 END + * CASE t.depth + 1 WHEN 1 THEN 1.0 WHEN 2 THEN 0.8 ELSE 0.6 END ) > %s ) SELECT DISTINCT ON (id) @@ -65,7 +70,8 @@ created_at, updated_at, accessed_at, access_count, metadata, depth AS min_depth, accumulated_relevance AS relevance, - via_relation_id, via_relation_type, via_description, via_notes + via_relation_id, via_relation_type, via_description, via_notes, + seed_origin_id FROM traversal ORDER BY id, depth, accumulated_relevance DESC """ From 8fa45316668b33a867fc7f85878a8264c4c4b07b Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 31 May 2026 18:09:37 +0100 Subject: [PATCH 03/10] chore(prompts): describe view_tree as a capability, not a rigid trigger MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-1 elevated view_tree under a category label ("default for entity- driven neighbourhood exploration"). Benchmarking showed both Claude (via the curl skill) and the in-house Qwen agent picked tree zero times across 5 questions x 2 paths = 10 runs. The category framing was too abstract. This round describes WHAT tree does (capability) with a SUGGESTIVE "when" hint, not a rigid trigger. The agent keeps full judgment about whether to use it — we just make the value clearer: reveals an entity's connections in one call: relations + 1-N hop neighbours + edge scores. Especially useful when you have an entity ID (from a previous result) and want its graph context. No "INSTEAD OF" commands. No decision-rule blocks. No examples. The shape mirrors what already worked for the search_sql demotion (capability + bounded use, agent decides). Net token impact on the agent's system prompt: ~-20 to -50 tokens (this is a shrink, not a bloat). Same edit applied across: - braindb/agent/tools.py (view_tree docstring — in agent prompt) - braindb/agent/prompts/system_prompt.md (TOOL PRIORITY block entry #2) - skills/braindb/SKILL.md (user-facing skill) - skills/braindb-agent/SKILL.md (user-facing skill) - BRAINDB_GUIDE.md (reference guide top block) No code changes. No behavioural change in any code path. Pure messaging. Verified separately on the same benchmark question set next. --- BRAINDB_GUIDE.md | 6 +++--- braindb/agent/prompts/system_prompt.md | 10 ++++------ braindb/agent/tools.py | 7 ++++--- skills/braindb-agent/SKILL.md | 8 ++++---- skills/braindb/SKILL.md | 11 +++++------ 5 files changed, 20 insertions(+), 22 deletions(-) diff --git a/BRAINDB_GUIDE.md b/BRAINDB_GUIDE.md index 11d0a65..0bfcbc5 100644 --- a/BRAINDB_GUIDE.md +++ b/BRAINDB_GUIDE.md @@ -13,9 +13,9 @@ fall back to flat SQL. 1. **`POST /api/v1/memory/context`** — default for **query-driven** recall, discovery, understanding ("what do we know about X?"). Keyword-mediated fuzzy + embeddings + graph + ranking. -2. **`GET /api/v1/memory/tree/?max_depth=N`** — default for **entity-driven** - neighbourhood exploration ("what's around entity Y?"). Returns the chain - with relation types and edge scores in one call. +2. **`GET /api/v1/memory/tree/?max_depth=N`** — reveals an entity's + connections in one call (relations + 1-N hop neighbours + edge scores). + Especially useful when you have an entity ID and want its graph context. 3. **`POST /api/v1/agent/query`** ("delegate to a subagent") — for multi-step investigation / disambiguation. 4. `GET /api/v1/entities…` and `/entities//relations` — direct lookups. diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index 87437e3..bef2d8c 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -54,12 +54,10 @@ fall back to flat SQL. and understanding: multi-query fuzzy + full-text + **keyword-embedding** + graph traversal + decay + ranking. Use first when you don't yet know which specific entity you want — "what do we know about X." -2. **`view_tree(entity_id, max_depth=N)`** — default for **entity-driven** - neighbourhood exploration. When you already have an entity ID (from a - previous `recall_memory` result or known beforehand) and want to see what - surrounds it 1-N hops out, with relation types and edge scores. Beats - `search_sql` for any "what's around this entity" question — SQL can't show - the chain in one call. +2. **`view_tree(, max_depth=N)`** — reveals an entity's connections in + one call: relations + 1-N hop neighbours + edge scores. Especially useful + when you have an entity ID and want its graph context — often sharper + than another `recall_memory` about the same entity. 3. **`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 diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index 87d177b..d289823 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -634,9 +634,10 @@ async def delete_relation(relation_id: str) -> str: @function_tool @_verbose("view_tree") async def view_tree(entity_id: str, max_depth: int = 2) -> str: - """⭐ Use when you already have an entity ID and want its 1-N hop neighbourhood - with relation types and edge scores. Preferred over search_sql for any - "what's around this entity" question — SQL can't show the chain in one call. + """⭐ Reveals an entity's connections in one call: relations + 1-N hop + neighbours + edge scores. Especially useful when you have an entity ID + (from a previous result) and want its graph context — often a sharper + choice than another `recall_memory` about the same entity. Args: entity_id: UUID of the root entity. diff --git a/skills/braindb-agent/SKILL.md b/skills/braindb-agent/SKILL.md index 1f42658..0df1066 100644 --- a/skills/braindb-agent/SKILL.md +++ b/skills/braindb-agent/SKILL.md @@ -35,10 +35,10 @@ anything to do with recall or understanding. 1. **Query-driven recall** — *"what do we know about X?"* → the agent calls `/memory/context` (keyword-mediated fuzzy + embedding + graph + ranking, with diversity quotas). The default for ALL discovery and understanding. -2. **Entity-driven neighbourhood** — *"what's connected to entity Y?"* (once - you have an ID from a previous recall) → the agent calls `/memory/tree/`. - Returns the multi-hop neighbourhood with relation types and edge scores in - one call. Use this instead of SQL for "around this entity" questions. +2. **Entity-driven neighbourhood** — the agent's `/memory/tree/` reveals + an entity's connections in one call (relations + 1-N hop neighbours + edge + scores). Especially useful when an entity ID is already in hand — often + sharper than another query about the same entity. 3. **Multi-step investigation** — *"investigate / disambiguate / resolve X"* → the agent delegates to a subagent. Keeps the main context clean. 4. **Direct lookups** — `view_entity_relations`, `get_entity`, `list_entities` diff --git a/skills/braindb/SKILL.md b/skills/braindb/SKILL.md index ecce359..c89687c 100644 --- a/skills/braindb/SKILL.md +++ b/skills/braindb/SKILL.md @@ -97,12 +97,11 @@ to flat SQL. `tagged_with`). A two-level diversity quota (per-search-term + per-keyword halving) keeps results balanced. Then graph traversal + decay + ranking. -2. **`GET /api/v1/memory/tree/?max_depth=N`** — the default for - **entity-driven** neighbourhood exploration ("what's around entity Y?" - when you already have an ID from a previous recall). Returns the multi-hop - neighbourhood in ONE call, with relation types and edge scores. Beats - `/memory/sql` for any "what's connected to X" question — SQL can't show - the chain in a single call. +2. **`GET /api/v1/memory/tree/?max_depth=N`** — reveals an entity's + connections in one call: relations + 1-N hop neighbours + edge scores. + Especially useful when you have an entity ID (from a previous recall) + and want its graph context — often a sharper choice than another + `/memory/context` call about the same entity. 3. **`POST /api/v1/agent/query` with "delegate to a subagent…"** — for multi-step investigation/disambiguation; the agent researches and returns a summary. From bccf2b4ff033750521b896b48893896ac303efb1 Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 31 May 2026 18:35:27 +0100 Subject: [PATCH 04/10] fix(tools): view_tree honours max_depth, groups by depth, wiki-name labels Three minimal fixes to the agent's view_tree tool. Round-2a benchmarks showed Claude started using tree (0/5 -> 3/5) but the Qwen agent didn't, and even on Path A one of three tree calls (q4) didn't pay off. Looking at the actual implementation revealed why: tree was advertised with a max_depth argument but ignored it (single-hop SQL), so an agent asking view_tree(id, max_depth=2) only got depth-1 connections. Fixes: 1. max_depth respected. Single-hop SQL replaced with a recursive CTE that walks bidirectionally (as the single-hop already did via the OR clause) and stops at the requested depth. A cycle-visited array prevents loops. 2. Depth grouping in output. "DEPTH N (count):" headers between sections. Within a depth, rows sorted by edge_score desc. Same line shape as before; only headers are new. 3. Wiki labels use canonical_name. The wiki:meta comment header was being truncated as if it were content body. Extract canonical_name via a small regex; everything else keeps the existing 80-char content truncation. No graph.py change. No system_prompt change (already committed in round-2a). No schema change. ~35 lines net in tools.py. Verified separately on the same 5 benchmark questions next. --- braindb/agent/tools.py | 71 +++++++++++++++++++++++++++++++++++------- 1 file changed, 59 insertions(+), 12 deletions(-) diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index d289823..00ab78b 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -643,31 +643,78 @@ async def view_tree(entity_id: str, max_depth: int = 2) -> str: entity_id: UUID of the root entity. max_depth: How far to traverse (1-3, default 2). """ + if max_depth < 1 or max_depth > 3: + return _err("max_depth must be 1-3") try: with get_conn() as conn: with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: cur.execute( - """SELECT e.*, r.relation_type, r.relevance_score, r.description AS rel_desc, - CASE WHEN r.from_entity_id = %s THEN 'out' ELSE 'in' END AS dir - FROM relations r - JOIN entities e ON e.id = CASE WHEN r.from_entity_id = %s THEN r.to_entity_id ELSE r.from_entity_id END - WHERE r.from_entity_id = %s OR r.to_entity_id = %s""", - (entity_id, entity_id, entity_id, entity_id), + "SELECT id, entity_type, content FROM entities WHERE id = %s", + (entity_id,), + ) + root = cur.fetchone() + if not root: + return _err(f"Entity not found: {entity_id}") + # Recursive bidirectional walk, respects max_depth. + cur.execute( + """WITH RECURSIVE traversal AS ( + SELECT e.id, e.entity_type, e.content, + 0 AS depth, ARRAY[e.id] AS visited, + NULL::TEXT AS via_rel, + NULL::FLOAT AS edge_score, + NULL::TEXT AS direction + FROM entities e WHERE e.id = %s + UNION ALL + SELECT target.id, target.entity_type, target.content, + t.depth + 1, t.visited || target.id, + r.relation_type, r.relevance_score, + CASE WHEN r.from_entity_id = t.id THEN 'out' ELSE 'in' END + FROM traversal t + JOIN relations r ON r.from_entity_id = t.id OR r.to_entity_id = t.id + JOIN entities target ON target.id = CASE + WHEN r.from_entity_id = t.id THEN r.to_entity_id + ELSE r.from_entity_id END + WHERE t.depth < %s AND NOT (target.id = ANY(t.visited)) + ) + SELECT DISTINCT ON (id) id, entity_type, content, depth, + via_rel, edge_score, direction + FROM traversal WHERE depth > 0 + ORDER BY id, depth, edge_score DESC NULLS LAST + """, + (entity_id, max_depth), ) rows = [dict(r) for r in cur.fetchall()] + out = [f"ROOT [{root['entity_type']}] {_tree_label(root)} (id: {root['id']})"] if not rows: - return "No connections." - lines = [f"{len(rows)} connections from {entity_id}:"] + out.append("\n(no connections)") + return _truncate("\n".join(out)) + rows.sort(key=lambda x: (x['depth'], -(x['edge_score'] or 0))) + last_depth = None for r in rows: - lines.append( - f" [{r['dir']}] {r['relation_type']} (rel={r['relevance_score']})\n" - f" [{r['entity_type']}] {r['content'][:80]} (id: {r['id']})" + if r['depth'] != last_depth: + same_depth_n = sum(1 for x in rows if x['depth'] == r['depth']) + out.append(f"\nDEPTH {r['depth']} ({same_depth_n}):") + last_depth = r['depth'] + out.append( + f" [{r['direction']}] {r['via_rel']} (rel={r['edge_score']})\n" + f" [{r['entity_type']}] {_tree_label(r)} (id: {r['id']})" ) - return _truncate("\n".join(lines)) + return _truncate("\n".join(out)) except Exception as e: return _err(str(e)) +def _tree_label(entity: dict) -> str: + """Short label for tree output. Wikis use canonical_name; others truncate content.""" + import re as _r + content = (entity.get('content') or '').replace('\n', ' ') + if entity.get('entity_type') == 'wiki': + m = _r.search(r'canonical_name=([^\s]+)', content) + if m: + return m.group(1) + return content[:80] + + @function_tool @_verbose("search_sql") async def search_sql(query: str) -> str: From 7d90fa2c089864d1adc52f8521557188b0fc64ef Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 31 May 2026 18:36:14 +0100 Subject: [PATCH 05/10] chore(prompts): reframe view_tree as the "explore around this entity" tool Round-2a benchmarks showed the Qwen agent (Path B) still picked tree 0/5 times despite the new wording. Hypothesis: "graph context" reads as a niche specialist feature to a smaller model, so it falls back to recall. The reframing makes tree sound like the GENERAL tool for the thing the agent actually wants to do once it has an entity in hand: explore around it. "Explore around this entity" is the everyday framing; "graph context" is the jargon framing. Same line count, same shape, no bloat. Just a verb change in the system_prompt.md TOOL PRIORITY block entry #2. Verified separately in the next benchmark. --- braindb/agent/prompts/system_prompt.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index bef2d8c..37adb61 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -54,10 +54,11 @@ fall back to flat SQL. and understanding: multi-query fuzzy + full-text + **keyword-embedding** + graph traversal + decay + ranking. Use first when you don't yet know which specific entity you want — "what do we know about X." -2. **`view_tree(, max_depth=N)`** — reveals an entity's connections in - one call: relations + 1-N hop neighbours + edge scores. Especially useful - when you have an entity ID and want its graph context — often sharper - than another `recall_memory` about the same entity. +2. **`view_tree(, max_depth=N)`** — the efficient "explore around this + entity" tool. One call returns the entity's neighbours grouped by + relation type, plus their edge scores, 1-N hops out. When you have an + entity ID and want to know what's related to it, this is usually a + better next step than another `recall_memory` about the same entity. 3. **`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 From e6cf29b6dcd72a20c294fddba16768b09b3506de Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Sun, 31 May 2026 19:34:05 +0100 Subject: [PATCH 06/10] refactor: unify tree implementation (single shared service, no behaviour drift) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round-2c benchmarks revealed that "tree" was served by two divergent implementations: the HTTP endpoint (routers/memory.py::entity_tree) did a single-hop SQL that silently ignored max_depth, while the agent's view_tree tool ran a proper recursive CTE. Same name, same input, different behaviour. Path A (HTTP) and Path B (agent tool) were not looking at the same data. This commit extracts one source of truth: - braindb/services/tree.py NEW: build_entity_tree(conn, entity_id, max_depth) recursive CTE walks bidirectionally and respects max_depth. Returns {"root": {...}, "connections": [...]} with the same shape the HTTP endpoint always advertised — the frontend Graph tab keeps reading it unchanged. - routers/memory.py::entity_tree shrinks from ~60 lines to 8: just calls build_entity_tree. - agent/tools.py::view_tree shrinks: drops its own recursive CTE (added in bccf2b4), calls build_entity_tree, keeps only the text rendering (depth headers, [out]/[in] arrows, _tree_label for wikis). Behavioural effects: - HTTP /memory/tree/?max_depth=N now actually walks N hops. Quick spot check: tree on the "value-investing" keyword used to return ~20 depth-1 connections; now returns 156 connections (20 d1 + 136 d2). - Frontend Graph tab: same field names, same direction values ("outgoing"/"incoming"), more nodes visible at depth 2. No JS change needed. - Agent view_tree tool: returns the same text shape we shipped in bccf2b4; underlying data now comes from the shared service. Tests: tests/test_search.py — all 6 tests pass (shape-agnostic check on /memory/tree was already there; refactor preserves the shape). Net diff: +60 / -90 across 3 files. Code SHRINK. --- braindb/agent/tools.py | 64 ++++++--------------- braindb/routers/memory.py | 70 ++++------------------- braindb/services/tree.py | 115 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 143 insertions(+), 106 deletions(-) create mode 100644 braindb/services/tree.py diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index 00ab78b..665213b 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -646,58 +646,30 @@ async def view_tree(entity_id: str, max_depth: int = 2) -> str: if max_depth < 1 or max_depth > 3: return _err("max_depth must be 1-3") try: + from braindb.services.tree import build_entity_tree with get_conn() as conn: - with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: - cur.execute( - "SELECT id, entity_type, content FROM entities WHERE id = %s", - (entity_id,), - ) - root = cur.fetchone() - if not root: - return _err(f"Entity not found: {entity_id}") - # Recursive bidirectional walk, respects max_depth. - cur.execute( - """WITH RECURSIVE traversal AS ( - SELECT e.id, e.entity_type, e.content, - 0 AS depth, ARRAY[e.id] AS visited, - NULL::TEXT AS via_rel, - NULL::FLOAT AS edge_score, - NULL::TEXT AS direction - FROM entities e WHERE e.id = %s - UNION ALL - SELECT target.id, target.entity_type, target.content, - t.depth + 1, t.visited || target.id, - r.relation_type, r.relevance_score, - CASE WHEN r.from_entity_id = t.id THEN 'out' ELSE 'in' END - FROM traversal t - JOIN relations r ON r.from_entity_id = t.id OR r.to_entity_id = t.id - JOIN entities target ON target.id = CASE - WHEN r.from_entity_id = t.id THEN r.to_entity_id - ELSE r.from_entity_id END - WHERE t.depth < %s AND NOT (target.id = ANY(t.visited)) - ) - SELECT DISTINCT ON (id) id, entity_type, content, depth, - via_rel, edge_score, direction - FROM traversal WHERE depth > 0 - ORDER BY id, depth, edge_score DESC NULLS LAST - """, - (entity_id, max_depth), - ) - rows = [dict(r) for r in cur.fetchall()] + tree = build_entity_tree(conn, entity_id, max_depth) + if tree is None: + return _err(f"Entity not found: {entity_id}") + root = tree["root"] + conns = tree["connections"] out = [f"ROOT [{root['entity_type']}] {_tree_label(root)} (id: {root['id']})"] - if not rows: + if not conns: out.append("\n(no connections)") return _truncate("\n".join(out)) - rows.sort(key=lambda x: (x['depth'], -(x['edge_score'] or 0))) + # Group output by depth last_depth = None - for r in rows: - if r['depth'] != last_depth: - same_depth_n = sum(1 for x in rows if x['depth'] == r['depth']) - out.append(f"\nDEPTH {r['depth']} ({same_depth_n}):") - last_depth = r['depth'] + for c in conns: + d = c["depth"] + if d != last_depth: + same_n = sum(1 for x in conns if x["depth"] == d) + out.append(f"\nDEPTH {d} ({same_n}):") + last_depth = d + ent = c["entity"] + dir_short = "out" if c.get("direction") == "outgoing" else "in" out.append( - f" [{r['direction']}] {r['via_rel']} (rel={r['edge_score']})\n" - f" [{r['entity_type']}] {_tree_label(r)} (id: {r['id']})" + f" [{dir_short}] {c['via_relation_type']} (rel={c['relevance']})\n" + f" [{ent['entity_type']}] {_tree_label(ent)} (id: {ent['id']})" ) return _truncate("\n".join(out)) except Exception as e: diff --git a/braindb/routers/memory.py b/braindb/routers/memory.py index c680d39..57b11d6 100644 --- a/braindb/routers/memory.py +++ b/braindb/routers/memory.py @@ -80,67 +80,17 @@ def get_rules(): @router.get("/tree/{entity_id}") def entity_tree(entity_id: UUID, max_depth: int = Query(default=2, ge=1, le=3)): - """Return an entity and its graph connections organized by depth.""" - with get_conn() as conn: - # Fetch root entity - with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: - cur.execute("SELECT * FROM entities WHERE id = %s", (str(entity_id),)) - root_row = cur.fetchone() - if not root_row: - raise HTTPException(404, "Entity not found") - root_row = dict(root_row) - - # Fetch root extension fields - root_ext = fetch_ext(conn, [root_row]) - root_data = { - "id": root_row["id"], "entity_type": root_row["entity_type"], - "title": root_row.get("title"), "content": root_row["content"], - "keywords": root_row.get("keywords") or [], - "importance": root_row["importance"], - "notes": root_row.get("notes"), - "ext": root_ext.get(root_row["id"], {}), - } + """Return an entity and its graph connections organized by depth. - # Find all connections via relations (both directions) - with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: - cur.execute(""" - SELECT e.*, r.relation_type, r.relevance_score, r.description AS rel_description, - CASE WHEN r.from_entity_id = %s THEN 'outgoing' ELSE 'incoming' END AS direction - FROM relations r - JOIN entities e ON e.id = CASE - WHEN r.from_entity_id = %s THEN r.to_entity_id - ELSE r.from_entity_id - END - WHERE r.from_entity_id = %s OR r.to_entity_id = %s - """, (str(entity_id), str(entity_id), str(entity_id), str(entity_id))) - direct_rows = [dict(r) for r in cur.fetchall()] - - conn_ext = fetch_ext(conn, direct_rows) if direct_rows else {} - - connections = [] - seen_ids = {entity_id} - for row in direct_rows: - eid = row["id"] - if eid in seen_ids: - continue - seen_ids.add(eid) - connections.append({ - "entity": { - "id": eid, "entity_type": row["entity_type"], - "title": row.get("title"), "content": row["content"], - "keywords": row.get("keywords") or [], - "importance": row["importance"], - "ext": conn_ext.get(eid, {}), - }, - "depth": 1, - "relevance": row.get("relevance_score", 1.0), - "via_relation_type": row.get("relation_type"), - "via_description": row.get("rel_description"), - "direction": row.get("direction"), - }) - - connections.sort(key=lambda c: (-c["relevance"], c["depth"])) - return {"root": root_data, "connections": connections} + Uses the shared `build_entity_tree` service (also used by the agent's + `view_tree` tool) so HTTP callers and the agent see the same data. + """ + from braindb.services.tree import build_entity_tree + with get_conn() as conn: + tree = build_entity_tree(conn, str(entity_id), max_depth) + if tree is None: + raise HTTPException(404, "Entity not found") + return tree @router.get("/log") diff --git a/braindb/services/tree.py b/braindb/services/tree.py new file mode 100644 index 0000000..c0467be --- /dev/null +++ b/braindb/services/tree.py @@ -0,0 +1,115 @@ +"""Shared tree-build service. + +Single source of truth for "walk the relation graph outward from an entity". +Used by both the HTTP endpoint (`/api/v1/memory/tree/`) and the agent's +`view_tree` tool. Walks bidirectionally and respects `max_depth`. +""" +from __future__ import annotations + +import psycopg2.extras + +from braindb.services.context import fetch_ext + + +_TREE_SQL = """ +WITH RECURSIVE traversal AS ( + SELECT e.id, e.entity_type, e.title, e.content, e.keywords, + e.importance, e.notes, + 0 AS depth, + ARRAY[e.id] AS visited, + NULL::TEXT AS via_relation_type, + NULL::TEXT AS via_description, + NULL::FLOAT AS relevance_score, + NULL::TEXT AS direction + FROM entities e + WHERE e.id = %s + + UNION ALL + + SELECT target.id, target.entity_type, target.title, target.content, + target.keywords, target.importance, target.notes, + t.depth + 1, + t.visited || target.id, + r.relation_type, + r.description, + r.relevance_score, + CASE WHEN r.from_entity_id = t.id THEN 'outgoing' ELSE 'incoming' END + FROM traversal t + JOIN relations r ON r.from_entity_id = t.id OR r.to_entity_id = t.id + JOIN entities target ON target.id = CASE + WHEN r.from_entity_id = t.id THEN r.to_entity_id + ELSE r.from_entity_id + END + WHERE t.depth < %s + AND NOT (target.id = ANY(t.visited)) +) +SELECT DISTINCT ON (id) + id, entity_type, title, content, keywords, importance, notes, + depth, via_relation_type, via_description, relevance_score, direction +FROM traversal +WHERE depth > 0 +ORDER BY id, depth, relevance_score DESC NULLS LAST +""" + + +def build_entity_tree(conn, entity_id: str, max_depth: int = 2) -> dict | None: + """Walk the relation graph bidirectionally from `entity_id` up to + `max_depth` hops. Returns: + + {"root": , "connections": [, ...]} + + or ``None`` if the root entity is not found. + + Each connection dict has keys: + entity, depth, relevance, via_relation_type, via_description, direction + where `direction` is "outgoing" or "incoming" relative to the root path. + """ + eid = str(entity_id) + with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: + cur.execute("SELECT * FROM entities WHERE id = %s", (eid,)) + root_row = cur.fetchone() + if not root_row: + return None + root_row = dict(root_row) + + cur.execute(_TREE_SQL, (eid, max_depth)) + rows = [dict(r) for r in cur.fetchall()] + + # Extension fields for root + all connection entities (single batched call) + ext_map = fetch_ext(conn, [root_row] + rows) + + root_data = { + "id": root_row["id"], + "entity_type": root_row["entity_type"], + "title": root_row.get("title"), + "content": root_row["content"], + "keywords": root_row.get("keywords") or [], + "importance": root_row["importance"], + "notes": root_row.get("notes"), + "ext": ext_map.get(root_row["id"], {}), + } + + connections = [] + for row in rows: + rid = row["id"] + connections.append({ + "entity": { + "id": rid, + "entity_type": row["entity_type"], + "title": row.get("title"), + "content": row["content"], + "keywords": row.get("keywords") or [], + "importance": row["importance"], + "ext": ext_map.get(rid, {}), + }, + "depth": row["depth"], + "relevance": row.get("relevance_score", 1.0) if row.get("relevance_score") is not None else 1.0, + "via_relation_type": row.get("via_relation_type"), + "via_description": row.get("via_description"), + "direction": row.get("direction"), + }) + + # Sort by depth asc, then relevance desc within depth + connections.sort(key=lambda c: (c["depth"], -c["relevance"])) + + return {"root": root_data, "connections": connections} From 8bccabf1fb9e28b16e2c33525c255f7edffe159b Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Mon, 1 Jun 2026 19:17:20 +0100 Subject: [PATCH 07/10] fix(scoring): per-edge LLM judgment in create_relation + graph CTE uses both scores * tools.py: create_relation gains importance_score parameter; INSERT writes both relevance_score and importance_score (column was NULL for all agent-created rows since day one). * ingest_watcher.py: stripped dictated certainty/importance/relevance_score literals from chunk-extraction and central_review prompts -- LLM judges per the tool docstring. AGENT_TIMEOUT now env-overridable, default 1200. * services/graph.py: per-hop multiplier now r.relevance_score * COALESCE(r.importance_score, 0.5) * depth_penalty; is_bidirectional dropped from the JOIN (always walks both directions, matches tree.py). * system_prompt.md: importance_score added to create_relation param list. * 4 new tests lock the behaviour: persistence, watcher-no-dictation, importance_score moves rank, unidirectional edges walk backwards. All 142 tests pass. Path A bench 5/5 PASS in 14s (zero view_tree). Path B bench 5/5 PASS in 1090s at 1200s timeout. Variance verified on live ingest of the AI Dark Output article. --- braindb/agent/prompts/system_prompt.md | 2 +- braindb/agent/tools.py | 12 ++-- braindb/ingest_watcher.py | 17 +++--- braindb/services/graph.py | 7 +-- tests/test_create_relation_persistence.py | 57 +++++++++++++++++++ tests/test_graph_scoring_uses_importance.py | 63 +++++++++++++++++++++ tests/test_graph_walks_both_directions.py | 57 +++++++++++++++++++ tests/test_ingest_watcher_no_dictation.py | 55 ++++++++++++++++++ 8 files changed, 253 insertions(+), 17 deletions(-) create mode 100644 tests/test_create_relation_persistence.py create mode 100644 tests/test_graph_scoring_uses_importance.py create mode 100644 tests/test_graph_walks_both_directions.py create mode 100644 tests/test_ingest_watcher_no_dictation.py diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index 37adb61..c52fb02 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -26,7 +26,7 @@ CRITICAL — every assistant message MUST be a tool call; never plain prose. The - `delete_entity(entity_id)` **Relations:** -- `create_relation(from_entity_id, to_entity_id, relation_type, relevance_score, description)` +- `create_relation(from_entity_id, to_entity_id, relation_type, relevance_score, importance_score, description)` - `view_entity_relations(entity_id)` - `delete_relation(relation_id)` diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index 665213b..12a5dee 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -542,7 +542,8 @@ async def create_relation( from_entity_id: str, to_entity_id: str, relation_type: str, - relevance_score: float = 0.7, + relevance_score: float = 0.5, + importance_score: float = 0.5, description: Optional[str] = None, ) -> str: """Create a relation between two entities. @@ -551,7 +552,8 @@ async def create_relation( from_entity_id: Source entity UUID. to_entity_id: Target entity UUID. relation_type: One of: supports, contradicts, elaborates, refers_to, derived_from, similar_to, is_example_of, challenges, tagged_with. - relevance_score: 0-1 (default 0.7). + relevance_score: 0-1 — how tight the semantic link is (0.9 = strong, 0.5 = neutral / "didn't judge", 0.2 = weak). + importance_score: 0-1 — how much losing this edge would degrade recall (0.9 = critical, 0.5 = neutral / "didn't judge", 0.2 = trivial). description: Why this relation exists. """ try: @@ -559,9 +561,9 @@ async def create_relation( with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur: try: cur.execute( - """INSERT INTO relations (from_entity_id, to_entity_id, relation_type, relevance_score, description) - VALUES (%s, %s, %s, %s, %s) RETURNING id""", - (from_entity_id, to_entity_id, relation_type, relevance_score, description), + """INSERT INTO relations (from_entity_id, to_entity_id, relation_type, relevance_score, importance_score, description) + VALUES (%s, %s, %s, %s, %s, %s) RETURNING id""", + (from_entity_id, to_entity_id, relation_type, relevance_score, importance_score, description), ) rid = cur.fetchone()["id"] except psycopg2.errors.UniqueViolation: diff --git a/braindb/ingest_watcher.py b/braindb/ingest_watcher.py index 4a159e1..cb3fea2 100644 --- a/braindb/ingest_watcher.py +++ b/braindb/ingest_watcher.py @@ -46,7 +46,7 @@ CHUNK_OVERLAP = 75 # words of overlap between adjacent chunks — catches facts that span a boundary INGEST_TIMEOUT = 60 -AGENT_TIMEOUT = 600 # NIM free tier is slow/flaky on gemma-4-31b; generous timeout gives retries room to succeed +AGENT_TIMEOUT = int(os.getenv("AGENT_TIMEOUT", "1200")) # env-overridable; matches wiki_scheduler pattern. Generous default for slow / deep-exploring LLM runs. UUID_RE = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}") logging.basicConfig( @@ -157,13 +157,15 @@ def extract_facts_from_chunk(ds_id: str, title: str, idx: int, total: int, chunk f"numbers, events, named decisions. Ignore filler, opinion, and generic\n" f"statements. Aim for quality over quantity: typically 3-8 facts per chunk.\n\n" f"For EACH fact:\n" - f' a) Call save_fact(content="", certainty=0.8,\n' - f' source="document", keywords=[2-4 precise tags], importance=0.6,\n' - f' notes="Extracted from {title} chunk {idx}/{total}"). Record the\n' - f" returned fact id.\n" + f' a) Call save_fact(content="",\n' + f' source="document", keywords=[2-4 precise tags],\n' + f' notes="Extracted from {title} chunk {idx}/{total}"). Judge\n' + f" certainty and importance per the tool's docstring guidance.\n" + f" Record the returned fact id.\n" f" b) Call create_relation(from_entity_id=, " f'to_entity_id="{ds_id}", relation_type="derived_from",\n' - f' relevance_score=0.9, description="Fact extracted from {title}").\n\n' + f' description="Fact extracted from {title}"). Judge\n' + f" relevance_score and importance_score per the tool's docstring.\n\n" f"Do NOT call get_entity. Do NOT call update_entity on the datasource.\n" f"Do NOT touch the datasource content — it is read-only.\n\n" f"When all facts in this chunk are processed, call final_answer with\n" @@ -210,7 +212,8 @@ def central_review(ds_id: str, title: str, fact_ids: list[str]) -> None: f" or is_example_of. Only create relations that are genuinely meaningful.\n" f"2. If the set as a whole suggests a broader observation or inference that\n" f" none of the individual facts capture, call save_thought with that\n" - f" thought (certainty=0.6-0.8, source='agent-inference') and then\n" + f" thought (source='agent-inference'; judge certainty + importance per\n" + f" the tool docstring) and then\n" f' create_relation(thought_id, "{ds_id}", "elaborates").\n' f"3. Optionally run recall_memory with 1-2 queries derived from the facts\n" f" to find related EXISTING entities in memory. If any are clearly\n" diff --git a/braindb/services/graph.py b/braindb/services/graph.py index 63dc591..bcd1eaa 100644 --- a/braindb/services/graph.py +++ b/braindb/services/graph.py @@ -34,6 +34,7 @@ ( t.accumulated_relevance * r.relevance_score + * COALESCE(r.importance_score, 0.5) * CASE t.depth + 1 WHEN 1 THEN 1.0 WHEN 2 THEN 0.8 @@ -47,10 +48,7 @@ r.notes, t.seed_origin_id FROM traversal t - JOIN relations r ON ( - r.from_entity_id = t.id - OR (r.is_bidirectional AND r.to_entity_id = t.id) - ) + JOIN relations r ON r.from_entity_id = t.id OR r.to_entity_id = t.id JOIN entities target ON target.id = CASE WHEN r.from_entity_id = t.id THEN r.to_entity_id @@ -61,6 +59,7 @@ AND ( t.accumulated_relevance * r.relevance_score + * COALESCE(r.importance_score, 0.5) * CASE t.depth + 1 WHEN 1 THEN 1.0 WHEN 2 THEN 0.8 ELSE 0.6 END ) > %s ) diff --git a/tests/test_create_relation_persistence.py b/tests/test_create_relation_persistence.py new file mode 100644 index 0000000..a0dba55 --- /dev/null +++ b/tests/test_create_relation_persistence.py @@ -0,0 +1,57 @@ +""" +Locks in the importance_score wiring on relations. + +For agent-created rows the column was NULL until this round; now both +relevance_score and importance_score must persist whatever the caller +sets, and a fresh GET must return the same values. +""" +import requests + + +def test_relation_persists_both_scores(api, make_fact): + a = make_fact("Anchor entity A for persistence check.") + b = make_fact("Anchor entity B for persistence check.") + + body = { + "from_entity_id": a["id"], + "to_entity_id": b["id"], + "relation_type": "supports", + "relevance_score": 0.82, + "importance_score": 0.71, + "description": "Persistence round-trip", + } + r = requests.post(f"{api}/api/v1/relations", json=body, timeout=30) + assert r.status_code == 201, f"create relation failed: {r.status_code} {r.text}" + rel = r.json() + + # POST response carries both scores. + assert rel["relevance_score"] == 0.82 + assert rel["importance_score"] == 0.71 + + # And a fresh GET reads them back identically — proves the column is + # actually written, not just echoed in the response. + r = requests.get(f"{api}/api/v1/relations/{rel['id']}", timeout=10) + assert r.status_code == 200 + fresh = r.json() + assert fresh["relevance_score"] == 0.82 + assert fresh["importance_score"] == 0.71 + + +def test_relation_importance_score_not_null_by_default(api, make_fact): + """If the caller omits importance_score, it falls back to the schema + default (0.5) — NOT NULL. This locks the schema/router contract.""" + a = make_fact("A for default-import.") + b = make_fact("B for default-import.") + + body = { + "from_entity_id": a["id"], + "to_entity_id": b["id"], + "relation_type": "supports", + } + r = requests.post(f"{api}/api/v1/relations", json=body, timeout=30) + assert r.status_code == 201 + rel = r.json() + + assert rel["importance_score"] is not None + # Default is the schema's neutral 0.5. + assert rel["importance_score"] == 0.5 diff --git a/tests/test_graph_scoring_uses_importance.py b/tests/test_graph_scoring_uses_importance.py new file mode 100644 index 0000000..42ab602 --- /dev/null +++ b/tests/test_graph_scoring_uses_importance.py @@ -0,0 +1,63 @@ +""" +Locks in that the recall graph CTE multiplies BOTH per-edge scores +(relevance_score AND importance_score) into the per-hop accumulated +relevance. Before this round, importance_score sat in the column +unused; the LLM's judgment of edge importance didn't affect ranking. + +Test approach: drive `graph_expand` directly (the function the recall +pipeline calls). This isolates the behavior we want to lock without +the noise of full-text discovery fallback + diversity-quota filtering +that /memory/context layers on top. +""" +import uuid + +import requests +from braindb.db import get_conn +from braindb.services.graph import graph_expand + + +def test_importance_score_moves_per_hop_relevance(api, make_fact): + """Two relations from the same seed, identical relation_type and + relevance_score, ONLY differing in importance_score. The hop's + accumulated_relevance from graph_expand must reflect the difference.""" + tag = uuid.uuid4().hex + seed = make_fact(f"Seed for {tag}", keywords=[tag]) + hi_target = make_fact("Generic hi target.") + lo_target = make_fact("Generic lo target.") + + for target_id, imp in ((hi_target["id"], 0.9), (lo_target["id"], 0.2)): + body = { + "from_entity_id": seed["id"], + "to_entity_id": target_id, + "relation_type": "elaborates", + "relevance_score": 0.7, + "importance_score": imp, + "description": "Test edge", + } + r = requests.post(f"{api}/api/v1/relations", json=body, timeout=30) + assert r.status_code == 201, r.text + + with get_conn() as conn: + rows = graph_expand(conn, [seed["id"]], max_depth=1, min_relevance=0.01) + + by_id = {str(r["id"]): r for r in rows} + assert hi_target["id"] in by_id, "hi_target not reached by graph_expand" + assert lo_target["id"] in by_id, "lo_target not reached by graph_expand" + + hi_rel = by_id[hi_target["id"]]["relevance"] + lo_rel = by_id[lo_target["id"]]["relevance"] + + # Both reached via the same one-hop path; only importance_score differs. + # With the fix in graph.py the per-hop multiplier multiplies by + # COALESCE(r.importance_score, 0.5), so hi (imp 0.9) > lo (imp 0.2). + # Per-hop math: 1.0 * 0.7 * imp * depth_penalty(1.0) + assert hi_rel > lo_rel, ( + f"importance_score does not move per-hop relevance: " + f"hi={hi_rel:.4f} lo={lo_rel:.4f}" + ) + # Sanity: ratio hi/lo ~= 0.9 / 0.2 = 4.5 (within rounding). + ratio = hi_rel / lo_rel + assert 4.0 < ratio < 5.0, ( + f"per-hop relevance ratio should track importance_score ratio " + f"(0.9/0.2=4.5); got hi/lo={ratio:.3f}" + ) diff --git a/tests/test_graph_walks_both_directions.py b/tests/test_graph_walks_both_directions.py new file mode 100644 index 0000000..db0e449 --- /dev/null +++ b/tests/test_graph_walks_both_directions.py @@ -0,0 +1,57 @@ +""" +Locks in that the graph CTE walks relations in BOTH directions +regardless of the `is_bidirectional` flag. + +Before this round, the CTE only walked `from -> to` unless +`is_bidirectional=true` was set on the edge. That diverged from +`services/tree.py` (which always walks both ways) and meant recall +seeded from the `to_entity_id` could not reach the `from_entity_id` +through a unidirectional edge — the bloat the user explicitly flagged. + +Test approach: drive `graph_expand` directly. Avoids the noise of the +full recall pipeline's discovery fallback that complicates HTTP-based +tests with arbitrary DB state. +""" +import uuid + +import requests +from braindb.db import get_conn +from braindb.services.graph import graph_expand + + +def test_graph_expand_walks_unidirectional_edge_backwards(api, make_fact): + tag = uuid.uuid4().hex + + # A is the "from" side. B is the "to" side and is the SEED. + # Edge: A -> B with is_bidirectional=false (default). + # graph_expand seeded from [B] must still reach A by walking the + # edge backwards. + a = make_fact(f"Anchor A for {tag}") + b = make_fact(f"Anchor B for {tag}", keywords=[tag]) + + body = { + "from_entity_id": a["id"], + "to_entity_id": b["id"], + "relation_type": "supports", + "relevance_score": 0.8, + "importance_score": 0.6, + "is_bidirectional": False, + "description": "Unidirectional A -> B for bidirectional-walk test", + } + r = requests.post(f"{api}/api/v1/relations", json=body, timeout=30) + assert r.status_code == 201, r.text + + with get_conn() as conn: + rows = graph_expand(conn, [b["id"]], max_depth=1, min_relevance=0.05) + + by_id = {str(r["id"]): r for r in rows} + assert b["id"] in by_id, "B (the seed) is missing from graph_expand result" + assert a["id"] in by_id, ( + "A (the from-side of a unidirectional edge) was NOT reached by " + "graph_expand seeded from B. The CTE is not walking edges backwards." + ) + a_row = by_id[a["id"]] + assert a_row["min_depth"] == 1, ( + f"A should be reached at depth=1, got depth={a_row['min_depth']}" + ) + assert a_row["via_relation_type"] == "supports" diff --git a/tests/test_ingest_watcher_no_dictation.py b/tests/test_ingest_watcher_no_dictation.py new file mode 100644 index 0000000..4910b80 --- /dev/null +++ b/tests/test_ingest_watcher_no_dictation.py @@ -0,0 +1,55 @@ +""" +Locks in that the watcher's prompts NO LONGER dictate parameter values +to the LLM. Before this round, every fact extracted from every document +got the same `certainty=0.8`, `importance=0.6`, `relevance_score=0.9` +because the watcher's prompt copy-pasted those numbers into the +example call the LLM was told to make. That hid the LLM's judgment. + +This is a static regex check on the prompt-builder source — no LLM +call, no DB hit, deterministic. The live variance check is performed +manually with a user-chosen file after deploy. +""" +from pathlib import Path + +import re + + +WATCHER_PATH = Path(__file__).resolve().parents[1] / "braindb" / "ingest_watcher.py" + + +def _watcher_source() -> str: + return WATCHER_PATH.read_text(encoding="utf-8") + + +def test_chunk_extraction_prompt_does_not_dictate_certainty(): + src = _watcher_source() + # Locate the chunk-extraction prompt block. + # The bug was a literal "certainty=0.8" embedded in the prompt string. + matches = re.findall(r'["\'][^"\']*certainty=0\.[0-9][^"\']*["\']', src) + assert not matches, ( + "Watcher prompt still dictates a certainty literal to the LLM: " + f"{matches}. The LLM should judge certainty per save_fact / " + "save_thought docstring instead." + ) + + +def test_chunk_extraction_prompt_does_not_dictate_importance(): + src = _watcher_source() + matches = re.findall(r'["\'][^"\']*importance=0\.[0-9][^"\']*["\']', src) + # Note: the watcher's datasource ingest body (importance=0.6) is fine + # — that's a Python dict literal, NOT inside a prompt string. The + # regex above looks specifically for the pattern inside a quoted + # prompt-text segment. + assert not matches, ( + "Watcher prompt still dictates an importance literal to the LLM: " + f"{matches}. The LLM should judge importance per the tool docstring." + ) + + +def test_chunk_extraction_prompt_does_not_dictate_relevance_score(): + src = _watcher_source() + matches = re.findall(r'["\'][^"\']*relevance_score=0\.[0-9][^"\']*["\']', src) + assert not matches, ( + "Watcher prompt still dictates a relevance_score literal to the LLM: " + f"{matches}. The LLM should judge relevance_score per create_relation's docstring." + ) From 92058ef324e4c0c6d29ff9688f6f224803310a7d Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Tue, 2 Jun 2026 13:35:06 +0100 Subject: [PATCH 08/10] test: add pytest-asyncio dev dep + fix ingest test isolation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * pyproject.toml: add pytest-asyncio==0.23.7 to [dev]. Existing tests use @pytest.mark.asyncio decorators (test_handoff_hooks, test_runhooks_countdown, test_final_answer_rename) but the plugin was not listed in deps, so `pip install -e ".[dev]"` left them skipped silently on a clean install. * tests/test_ingest.py: the three datasource-ingest tests used fixed content strings, so a previous run's row in the DB caused dedup-by-hash to fire and the 201 assertion to fail on subsequent runs. Prepend a per-run uuid to the content so each invocation is genuinely fresh. No production-code change. 134/134 pass (8 wiki_jobs_grouping deselected — those use the host port mapping; they run from the host per tests/README.md, not from inside the api container). Co-Authored-By: Claude Opus 4.7 --- pyproject.toml | 1 + tests/test_ingest.py | 10 +++++++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index cb01094..920fc2a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,6 +52,7 @@ dependencies = [ dev = [ "pytest==9.0.3", "pytest-timeout==2.4.0", + "pytest-asyncio==0.23.7", ] [tool.hatch.build.targets.wheel] diff --git a/tests/test_ingest.py b/tests/test_ingest.py index e1b0b7a..45b8261 100644 --- a/tests/test_ingest.py +++ b/tests/test_ingest.py @@ -11,6 +11,7 @@ end-to-end ingestion: the Smart Sand article's 10,357-char content was preserved across three full runs of the watcher's chunked pipeline. """ +import uuid from pathlib import Path import requests @@ -56,7 +57,8 @@ def _find_datasource_by_title(api: str, title: str) -> dict | None: def test_ingest_new_returns_201(api, created_entities): """First ingest of a fresh file returns 201.""" - content = "A unique pytest ingest-test body. " * 10 + # Unique per-run content so prior runs' rows in the DB can't dedup-fire on us. + content = f"A unique pytest ingest-test body {uuid.uuid4().hex}. " * 10 _write_sample(content) try: r = requests.post( @@ -80,7 +82,8 @@ def test_ingest_new_returns_201(api, created_entities): def test_ingest_duplicate_returns_200(api, created_entities): """Second ingest with the same bytes returns 200 (idempotent).""" - content = "Idempotency pytest body. " * 15 + # Unique per-run content so prior runs' rows in the DB can't dedup-fire on us. + content = f"Idempotency pytest body {uuid.uuid4().hex}. " * 15 _write_sample(content) try: # First call — 201 @@ -120,7 +123,8 @@ def test_ingest_dup_preserves_first_seen_metadata(api, created_entities): with new metadata). A second ingest with different keywords must not overwrite the first-seen keywords or swap the id. """ - content = "Dup-metadata pytest body. " * 20 + # Unique per-run content so prior runs' rows in the DB can't dedup-fire on us. + content = f"Dup-metadata pytest body {uuid.uuid4().hex}. " * 20 _write_sample(content) try: r1 = requests.post( From 94411867270786a6236a2fcd3ef33b0c4d5901ca Mon Sep 17 00:00:00 2001 From: dimknaf <136385722+dimknaf@users.noreply.github.com> Date: Tue, 2 Jun 2026 13:35:28 +0100 Subject: [PATCH 09/10] refactor(tree): nested JSON shape, one builder, filter keywords + retired wikis MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * services/tree.py: one build_entity_tree function. Recursive CTE carries parent_id + accumulated_score (relevance × COALESCE( importance_score, 0.5) × depth_penalty, same formula as graph.py). DISTINCT ON (id) ORDER BY id, accumulated_score DESC -- multi-path first-wins by best score. Skip tagged_with edges + target.entity_type='keyword' by default; skip wikis_ext.retired_at IS NOT NULL. New shape: root keyed by entity_type, children arrays per node, _truncated last-child marker. * routers/memory.py: /memory/tree/ returns the nested shape; new query params include_keywords, top_k (default 40), min_path_score. * agent/tools.py::view_tree: returns json.dumps(tree) directly; _tree_label helper removed. * system_prompt.md: view_tree blurb updated to describe nested JSON. Path A 5/5 PASS, view_tree 0/5 -> 1/5 (the agent reaches for tree now that the shape is structured). Path B 5/5 PASS, 1090s -> 773s (-29%), 54 tool calls -> 40 (-26%), zero delegate calls on q4 (was 2). Two latent bugs caught by the new shape and fixed in this commit: keyword children leaking through non-tagged_with edges; duplicate retired-wiki siblings. Frontend Graph tab will be broken until graph.js consumes the new shape -- follow-up commit. Co-Authored-By: Claude Opus 4.7 --- braindb/agent/prompts/system_prompt.md | 10 +- braindb/agent/tools.py | 46 ++---- braindb/routers/memory.py | 30 +++- braindb/services/tree.py | 188 ++++++++++++++++--------- 4 files changed, 163 insertions(+), 111 deletions(-) diff --git a/braindb/agent/prompts/system_prompt.md b/braindb/agent/prompts/system_prompt.md index c52fb02..7d289d3 100644 --- a/braindb/agent/prompts/system_prompt.md +++ b/braindb/agent/prompts/system_prompt.md @@ -55,10 +55,12 @@ fall back to flat SQL. graph traversal + decay + ranking. Use first when you don't yet know which specific entity you want — "what do we know about X." 2. **`view_tree(, max_depth=N)`** — the efficient "explore around this - entity" tool. One call returns the entity's neighbours grouped by - relation type, plus their edge scores, 1-N hops out. When you have an - entity ID and want to know what's related to it, this is usually a - better next step than another `recall_memory` about the same entity. + entity" tool. Returns a nested JSON tree: root keyed by entity_type, + `children` arrays per node, 1-N hops out, keyword/retired-wiki noise + filtered, `_truncated` marker if more remain. When you have an entity + ID and want what's connected, this is usually a better next step than + another `recall_memory`. On hub entities (wikis), pass `max_depth=3` + to see narrative chains. 3. **`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 diff --git a/braindb/agent/tools.py b/braindb/agent/tools.py index 12a5dee..b2fca19 100644 --- a/braindb/agent/tools.py +++ b/braindb/agent/tools.py @@ -636,10 +636,13 @@ async def delete_relation(relation_id: str) -> str: @function_tool @_verbose("view_tree") async def view_tree(entity_id: str, max_depth: int = 2) -> str: - """⭐ Reveals an entity's connections in one call: relations + 1-N hop - neighbours + edge scores. Especially useful when you have an entity ID - (from a previous result) and want its graph context — often a sharper - choice than another `recall_memory` about the same entity. + """⭐ Reveals an entity's neighbourhood as a nested JSON tree: + root keyed by ``entity_type``, ``children`` arrays per node, multi-path + first-wins, keyword/retired-wiki noise filtered, ``_truncated`` marker + when more remain. Especially useful when you have an entity ID (from a + previous result) and want its graph context — often a sharper choice + than another `recall_memory` about the same entity. Pass `max_depth=3` + on hub entities (wikis with many connections) to see narrative chains. Args: entity_id: UUID of the root entity. @@ -650,45 +653,14 @@ async def view_tree(entity_id: str, max_depth: int = 2) -> str: try: from braindb.services.tree import build_entity_tree with get_conn() as conn: - tree = build_entity_tree(conn, entity_id, max_depth) + tree = build_entity_tree(conn, entity_id, max_depth=max_depth) if tree is None: return _err(f"Entity not found: {entity_id}") - root = tree["root"] - conns = tree["connections"] - out = [f"ROOT [{root['entity_type']}] {_tree_label(root)} (id: {root['id']})"] - if not conns: - out.append("\n(no connections)") - return _truncate("\n".join(out)) - # Group output by depth - last_depth = None - for c in conns: - d = c["depth"] - if d != last_depth: - same_n = sum(1 for x in conns if x["depth"] == d) - out.append(f"\nDEPTH {d} ({same_n}):") - last_depth = d - ent = c["entity"] - dir_short = "out" if c.get("direction") == "outgoing" else "in" - out.append( - f" [{dir_short}] {c['via_relation_type']} (rel={c['relevance']})\n" - f" [{ent['entity_type']}] {_tree_label(ent)} (id: {ent['id']})" - ) - return _truncate("\n".join(out)) + return _truncate(json.dumps(tree, indent=2, default=str, ensure_ascii=False)) except Exception as e: return _err(str(e)) -def _tree_label(entity: dict) -> str: - """Short label for tree output. Wikis use canonical_name; others truncate content.""" - import re as _r - content = (entity.get('content') or '').replace('\n', ' ') - if entity.get('entity_type') == 'wiki': - m = _r.search(r'canonical_name=([^\s]+)', content) - if m: - return m.group(1) - return content[:80] - - @function_tool @_verbose("search_sql") async def search_sql(query: str) -> str: diff --git a/braindb/routers/memory.py b/braindb/routers/memory.py index 57b11d6..8638483 100644 --- a/braindb/routers/memory.py +++ b/braindb/routers/memory.py @@ -79,15 +79,33 @@ def get_rules(): @router.get("/tree/{entity_id}") -def entity_tree(entity_id: UUID, max_depth: int = Query(default=2, ge=1, le=3)): - """Return an entity and its graph connections organized by depth. - - Uses the shared `build_entity_tree` service (also used by the agent's - `view_tree` tool) so HTTP callers and the agent see the same data. +def entity_tree( + entity_id: UUID, + max_depth: int = Query(default=2, ge=1, le=3), + include_keywords: bool = Query(default=False), + top_k: int = Query(default=40, ge=1, le=500), + min_path_score: float = Query(default=0.0, ge=0.0, le=1.0), +): + """Return an entity and its graph neighbourhood as a nested JSON tree. + + Round-2f shape: root keyed by ``entity_type``; ``children`` array of + typed nodes (each keyed by its own ``entity_type`` and labelled to + its type — wiki=title, fact/thought=content, source=filename, etc.); + multi-path first-wins by accumulated path score; ``tagged_with`` + keyword edges skipped by default (root's ``keywords`` array is the + one-liner instead). Top-``top_k`` connections are kept; the rest are + summarised with a single ``_truncated`` marker. """ from braindb.services.tree import build_entity_tree with get_conn() as conn: - tree = build_entity_tree(conn, str(entity_id), max_depth) + tree = build_entity_tree( + conn, + str(entity_id), + max_depth=max_depth, + include_keywords=include_keywords, + top_k=top_k, + min_path_score=min_path_score, + ) if tree is None: raise HTTPException(404, "Entity not found") return tree diff --git a/braindb/services/tree.py b/braindb/services/tree.py index c0467be..01c9216 100644 --- a/braindb/services/tree.py +++ b/braindb/services/tree.py @@ -1,68 +1,126 @@ """Shared tree-build service. -Single source of truth for "walk the relation graph outward from an entity". -Used by both the HTTP endpoint (`/api/v1/memory/tree/`) and the agent's -`view_tree` tool. Walks bidirectionally and respects `max_depth`. +`build_entity_tree` walks the relation graph bidirectionally from a root +entity and returns a nested JSON tree. The same function is used by the +HTTP endpoint `/api/v1/memory/tree/` and the agent's `view_tree` +tool — one shape, no behaviour drift. + +Shape (root keyed by ``entity_type``): + + { + "": "