-
-
-
-
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py b/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py
index 6c773c1c8..4a506c2b1 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/bridge.py
@@ -7,8 +7,10 @@
from __future__ import annotations
-from typing import Any
-
+from shared.services.retrieval.agent_explore.evidence_pool import (
+ Candidate,
+ pool_trace_records,
+)
from shared.services.retrieval.agent_explore.types import AgentStep
from shared.services.retrieval.trace import DecisionTraceStep
@@ -28,53 +30,49 @@ def build_decision_trace(steps: list[AgentStep]) -> list[DecisionTraceStep]:
observation: dict[str, object] = {"observation_text": observation_text}
if step.ref_status is not None:
observation["ref_status"] = step.ref_status
+ if step.candidates is not None:
+ observation["candidates"] = step.candidates
+ decision: dict[str, object] = {
+ "action": step.tool_name or "no_tool_call",
+ "args": step.tool_args,
+ "round_index": step.round_index,
+ }
+ result: dict[str, object] = {
+ "status": "error" if step.error else "ok",
+ "error": step.error,
+ }
+ if step.picked is not None:
+ result["picked"] = step.picked
+ if step.pick_rejected is not None:
+ result["pick_rejected"] = step.pick_rejected
trace_steps.append(
DecisionTraceStep(
step_index=step.step_index,
agent="agent_explore",
phase=phase,
observation=observation,
- decision={
- "action": step.tool_name or "no_tool_call",
- "args": step.tool_args,
- },
- result={
- "status": "error" if step.error else "ok",
- "error": step.error,
- },
- budget={
- "tokens_used_delta": step.tokens_used_delta,
- "tokens_used_total": step.tokens_used_total,
- },
+ decision=decision,
+ result=result,
+ budget={},
elapsed_ms=step.elapsed_ms,
)
)
return trace_steps
-def attach_ref_provenance(
+def attach_evidence_pool(
steps: list[DecisionTraceStep],
- *,
- agent_selected_refs: list[dict[str, Any]] | None,
- fallback_refs: list[dict[str, Any]],
- resolved_refs: list[dict[str, Any]],
- dropped_refs: list[dict[str, Any]],
+ pool: list[Candidate],
) -> list[DecisionTraceStep]:
- """Write ref provenance onto the existing finish TRACE step.
+ """Write the evidence pool onto the existing finish TRACE step.
- Reuses ``DecisionTraceStep.result``. If the episode never called finish,
- append one finish step so the same public TRACE shape still holds the
- four lists.
+ If the episode never called finish, append one finish step so the same
+ public TRACE shape still holds the pool.
"""
- provenance = {
- "agent_selected_refs": agent_selected_refs,
- "fallback_refs": fallback_refs,
- "resolved_refs": resolved_refs,
- "dropped_refs": dropped_refs,
- }
+ records = pool_trace_records(pool)
for step in reversed(steps):
if step.phase == "finish":
- step.result = {**step.result, **provenance}
+ step.result = {**step.result, "evidence_pool": records}
return steps
steps.append(
DecisionTraceStep(
@@ -82,8 +80,8 @@ def attach_ref_provenance(
agent="agent_explore",
phase="finish",
observation={"observation_text": "finish was not called"},
- decision={"action": "finish", "args": {"refs": None}},
- result={"status": "ok", "error": None, **provenance},
+ decision={"action": "finish", "args": {}},
+ result={"status": "ok", "error": None, "evidence_pool": records},
)
)
return steps
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/budget.py b/packages/shared-python/shared/services/retrieval/agent_explore/budget.py
index 47a0c6067..652f559d2 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/budget.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/budget.py
@@ -2,17 +2,13 @@
Standalone from ``nav/nav_token_budget.py`` on purpose — this package must
have no runtime dependency on ``nav/`` (see ``config.py``'s module
-docstring). Re-reads the same ``RETRIEVAL_NAV_TOKEN_LIMIT`` env var the plan
-calls for, so operators keep one token-limit knob across both agentic
-routes, but the counting mechanism is a fresh, request-scoped object (this
-episode runs as one async function, not a separate thread with recursive
-calls), not the nav contextvar machinery ``nav_token_budget`` needed for its
-own call shape.
+docstring). Exploration is limited by steps; the wall clock is a
+service-level guard. Tokens are not limited: ``tokens_used`` only records the
+episode total.
"""
from __future__ import annotations
-import os
import time
from dataclasses import dataclass, field
from typing import Any
@@ -22,26 +18,13 @@
AGENT_EXPLORE_WALL_CLOCK_SECONDS,
)
-_ENV_TOKEN_LIMIT = "RETRIEVAL_NAV_TOKEN_LIMIT"
-_DEFAULT_TOKEN_LIMIT = 100_000
-
-StopReason = str # one of: "token_limit" | "max_steps" | "wall_clock"
-
-
-def resolve_token_limit() -> int:
- """Always a positive limit: env override, else the shared default."""
- try:
- limit = int(os.environ.get(_ENV_TOKEN_LIMIT, "").strip())
- except ValueError:
- limit = 0
- return limit if limit > 0 else _DEFAULT_TOKEN_LIMIT
+StopReason = str # one of: "max_steps" | "wall_clock"
@dataclass
class EpisodeBudget:
- """Tracks one episode's LLM-token / step / wall-clock spend."""
+ """Tracks one episode's step / wall-clock spend and its token total."""
- token_limit: int = field(default_factory=resolve_token_limit)
max_steps: int = AGENT_EXPLORE_MAX_STEPS
wall_clock_seconds: float = AGENT_EXPLORE_WALL_CLOCK_SECONDS
tokens_used: int = 0
@@ -59,22 +42,13 @@ def record_usage(self, usage: dict[str, Any] | None) -> None:
def record_step(self) -> None:
self.steps_used += 1
+ def remaining_seconds(self) -> float:
+ return self.wall_clock_seconds - (time.monotonic() - self._started_at)
+
def exhausted(self) -> StopReason | None:
"""Return which budget dimension is exceeded, if any, else None."""
- if self.tokens_used >= self.token_limit:
- return "token_limit"
if self.steps_used >= self.max_steps:
return "max_steps"
- if time.monotonic() - self._started_at >= self.wall_clock_seconds:
+ if self.remaining_seconds() <= 0:
return "wall_clock"
return None
-
- def snapshot(self) -> dict[str, Any]:
- return {
- "token_limit": self.token_limit,
- "tokens_used": self.tokens_used,
- "max_steps": self.max_steps,
- "steps_used": self.steps_used,
- "wall_clock_seconds": self.wall_clock_seconds,
- "elapsed_seconds": round(time.monotonic() - self._started_at, 3),
- }
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/config.py b/packages/shared-python/shared/services/retrieval/agent_explore/config.py
index 258dfca1c..4cb318acb 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/config.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/config.py
@@ -8,8 +8,6 @@
from __future__ import annotations
-from shared.services.retrieval.agent_tools.registry import REF_ADDRESS_ONE_OF, REF_ADDRESS_RULE
-
AGENT_EXPLORE_MODEL = "deepseek-v4-flash"
# Model for AGENT_EXPLORE_HARNESS=cursor_sdk (harness/cursor_harness.py) —
@@ -22,11 +20,11 @@
AGENT_EXPLORE_CURSOR_MODEL = "composer-2.5"
# One LLM turn = one round-trip that may contain several parallel tool calls
-# (see episode.py).
-AGENT_EXPLORE_MAX_STEPS = 12
+# (see episode.py). Both harnesses count one explore turn as one step.
+AGENT_EXPLORE_MAX_STEPS = 15
# Wall-clock ceiling for the whole episode (LLM round-trips + tool
-# dispatch), independent of the token budget. New constant, not tuned yet.
+# dispatch). A service-level guard, never shown to the model.
AGENT_EXPLORE_WALL_CLOCK_SECONDS = 180.0
# Max tokens requested per LLM completion turn (thinking disabled — see
@@ -38,104 +36,43 @@
FINISH_TOOL_SCHEMA: dict[str, object] = {
"type": "object",
"properties": {
- "refs": {
- "type": "array",
- "description": (
- "Final cited evidence, in priority order. Each item "
- "identifies one section or chunk you have already looked at "
- "via corpus.read, corpus.assets, or corpus.query_table."
- ),
- "items": {
- "type": "object",
- "properties": {
- "document_id": {
- "type": "string",
- "description": "Document that owns the cited section or chunk.",
- },
- "section_path": {
- "type": "string",
- "description": "Cited section. Omit when chunk_id is set.",
- },
- "chunk_id": {
- "type": "string",
- "description": "Cited chunk. Omit when section_path is set.",
- },
- },
- "required": ["document_id"],
- "oneOf": REF_ADDRESS_ONE_OF,
- "additionalProperties": False,
- "description": REF_ADDRESS_RULE,
- },
- },
"notes": {
"type": "string",
"description": (
- "Optional short note on why these refs answer the query, "
- "or why none were found."
+ "Optional short note for the answer writer, or why nothing "
+ "was found."
),
},
},
- "required": ["refs"],
+ "required": [],
"additionalProperties": False,
}
FINISH_TOOL_DESCRIPTION = (
- "Call this when you have gathered enough evidence to answer the query, "
- "or when you are certain the corpus does not contain an answer. This "
- "ends the exploration — do not answer in plain text; the final answer "
- "is synthesized downstream from the refs you cite here."
+ "End the exploration now. The evidence pool becomes the final evidence."
)
-# Appended after the verbatim CORPUS_SCHEMA.md text (schema_doc.py) to form
-# this harness's system prompt. Kept out of CORPUS_SCHEMA.md itself because
-# the finish-tool loop contract is agent_explore-specific, not something the
-# MCP-facing harnesses (Cursor/Codex/Claude) need — see that file's own
-# "single source... do not duplicate" header.
-LOOP_CONTRACT_SUFFIX = f"""
+# Harness-level tool, not in agent_tools.REGISTRY, so MCP never exposes it.
+PICK_TOOL_NAME = "corpus.pick"
----
-
-## Exploration loop contract
-
-You are exploring this corpus autonomously to answer one query. Tool
-names in this prompt are the registered names (`corpus.read`). This
-function-calling loop exposes the same tools with `.` replaced by `_`
-(`corpus_read`); use that underscore form when calling. Logs keep the
-dotted name.
-
-Use the tools above to navigate; you may call several tools in one turn
-when they are independent. If a call's arguments depend on another
-call's result, do not issue them in the same turn. In particular: do
-not call `corpus.grep` in the same turn as `corpus.recall`,
-`corpus.read`, `corpus.list_documents`, `corpus.outline`,
-`corpus.node_filter`, or `corpus.assets` unless `pattern` / `patterns`
-is already known — those calls produce the term; a grep with no term is
-an empty call. Wait for the result, then grep. `corpus.list_documents`
-enumerates the entire namespace: use it only if the user explicitly
-asks to list or inventory the corpus's documents, never as the starting
-step for a content question. When you have enough evidence, call
-`{FINISH_TOOL_NAME}`
-with the `refs` you want cited as the answer — do not write the final answer
-as plain text yourself, it is synthesized downstream from your cited refs.
-If you exhaust your tool budget without a confident answer, call
-`{FINISH_TOOL_NAME}` with your best-effort `refs` (or an empty list plus a
-`notes` explanation of why nothing was found) rather than continuing to
-call other tools.
-
-`refs` is REQUIRED and must not be omitted or left empty if you called
-`corpus.read`, `corpus.assets`, or `corpus.query_table` even once during
-this exploration: copy the `document_id` and `chunk_id`/`section_path` of
-every section/chunk you looked at that supports your answer into `refs`
-before calling `{FINISH_TOOL_NAME}`. Calling `{FINISH_TOOL_NAME}` with no
-`refs` after having already read relevant content discards that evidence.
+PICK_TOOL_SCHEMA: dict[str, object] = {
+ "type": "object",
+ "properties": {
+ "pick": {
+ "type": "array",
+ "items": {"type": "string", "pattern": r"^(R\d+\.\d+|O\d+)$"},
+ "uniqueItems": True,
+ "description": (
+ "Ids from your latest corpus.read / corpus.outline results to "
+ "keep. Pass an empty list to keep none."
+ ),
+ },
+ },
+ "required": ["pick"],
+ "additionalProperties": False,
+}
-Before calling `{FINISH_TOOL_NAME}`, if you have not called `corpus.read`,
-`corpus.assets`, or `corpus.query_table` even once this exploration, you
-have not actually verified anything yet — a search tool returning
-candidates is not the same as having read them. In that case, either read
-your best candidate section first, or — only if you have positively
-confirmed there is nothing to read (e.g. a structural check came back
-with zero matching sections) — say so explicitly in `notes`. Repeatedly
-rephrasing the same search instead of reading a candidate you already
-found is not a substitute for reading it.
-"""
+PICK_TOOL_DESCRIPTION = (
+ "Pick which of your latest corpus.read / corpus.outline results join the "
+ "evidence pool. Body text you read but do not pick cannot be read again."
+)
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py
index f002472bd..febf2da20 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/dispatch.py
@@ -20,14 +20,21 @@
from shared.services.retrieval.document_scope import DocumentScope
-from collections.abc import Callable
+from collections.abc import Callable, Mapping
from contextlib import AbstractAsyncContextManager
from typing import Any
from sqlalchemy.ext.asyncio import AsyncSession
from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext
-from shared.services.retrieval.agent_tools import REGISTRY, ToolBudget, ToolContext, ToolResult
+from shared.services.retrieval.agent_tools import (
+ REGISTRY,
+ Decision,
+ ReadableAddresses,
+ ToolBudget,
+ ToolContext,
+ ToolResult,
+)
DbFactory = Callable[[], AbstractAsyncContextManager[AsyncSession]]
@@ -42,6 +49,8 @@ async def dispatch_tool_call(
document_scope: DocumentScope = DocumentScope(),
budget: ToolBudget | None = None,
query: str = "",
+ readable: ReadableAddresses | None = None,
+ decided: Mapping[tuple[str, str], Decision] | None = None,
) -> ToolResult:
"""Run one ``REGISTRY`` tool call against a fresh, call-scoped DB session.
@@ -52,7 +61,8 @@ async def dispatch_tool_call(
into the (call-scoped, already-closed) ``ToolContext``. ``query`` is the
end user's original query for this episode (see ``ToolContext.query``'s
docstring) — both harnesses' ``run_episode`` already receive it, so this
- is a pass-through, not a new source of truth.
+ is a pass-through, not a new source of truth. ``readable``/``decided`` are
+ the episode's ``corpus.read`` gates, passed through unchanged.
"""
try:
async with db_factory() as db:
@@ -65,6 +75,8 @@ async def dispatch_tool_call(
budget=budget or ToolBudget(),
document_scope=document_scope.narrow(list(pins)) if pins is not None else document_scope,
query=query,
+ readable=readable,
+ decided=decided,
)
return await REGISTRY.dispatch(name, tool_ctx, args)
except Exception as exc: # noqa: BLE001 - one broken tool must not kill the episode
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/evidence_pool.py b/packages/shared-python/shared/services/retrieval/agent_explore/evidence_pool.py
new file mode 100644
index 000000000..b13a53161
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/evidence_pool.py
@@ -0,0 +1,404 @@
+"""Per-round addresses, pick ids, and the episode evidence pool.
+
+Results of one round become usable only when the next round begins: their
+addresses become readable, and their pick ids open a pick phase.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from types import MappingProxyType
+from typing import Any, Literal
+
+from shared.services.retrieval.agent_explore.budget import EpisodeBudget
+from shared.services.retrieval.agent_tools import Decision, ReadableAddresses, ToolResult
+from shared.services.retrieval.hydration.evidence_compose import group_evidence_units
+from shared.services.retrieval.hydration.row_utils import (
+ iter_connected_target_ids,
+ normalize_chunk_type,
+)
+from shared.services.retrieval.settings import ASSET_CHUNK_TYPES
+
+
+@dataclass(frozen=True)
+class Candidate:
+ handle: str
+ kind: Literal["read", "outline"]
+ document_id: str
+ source_file_name: str
+ section_path: str | None = None
+ chunk_type: str | None = None
+ chunk_ids: tuple[str, ...] = ()
+ summary: str = ""
+ extra_sections: int = 0
+ outline_lines: tuple[str, ...] = ()
+
+
+@dataclass
+class PickOutcome:
+ added: list[str] = field(default_factory=list)
+ rejected: list[dict[str, str]] = field(default_factory=list)
+
+
+class EvidencePool:
+ def __init__(self) -> None:
+ self.entries: list[Candidate] = []
+ self.round_index = 0
+ self.pick_phase = False
+ self._read_calls = 0
+ self._outline_calls = 0
+ self._chunk_keys: set[tuple[str, str]] = set()
+ self._seen_active: set[tuple[str, str]] = set()
+ self._seen_staged: set[tuple[str, str]] = set()
+ self._batch: dict[str, Candidate] = {}
+ self._batch_staged: dict[str, Candidate] = {}
+ self._batch_round = 0
+ self._decided: dict[tuple[str, str], Decision] = {}
+ self._queried_tables: dict[tuple[str, str], str] = {}
+
+ def begin_round(self) -> None:
+ self.round_index += 1
+ self._seen_active |= self._seen_staged
+ self._seen_staged = set()
+ if self._batch_staged:
+ self._batch.update(self._batch_staged)
+ self._batch_staged = {}
+ self._batch_round = self.round_index - 1
+ self.pick_phase = bool(self._batch)
+
+ def issue(self, tool_name: str, result: ToolResult) -> list[Candidate]:
+ if result.error:
+ return []
+ self._collect_seen(result)
+ self._record_queried_table(tool_name, result)
+ if tool_name == "corpus.read":
+ return self._issue_read(result)
+ if tool_name == "corpus.outline":
+ return self._issue_outline(result)
+ return []
+
+ def apply_pick(self, handles: list[str]) -> PickOutcome:
+ outcome = PickOutcome()
+ seen: set[str] = set()
+ for handle in handles:
+ if handle in seen:
+ continue
+ seen.add(handle)
+ candidate = self._batch.get(handle)
+ if candidate is None:
+ outcome.rejected.append({"handle": handle, "reason": "unknown id"})
+ continue
+ if candidate.kind == "read" and self._all_chunks_in_pool(candidate):
+ outcome.rejected.append({"handle": handle, "reason": "already in pool"})
+ continue
+ self.entries.append(candidate)
+ for chunk_id in candidate.chunk_ids:
+ self._chunk_keys.add((candidate.document_id, chunk_id))
+ outcome.added.append(handle)
+ picked = set(outcome.added)
+ for candidate in self._batch.values():
+ if candidate.kind != "read":
+ continue
+ picked_handle = candidate.handle if candidate.handle in picked else None
+ for chunk_id in candidate.chunk_ids:
+ self._decided[(candidate.document_id, chunk_id)] = Decision(
+ read_round=self._batch_round,
+ picked_handle=picked_handle,
+ handle=candidate.handle,
+ )
+ self._batch = {}
+ self.pick_phase = False
+ return outcome
+
+ def readable(self) -> ReadableAddresses:
+ return frozenset(self._seen_active)
+
+ def decided(self) -> Mapping[tuple[str, str], Decision]:
+ return MappingProxyType(dict(self._decided))
+
+ def queried_tables(self) -> Mapping[tuple[str, str], str]:
+ return MappingProxyType(self._queried_tables)
+
+ def render_candidates(self, candidates: list[Candidate]) -> str:
+ if not candidates:
+ return ""
+ if len(candidates) == 1 and candidates[0].kind == "outline":
+ return f"[pick id for this result: {candidates[0].handle} = this outline]"
+ parts: list[str] = []
+ for candidate in candidates:
+ label = candidate.section_path or candidate.handle
+ parts.append(f"{candidate.handle} = {label}")
+ return f"[pick ids for this result: {', '.join(parts)}]"
+
+ def render_pool(self) -> str:
+ if not self.entries:
+ return "evidence pool: empty"
+ lines = [f"evidence pool ({len(self.entries)}):"]
+ for candidate in self.entries:
+ if candidate.kind == "outline":
+ lines.extend(candidate.outline_lines)
+ continue
+ path = candidate.section_path or ""
+ kind = candidate.chunk_type or "text"
+ lines.append(
+ f" [{candidate.handle}] {candidate.source_file_name} / {path} [{kind}]"
+ )
+ if candidate.summary:
+ lines.append(f" summary: {candidate.summary}")
+ return "\n".join(lines)
+
+ def chunk_refs(self) -> list[dict[str, str]]:
+ return pool_chunk_refs(self.entries)
+
+ def _all_chunks_in_pool(self, candidate: Candidate) -> bool:
+ if not candidate.chunk_ids:
+ return False
+ return all(
+ (candidate.document_id, chunk_id) in self._chunk_keys
+ for chunk_id in candidate.chunk_ids
+ )
+
+ def _record_queried_table(self, tool_name: str, result: ToolResult) -> None:
+ if tool_name != "corpus.query_table":
+ return
+ document_id = str(result.payload.get("document_id") or "").strip()
+ chunk_id = str(result.payload.get("chunk_id") or "").strip()
+ table_html = result.payload.get("table_html")
+ if not document_id or not chunk_id or not isinstance(table_html, str):
+ return
+ self._queried_tables[(document_id, chunk_id)] = table_html
+
+ def _collect_seen(self, result: ToolResult) -> None:
+ for key in ("rows", "chunks"):
+ rows = result.payload.get(key)
+ if not isinstance(rows, list):
+ continue
+ for row in rows:
+ if not isinstance(row, dict):
+ continue
+ document_id = str(row.get("document_id") or "")
+ if not document_id:
+ continue
+ section_path = str(row.get("section_path") or "").strip()
+ if section_path:
+ self._seen_staged.add((document_id, section_path))
+ chunk_ids = [
+ row.get("chunk_id"),
+ *(row.get("mounted_chunk_ids") or []),
+ *iter_connected_target_ids(row),
+ ]
+ for chunk_id in chunk_ids:
+ if chunk_id:
+ self._seen_staged.add((document_id, str(chunk_id)))
+
+ def _issue_read(self, result: ToolResult) -> list[Candidate]:
+ self._read_calls += 1
+ statuses = result.payload.get("refs")
+ chunks = result.payload.get("chunks")
+ if not isinstance(statuses, list):
+ return []
+ chunk_rows = chunks if isinstance(chunks, list) else []
+ rows_by_id = {
+ str(row.get("chunk_id") or ""): row
+ for row in chunk_rows
+ if isinstance(row, dict) and row.get("chunk_id")
+ }
+ issued: list[Candidate] = []
+ for index, entry in enumerate(statuses, start=1):
+ if not isinstance(entry, dict) or entry.get("status") != "ok":
+ continue
+ chunk_ids = tuple(
+ str(item) for item in (entry.get("chunk_ids") or []) if str(item)
+ )
+ first = rows_by_id.get(chunk_ids[0]) if chunk_ids else None
+ document_id = str(entry.get("document_id") or "")
+ section_path = str(entry.get("section_path") or "").strip() or None
+ if first and not section_path:
+ section_path = str(first.get("section_path") or "").strip() or None
+ chunk_type = normalize_chunk_type((first or {}).get("chunk_type"))
+ summary = _read_summary(first, chunk_type)
+ candidate = Candidate(
+ handle=f"R{self._read_calls}.{index}",
+ kind="read",
+ document_id=document_id,
+ source_file_name=str((first or {}).get("source_file_name") or ""),
+ section_path=section_path,
+ chunk_type=chunk_type or None,
+ chunk_ids=chunk_ids,
+ summary=summary,
+ extra_sections=max(len(chunk_ids) - 1, 0),
+ )
+ issued.append(candidate)
+ self._batch_staged[candidate.handle] = candidate
+ return issued
+
+ def _issue_outline(self, result: ToolResult) -> list[Candidate]:
+ self._outline_calls += 1
+ rows = result.payload.get("rows")
+ details = result.payload.get("details")
+ if not isinstance(rows, list) or not rows:
+ return []
+ documents = details.get("documents") if isinstance(details, dict) else None
+ names: Mapping[str, Any] = documents if isinstance(documents, dict) else {}
+ handle = f"O{self._outline_calls}"
+ first = rows[0] if isinstance(rows[0], dict) else {}
+ document_id = str(first.get("document_id") or "")
+ source_file_name = str(names.get(document_id) or first.get("title") or "")
+ lines = _outline_title_tree(rows, names, handle=handle)
+ candidate = Candidate(
+ handle=handle,
+ kind="outline",
+ document_id=document_id,
+ source_file_name=source_file_name,
+ outline_lines=tuple(lines),
+ )
+ self._batch_staged[candidate.handle] = candidate
+ return [candidate]
+
+
+def render_budget(budget: EpisodeBudget) -> str:
+ return f"budget: steps {budget.steps_used}/{budget.max_steps}"
+
+
+def render_pick_outcome(outcome: PickOutcome) -> str:
+ text = f"picked: {', '.join(outcome.added) or 'none'}"
+ if outcome.rejected:
+ rejected = ", ".join(
+ f"{item['handle']} ({item['reason']})" for item in outcome.rejected
+ )
+ text += f"; rejected: {rejected}"
+ return text
+
+
+def render_trace_line(
+ index: int,
+ wire_name: str,
+ args: dict[str, Any],
+ status: str,
+) -> str:
+ encoded = json.dumps(args, ensure_ascii=False, default=str)
+ return f"{index}. {wire_name} {encoded} -> {status}"
+
+
+def pool_chunk_refs(entries: list[Candidate]) -> list[dict[str, str]]:
+ refs: list[dict[str, str]] = []
+ for candidate in entries:
+ if candidate.kind != "read":
+ continue
+ for chunk_id in candidate.chunk_ids:
+ refs.append({"document_id": candidate.document_id, "chunk_id": chunk_id})
+ return refs
+
+
+def compose_pool_evidence(
+ entries: list[Candidate],
+ assembled_rows: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+ rows_by_id = {
+ str(row.get("chunk_id") or ""): row
+ for row in assembled_rows
+ if row.get("chunk_id")
+ }
+ units: list[dict[str, Any]] = []
+ for candidate in entries:
+ if candidate.kind == "outline":
+ text = "\n".join(candidate.outline_lines)
+ if not text:
+ continue
+ units.append(
+ {
+ "document_id": candidate.document_id,
+ "source_file_name": candidate.source_file_name,
+ "section_path": "",
+ "sort_order": None,
+ "parts": [{"type": "text", "text": text}],
+ "kind": "outline",
+ }
+ )
+ continue
+ chunk_rows = sorted(
+ (rows_by_id[chunk_id] for chunk_id in candidate.chunk_ids if chunk_id in rows_by_id),
+ key=lambda row: row["sort_order"],
+ )
+ parts = [part for row in chunk_rows for part in row["composed"]]
+ if not parts:
+ continue
+ units.append(
+ {
+ "document_id": candidate.document_id,
+ "source_file_name": candidate.source_file_name,
+ "section_path": candidate.section_path or "",
+ "sort_order": chunk_rows[0]["sort_order"],
+ "parts": parts,
+ "kind": "read",
+ }
+ )
+ return group_evidence_units(units)
+
+
+def pool_trace_records(entries: list[Candidate]) -> list[dict[str, Any]]:
+ records: list[dict[str, Any]] = []
+ for candidate in entries:
+ records.append(
+ {
+ "handle": candidate.handle,
+ "kind": candidate.kind,
+ "document_id": candidate.document_id,
+ "section_path": candidate.section_path,
+ "chunk_ids": list(candidate.chunk_ids),
+ }
+ )
+ return records
+
+
+def trace_status(result: ToolResult) -> str:
+ if result.error:
+ return f"failed: {result.error}"
+ statuses = result.payload.get("refs")
+ if isinstance(statuses, list):
+ failures: list[str] = []
+ for index, entry in enumerate(statuses, start=1):
+ if isinstance(entry, dict) and entry.get("status") == "failed":
+ reason = str(entry.get("reason") or "failed")
+ failures.append(f"ref {index} failed: {reason}")
+ if failures:
+ return "partial: " + "; ".join(failures)
+ return "ok"
+
+
+def _read_summary(row: dict[str, Any] | None, chunk_type: str) -> str:
+ if not row:
+ return ""
+ if chunk_type in ASSET_CHUNK_TYPES:
+ metadata = row.get("chunk_metadata") or {}
+ if isinstance(metadata, dict):
+ return str(metadata.get("summary") or "").strip()
+ return ""
+ return str(row.get("section_summary") or "").strip()
+
+
+def _outline_title_tree(
+ rows: list[Any], names: Mapping[str, Any], *, handle: str
+) -> list[str]:
+ lines: list[str] = []
+ current_doc = ""
+ first_doc = True
+ for row in rows:
+ if not isinstance(row, dict):
+ continue
+ document_id = str(row.get("document_id") or "")
+ if document_id != current_doc:
+ current_doc = document_id
+ file_name = str(names.get(document_id) or row.get("title") or document_id)
+ if first_doc:
+ lines.append(f" [{handle}] {file_name}")
+ first_doc = False
+ else:
+ lines.append(f" {file_name}")
+ title = str(row.get("title") or "").strip()
+ if title:
+ depth = int(row.get("depth") or 0)
+ lines.append(f"{' ' * (depth + 1)}{title}")
+ return lines
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py
index ba2a3b4bd..591dba905 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/cursor_harness.py
@@ -14,50 +14,35 @@
guarded import in ``_require_cursor_sdk`` raises a clear error at episode
start if ``cursor-sdk`` isn't installed.
-Architectural difference from ``openai_harness.py`` that budget enforcement
-has to work around: this harness does not control the LLM turn loop.
-``agent.send(...)`` + ``await run.wait()`` hands the *entire* multi-turn
-tool-calling loop to the Cursor SDK; the host process only sees (a)
-``execute`` callbacks for each ``corpus.*``/``finish`` tool call — run
-synchronously off the SDK's own thread and bridged back onto this event
-loop via ``asyncio.run_coroutine_threadsafe`` (mirrors the PoC's
-``_dispatch_sync``) — and (b) the terminal ``RunResult`` once ``wait()``
-returns. There is no per-turn hook to inspect budget mid-turn and force
-``tool_choice`` the way ``openai_harness.py`` does. Each budget dimension is
-therefore enforced (or explicitly not) differently here:
-
-- **``max_steps``**: a plain counter (``budget.steps_used``, incremented once
- per ``corpus.*`` dispatch — ``finish`` does not count, it ends the episode
- on its own) checked in ``_dispatch_sync`` *before* dispatching. Once the
- counter reaches ``budget.max_steps``, further ``corpus.*`` calls are not
- forwarded to ``REGISTRY`` at all — the callback returns a fixed
- "budget exhausted, call finish now" string instead. This is a real,
- synchronous cutoff (unlike the two dimensions below): no reliance on
- cancelling the SDK run from a background task.
-- **``wall_clock``**: ``asyncio.wait_for(run.wait(), timeout=budget.wall_clock_seconds)``,
- per the plan. On timeout, best-effort ``await run.cancel()`` (own short
- timeout, so a hung cancel RPC can't hang this call forever) so the
- underlying agent run actually stops server-side instead of merely being
- abandoned by this process, then the episode result is built from whatever
- ``finish``/tool-trajectory refs were captured before the timeout.
-- **``token_limit``**: **not actively enforced mid-run.** Verified by reading
- the installed ``cursor-sdk`` package's source
- (``cursor_sdk._run_base._RunBase.usage``, 2026-09) that cumulative token
- usage is incrementally accumulated from ``SDKUsageMessage`` stream events
- as ``run.wait()`` consumes them, and exposed as a live ``run.usage``
- property — so a mid-run cutoff (e.g. a polling task calling
- ``run.cancel()``) is *plausible*. It was **not** exercised against a real
- ``CURSOR_API_KEY`` run in this change (no key available in the
- implementation environment), so building an active cutoff on top of an
- unverified live-update assumption was judged higher-risk than shipping.
- ``EpisodeBudget.exhausted()``'s ``token_limit`` dimension is therefore left
- unchecked here by design; ``budget.tokens_used`` is only populated
- *post-hoc* from the terminal ``RunResult.usage`` for observability
- (``EpisodeResult.tokens_used``), after the episode has already finished.
- If ``eval-cursor-harness`` confirms ``run.usage`` does update live against
- a real key, promoting this to an active cutoff (mirroring the
- ``max_steps``/``wall_clock`` pattern above) is straightforward follow-up
- work — deliberately not done speculatively here.
+This harness does not control the LLM turn loop. ``agent.send(...)`` +
+``await run.wait()`` hands the multi-turn tool-calling loop to the Cursor
+SDK. The host process only sees ``execute`` callbacks, ``on_delta`` updates
+(driven by ``run.wait()``), and the terminal ``RunResult``.
+
+Rounds: calls the model issues in one turn share one ``model_call_id``,
+carried by the ``tool-call-started`` update. That update and the tool
+callback arrive almost together, so each callback first waits for its own
+tag (at most the remaining wall clock; without it the call fails with
+``round tag missing``, unexecuted and uncounted). A new ``model_call_id``
+begins a new pool round.
+
+Pick phase: the Cursor tool list is fixed when the agent is created, so a
+round after a successful read or outline accepts only ``corpus_pick``; every
+other call (including ``finish``) is rejected at once, unexecuted and
+uncounted. ``corpus_pick`` itself never counts as a step.
+
+Budget dimensions:
+
+- **``max_steps``**: one explore round (shared ``model_call_id``) is one
+ step, matching the OpenAI harness. Checked when a new explore round
+ starts. Once ``budget.steps_used`` reaches ``budget.max_steps``, further
+ explore rounds are not forwarded; the callback returns a fixed exhausted
+ message plus pool and budget. Same-round sibling ``corpus.*`` calls share
+ that one step and still dispatch.
+- **``wall_clock``**: ``asyncio.wait_for(run.wait(), timeout=budget.wall_clock_seconds)``.
+ On timeout, best-effort ``await run.cancel()``.
+- **tokens**: not limited; ``budget.tokens_used`` is only the episode total
+ from the terminal ``RunResult.usage``.
"""
from __future__ import annotations
@@ -78,43 +63,52 @@
FINISH_TOOL_DESCRIPTION,
FINISH_TOOL_NAME,
FINISH_TOOL_SCHEMA,
- LOOP_CONTRACT_SUFFIX,
+ PICK_TOOL_DESCRIPTION,
+ PICK_TOOL_NAME,
+ PICK_TOOL_SCHEMA,
)
from shared.services.retrieval.agent_explore.dispatch import DbFactory, dispatch_tool_call
+from shared.services.retrieval.agent_explore.evidence_pool import (
+ EvidencePool,
+ render_budget,
+ render_pick_outcome,
+)
+from shared.services.retrieval.agent_explore.prompt import AGENT_SYSTEM_PROMPT
from shared.services.retrieval.agent_explore.shared import (
- EVIDENCE_TOOL_NAMES,
- budget_status_line,
cursor_execute_content,
- finish_refs_from_args,
invalid_finish_message,
+ invalid_pick_message,
read_ref_status,
- select_episode_refs,
tool_message_content,
validate_finish_args,
+ validate_pick_args,
wire_safe_tool_name,
)
from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult
-from shared.services.retrieval.agent_tools import (
- REGISTRY,
- ToolBudget,
- ToolResult,
- load_corpus_schema_text,
-)
+from shared.services.retrieval.agent_tools import REGISTRY, ToolBudget, ToolResult
# Grace period for the underlying agent run to actually stop, after a
# best-effort run.cancel() following a wall_clock timeout, before this
-# process gives up waiting for a terminal RunResult and falls back to
-# whatever refs the trajectory already captured. Not tuned against real
-# cancel-RPC latency yet (no CURSOR_API_KEY in the implementation
-# environment) — revisit with eval-cursor-harness data.
+# process gives up waiting for a terminal RunResult.
_CANCEL_GRACE_SECONDS = 30.0
_BUDGET_EXHAUSTED_MESSAGE = (
"error: step budget exhausted for this episode — do not call any more "
- "corpus.* tools; call finish now with whatever refs you already have "
- "(or an empty list plus a notes explanation)."
+ "corpus.* tools; call finish now."
+)
+
+_PICK_PHASE_REJECTION = (
+ "Pick phase: only corpus_pick is accepted now. Pick from R/O ids above, "
+ "or pass an empty list."
+)
+
+_PICK_PHASE_INSTRUCTION = (
+ "Pick phase: you can only call corpus_pick now. Pick the ids worth "
+ "keeping, or pass an empty list."
)
+_ROUND_TAG_MISSING = "round tag missing"
+
def _require_cursor_sdk() -> Any:
try:
@@ -129,6 +123,10 @@ def _require_cursor_sdk() -> Any:
return cursor_sdk
+def _state_tail(pool: EvidencePool, budget: EpisodeBudget) -> str:
+ return "\n\n".join([pool.render_pool(), render_budget(budget)])
+
+
class CursorHarness:
"""``Harness`` implementation over the Cursor SDK local agent."""
@@ -156,42 +154,82 @@ async def run_episode(
tool_budget = ToolBudget()
loop = asyncio.get_running_loop()
- # The Cursor SDK delivers same-turn parallel tool calls as concurrent
- # HTTP requests on separate threads (ThreadingHTTPServer in
- # cursor_sdk._tool_callback) — CORPUS_SCHEMA.md's loop contract
- # explicitly invites the model to call several independent tools in
- # one turn. Without this lock, the check-then-increment on
- # budget.steps_used below is a race: several concurrently-dispatched
- # calls can each read "not exhausted yet" before any of them
- # increments, letting more calls through than max_steps allows.
- budget_lock = threading.Lock()
+ state = threading.Condition()
+ pool = EvidencePool()
+ round_of: dict[str, str] = {}
+ current_round: dict[str, Any] = {"id": None, "pick_phase": False}
steps: list[AgentStep] = []
- trajectory_refs: list[dict[str, Any]] = []
- # None until finish specifies refs. Explicit ``refs: []`` stays [].
- # Omitted refs still fall back via select_episode_refs.
- finish_state: dict[str, Any] = {"refs": None, "notes": ""}
stop_reason = "finished"
+ def _on_delta(update: Any) -> None:
+ if getattr(update, "type", None) != "tool-call-started":
+ return
+ with state:
+ round_of[update.call_id] = update.model_call_id
+ state.notify_all()
+
+ def _enter_round(ctx: Any) -> tuple[bool | None, bool]:
+ """Wait for this call's round tag (caller holds ``state``).
+
+ Returns ``(pick_phase, is_new_round)``. ``pick_phase`` is
+ ``None`` when the tag never arrived.
+ """
+ call_id = ctx.tool_call_id
+ if not state.wait_for(
+ lambda: call_id in round_of, timeout=budget.remaining_seconds()
+ ):
+ return None, False
+ model_call_id = round_of[call_id]
+ if model_call_id != current_round["id"]:
+ pool.begin_round()
+ current_round["id"] = model_call_id
+ current_round["pick_phase"] = pool.pick_phase
+ return current_round["pick_phase"], True
+ return current_round["pick_phase"], False
+
+ def _append_step(**fields: Any) -> None:
+ steps.append(AgentStep(step_index=len(steps), **fields))
+
+ def _reject(tool_name: str, args: dict[str, Any], message: str) -> None:
+ _append_step(
+ tool_name=tool_name,
+ tool_args=args,
+ observation_text=message,
+ error=message,
+ elapsed_ms=0,
+ round_index=pool.round_index,
+ )
+
def _dispatch_sync(
- tool_name: str, args: dict[str, Any]
+ tool_name: str, args: dict[str, Any], ctx: Any
) -> str | list[dict[str, Any]]:
- with budget_lock:
- if budget.steps_used >= budget.max_steps:
- steps.append(
- AgentStep(
- step_index=len(steps),
+ with state:
+ pick_round, new_round = _enter_round(ctx)
+ round_index = pool.round_index
+ if pick_round is None:
+ _reject(tool_name, args, _ROUND_TAG_MISSING)
+ return f"error: {_ROUND_TAG_MISSING}"
+ if pick_round:
+ _reject(tool_name, args, _PICK_PHASE_REJECTION)
+ return _PICK_PHASE_REJECTION
+ if new_round:
+ if budget.steps_used >= budget.max_steps:
+ observation = (
+ _BUDGET_EXHAUSTED_MESSAGE + "\n" + _state_tail(pool, budget)
+ )
+ _append_step(
tool_name=tool_name,
tool_args=args,
- observation_text=_BUDGET_EXHAUSTED_MESSAGE,
+ observation_text=observation,
error="budget_max_steps",
elapsed_ms=0,
- tokens_used_delta=0,
- tokens_used_total=budget.tokens_used,
+ round_index=round_index,
)
- )
- return _BUDGET_EXHAUSTED_MESSAGE
- budget.record_step()
+ return observation
+ budget.record_step()
+ readable = pool.readable()
+ decided = pool.decided()
tool_started = time.perf_counter()
future = asyncio.run_coroutine_threadsafe(
dispatch_tool_call(
@@ -203,6 +241,8 @@ def _dispatch_sync(
document_scope=document_scope,
budget=tool_budget,
query=query,
+ readable=readable,
+ decided=decided,
),
loop,
)
@@ -216,29 +256,28 @@ def _dispatch_sync(
tool_name=tool_name,
max_chars=tool_budget.max_chars,
)
- # Appended per call, unlike openai_harness.py's once-per-turn
- # placement — this harness has no batched-turn concept exposed to
- # the host process (see module docstring): each corpus.* dispatch
- # is the only per-step hook available to surface budget state.
- content_with_budget = content + "\n" + budget_status_line(budget)
- observation = cursor_execute_content(
- tool_result, text=content_with_budget
- )
- steps.append(
- AgentStep(
- step_index=len(steps),
+ with state:
+ candidates = pool.issue(tool_name, tool_result)
+ tail = _state_tail(pool, budget)
+ text = content
+ if candidates:
+ text = (
+ f"{content}\n{pool.render_candidates(candidates)}\n"
+ f"{_PICK_PHASE_INSTRUCTION}"
+ )
+ text = text + "\n\n" + tail
+ observation = cursor_execute_content(tool_result, text=text)
+ with state:
+ _append_step(
tool_name=tool_name,
tool_args=args,
- observation_text=content_with_budget,
+ observation_text=text,
error=tool_result.error,
elapsed_ms=elapsed_ms,
- tokens_used_delta=0,
- tokens_used_total=budget.tokens_used,
+ round_index=round_index,
ref_status=read_ref_status(tool_name, tool_result),
+ candidates=[item.handle for item in candidates] or None,
)
- )
- if tool_name in EVIDENCE_TOOL_NAMES and not tool_result.error:
- trajectory_refs.extend(tool_result.refs)
return observation
custom_tools: dict[str, Any] = {}
@@ -247,9 +286,9 @@ def _dispatch_sync(
def _make_execute(resolved_name: str):
def execute(
- args: dict[str, Any], _ctx: Any
+ args: dict[str, Any], ctx: Any
) -> str | list[dict[str, Any]]:
- return _dispatch_sync(resolved_name, dict(args or {}))
+ return _dispatch_sync(resolved_name, dict(args or {}), ctx)
return execute
@@ -259,45 +298,65 @@ def execute(
input_schema=spec.json_schema,
)
- def finish_execute(args: dict[str, Any], _ctx: Any) -> str:
+ def pick_execute(args: dict[str, Any], ctx: Any) -> str:
args = dict(args or {})
- validation_error = validate_finish_args(args)
- if validation_error is not None:
- with budget_lock:
- steps.append(
- AgentStep(
- step_index=len(steps),
- tool_name=FINISH_TOOL_NAME,
- tool_args=args,
- observation_text=validation_error,
- error=validation_error,
- elapsed_ms=0,
- tokens_used_delta=0,
- tokens_used_total=budget.tokens_used,
- )
- )
- return json.dumps(
- {"status": "error", "error": invalid_finish_message(validation_error)}
+ with state:
+ pick_round, _new_round = _enter_round(ctx)
+ if pick_round is None:
+ _reject(PICK_TOOL_NAME, args, _ROUND_TAG_MISSING)
+ return f"error: {_ROUND_TAG_MISSING}"
+ validation_error = validate_pick_args(args)
+ if validation_error is not None:
+ message = invalid_pick_message(validation_error)
+ _reject(PICK_TOOL_NAME, args, message)
+ return f"error: {message}"
+ outcome = pool.apply_pick(list(args["pick"]))
+ observation = render_pick_outcome(outcome)
+ _append_step(
+ tool_name=PICK_TOOL_NAME,
+ tool_args=args,
+ observation_text=observation,
+ error=None,
+ elapsed_ms=0,
+ round_index=pool.round_index,
+ picked=outcome.added,
+ pick_rejected=outcome.rejected,
)
- selected = finish_refs_from_args(args)
- cited = selected if selected is not None else []
- notes = str(args.get("notes") or "")
- with budget_lock:
- finish_state["refs"] = selected
- finish_state["notes"] = notes
- steps.append(
- AgentStep(
- step_index=len(steps),
- tool_name=FINISH_TOOL_NAME,
- tool_args=args,
- observation_text=f"refs={len(cited)} notes={notes!r}",
- error=None,
- elapsed_ms=0,
- tokens_used_delta=0,
- tokens_used_total=budget.tokens_used,
+ return observation + "\n\n" + _state_tail(pool, budget)
+
+ custom_tools[wire_safe_tool_name(PICK_TOOL_NAME)] = cursor_sdk.CustomTool(
+ execute=pick_execute,
+ description=PICK_TOOL_DESCRIPTION,
+ input_schema=PICK_TOOL_SCHEMA,
+ )
+
+ def finish_execute(args: dict[str, Any], ctx: Any) -> str:
+ args = dict(args or {})
+ with state:
+ pick_round, _new_round = _enter_round(ctx)
+ if pick_round is None or pick_round:
+ message = _ROUND_TAG_MISSING if pick_round is None else _PICK_PHASE_REJECTION
+ _reject(FINISH_TOOL_NAME, args, message)
+ return json.dumps({"status": "error", "error": message})
+ validation_error = validate_finish_args(args)
+ if validation_error is not None:
+ _reject(FINISH_TOOL_NAME, args, validation_error)
+ return json.dumps(
+ {
+ "status": "error",
+ "error": invalid_finish_message(validation_error),
+ }
)
+ notes = str(args.get("notes") or "")
+ _append_step(
+ tool_name=FINISH_TOOL_NAME,
+ tool_args=args,
+ observation_text=f"pool={len(pool.entries)} notes={notes!r}",
+ error=None,
+ elapsed_ms=0,
+ round_index=pool.round_index,
)
- return json.dumps({"status": "finished", "refs": len(cited)})
+ return json.dumps({"status": "finished", "pool": len(pool.entries)})
custom_tools[FINISH_TOOL_NAME] = cursor_sdk.CustomTool(
execute=finish_execute,
@@ -305,13 +364,7 @@ def finish_execute(args: dict[str, Any], _ctx: Any) -> str:
input_schema=FINISH_TOOL_SCHEMA,
)
- system_prompt = load_corpus_schema_text() + LOOP_CONTRACT_SUFFIX
- user_prompt = (
- f"{system_prompt}\n\n---\n\nUser query:\n{query}\n\n"
- "Tool names on the wire use underscores "
- f"({', '.join(sorted(custom_tools))}). Explore with those tools, "
- "then call finish with cited refs."
- )
+ user_prompt = f"{AGENT_SYSTEM_PROMPT}\n\n---\n\nUser query:\n{query}"
result: Any = None
async with await cursor_sdk.AsyncClient.launch_bridge(
@@ -327,7 +380,9 @@ def finish_execute(args: dict[str, Any], _ctx: Any) -> str:
),
)
) as agent:
- run = await agent.send(user_prompt)
+ run = await agent.send(
+ user_prompt, cursor_sdk.SendOptions(on_delta=_on_delta)
+ )
try:
result = await asyncio.wait_for(
run.wait(), timeout=budget.wall_clock_seconds
@@ -346,18 +401,11 @@ def finish_execute(args: dict[str, Any], _ctx: Any) -> str:
result_usage_total_tokens = result.usage.total_tokens
budget.record_usage({"total_tokens": result_usage_total_tokens})
- selection = select_episode_refs(
- finish_state["refs"],
- trajectory_refs,
- str(finish_state["notes"] or ""),
- )
return EpisodeResult(
- refs=selection.refs,
- notes=selection.notes,
+ pool=list(pool.entries),
steps=steps,
stop_reason=stop_reason,
tokens_used=budget.tokens_used,
model_name=self._model,
- agent_selected_refs=selection.agent_selected_refs,
- fallback_refs=selection.fallback_refs,
+ queried_tables=dict(pool.queried_tables()),
)
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py
index 1d84fdabd..1db56ac54 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/harness/openai_harness.py
@@ -7,45 +7,19 @@
(2026-09-08) against ``deepseek-v4-flash`` via this codebase's client:
single tool call, parallel tool calls in one turn, tool-result feedback +
final synthesis, and forced ``tool_choice`` (used for the budget-exhaustion
-cutoff below) all work.
+cutoff and the pick phase below) all work.
+
+Each turn is one round and sends only a system message plus one rebuilt
+user message (query, one-line call trace, latest-turn results with pick ids,
+evidence pool, steps). Native tool-call history is not kept.
+
+A turn after a successful read or outline is a pick phase: only
+``corpus_pick`` is offered and forced, and the turn does not count as a
+step. Tokens are only summed for the episode total.
Tool calls within one turn are dispatched sequentially through
-``dispatch.dispatch_tool_call`` (fresh DB session per call — safe for any
-harness, not just this one), not via ``asyncio.gather``: batching several
-tool calls into one LLM turn already removes the LLM round-trip per tool
-(the dominant cost); true DB-level concurrency within a turn is not
-implemented here since this provider's tool calls are handled one at a time
-by design, not because concurrent dispatch would be unsafe (it no longer is
-— see ``dispatch.py``).
-
-No import from archived map-nav modules.
-
-Two Phase 4 fixes (audited live against the eval fixture in
-``apps/worker/scripts/fixtures/changheba_archive_eval_queries.json``), both
-now shared with any other harness via ``shared.py`` except stale-message
-collapsing (fix 1), which stays here — it mutates this harness's own
-``messages: list[dict]`` history, a mechanism the Cursor SDK harness has no
-equivalent hook for:
-
-1. **Stale discovery-message collapsing** (``_TOOL_MESSAGE_FRESH_TURNS``,
- ``_collapse_stale_tool_messages``): ``messages`` only ever appended, so a
- single ``corpus.outline``/``corpus.node_filter`` call (each capped at
- ``ToolBudget.max_chars`` — currently ``EVIDENCE_TEXT_CHAR_BUDGET=12_000``,
- see ``registry.py``) was resent in full on every later turn. Verified
- live: two independent queries (q04, q06 in the eval fixture) hit
- ``RETRIEVAL_NAV_TOKEN_LIMIT`` (100k default) within 7-8 LLM turns from
- this resend alone, not from query difficulty — per-turn token cost grew
- monotonically (q04: 4.4k -> 4.8k -> 19.8k -> 21.5k -> 24.2k -> 29.6k).
- Successful evidence reads remain visible until finish. Removing them after
- two turns caused a real SpaceX query to alternate between complementary
- sections until its token budget ran out. Per-result text caps and the
- episode budget still apply.
-2. **Trajectory refs fallback** (``shared.select_episode_refs``): verified
- live that ``finish`` can be called with no ``refs`` key at all (raw
- ``function.arguments`` was literally ``'{}'``) even after the model had
- already read clearly relevant sections via ``corpus.read``. Omitted refs
- (``None``) still fall back to ``corpus.read``/``corpus.assets`` trajectory
- refs. An explicit ``refs: []`` is respected and does not fall back.
+``dispatch.dispatch_tool_call`` (fresh DB session per call). No import from
+archived map-nav modules.
"""
from __future__ import annotations
@@ -64,38 +38,38 @@
FINISH_TOOL_DESCRIPTION,
FINISH_TOOL_NAME,
FINISH_TOOL_SCHEMA,
- LOOP_CONTRACT_SUFFIX,
+ PICK_TOOL_DESCRIPTION,
+ PICK_TOOL_NAME,
+ PICK_TOOL_SCHEMA,
)
from shared.services.retrieval.agent_explore.dispatch import DbFactory, dispatch_tool_call
+from shared.services.retrieval.agent_explore.evidence_pool import (
+ EvidencePool,
+ render_budget,
+ render_pick_outcome,
+ render_trace_line,
+ trace_status,
+)
+from shared.services.retrieval.agent_explore.prompt import AGENT_SYSTEM_PROMPT
from shared.services.retrieval.agent_explore.shared import (
- EVIDENCE_TOOL_NAMES,
- budget_status_line,
build_wire_tool_name_map,
- finish_refs_from_args,
- https_image_parts,
invalid_finish_message,
- model_accepts_images,
+ invalid_pick_message,
read_ref_status,
- select_episode_refs,
tool_message_content,
validate_finish_args,
+ validate_pick_args,
wire_safe_tool_name,
)
from shared.services.retrieval.agent_explore.types import AgentStep, EpisodeResult
-from shared.services.retrieval.agent_tools import (
- REGISTRY,
- ToolBudget,
- ToolResult,
- load_corpus_schema_text,
-)
+from shared.services.retrieval.agent_tools import REGISTRY, ToolBudget, ToolResult
+
+_PICK_WIRE_NAME = wire_safe_tool_name(PICK_TOOL_NAME)
-# A discovery or failed tool message is kept in full for the turn it was produced plus this
-# many additional turns, then collapsed to a placeholder — see module
-# docstring point 1. Not tuned against a real recall-vs-token tradeoff yet;
-# 2 was chosen so a result stays fully visible for one full turn after the
-# one it was produced in (enough for the model to act on it immediately),
-# revisit with more Phase 4 data.
-_TOOL_MESSAGE_FRESH_TURNS = 2
+_PICK_PHASE_INSTRUCTION = (
+ "Pick phase: you can only call corpus_pick now. Pick the ids worth "
+ "keeping, or pass an empty list."
+)
def _resolve_client_and_model() -> tuple[Any, str]:
@@ -141,6 +115,17 @@ def _build_openai_tools() -> tuple[list[dict[str, Any]], dict[str, str]]:
return tools, name_map
+def _build_pick_tool() -> dict[str, Any]:
+ return {
+ "type": "function",
+ "function": {
+ "name": _PICK_WIRE_NAME,
+ "description": PICK_TOOL_DESCRIPTION,
+ "parameters": PICK_TOOL_SCHEMA,
+ },
+ }
+
+
def _parse_tool_arguments(raw: str | None) -> tuple[dict[str, Any], str | None]:
"""Parse one tool call's JSON arguments, or return the parse error.
@@ -160,31 +145,27 @@ def _parse_tool_arguments(raw: str | None) -> tuple[dict[str, Any], str | None]:
return parsed, None
-def _collapse_stale_tool_messages(
- messages: list[dict[str, Any]],
- tool_message_log: list[dict[str, Any]],
+def _compose_user_content(
*,
- current_turn: int,
- fresh_turns: int,
-) -> None:
- """Replace stale discovery/error messages while preserving read evidence.
-
- ``messages`` only ever grows within one episode (see module docstring
- point 1); this is what keeps that growth bounded instead of resending
- every past discovery result on every later turn. Successful evidence must
- remain visible so complementary reads can be synthesized without rereading.
- """
- for entry in tool_message_log:
- if entry["collapsed"] or entry.get("has_evidence", False):
- continue
- if current_turn - entry["turn_index"] < fresh_turns:
- continue
- messages[entry["message_index"]]["content"] = (
- f"[collapsed: {entry['tool_name']} result from turn "
- f"{entry['turn_index']} was {entry['original_chars']} chars — "
- "call the tool again if you need it back in view]"
+ query: str,
+ trace_lines: list[str],
+ latest_results: list[str],
+ latest_turn: int,
+ pool: EvidencePool,
+ budget: EpisodeBudget,
+ instruction: str,
+) -> str:
+ parts = [f"User query: {query}"]
+ if trace_lines:
+ parts.append("trace:\n" + "\n".join(f" {line}" for line in trace_lines))
+ if latest_results:
+ parts.append(
+ f"latest results (turn {latest_turn}):\n" + "\n\n".join(latest_results)
)
- entry["collapsed"] = True
+ parts.append(pool.render_pool())
+ parts.append(render_budget(budget))
+ parts.append(instruction)
+ return "\n\n".join(part for part in parts if part)
class OpenAIHarness:
@@ -203,38 +184,44 @@ async def run_episode(
tool_budget = ToolBudget()
client, model = _resolve_client_and_model()
openai_tools, tool_name_map = _build_openai_tools()
-
- system_prompt = load_corpus_schema_text() + LOOP_CONTRACT_SUFFIX
- messages: list[dict[str, Any]] = [
- {"role": "system", "content": system_prompt},
- {"role": "user", "content": query},
- ]
-
+ pick_tool = _build_pick_tool()
+ pool = EvidencePool()
+ trace_lines: list[str] = []
+ latest_results: list[str] = []
+ latest_turn = 0
steps: list[AgentStep] = []
- finish_refs: list[dict[str, Any]] | None = None
result_notes = ""
- # Refs from every corpus.read/corpus.assets/corpus.query_table call this episode, in call
- # order — the fallback source when finish's own refs end up empty (see
- # module docstring point 2).
- trajectory_refs: list[dict[str, Any]] = []
- # One entry per appended tool-role message: {message_index, turn_index,
- # tool_name, original_chars, has_evidence, collapsed} — see _collapse_stale_tool_messages.
- tool_message_log: list[dict[str, Any]] = []
turn_index = 0
+ stop_reason = "finished"
while True:
turn_index += 1
- _collapse_stale_tool_messages(
- messages,
- tool_message_log,
- current_turn=turn_index,
- fresh_turns=_TOOL_MESSAGE_FRESH_TURNS,
- )
-
- forced_reason = budget.exhausted()
- tool_choice: Any = "auto"
- if forced_reason is not None:
- tool_choice = {"type": "function", "function": {"name": FINISH_TOOL_NAME}}
+ pool.begin_round()
+ pick_turn = pool.pick_phase
+ forced_reason = None if pick_turn else budget.exhausted()
+ if pick_turn:
+ tools = [pick_tool]
+ tool_choice: Any = {"type": "function", "function": {"name": _PICK_WIRE_NAME}}
+ else:
+ tools = openai_tools
+ tool_choice = "auto"
+ if forced_reason is not None:
+ tool_choice = {"type": "function", "function": {"name": FINISH_TOOL_NAME}}
+ messages = [
+ {"role": "system", "content": AGENT_SYSTEM_PROMPT},
+ {
+ "role": "user",
+ "content": _compose_user_content(
+ query=query,
+ trace_lines=trace_lines,
+ latest_results=latest_results,
+ latest_turn=latest_turn,
+ pool=pool,
+ budget=budget,
+ instruction=_PICK_PHASE_INSTRUCTION if pick_turn else "",
+ ),
+ },
+ ]
turn_started = time.perf_counter()
response, usage = await asyncio.to_thread(
@@ -243,13 +230,13 @@ async def run_episode(
model=model,
temperature=0.0,
max_tokens=AGENT_EXPLORE_MAX_COMPLETION_TOKENS,
- tools=openai_tools,
+ tools=tools,
tool_choice=tool_choice,
)
budget.record_usage(usage)
- budget.record_step()
+ if not pick_turn:
+ budget.record_step()
turn_elapsed_ms = int((time.perf_counter() - turn_started) * 1000)
- turn_tokens = int((usage or {}).get("total_tokens", 0) or 0)
message = response.choices[0].message
tool_calls = list(message.tool_calls or [])
@@ -265,47 +252,102 @@ async def run_episode(
observation_text=result_notes,
error=None,
elapsed_ms=turn_elapsed_ms,
- tokens_used_delta=turn_tokens,
- tokens_used_total=budget.tokens_used,
+ round_index=pool.round_index,
)
)
break
- finish_call = next(
- (tc for tc in tool_calls if tc.function.name == FINISH_TOOL_NAME), None
+ if pick_turn:
+ for call_index, tool_call in enumerate(tool_calls):
+ args, parse_error = _parse_tool_arguments(tool_call.function.arguments)
+ pick_error = parse_error or validate_pick_args(args)
+ if pick_error is not None:
+ message = invalid_pick_message(pick_error)
+ trace_lines.append(
+ render_trace_line(
+ len(trace_lines) + 1,
+ _PICK_WIRE_NAME,
+ args,
+ trace_status(ToolResult(text="", error=message)),
+ )
+ )
+ steps.append(
+ AgentStep(
+ step_index=len(steps),
+ tool_name=PICK_TOOL_NAME,
+ tool_args=args,
+ observation_text=message,
+ error=message,
+ elapsed_ms=turn_elapsed_ms if call_index == 0 else 0,
+ round_index=pool.round_index,
+ )
+ )
+ continue
+ outcome = pool.apply_pick(list(args["pick"]))
+ observation = render_pick_outcome(outcome)
+ trace_lines.append(
+ render_trace_line(
+ len(trace_lines) + 1, _PICK_WIRE_NAME, args, observation
+ )
+ )
+ steps.append(
+ AgentStep(
+ step_index=len(steps),
+ tool_name=PICK_TOOL_NAME,
+ tool_args=args,
+ observation_text=observation,
+ error=None,
+ elapsed_ms=turn_elapsed_ms if call_index == 0 else 0,
+ round_index=pool.round_index,
+ picked=outcome.added,
+ pick_rejected=outcome.rejected,
+ )
+ )
+ continue
+
+ parsed_calls: list[tuple[str, dict[str, Any], str | None]] = [
+ (
+ str(tool_call.function.name or ""),
+ *_parse_tool_arguments(tool_call.function.arguments),
+ )
+ for tool_call in tool_calls
+ ]
+
+ finish_entry = next(
+ (item for item in parsed_calls if item[0] == FINISH_TOOL_NAME),
+ None,
)
finish_error: str | None = None
- if finish_call is not None:
- args, parse_error = _parse_tool_arguments(finish_call.function.arguments)
+ finish_args: dict[str, Any] = {}
+ finish_parse_error: str | None = None
+ if finish_entry is not None:
+ _name, finish_args, finish_parse_error = finish_entry
+ parse_error = finish_parse_error
if parse_error is None:
- parse_error = validate_finish_args(args)
+ parse_error = validate_finish_args(finish_args)
if parse_error is not None and forced_reason is None:
finish_error = parse_error
- if finish_call is not None and finish_error is None:
- finish_refs = finish_refs_from_args(args) if parse_error is None else None
- cited = finish_refs if finish_refs is not None else []
+ if finish_entry is not None and finish_error is None:
result_notes = (
- parse_error if parse_error is not None else str(args.get("notes") or "")
+ finish_parse_error
+ if finish_parse_error is not None
+ else str(finish_args.get("notes") or "")
)
stop_reason = f"budget_{forced_reason}" if forced_reason else "finished"
steps.append(
AgentStep(
step_index=len(steps),
tool_name=FINISH_TOOL_NAME,
- tool_args=args,
- observation_text=f"refs={len(cited)} notes={result_notes!r}",
- error=parse_error,
+ tool_args=finish_args,
+ observation_text=f"pool={len(pool.entries)} notes={result_notes!r}",
+ error=finish_parse_error,
elapsed_ms=turn_elapsed_ms,
- tokens_used_delta=turn_tokens,
- tokens_used_total=budget.tokens_used,
+ round_index=pool.round_index,
)
)
break
if forced_reason is not None:
- # Forced tool_choice=finish but the provider returned a
- # different tool anyway (not observed in verification, but a
- # budget cutoff must never loop past). Stop here regardless.
stop_reason = f"budget_{forced_reason}"
result_notes = str(message.content or "") or (
"budget exhausted; provider did not return finish"
@@ -318,131 +360,85 @@ async def run_episode(
observation_text=result_notes,
error="forced_finish_not_honored",
elapsed_ms=turn_elapsed_ms,
- tokens_used_delta=turn_tokens,
- tokens_used_total=budget.tokens_used,
+ round_index=pool.round_index,
)
)
break
- messages.append(
- {
- "role": "assistant",
- "content": message.content or "",
- "tool_calls": [
- {
- "id": tc.id,
- "type": "function",
- "function": {
- "name": tc.function.name,
- "arguments": tc.function.arguments,
- },
- }
- for tc in tool_calls
- ],
- }
- )
- first_tool_tokens_recorded = False
- for tc in tool_calls:
+ readable = pool.readable()
+ decided = pool.decided()
+ first_step_recorded = False
+ turn_results: list[str] = []
+ for requested_name, call_args, parse_error in parsed_calls:
tool_started = time.perf_counter()
- args, parse_error = _parse_tool_arguments(tc.function.arguments)
- requested_name = str(tc.function.name or "")
canonical_name = tool_name_map.get(requested_name, requested_name)
- if tc is finish_call:
+ if requested_name == FINISH_TOOL_NAME:
tool_result = ToolResult(
text="", error=invalid_finish_message(str(finish_error))
)
+ trace_result = ToolResult(text="", error=str(finish_error))
elif parse_error is not None:
tool_result = ToolResult(text="", error=parse_error)
+ trace_result = tool_result
else:
tool_result = await dispatch_tool_call(
canonical_name,
- args,
+ call_args,
db_factory=db_factory,
user_id=user_id,
namespace=namespace,
document_scope=document_scope,
budget=tool_budget,
query=query,
+ readable=readable,
+ decided=decided,
)
+ trace_result = tool_result
tool_elapsed_ms = int((time.perf_counter() - tool_started) * 1000)
content = tool_message_content(
tool_result,
tool_name=canonical_name,
max_chars=tool_budget.max_chars,
)
- messages.append(
- {"role": "tool", "tool_call_id": tc.id, "content": content}
- )
- tool_message_index = len(messages) - 1
- if model_accepts_images(model):
- image_parts = https_image_parts(tool_result)
- if image_parts:
- messages.append(
- {"role": "user", "content": image_parts}
- )
- tool_message_log.append(
- {
- "message_index": tool_message_index,
- "turn_index": turn_index,
- "tool_name": canonical_name,
- "original_chars": len(content),
- "has_evidence": (
- canonical_name in EVIDENCE_TOOL_NAMES
- and not tool_result.error
- and bool(tool_result.refs)
- ),
- "collapsed": False,
- }
+ candidates = pool.issue(canonical_name, tool_result)
+ if candidates:
+ content = f"{content}\n{pool.render_candidates(candidates)}"
+ turn_results.append(f"{requested_name}:\n{content}")
+ trace_lines.append(
+ render_trace_line(
+ len(trace_lines) + 1,
+ requested_name,
+ call_args,
+ trace_status(trace_result),
+ )
)
- if canonical_name in EVIDENCE_TOOL_NAMES and not tool_result.error:
- trajectory_refs.extend(tool_result.refs)
- # Turn-level token usage is attributed to the first tool step in
- # this turn (the completion that decided all calls in it); the
- # rest are 0 to avoid double-counting the same LLM usage.
steps.append(
AgentStep(
step_index=len(steps),
tool_name=canonical_name,
- tool_args=args,
+ tool_args=call_args,
observation_text=content,
error=tool_result.error,
elapsed_ms=(
tool_elapsed_ms
- if first_tool_tokens_recorded
+ if first_step_recorded
else turn_elapsed_ms + tool_elapsed_ms
),
- tokens_used_delta=0 if first_tool_tokens_recorded else turn_tokens,
- tokens_used_total=budget.tokens_used,
+ round_index=pool.round_index,
ref_status=read_ref_status(canonical_name, tool_result),
+ candidates=[item.handle for item in candidates] or None,
)
)
- first_tool_tokens_recorded = True
+ first_step_recorded = True
- # Appended once per turn, to the last tool message only (not
- # every AgentStep's recorded observation_text above) — the model
- # only needs to see current remaining budget once before its next
- # completion call, not once per parallel tool call in this turn.
- last_tool = next(
- (
- message
- for message in reversed(messages)
- if message.get("role") == "tool"
- ),
- None,
- )
- if last_tool is not None:
- last_tool["content"] = (
- str(last_tool["content"]) + "\n" + budget_status_line(budget)
- )
+ latest_results = turn_results
+ latest_turn = turn_index
- selection = select_episode_refs(finish_refs, trajectory_refs, result_notes)
return EpisodeResult(
- refs=selection.refs,
- notes=selection.notes,
+ pool=list(pool.entries),
steps=steps,
stop_reason=stop_reason,
tokens_used=budget.tokens_used,
model_name=model,
- agent_selected_refs=selection.agent_selected_refs,
- fallback_refs=selection.fallback_refs,
+ queried_tables=dict(pool.queried_tables()),
)
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py b/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py
new file mode 100644
index 000000000..c84f6c17b
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/prompt.py
@@ -0,0 +1,45 @@
+"""Agent-explore system prompt.
+
+Tool-specific usage stays on each tool's own description. This module only
+holds the shared overview plus the exploration-loop contract.
+"""
+
+from __future__ import annotations
+
+from shared.services.retrieval.agent_tools import load_corpus_overview_text
+
+LOOP_RULES = """\
+## How this exploration works
+
+You collect evidence for one user question; you do not write the answer.
+Each turn, either call tools or call finish.
+
+Reading:
+- corpus.read only accepts addresses (document_id plus section_path or
+ chunk_id) that appeared in results from your earlier turns. On your first
+ turn you have no results yet, so do not call corpus.read.
+- Calls in the same turn cannot use each other's results.
+
+Evidence pool: only what you pick reaches the answer writer.
+- A corpus.read result gives each successfully read ref a pick id (R3.1,
+ R3.2, ...). A corpus.outline result gives the whole outline one pick id (O2).
+ Failed refs get no id. Other tools give no ids: to pick an image or table,
+ read it by chunk_id first.
+- After a turn where corpus.read or corpus.outline succeeds, the next turn is
+ a pick phase: you can only call corpus.pick. Pass the ids worth keeping, or
+ an empty list. Picked ids join the evidence pool.
+- Body text you read but did not pick cannot be read again.
+- Each turn you see: a trace of your past calls (tool, arguments, ok or why
+ it failed), the full results of your latest calls, the evidence pool, and
+ the steps used so far.
+
+Call finish when the pool is enough to answer, or when you are sure the
+corpus has no answer. finish ends the exploration at once.
+
+If a call needs another call's result, wait for that result instead of
+issuing both in the same turn.
+
+Tool names on the wire use "_" instead of "." (corpus_read = corpus.read).
+"""
+
+AGENT_SYSTEM_PROMPT = load_corpus_overview_text() + "\n\n" + LOOP_RULES
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py b/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py
deleted file mode 100644
index d244e41ea..000000000
--- a/packages/shared-python/shared/services/retrieval/agent_explore/ref_resolution.py
+++ /dev/null
@@ -1,154 +0,0 @@
-"""Resolve ``finish.refs`` into the chunk_id-bearing shape ``resolve_workflow_references`` requires.
-
-Verified live: ``resolve_workflow_references`` -> ``hydrate_referenced_chunk_rows``
-drops any ref whose lookup key has an empty ``chunk_id`` (``row_utils.
-build_reference_lookup_key`` + the ``ref_keys = [k for k in ref_keys if k[0]
-and k[1]]`` guard in ``hydration/reference.py``) — a ``{document_id,
-section_path}``-only ref silently resolves to zero ``referenced_chunks``,
-which is exactly the shape ``agent_tools.CORPUS_SCHEMA.md``/``corpus.read``
-teaches the agent to cite (``corpus.read``'s rendered ``text`` — the only
-thing the LLM ever sees — shows ``section_path``, never ``chunk_id``;
-``chunk_id`` only appears in its structured ``payload``/``refs``, which the
-LLM does not see). This module closes that gap at the harness boundary
-instead of changing what the agent is taught to cite: for any ref missing
-``chunk_id``, resolve that section's own body chunk — the same "one section,
-one body chunk" lookup ``corpus.read``'s ``section_path`` branch performs
-(``agent_tools/section_path_lookup.py``), including the same suffix fallback
-when the agent cites a path without ancestor prefixes.
-"""
-
-from __future__ import annotations
-
-from shared.services.retrieval.corpus_revision_context import CorpusRevisionContext
-
-from shared.services.retrieval.corpus_storage import CorpusStorage
-
-from shared.services.retrieval.document_scope import DocumentScope
-
-from dataclasses import dataclass, field
-from typing import Any
-
-from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from shared.services.retrieval.agent_tools.section_path_lookup import (
- resolve_section_path_anchor,
-)
-
-_BODY_CHUNK_TYPES = ("text", "page")
-
-
-@dataclass
-class FinishRefResolution:
- resolved: list[dict[str, Any]] = field(default_factory=list)
- dropped: list[dict[str, Any]] = field(default_factory=list)
-
-
-async def resolve_finish_refs(
- db: AsyncSession,
- *,
- user_id: str,
- namespace: str,
- refs: list[dict[str, Any]],
- document_scope: DocumentScope = DocumentScope(),
-) -> FinishRefResolution:
- """Return refs with ``chunk_id`` populated; dropped refs keep a reason."""
- corpusStorage: CorpusStorage = CorpusStorage.resolve_namespace(namespace)
- resolution = FinishRefResolution()
- scoped: list[dict[str, Any]] = []
- for ref in refs:
- document_id = str(ref.get("document_id") or "").strip()
- if not document_id:
- resolution.dropped.append({"ref": ref, "reason": "missing document_id"})
- continue
- if not document_scope.allows(document_id):
- resolution.dropped.append(
- {"ref": ref, "reason": f"document_id out of scope: {document_id}"}
- )
- continue
- scoped.append(ref)
-
- document_ids = {
- str(ref.get("document_id") or "").strip() for ref in scoped if ref.get("document_id")
- }
- if not document_ids:
- return resolution
-
- documents = (
- (
- await db.execute(
- select(corpusStorage.Document)
- .where(corpusStorage.Document.document_id.in_(document_ids))
- .where(corpusStorage.Document.user_id == corpusStorage.resolve_owner(user_id))
- .where(corpusStorage.Document.namespace == namespace)
- .where(corpusStorage.Document.status == "active")
- .where(document_scope.predicate(corpusStorage.Document.document_id))
- )
- )
- .scalars()
- .all()
- )
- revision_by_doc = {
- d.document_id: revision for d in documents if (revision := CorpusRevisionContext.resolve_revision(d.document_id, d.current_job_result_id))
- }
-
- for ref in scoped:
- document_id = str(ref.get("document_id") or "").strip()
- chunk_id = str(ref.get("chunk_id") or "").strip()
- if document_id and chunk_id:
- resolution.resolved.append({"document_id": document_id, "chunk_id": chunk_id})
- continue
-
- section_path = str(ref.get("section_path") or "").strip()
- job_result_id = revision_by_doc.get(document_id)
- if not job_result_id:
- resolution.dropped.append(
- {"ref": ref, "reason": f"unknown document_id: {document_id}"}
- )
- continue
- if not section_path:
- resolution.dropped.append(
- {"ref": ref, "reason": f"ref for {document_id} needs section_path or chunk_id"}
- )
- continue
-
- resolved_path, path_error = await resolve_section_path_anchor(
- db,
- document_id=document_id,
- job_result_id=job_result_id,
- section_path=section_path,
- )
- if path_error or not resolved_path:
- resolution.dropped.append(
- {
- "ref": ref,
- "reason": path_error or f"unknown section_path for {document_id}",
- }
- )
- continue
-
- row = (
- await db.execute(
- select(corpusStorage.DocumentChunk.chunk_id)
- .select_from(corpusStorage.DocumentChunk)
- .join(
- corpusStorage.DocumentSection,
- corpusStorage.DocumentSection.section_id == corpusStorage.DocumentChunk.section_id,
- )
- .where(corpusStorage.DocumentChunk.document_id == document_id)
- .where(corpusStorage.DocumentChunk.job_result_id == job_result_id)
- .where(corpusStorage.DocumentSection.section_path == resolved_path)
- .where(corpusStorage.DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES))
- )
- ).first()
- if row is None:
- resolution.dropped.append(
- {
- "ref": ref,
- "reason": f"no body chunk for {document_id}: {resolved_path}",
- }
- )
- continue
- resolution.resolved.append({"document_id": document_id, "chunk_id": str(row[0])})
-
- return resolution
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/shared.py b/packages/shared-python/shared/services/retrieval/agent_explore/shared.py
index 47fbc60ef..0294b5f92 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/shared.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/shared.py
@@ -5,34 +5,21 @@
and, as a debug-script PoC, had started duplicating and drifting from it
instead of sharing it — see the Phase 3.5 section of
``.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md``.
-
-Deliberately excludes anything that assumes an editable ``messages:
-list[dict]`` conversation history (that's OpenAI-harness-specific — see
-``harness/openai_harness.py``'s ``_collapse_stale_tool_messages``, which is
-NOT here because the Cursor SDK manages its own context with no equivalent
-hook exposed to the host process).
"""
from __future__ import annotations
-from dataclasses import dataclass
from typing import Any
import jsonschema
import jsonschema.validators
-from shared.services.retrieval.agent_explore.budget import EpisodeBudget
-from shared.services.retrieval.agent_explore.config import FINISH_TOOL_SCHEMA
-from shared.services.retrieval.agent_tools import ToolResult
-
-# Tools whose ToolResult.refs point at evidence the agent has actually looked
-# at (full body content), as opposed to candidate/listing refs from
-# list_documents/outline/node_filter/recall/grep — those describe *where
-# things are*, not *what was read*, and would inject unread noise into the
-# trajectory-refs fallback below if included.
-EVIDENCE_TOOL_NAMES = frozenset(
- {"corpus.read", "corpus.assets", "corpus.query_table"}
+from shared.services.retrieval.agent_explore.config import (
+ FINISH_TOOL_SCHEMA,
+ PICK_TOOL_NAME,
+ PICK_TOOL_SCHEMA,
)
+from shared.services.retrieval.agent_tools import ToolResult
# The two map-narrowing tools bound their own text at MAP_TOOL_CHAR_BUDGET by
# folding whole subtrees (agent_tools.map_render) or failing the call, so
@@ -112,6 +99,7 @@ def tool_message_content(result: ToolResult, *, tool_name: str, max_chars: int)
if tool_name in MAP_TOOL_NAMES or len(text) <= max_chars:
return text
omitted = len(text) - max_chars
+ # TODO: return a fragment for oversized corpus.read bodies, the way oversized tables do, instead of this head cut.
return (
text[:max_chars]
+ f"\n...[truncated, {omitted} more chars — narrow the scope "
@@ -119,54 +107,6 @@ def tool_message_content(result: ToolResult, *, tool_name: str, max_chars: int)
)
-@dataclass(frozen=True)
-class EpisodeRefSelection:
- """Final episode refs plus how they were chosen.
-
- ``agent_selected_refs`` is ``None`` only when finish was never called.
- An explicit empty finish stays empty and does not take the trajectory
- fallback.
- """
-
- refs: list[dict[str, Any]]
- notes: str
- agent_selected_refs: list[dict[str, Any]] | None
- fallback_refs: list[dict[str, Any]]
-
-
-def select_episode_refs(
- finish_refs: list[dict[str, Any]] | None,
- trajectory_refs: list[dict[str, Any]],
- notes: str,
-) -> EpisodeRefSelection:
- """Choose episode refs without treating ``[]`` and ``None`` as the same."""
- if finish_refs is None:
- fallback_refs = dedup_refs(trajectory_refs)
- if not fallback_refs:
- return EpisodeRefSelection(
- refs=[],
- notes=notes,
- agent_selected_refs=None,
- fallback_refs=[],
- )
- suffix = (
- "[refs auto-filled from corpus.read/corpus.assets/"
- "corpus.query_table trajectory; finish was not called]"
- )
- return EpisodeRefSelection(
- refs=fallback_refs,
- notes=(notes + " " if notes else "") + suffix,
- agent_selected_refs=None,
- fallback_refs=fallback_refs,
- )
- return EpisodeRefSelection(
- refs=list(finish_refs),
- notes=notes,
- agent_selected_refs=list(finish_refs),
- fallback_refs=[],
- )
-
-
def validate_finish_args(args: dict[str, Any]) -> str | None:
"""Validate ``finish``'s own args against ``FINISH_TOOL_SCHEMA``; ``None`` = valid.
@@ -176,10 +116,26 @@ def validate_finish_args(args: dict[str, Any]) -> str | None:
is unverified (see ``config.py``'s ``FINISH_TOOL_SCHEMA`` docstring
context), so this is the real enforcement point.
"""
+ return _schema_args_error("finish", FINISH_TOOL_SCHEMA, args)
+
+
+def validate_pick_args(args: dict[str, Any]) -> str | None:
+ """Validate ``corpus.pick`` args against ``PICK_TOOL_SCHEMA``; ``None`` = valid.
+
+ The Cursor SDK does not enforce ``input_schema`` (a live run sent
+ ``pick_ids``), so a malformed call must fail here instead of being
+ applied as an empty pick that closes the pick phase.
+ """
+ return _schema_args_error(PICK_TOOL_NAME, PICK_TOOL_SCHEMA, args)
+
+
+def _schema_args_error(
+ tool_name: str, schema: dict[str, object], args: dict[str, Any]
+) -> str | None:
if not isinstance(args, dict):
- return f"finish: arguments must be a JSON object, got {type(args).__name__}"
- validator_cls = jsonschema.validators.validator_for(FINISH_TOOL_SCHEMA)
- validator = validator_cls(FINISH_TOOL_SCHEMA)
+ return f"{tool_name}: arguments must be a JSON object, got {type(args).__name__}"
+ validator_cls = jsonschema.validators.validator_for(schema)
+ validator = validator_cls(schema)
errors = sorted(
validator.iter_errors(args), key=lambda error: [str(p) for p in error.path]
)
@@ -190,7 +146,7 @@ def validate_finish_args(args: dict[str, Any]) -> str | None:
if error.validator == "oneOf" and isinstance(error.schema, dict):
message = str(error.schema.get("description") or message)
location = "".join(f"[{p}]" if isinstance(p, int) else f".{p}" for p in error.path)
- return f"finish: invalid arguments{f' at {location}' if location else ''}: {message}"
+ return f"{tool_name}: invalid arguments{f' at {location}' if location else ''}: {message}"
def read_ref_status(tool_name: str, result: ToolResult) -> list[dict[str, Any]] | None:
@@ -205,65 +161,14 @@ def invalid_finish_message(error: str) -> str:
"""Model-facing text for a rejected ``finish`` call; the episode goes on."""
return (
f"{error}. finish was not accepted and the episode continues. Correct "
- 'form: {"refs": [{"document_id": "...", "section_path": "..."} or '
- '{"document_id": "...", "chunk_id": "..."}], "notes": "..."} — each '
- "ref names exactly one of section_path or chunk_id. Call finish again "
+ 'form: {"notes": "..."} — Call finish again '
"with fixed arguments, or keep exploring with the corpus tools."
)
-def finish_refs_from_args(args: dict[str, Any] | None) -> list[dict[str, Any]] | None:
- """Return cited refs, or ``None`` when finish omitted the ``refs`` key.
-
- ``None`` means the agent did not specify refs (never called finish, or
- called finish with ``{}``). An explicit ``refs: []`` stays an empty list
- and must not be collapsed into ``None``.
- """
- if not isinstance(args, dict) or "refs" not in args or args.get("refs") is None:
- return None
- return normalize_finish_refs(args.get("refs"))
-
-
-def normalize_finish_refs(raw: Any) -> list[dict[str, Any]]:
- """Keep only dict items with a non-empty ``document_id`` from a raw ``finish.refs``."""
- if not isinstance(raw, list):
- return []
- normalized: list[dict[str, Any]] = []
- for item in raw:
- if isinstance(item, dict) and str(item.get("document_id") or "").strip():
- normalized.append(item)
- return normalized
-
-
-def budget_status_line(budget: EpisodeBudget) -> str:
- """One-line remaining-budget summary appended to a tool observation.
-
- Neither harness previously surfaced ``EpisodeBudget``'s own counters
- (``steps_used``/``max_steps``, ``tokens_used``/``token_limit``,
- elapsed/wall_clock) to the model at all — it had no way to tell "I'm on
- step 3 of 12" from "I'm on step 11 of 12", so it could not self-regulate
- when to stop exploring and call ``finish``. This exposes the same
- ``EpisodeBudget.snapshot()`` data already used for the hard cutoff,
- reused as a soft signal the model can read every turn.
- """
- snap = budget.snapshot()
+def invalid_pick_message(error: str) -> str:
+ """Model-facing text for a rejected ``corpus.pick`` call; the pick phase goes on."""
return (
- f"[budget: steps {snap['steps_used']}/{snap['max_steps']}, "
- f"tokens {snap['tokens_used']}/{snap['token_limit']}, "
- f"elapsed {snap['elapsed_seconds']:.0f}s/{snap['wall_clock_seconds']:.0f}s]"
+ f"{error}. Nothing was picked and the pick phase continues. Correct "
+ 'form: {"pick": ["", ...]} or {"pick": []} — call corpus_pick again.'
)
-
-
-def dedup_refs(refs: list[dict[str, Any]]) -> list[dict[str, Any]]:
- """Dedup by ``(document_id, chunk_id)``, keeping first-seen order."""
- seen: set[tuple[str, str]] = set()
- deduped: list[dict[str, Any]] = []
- for ref in refs:
- document_id = str(ref.get("document_id") or "").strip()
- chunk_id = str(ref.get("chunk_id") or "").strip()
- key = (document_id, chunk_id)
- if not document_id or not chunk_id or key in seen:
- continue
- seen.add(key)
- deduped.append(ref)
- return deduped
diff --git a/packages/shared-python/shared/services/retrieval/agent_explore/types.py b/packages/shared-python/shared/services/retrieval/agent_explore/types.py
index 7d16092d7..beb71ad98 100644
--- a/packages/shared-python/shared/services/retrieval/agent_explore/types.py
+++ b/packages/shared-python/shared/services/retrieval/agent_explore/types.py
@@ -5,6 +5,8 @@
from dataclasses import dataclass, field
from typing import Any
+from shared.services.retrieval.agent_explore.evidence_pool import Candidate
+
@dataclass
class AgentStep:
@@ -16,22 +18,23 @@ class AgentStep:
observation_text: str
error: str | None
elapsed_ms: int
- tokens_used_delta: int
- tokens_used_total: int
+ round_index: int
# corpus.read only: one {ref, status, reason?} entry per requested ref,
# kept whole in the trace (observation_text is capped there).
ref_status: list[dict[str, Any]] | None = None
+ candidates: list[str] | None = None
+ # corpus.pick only.
+ picked: list[str] | None = None
+ pick_rejected: list[dict[str, str]] | None = None
@dataclass
class EpisodeResult:
"""Everything ``bridge.py`` / the route need after the episode ends."""
- refs: list[dict[str, Any]]
- notes: str
+ pool: list[Candidate]
steps: list[AgentStep] = field(default_factory=list)
stop_reason: str = "finished"
tokens_used: int = 0
model_name: str = ""
- agent_selected_refs: list[dict[str, Any]] | None = None
- fallback_refs: list[dict[str, Any]] = field(default_factory=list)
+ queried_tables: dict[tuple[str, str], str] = field(default_factory=dict)
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md
new file mode 100644
index 000000000..09fd4de50
--- /dev/null
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_OVERVIEW.md
@@ -0,0 +1,27 @@
+# Knowhere corpus
+
+A namespace holds documents. Each document is a tree of sections that follows
+its headings. A section is addressed by document_id + section_path;
+section_path joins the heading titles from the top level down with " / ".
+
+A section can own one body chunk:
+- text: the section's own text (Word, Markdown, Excel, ...)
+- page: PDF pages belonging to the section, as text plus a page image
+A section with child sections may have no body of its own; its content is in
+its children.
+
+Images and tables are separate chunks, addressed by document_id + chunk_id.
+Result rows show the section that contains them.
+
+Example:
+
+doc_7f3a Hypertension_Guideline.pdf
+├── 1 Overview [text]
+├── 2 Treatment (no body; see children)
+│ ├── 2 Treatment / 2.1 Lifestyle [text]
+│ └── 2 Treatment / 2.2 Drugs [text] contains table chunk_id=tbl_12
+└── 3 Follow-up [page p.14-15]
+
+Addresses from this example:
+- {document_id: "doc_7f3a", section_path: "2 Treatment / 2.2 Drugs"}
+- {document_id: "doc_7f3a", chunk_id: "tbl_12"}
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md b/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md
deleted file mode 100644
index 225848378..000000000
--- a/packages/shared-python/shared/services/retrieval/agent_tools/CORPUS_SCHEMA.md
+++ /dev/null
@@ -1,216 +0,0 @@
-# Knowhere Corpus Schema (agent-facing)
-
-Single source of truth for how an agent should understand and explore a
-Knowhere corpus through the tools in this package. This text is intended to
-be shipped verbatim as both the API `/mcp` server instructions and the
-`agent_explore` system prompt — do not duplicate it elsewhere; edit here.
-
-This describes the **published, DB-served corpus** (`documents`,
-`document_sections`, `document_chunks`, `graph_nodes`, `graph_edges`) that
-the tools below are designed to query. It is not the on-disk parse artifact schema
-(`chunks.json` / `doc_nav.json` on disk use different join keys and a
-different path separator) — that schema is for the parser pipeline, not for
-this tool set.
-
----
-
-## 1. Corpus model
-
-```text
-Namespace
- └─ Document (document_id, source_file_name, parse_track)
- └─ Section (section_id, parent_section_id, section_path, section_level, summary)
- └─ Chunk (chunk_id, chunk_type, content, chunk_metadata)
-```
-
-- **Namespace**: the retrieval scope. One namespace can hold documents parsed
- by different tracks (see §2) — do not assume a namespace is uniform.
-- **Document**: `document_id` is the stable identifier for every other tool
- call. `parse_track` is `page_memory` or `chunk` (see §2). Tool hit lines
- print `document_id` in parentheses after the filename (or alone, for
- `node_filter`); copy that value into later calls. Never pass
- `source_file_name` as `document_id`.
-- **Section**: the navigable tree. `section_path` segments are joined with
- `" / "` (space-slash-space), one segment per heading/synthetic level.
-- **Chunk**: the retrieval unit. `chunk_type` is `text`, `page`, `image`, or
- `table`. A section owns **at most one body chunk** (`text` or `page`) —
- but which sections get one is track-dependent (see §2). `image`/`table`
- chunks are **not** attached to the section they are conceptually "in" —
- see §3.
-
-Document-level relationships (`related` edges, including entity-overlap
-scoring) exist in a separate graph — see §4.
-
-## 2. Body chunks: `text` vs `page`, and the `SAME-AS` pointer
-
-A namespace can mix both tracks across documents (e.g. one PDF on the
-`page_memory` track next to one DOCX/XLSX on the `chunk` track). Treat
-`text` and `page` as the same *role* — both are a section's own body — but
-they differ in shape, and in **which sections get a body chunk at all**:
-
-- **`chunk` track** (`text` chunks, e.g. DOCX/XLSX/MD): a section owns a
- body chunk if it has any content of its own directly under its heading,
- before any child heading — this can be **any section, leaf or not**. Do
- not assume a section with children has no body text of its own; check
- whether it owns a chunk instead of assuming from tree position.
-- **`page_memory` track** (`page` chunks, PDF): only **leaf** sections own a
- body chunk. Internal/structural sections never carry a `page` chunk
- themselves — their summaries aggregate from their leaf descendants.
-
-A section owns **at most one** body chunk either way (verified against
-real published data: a section with children and a section without both
-had exactly 0 or 1, never more). So a "chunk" and "the section that owns
-it" are the same unit, not two different granularities — there is no
-separate finer-or-coarser level to choose between. The only structural
-sections with no chunk of their own are `chunk`-track sections whose
-entire content lives in their descendants.
-
-For `page` chunks specifically: one leaf section's body may span one or
-more physical pages. A page's text is stored **once**, under whichever leaf
-is first in reading order to cover that page (the "owner"). Every other
-leaf section that also covers that physical page has, in place of the
-text, a literal marker:
-
- ```text
- [SAME-AS p]
- ```
-
- This is a pointer, not a preview. If you need that page's actual text,
- resolve the marker by reading the owner section's chunk at that page
- number — do not treat the marker's absence of text as "this section has
- no content there." `page` chunks also carry `page_nums` (all physical
- pages they cover) and `page_assets` (rendered page-citation screenshots —
- these are references for citation, not separate `image` chunks).
-
- **Format trap**: `` inside the marker is written
- verbatim by the parser and stored as-is — it is the on-disk path
- (`"///..."`, plain `/`, filename
- included), **not** the DB `section_path` you get back from
- `corpus.outline` / `corpus.node_filter` / `corpus.recall` (which is
- `" / "`-joined and excludes the filename). Do not string-match the
- marker directly against a DB `section_path`. Convert it first —
- `section_path_from_chunk_path()` in `search/lexical_text.py` already
- does this conversion and is the function to reuse when implementing
- marker resolution, not a new one.
-
-## 3. Asset chunks: `image` / `table`, and `connect_to`
-
-`image` and `table` chunks are **not children of the section they visually
-belong to**. In the DB they are parked under their document's synthetic
-`Root` section. The real association to a body section is the `connect_to`
-list on the **body chunk**, not a location on the asset:
-
-- `relation: "embeds"` — the body chunk that owns/embeds this asset inline.
-- `relation: "related"` — another body chunk that shares the same source
- page as a page-track asset, without owning it (`same_as_owner` may name
- the owning section).
-
-This link is **one-directional** (body → asset). There is no stored
-asset → body back-link. `corpus.grep`, `corpus.recall`, and `corpus.assets`
-resolve the host through that body link and print the hosting
-`section_path` on the asset row. If no host is found, the row keeps
-`Root` and is marked as having no host section. Use `corpus.assets`
-`host_of` to list every host, rather than reading the asset's stored
-path.
-
-## 4. Document graph
-
-Today the graph only has **document-level** nodes (`node_kind='document'`)
-and undirected `related` edges between documents. Edge scoring prefers
-**typed-entity overlap** between the two documents' aggregated `entities`
-first; it only falls back to free-form TF-IDF keyword overlap when either
-document lacks entities. Either way the edge carries which terms matched
-(`properties.shared_entities` or `properties.shared_keywords`). There are
-no section-level or entity-level graph *nodes* — an entity is not itself a
-queryable node, and there is no entity-to-entity or entity-to-chunk edge,
-only this document-to-document rollup. Use `neighbors` for "what else is
-like this document" (and to see which shared terms justify that link), not
-for anything finer-grained than a document pair.
-
-## 5. Reserved / not yet available
-
-- **Vector**: not available yet. `recall` today is path + content BM25.
- Exact-string lookup is `grep`.
-
-## 6. Tools and how they work together
-
-Every tool exists to find the `section_path`s (or asset `chunk_id`s) that
-answer the query, then hand them to `corpus.read`. Names below are the
-registered names (`corpus.outline`). MCP clients call those names as-is.
-Exact call parameters live only in each tool's own schema/description
-(ask for that, do not memorize names here) — this section is the
-collaboration map: which tools narrow a search space, which produce hits,
-and what to do with either kind of result.
-
-```mermaid
-flowchart LR
- subgraph MapNarrowing["Narrow a map (overview / structural predicate)"]
- outline["corpus.outline"]
- node_filter["corpus.node_filter"]
- end
- subgraph LeafHits["Find hits (exact string / fuzzy / asset listing)"]
- grep["corpus.grep"]
- recall["corpus.recall"]
- assets["corpus.assets"]
- end
- outline -->|narrows scope for| LeafHits
- node_filter -->|narrows scope for| LeafHits
- MapNarrowing -->|section_path| read["corpus.read"]
- LeafHits -->|section_path or chunk_id| read
- read -->|table too large| query_table["corpus.query_table"]
- read --> finish
- query_table --> finish
- list_documents["corpus.list_documents (namespace inventory, standalone)"]
- neighbors["corpus.neighbors (related documents, standalone)"]
-```
-
-**`corpus.outline` and `corpus.node_filter`** are the two map-narrowing
-tools: both take one or more scope targets (a whole document, or a section
-and everything under it) and return the *same* row shape — every
-`section_path` in that scope, indented by level, with title/summary.
-`corpus.outline` returns that unconditionally (the outline of what is
-there); `corpus.node_filter` returns it only for the branches around a
-structural predicate match (marked as a hit), still shown as full context
-(ancestors + the entire matched subtree), not a bare list of matches. Use
-a returned `section_path` to `corpus.read` a branch, or as the scope of a
-`corpus.grep` / `corpus.recall` / `corpus.assets` call. Neither is
-callable on a single leaf section with nothing under it; `corpus.read`
-that directly instead.
-
-**`corpus.grep`, `corpus.recall`, and `corpus.assets`** are the leaf-hit
-tools: they search *within* a scope (the whole corpus, or one narrowed by
-a prior `corpus.outline`/`corpus.node_filter` call) and return rows
-pointing at specific sections/chunks, not a map. `corpus.grep` is exact
-string (several terms are any-match / OR, rows stay in document order,
-not relevance order); `corpus.recall` is fuzzy ranked search;
-`corpus.assets` lists image/table chunks by type — on its own it is an
-unfiltered listing with no way to judge relevance, so pair it with a
-prior `corpus.grep`/`corpus.recall` hit (scope to that hit's section, or
-reverse-resolve the hit's own asset references) rather than browsing
-every asset in a document. Grep does not scan table-cell HTML — a table
-only surfaces from `corpus.grep`/`corpus.recall` via its published
-summary/keywords/caption, never by a literal cell value; find the table
-via `corpus.assets` (or a hit on the surrounding body text) and use
-`corpus.read`/`corpus.query_table` to inspect its cells. Every hit row
-names a `document_id` plus either a `section_path` (body hits) or a
-`chunk_id` (image/table hits). Image/table rows show the hosting
-section, not the asset's stored `Root` path — read whichever address
-applies with `corpus.read` next, then decide pick or no-pick for that
-ref immediately from `corpus.read`'s own per-ref status.
-
-**`corpus.read`** is where body content actually gets consumed: it accepts
-many refs at once (mixing map-tool section_paths and leaf-hit chunk_ids),
-and reports each ref's outcome separately so a failed ref does not hide a
-successful one. A table `corpus.read` reports as too large points at
-**`corpus.query_table`** for cell-level `SELECT`s over that same table.
-
-**`corpus.list_documents`** and **`corpus.neighbors`** stand outside this
-flow: `corpus.list_documents` is a namespace-wide inventory (only for an
-explicit "list/inventory the documents" request, never a cold-start for a
-content question); `corpus.neighbors` returns document-level related
-edges (§4), unscoped by section.
-
-General rule: narrow with `corpus.outline`/`corpus.node_filter` before
-searching a large corpus; prefer `corpus.grep` over `corpus.recall` when
-you know the exact string you are looking for, and fall back to
-`corpus.recall` only for genuinely fuzzy questions.
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py b/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py
index 5583243f4..42b2e3573 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/__init__.py
@@ -2,7 +2,7 @@
See ``.cursor/plans/agentic_corpus_explore_retrieval_c2c4ea21.plan.md`` (Phase 2)
for the design. Tools in this package query the published, DB-served corpus
-described in ``CORPUS_SCHEMA.md`` (``documents`` / ``document_sections`` /
+described in ``CORPUS_OVERVIEW.md`` (``documents`` / ``document_sections`` /
``document_chunks`` / ``graph_nodes`` / ``graph_edges``) — not the on-disk
parse artifacts.
@@ -18,6 +18,8 @@
from shared.services.retrieval.agent_tools.registry import (
REGISTRY,
+ Decision,
+ ReadableAddresses,
ToolBudget,
ToolContext,
ToolRegistry,
@@ -26,7 +28,7 @@
capped_limit,
register_tool,
)
-from shared.services.retrieval.agent_tools.schema_doc import load_corpus_schema_text
+from shared.services.retrieval.agent_tools.schema_doc import load_corpus_overview_text
# Register every ``corpus.*`` tool into ``REGISTRY`` as an import side effect:
# each module under ``tools/`` calls ``@register_tool`` at import time. Kept
@@ -40,12 +42,14 @@
__all__ = [
"_tools",
"REGISTRY",
+ "Decision",
+ "ReadableAddresses",
"ToolBudget",
"ToolContext",
"ToolRegistry",
"ToolResult",
"ToolSpec",
"capped_limit",
- "load_corpus_schema_text",
+ "load_corpus_overview_text",
"register_tool",
]
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py b/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py
index 337793b97..55d9ee92e 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/asset_hosts.py
@@ -2,7 +2,7 @@
Image/table chunks are stored under their document's ``Root`` section; the
real placement is ``chunk_metadata.connect_to`` on the body chunk that embeds
-them (``CORPUS_SCHEMA.md`` §3). Hosts are keyed by ``(document_id,
+them. Hosts are keyed by ``(document_id,
asset chunk_id)`` — chunk ids are content hashes and repeat across
documents — and read only from each document's current revision. When an
asset has several hosts they are kept in body order, so every tool that
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py b/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py
index c4cbaa050..b4282ac28 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/explore_mount.py
@@ -1,8 +1,9 @@
-"""Shared explore-phase mounting for ``corpus.grep`` and ``corpus.recall`` hits.
+"""Shared explore-phase mounting for ``corpus.recall`` hits.
Table/image hits are rendered with the same explore table/image functions
``corpus.read`` uses. Body hits that already list ``connect_to`` targets get
those assets rendered and labeled as mounted, not as the body match itself.
+Each returned row lists the asset ids it rendered in ``mounted_chunk_ids``.
"""
from __future__ import annotations
@@ -101,6 +102,7 @@ async def mount_explore_hits(
chunk_id = str(row.get("chunk_id") or "")
chunk_type = normalize_chunk_type(row.get("chunk_type"))
enriched_self = by_id.get(chunk_id, row)
+ mounted_ids: list[str] = []
if chunk_type == "table":
row["rendered"] = render_explore_table(
enriched_self, char_budget=char_budget
@@ -119,11 +121,14 @@ async def mount_explore_hits(
body = render_explore_table(target, char_budget=char_budget)
if body:
parts.append(f"mounted table chunk_id={target_id}:\n{body}")
+ mounted_ids.append(target_id)
elif target_type == "image":
body = _image_display_content(target)
if body:
parts.append(f"mounted image chunk_id={target_id}:\n{body}")
+ mounted_ids.append(target_id)
_append_image_media(target, media, seen_urls)
row["rendered"] = "\n\n".join(parts)
+ row["mounted_chunk_ids"] = mounted_ids
mounted.append(row)
return mounted, media
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py
index 282302f60..762571d61 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/registry.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/registry.py
@@ -14,7 +14,7 @@
from __future__ import annotations
import logging
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Mapping, Set
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass, field, replace
from typing import Any
@@ -73,14 +73,14 @@ class ToolBudget:
below, and notes it in ``ToolResult.text`` when the request was clamped.
Tools that promise a complete, non-truncated set by contract
(``node_filter``, ``outline``) do not apply this budget to their
- matched-set cardinality — see ``CORPUS_SCHEMA.md`` §6.
+ matched-set cardinality.
``max_chars`` caps the rendered ``ToolResult.text`` before it enters LLM
context. Applied in ``agent_explore.shared.tool_message_content`` (not
inside individual tools) so ``read`` can return full body text from the
tool while the harness still bounds what the model sees per turn. Aligned
with final evidence packing via ``EVIDENCE_TEXT_CHAR_BUDGET``
- (12_000).
+ (15_000).
"""
max_chars: int = EVIDENCE_TEXT_CHAR_BUDGET
@@ -96,6 +96,20 @@ def capped_limit(requested: int, budget: ToolBudget) -> int:
return min(requested, budget.max_items)
+# ``(document_id, section_path)`` and ``(document_id, chunk_id)`` addresses
+# an ``agent_explore`` episode received in results of earlier rounds.
+ReadableAddresses = Set[tuple[str, str]]
+
+
+@dataclass(frozen=True)
+class Decision:
+ """One body chunk an ``agent_explore`` episode read and then picked or not."""
+
+ read_round: int
+ picked_handle: str | None
+ handle: str
+
+
@dataclass
class ToolContext:
"""Per-call execution context. One instance is built per tool dispatch."""
@@ -116,6 +130,10 @@ class ToolContext:
# dropping summaries does not fit then fails that one call as "too
# large to map" instead of guessing.
query: str = ""
+ # ``corpus.read`` gates, set only by ``agent_explore``. ``None`` means no
+ # gate: an external MCP client calling one tool has no episode.
+ readable: ReadableAddresses | None = None
+ decided: Mapping[tuple[str, str], Decision] | None = None
@dataclass
@@ -208,7 +226,7 @@ def validate_tool_args(spec: ToolSpec, args: dict[str, Any]) -> str | None:
Every tool schema declares ``additionalProperties: false`` at every
object level (registered tools are expected to keep this true — see
- ``CORPUS_SCHEMA.md`` §6), so an unknown key anywhere in the argument
+ each tool schema), so an unknown key anywhere in the argument
tree fails here instead of being silently dropped or ignored downstream.
The returned message names the offending path, the tool's legal
top-level argument names, and — for an unknown-key error — the exact
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py b/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py
index 37569589f..ffd2983ab 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/schema_doc.py
@@ -1,10 +1,7 @@
-"""Single loader for ``CORPUS_SCHEMA.md`` — the agent-facing corpus schema text.
+"""Loader for ``CORPUS_OVERVIEW.md`` — the agent-facing corpus overview.
-Both harnesses (Phase 3) read through this function instead of the file
-directly, so there is exactly one place that resolves the path: the API
-``/mcp`` server's ``instructions`` and ``agent_explore``'s system prompt must
-stay byte-identical for the shared schema portion (see the module docstring
-at the top of ``CORPUS_SCHEMA.md`` — "do not duplicate it elsewhere").
+The API ``/mcp`` server instructions and ``agent_explore`` system prompt
+share this text. Tool-specific usage lives on each tool's own description.
"""
from __future__ import annotations
@@ -12,10 +9,10 @@
from functools import lru_cache
from pathlib import Path
-_SCHEMA_PATH = Path(__file__).with_name("CORPUS_SCHEMA.md")
+_OVERVIEW_PATH = Path(__file__).with_name("CORPUS_OVERVIEW.md")
@lru_cache(maxsize=1)
-def load_corpus_schema_text() -> str:
- """Return the verbatim contents of ``CORPUS_SCHEMA.md``."""
- return _SCHEMA_PATH.read_text(encoding="utf-8")
+def load_corpus_overview_text() -> str:
+ """Return the verbatim contents of ``CORPUS_OVERVIEW.md``."""
+ return _OVERVIEW_PATH.read_text(encoding="utf-8")
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py b/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py
index f15735181..54858a2b4 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/section_path_lookup.py
@@ -1,6 +1,6 @@
"""Shared section_path resolution for agent tools and harness bridge.
-``corpus.read`` and ``resolve_finish_refs`` both need to turn an agent-supplied
+``corpus.read`` needs to turn an agent-supplied
``section_path`` into one canonical DB path. Agents often cite a suffix (e.g.
``3 工程地质 / 3.2 覆盖层``) while the stored path includes ancestors
(``附件目录 / 3 工程地质 / 3.2 覆盖层``). Exact match alone fails silently
@@ -15,7 +15,29 @@
from sqlalchemy import or_, select
from sqlalchemy.ext.asyncio import AsyncSession
-from shared.services.retrieval.search.lexical_text import normalize_section_path
+from shared.services.retrieval.agent_tools.registry import ReadableAddresses
+from shared.services.retrieval.search.lexical_text import (
+ normalize_section_path,
+ split_section_path,
+)
+
+
+def section_path_received(
+ readable: ReadableAddresses, document_id: str, section_path: str
+) -> bool:
+ """Whether ``section_path`` runs contiguously inside a received path of the document.
+
+ Ancestor and suffix refs pass, matching how ``corpus.read`` resolves paths.
+ """
+ wanted = split_section_path(section_path)
+ size = len(wanted)
+ for seen_document_id, seen_path in readable:
+ if seen_document_id != document_id:
+ continue
+ parts = split_section_path(seen_path)
+ if any(parts[i : i + size] == wanted for i in range(len(parts) - size + 1)):
+ return True
+ return False
def paths_matching_section_ref(normalized: str, candidate_paths: list[str]) -> list[str]:
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py b/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py
index dd7619f90..f5284858b 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/snippet.py
@@ -122,9 +122,11 @@ def format_row(row: Mapping[str, Any]) -> str:
header = f"{indent}- [{row['kind']}]"
if row["title"]:
header += f" {row['title']}"
- header += f" | document_id={row['document_id']} section_path={row['section_path']}"
+ header += f" | document_id={row['document_id']}"
if row["hosted"] is False:
- header += " (no host section)"
+ header += " (not in any section)"
+ else:
+ header += f" section_path={row['section_path']}"
if row["chunk_id"]:
header += f" chunk_id={row['chunk_id']}"
if row["score"] is not None:
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py
index d931b9389..83c8d9b7f 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/assets.py
@@ -1,7 +1,7 @@
"""``corpus.assets`` — forward asset search and reverse asset -> hosts lookup.
Image/table chunks are parked under their document's synthetic ``Root``
-section in the DB (§3 of ``CORPUS_SCHEMA.md``); the real association to a
+section in the DB; the real association to a
body section lives in ``chunk_metadata.connect_to`` on the *body* chunk.
Both directions resolve hosts through ``agent_tools.asset_hosts`` (keyed by
document + chunk, current revision only), the same resolver ``corpus.grep``
@@ -54,16 +54,13 @@
@register_tool(
name="corpus.assets",
description=(
- "Forward search for image/table chunks by type/query, or reverse "
- "lookup: given asset chunk_ids (host_of), find which body "
- "section(s) embed or reference them via connect_to. Rows show the "
- "hosting section_path; an asset no body section embeds keeps Root "
- "and is marked '(no host section)'. A forward search alone only "
- "lists candidate assets — it does not tell you which ones matter "
- "for the current question. Pair it with a prior corpus.recall/"
- "corpus.grep hit: scope this call to that hit's section, or "
- "reverse-resolve the hit's own connect_to targets via host_of "
- "instead of browsing every asset in a document."
+ "Find image and table chunks by type or query, or, given asset "
+ "chunk_ids (host_of), list the sections that contain them. Rows "
+ "show the containing section_path; an asset that no section "
+ "contains is marked (not in any section). A plain listing does not "
+ "tell you which assets matter: scope it to a section from a prior "
+ "corpus.grep/corpus.recall hit, or use host_of on chunk_ids you "
+ "already have."
),
json_schema={
"type": "object",
@@ -243,7 +240,7 @@ async def _reverse_lookup(
lines = [format_row(row) for row in rows]
lines.extend(
- f"- {target_id}: no host section found in this scope (asset stays under Root)"
+ f"- {target_id}: not contained in any section within this scope"
for target_id in unhosted
)
return ToolResult(
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py
index 8789ddde9..2815cc0cc 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/grep.py
@@ -4,12 +4,13 @@
current revision. That expression is the existing term trigram index.
The field is already published (body + filename + section path; tables
use summary/keywords/caption; images use their description).
-Reports a total match count (over the full in-scope corpus, not just the
-returned page) alongside capped snippets.
-Matching table/image chunks, and body chunks that list ``connect_to``
-targets, are mounted through the shared explore mount so the tool result
-includes rendered table/image content.
+Several terms are matched in one query (OR). Hits are then split back
+into one block per term, one row per matching section (first chunk in
+document / sort order). Terms with fewer matching sections are listed
+first so a rare term is not buried by a generic one. Visible rows stop
+at the existing ``limit``; leftover sections become one fold line per
+document with a reusable ``scope``.
Snippets are built by the shared ``agent_tools.snippet.build_snippet`` (head
+ first-match window + tail, ``...``-joined, overlap-merged). Rows render
@@ -25,13 +26,12 @@
from shared.services.retrieval.corpus_storage import CorpusStorage
import re
+from collections import defaultdict
from typing import Any
from sqlalchemy import func, or_, select
-from shared.models.database.job_result import JobResult
from shared.services.retrieval.agent_tools.asset_hosts import host_paths_for_hits
-from shared.services.retrieval.agent_tools.explore_mount import mount_explore_hits
from shared.services.retrieval.agent_tools.registry import (
ToolContext,
ToolResult,
@@ -91,23 +91,78 @@ def _term_search(terms: list[str], namespace: str = "default") -> tuple[re.Patte
)
+def _term_span(text: str, term: str) -> tuple[int, int] | None:
+ match = re.search(re.escape(term), text, flags=re.IGNORECASE)
+ return match.span() if match else None
+
+
+def _sections_for_term(
+ hits: list[dict[str, Any]], term: str, context_chars: int
+) -> list[dict[str, Any]]:
+ """Keep the first hit per ``(document_id, section_path)`` for this term.
+
+ ``hits`` must already be in document / sort order.
+ """
+ kept: dict[tuple[str, str], dict[str, Any]] = {}
+ order: list[tuple[str, str]] = []
+ for hit in hits:
+ text = str(hit["term_search_text"] or "")
+ span = _term_span(text, term)
+ if span is None:
+ continue
+ key = (str(hit["document_id"]), str(hit["section_path"] or ""))
+ if key in kept:
+ continue
+ kept[key] = {**hit, "snippet": build_snippet(text, span, hit_context=context_chars)}
+ order.append(key)
+ return [kept[key] for key in order]
+
+
+def _fold_line(source_file_name: str, document_id: str, leftover: int) -> str:
+ return (
+ f"{source_file_name} 还有 {leftover} 节命中 — "
+ f"scope=[{{document_id: {document_id}}}]"
+ )
+
+
+def _build_hit_row(
+ hit: dict[str, Any],
+ host_paths: dict[tuple[str, str], tuple[str, bool]],
+) -> dict[str, Any]:
+ chunk_type = str(hit["chunk_type"] or "").strip()
+ is_asset = chunk_type in ASSET_CHUNK_TYPES
+ key = (str(hit["document_id"]), str(hit["chunk_id"]))
+ section_path, hosted = (
+ host_paths[key] if is_asset and key in host_paths else (hit["section_path"], None)
+ )
+ return build_row(
+ kind=chunk_type or "text",
+ document_id=hit["document_id"],
+ section_path=section_path,
+ title=hit["source_file_name"],
+ chunk_id=hit["chunk_id"] if is_asset else None,
+ snippet=hit["snippet"],
+ hosted=hosted if is_asset else None,
+ )
+
+
@register_tool(
name="corpus.grep",
description=(
- "Exact string search against published term_search_text "
- "(body text, image descriptions, table summaries/keywords, plus "
- "filename and section path). Returns the total number of matching "
- "chunks plus a capped list of snippet rows (score-free — this is an "
- "exact match, not a ranked search; use corpus.recall for ranking). "
- "Table and image hits include rendered content; body hits include "
- "any connected table or image — read the hit's document_id + "
- "section_path with corpus.read next. Provide 'pattern' for one "
- "term, or 'patterns' for several candidate terms OR'd together in "
- "this single call (e.g. synonyms) — issue one call with multiple "
- "terms instead of several parallel corpus.grep calls for different "
- "terms in the same turn. Several terms are any-match (OR), and "
- "rows stay in document order — this tool does not rank by "
- "relevance. At least one of pattern/patterns is required."
+ "Exact string search over body text, image descriptions, table "
+ "summaries and keywords, file names and section paths. Table cell "
+ "values are not searched. Requires pattern or patterns — do not "
+ "call without a term. Each term is a separate block of section "
+ "rows (score-free — this is an exact match, not a ranked search; "
+ "use corpus.recall for ranking). Terms with fewer matching "
+ "sections are listed first. Rows past limit are folded per "
+ "document with a reusable scope — copy that scope onto a later "
+ "corpus.grep call to search only that document. Table and image "
+ "hits are address + snippet rows; read the hit's document_id + "
+ "section_path (or chunk_id for image/table) with corpus.read "
+ "next. Provide 'pattern' for one term, or 'patterns' for several "
+ "terms in this single call instead of several parallel "
+ "corpus.grep calls. At least one of pattern/patterns is required."
),
json_schema={
"type": "object",
@@ -120,9 +175,10 @@ def _term_search(terms: list[str], namespace: str = "default") -> tuple[re.Patte
"type": "array",
"items": {"type": "string"},
"description": (
- "Several search terms OR'd together in this one call. "
- "A chunk matches if any term hits. Results stay in "
- "document order, not relevance order."
+ "Several search terms in this one call. Each term is "
+ "listed as its own block. Terms with fewer matching "
+ "sections are shown first. This tool does not rank "
+ "by relevance."
),
},
"scope": SCOPE_SCHEMA,
@@ -139,7 +195,10 @@ def _term_search(terms: list[str], namespace: str = "default") -> tuple[re.Patte
"type": "integer",
"minimum": 1,
"default": _DEFAULT_LIMIT,
- "description": "Max rows to return. The total match count is still reported.",
+ "description": (
+ "Max section rows to show. Further matching sections "
+ "are folded per document with a reusable scope."
+ ),
},
},
"anyOf": [
@@ -179,7 +238,7 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
section_path_col=corpusStorage.DocumentSection.section_path,
)
- compiled, term_filter = _term_search(terms, ctx.namespace)
+ _compiled, term_filter = _term_search(terms, ctx.namespace)
# Match the indexed haystack first so PostgreSQL can use
# idx_document_chunks_term_trgm. Section is joined here (not later) so a
# scope's section-subtree condition can apply inside this same CTE.
@@ -191,9 +250,6 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
corpusStorage.DocumentChunk.job_result_id,
corpusStorage.DocumentChunk.chunk_type,
corpusStorage.DocumentChunk.term_search_text,
- corpusStorage.DocumentChunk.content,
- corpusStorage.DocumentChunk.file_path,
- corpusStorage.DocumentChunk.chunk_metadata,
corpusStorage.DocumentChunk.section_id,
corpusStorage.DocumentChunk.sort_order,
corpusStorage.DocumentSection.section_path,
@@ -222,113 +278,86 @@ async def grep(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
matched.c.document_id,
matched.c.chunk_type,
matched.c.term_search_text,
- matched.c.content,
- matched.c.file_path,
- matched.c.chunk_metadata,
- matched.c.job_result_id,
- JobResult.job_id,
- JobResult.document_metadata["result_raw_prefix"].as_string(),
+ matched.c.sort_order,
matched.c.section_path,
corpusStorage.Document.source_file_name,
- func.count().over().label("total_matches"),
)
.select_from(matched)
.join(corpusStorage.Document, corpusStorage.Document.document_id == matched.c.document_id)
- .outerjoin(JobResult, JobResult.id == CorpusRevisionContext.build_revision_column(corpusStorage.Document))
.where(
corpusStorage.Document.user_id == corpusStorage.resolve_owner(ctx.user_id),
corpusStorage.Document.namespace == ctx.namespace,
corpusStorage.Document.status == "active",
CorpusRevisionContext.build_revision_column(corpusStorage.Document) == matched.c.job_result_id,
)
- .order_by(matched.c.document_id, matched.c.sort_order)
- .limit(limit)
+ .order_by(matched.c.document_id, matched.c.sort_order, matched.c.chunk_id)
)
matched_rows = (await ctx.db.execute(rows_stmt)).all()
- total_matches = int(matched_rows[0][-1]) if matched_rows else 0
- results: list[dict[str, Any]] = []
+ hits: list[dict[str, Any]] = []
for (
chunk_id,
document_id,
chunk_type,
term_search_text,
- content,
- file_path,
- chunk_metadata,
- job_result_id,
- job_id,
- result_raw_prefix,
+ _sort_order,
section_path,
source_file_name,
- _total_matches,
) in matched_rows:
- text = str(term_search_text or "")
- match = compiled.search(text)
- snippet = build_snippet(
- text, match.span() if match else None, hit_context=context_chars
- )
- results.append(
+ hits.append(
{
"document_id": document_id,
"source_file_name": source_file_name,
"chunk_id": chunk_id,
"chunk_type": chunk_type,
"section_path": section_path,
- "snippet": snippet,
- "content": content,
- "file_path": file_path,
- "chunk_metadata": chunk_metadata or {},
- "job_result_id": job_result_id,
- "job_id": job_id,
- "result_raw_prefix": result_raw_prefix,
+ "term_search_text": term_search_text,
}
)
- media: list[dict[str, str]] = []
- if results:
- results, media = await mount_explore_hits(
- ctx, results, char_budget=ctx.budget.max_chars
- )
- host_paths = await host_paths_for_hits(ctx, results, scope)
+ blocks: list[tuple[int, int, str, list[dict[str, Any]]]] = []
+ for index, term in enumerate(terms):
+ sections = _sections_for_term(hits, term, context_chars)
+ if sections:
+ blocks.append((len(sections), index, term, sections))
+ blocks.sort()
- rows: list[dict[str, Any]] = []
- lines = [f"total_matches={total_matches} returned={len(results)}"]
+ shown_hits: list[dict[str, Any]] = []
+ remaining = limit
+ allocated: list[tuple[str, list[dict[str, Any]], list[dict[str, Any]]]] = []
+ for _count, _index, term, sections in blocks:
+ visible = sections[:remaining]
+ leftover = sections[remaining:]
+ remaining -= len(visible)
+ shown_hits.extend(visible)
+ allocated.append((term, visible, leftover))
+
+ host_paths = await host_paths_for_hits(ctx, shown_hits, scope)
+
+ lines: list[str] = []
if requested_limit > limit:
lines.append(f"note: capped to budget.max_items={ctx.budget.max_items}")
- for r in results:
- chunk_type = str(r["chunk_type"] or "").strip()
- is_asset = chunk_type in ASSET_CHUNK_TYPES
- key = (str(r["document_id"]), str(r["chunk_id"]))
- section_path, hosted = (
- host_paths[key] if is_asset and key in host_paths else (r["section_path"], None)
- )
- row = build_row(
- kind=chunk_type or "text",
- document_id=r["document_id"],
- section_path=section_path,
- title=r["source_file_name"],
- chunk_id=r["chunk_id"] if is_asset else None,
- snippet=r["snippet"],
- hosted=hosted if is_asset else None,
- )
- rows.append(row)
- lines.append(format_row(row))
- rendered = str(r.get("rendered") or "").strip()
- if rendered:
- lines.append(rendered)
+ rows: list[dict[str, Any]] = []
+ for term, visible, leftover in allocated:
+ lines.append(term)
+ for hit in visible:
+ row = _build_hit_row(hit, host_paths)
+ rows.append(row)
+ lines.append(format_row(row))
+ leftover_by_doc: dict[str, list[dict[str, Any]]] = defaultdict(list)
+ for hit in leftover:
+ leftover_by_doc[str(hit["document_id"])].append(hit)
+ for document_id in leftover_by_doc:
+ group = leftover_by_doc[document_id]
+ lines.append(
+ _fold_line(str(group[0]["source_file_name"] or ""), document_id, len(group))
+ )
return ToolResult(
text="\n".join(lines),
- payload={
- "rows": rows,
- "details": {
- "total_matches": total_matches,
- "capped": requested_limit > limit,
- },
- },
+ payload={"rows": rows, "details": {}},
refs=[
- {"document_id": r["document_id"], "chunk_id": r["chunk_id"]} for r in results
+ {"document_id": hit["document_id"], "chunk_id": hit["chunk_id"]}
+ for hit in shown_hits
],
- media=media,
)
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py
index a742ee18f..a312a8ac3 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/list_documents.py
@@ -1,7 +1,7 @@
"""``corpus.list_documents`` — namespace-level document overview.
-Joins ``documents`` with the document-level ``graph_nodes`` row (§4 of
-``CORPUS_SCHEMA.md``) to surface per-document keywords/summary/type-mix
+Joins ``documents`` with the document-level ``graph_nodes`` row
+to surface per-document keywords/summary/type-mix
without reading any chunk content.
"""
@@ -24,9 +24,8 @@
@register_tool(
name="corpus.list_documents",
description=(
- "List every active document in the namespace with its parse_track "
- "and, when available, document-level graph metadata (top_keywords, "
- "top_summary, chunk type mix). Use only when explicitly asked to "
+ "List every active document in the namespace with its summary, "
+ "top keywords and chunk count. Use only when explicitly asked to "
"inventory the namespace's documents, not as a first step for "
"question answering."
),
@@ -63,7 +62,6 @@ async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult:
"document_id": document.document_id,
"job_result_id": CorpusRevisionContext.resolve_revision(document.document_id, document.current_job_result_id),
"source_file_name": document.source_file_name,
- "parse_track": document.parse_track,
"top_keywords": props.get("top_keywords") or [],
"top_summary": props.get("top_summary") or "",
"types": props.get("types") or {},
@@ -72,7 +70,7 @@ async def list_documents(ctx: ToolContext, _args: dict[str, Any]) -> ToolResult:
documents.append(entry)
summary_line = (
f"- {entry['source_file_name']} ({entry['document_id']}, "
- f"track={entry['parse_track']}, chunks={entry['chunks_count']})"
+ f"chunks={entry['chunks_count']})"
)
if entry["top_summary"]:
summary_line += f"\n summary: {entry['top_summary']}"
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py
index e70ce9a5e..f37f1f7fc 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/neighbors.py
@@ -1,6 +1,6 @@
"""``corpus.neighbors`` — document-level ``related`` graph edges.
-Document-to-document only (§4 of ``CORPUS_SCHEMA.md``): no section- or
+Document-to-document only: no section- or
entity-level graph nodes exist yet. Edges are undirected and were written by
``DocumentGraphService.publish_document_graph`` with ``shared_entities`` (typed
entity overlap, preferred) or ``shared_keywords`` (TF-IDF fallback).
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py
index 7df3d3459..2b84d3982 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/node_filter.py
@@ -106,8 +106,8 @@
},
"chunk_types": chunk_types_schema(
"Narrow matched sections to those owning a chunk of one of "
- "these types (e.g. ['page'] to filter to page-track leaves "
- "only). Omit for no narrowing."
+ "these types (e.g. ['page'] to keep only sections whose "
+ "body is PDF pages). Omit for no narrowing."
),
},
"required": ["scope", "predicates"],
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py
index 4c1246315..b6ac97a1a 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/outline.py
@@ -204,7 +204,10 @@ def header(visible: int, hidden: int) -> str:
text=rendered.text,
payload={
"rows": visible_rows,
- "details": {"hidden_count": len(order) - len(visible_rows)},
+ "details": {
+ "hidden_count": len(order) - len(visible_rows),
+ "documents": source_file_name_by_doc,
+ },
},
refs=[
{"document_id": row["document_id"], "section_path": row["section_path"]}
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py
index f2b1a6805..fad4f97c2 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/query_table.py
@@ -145,10 +145,11 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
connection.close()
result_grid = [result_headers, *result_rows] if result_headers else []
+ table_html = grid_to_html(result_grid)
text = (
f"### {doc.source_file_name} ({document_id}) / {section_path} [table]\n"
f"columns={', '.join(columns)}\n"
- f"{grid_to_html(result_grid)}"
+ f"{table_html}"
)
return ToolResult(
text=text,
@@ -158,6 +159,7 @@ async def query_table(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
"columns": columns,
"headers": result_headers,
"rows": result_rows,
+ "table_html": table_html,
},
refs=[{"document_id": document_id, "chunk_id": chunk_id}],
)
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py
index 39c36480d..d36dd841e 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/read.py
@@ -3,10 +3,10 @@
Unlike ``hydration.result_assembly.assemble_retrieval_results`` (which
down-weights ``page`` chunks to their summary — see ``page_summary``, a
deliberate trade-off for the retrieval-answer surface),
-``read`` returns the page chunk's full body content, with ``[SAME-AS
-p]`` markers resolved to the owner section's text (§2 of
-``CORPUS_SCHEMA.md``) rather than stripped or summarized. ``connect_to``
-assets are still inlined via the same placeholder mechanism as retrieval.
+``read`` returns the page chunk's full body content, with shared-page
+markers resolved to the owner section's text rather than stripped or
+summarized. Embedded images and tables are still inlined via the same
+placeholder mechanism as retrieval.
Table chunks load the stored HTML: small tables return that HTML; large
tables return row/column headers only and point at ``corpus.query_table``
(no window — GREP/recall never scan table-cell HTML, so there is no real
@@ -23,6 +23,10 @@
exact match first, then a unique segment-bound suffix match when the agent
omits ancestor segments; ambiguous suffix matches return an error listing
candidate full paths instead of picking one silently.
+
+When ``agent_explore`` sets ``ToolContext.readable``/``decided``, a ref must
+name an address from an earlier round's results, and body chunks the episode
+already picked or passed on are not read again.
"""
from __future__ import annotations
@@ -44,6 +48,7 @@
from shared.services.retrieval.agent_tools.registry import (
REF_ADDRESS_ONE_OF,
REF_ADDRESS_RULE,
+ Decision,
ToolContext,
ToolResult,
register_tool,
@@ -62,22 +67,72 @@
from shared.services.retrieval.agent_tools.section_path_lookup import (
resolve_section_path_anchor,
section_path_anchor_filter,
+ section_path_received,
section_path_subtree_filter,
)
from shared.services.retrieval.search.lexical_text import section_path_from_chunk_path
_SAME_AS_MARKER_RE = re.compile(r"\[SAME-AS (.+?) p(\d+)\]")
_BODY_CHUNK_TYPES = ("text", "page")
-# Repeated after the ref list and after the body: a long body can be cut by
-# the per-turn text cap, and the ref list alone may scroll out of view.
-_PICK_REMINDER = (
- "[pick: decide now, for each [ok] ref, whether to cite it in finish; "
- "[failed] refs cannot be cited]"
+_LOCATE_HINT = (
+ " Copy document_id and section_path or chunk_id exactly from a prior "
+ "result row, or locate the section first with corpus.outline, "
+ "corpus.grep, corpus.recall, or corpus.assets."
+)
+_NOT_RECEIVED_REASON = (
+ "not in any earlier result you have received; results from calls in "
+ "this same turn are not available yet"
)
+def _with_locate_hint(reason: str) -> str:
+ if _LOCATE_HINT.strip() in reason:
+ return reason
+ return reason + _LOCATE_HINT
+
+
+def _ref_received(ctx: ToolContext, document_id: str, ref: dict[str, Any]) -> bool:
+ if ctx.readable is None:
+ return True
+ chunk_id = str(ref.get("chunk_id") or "").strip()
+ if chunk_id:
+ return (document_id, chunk_id) in ctx.readable
+ section_path = str(ref.get("section_path") or "").strip()
+ return section_path_received(ctx.readable, document_id, section_path)
+
+
+def _decided_chunks(
+ ctx: ToolContext, document_id: str, chunk_ids: list[str]
+) -> dict[str, Decision]:
+ if ctx.decided is None:
+ return {}
+ return {
+ chunk_id: ctx.decided[(document_id, chunk_id)]
+ for chunk_id in chunk_ids
+ if (document_id, chunk_id) in ctx.decided
+ }
+
+
+def _already_read_reason(decisions: list[Decision]) -> str:
+ parts: list[str] = []
+ for decision in dict.fromkeys(decisions):
+ state = (
+ f"picked as {decision.picked_handle}"
+ if decision.picked_handle
+ else "not picked"
+ )
+ parts.append(f"round {decision.read_round} ({state})")
+ return f"already read in {', '.join(parts)}; cannot read again"
+
+
def _ref_status_line(entry: dict[str, Any]) -> str:
- tag = "[ok]" if entry["status"] == "ok" else f"[failed: {entry['reason']}]"
+ if entry["status"] == "ok":
+ omitted = entry.get("omitted_chunk_ids")
+ tag = (
+ f"[ok, {len(omitted)} already-read chunks omitted]" if omitted else "[ok]"
+ )
+ else:
+ tag = f"[failed: {entry['reason']}]"
return (
f"{tag} document_id={entry['document_id']} "
f"section_path={entry['section_path']} chunk_id={entry['chunk_id']}"
@@ -159,9 +214,9 @@ async def _resolve_same_as_markers(
for marker_text, owner_path in row_matches:
resolved = owner_content.get((document_id, owner_path), "")
if resolved:
- replacement = f"(SAME-AS {owner_path} resolved)\n{resolved}"
+ replacement = f"(page text shared with section {owner_path})\n{resolved}"
else:
- replacement = f"(SAME-AS {owner_path} — page not found)"
+ replacement = f"(page text shared with section {owner_path}: not found)"
content = content.replace(marker_text, replacement, 1)
row["content"] = content
@@ -257,15 +312,14 @@ def _add(row: dict[str, Any]) -> None:
@register_tool(
name="corpus.read",
description=(
- "Read full body content for already-located sections or chunks. "
- "Resolves page-track SAME-AS pointers to the owner section's text, "
- "inlines connect_to assets, loads table HTML (small tables in full; "
- "large tables as row/column headers plus a pointer to "
- "corpus.query_table), and converts asset/page_assets references to "
- "URLs. Use after outline/node_filter/recall/grep have located where "
- "to look. Each ref's outcome (ok, or failed with a reason) is "
- "reported separately — decide pick/no-pick for each ref right "
- "after seeing its status; a failed ref defaults to not picked."
+ "Read full content for sections (document_id + section_path) or "
+ "chunks (document_id + chunk_id) that you already located. Copy "
+ "both values exactly from a prior result row. Page text shared "
+ "with another section is filled in automatically. Images and "
+ "tables the section contains are shown inline: small tables in "
+ "full, large tables as row and column headers (then use "
+ "corpus.query_table). Each ref's outcome (ok, or failed with a "
+ "reason) is reported separately."
),
json_schema={
"type": "object",
@@ -308,12 +362,7 @@ def _add(row: dict[str, Any]) -> None:
"include_assets": {
"type": "boolean",
"default": True,
- "description": "Inline connect_to images and tables into the body.",
- },
- "resolve_same_as": {
- "type": "boolean",
- "default": True,
- "description": "Replace page-track SAME-AS markers with the owner text.",
+ "description": "Show the images and tables this section contains inline.",
},
},
"required": ["refs"],
@@ -329,7 +378,6 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
if mode not in ("self", "descendants"):
return ToolResult(text="", error=f"unsupported mode: {mode}")
include_assets = bool(args.get("include_assets", True))
- resolve_same_as_flag = bool(args.get("resolve_same_as", True))
char_budget = ctx.budget.max_chars
document_ids = {
@@ -387,10 +435,16 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
emit_items: list[tuple[str, Any]] = []
for index, ref in enumerate(refs):
document_id = str(ref.get("document_id") or "").strip()
+ if not _ref_received(ctx, document_id, ref):
+ ref_status[index]["status"] = "failed"
+ ref_status[index]["reason"] = _NOT_RECEIVED_REASON
+ continue
job_result_id = revision_by_doc.get(document_id)
if not job_result_id:
ref_status[index]["status"] = "failed"
- ref_status[index]["reason"] = f"unknown document_id: {document_id}"
+ ref_status[index]["reason"] = _with_locate_hint(
+ f"unknown document_id: {document_id}"
+ )
continue
chunk_id = str(ref.get("chunk_id") or "").strip()
section_path = str(ref.get("section_path") or "").strip()
@@ -400,7 +454,11 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
if chunk_id:
row = (
await ctx.db.execute(
- select(corpusStorage.DocumentChunk, corpusStorage.DocumentSection.section_path)
+ select(
+ corpusStorage.DocumentChunk,
+ corpusStorage.DocumentSection.section_path,
+ corpusStorage.DocumentSection.summary,
+ )
.select_from(corpusStorage.DocumentChunk)
.outerjoin(
corpusStorage.DocumentSection,
@@ -413,9 +471,16 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
).first()
if row is None:
ref_status[index]["status"] = "failed"
- ref_status[index]["reason"] = f"unknown chunk_id: {chunk_id} in {document_id}"
+ ref_status[index]["reason"] = _with_locate_hint(
+ f"unknown chunk_id: {chunk_id} in {document_id}"
+ )
+ continue
+ chunk, resolved_section_path, section_summary = row
+ decided = _decided_chunks(ctx, document_id, [chunk.chunk_id])
+ if decided:
+ ref_status[index]["status"] = "failed"
+ ref_status[index]["reason"] = _already_read_reason(list(decided.values()))
continue
- chunk, resolved_section_path = row
ref_status[index]["status"] = "ok"
ref_status[index]["chunk_ids"] = [chunk.chunk_id]
emit_items.append(
@@ -435,6 +500,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
"content": chunk.content,
"chunk_metadata": chunk.chunk_metadata or {},
"file_path": chunk.file_path,
+ "section_summary": str(section_summary or "").strip(),
}
],
)
@@ -443,7 +509,9 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
if not section_path:
ref_status[index]["status"] = "failed"
- ref_status[index]["reason"] = f"ref for {document_id} needs section_path or chunk_id"
+ ref_status[index]["reason"] = _with_locate_hint(
+ f"ref for {document_id} needs section_path or chunk_id"
+ )
continue
resolved_path, path_error = await resolve_section_path_anchor(
@@ -454,7 +522,9 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
)
if path_error or not resolved_path:
ref_status[index]["status"] = "failed"
- ref_status[index]["reason"] = path_error or f"unknown section_path for {document_id}"
+ ref_status[index]["reason"] = _with_locate_hint(
+ path_error or f"unknown section_path for {document_id}"
+ )
continue
emit_items.append(
(
@@ -513,12 +583,9 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
corpusStorage.DocumentChunk.section_id.in_(section_ids),
# Body chunks only (text/page). image/table chunks share a
# section_id with whichever section happens to store them
- # in the DB (always Root — CORPUS_SCHEMA.md §3), which is
- # not the same as "belonging" to that section; their real
- # association is connect_to on the body chunk, resolved
- # below via hydrate_connected_target_rows. Without this
- # filter, reading Root would return every still-unmounted
- # asset in the document as spurious top-level entries.
+ # in the DB, which is not the same as belonging to that
+ # section; their real association is on the body chunk,
+ # resolved below via hydrate_connected_target_rows.
corpusStorage.DocumentChunk.chunk_type.in_(_BODY_CHUNK_TYPES),
)
)
@@ -550,9 +617,26 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
for section in section_matches
if section.section_id in matched_ids
}
+ section_summary_by_id = {
+ section.section_id: str(section.summary or "").strip()
+ for section in section_matches
+ if section.section_id in matched_ids
+ }
+ job_chunks = [
+ chunk for chunk in batched_chunks if chunk.section_id in matched_ids
+ ]
+ decided = _decided_chunks(
+ ctx, job["document_id"], [chunk.chunk_id for chunk in job_chunks]
+ )
+ if job_chunks and all(chunk.chunk_id in decided for chunk in job_chunks):
+ ref_status[ref_index]["status"] = "failed"
+ ref_status[ref_index]["reason"] = _already_read_reason(list(decided.values()))
+ continue
+ if decided:
+ ref_status[ref_index]["omitted_chunk_ids"] = list(decided)
job_chunk_ids: list[str] = []
for chunk in batched_chunks:
- if chunk.section_id not in matched_ids:
+ if chunk.section_id not in matched_ids or chunk.chunk_id in decided:
continue
job_chunk_ids.append(chunk.chunk_id)
base_rows.append(
@@ -573,6 +657,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
"content": chunk.content,
"chunk_metadata": chunk.chunk_metadata or {},
"file_path": chunk.file_path,
+ "section_summary": section_summary_by_id.get(chunk.section_id, ""),
}
)
if job_chunk_ids:
@@ -580,7 +665,7 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
ref_status[ref_index]["chunk_ids"] = job_chunk_ids
else:
ref_status[ref_index]["status"] = "failed"
- ref_status[ref_index]["reason"] = (
+ ref_status[ref_index]["reason"] = _with_locate_hint(
f"no body chunk found for {job['resolved_path']} in "
f"{job['document_id']} (image/table-only or empty section)"
)
@@ -593,19 +678,15 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
error=(
"read: every ref failed, nothing was read:\n"
+ "\n".join(status_lines)
- + "\nNone of these can be picked for finish. Fix each ref "
- "(copy document_id + section_path/chunk_id exactly from an "
- "outline/node_filter/grep/recall/assets row) and retry."
),
)
- if resolve_same_as_flag:
- await _resolve_same_as_markers(
- ctx.db,
- base_rows,
- revision_by_doc=revision_by_doc,
- source_file_name_by_doc=source_file_name_by_doc,
- )
+ await _resolve_same_as_markers(
+ ctx.db,
+ base_rows,
+ revision_by_doc=revision_by_doc,
+ source_file_name_by_doc=source_file_name_by_doc,
+ )
connected_rows: list[dict[str, Any]] = []
if include_assets:
@@ -653,14 +734,13 @@ async def read(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
composed["content"] = _image_display_content(row)
assembled.append(composed)
- lines = ["refs:", *(f" {line}" for line in status_lines), _PICK_REMINDER]
+ lines = ["refs:", *(f" {line}" for line in status_lines)]
for row in assembled:
lines.append(
f"### {row.get('source_file_name')} ({row.get('document_id')}) / "
f"{row.get('section_path')} [{row.get('chunk_type')}]"
)
lines.append(str(row.get("content") or ""))
- lines.append(_PICK_REMINDER)
return ToolResult(
text="\n".join(lines),
diff --git a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py
index 6ac6bf72a..d265f163b 100644
--- a/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py
+++ b/packages/shared-python/shared/services/retrieval/agent_tools/tools/recall.py
@@ -60,10 +60,11 @@ def _identifier_snippet(row: dict[str, Any]) -> str:
name="corpus.recall",
description=(
"Fuzzy ranked candidate search for a question when you don't know "
- "where the answer lives. Scores path and content with BM25. "
+ "where the answer lives. Ranks by how well section path and text "
+ "match the query. "
"Returns candidates with chunk_type, document_id, path and snippet. "
- "Table and image hits include rendered content; body hits include "
- "any connected table or image. Exact string / identifier lookup is "
+ "Table and image hits show their content; body hits also show the "
+ "images and tables they contain. Exact string / identifier lookup is "
"corpus.grep, not this tool. Read the candidate's document_id + "
"section_path (or chunk_id, for image/table hits) with corpus.read "
"next, not the filename."
@@ -147,7 +148,7 @@ async def recall(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
# recall regardless of what a better next step happens to be for
# this corpus/query, so it nudges the agent to change approach
# without prescribing which other tool to reach for (that's already
- # covered generically in CORPUS_SCHEMA.md §6's tool-selection table).
+ # covered generically in each other tool's own description).
lines.append(
"note: few or no candidates for this phrasing — rephrasing the "
"query and calling recall again rarely surfaces more; a "
@@ -177,6 +178,7 @@ async def recall(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
score=float(raw_score) if raw_score is not None else None,
hosted=hosted if is_asset else None,
)
+ row["mounted_chunk_ids"] = hit["mounted_chunk_ids"]
payload_rows.append(row)
lines.append(format_row(row))
rendered = str(hit.get("rendered") or "").strip()
diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py
index 4f206196a..6cba8f9f8 100644
--- a/packages/shared-python/shared/services/retrieval/execution/routes.py
+++ b/packages/shared-python/shared/services/retrieval/execution/routes.py
@@ -14,8 +14,7 @@
RetrievalRouteOutcome,
)
from shared.services.retrieval.hydration.evidence_compose import (
- collect_evidence,
- flatten_parts,
+ group_evidence_units,
)
from shared.services.retrieval.hydration.result_assembly import (
assemble_retrieval_results,
@@ -50,12 +49,24 @@ def open_agent_explore_database_context() -> AbstractAsyncContextManager[AsyncSe
def _evidence_fields(rows: list[dict]) -> dict:
- # TODO: 后面用 TypeSafe JEV 补结果重排/筛选。现在没有这一步。
- # 无论怎么做，发出去的 evidence 和 results 都是已经筛过或重排过的完整列表，不在 Knowhere 外面做。
- evidence = collect_evidence(rows)
+ # TODO: later use TypeSafe JEV for result rerank/filter. This step does not exist yet.
+ evidence = group_evidence_units(
+ [
+ {
+ "document_id": row["document_id"],
+ "source_file_name": row["source_file_name"],
+ "section_path": row["section_path"],
+ "sort_order": row["sort_order"],
+ "parts": row["composed"],
+ "kind": "read",
+ }
+ for row in rows
+ if row["composed"]
+ ]
+ )
return {
"evidence": evidence,
- "evidence_text": flatten_parts(evidence),
+ "evidence_text": "",
}
@@ -202,14 +213,15 @@ async def _run_agent_explore_route(
``AGENT_EXPLORE_HARNESS`` switch resolved by ``resolve_harness()``.
"""
from shared.services.retrieval.agent_explore.bridge import (
- attach_ref_provenance,
+ attach_evidence_pool,
build_decision_trace,
)
from shared.services.retrieval.agent_explore.budget import EpisodeBudget
- from shared.services.retrieval.agent_explore.harness import resolve_harness
- from shared.services.retrieval.agent_explore.ref_resolution import (
- resolve_finish_refs,
+ from shared.services.retrieval.agent_explore.evidence_pool import (
+ compose_pool_evidence,
+ pool_chunk_refs,
)
+ from shared.services.retrieval.agent_explore.harness import resolve_harness
from shared.services.retrieval.trace import TraceRecorder
# End any read transaction created by the route's pre-episode work before
@@ -229,37 +241,21 @@ async def _run_agent_explore_route(
document_scope=context.document_scope,
)
logger.info(
- "retrieval agent_explore stage=episode seconds={:.3f} refs={} "
+ "retrieval agent_explore stage=episode seconds={:.3f} pool={} "
"steps={} tokens={} stop_reason={}".format(
time.perf_counter() - episode_started,
- len(episode.refs),
+ len(episode.pool),
len(episode.steps),
episode.tokens_used,
episode.stop_reason,
)
)
- # episode.refs are document_id + section_path (what the agent actually
- # sees in tool text); resolve_workflow_references requires chunk_id —
- # see ref_resolution.py's module docstring for why this bridge exists.
decision_steps = build_decision_trace(episode.steps)
+ chunk_refs = pool_chunk_refs(episode.pool)
+ decision_steps = attach_evidence_pool(decision_steps, episode.pool)
async with open_fresh_database_context() as final_db:
- finish_resolution = await resolve_finish_refs(
- final_db,
- user_id=context.user_id,
- namespace=context.namespace,
- refs=episode.refs,
- document_scope=context.document_scope,
- )
- chunk_refs = finish_resolution.resolved
- decision_steps = attach_ref_provenance(
- decision_steps,
- agent_selected_refs=episode.agent_selected_refs,
- fallback_refs=episode.fallback_refs,
- resolved_refs=chunk_refs,
- dropped_refs=finish_resolution.dropped,
- )
resolved = await resolve_workflow_references(
db=final_db,
user_id=context.user_id,
@@ -275,7 +271,13 @@ async def _run_agent_explore_route(
exclude_sections=context.exclude_sections,
allowed_chunk_types=context.allowed_chunk_types,
revision_pins=context.revision_pins,
+ queried_tables=episode.queried_tables,
)
+ evidence = compose_pool_evidence(episode.pool, assembled_rows)
+ evidence_fields = {
+ "evidence": evidence,
+ "evidence_text": "",
+ }
selected_doc_ids = list(
{row.get("document_id", "") for row in resolved.rows if row.get("document_id")}
@@ -303,7 +305,6 @@ async def _run_agent_explore_route(
)
decision_trace = [step.to_dict() for step in decision_steps]
- evidence_fields = _evidence_fields(assembled_rows)
response = {
"namespace": context.namespace,
"query": context.query,
@@ -321,6 +322,6 @@ async def _run_agent_explore_route(
completion_label="AGENT EXPLORE RETRIEVAL",
completion_count=len(resolved.refs),
completion_detail=(
- f"chunks | evidence={len(evidence_fields['evidence_text'])} chars | router=agent_explore"
+ f"chunks | evidence={len(evidence)} parts | router=agent_explore"
),
)
diff --git a/packages/shared-python/shared/services/retrieval/hydration/asset_inline.py b/packages/shared-python/shared/services/retrieval/hydration/asset_inline.py
index 7dd7b046e..b8d98ff7e 100644
--- a/packages/shared-python/shared/services/retrieval/hydration/asset_inline.py
+++ b/packages/shared-python/shared/services/retrieval/hydration/asset_inline.py
@@ -11,12 +11,13 @@
from collections.abc import Mapping, Sequence
from typing import Any
-_PATH_REF_RE = re.compile(r"\[(?:images|tables)/[^\]\n]+\]")
+from shared.utils.chunk_refs import RESOURCE_PATH_REF_RE, match_resource_path_ref
+
_SAME_AS_RE = re.compile(r"\[SAME-AS [^\]]+\]")
def remove_path_placeholders(content: str) -> str:
- text = _PATH_REF_RE.sub("", content)
+ text = RESOURCE_PATH_REF_RE.sub("", content)
return _SAME_AS_RE.sub("", text)
@@ -51,12 +52,11 @@ def inline_assets_at_placeholders(
ref = str(item.get("ref") or "").strip()
placed = False
- for candidate in _ref_candidates(ref):
- if candidate and candidate in text:
- text = text.replace(candidate, f"\n{body}\n", 1)
- embedded.add(target_id)
- placed = True
- break
+ match = match_resource_path_ref(text, ref)
+ if match is not None:
+ text = f"{text[: match.start()]}\n{body}\n{text[match.end() :]}"
+ embedded.add(target_id)
+ placed = True
if not placed:
pending_append.append((target_id, body))
@@ -70,19 +70,3 @@ def inline_assets_at_placeholders(
embedded.add(target_id)
return strip_path_placeholders(text), embedded
-
-
-def _ref_candidates(ref: str) -> list[str]:
- raw = str(ref or "").strip()
- if not raw:
- return []
- out: list[str] = [raw]
- if raw.startswith("[") and raw.endswith("]"):
- inner = raw[1:-1].strip()
- if inner and inner not in out:
- out.append(inner)
- else:
- bracketed = f"[{raw}]"
- if bracketed not in out:
- out.append(bracketed)
- return out
diff --git a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py
index 3f12020be..d0f0bd713 100644
--- a/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py
+++ b/packages/shared-python/shared/services/retrieval/hydration/evidence_compose.py
@@ -10,6 +10,7 @@
import base64
import tempfile
from pathlib import Path
+from collections.abc import Mapping
from typing import Any
from loguru import logger
@@ -17,6 +18,7 @@
from shared.services.retrieval.hydration.asset_inline import (
remove_path_placeholders,
)
+from shared.utils.chunk_refs import match_resource_path_ref
from shared.services.retrieval.hydration.row_utils import (
extract_page_nums,
normalize_chunk_type,
@@ -26,6 +28,7 @@
TableDownloadError,
load_table_html,
)
+from shared.services.retrieval.search.lexical_text import split_section_path
from shared.services.storage.raw_prefix_arguments import RawPrefixArguments
from shared.services.storage.result_storage import get_result_storage
@@ -41,47 +44,89 @@
def compose_evidence_parts(
row: dict[str, Any],
rows_by_chunk_id: dict[str, dict[str, Any]],
+ queried_tables: Mapping[tuple[str, str], str] | None = None,
) -> list[dict[str, Any]]:
chunk_type = normalize_chunk_type(row.get("chunk_type"))
if chunk_type == "page":
return _compose_page_parts(row)
if chunk_type == "table":
- return _compose_standalone_table_parts(row)
+ return _compose_standalone_table_parts(row, queried_tables)
if chunk_type == "image":
return _compose_standalone_image_parts(row)
- return _compose_text_parts(row, rows_by_chunk_id)
+ return _compose_text_parts(row, rows_by_chunk_id, queried_tables)
-def collect_evidence(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
- parts: list[dict[str, Any]] = []
- for row in rows:
- composed = row.get("composed")
- if isinstance(composed, list):
- revisionId: object = row.get("job_result_id")
- parts.extend(
- {**part, "job_result_id": revisionId} if revisionId else part
- for part in composed
- )
- return parts
+_GROUP_BODY_INDENT = " "
-def flatten_parts(parts: list[dict[str, Any]] | None) -> str:
- texts: list[str] = []
- for part in parts or []:
- if not isinstance(part, dict):
- continue
- if part.get("type") == "text":
- text = str(part.get("text") or "")
- if text:
- texts.append(text)
- continue
- if part.get("type") != "image":
+def group_evidence_units(units: list[dict[str, Any]]) -> list[dict[str, Any]]:
+ """Group units by document and the direct parent of ``section_path``.
+
+ Each unit is ``document_id``, ``source_file_name``, ``section_path``,
+ ``sort_order``, ``parts``, and ``kind`` (``read`` or ``outline``).
+ Outline units ignore ``section_path`` and ``sort_order``.
+ """
+ documents: dict[str, dict[str, Any]] = {}
+ for index, unit in enumerate(units):
+ document = documents.setdefault(
+ unit["document_id"], {"outlines": [], "reads": {}}
+ )
+ if unit["kind"] == "outline":
+ document["outlines"].append(
+ (f"[§ {unit['source_file_name']}]", [unit["parts"]])
+ )
continue
- media_type = str(part.get("media_type") or "").strip() or "application/octet-stream"
- data = str(part.get("data") or "").strip()
- if data:
- texts.append(f"data:{media_type};base64,{data}")
- return "".join(texts)
+ parent_path = _direct_parent_path(unit["section_path"])
+ group = document["reads"].setdefault(
+ parent_path,
+ {
+ "title": f"[§ {unit['source_file_name']} / {parent_path}]",
+ "members": [],
+ },
+ )
+ group["members"].append((unit["sort_order"], index, unit["parts"]))
+
+ evidence: list[dict[str, Any]] = []
+ group_number = 0
+ for document in documents.values():
+ read_groups = sorted(
+ (
+ (min(group["members"]), group["title"], sorted(group["members"]))
+ for group in document["reads"].values()
+ ),
+ key=lambda item: item[0],
+ )
+ ordered = [*document["outlines"]] + [
+ (title, [member[2] for member in members])
+ for _first, title, members in read_groups
+ ]
+ for title, member_parts in ordered:
+ if group_number:
+ evidence.append(_text_part("\n"))
+ group_number += 1
+ evidence.append(_text_part(f"[E{group_number}] {title}"))
+ for parts in member_parts:
+ evidence.extend(_indent_member_part(part) for part in parts)
+ return evidence
+
+
+def _direct_parent_path(section_path: str) -> str:
+ parts = split_section_path(section_path)
+ if len(parts) > 1:
+ return " / ".join(parts[:-1])
+ return section_path
+
+
+def _indent_member_part(part: dict[str, Any]) -> dict[str, Any]:
+ if part["type"] != "text":
+ return part
+ return {
+ "type": "text",
+ "text": "".join(
+ (_GROUP_BODY_INDENT + line if line.strip() else line)
+ for line in part["text"].splitlines(keepends=True)
+ ),
+ }
def _compose_page_parts(row: dict[str, Any]) -> list[dict[str, Any]]:
@@ -97,8 +142,11 @@ def _compose_page_parts(row: dict[str, Any]) -> list[dict[str, Any]]:
return parts
-def _compose_standalone_table_parts(row: dict[str, Any]) -> list[dict[str, Any]]:
- html = _try_read_table_html(row)
+def _compose_standalone_table_parts(
+ row: dict[str, Any],
+ queried_tables: Mapping[tuple[str, str], str] | None = None,
+) -> list[dict[str, Any]]:
+ html = _try_read_table_html(row, queried_tables)
if html is None:
return []
return [_text_part(f"\n{html}\n")]
@@ -118,11 +166,12 @@ def _compose_standalone_image_parts(row: dict[str, Any]) -> list[dict[str, Any]]
def _compose_text_parts(
row: dict[str, Any],
rows_by_chunk_id: dict[str, dict[str, Any]],
+ queried_tables: Mapping[tuple[str, str], str] | None = None,
) -> list[dict[str, Any]]:
content = str(row.get("content") or "")
tables, images = _embed_targets(row, rows_by_chunk_id)
for _target_id, target_row, ref in tables:
- html = _try_read_table_html(target_row)
+ html = _try_read_table_html(target_row, queried_tables)
content, placed = _replace_placeholder(content, ref, "" if html is None else f"\n{html}\n")
if html is not None and not placed:
_warn_skipped(target_row, "table", "placeholder not found")
@@ -194,10 +243,10 @@ def _connections(row: dict[str, Any]) -> list[dict[str, Any]]:
def _replace_placeholder(text: str, ref: str, replacement: str) -> tuple[str, bool]:
- for candidate in _ref_candidates(ref):
- if candidate and candidate in text:
- return text.replace(candidate, replacement, 1), True
- return text, False
+ match = match_resource_path_ref(text, ref)
+ if match is None:
+ return text, False
+ return text[: match.start()] + replacement + text[match.end() :], True
def _earliest_image_placeholder(
@@ -206,42 +255,31 @@ def _earliest_image_placeholder(
) -> tuple[str, str, dict[str, Any]] | None:
best: tuple[int, int, dict[str, Any]] | None = None
for _target_id, target_row, ref in images:
- for candidate in _ref_candidates(ref):
- if not candidate:
- continue
- index = text.find(candidate)
- if index < 0:
- continue
- length = len(candidate)
- if (
- best is None
- or index < best[0]
- or (index == best[0] and length > best[1])
- ):
- best = (index, length, target_row)
+ match = match_resource_path_ref(text, ref)
+ if match is None:
+ continue
+ index = match.start()
+ length = match.end() - match.start()
+ if (
+ best is None
+ or index < best[0]
+ or (index == best[0] and length > best[1])
+ ):
+ best = (index, length, target_row)
if best is None:
return None
index, length, target_row = best
return text[:index], text[index + length :], target_row
-def _ref_candidates(ref: str) -> list[str]:
- raw = str(ref or "").strip()
- if not raw:
- return []
- out = [raw]
- if raw.startswith("[") and raw.endswith("]"):
- inner = raw[1:-1].strip()
- if inner and inner not in out:
- out.append(inner)
- else:
- bracketed = f"[{raw}]"
- if bracketed not in out:
- out.append(bracketed)
- return out
-
-
-def _try_read_table_html(row: dict[str, Any]) -> str | None:
+def _try_read_table_html(
+ row: dict[str, Any],
+ queried_tables: Mapping[tuple[str, str], str] | None = None,
+) -> str | None:
+ if queried_tables is not None:
+ key = (str(row.get("document_id") or ""), str(row.get("chunk_id") or ""))
+ if key in queried_tables:
+ return queried_tables[key].strip()
try:
html = load_table_html(row).strip()
except TableDownloadError as exc:
diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py
index 1cc3219fa..f62f6745c 100644
--- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py
+++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py
@@ -27,6 +27,7 @@ async def assemble_retrieval_results(
document_scope: DocumentScope = DocumentScope(),
allowed_chunk_types: set[str] | None = None,
revision_pins: Mapping[str, str] | None = None,
+ queried_tables: Mapping[tuple[str, str], str] | None = None,
) -> list[dict[str, Any]]:
scoped_rows = filter_excluded_rows(
rows,
@@ -78,6 +79,7 @@ async def assemble_retrieval_results(
assembled_row['composed'] = compose_evidence_parts(
assembled_row,
rows_by_chunk_id,
+ queried_tables=queried_tables,
)
assembled.append(assembled_row)
return assembled
diff --git a/packages/shared-python/shared/services/retrieval/hydration/table_grid.py b/packages/shared-python/shared/services/retrieval/hydration/table_grid.py
index 668d041ea..482bb8642 100644
--- a/packages/shared-python/shared/services/retrieval/hydration/table_grid.py
+++ b/packages/shared-python/shared/services/retrieval/hydration/table_grid.py
@@ -196,6 +196,8 @@ def normalize_select_sql(sql: str) -> str | None:
return text
+# TODO: later adapt nested / multi-level row and column headers.
+# Current mapping is first row as columns and the first cell of each data row as row_header.
def sql_columns(grid: list[list[str]]) -> list[str]:
if not grid:
return ["row_header"]
diff --git a/packages/shared-python/shared/services/retrieval/settings.py b/packages/shared-python/shared/services/retrieval/settings.py
index 4768177a1..85ae9dc34 100644
--- a/packages/shared-python/shared/services/retrieval/settings.py
+++ b/packages/shared-python/shared/services/retrieval/settings.py
@@ -3,9 +3,9 @@
INTERNAL_RECALL_K_MULTIPLIER = 2
DEFAULT_TOP_K = 10
-# Final evidence / tool-observation text budget (characters). Used by
+# Agent-visible tool-observation text budget (characters). Used by
# agent tool-loop harness caps (``ToolBudget.max_chars``).
-EVIDENCE_TEXT_CHAR_BUDGET = 12_000
+EVIDENCE_TEXT_CHAR_BUDGET = 15_000
# Explore-phase table policy (user-specified). A table is large when either
# axis meets this size; the same number caps query_table SELECT rows.
diff --git a/packages/shared-python/shared/services/retrieval/trace/recorder.py b/packages/shared-python/shared/services/retrieval/trace/recorder.py
index 8b5156a3b..93821265a 100644
--- a/packages/shared-python/shared/services/retrieval/trace/recorder.py
+++ b/packages/shared-python/shared/services/retrieval/trace/recorder.py
@@ -107,8 +107,6 @@ async def create_run(self) -> None:
def record_decision_trace_step(self, step: DecisionTraceStep) -> None:
"""Buffer a DB trace row derived from the public decision trace step."""
result_status = str(step.result.get("status") or "unknown")
- budget = step.budget or {}
- tokens = int(budget.get("tokens_used_delta") or 0)
model = (
step.result.get("model")
or (step.decision or {}).get("model")
@@ -145,7 +143,6 @@ def record_decision_trace_step(self, step: DecisionTraceStep) -> None:
"observation_payload_keys": list(step.observation.keys()),
"latency_ms": int(step.elapsed_ms or 0),
"error": step.result.get("error"),
- "tokens_used": tokens,
"selected_paths": selected_paths,
"selected_doc_ids": selected_docs,
"model_name": str(model).strip() if model else None,
@@ -164,7 +161,6 @@ def record_budget_stop(self, reason: str) -> None:
"observation_payload_keys": [],
"latency_ms": 0,
"error": None,
- "tokens_used": 0,
"selected_paths": None,
"selected_doc_ids": None,
"model_name": None,
@@ -188,10 +184,7 @@ async def complete(
return
total_latency = int((time.monotonic() - self._start_time) * 1000)
- summed_tokens = sum(int(step.get("tokens_used", 0) or 0) for step in self._steps)
- run_token_count = (
- int(token_count) if token_count is not None else summed_tokens
- )
+ run_token_count = int(token_count) if token_count is not None else None
try:
from shared.models.database.document import RetrievalStep
@@ -206,12 +199,11 @@ async def complete(
observation={
"status": step_data["observation_status"],
"payload_keys": step_data["observation_payload_keys"],
- "tokens_used": step_data.get("tokens_used", 0),
},
selected_paths=step_data.get("selected_paths"),
selected_doc_ids=step_data.get("selected_doc_ids"),
latency_ms=step_data["latency_ms"],
- token_count=step_data.get("tokens_used", 0),
+ token_count=None,
model_name=step_data.get("model_name") or model_name,
error=step_data.get("error"),
created_at=step_data["created_at"],
diff --git a/packages/shared-python/shared/tests/test_agent_explore_evidence_pool.py b/packages/shared-python/shared/tests/test_agent_explore_evidence_pool.py
new file mode 100644
index 000000000..d4783182d
--- /dev/null
+++ b/packages/shared-python/shared/tests/test_agent_explore_evidence_pool.py
@@ -0,0 +1,572 @@
+"""Evidence pool rounds / issue / pick / render / compose."""
+
+from __future__ import annotations
+
+import os
+
+os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test")
+os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test")
+os.environ.setdefault("S3_BUCKET_NAME", "test-uploads")
+os.environ.setdefault("S3_ACCESS_KEY_ID", "test")
+os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test")
+os.environ.setdefault("S3_TEMP_PATH", "/tmp")
+
+from shared.services.retrieval.agent_explore.budget import EpisodeBudget
+from shared.services.retrieval.agent_explore.evidence_pool import (
+ Candidate,
+ EvidencePool,
+ PickOutcome,
+ compose_pool_evidence,
+ render_budget,
+ render_pick_outcome,
+ render_trace_line,
+ trace_status,
+)
+from shared.services.retrieval.agent_tools import Decision, ToolResult
+from shared.services.retrieval.hydration.evidence_compose import compose_evidence_parts
+from shared.services.retrieval.agent_tools.section_path_lookup import (
+ section_path_received,
+)
+
+
+def _ok_read(*, refs: list[dict], chunks: list[dict]) -> ToolResult:
+ return ToolResult(text="read", payload={"refs": refs, "chunks": chunks})
+
+
+def _read_one(section_path: str, chunk_ids: list[str]) -> ToolResult:
+ return _ok_read(
+ refs=[
+ {
+ "status": "ok",
+ "document_id": "doc_a",
+ "section_path": section_path,
+ "chunk_ids": chunk_ids,
+ }
+ ],
+ chunks=[
+ {
+ "chunk_id": chunk_id,
+ "document_id": "doc_a",
+ "source_file_name": "guide.pdf",
+ "section_path": section_path,
+ "chunk_type": "text",
+ "section_summary": section_path,
+ }
+ for chunk_id in chunk_ids
+ ],
+ )
+
+
+def _outline() -> ToolResult:
+ return ToolResult(
+ text="outline",
+ payload={
+ "rows": [
+ {
+ "document_id": "doc_a",
+ "section_path": "guide.pdf / 1 Overview",
+ "title": "1 Overview",
+ "depth": 0,
+ }
+ ],
+ "details": {"documents": {"doc_a": "guide.pdf"}},
+ },
+ )
+
+
+def test_issue_read_only_numbers_successful_refs_at_original_index() -> None:
+ pool = EvidencePool()
+ result = _ok_read(
+ refs=[
+ {
+ "status": "failed",
+ "document_id": "doc_a",
+ "section_path": "missing",
+ "chunk_ids": [],
+ },
+ {
+ "status": "ok",
+ "document_id": "doc_a",
+ "section_path": "2 Treatment",
+ "chunk_ids": ["c2"],
+ },
+ ],
+ chunks=[
+ {
+ "chunk_id": "c2",
+ "document_id": "doc_a",
+ "source_file_name": "guide.pdf",
+ "section_path": "2 Treatment",
+ "chunk_type": "text",
+ "section_summary": "First-line drugs",
+ }
+ ],
+ )
+ issued = pool.issue("corpus.read", result)
+ assert [item.handle for item in issued] == ["R1.2"]
+ assert issued[0].summary == "First-line drugs"
+ assert issued[0].chunk_ids == ("c2",)
+
+
+def test_ids_issued_this_round_take_effect_next_round() -> None:
+ pool = EvidencePool()
+ pool.begin_round()
+ pool.issue("corpus.read", _read_one("A", ["c1"]))
+ assert pool.pick_phase is False
+ same_round = pool.apply_pick(["R1.1"])
+ assert same_round.added == []
+ assert same_round.rejected == [{"handle": "R1.1", "reason": "unknown id"}]
+
+ pool.begin_round()
+ assert pool.round_index == 2
+ assert pool.pick_phase is True
+ next_round = pool.apply_pick(["R1.1"])
+ assert next_round.added == ["R1.1"]
+
+
+def test_pick_phase_lasts_until_pick_is_called() -> None:
+ pool = EvidencePool()
+ pool.begin_round()
+ pool.issue("corpus.outline", _outline())
+ pool.begin_round()
+ assert pool.pick_phase is True
+ pool.begin_round()
+ assert pool.pick_phase is True
+ pool.apply_pick([])
+ assert pool.pick_phase is False
+ pool.begin_round()
+ assert pool.pick_phase is False
+
+
+def test_apply_pick_rejects_unknown_duplicate_and_already_in_pool() -> None:
+ pool = EvidencePool()
+ pool.begin_round()
+ pool.issue(
+ "corpus.read",
+ _ok_read(
+ refs=[
+ {
+ "status": "ok",
+ "document_id": "doc_a",
+ "section_path": "A",
+ "chunk_ids": ["c1"],
+ },
+ {
+ "status": "ok",
+ "document_id": "doc_a",
+ "section_path": "A",
+ "chunk_ids": ["c1"],
+ },
+ ],
+ chunks=[
+ {
+ "chunk_id": "c1",
+ "document_id": "doc_a",
+ "source_file_name": "guide.pdf",
+ "section_path": "A",
+ "chunk_type": "text",
+ "section_summary": "A",
+ }
+ ],
+ ),
+ )
+ pool.begin_round()
+ outcome = pool.apply_pick(["R1.1", "R1.1", "R9.9", "R1.2"])
+ assert outcome.added == ["R1.1"]
+ assert outcome.rejected == [
+ {"handle": "R9.9", "reason": "unknown id"},
+ {"handle": "R1.2", "reason": "already in pool"},
+ ]
+
+
+def test_unpicked_read_chunks_are_decided_and_unpicked_outline_is_dropped() -> None:
+ pool = EvidencePool()
+ pool.begin_round()
+ pool.issue("corpus.read", _read_one("A", ["c1", "c2"]))
+ pool.issue("corpus.read", _read_one("B", ["c3"]))
+ pool.issue("corpus.outline", _outline())
+ pool.begin_round()
+ outcome = pool.apply_pick(["R2.1"])
+ assert outcome.added == ["R2.1"]
+ assert [item.handle for item in pool.entries] == ["R2.1"]
+ assert dict(pool.decided()) == {
+ ("doc_a", "c1"): Decision(read_round=1, picked_handle=None, handle="R1.1"),
+ ("doc_a", "c2"): Decision(read_round=1, picked_handle=None, handle="R1.1"),
+ ("doc_a", "c3"): Decision(read_round=1, picked_handle="R2.1", handle="R2.1"),
+ }
+ pool.begin_round()
+ assert pool.pick_phase is False
+ assert pool.apply_pick(["O1"]).rejected == [{"handle": "O1", "reason": "unknown id"}]
+
+
+def test_seen_addresses_become_readable_when_next_round_begins() -> None:
+ pool = EvidencePool()
+ pool.begin_round()
+ pool.issue(
+ "corpus.grep",
+ ToolResult(
+ text="hits",
+ payload={
+ "rows": [
+ {
+ "document_id": "doc_a",
+ "section_path": "guide.pdf / 2 Treatment",
+ "chunk_id": None,
+ "mounted_chunk_ids": ["tbl_1"],
+ },
+ {
+ "document_id": "doc_a",
+ "section_path": "guide.pdf / Root",
+ "chunk_id": "img_1",
+ "mounted_chunk_ids": [],
+ },
+ ]
+ },
+ ),
+ )
+ pool.issue(
+ "corpus.read",
+ _ok_read(
+ refs=[
+ {
+ "status": "ok",
+ "document_id": "doc_b",
+ "section_path": "notes.pdf / Intro",
+ "chunk_ids": ["c9"],
+ }
+ ],
+ chunks=[
+ {
+ "chunk_id": "c9",
+ "document_id": "doc_b",
+ "section_path": "notes.pdf / Intro",
+ "chunk_type": "text",
+ "chunk_metadata": {"connect_to": [{"target": "img_9"}]},
+ }
+ ],
+ ),
+ )
+ assert pool.readable() == frozenset()
+ pool.begin_round()
+ assert pool.readable() == {
+ ("doc_a", "guide.pdf / 2 Treatment"),
+ ("doc_a", "tbl_1"),
+ ("doc_a", "guide.pdf / Root"),
+ ("doc_a", "img_1"),
+ ("doc_b", "notes.pdf / Intro"),
+ ("doc_b", "c9"),
+ ("doc_b", "img_9"),
+ }
+
+
+def test_section_path_received_accepts_ancestor_and_suffix_paths() -> None:
+ readable = {("doc_a", "guide.pdf / 1 Overview / 1.1 Findings")}
+ assert section_path_received(readable, "doc_a", "guide.pdf / 1 Overview")
+ assert section_path_received(readable, "doc_a", "1 Overview / 1.1 Findings")
+ assert section_path_received(readable, "doc_a", "1.1 Findings")
+ assert not section_path_received(readable, "doc_a", "guide.pdf / 1.1 Findings")
+ assert not section_path_received(readable, "doc_a", "2 Treatment")
+ assert not section_path_received(readable, "doc_b", "1 Overview")
+
+
+def test_outline_tree_is_titles_and_indent_only() -> None:
+ pool = EvidencePool()
+ issued = pool.issue(
+ "corpus.outline",
+ ToolResult(
+ text="outline",
+ payload={
+ "rows": [
+ {
+ "document_id": "doc_a",
+ "title": "1 Overview",
+ "depth": 0,
+ },
+ {
+ "document_id": "doc_a",
+ "title": "2 Treatment",
+ "depth": 0,
+ },
+ {
+ "document_id": "doc_a",
+ "title": "2.1 Lifestyle",
+ "depth": 1,
+ },
+ ],
+ "details": {"documents": {"doc_a": "Hypertension_Guideline.pdf"}},
+ },
+ ),
+ )
+ assert issued[0].handle == "O1"
+ assert issued[0].outline_lines == (
+ " [O1] Hypertension_Guideline.pdf",
+ " 1 Overview",
+ " 2 Treatment",
+ " 2.1 Lifestyle",
+ )
+ rendered = pool.render_candidates(issued)
+ assert rendered == "[pick id for this result: O1 = this outline]"
+
+
+def test_descendants_keeps_anchor_summary_and_extra_sections() -> None:
+ pool = EvidencePool()
+ issued = pool.issue(
+ "corpus.read",
+ _ok_read(
+ refs=[
+ {
+ "status": "ok",
+ "document_id": "doc_a",
+ "section_path": "2 Treatment",
+ "chunk_ids": ["c_anchor", "c_child"],
+ }
+ ],
+ chunks=[
+ {
+ "chunk_id": "c_anchor",
+ "document_id": "doc_a",
+ "source_file_name": "guide.pdf",
+ "section_path": "2 Treatment",
+ "chunk_type": "text",
+ "section_summary": "anchor summary",
+ },
+ {
+ "chunk_id": "c_child",
+ "document_id": "doc_a",
+ "source_file_name": "guide.pdf",
+ "section_path": "2 Treatment / 2.1 Lifestyle",
+ "chunk_type": "text",
+ "section_summary": "child summary",
+ },
+ ],
+ ),
+ )
+ assert issued[0].summary == "anchor summary"
+ assert issued[0].extra_sections == 1
+ assert issued[0].chunk_ids == ("c_anchor", "c_child")
+
+
+def test_render_pool_and_empty() -> None:
+ pool = EvidencePool()
+ assert pool.render_pool() == "evidence pool: empty"
+ outline = pool.issue(
+ "corpus.outline",
+ ToolResult(
+ text="outline",
+ payload={
+ "rows": [{"document_id": "doc_a", "title": "1 Overview", "depth": 0}],
+ "details": {"documents": {"doc_a": "Hypertension_Guideline.pdf"}},
+ },
+ ),
+ )
+ pool.begin_round()
+ pool.apply_pick(["O1"])
+ text = pool.render_pool()
+ assert text.startswith("evidence pool (1):")
+ assert " [O1] Hypertension_Guideline.pdf" in text
+ assert " 1 Overview" in text
+ assert pool.render_candidates(outline) == "[pick id for this result: O1 = this outline]"
+
+
+def test_render_budget_and_trace_line_ok_partial_failed() -> None:
+ budget = EpisodeBudget(max_steps=12, wall_clock_seconds=180)
+ budget.record_step()
+ budget.record_usage({"total_tokens": 12034})
+ assert render_budget(budget) == "budget: steps 1/12"
+ assert budget.tokens_used == 12034
+
+ ok = render_trace_line(
+ 1,
+ "corpus_outline",
+ {"scope": [{"document_id": "doc_a"}]},
+ trace_status(ToolResult(text="ok")),
+ )
+ assert ok == '1. corpus_outline {"scope": [{"document_id": "doc_a"}]} -> ok'
+
+ failed = render_trace_line(
+ 3,
+ "corpus_grep",
+ {"patterns": ["ACEI"]},
+ trace_status(
+ ToolResult(text="", error="corpus.grep: unknown argument(s) ['max_results']")
+ ),
+ )
+ assert failed.endswith(" -> failed: corpus.grep: unknown argument(s) ['max_results']")
+
+ partial = render_trace_line(
+ 2,
+ "corpus_read",
+ {"refs": [{"document_id": "doc_a", "section_path": "A"}]},
+ trace_status(
+ ToolResult(
+ text="partial",
+ payload={
+ "refs": [
+ {"status": "ok"},
+ {"status": "failed", "reason": "unknown section_path for doc_a"},
+ ]
+ },
+ )
+ ),
+ )
+ assert " -> partial: ref 2 failed: unknown section_path for doc_a" in partial
+
+
+def test_render_pick_outcome() -> None:
+ assert render_pick_outcome(PickOutcome()) == "picked: none"
+ assert (
+ render_pick_outcome(
+ PickOutcome(
+ added=["R1.1", "O2"],
+ rejected=[{"handle": "R9.9", "reason": "unknown id"}],
+ )
+ )
+ == "picked: R1.1, O2; rejected: R9.9 (unknown id)"
+ )
+
+
+def test_compose_pool_evidence_keeps_pick_order_and_outline_fragment() -> None:
+ entries = [
+ Candidate(
+ handle="O1",
+ kind="outline",
+ document_id="doc_a",
+ source_file_name="guide.pdf",
+ outline_lines=(" [O1] guide.pdf", " 1 Overview"),
+ ),
+ Candidate(
+ handle="R2.1",
+ kind="read",
+ document_id="doc_a",
+ source_file_name="guide.pdf",
+ section_path="2 Treatment",
+ chunk_ids=("c1", "missing"),
+ ),
+ ]
+ assembled = [
+ {
+ "chunk_id": "c1",
+ "sort_order": 4,
+ "composed": [{"type": "text", "text": "treatment body"}],
+ }
+ ]
+ evidence = compose_pool_evidence(entries, assembled)
+ assert evidence == [
+ {"type": "text", "text": "[E1] [§ guide.pdf]"},
+ {"type": "text", "text": " [O1] guide.pdf\n 1 Overview"},
+ {"type": "text", "text": "\n"},
+ {"type": "text", "text": "[E2] [§ guide.pdf / 2 Treatment]"},
+ {"type": "text", "text": " treatment body"},
+ ]
+
+
+def test_issue_ignores_other_tools_and_errors() -> None:
+ pool = EvidencePool()
+ assert pool.issue("corpus.grep", ToolResult(text="hits")) == []
+ assert pool.issue("corpus.read", ToolResult(text="", error="bad args")) == []
+ assert pool.issue("corpus.assets", ToolResult(text="assets")) == []
+
+
+def test_query_table_records_last_successful_html_and_issues_nothing() -> None:
+ pool = EvidencePool()
+ first = ToolResult(
+ text="first",
+ payload={
+ "document_id": "doc_a",
+ "chunk_id": "t1",
+ "table_html": "
"
+ },
+ )
+ composed_text = "".join(part["text"] for part in parts if part["type"] == "text")
+ assert "tablet" in composed_text
+ assert "FULL TABLE SHOULD NOT APPEAR" not in composed_text
+ assert "[tables/" not in composed_text
+
+
def test_standalone_image_is_bytes(monkeypatch) -> None:
monkeypatch.setattr(
"shared.services.retrieval.hydration.evidence_compose._try_read_image_artifact",
@@ -176,6 +218,43 @@ async def test_text_result_keeps_placeholders_and_composes_assets(monkeypatch) -
assert "[images/" not in composed_text
+@pytest.mark.asyncio
+async def test_assemble_uses_queried_subtable_and_keeps_full_without_map() -> None:
+ full_html = "