From 73aff1695f345686a900cc7d687b47e467503a2f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 22 Sep 2025 11:24:32 +0300 Subject: [PATCH] Update LLM prompts/tools to reflect we will pass title (if available) or URI --- src/haiku/rag/qa/agent.py | 6 +++-- src/haiku/rag/qa/prompts.py | 4 ++-- src/haiku/rag/research/models.py | 4 ++-- src/haiku/rag/research/nodes/search.py | 4 +++- src/haiku/rag/research/prompts.py | 7 +++--- tests/test_search.py | 32 ++++++++++++++++++++++++++ 6 files changed, 47 insertions(+), 10 deletions(-) diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index 3493b0ce..b3ad2248 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -12,7 +12,9 @@ from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIO class SearchResult(BaseModel): content: str = Field(description="The document text content") score: float = Field(description="Relevance score (higher is more relevant)") - document_uri: str = Field(description="Source URI/path of the document") + document_uri: str = Field( + description="Source title (if available) or URI/path of the document" + ) class Dependencies(BaseModel): @@ -59,7 +61,7 @@ class QuestionAnswerAgent: SearchResult( content=chunk.content, score=score, - document_uri=chunk.document_uri or "", + document_uri=(chunk.document_title or chunk.document_uri or ""), ) for chunk, score in expanded_results ] diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index 19896d2c..b25b36af 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/src/haiku/rag/qa/prompts.py @@ -44,9 +44,9 @@ Guidelines: Citation Format: After your answer, include a "Citations:" section that lists: -- The document URI from each search result used +- The document title (if available) or URI from each search result used - A brief excerpt (first 50-100 characters) of the content that supported your answer -- Format: "Citations:\n- [document_uri]: [content_excerpt]..." +- Format: "Citations:\n- [document title or URI]: [content_excerpt]..." Example response format: [Your answer here] diff --git a/src/haiku/rag/research/models.py b/src/haiku/rag/research/models.py index 3f789842..30cb97d1 100644 --- a/src/haiku/rag/research/models.py +++ b/src/haiku/rag/research/models.py @@ -19,8 +19,8 @@ class SearchAnswer(BaseModel): ) sources: list[str] = Field( description=( - "Document URIs corresponding to the snippets actually used in the" - " answer (one URI per snippet; omit if none)" + "Document titles (if available) or URIs corresponding to the" + " snippets actually used in the answer (one per snippet; omit if none)" ), default_factory=list, ) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py index 0943338a..ae863d37 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/research/nodes/search.py @@ -59,7 +59,9 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): { "text": chunk.content, "score": score, - "document_uri": (chunk.document_uri or ""), + "document_uri": ( + chunk.document_title or chunk.document_uri or "" + ), } for chunk, score in expanded ] diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index 6d9aad48..29e10424 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -27,13 +27,14 @@ Tasks: Tool usage: - Always call search_and_answer before drafting any answer. - The tool returns snippets with verbatim `text`, a relevance `score`, and the - originating `document_uri`. + originating document identifier (document title if available, otherwise URI). - You may call the tool multiple times to refine or broaden context, but do not exceed 3 total calls. Favor precision over volume. - Use scores to prioritize evidence, but include only the minimal subset of snippet texts (verbatim) in SearchAnswer.context (typically 1‑4). -- Set SearchAnswer.sources to the corresponding document_uris for the snippets - you used (one URI per snippet; same order as context). Context must be text‑only. +- Set SearchAnswer.sources to the corresponding document identifiers for the + snippets you used (title if available, otherwise URI; one per snippet; same + order as context). Context must be text‑only. - If no relevant information is found, clearly say so and return an empty context list and sources list. diff --git a/tests/test_search.py b/tests/test_search.py index f97086e9..683649c0 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -106,6 +106,38 @@ async def test_chunks_include_document_info(temp_db_path): store.close() +@pytest.mark.asyncio +async def test_chunks_include_document_title(temp_db_path): + """Test that search results include the parent document title when present.""" + store = Store(temp_db_path) + doc_repo = DocumentRepository(store) + chunk_repo = ChunkRepository(store) + + # Create a document with URI and title + document = Document( + content="This is a test document with a custom title to verify enrichment.", + uri="file:///tmp/title-test.md", + title="My Custom Title", + ) + + # Create the document with chunks + from haiku.rag.utils import text_to_docling_document + + dl = text_to_docling_document(document.content, name="title-test.md") + await doc_repo._create_with_docling(document, dl) + + # Perform a search that should find this document + results = await chunk_repo.search("custom title", limit=3, search_type="hybrid") + + assert results, "Expected at least one search result" + for chunk, _ in results: + # All returned chunks for this doc should carry the document title + if chunk.document_uri == "file:///tmp/title-test.md": + assert chunk.document_title == "My Custom Title" + + store.close() + + @pytest.mark.asyncio async def test_search_score_types(temp_db_path): """Test that different search types return appropriate score ranges."""