Update LLM prompts/tools to reflect we will pass title (if available) or URI
This commit is contained in:
parent
2c1e81dc6e
commit
73aff1695f
6 changed files with 47 additions and 10 deletions
|
|
@ -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
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -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]
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
)
|
)
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
]
|
]
|
||||||
|
|
|
||||||
|
|
@ -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 1‑4).
|
snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
|
||||||
- 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 text‑only.
|
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
|
- If no relevant information is found, clearly say so and return an empty
|
||||||
context list and sources list.
|
context list and sources list.
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue