diff --git a/CHANGELOG.md b/CHANGELOG.md index d798a5fb..02f4ab89 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,11 @@ - `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`). +### 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 - 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. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index c2f0a10a..a712ac40 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -139,14 +139,23 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult]) ) if not bytes_by_ref: 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: attached: dict[str, str] = {} + captions: dict[str, str] = {} for ref in r.doc_item_refs: blob = bytes_by_ref.get(ref) if blob: attached[ref] = base64.b64encode(blob).decode("ascii") + caption = captions_by_ref.get(ref) + if caption: + captions[ref] = caption if attached: r.image_data = attached + if captions: + r.picture_captions = captions async def expand_context( diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 13436e60..bf7cb4e6 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -215,7 +215,7 @@ class ProcessingConfig(BaseModel): class SearchConfig(BaseModel): - limit: int = 10 + limit: int = 5 max_context_chars: int = 10000 vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine" vector_refine_factor: int = 30 diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index 5e86ec80..d6d3ba6c 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -138,6 +138,7 @@ class SearchResult(BaseModel): headings: list[str] | None = None labels: list[str] = [] image_data: dict[str, str] | None = None + picture_captions: dict[str, str] = {} @classmethod def from_chunk( @@ -198,6 +199,15 @@ class SearchResult(BaseModel): if 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 parts.append(f"Content:\n{self.content}") diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 7dde27b5..09052e98 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -206,3 +206,32 @@ class DocumentItemRepository: if data: result[row["self_ref"]] = data 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 diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index b6741834..145fe0fc 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -502,6 +502,55 @@ class TestPictureDataStorage: # Empty refs returns empty dict 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): """Light read paths must NOT pull picture_data into memory.""" async with HaikuRAG(temp_db_path, create=True) as rag: diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 4d37945e..f0679d1f 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -279,6 +279,50 @@ def test_search_result_format_for_agent_with_rank(): 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(): """Test format_for_agent with rank but no total.""" result = SearchResult(