Update LLM prompts/tools to reflect we will pass title (if available) or URI

This commit is contained in:
Yiorgis Gozadinos 2025-09-22 11:24:32 +03:00
parent 2c1e81dc6e
commit 73aff1695f
No known key found for this signature in database
6 changed files with 47 additions and 10 deletions

View file

@ -12,7 +12,9 @@ from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIO
class SearchResult(BaseModel): class SearchResult(BaseModel):
content: str = Field(description="The document text content") content: str = Field(description="The document text content")
score: float = Field(description="Relevance score (higher is more relevant)") 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): class Dependencies(BaseModel):
@ -59,7 +61,7 @@ class QuestionAnswerAgent:
SearchResult( SearchResult(
content=chunk.content, content=chunk.content,
score=score, score=score,
document_uri=chunk.document_uri or "", document_uri=(chunk.document_title or chunk.document_uri or ""),
) )
for chunk, score in expanded_results for chunk, score in expanded_results
] ]

View file

@ -44,9 +44,9 @@ Guidelines:
Citation Format: Citation Format:
After your answer, include a "Citations:" section that lists: 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 - 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: Example response format:
[Your answer here] [Your answer here]

View file

@ -19,8 +19,8 @@ class SearchAnswer(BaseModel):
) )
sources: list[str] = Field( sources: list[str] = Field(
description=( description=(
"Document URIs corresponding to the snippets actually used in the" "Document titles (if available) or URIs corresponding to the"
" answer (one URI per snippet; omit if none)" " snippets actually used in the answer (one per snippet; omit if none)"
), ),
default_factory=list, default_factory=list,
) )

View file

@ -59,7 +59,9 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
{ {
"text": chunk.content, "text": chunk.content,
"score": score, "score": score,
"document_uri": (chunk.document_uri or ""), "document_uri": (
chunk.document_title or chunk.document_uri or ""
),
} }
for chunk, score in expanded for chunk, score in expanded
] ]

View file

@ -27,13 +27,14 @@ Tasks:
Tool usage: Tool usage:
- Always call search_and_answer before drafting any answer. - Always call search_and_answer before drafting any answer.
- The tool returns snippets with verbatim `text`, a relevance `score`, and the - 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 - You may call the tool multiple times to refine or broaden context, but do not
exceed 3 total calls. Favor precision over volume. exceed 3 total calls. Favor precision over volume.
- Use scores to prioritize evidence, but include only the minimal subset of - Use scores to prioritize evidence, but include only the minimal subset of
snippet texts (verbatim) in SearchAnswer.context (typically 14). snippet texts (verbatim) in SearchAnswer.context (typically 14).
- Set SearchAnswer.sources to the corresponding document_uris for the snippets - Set SearchAnswer.sources to the corresponding document identifiers for the
you used (one URI per snippet; same order as context). Context must be textonly. snippets you used (title if available, otherwise URI; one per snippet; same
order as context). Context must be textonly.
- If no relevant information is found, clearly say so and return an empty - If no relevant information is found, clearly say so and return an empty
context list and sources list. context list and sources list.

View file

@ -106,6 +106,38 @@ async def test_chunks_include_document_info(temp_db_path):
store.close() 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 @pytest.mark.asyncio
async def test_search_score_types(temp_db_path): async def test_search_score_types(temp_db_path):
"""Test that different search types return appropriate score ranges.""" """Test that different search types return appropriate score ranges."""