danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Give the conversational agent a list documents and a summarize tools

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-27 16:39:48 +02:00
parent 8a587ec554
commit cca03dbe61
No known key found for this signature in database
10 changed files with 2532 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
CHANGELOG.md
View file

@ -3,6 +3,9 @@
### Added
- **Chat Agent Document Awareness Tools**: Two new tools for browsing and understanding the knowledge base
- `list_documents` — Browse available documents with title, URI, and creation date; respects session document filter
- `summarize_document` — Generate LLM-powered summaries of specific documents
- **Read-Only Initial Context**: Initial context is now locked after the first message, providing consistent session context
- Chat TUI: `--initial-context` CLI option sets background context for the session
- Context can be edited via command palette before the first message is sent

6
docs/agents.md
View file

@ -61,11 +61,13 @@ Key features:
### Tools
The chat agent uses three tools:
The chat agent uses five tools:
- `list_documents` — Browse available documents in the knowledge base
- `summarize_document` — Generate a summary of a specific document
- `get_document` — Retrieve a specific document by title or URI
- `search` — Hybrid search with optional document filter
- `ask` — Answer questions using the conversational research graph (automatically recalls prior answers)
- `get_document` — Retrieve a specific document by title or URI
The `ask` tool automatically checks conversation history before running research. It uses embedding similarity (0.7 cosine threshold) to find semantically matching prior answers, which are passed to the research planner as context. When prior answers are sufficient, the planner can skip searching entirely.

