Name the collection on a retrieved image, not only in the text
Search results and capsule entries name the collection they came from; the images attached beside them carried only the chunk id and reference. Two collections can return the same picture of the same document, so the two labels were identical and the model could place neither. The decision is the one already made for the text: `covers_multiple` at the generic search tool, and the flag `search_corpus` computed for the capability tools, which it now returns.
This commit is contained in:
Yiorgis Gozadinos
parent
bfd09f6880
commit
fa319596cc
8 changed files with 204 additions and 26 deletions
|
|
@ -505,7 +505,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
|||
"the results you already have."
|
||||
)
|
||||
async with self.rag_lock:
|
||||
formatted, results = await search_corpus(
|
||||
formatted, results, include_collection = await search_corpus(
|
||||
await self._ensure_rag(),
|
||||
query,
|
||||
limit=limit,
|
||||
|
|
@ -517,7 +517,11 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
|||
# narrower return must not drop what the wider one already showed it.
|
||||
merge_results(state.searches.setdefault(query, []), results)
|
||||
self._note_evidence()
|
||||
if self.vision and (parts := build_image_content_from_results(results)):
|
||||
if self.vision and (
|
||||
parts := build_image_content_from_results(
|
||||
results, include_collection=include_collection
|
||||
)
|
||||
):
|
||||
return ToolReturn(return_value=formatted, content=parts)
|
||||
return formatted
|
||||
|
||||
|
|
|
|||
|
|
@ -19,8 +19,8 @@ async def search_corpus(
|
|||
limit: int | None = None,
|
||||
document_filter: str | None = None,
|
||||
sources: list[str] | None = None,
|
||||
) -> tuple[str, list[SearchResult]]:
|
||||
"""Search and context-expand results for a capability tool."""
|
||||
) -> tuple[str, list[SearchResult], bool]:
|
||||
"""Search and context-expand results, and whether they name their collection."""
|
||||
results = await rag.search(
|
||||
query, limit=limit, filter=document_filter, sources=sources
|
||||
)
|
||||
|
|
@ -35,7 +35,7 @@ async def search_corpus(
|
|||
)
|
||||
for index, result in enumerate(results)
|
||||
)
|
||||
return formatted or "No results found.", list(results)
|
||||
return formatted or "No results found.", list(results), include_collection
|
||||
|
||||
|
||||
def merge_results(
|
||||
|
|
|
|||
|
|
@ -49,10 +49,12 @@ def group_label(position: int) -> str:
|
|||
return f"[Cited evidence group {position}]"
|
||||
|
||||
|
||||
def picture_label(chunk_id: str, self_ref: str) -> str:
|
||||
def picture_label(chunk_id: str, self_ref: str, collection: str | None = None) -> str:
|
||||
named = f"Collection: {collection}. " if collection else ""
|
||||
return (
|
||||
f"Page image retrieved from the knowledge base for cited evidence "
|
||||
f"[{chunk_id}] ({self_ref}). Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
||||
f"[{chunk_id}] ({self_ref}). {named}"
|
||||
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -191,7 +193,11 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule:
|
|||
chunk_id=entry.chunk_id,
|
||||
document_id=entry.citation.document_id,
|
||||
self_ref=self_ref,
|
||||
label=picture_label(entry.chunk_id, self_ref),
|
||||
label=picture_label(
|
||||
entry.chunk_id,
|
||||
self_ref,
|
||||
entry.citation.source if include_collection else None,
|
||||
),
|
||||
source=entry.citation.source,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
|
|||
|
||||
def build_image_content_from_results(
|
||||
results: list[SearchResult],
|
||||
include_collection: bool = False,
|
||||
) -> list[str | BinaryContent]:
|
||||
"""Decode and validate picture bytes attached to search results, labelled.
|
||||
|
||||
|
|
@ -57,7 +58,7 @@ def build_image_content_from_results(
|
|||
``BinaryContent.identifier`` cannot do — it does not survive serialization
|
||||
to the vision API.
|
||||
"""
|
||||
collected: list[tuple[str | None, str, BinaryContent]] = []
|
||||
collected: list[tuple[str | None, str | None, str, BinaryContent]] = []
|
||||
seen: set[tuple[str | None, str | None, str]] = set()
|
||||
for result in results:
|
||||
if not result.image_data:
|
||||
|
|
@ -69,15 +70,16 @@ def build_image_content_from_results(
|
|||
picture = decode_picture(base64.b64decode(b64), self_ref)
|
||||
if picture is None:
|
||||
continue
|
||||
collected.append((result.chunk_id, self_ref, picture))
|
||||
collected.append((result.source, result.chunk_id, self_ref, picture))
|
||||
seen.add(key)
|
||||
|
||||
content: list[str | BinaryContent] = []
|
||||
total = len(collected)
|
||||
for position, (chunk_id, self_ref, picture) in enumerate(collected, 1):
|
||||
for position, (source, chunk_id, self_ref, picture) in enumerate(collected, 1):
|
||||
collection = f"Collection: {source}. " if include_collection and source else ""
|
||||
content.append(
|
||||
f"Page image {position} of {total}, retrieved from the knowledge base "
|
||||
f"for search result [{chunk_id}] ({self_ref}). "
|
||||
f"for search result [{chunk_id}] ({self_ref}). {collection}"
|
||||
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
|
||||
)
|
||||
content.append(picture)
|
||||
|
|
@ -172,7 +174,9 @@ def create_search_toolset(
|
|||
if not config.qa.model.vision:
|
||||
return text
|
||||
|
||||
image_content = build_image_content_from_results(results_list)
|
||||
image_content = build_image_content_from_results(
|
||||
results_list, include_collection=include_collection
|
||||
)
|
||||
if image_content:
|
||||
return ToolReturn(return_value=text, content=image_content)
|
||||
return text
|
||||
|
|
|
|||
|
|
@ -425,6 +425,59 @@ def _stub_client(*batches: list[SearchResult]) -> AsyncMock:
|
|||
return client
|
||||
|
||||
|
||||
def _picture_result(source: str) -> SearchResult:
|
||||
import base64
|
||||
from io import BytesIO
|
||||
|
||||
from PIL import Image as PILImage
|
||||
|
||||
buffer = BytesIO()
|
||||
PILImage.new("RGB", (4, 4), "red").save(buffer, format="PNG")
|
||||
return SearchResult(
|
||||
content="body",
|
||||
score=0.9,
|
||||
source=source,
|
||||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
image_data={"#/pictures/0": base64.b64encode(buffer.getvalue()).decode()},
|
||||
)
|
||||
|
||||
|
||||
async def _labels_of_search(temp_db_path, *sources: str) -> list[str]:
|
||||
"""The labels the search tool attaches to its images, for a client covering
|
||||
`sources`."""
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig(), vision=True)
|
||||
capability.state = RAGState()
|
||||
client = _stub_client([_picture_result(sources[0])])
|
||||
client.source_names = sources
|
||||
|
||||
with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)):
|
||||
returned = await capability._search("cats", None)
|
||||
|
||||
assert isinstance(returned, ToolReturn)
|
||||
assert returned.content is not None
|
||||
return [item for item in returned.content if isinstance(item, str)]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_search_spanning_collections_names_them_on_its_images(temp_db_path):
|
||||
"""Images travel beside the results and are labelled the same way."""
|
||||
labels = await _labels_of_search(temp_db_path, "alpha", "beta")
|
||||
|
||||
assert "Collection: alpha." in labels[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_search_over_one_collection_does_not_name_it_on_its_images(
|
||||
temp_db_path,
|
||||
):
|
||||
labels = await _labels_of_search(temp_db_path, "alpha")
|
||||
|
||||
assert not [label for label in labels if "Collection" in label]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_fruitless_search_says_so(temp_db_path):
|
||||
"""A blank tool return reads as a broken tool, not as an empty corpus."""
|
||||
|
|
|
|||
|
|
@ -161,6 +161,50 @@ def test_a_picture_in_two_collections_is_retained_from_each():
|
|||
("papers", "#/pictures/0"),
|
||||
("wiki", "#/pictures/0"),
|
||||
]
|
||||
# Nothing else tells the two apart once they are attached.
|
||||
assert "Collection: papers." in capsule.pictures[0].label
|
||||
assert "Collection: wiki." in capsule.pictures[1].label
|
||||
|
||||
|
||||
def test_one_chunk_id_cited_from_two_collections_labels_each_picture():
|
||||
"""Both capabilities can cite the same id from different collections, where
|
||||
the reference and the document are the same too."""
|
||||
found = [
|
||||
replace(
|
||||
discovered(
|
||||
capability=capability,
|
||||
cited={"c1": [2]},
|
||||
pictures={"c1": ["#/pictures/0"]},
|
||||
),
|
||||
citations={
|
||||
"c1": replace_citation(
|
||||
citation("c1", pictures=["#/pictures/0"]),
|
||||
document_id="shared",
|
||||
source=source,
|
||||
)
|
||||
},
|
||||
)
|
||||
for capability, source in (("analysis", "papers"), ("rag", "wiki"))
|
||||
]
|
||||
|
||||
capsule = build_capsule(found)
|
||||
|
||||
labels = [picture.label for picture in capsule.pictures]
|
||||
assert all("[c1] (#/pictures/0)" in label for label in labels)
|
||||
assert "Collection: papers." in labels[0]
|
||||
assert "Collection: wiki." in labels[1]
|
||||
|
||||
|
||||
def test_a_picture_from_one_collection_is_not_labelled_with_it():
|
||||
found = discovered(cited={"c1": [2]}, pictures={"c1": ["#/pictures/0"]})
|
||||
found = replace(
|
||||
found,
|
||||
citations={"c1": replace_citation(found.citations["c1"], source="papers")},
|
||||
)
|
||||
|
||||
[picture] = build_capsule([found]).pictures
|
||||
|
||||
assert "Collection" not in picture.label
|
||||
|
||||
|
||||
def test_nothing_cited_produces_no_capsule():
|
||||
|
|
|
|||
|
|
@ -306,11 +306,13 @@ class TestWhenTheModelIsToldTheCollection:
|
|||
async with HaikuRAG(config=config) as rag:
|
||||
monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
|
||||
|
||||
spanning, _ = await search_corpus(rag, "cats")
|
||||
narrowed, _ = await search_corpus(rag, "cats", sources=["alpha"])
|
||||
spanning, _, spans = await search_corpus(rag, "cats")
|
||||
narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"])
|
||||
|
||||
assert "Collection: alpha" in spanning
|
||||
assert "Collection" not in narrowed
|
||||
# Images travel beside the results and are labelled the same way.
|
||||
assert (spans, narrows) == (True, False)
|
||||
|
||||
|
||||
class TestActionableFailures:
|
||||
|
|
|
|||
|
|
@ -42,7 +42,7 @@ class TestNamingTheCollection:
|
|||
what the search spans."""
|
||||
|
||||
@staticmethod
|
||||
def _client(covers_multiple: bool, source: str | None):
|
||||
def _client(covers_multiple: bool, source: str | None, *, pictures: bool = False):
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
results = [
|
||||
|
|
@ -53,6 +53,7 @@ class TestNamingTheCollection:
|
|||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
document_title="Report",
|
||||
image_data={"#/pictures/0": _png_b64()} if pictures else None,
|
||||
)
|
||||
]
|
||||
return SimpleNamespace(
|
||||
|
|
@ -79,6 +80,42 @@ class TestNamingTheCollection:
|
|||
|
||||
assert "Collection" not in text
|
||||
|
||||
@staticmethod
|
||||
def _seeing_config(search_config):
|
||||
config = search_config.model_copy(deep=True)
|
||||
config.qa.model.vision = True
|
||||
return config
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_client_covering_a_set_names_each_image(self, search_config):
|
||||
"""Images travel beside the results and are labelled the same way."""
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
toolset = create_search_toolset(self._seeing_config(search_config))
|
||||
client = self._client(covers_multiple=True, source="alpha", pictures=True)
|
||||
|
||||
returned = await toolset.tools["search"].function(make_ctx(client), "cats")
|
||||
|
||||
assert isinstance(returned, ToolReturn)
|
||||
labels = [item for item in returned.content if isinstance(item, str)]
|
||||
assert "Collection: alpha." in labels[0]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_named_collection_is_not_named_on_an_image(self, search_config):
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
|
||||
toolset = create_search_toolset(self._seeing_config(search_config))
|
||||
client = self._client(covers_multiple=False, source="alpha", pictures=True)
|
||||
|
||||
returned = await toolset.tools["search"].function(make_ctx(client), "cats")
|
||||
|
||||
assert isinstance(returned, ToolReturn)
|
||||
assert not [
|
||||
item
|
||||
for item in returned.content
|
||||
if isinstance(item, str) and "Collection" in item
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.vcr()
|
||||
class TestSearchToolExecution:
|
||||
|
|
@ -304,14 +341,11 @@ class TestBuildImageContentFromResults:
|
|||
images = [item for item in content if isinstance(item, BinaryContent)]
|
||||
assert len(images) == 1
|
||||
|
||||
def test_the_same_picture_in_two_collections_is_attached_from_each(self):
|
||||
"""A document copied into another collection keeps its id and picture refs."""
|
||||
from pydantic_ai.messages import BinaryContent
|
||||
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
@staticmethod
|
||||
def _one_picture_in_two_collections():
|
||||
"""A document copied into another collection keeps its ids and refs."""
|
||||
shared = {"#/pictures/0": _png_b64()}
|
||||
results = [
|
||||
return [
|
||||
SearchResult(
|
||||
content="a",
|
||||
score=0.9,
|
||||
|
|
@ -321,20 +355,51 @@ class TestBuildImageContentFromResults:
|
|||
image_data=shared,
|
||||
),
|
||||
SearchResult(
|
||||
content="b",
|
||||
content="a",
|
||||
score=0.8,
|
||||
chunk_id="c2",
|
||||
chunk_id="c1",
|
||||
document_id="doc-1",
|
||||
source="wiki",
|
||||
image_data=shared,
|
||||
),
|
||||
]
|
||||
|
||||
content = build_image_content_from_results(results)
|
||||
def test_the_same_picture_in_two_collections_is_attached_from_each(self):
|
||||
from pydantic_ai.messages import BinaryContent
|
||||
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
content = build_image_content_from_results(
|
||||
self._one_picture_in_two_collections()
|
||||
)
|
||||
|
||||
images = [item for item in content if isinstance(item, BinaryContent)]
|
||||
assert len(images) == 2
|
||||
|
||||
def test_each_image_is_labelled_with_the_collection_it_came_from(self):
|
||||
"""Nothing else tells the two apart: same chunk id, same document, same
|
||||
reference."""
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
content = build_image_content_from_results(
|
||||
self._one_picture_in_two_collections(), include_collection=True
|
||||
)
|
||||
|
||||
labels = [item for item in content if isinstance(item, str)]
|
||||
assert "Collection: papers." in labels[0]
|
||||
assert "Collection: wiki." in labels[1]
|
||||
|
||||
def test_an_unasked_for_collection_is_not_named_on_an_image(self):
|
||||
from haiku.rag.tools.search import build_image_content_from_results
|
||||
|
||||
content = build_image_content_from_results(
|
||||
self._one_picture_in_two_collections()
|
||||
)
|
||||
|
||||
assert not [
|
||||
item for item in content if isinstance(item, str) and "Collection" in item
|
||||
]
|
||||
|
||||
def test_each_image_is_labelled_with_the_result_it_belongs_to(self):
|
||||
"""Label every picture, not just the batch.
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue