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

Merge pull request #262 from ggozad/feat/chat-agent-improvements

Browse source 
Conversational agent improvements
This commit is contained in:
Yiorgis Gozadinos 2026-01-28 14:05:43 +02:00 committed by GitHub
commit 27532eb657
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 8295 additions and 4790 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
CHANGELOG.md
View file

@ -1,6 +1,27 @@
# Changelog
## [Unreleased]
### Added
- **Chat Agent Document Awareness Tools**: Two new tools for browsing and understanding the knowledge base
- `list_documents` — Returns `DocumentListResponse` with paginated documents (50 per page), page number, total pages, and total count; respects session document filter
- `summarize_document` — Generate LLM-powered summaries of specific documents
- **Document Count API**: New `count_documents(filter)` method on `HaikuRAG` client for efficient document counting
- **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
- After first message, context becomes read-only (view only)
- Clearing chat resets context to CLI value and unlocks editing
- Web app: Memory panel now serves dual purpose - edit initial context before first message, view session context after
- Agent uses `initial_context` as fallback when `session_context` is empty
### Changed
- **Selective Citation Filtering**: Synthesis steps now select only relevant citations instead of including all
- LLM receives `<available_citations>` with chunk IDs and content previews
- LLM populates `cited_chunks` with only chunks that directly support the answer
- `ResearchReport` now has `cited_chunks` and `citations` fields; removed `sources_summary`
## [0.27.1] - 2026-01-27
### Added

3
app/backend/main.py
View file

@ -83,6 +83,7 @@ async def stream_chat(request: Request) -> Response:
initial_qa_history: list[QAResponse] = []
session_id: str | None = None
document_filter: list[str] = []
initial_context: str | None = None
state = getattr(run_input, "state", None)
if state:
chat_state = state.get(AGUI_STATE_KEY, state)
@ -92,6 +93,7 @@ async def stream_chat(request: Request) -> Response:
]
session_id = chat_state.get("session_id")
document_filter = chat_state.get("document_filter", [])
initial_context = chat_state.get("initial_context")
deps = ChatDeps(
client=get_client(db_path),
@ -99,6 +101,7 @@ async def stream_chat(request: Request) -> Response:
session_state=ChatSessionState(
qa_history=initial_qa_history,
document_filter=document_filter,
initial_context=initial_context,
**({"session_id": session_id} if session_id else {}),
),
state_key=AGUI_STATE_KEY,

41
app/frontend/components/Chat.tsx
View file

@ -42,6 +42,7 @@ interface SessionContext {
interface ChatSessionState {
session_id: string;
initial_context: string | null;
citations: Citation[];
qa_history: QAResponse[];
session_context: SessionContext | null;
@ -398,6 +399,7 @@ function ChatContentInner() {
initialState: {
[AGUI_STATE_KEY]: {
session_id: "",
initial_context: null,
citations: [],
qa_history: [],
session_context: null,
@ -407,9 +409,14 @@ function ChatContentInner() {
},
);
// Extract session context and document filter from agent state
// Extract session context, document filter, and initial context from agent state
const sessionContext = agentState?.[AGUI_STATE_KEY]?.session_context ?? null;
const documentFilter = agentState?.[AGUI_STATE_KEY]?.document_filter ?? [];
const initialContext = agentState?.[AGUI_STATE_KEY]?.initial_context ?? "";
// Context is locked after first message (qa_history has entries)
const isContextLocked =
(agentState?.[AGUI_STATE_KEY]?.qa_history?.length ?? 0) > 0;
const handleFilterApply = (selected: string[]) => {
setAgentState({
@ -417,6 +424,7 @@ function ChatContentInner() {
[AGUI_STATE_KEY]: {
...agentState?.[AGUI_STATE_KEY],
session_id: agentState?.[AGUI_STATE_KEY]?.session_id ?? "",
initial_context: agentState?.[AGUI_STATE_KEY]?.initial_context ?? null,
citations: agentState?.[AGUI_STATE_KEY]?.citations ?? [],
qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [],
session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null,
@ -425,6 +433,22 @@ function ChatContentInner() {
});
};
const handleInitialContextChange = (value: string) => {
if (isContextLocked) return;
setAgentState({
...agentState,
[AGUI_STATE_KEY]: {
...agentState?.[AGUI_STATE_KEY],
session_id: agentState?.[AGUI_STATE_KEY]?.session_id ?? "",
initial_context: value || null,
citations: agentState?.[AGUI_STATE_KEY]?.citations ?? [],
qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [],
session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null,
document_filter: agentState?.[AGUI_STATE_KEY]?.document_filter ?? [],
},
});
};
useCoAgentStateRender<AgentState>({
name: "chat_agent",
render: ({ state }) => {
@ -568,12 +592,16 @@ function ChatContentInner() {
</button>
<button
type="button"
className={`header-btn ${sessionContext?.summary ? "has-content" : ""}`}
className={`header-btn ${initialContext || sessionContext?.summary ? "has-content" : ""}`}
onClick={() => setContextOpen(true)}
title={
sessionContext?.summary
? "View session context"
: "No session context yet"
isContextLocked
? sessionContext?.summary
? "View session context"
: "No session context yet"
: initialContext
? "Edit initial context"
: "Set initial context"
}
>
<BrainIcon />
@ -596,6 +624,9 @@ function ChatContentInner() {
isOpen={contextOpen}
onClose={() => setContextOpen(false)}
sessionContext={sessionContext}
initialContext={initialContext}
onInitialContextChange={handleInitialContextChange}
isLocked={isContextLocked}
/>
<DocumentFilter
isOpen={filterOpen}

101
app/frontend/components/ContextPanel.tsx
View file

@ -1,6 +1,6 @@
"use client";
import { useCallback, useId } from "react";
import { useCallback, useEffect, useId, useState } from "react";
interface SessionContext {
summary: string;
@ -11,6 +11,9 @@ interface ContextPanelProps {
isOpen: boolean;
onClose: () => void;
sessionContext: SessionContext | null;
initialContext?: string;
onInitialContextChange?: (value: string) => void;
isLocked?: boolean;
}
function formatRelativeTime(isoString: string): string {
@ -62,8 +65,18 @@ export default function ContextPanel({
isOpen,
onClose,
sessionContext,
initialContext = "",
onInitialContextChange,
isLocked = false,
}: ContextPanelProps) {
const titleId = useId();
const [localValue, setLocalValue] = useState(initialContext);
useEffect(() => {
if (isOpen) {
setLocalValue(initialContext);
}
}, [isOpen, initialContext]);
const handleKeyDown = useCallback(
(e: React.KeyboardEvent) => {
@ -74,11 +87,18 @@ export default function ContextPanel({
[onClose],
);
const handleSave = useCallback(() => {
onInitialContextChange?.(localValue);
onClose();
}, [localValue, onInitialContextChange, onClose]);
if (!isOpen) {
return null;
}
const hasContext = sessionContext?.summary && sessionContext.summary.trim();
const hasSessionContext = sessionContext?.summary?.trim();
// Show edit mode when: not locked AND no session context yet
const isEditMode = !isLocked && !hasSessionContext;
return (
<>
@ -146,6 +166,24 @@ export default function ContextPanel({
max-height: 400px;
overflow-y: auto;
}
.context-textarea {
width: 100%;
min-height: 200px;
padding: 1rem;
font-size: 0.875rem;
line-height: 1.6;
color: #334155;
background: white;
border: 1px solid #e2e8f0;
border-radius: 8px;
resize: vertical;
font-family: inherit;
}
.context-textarea:focus {
outline: none;
border-color: #3b82f6;
box-shadow: 0 0 0 3px rgba(59, 130, 246, 0.1);
}
.context-empty {
display: flex;
flex-direction: column;
@ -175,13 +213,15 @@ export default function ContextPanel({
font-size: 0.75rem;
color: #94a3b8;
}
.context-btn-close {
.context-btn {
padding: 0.5rem 1rem;
font-size: 0.875rem;
font-weight: 500;
border-radius: 6px;
cursor: pointer;
transition: all 0.15s;
}
.context-btn-close {
background: white;
color: #475569;
border: 1px solid #e2e8f0;
@ -190,6 +230,20 @@ export default function ContextPanel({
background: #f8fafc;
border-color: #cbd5e1;
}
.context-btn-save {
background: #3b82f6;
color: white;
border: 1px solid #3b82f6;
margin-left: 0.5rem;
}
.context-btn-save:hover {
background: #2563eb;
border-color: #2563eb;
}
.context-footer-buttons {
display: flex;
gap: 0.5rem;
}
`}</style>
<div
className="context-modal-overlay"
@ -210,14 +264,22 @@ export default function ContextPanel({
<BrainIcon />
</div>
<h2 id={titleId} className="context-modal-title">
Session Context
{isEditMode ? "Initial Context" : "Session Context"}
</h2>
</div>
<p className="context-modal-description">
This is what the assistant has learned from your conversation so
far. It uses this context to provide more relevant answers.
{isEditMode
? "Set background context to guide the conversation. This will be locked after you send your first message."
: "This is what the assistant has learned from your conversation so far. It uses this context to provide more relevant answers."}
</p>
{hasContext ? (
{isEditMode ? (
<textarea
className="context-textarea"
placeholder="Enter any background context or instructions for the assistant..."
value={localValue}
onChange={(e) => setLocalValue(e.target.value)}
/>
) : hasSessionContext ? (
<div className="context-content">{sessionContext.summary}</div>
) : (
<div className="context-empty">
@ -235,13 +297,24 @@ export default function ContextPanel({
? `Last updated: ${formatRelativeTime(sessionContext.last_updated)}`
: ""}
</span>
<button
type="button"
className="context-btn-close"
onClick={onClose}
>
Close
</button>
<div className="context-footer-buttons">
<button
type="button"
className="context-btn context-btn-close"
onClick={onClose}
>
{isEditMode ? "Cancel" : "Close"}
</button>
{isEditMode && (
<button
type="button"
className="context-btn context-btn-save"
onClick={handleSave}
>
Save
</button>
)}
</div>
</div>
</div>
</div>

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.

9
docs/apps.md
View file

@ -31,8 +31,8 @@ Press `Ctrl+P` to open the command palette:
| Command | Description |
|---------|-------------|
| Memory | Edit initial context (before first message) or view session context (after) |
| Filter documents | Select documents to restrict searches |
| Show context | View current session context |
| Show database info | View document/chunk counts and storage info |
| Visual grounding | View chunk source location in document |
| Clear chat | Clear chat history and reset session |
@ -43,7 +43,9 @@ Press `Ctrl+P` to open the command palette:
- Previous Q/A pairs are used as context for follow-up questions
- Citations are tracked per response and can be inspected
- Document filter restricts all searches to selected documents
- Clearing chat resets session state but preserves document filter
- Initial context can be set via CLI (`--initial-context`) or command palette
- Initial context is editable until the first message is sent, then becomes read-only
- Clearing chat resets session state, restores CLI-provided context, and unlocks editing
## Web Application
@ -55,8 +57,7 @@ Browser-based conversational RAG with a CopilotKit frontend.
- Expandable citations with source documents, pages, and headings
- Visual grounding to view chunk source locations in documents
- Document filter to restrict searches to selected documents
- Session context that summarizes conversation history
- Settings panel for background context configuration
- Memory panel: set initial context before first message, view session context after
### Quick Start

17
docs/cli.md
View file

@ -178,10 +178,9 @@ haiku-rag chat
haiku-rag chat --db /path/to/database.lancedb
```
Provide background context for the conversation:
Provide initial background context for the conversation:
```bash
haiku-rag chat --context "Focus on Python programming concepts"
haiku-rag chat --context-file domain-context.txt
haiku-rag chat --initial-context "Focus on Python programming concepts"
```
!!! note
@ -193,12 +192,18 @@ The chat interface provides:
- Expandable citations with source metadata
- Session memory for context-aware follow-up questions
- Visual grounding to inspect chunk source locations
- Background context that persists across the entire conversation
- Initial context that can be edited before the first message
**Initial Context Behavior:**
- Edit initial context via command palette before sending your first message
- Once you send a message, initial context becomes read-only
- The agent uses initial context as a starting point for session summarization
- Clearing chat resets to the CLI-provided context and unlocks editing
Flags:
- `--context`: Background context for the conversation
- `--context-file`: Path to a file containing background context
- `--initial-context`: Initial background context for the conversation (editable until first message)
See [Applications](apps.md#chat-tui) for keyboard shortcuts and features.

5
docs/configuration/prompts.md
View file

@ -71,7 +71,7 @@ prompts:
Replace the research report synthesis prompt by setting `prompts.synthesis`. This controls how the multi-agent research workflow generates its final report.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `sources_summary`.
The prompt should produce a `ResearchReport` with: `title`, `executive_summary`, `main_findings`, `conclusions`, `recommendations`, `limitations`, and `cited_chunks`.
**Example:**
@ -87,12 +87,13 @@ prompts:
- conclusions: 2-4 bullet points
- recommendations: 2-5 actionable recommendations
- limitations: 1-3 limitations or gaps
- sources_summary: Brief description of sources used
- cited_chunks: List of chunk IDs that directly support the report
Guidelines:
- Base all content strictly on collected evidence
- Be specific and objective
- Avoid meta-commentary like "This report covers..."
- Only include chunks in cited_chunks that directly support claims in the report
```
## Picture Description Prompt

9
docs/python.md
View file

@ -151,6 +151,15 @@ docs = await client.list_documents(
)
```
Count documents:
```python
# Count all documents
total = await client.count_documents()
# Count with filter
pdf_count = await client.count_documents(filter="uri LIKE '%.pdf'")
```
### Updating Documents
```python

4
haiku_rag_slim/haiku/rag/agents/chat/__init__.py
View file

@ -8,6 +8,8 @@ from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatDeps,
ChatSessionState,
DocumentInfo,
DocumentListResponse,
QAResponse,
SearchDeps,
SessionContext,
@ -20,6 +22,8 @@ __all__ = [
"SearchAgent",
"ChatDeps",
"ChatSessionState",
"DocumentInfo",
"DocumentListResponse",
"QAResponse",
"SearchDeps",
"SessionContext",

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

@ -8,12 +8,14 @@ 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,
ChatDeps,
ChatSessionState,
DocumentInfo,
DocumentListResponse,
QAResponse,
build_document_filter,
build_multi_document_filter,
@ -23,6 +25,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 +372,77 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
],
)
@agent.tool
async def list_documents(
ctx: RunContext[ChatDeps],
page: int = 1,
) -> DocumentListResponse:
"""List available documents in the knowledge base.
Use this when the user wants to browse or see what documents are available.
Args:
page: Page number (default: 1, 50 documents per page)
"""
page_size = 50
offset = (page - 1) * page_size
# 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=page_size, offset=offset, filter=doc_filter
)
total = await ctx.deps.client.count_documents(filter=doc_filter)
total_pages = (total + page_size - 1) // page_size if total > 0 else 1
return DocumentListResponse(
documents=[
DocumentInfo(
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
)
for doc in docs
],
page=page,
total_pages=total_pages,
total_documents=total,
)
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 +455,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 +463,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}"""

23
haiku_rag_slim/haiku/rag/agents/chat/state.py
View file

@ -42,6 +42,23 @@ class QAResponse(BaseModel):
)
class DocumentInfo(BaseModel):
"""Document info for list_documents response."""
title: str
uri: str
created: str
class DocumentListResponse(BaseModel):
"""Response from list_documents tool."""
documents: list[DocumentInfo]
page: int
total_pages: int
total_documents: int
class SessionContext(BaseModel):
"""Compressed summary of conversation history for research graph."""
@ -133,8 +150,10 @@ class ChatDeps:
)
if "citation_registry" in state_data:
self.session_state.citation_registry = state_data["citation_registry"]
# NOTE: session_context intentionally NOT updated from client
# The agent owns this via server-side cache
if "initial_context" in state_data:
self.session_state.initial_context = state_data.get("initial_context")
# NOTE: session_context is server-managed; we don't accept it from the client
# to maintain server-side ownership of conversation summarization
@dataclass

6
haiku_rag_slim/haiku/rag/agents/qa/prompts.py
View file

@ -26,7 +26,11 @@ Each result includes:
- Type: content type like paragraph, table, code, list_item (when available)
- Content: the actual text
In your response, include the chunk IDs you used in cited_chunks.
Citation guidelines:
- In cited_chunks, include ONLY chunk IDs that directly support your answer.
- Do NOT cite chunks that are merely related or that you reviewed but did not use.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT, COMPLETE chunk IDs (full UUIDs).
Guidelines:
- Base answers strictly on retrieved content - do not use external knowledge

73
haiku_rag_slim/haiku/rag/agents/research/graph.py
View file

@ -32,6 +32,7 @@ from haiku.rag.utils import build_prompt, get_model
def format_context_for_prompt(
context: ResearchContext,
include_pending_questions: bool = True,
include_citations: bool = False,
) -> str:
"""Format the research context as XML for prompts.
@ -39,6 +40,8 @@ def format_context_for_prompt(
context: The research context to format.
include_pending_questions: Whether to include pending sub-questions.
Set to False for synthesis prompts where pending questions aren't relevant.
include_citations: Whether to include available citations for selection.
Set to True for synthesis prompts where the LLM should select relevant citations.
"""
context_data: dict[str, object] = {}
@ -61,6 +64,26 @@ def format_context_for_prompt(
for qa in context.qa_responses
]
if include_citations and context.qa_responses:
seen_chunks: set[str] = set()
available_citations: list[dict[str, str]] = []
for qa in context.qa_responses:
for c in qa.citations:
if c.chunk_id not in seen_chunks:
seen_chunks.add(c.chunk_id)
content_preview = (
c.content[:500] + "..." if len(c.content) > 500 else c.content
)
available_citations.append(
{
"chunk_id": c.chunk_id,
"document": c.document_title or c.document_uri,
"content": content_preview,
}
)
if available_citations:
context_data["available_citations"] = available_citations
return format_as_xml(context_data, root_tag="context")
@ -350,7 +373,10 @@ def build_research_graph(
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
# Include available citations for the LLM to select from
context_xml = format_context_for_prompt(
state.context, include_pending_questions=False, include_citations=True
)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
@ -361,7 +387,21 @@ def build_research_graph(
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
return result.output
report = result.output
citation_lookup: dict[str, Citation] = {}
for qa in state.context.qa_responses:
for c in qa.citations:
if c.chunk_id not in citation_lookup:
citation_lookup[c.chunk_id] = c
resolved_citations: list[Citation] = []
for chunk_id in report.cited_chunks:
if chunk_id in citation_lookup:
resolved_citations.append(citation_lookup[chunk_id])
report.citations = resolved_citations
return report
# Build the graph structure
collect_answers = g.join(
@ -479,17 +519,19 @@ def build_conversational_graph(
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[assignment]
# Use RawSearchAnswer so LLM can select which chunks to cite
agent: Agent[ResearchDependencies, RawSearchAnswer] = Agent( # type: ignore[assignment]
model=get_model(config.research.model, config),
output_type=ConversationalAnswer,
output_type=RawSearchAnswer,
instructions=conversational_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
# Include available citations for the LLM to select from
context_xml = format_context_for_prompt(
state.context, include_pending_questions=False
state.context, include_pending_questions=False, include_citations=True
)
prompt = f"Answer the question based on the gathered evidence.\n\n{context_xml}"
agent_deps = ResearchDependencies(
@ -497,20 +539,23 @@ def build_conversational_graph(
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
raw_answer = result.output
# Collect unique citations from qa_responses (dedupe by chunk_id)
seen_chunks: set[str] = set()
unique_citations: list[Citation] = []
citation_lookup: dict[str, Citation] = {}
for qa in state.context.qa_responses:
for c in qa.citations:
if c.chunk_id not in seen_chunks:
seen_chunks.add(c.chunk_id)
unique_citations.append(c)
if c.chunk_id not in citation_lookup:
citation_lookup[c.chunk_id] = c
filtered_citations: list[Citation] = []
for chunk_id in raw_answer.cited_chunks:
if chunk_id in citation_lookup:
filtered_citations.append(citation_lookup[chunk_id])
return ConversationalAnswer(
answer=result.output.answer,
citations=unique_citations,
confidence=result.output.confidence,
answer=raw_answer.answer,
citations=filtered_citations,
confidence=raw_answer.confidence,
)
# Build the graph structure (simplified: plan → search → synthesize)

9
haiku_rag_slim/haiku/rag/agents/research/models.py
View file

@ -163,6 +163,11 @@ class ResearchReport(BaseModel):
recommendations: list[str] = Field(
description="Actionable recommendations based on findings", default=[]
)
sources_summary: str = Field(
description="Summary of sources used and their reliability"
cited_chunks: list[str] = Field(
default_factory=list,
description="Chunk IDs selected by synthesis as directly supporting the report",
)
citations: list[Citation] = Field(
default_factory=list,
description="Resolved citations with full metadata",
)

21
haiku_rag_slim/haiku/rag/agents/research/prompts.py
View file

@ -116,6 +116,7 @@ Goals:
2. Present findings clearly and concisely.
3. Draw evidence-based conclusions and recommendations.
4. State limitations and uncertainties transparently.
5. Select only the citations that directly support your final answer.
Report guidelines (map to output fields):
- title: concise (5-12 words), informative.
@ -127,10 +128,17 @@ Report guidelines (map to output fields):
- conclusions: list of plain strings, 2-4 bullets following logically from findings.
- recommendations: list of plain strings, 2-5 actionable bullets tied to findings.
- limitations: list of plain strings, 1-3 bullets describing constraints or uncertainties.
- sources_summary: single string listing sources with document paths and page numbers.
- cited_chunks: list of chunk IDs that DIRECTLY support your report.
All list fields must contain plain strings only, not objects.
Citation selection:
- Review the <available_citations> section in the context.
- Include ONLY chunk IDs whose content directly supports specific claims in your report.
- Do NOT include chunks that are merely related, tangential, or were reviewed but unused.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT chunk IDs from the available_citations (full UUIDs).
Style:
- Base all content solely on the collected evidence.
- Be professional, objective, and specific.
@ -141,9 +149,11 @@ CONVERSATIONAL_SYNTHESIS_PROMPT = """Generate a direct, conversational answer
to the question based on the gathered evidence.
Output:
- query: Echo the original question being answered.
- answer: Direct, comprehensive answer with a natural, helpful tone.
Write the actual answer, not a description of what you found.
Use as many sentences as needed to fully address the question.
- cited_chunks: List of chunk IDs that DIRECTLY support your answer.
- confidence: Score from 0.0 to 1.0 indicating answer quality.
Guidelines:
@ -153,4 +163,11 @@ Guidelines:
- Use formatting (bullet points, numbered lists) when it improves clarity.
- Do NOT use meta-commentary like "Based on the research..." or "The evidence shows..."
Instead, directly state the information.
- If the evidence is incomplete, acknowledge limitations briefly."""
- If the evidence is incomplete, acknowledge limitations briefly.
Citation selection:
- Review the <available_citations> section in the context.
- Include ONLY chunk IDs whose content directly supports specific statements in your answer.
- Do NOT include chunks that are merely related, tangential, or were reviewed but unused.
- Quality over quantity: fewer relevant citations are better than many marginal ones.
- Use the EXACT chunk IDs from available_citations (full UUIDs)."""

15
haiku_rag_slim/haiku/rag/app.py
View file

@ -416,10 +416,9 @@ class HaikuRAGApp:
self.console.print("[bold cyan]Key Findings:[/bold cyan]")
for finding in report.main_findings:
self.console.print(f"• {finding}")
if report.sources_summary:
self.console.print()
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
if report.citations:
for renderable in format_citations_rich(report.citations):
self.console.print(renderable)
else:
self.console.print("[yellow]No answer generated.[/yellow]")
else:
@ -513,10 +512,10 @@ class HaikuRAGApp:
self.console.print(f"• {limitation}")
self.console.print()
# Sources Summary
if report.sources_summary:
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
# Sources
if report.citations:
for renderable in format_citations_rich(report.citations):
self.console.print(renderable)
async def rebuild(self, mode: RebuildMode = RebuildMode.FULL):
async with HaikuRAG(

3
haiku_rag_slim/haiku/rag/chat/__init__.py
View file

@ -6,6 +6,7 @@ def run_chat(
db_path: Path | None = None,
read_only: bool = False,
before: datetime | None = None,
initial_context: str | None = None,
) -> None:
"""Run the chat TUI.
@ -13,6 +14,7 @@ def run_chat(
db_path: Path to the LanceDB database. If None, uses default from config.
read_only: Whether to open the database in read-only mode.
before: Query database as it existed before this datetime.
initial_context: Initial background context to provide to the conversation.
"""
try:
from haiku.rag.chat.app import ChatApp
@ -31,5 +33,6 @@ def run_chat(
db_path,
read_only=read_only,
before=before,
initial_context=initial_context,
)
app.run()

26
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -88,11 +88,14 @@ class ChatApp(App):
db_path: Path,
read_only: bool = False,
before: datetime | None = None,
initial_context: str | None = None,
) -> None:
super().__init__()
self.db_path = db_path
self.read_only = read_only
self.before = before
self._initial_context = initial_context
self._context_locked = False
self.client: HaikuRAG | None = None
self.config = get_config()
self.agent: Agent[ChatDeps, str] | None = None
@ -135,8 +138,8 @@ class ChatApp(App):
self.action_show_info,
)
yield SystemCommand(
"Session context",
"Show current session context",
"Memory",
"View/edit context (editable before first message)",
self.action_show_context,
)
@ -153,6 +156,7 @@ class ChatApp(App):
# Create agent and session state
self.agent = create_chat_agent(self.config)
self.session_state = ChatSessionState(
initial_context=self._initial_context,
document_filter=self._document_filter,
)
@ -216,6 +220,9 @@ class ChatApp(App):
if not self.client or not self.agent:
return
# Lock context after first message
self._context_locked = True
# Clear the input
event.input.clear()
@ -297,8 +304,10 @@ class ChatApp(App):
await chat_history.clear_messages()
self._last_citations.clear()
self._message_history.clear()
# Reset session state for fresh conversation (preserve document filter)
# Reset context lock and session state (reset to CLI value)
self._context_locked = False
self.session_state = ChatSessionState(
initial_context=self._initial_context,
document_filter=self._document_filter,
)
@ -349,10 +358,17 @@ class ChatApp(App):
await self.push_screen(InfoModal(self.client, self.db_path))
async def action_show_context(self) -> None:
"""Show current session context in a modal."""
"""Show context modal (edit initial context or view session context)."""
from haiku.rag.chat.widgets.context_modal import ContextModal
await self.push_screen(ContextModal(self.session_state))
await self.push_screen(
ContextModal(self.session_state, is_locked=self._context_locked)
)
def on_context_modal_context_updated(self, event: Any) -> None:
"""Handle context updates from modal."""
if self.session_state and not self._context_locked:
self.session_state.initial_context = event.context or None
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
"""Handle citation selection."""

111
haiku_rag_slim/haiku/rag/chat/widgets/context_modal.py
View file

@ -2,20 +2,25 @@ from typing import TYPE_CHECKING
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Vertical, VerticalScroll
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.message import Message
from textual.screen import ModalScreen
from textual.widgets import Markdown, Static
from textual.widgets import Button, Markdown, Static, TextArea
if TYPE_CHECKING:
from haiku.rag.agents.chat.state import ChatSessionState
class ContextModal(ModalScreen): # pragma: no cover
"""Modal screen for displaying session context."""
"""Modal screen for viewing/editing context.
Before first message (not locked, no session context): Edit initial context
After first message (locked or has session context): View session context
"""
BINDINGS = [
Binding("escape", "dismiss", "Close", show=True),
Binding("ctrl+o", "dismiss", "Close", show=True),
Binding("escape", "cancel", "Close", show=False),
Binding("ctrl+o", "cancel", "Close", show=False),
]
CSS = """
@ -25,11 +30,9 @@ class ContextModal(ModalScreen): # pragma: no cover
}
#context-container {
width: auto;
min-width: 40;
max-width: 80;
width: 70;
height: auto;
max-height: 20;
max-height: 32;
background: $surface;
border: tall $primary;
padding: 1 2;
@ -40,23 +43,87 @@ class ContextModal(ModalScreen): # pragma: no cover
margin-bottom: 1;
}
#context-description {
height: auto;
margin-bottom: 1;
color: $text-muted;
}
#context-editor {
height: 12;
min-height: 8;
max-height: 16;
}
#context-content {
height: 1fr;
max-height: 16;
scrollbar-gutter: stable;
}
#button-row {
height: auto;
margin-top: 1;
align: right middle;
}
#button-row Button {
margin-left: 1;
}
"""
def __init__(self, session_state: "ChatSessionState | None"):
class ContextUpdated(Message):
"""Emitted when the context is saved."""
def __init__(self, context: str) -> None:
super().__init__()
self.context = context
def __init__(
self, session_state: "ChatSessionState | None", is_locked: bool = False
) -> None:
super().__init__()
self.session_state = session_state
self._is_locked = is_locked
@property
def _is_edit_mode(self) -> bool:
"""Edit mode when not locked and no session context yet."""
has_session_context = (
self.session_state
and self.session_state.session_context
and self.session_state.session_context.summary
)
return not self._is_locked and not has_session_context
def compose(self) -> ComposeResult:
with Vertical(id="context-container"):
yield Static("[bold]Session Context[/bold]", id="context-header")
with VerticalScroll(id="context-content"):
yield Markdown(self._get_content())
if self._is_edit_mode:
yield Static("[bold]Initial Context[/bold]", id="context-header")
yield Static(
"Set background context to guide the conversation. "
"This will be locked after you send your first message.",
id="context-description",
)
initial_value = ""
if self.session_state and self.session_state.initial_context:
initial_value = self.session_state.initial_context
yield TextArea(initial_value, id="context-editor")
with Horizontal(id="button-row"):
yield Button("Cancel", id="cancel-btn", variant="default")
yield Button("Save", id="save-btn", variant="primary")
else:
yield Static("[bold]Session Context[/bold]", id="context-header")
yield Static(
"What the assistant has learned from your conversation.",
id="context-description",
)
with VerticalScroll(id="context-content"):
yield Markdown(self._get_session_content())
with Horizontal(id="button-row"):
yield Button("Close", id="cancel-btn", variant="primary")
def _get_content(self) -> str:
def _get_session_content(self) -> str:
if not self.session_state:
return "*No session state.*"
@ -72,5 +139,19 @@ class ContextModal(ModalScreen): # pragma: no cover
return f"**Last updated:** {updated}\n\n---\n\n{ctx.summary}"
async def action_dismiss(self, result=None) -> None:
def on_button_pressed(self, event: Button.Pressed) -> None:
"""Handle button presses."""
if event.button.id == "cancel-btn":
self.action_cancel()
elif event.button.id == "save-btn":
self.action_save()
def action_cancel(self) -> None:
"""Cancel and close without saving."""
self.app.pop_screen()
def action_save(self) -> None:
"""Save context and close."""
editor = self.query_one("#context-editor", TextArea)
self.post_message(self.ContextUpdated(editor.text))
self.app.pop_screen()

6
haiku_rag_slim/haiku/rag/cli.py
View file

@ -593,6 +593,11 @@ def chat(
"--db",
help="Path to the LanceDB database file",
),
initial_context: str | None = typer.Option(
None,
"--initial-context",
help="Initial background context to provide to the conversation",
),
):
"""Launch the chat TUI for conversational RAG."""
from haiku.rag.chat import run_chat
@ -603,6 +608,7 @@ def chat(
db_path,
read_only=_read_only,
before=_before,
initial_context=initial_context,
)

11
haiku_rag_slim/haiku/rag/client.py
View file

@ -854,6 +854,17 @@ class HaikuRAG:
limit=limit, offset=offset, filter=filter
)
async def count_documents(self, filter: str | None = None) -> int:
"""Count documents with optional filtering.
Args:
filter: Optional SQL WHERE clause to filter documents.
Returns:
Number of documents matching the criteria.
"""
return await self.document_repository.count(filter=filter)
async def search(
self,
query: str,

11
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -167,6 +167,17 @@ class DocumentRepository:
results = list(query.to_pydantic(DocumentRecord))
return [self._record_to_document(doc) for doc in results]
async def count(self, filter: str | None = None) -> int:
"""Count documents with optional filtering.
Args:
filter: Optional SQL WHERE clause to filter documents.
Returns:
Number of documents matching the criteria.
"""
return self.store.documents_table.count_rows(filter=filter)
async def get_by_uri(self, uri: str) -> Document | None:
"""Get a document by its URI."""
escaped_uri = _escape_sql_string(uri)

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

@ -966,6 +966,203 @@ 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
# =============================================================================
# count_documents Tests
# =============================================================================
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_count_documents(temp_db_path):
"""Test count_documents method."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Empty database
assert await client.count_documents() == 0
# Add documents
await client.create_document(content="Doc 1", uri="test/doc1.pdf")
await client.create_document(content="Doc 2", uri="test/doc2.pdf")
await client.create_document(content="Doc 3", uri="other/doc3.txt")
# Count all
assert await client.count_documents() == 3
# Count with filter
assert await client.count_documents(filter="uri LIKE '%.pdf'") == 2
assert await client.count_documents(filter="uri LIKE '%.txt'") == 1
def test_citation_index_fallback_without_session_state():
"""Test that citation indices fall back to sequential numbering without session_state.

43
tests/agents/research/test_research_graph.py
View file

@ -184,3 +184,46 @@ def test_format_context_for_prompt_with_prior_answers():
assert "<prior_answers>" in result
assert "Sub question?" in result
assert "The answer is here." in result
def test_format_context_for_prompt_with_citations():
"""Test format_context_for_prompt includes available_citations when requested."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
from haiku.rag.agents.research.models import Citation, SearchAnswer
context = ResearchContext(original_question="Main question?")
context.add_qa_response(
SearchAnswer(
query="Sub question?",
answer="The answer is here.",
confidence=0.9,
cited_chunks=["chunk-123", "chunk-456"],
citations=[
Citation(
document_id="doc-1",
chunk_id="chunk-123",
document_uri="test://doc1",
document_title="Test Document",
content="This is the chunk content.",
),
Citation(
document_id="doc-1",
chunk_id="chunk-456",
document_uri="test://doc1",
document_title="Test Document",
content="More chunk content here.",
),
],
)
)
result_without = format_context_for_prompt(context, include_citations=False)
assert "<available_citations>" not in result_without
result_with = format_context_for_prompt(context, include_citations=True)
assert "<available_citations>" in result_with
assert "chunk-123" in result_with
assert "chunk-456" in result_with
assert "Test Document" in result_with
assert "This is the chunk content." in result_with

1914
tests/cassettes/test_chat_agent/test_chat_agent_ask_adds_citations.yaml
View file

File diff suppressed because one or more lines are too long

2994
tests/cassettes/test_chat_agent/test_chat_agent_ask_triggers_background_summarization.yaml
View file

File diff suppressed because one or more lines are too long

4893
tests/cassettes/test_chat_agent/test_chat_agent_ask_with_prior_answer_retrieval.yaml
View file

File diff suppressed because one or more lines are too long

122
tests/cassettes/test_chat_agent/test_count_documents.yaml Normal file
View file

File diff suppressed because one or more lines are too long

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

2
tests/test_app.py
View file

@ -349,7 +349,6 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()
@ -387,7 +386,6 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
executive_summary="Deep research answer",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Sources",
)
mock_graph = AsyncMock()

1
tests/test_mcp.py
View file

@ -277,7 +277,6 @@ async def test_mcp_research_question():
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
recommendations=["Recommendation 1"],
sources_summary="Sources used",
)
with (