125
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -8,7 +8,7 @@ from haiku.rag.agents.chat.context import (
get_cached_session_context,
update_session_context,
)
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT, DOCUMENT_SUMMARY_PROMPT
from haiku.rag.agents.chat.search import SearchAgent
from haiku.rag.agents.chat.state import (
MAX_QA_HISTORY,
@ -23,6 +23,7 @@ from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_conversational_graph
from haiku.rag.agents.research.models import Citation
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.embeddings import get_embedder
from haiku.rag.utils import get_model
@ -369,6 +370,72 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
],
)
@agent.tool
async def list_documents(
ctx: RunContext[ChatDeps],
limit: int | None = None,
offset: int | None = None,
) -> str:
"""List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available.
Args:
limit: Maximum number of documents to return
offset: Number of documents to skip (for pagination)
"""
# Build session filter from document_filter
doc_filter = None
if ctx.deps.session_state and ctx.deps.session_state.document_filter:
doc_filter = build_multi_document_filter(
ctx.deps.session_state.document_filter
)
docs = await ctx.deps.client.list_documents(
limit=limit, offset=offset, filter=doc_filter
)
if not docs:
return "No documents found in the knowledge base."
lines = [f"Found {len(docs)} document(s):\n"]
for doc in docs:
title = doc.title or "Untitled"
uri = doc.uri or "N/A"
created = doc.created_at.strftime("%Y-%m-%d")
lines.append(f"- **{title}** (URI: {uri}, Created: {created})")
return "\n".join(lines)
async def _find_document(client: HaikuRAG, query: str):
"""Find a document by exact URI, partial URI, or partial title match."""
# Try exact URI match first
doc = await client.get_document_by_uri(query)
if doc is not None:
return doc
escaped_query = query.replace("'", "''")
# Also try without spaces for matching "TB MED 593" to "tbmed593"
no_spaces = escaped_query.replace(" ", "")
# Try partial URI match (with and without spaces)
docs = await client.list_documents(
limit=1,
filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%') OR LOWER(uri) LIKE LOWER('%{no_spaces}%')",
)
if docs:
return docs[0]
# Try partial title match (with and without spaces)
docs = await client.list_documents(
limit=1,
filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')",
)
if docs:
return docs[0]
return None
@agent.tool
async def get_document(
ctx: RunContext[ChatDeps],
@ -381,30 +448,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
Args:
query: The document title or URI to look up
"""
# Try exact URI match first
doc = await ctx.deps.client.get_document_by_uri(query)
escaped_query = query.replace("'", "''")
# Also try without spaces for matching "TB MED 593" to "tbmed593"
no_spaces = escaped_query.replace(" ", "")
# If not found, try partial URI match (with and without spaces)
if doc is None:
docs = await ctx.deps.client.list_documents(
limit=1,
filter=f"LOWER(uri) LIKE LOWER('%{escaped_query}%') OR LOWER(uri) LIKE LOWER('%{no_spaces}%')",
)
if docs:
doc = docs[0]
# If still not found, try partial title match (with and without spaces)
if doc is None:
docs = await ctx.deps.client.list_documents(
limit=1,
filter=f"LOWER(title) LIKE LOWER('%{escaped_query}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')",
)
if docs:
doc = docs[0]
doc = await _find_document(ctx.deps.client, query)
if doc is None:
return f"Document not found: {query}"
@ -412,9 +456,38 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
return (
f"**{doc.title or 'Untitled'}**\n\n"
f"- ID: {doc.id}\n"
f"- URI: {doc.uri or 'N/A'}\n"
f"- URI: {doc.uri}\n"
f"- Created: {doc.created_at.strftime('%Y-%m-%d %H:%M')}\n\n"
f"**Content:**\n{doc.content}"
)
@agent.tool
async def summarize_document(
ctx: RunContext[ChatDeps],
query: str,
) -> str:
"""Generate a summary of a specific document.
Use this when the user wants an overview or summary of a document's content.
Args:
query: The document title or URI to summarize
"""
doc = await _find_document(ctx.deps.client, query)
if doc is None:
return f"Document not found: {query}"
# Use LLM to generate summary
summary_model = get_model(ctx.deps.config.qa.model, ctx.deps.config)
summary_agent: Agent[None, str] = Agent(
summary_model,
output_type=str,
)
result = await summary_agent.run(
DOCUMENT_SUMMARY_PROMPT.format(content=doc.content or "")
)
return f"**Summary of {doc.title or doc.uri}:**\n\n{result.output}"
return agent

17
haiku_rag_slim/haiku/rag/agents/chat/prompts.py
View file

@ -10,7 +10,9 @@ CRITICAL RULES:
5. NEVER make up information - always use tools to get facts from the knowledge base
How to decide which tool to use:
- "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
@ -58,3 +60,16 @@ Rules:
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself."""
DOCUMENT_SUMMARY_PROMPT = """Generate a summary of the document content provided below.
Start with a one-paragraph overview, then list the main topics covered, and highlight any key findings or conclusions.
Guidelines:
- Aim for 1-2 paragraphs for short documents, 3-4 paragraphs for longer ones
- Focus on factual content and key information
- Do not include meta-commentary like "This document discusses..." or "The document covers..."
- Do not speculate beyond what's in the content
Document content:
{content}"""

171
tests/agents/chat/test_chat_agent.py
View file

@ -966,6 +966,177 @@ async def test_summarization_task_cancellation():
_summarization_tasks.clear()
# =============================================================================
# list_documents Tool Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_basic(allow_model_requests, temp_db_path):
"""Test that list_documents tool returns available documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
agent = create_chat_agent(Config)
deps = ChatDeps(
client=client,
config=Config,
)
# Ask to list documents
result = await agent.run(
"What documents are available in the knowledge base?",
deps=deps,
)
assert result.output is not None
# Should mention both documents
assert "DocLayNet" in result.output
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_with_session_filter(allow_model_requests, temp_db_path):
"""Test that list_documents respects session document_filter."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
agent = create_chat_agent(Config)
# Set session filter to only include the labels document
session_state = ChatSessionState(
session_id="test-list-filter",
document_filter=["DocLayNet Class Labels"],
)
deps = ChatDeps(
client=client,
config=Config,
session_state=session_state,
)
# Ask to list documents - should only show filtered documents
result = await agent.run(
"Show me what documents are available",
deps=deps,
)
assert result.output is not None
# Should only mention the Labels document, not Sources
assert "Labels" in result.output or "labels" in result.output
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_list_documents_pagination(allow_model_requests, temp_db_path):
"""Test that list_documents supports pagination via limit/offset."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add multiple test documents
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
await client.create_document(
content=DOCLAYNET_ANNOTATION,
uri="doclaynet-annotation",
title="DocLayNet Annotation",
)
await client.create_document(
content=DOCLAYNET_DATA_SOURCES,
uri="doclaynet-sources",
title="DocLayNet Sources",
)
agent = create_chat_agent(Config)
deps = ChatDeps(
client=client,
config=Config,
)
# Ask to list first 2 documents
result = await agent.run(
"List the first 2 documents available",
deps=deps,
)
assert result.output is not None
# =============================================================================
# summarize_document Tool Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_document_found(allow_model_requests, temp_db_path):
"""Test that summarize_document generates a summary for a found document."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Add a test document
await client.create_document(
content=DOCLAYNET_CLASS_LABELS,
uri="doclaynet-labels",
title="DocLayNet Class Labels",
)
agent = create_chat_agent(Config)
deps = ChatDeps(
client=client,
config=Config,
)
# Ask to summarize a specific document
result = await agent.run(
"Summarize the DocLayNet Class Labels document",
deps=deps,
)
assert result.output is not None
# Should contain summary content about class labels
assert len(result.output) > 50
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_document_not_found(allow_model_requests, temp_db_path):
"""Test that summarize_document handles not found documents gracefully."""
async with HaikuRAG(temp_db_path, create=True) as client:
agent = create_chat_agent(Config)
deps = ChatDeps(
client=client,
config=Config,
)
# Ask to summarize a document that doesn't exist
result = await agent.run(
"Summarize the nonexistent document",
deps=deps,
)
assert result.output is not None
# Should indicate the document wasn't found
def test_citation_index_fallback_without_session_state():
"""Test that citation indices fall back to sequential numbering without session_state.

476
tests/cassettes/test_chat_agent/test_list_documents_basic.yaml Normal file
View file

File diff suppressed because one or more lines are too long

524
tests/cassettes/test_chat_agent/test_list_documents_pagination.yaml Normal file
View file

File diff suppressed because one or more lines are too long

480
tests/cassettes/test_chat_agent/test_list_documents_with_session_filter.yaml Normal file
View file

File diff suppressed because one or more lines are too long

563
tests/cassettes/test_chat_agent/test_summarize_document_found.yaml Normal file
View file

File diff suppressed because one or more lines are too long

196
tests/cassettes/test_chat_agent/test_summarize_document_not_found.yaml Normal file
View file

@ -0,0 +1,196 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '5368'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a helpful research assistant powered by haiku.rag, a knowledge base system.
You have access to a knowledge base of documents. Use your tools to search and answer questions.
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context
3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
4. NEVER call the same tool multiple times for a single user message
5. NEVER make up information - always use tools to get facts from the knowledge base
How to decide which tool to use:
- "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
- "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
- "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
- "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
IMPORTANT - When user mentions a document in search/ask:
- If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
- Extract the TOPIC as `query`/`question`
- Extract the DOCUMENT NAME as `document_name`
- Examples for search:
- "search for embeddings in the ML paper" → query="embeddings", document_name="ML paper"
- "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566"
- Examples for ask:
- "what does the ML paper say about embeddings?" → question="what are the embedding methods?", document_name="ML paper"
- "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566"
Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user.
role: system
- content: Summarize the nonexistent document
role: user
model: gpt-oss
reasoning_effort: low
stream: false
tool_choice: auto
tools:
- function:
description: |-
Search the knowledge base for relevant documents.
Use this when you need to find documents or explore the knowledge base.
Results are displayed to the user - just list the titles found.
name: search
parameters:
additionalProperties: false
properties:
document_name:
anyOf:
- type: string
- type: 'null'
default: null
description: Optional document name/title to search within
limit:
anyOf:
- type: integer
- type: 'null'
default: null
description: 'Number of results to return (default: 5)'
query:
description: The search query (what to search for)
type: string
required:
- query
type: object
type: function
- function:
description: |-
Answer a specific question using the knowledge base.
Use this for direct questions that need a focused answer with citations.
Uses a research graph for planning, searching, and synthesis.
name: ask
parameters:
additionalProperties: false
properties:
document_name:
anyOf:
- type: string
- type: 'null'
default: null
description: Optional document name/title to search within (e.g., "tbmed593", "army manual")
question:
description: The question to answer
type: string
required:
- question
type: object
type: function
- function:
description: |-
List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available.
name: list_documents
parameters:
additionalProperties: false
properties:
limit:
anyOf:
- type: integer
- type: 'null'
default: null
description: Maximum number of documents to return
offset:
anyOf:
- type: integer
- type: 'null'
default: null
description: Number of documents to skip (for pagination)
type: object
type: function
- function:
description: |-
Retrieve a specific document by title or URI.
Use this when the user wants to fetch/get/retrieve a specific document.
name: get_document
parameters:
additionalProperties: false
properties:
query:
description: The document title or URI to look up
type: string
required:
- query
type: object
strict: true
type: function
- function:
description: |-
Generate a summary of a specific document.
Use this when the user wants an overview or summary of a document's content.
name: summarize_document
parameters:
additionalProperties: false
properties:
query:
description: The document title or URI to summarize
type: string
required:
- query
type: object
strict: true
type: function
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '655'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: I’m sorry, but I couldn’t find a document with that name. If you have the exact title or a related keyword,
let me know and I’ll try again.
reasoning: User asks to summarize nonexistent document. According to rule, for summary use summarize_document tool,
but if document doesn't exist? We must search? Likely we respond that document not found. No tool needed.
role: assistant
created: 1769523898
id: chatcmpl-69
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 86
prompt_tokens: 1039
total_tokens: 1125
status:
code: 200
message: OK
version: 1