surface figure captions in search results, lower search.limit to 5
This commit is contained in:
parent
7d99be471f
commit
514b0c51f8
7 changed files with 147 additions and 1 deletions
|
|
@ -13,6 +13,11 @@
|
||||||
- `llm()` from the analysis sandbox. Sandbox externals are now `search` and `list_documents` only.
|
- `llm()` from the analysis sandbox. Sandbox externals are now `search` and `list_documents` only.
|
||||||
- `list_documents` top-level tool from the analysis skill (still available as `await list_documents()` inside `execute_code`).
|
- `list_documents` top-level tool from the analysis skill (still available as `await list_documents()` inside `execute_code`).
|
||||||
|
|
||||||
|
### Changed
|
||||||
|
|
||||||
|
- `search.limit` default lowered from `10` to `5`. Reduces text + binary noise in vision-tool returns (picture count tracks result count after expansion + dedup); the cite path still selects from all returned chunks.
|
||||||
|
- Search result formatter surfaces picture captions on a labelled line when a chunk's expanded refs include pictures. The OpenAI vision API has no identifier field for binary parts, so the caption is the only signal a model can use to map a description to the figure it sees.
|
||||||
|
|
||||||
### Fixed
|
### Fixed
|
||||||
|
|
||||||
- Chat TUI's state-edit screen now syntax-highlights JSON instead of falling back to plain text. Adds `tree-sitter` + `tree-sitter-json` to the `[tui]` extra.
|
- Chat TUI's state-edit screen now syntax-highlights JSON instead of falling back to plain text. Adds `tree-sitter` + `tree-sitter-json` to the `[tui]` extra.
|
||||||
|
|
|
||||||
|
|
@ -139,14 +139,23 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult])
|
||||||
)
|
)
|
||||||
if not bytes_by_ref:
|
if not bytes_by_ref:
|
||||||
continue
|
continue
|
||||||
|
captions_by_ref = await client.document_item_repository.get_captions_for_chunk(
|
||||||
|
doc_id, list(bytes_by_ref.keys())
|
||||||
|
)
|
||||||
for r in doc_results:
|
for r in doc_results:
|
||||||
attached: dict[str, str] = {}
|
attached: dict[str, str] = {}
|
||||||
|
captions: dict[str, str] = {}
|
||||||
for ref in r.doc_item_refs:
|
for ref in r.doc_item_refs:
|
||||||
blob = bytes_by_ref.get(ref)
|
blob = bytes_by_ref.get(ref)
|
||||||
if blob:
|
if blob:
|
||||||
attached[ref] = base64.b64encode(blob).decode("ascii")
|
attached[ref] = base64.b64encode(blob).decode("ascii")
|
||||||
|
caption = captions_by_ref.get(ref)
|
||||||
|
if caption:
|
||||||
|
captions[ref] = caption
|
||||||
if attached:
|
if attached:
|
||||||
r.image_data = attached
|
r.image_data = attached
|
||||||
|
if captions:
|
||||||
|
r.picture_captions = captions
|
||||||
|
|
||||||
|
|
||||||
async def expand_context(
|
async def expand_context(
|
||||||
|
|
|
||||||
|
|
@ -215,7 +215,7 @@ class ProcessingConfig(BaseModel):
|
||||||
|
|
||||||
|
|
||||||
class SearchConfig(BaseModel):
|
class SearchConfig(BaseModel):
|
||||||
limit: int = 10
|
limit: int = 5
|
||||||
max_context_chars: int = 10000
|
max_context_chars: int = 10000
|
||||||
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
|
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
|
||||||
vector_refine_factor: int = 30
|
vector_refine_factor: int = 30
|
||||||
|
|
|
||||||
|
|
@ -138,6 +138,7 @@ class SearchResult(BaseModel):
|
||||||
headings: list[str] | None = None
|
headings: list[str] | None = None
|
||||||
labels: list[str] = []
|
labels: list[str] = []
|
||||||
image_data: dict[str, str] | None = None
|
image_data: dict[str, str] | None = None
|
||||||
|
picture_captions: dict[str, str] = {}
|
||||||
|
|
||||||
@classmethod
|
@classmethod
|
||||||
def from_chunk(
|
def from_chunk(
|
||||||
|
|
@ -198,6 +199,15 @@ class SearchResult(BaseModel):
|
||||||
if primary_label:
|
if primary_label:
|
||||||
parts.append(f"Type: {primary_label}")
|
parts.append(f"Type: {primary_label}")
|
||||||
|
|
||||||
|
# Surface picture captions when present. Order matches the binary
|
||||||
|
# attachments emitted by build_binary_parts_from_results, so the model
|
||||||
|
# can correlate caption ↔ attached image by position (BinaryContent
|
||||||
|
# identifiers don't survive serialization to the OpenAI vision API).
|
||||||
|
if self.picture_captions:
|
||||||
|
for self_ref, caption in self.picture_captions.items():
|
||||||
|
if caption:
|
||||||
|
parts.append(f"Figure caption ({self_ref}): {caption}")
|
||||||
|
|
||||||
# The actual content
|
# The actual content
|
||||||
parts.append(f"Content:\n{self.content}")
|
parts.append(f"Content:\n{self.content}")
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -206,3 +206,32 @@ class DocumentItemRepository:
|
||||||
if data:
|
if data:
|
||||||
result[row["self_ref"]] = data
|
result[row["self_ref"]] = data
|
||||||
return result
|
return result
|
||||||
|
|
||||||
|
async def get_captions_for_chunk(
|
||||||
|
self, document_id: str, refs: list[str]
|
||||||
|
) -> dict[str, str]:
|
||||||
|
"""Fetch caption text for multiple self_refs within a single document.
|
||||||
|
|
||||||
|
Returns ``{self_ref: text}`` for refs that have non-empty text. Used
|
||||||
|
alongside ``get_pictures_for_chunk`` to label figures in agent-facing
|
||||||
|
search results — the OpenAI vision message format has no identifier
|
||||||
|
field for binary parts, so the caption is the only signal a model can
|
||||||
|
use to correlate a description with the picture it sees.
|
||||||
|
"""
|
||||||
|
if not refs:
|
||||||
|
return {}
|
||||||
|
|
||||||
|
safe_id = escape_sql_string(document_id)
|
||||||
|
refs_sql = ", ".join(f"'{escape_sql_string(r)}'" for r in refs)
|
||||||
|
rows = await (
|
||||||
|
self.store.document_items_table.query()
|
||||||
|
.select(["self_ref", "text"])
|
||||||
|
.where(f"document_id = '{safe_id}' AND self_ref IN ({refs_sql})")
|
||||||
|
.to_list()
|
||||||
|
)
|
||||||
|
result: dict[str, str] = {}
|
||||||
|
for row in rows:
|
||||||
|
text = row.get("text") or ""
|
||||||
|
if text:
|
||||||
|
result[row["self_ref"]] = text
|
||||||
|
return result
|
||||||
|
|
|
||||||
|
|
@ -502,6 +502,55 @@ class TestPictureDataStorage:
|
||||||
# Empty refs returns empty dict
|
# Empty refs returns empty dict
|
||||||
assert await repo.get_pictures_for_chunk("doc-1", []) == {}
|
assert await repo.get_pictures_for_chunk("doc-1", []) == {}
|
||||||
|
|
||||||
|
async def test_get_captions_for_chunk(self, temp_db_path):
|
||||||
|
"""Captions are returned for refs whose text is non-empty.
|
||||||
|
|
||||||
|
In practice pictures carry their caption in the ``text`` field
|
||||||
|
(populated by the VLM picture-description pass during ingest); this
|
||||||
|
method surfaces that text alongside the picture bytes so the model can
|
||||||
|
correlate a description with the binary it sees.
|
||||||
|
"""
|
||||||
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||||
|
repo = DocumentItemRepository(rag.store)
|
||||||
|
await repo.create_items(
|
||||||
|
"doc-1",
|
||||||
|
[
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=0,
|
||||||
|
self_ref="#/pictures/0",
|
||||||
|
label="picture",
|
||||||
|
text="Figure 1. CCS generation over time.",
|
||||||
|
picture_data=b"\x89PNG\r\n\x1a\nfake",
|
||||||
|
),
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=1,
|
||||||
|
self_ref="#/pictures/1",
|
||||||
|
label="picture",
|
||||||
|
text="", # no VLM caption available
|
||||||
|
picture_data=b"\x89PNG\r\n\x1a\nfake2",
|
||||||
|
),
|
||||||
|
DocumentItem(
|
||||||
|
document_id="doc-1",
|
||||||
|
position=2,
|
||||||
|
self_ref="#/texts/0",
|
||||||
|
label="paragraph",
|
||||||
|
text="Inline prose.",
|
||||||
|
),
|
||||||
|
],
|
||||||
|
)
|
||||||
|
|
||||||
|
captions = await repo.get_captions_for_chunk(
|
||||||
|
"doc-1",
|
||||||
|
["#/pictures/0", "#/pictures/1", "#/texts/0", "#/pictures/999"],
|
||||||
|
)
|
||||||
|
assert captions == {
|
||||||
|
"#/pictures/0": "Figure 1. CCS generation over time.",
|
||||||
|
"#/texts/0": "Inline prose.",
|
||||||
|
}
|
||||||
|
assert await repo.get_captions_for_chunk("doc-1", []) == {}
|
||||||
|
|
||||||
async def test_hot_paths_exclude_picture_data(self, temp_db_path):
|
async def test_hot_paths_exclude_picture_data(self, temp_db_path):
|
||||||
"""Light read paths must NOT pull picture_data into memory."""
|
"""Light read paths must NOT pull picture_data into memory."""
|
||||||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||||
|
|
|
||||||
|
|
@ -279,6 +279,50 @@ def test_search_result_format_for_agent_with_rank():
|
||||||
assert "Content:\nThis is the chunk content about elections." in formatted
|
assert "Content:\nThis is the chunk content about elections." in formatted
|
||||||
|
|
||||||
|
|
||||||
|
def test_search_result_format_for_agent_picture_captions():
|
||||||
|
"""Picture captions render as labelled lines so the model can correlate them
|
||||||
|
with binary parts (BinaryContent.identifier doesn't survive serialization
|
||||||
|
to the OpenAI vision API; insertion order is the only reliable signal)."""
|
||||||
|
result = SearchResult(
|
||||||
|
content="...surrounding text...",
|
||||||
|
score=0.5,
|
||||||
|
chunk_id="chunk-xyz",
|
||||||
|
labels=["picture", "text"],
|
||||||
|
picture_captions={
|
||||||
|
"#/pictures/0": "Figure 1. Results from each model.",
|
||||||
|
"#/pictures/1": "Figure 2. Projected annual emissions.",
|
||||||
|
},
|
||||||
|
)
|
||||||
|
|
||||||
|
formatted = result.format_for_agent(rank=1, total=2)
|
||||||
|
|
||||||
|
lines = formatted.splitlines()
|
||||||
|
cap0 = next(
|
||||||
|
i for i, line in enumerate(lines) if "Figure caption (#/pictures/0)" in line
|
||||||
|
)
|
||||||
|
cap1 = next(
|
||||||
|
i for i, line in enumerate(lines) if "Figure caption (#/pictures/1)" in line
|
||||||
|
)
|
||||||
|
content_line = next(
|
||||||
|
i for i, line in enumerate(lines) if line.startswith("Content:")
|
||||||
|
)
|
||||||
|
assert cap0 < cap1 < content_line
|
||||||
|
assert "Figure 1. Results from each model." in formatted
|
||||||
|
assert "Figure 2. Projected annual emissions." in formatted
|
||||||
|
|
||||||
|
|
||||||
|
def test_search_result_format_for_agent_no_captions_no_line():
|
||||||
|
"""Without picture_captions, no caption lines appear (zero-overhead for text chunks)."""
|
||||||
|
result = SearchResult(
|
||||||
|
content="prose",
|
||||||
|
score=0.5,
|
||||||
|
chunk_id="chunk-abc",
|
||||||
|
labels=["text"],
|
||||||
|
)
|
||||||
|
formatted = result.format_for_agent(rank=1, total=1)
|
||||||
|
assert "Figure caption" not in formatted
|
||||||
|
|
||||||
|
|
||||||
def test_search_result_format_for_agent_rank_only():
|
def test_search_result_format_for_agent_rank_only():
|
||||||
"""Test format_for_agent with rank but no total."""
|
"""Test format_for_agent with rank but no total."""
|
||||||
result = SearchResult(
|
result = SearchResult(
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue