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

Merge pull request #282 from ggozad/feat/haiku-skills

Browse source 
Integrate haiku.skills: RAG & RLM skills, simplified toolsets, rebuilt TUI and web app
This commit is contained in:
Yiorgis Gozadinos 2026-02-20 19:00:04 +02:00 committed by GitHub
commit 82a5ce8519
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
135 changed files with 2671 additions and 31455 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
.github/workflows/test.yml vendored
View file

@ -68,6 +68,8 @@ jobs:
run: uv run python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B')"
- name: Run tests with coverage
run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
env:
HF_HUB_OFFLINE: "1"
- name: Upload coverage to Codecov
uses: codecov/codecov-action@v5
with:

21
CHANGELOG.md
View file

@ -1,6 +1,27 @@
# Changelog
## [Unreleased]
### Added
- **RAG skill** (`haiku.rag.skills.rag`): haiku.skills integration with search, list_documents, get_document, ask, and research tools plus managed `RAGState`
- **RLM skill** (`haiku.rag.skills.rlm`): haiku.skills integration with analyze tool for computational analysis via code execution
- **`HaikuRAG.research()`**: Client method for multi-agent research
- **haiku.skills entry points**: `rag = "haiku.rag.skills.rag:create_skill"`, `rag-rlm = "haiku.rag.skills.rlm:create_skill"`
### Changed
- **Chat TUI**: Rebuilt on RAG skill + haiku.skills `SkillToolset`
- **Web app backend**: Rebuilt on RAG skill + `AGUIAdapter`
- **Toolsets simplified**: Removed `ToolContext`, `SessionState`, `AgentDeps`, `Toolkit`; kept core `FunctionToolset` factories
- **Research graph**: Removed `session_context` and conversational output mode
### Removed
- **`agents/chat/`**: Entire chat agent module (replaced by RAG skill)
- **`--deep` flag**: Removed from `ask` CLI (use `research` command instead)
- **`--context`/`--context-file`**: Removed from `ask` CLI
- **`tools/` state machinery**: `ToolContext`, `ToolContextCache`, `SessionState`, `AgentDeps`, `Toolkit`, etc.
## [0.30.2] - 2026-02-19
### Fixed

5
README.md
View file

@ -59,9 +59,6 @@ haiku-rag search "attention mechanism"
# Ask questions with citations
haiku-rag ask "What datasets were used for evaluation?" --cite
# Deep QA — decomposes complex questions into sub-queries
haiku-rag ask "How does the proposed method compare to the baseline on MMLU?" --deep
# Research mode — iterative planning and search
haiku-rag research "What are the limitations of the approach?"
@ -140,7 +137,7 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
- [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - YAML configuration
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA, chat, and research agents
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA and research agents
- [RLM Agent](https://ggozad.github.io/haiku.rag/rlm/) - Complex analytical tasks via code execution
- [Applications](https://ggozad.github.io/haiku.rag/apps/) - Chat TUI, web app, and inspector
- [Server](https://ggozad.github.io/haiku.rag/server/) - File monitoring and MCP

4
app/README.md
View file

@ -46,7 +46,7 @@ A conversational RAG interface built with [CopilotKit](https://copilotkit.ai/) a
### haiku.rag.yaml
Configure the chat agent's LLM, embeddings, and search settings:
Configure the LLM, embeddings, and search settings:
```yaml
qa:
@ -100,7 +100,7 @@ docker compose -f docker-compose.dev.yml up -d --build
## Chat Capabilities
The chat agent can:
The chat can:
- **Search** your documents with hybrid vector + full-text search
- **Answer questions** with citations from your knowledge base

47
app/backend/main.py
View file

@ -3,8 +3,9 @@ import os
from pathlib import Path
from dotenv import find_dotenv, load_dotenv
from pydantic_ai import Agent
from pydantic_ai.ag_ui import AGUIAdapter
from pydantic_ai.ui import SSE_CONTENT_TYPE
from pydantic_ai.ui.ag_ui import AGUIAdapter
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
@ -12,22 +13,14 @@ from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.agents.chat import (
AGUI_STATE_KEY,
ChatDeps,
build_chat_toolkit,
create_chat_agent,
)
from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContextCache
from haiku.rag.skills.rag import AGENT_PREAMBLE, create_skill
from haiku.skills import SkillDeps, SkillToolset
load_dotenv(find_dotenv(usecwd=True))
# Cache ToolContext instances by thread_id across requests
context_cache = ToolContextCache()
# Configure logfire (only sends data if LOGFIRE_TOKEN is present)
try:
import logfire
@ -69,34 +62,26 @@ def get_client() -> HaikuRAG:
return _client
# Toolkit and agent are created once at module level
chat_toolkit = build_chat_toolkit(Config)
agent = create_chat_agent(Config, toolkit=chat_toolkit)
# Create skill, toolset, and agent
skill = create_skill(db_path=db_path, config=Config)
toolset = SkillToolset(skills=[skill])
agent = Agent(
os.getenv("HAIKU_CHAT_MODEL", "openai:gpt-4o"),
instructions=AGENT_PREAMBLE + toolset.system_prompt,
toolsets=[toolset],
deps_type=SkillDeps,
)
async def stream_chat(request: Request) -> Response:
"""Chat streaming endpoint with AG-UI protocol.
Uses ToolContextCache to maintain state across requests for the same thread.
AGUIAdapter restores client-sent state via ChatDeps.state setter.
"""
"""Chat streaming endpoint with AG-UI protocol."""
body = await request.body()
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body)
thread_id = getattr(run_input, "thread_id", None) or "default"
context, is_new = context_cache.get_or_create(thread_id)
if is_new:
chat_toolkit.prepare(context, state_key=AGUI_STATE_KEY)
deps = ChatDeps(
config=Config,
client=get_client(),
tool_context=context,
)
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
event_stream = adapter.run_stream(deps=deps)
event_stream = adapter.run_stream(deps=SkillDeps())
sse_event_stream = adapter.encode_stream(event_stream)
return StreamingResponse(

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

@ -17,27 +17,27 @@ import {
useMemo,
useState,
} from "react";
import { BrainIcon, FilterIcon } from "../lib/icons";
import type { ChatSessionState } from "../lib/sessionStorage";
import { FilterIcon } from "../lib/icons";
import type { RAGState } from "../lib/sessionStorage";
import {
createSession,
deriveCitationsHistory,
getActiveSessionId,
getSession,
normalizeChatState,
normalizeRAGState,
updateSessionMessages,
} from "../lib/sessionStorage";
import CitationBlock from "./CitationBlock";
import ContextPanel from "./ContextPanel";
import DbInfo from "./DbInfo";
import DocumentFilter from "./DocumentFilter";
import SessionManager from "./SessionManager";
// Must match AGUI_STATE_KEY from haiku.rag.agents.chat
const AGUI_STATE_KEY = "haiku.rag.chat";
// Must match state_namespace from haiku.rag.skills.rag
const AGUI_STATE_KEY = "rag";
// AG-UI state is namespaced under AGUI_STATE_KEY
interface AgentState {
[AGUI_STATE_KEY]?: ChatSessionState;
[AGUI_STATE_KEY]?: RAGState;
}
// biome-ignore lint/suspicious/noExplicitAny: CopilotKit message objects vary at runtime
@ -152,6 +152,8 @@ function ToolCallIndicator({
return <MessageIcon />;
case "get_document":
return <FileIcon />;
case "execute_skill":
return <MessageIcon />;
default:
return <SearchIcon />;
}
@ -165,6 +167,12 @@ function ToolCallIndicator({
return "Ask";
case "get_document":
return "Document";
case "execute_skill":
return "Skill";
case "analyze":
return "Analyze";
case "research":
return "Research";
default:
return toolName;
}
@ -172,38 +180,30 @@ function ToolCallIndicator({
const getDescription = () => {
switch (toolName) {
case "execute_skill": {
const skill = args.skill_name as string | undefined;
const request = args.request as string | undefined;
return (
<span className="tool-query">
{skill ? `${skill}: ` : ""}
{request ?? "Processing..."}
</span>
);
}
case "search": {
const query = args.query as string;
const docName = args.document_name as string | undefined;
return (
<>
<span className="tool-query">{query}</span>
{docName && (
<span className="tool-context">
{" "}
in <em>{docName}</em>
</span>
)}
</>
);
return <span className="tool-query">{query}</span>;
}
case "ask": {
const question = args.question as string;
const docName = args.document_name as string | undefined;
return (
<>
<span className="tool-query">{question}</span>
{docName && (
<span className="tool-context">
{" "}
from <em>{docName}</em>
</span>
)}
</>
);
return <span className="tool-query">{question}</span>;
}
case "get_document":
return <span className="tool-query">{args.query as string}</span>;
case "analyze":
return <span className="tool-query">{args.question as string}</span>;
case "research":
return <span className="tool-query">{args.question as string}</span>;
default:
return <span>Processing...</span>;
}
@ -231,7 +231,7 @@ function ToolCallIndicator({
}
// Context for sharing chat state with the message view
const ChatStateContext = createContext<ChatSessionState | null>(null);
const ChatStateContext = createContext<RAGState | null>(null);
// Wildcard tool call renderer for all server-side tools
const toolCallRenderers = [
@ -251,14 +251,15 @@ const toolCallRenderers = [
// Uses CopilotChatMessageView's children render prop to post-process the
// rendered message elements and inject citations at the right positions.
function MessageViewWithCitations({
messages,
isRunning,
messages = [],
isRunning = false,
}: {
// biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union
messages: any[];
isRunning: boolean;
messages?: any[];
isRunning?: boolean;
}) {
const chatState = useContext(ChatStateContext);
const ragState = useContext(ChatStateContext);
const citationsHistory = ragState ? deriveCitationsHistory(ragState) : [];
const cursor = isRunning ? (
<div key="cursor" className="streaming-cursor">
@ -271,7 +272,7 @@ function MessageViewWithCitations({
return (
<CopilotChatMessageView messages={messages} isRunning={isRunning}>
{({ messageElements }) => {
if (!chatState?.citations_history?.length) {
if (!citationsHistory.length) {
return (
<>
{messageElements}
@ -284,9 +285,9 @@ function MessageViewWithCitations({
// message (tool messages produce nothing). We correlate elements with
// messages to inject CitationBlocks after the right assistant responses.
//
// Both search and ask tools append to citations_history in order,
// Both search and ask tools append to citations via qa_history,
// so after each assistant text response that followed tool calls,
// we inject the next citations_history entry.
// we inject the next citations entry.
const result: React.ReactNode[] = [];
let citIdx = 0;
let seenToolCalls = false;
@ -317,10 +318,10 @@ function MessageViewWithCitations({
}
// After an assistant text response that followed tool calls,
// inject the next citations_history entry (one per turn)
// inject the next citations entry (one per turn)
if (msg.role === "assistant" && msg.content && seenToolCalls) {
if (citIdx < chatState.citations_history.length) {
const citations = chatState.citations_history[citIdx];
if (citIdx < citationsHistory.length) {
const citations = citationsHistory[citIdx];
if (citations?.length) {
result.push(
<CitationBlock
@ -350,6 +351,7 @@ function MessageViewWithCitations({
</CopilotChatMessageView>
);
}
MessageViewWithCitations.Cursor = CopilotChatMessageView.Cursor;
function ChatContentInner({
sessionId,
@ -358,8 +360,9 @@ function ChatContentInner({
sessionId: string;
onSessionChange: (id: string) => void;
}) {
const [contextOpen, setContextOpen] = useState(false);
const [filterOpen, setFilterOpen] = useState(false);
// Track selected document names locally (frontend-only)
const [selectedDocuments, setSelectedDocuments] = useState<string[]>([]);
const { agent } = useAgent({
agentId: "chat_agent",
@ -376,23 +379,10 @@ function ChatContentInner({
agent.threadId = sessionId;
}, [agent, sessionId]);
const chatState = normalizeChatState(
const ragState = normalizeRAGState(
(agent.state as AgentState)?.[AGUI_STATE_KEY],
);
const mergeChatState = (partial: Partial<ChatSessionState>) => {
const current = normalizeChatState(
(agent.state as AgentState)?.[AGUI_STATE_KEY],
);
agent.setState({
...agent.state,
[AGUI_STATE_KEY]: {
...current,
...partial,
},
});
};
// Restore session from localStorage when agent reference changes.
// useAgent returns a provisional agent initially, then the real agent
// after runtime connects — re-run restore each time so messages stick.
@ -400,9 +390,9 @@ function ChatContentInner({
if (agent.messages.length > 0) return;
const session = getSession(sessionId);
if (!session) return;
if (session.chatState) {
if (session.ragState) {
agent.setState({
[AGUI_STATE_KEY]: normalizeChatState(session.chatState),
[AGUI_STATE_KEY]: normalizeRAGState(session.ragState),
});
}
if (session.messages.length > 0) {
@ -412,21 +402,21 @@ function ChatContentInner({
}, [agent, sessionId]);
// Persist messages and state to localStorage.
// Read chatState from agent.state at effect time (not render time) so that
// Read ragState from agent.state at effect time (not render time) so that
// restore and persist effects in the same commit see consistent state.
// biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes
useEffect(() => {
if (sessionId && agent.messages.length > 0) {
const currentChatState = normalizeChatState(
const currentRagState = normalizeRAGState(
(agent.state as AgentState)?.[AGUI_STATE_KEY],
);
updateSessionMessages(
sessionId,
serializeMessages(agent.messages),
currentChatState,
currentRagState,
);
}
}, [JSON.stringify(agent.messages), chatState, sessionId]);
}, [JSON.stringify(agent.messages), ragState, sessionId]);
// biome-ignore lint/correctness/useExhaustiveDependencies: stable identity via agent ref
const messages = useMemo(
@ -458,24 +448,29 @@ function ChatContentInner({
}
}, [agent, ck]);
const sessionContext = chatState.session_context;
const documentFilter = chatState.document_filter;
const initialContext = chatState.initial_context ?? "";
// Context is locked after first message (qa_history has entries)
const isContextLocked = (chatState.qa_history?.length ?? 0) > 0;
const handleFilterApply = (selected: string[]) => {
mergeChatState({ document_filter: selected });
};
const handleInitialContextChange = (value: string) => {
if (isContextLocked) return;
mergeChatState({ initial_context: value || null });
setSelectedDocuments(selected);
// Convert selected document names to SQL filter for the backend
const filter =
selected.length > 0
? selected
.map(
(name) =>
`(title LIKE '%${name.replace(/'/g, "''")}%' OR uri LIKE '%${name.replace(/'/g, "''")}%')`,
)
.join(" OR ")
: null;
agent.setState({
...agent.state,
[AGUI_STATE_KEY]: {
...ragState,
document_filter: filter,
},
});
};
return (
<ChatStateContext.Provider value={chatState}>
<ChatStateContext.Provider value={ragState}>
<div className="chat-wrapper">
<div className="chat-container">
<div className="chat-header">
@ -485,36 +480,19 @@ function ChatContentInner({
/>
<button
type="button"
className={`header-btn ${documentFilter.length > 0 ? "has-content" : ""}`}
className={`header-btn ${selectedDocuments.length > 0 ? "has-content" : ""}`}
onClick={() => setFilterOpen(true)}
title={
documentFilter.length > 0
? `Filtering: ${documentFilter.length} document(s)`
selectedDocuments.length > 0
? `Filtering: ${selectedDocuments.length} document(s)`
: "Filter documents"
}
>
<FilterIcon />
{documentFilter.length > 0
? `Filter (${documentFilter.length})`
{selectedDocuments.length > 0
? `Filter (${selectedDocuments.length})`
: "Filter"}
</button>
<button
type="button"
className={`header-btn ${initialContext || sessionContext?.summary ? "has-content" : ""}`}
onClick={() => setContextOpen(true)}
title={
isContextLocked
? sessionContext?.summary
? "View session context"
: "No session context yet"
: initialContext
? "Edit initial context"
: "Set initial context"
}
>
<BrainIcon />
Memory
</button>
</div>
<div className="chat-content">
<CopilotChatView
@ -535,18 +513,10 @@ function ChatContentInner({
<DbInfo />
</div>
</div>
<ContextPanel
isOpen={contextOpen}
onClose={() => setContextOpen(false)}
sessionContext={sessionContext}
initialContext={initialContext}
onInitialContextChange={handleInitialContextChange}
isLocked={isContextLocked}
/>
<DocumentFilter
isOpen={filterOpen}
onClose={() => setFilterOpen(false)}
selected={documentFilter}
selected={selectedDocuments}
onApply={handleFilterApply}
/>
</ChatStateContext.Provider>

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

@ -1,131 +0,0 @@
"use client";
import { useCallback, useEffect, useId, useState } from "react";
import { formatRelativeTime } from "../lib/format";
import { BrainIcon } from "../lib/icons";
import type { SessionContext } from "../lib/sessionStorage";
interface ContextPanelProps {
isOpen: boolean;
onClose: () => void;
sessionContext: SessionContext | null;
initialContext?: string;
onInitialContextChange?: (value: string) => void;
isLocked?: boolean;
}
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) => {
if (e.key === "Escape") {
onClose();
}
},
[onClose],
);
const handleSave = useCallback(() => {
onInitialContextChange?.(localValue);
onClose();
}, [localValue, onInitialContextChange, onClose]);
if (!isOpen) {
return null;
}
const hasSessionContext = sessionContext?.summary?.trim();
// Show edit mode when: not locked AND no session context yet
const isEditMode = !isLocked && !hasSessionContext;
return (
<div
className="context-modal-overlay"
onClick={onClose}
onKeyDown={handleKeyDown}
role="dialog"
aria-modal="true"
aria-labelledby={titleId}
>
{/* biome-ignore lint/a11y/noStaticElementInteractions: modal content wrapper */}
<div
className="context-modal"
onClick={(e) => e.stopPropagation()}
onKeyDown={(e) => e.stopPropagation()}
>
<div className="context-modal-header">
<div className="context-modal-icon">
<BrainIcon size={24} strokeWidth={1.5} />
</div>
<h2 id={titleId} className="context-modal-title">
{isEditMode ? "Initial Context" : "Session Context"}
</h2>
</div>
<p className="context-modal-description">
{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>
{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">
<div className="context-empty-icon">
<BrainIcon size={24} strokeWidth={1.5} />
</div>
<div className="context-empty-text">
No context yet. Ask some questions to build context.
</div>
</div>
)}
<div className="context-footer">
<span className="context-timestamp">
{sessionContext?.last_updated
? `Last updated: ${formatRelativeTime(sessionContext.last_updated)}`
: ""}
</span>
<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>
);
}

57
app/frontend/lib/sessionStorage.ts
View file

@ -9,26 +9,33 @@ export interface Citation {
content: string;
}
export interface QAResponse {
export interface QAHistoryEntry {
question: string;
answer: string;
confidence: number;
citations: Citation[];
}
export interface SessionContext {
summary: string;
last_updated: string | null;
export interface DocumentInfo {
id: string;
title: string;
uri: string;
created: string;
}
export interface ChatSessionState {
initial_context: string | null;
export interface ResearchEntry {
question: string;
title: string;
executive_summary: string;
}
// Matches RAGState from the backend skill
export interface RAGState {
citations: Citation[];
citations_history: Citation[][];
qa_history: QAResponse[];
session_context: SessionContext | null;
document_filter: string[];
citation_registry: Record<string, number>;
qa_history: QAHistoryEntry[];
document_filter: string | null;
searches: Record<string, unknown[]>;
documents: DocumentInfo[];
reports: ResearchEntry[];
}
export interface StoredMessage {
@ -42,7 +49,7 @@ export interface StoredSession {
id: string;
title: string;
messages: StoredMessage[];
chatState: ChatSessionState;
ragState: RAGState;
createdAt: string;
updatedAt: string;
}
@ -50,18 +57,24 @@ export interface StoredSession {
const SESSIONS_KEY = "haiku.rag.sessions";
const ACTIVE_SESSION_KEY = "haiku.rag.activeSession";
export function normalizeChatState(state?: ChatSessionState): ChatSessionState {
export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
return {
initial_context: state?.initial_context ?? null,
citations: state?.citations ?? [],
citations_history: state?.citations_history ?? [],
qa_history: state?.qa_history ?? [],
session_context: state?.session_context ?? null,
document_filter: state?.document_filter ?? [],
citation_registry: state?.citation_registry ?? {},
document_filter: state?.document_filter ?? null,
searches: state?.searches ?? {},
documents: state?.documents ?? [],
reports: state?.reports ?? [],
};
}
// Derive per-turn citation arrays from qa_history
export function deriveCitationsHistory(state: RAGState): Citation[][] {
return state.qa_history
.filter((entry) => entry.citations?.length > 0)
.map((entry) => entry.citations);
}
export function getAllSessions(): StoredSession[] {
const raw = localStorage.getItem(SESSIONS_KEY);
if (!raw) return [];
@ -90,7 +103,7 @@ export function createSession(): StoredSession {
id: crypto.randomUUID(),
title: "New Session",
messages: [],
chatState: normalizeChatState(),
ragState: normalizeRAGState(),
createdAt: now,
updatedAt: now,
};
@ -115,7 +128,7 @@ export function saveSession(session: StoredSession): void {
export function updateSessionMessages(
id: string,
messages: StoredMessage[],
chatState: ChatSessionState,
ragState: RAGState,
): void {
const sessions = getAllSessions();
const idx = sessions.findIndex((s) => s.id === id);
@ -123,7 +136,7 @@ export function updateSessionMessages(
const session = sessions[idx];
session.messages = messages;
session.chatState = chatState;
session.ragState = ragState;
session.updatedAt = new Date().toISOString();
// Derive title from first user message

173
docs/agents/index.md
View file

@ -1,12 +1,13 @@
# Agents
Four agentic flows are provided by haiku.rag:
Three agentic flows are provided by haiku.rag:
- **Simple QA Agent** — a focused question answering agent
- **Chat Agent** — multi-turn conversational RAG with session memory
- **Research Graph** — a multi-step research workflow with question decomposition
- **RLM Agent** — complex analytical tasks via sandboxed Python code execution (see [RLM Agent](rlm.md))
For multi-turn conversational RAG, haiku.rag provides [skills](../skills/index.md) built on [haiku.skills](https://github.com/ggozad/haiku.skills). The skills bundle search, Q&A, analysis, and research tools with session state management.
See [QA and Research Configuration](../configuration/qa-research.md) for configuring model, iterations, concurrency, and other settings.
## Simple QA Agent
@ -26,9 +27,6 @@ haiku-rag ask "What is climate change?"
# With citations
haiku-rag ask "What is climate change?" --cite
# Deep mode (uses research graph with optimized settings)
haiku-rag ask "What are the main features of haiku.rag?" --deep
```
**Python usage:**
@ -48,121 +46,6 @@ async with HaikuRAG(path_to_db) as client:
print(answer)
```
## Chat Agent
The chat agent enables multi-turn conversational RAG. It is built from composable [toolsets](../tools.md) and maintains session state to improve follow-up answers.
Key features:
- **Composable toolsets**: Built from reusable `FunctionToolset` factories — see [Toolsets](../tools.md)
- **Semantic prior answer recall**: Similar prior Q/A pairs are retrieved and passed to the research planner, which can skip searching when they suffice
- **Background summarization**: After each `ask` call, the QA history is summarized into a compact session context for the next request
- **Document filtering**: Session-level or per-query document filtering
### CLI Usage
```bash
haiku-rag chat
haiku-rag chat --db /path/to/database.lancedb
```
See [Applications](../apps.md#chat-tui) for the full TUI interface guide.
### Python Usage
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.agents.chat import create_chat_agent, prepare_chat_context, ChatDeps
from haiku.rag.tools import ToolContext
agent = create_chat_agent(config)
async with HaikuRAG(path_to_db) as client:
context = ToolContext()
prepare_chat_context(context)
deps = ChatDeps(config=config, client=client, tool_context=context)
# First question
result = await agent.run("What is haiku.rag?", deps=deps)
print(result.output)
# Follow-up (uses session context)
result = await agent.run("How does it handle PDFs?", deps=deps)
print(result.output)
```
### Feature Selection
By default, `create_chat_agent` enables search, documents, and QA toolsets. You can customize which capabilities the agent has via the `features` parameter:
```python
from haiku.rag.agents.chat import (
create_chat_agent,
FEATURE_SEARCH,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_ANALYSIS,
)
# Search-only agent
agent = create_chat_agent(config, features=[FEATURE_SEARCH])
# All features including code analysis
agent = create_chat_agent(
config,
features=[FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS],
)
```
Available features:
| Feature | Constant | Tools added |
|---------|----------|-------------|
| Search | `FEATURE_SEARCH` | `search` |
| Documents | `FEATURE_DOCUMENTS` | `list_documents`, `get_document`, `summarize_document` |
| QA | `FEATURE_QA` | `ask` |
| Analysis | `FEATURE_ANALYSIS` | `analyze` |
The system prompt is automatically composed to match the selected features. See [Toolsets](../tools.md) for details on each toolset's parameters and behavior.
### Session State
Session state is managed through `ToolContext` — a namespace-based state container shared across all toolsets. The chat agent uses two namespaces:
**`SessionState`** (session management):
- `document_filter` — List of document titles/URIs to restrict searches
- `citation_registry` — Stable mapping of chunk IDs to citation indices
- `citations` — Citations from the current query
**`QASessionState`** (QA history and context):
- `qa_history` — List of previous Q/A pairs with embeddings
- `session_context` — Automatically maintained session context summary
For multi-session applications (e.g., web backends), use `ToolContextCache` to cache `ToolContext` instances by external session/thread ID:
```python
from haiku.rag.tools import ToolContext, ToolContextCache
cache = ToolContextCache() # TTL-based, defaults to 1 hour
context, _is_new = cache.get_or_create(thread_id)
```
**Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source.
### Conversational Memory
The chat agent maintains two layers of conversational memory:
**1. Semantic prior answer recall**
When the `ask` tool receives a question, it embeds the question and compares it against prior Q/A embeddings. Sufficiently similar prior answers are passed to the research planner, which can skip searching entirely if they already cover the question.
**2. Background session summarization**
After each `ask` call, a background task summarizes the full QA history into a compact session context. This summary is injected into the research planner on the next request, allowing it to resolve ambiguous references ("Tell me more about the authentication part") without having seen the full conversation.
## Research Graph
The research workflow is implemented as a typed pydantic-graph. It uses an iterative feedback loop where the planner proposes one question at a time, sees the answer, then decides whether to continue or synthesize.
@ -181,8 +64,7 @@ stateDiagram-v2
synthesize --> [*]
note right of plan_next
Receives session_context as background
and prior_answers from conversation history.
Uses prior_answers from previous iterations.
Uses a different prompt when prior answers exist.
end note
```
@ -190,8 +72,7 @@ stateDiagram-v2
The graph receives a `ResearchContext` containing:
- `original_question` — the user's question
- `session_context` — summary of conversation history (injected as `<background>` XML)
- `qa_responses` — prior answers from semantic matching or previous iterations (injected as `<prior_answers>` XML)
- `qa_responses` — prior answers from previous iterations (injected as `<prior_answers>` XML)
When prior answers are provided, the planner uses a context-aware prompt that evaluates whether existing evidence is sufficient. If it is, the planner marks `is_complete=True` and the graph skips directly to synthesis without any searches.
@ -201,20 +82,10 @@ When prior answers are provided, the planner uses a context-aware prompt that ev
- **search_one**: Answers a single question using the knowledge base (up to 3 search calls per question). Each answer is added to `ResearchContext.qa_responses` for the next planning iteration.
- **synthesize**: Generates the final output from all gathered evidence.
**Output modes:**
The graph supports two output modes via `build_research_graph(output_mode=...)`:
| Mode | Output type | Used by |
|------|-------------|---------|
| `"report"` | `ResearchReport` (title, executive summary, findings, conclusions, recommendations) | CLI `haiku-rag research`, Python API |
| `"conversational"` | `ConversationalAnswer` (answer, citations, confidence) | Chat agent's `ask` tool |
**Iterative flow:**
- Each iteration: planner evaluates context → proposes one question → search answers it → loop back
- Planner can decompose complex questions (e.g., "benefits and drawbacks" → start with "benefits")
- Session context resolves ambiguous references and informs planning
- Prior answers let the planner skip redundant searches
- Loop terminates when planner marks `is_complete=True` or `max_iterations` is reached
@ -277,40 +148,6 @@ async with HaikuRAG(path_to_db) as client:
report = await graph.run(state=state, deps=deps)
```
**Conversational mode with prior answers:**
```python
from haiku.rag.config import Config
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import SearchAnswer
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
# Conversational mode returns ConversationalAnswer instead of ResearchReport
graph = build_research_graph(config=Config, output_mode="conversational")
# Pass session context and prior answers from conversation history
context = ResearchContext(
original_question="How does it handle authentication?",
session_context="User is building a Python web app with FastAPI.",
qa_responses=[
SearchAnswer(
query="What authentication methods are supported?",
answer="JWT and OAuth2 are supported.",
confidence=0.95,
cited_chunks=["chunk-1"],
)
],
)
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
print(result.answer) # Direct conversational answer
print(result.confidence) # 0.0-1.0
print(result.citations) # Deduplicated citations from all searches
```
### Filtering Documents
Restrict searches to specific documents via the `search_filter` parameter:

10
docs/apps.md
View file

@ -37,7 +37,7 @@ Press `Ctrl+P` to open the command palette:
| Command | Description |
|---------|-------------|
| Memory | Edit initial context (before first message) or view session context (after) |
| View state | View the current session state |
| Filter documents | Select documents to restrict searches |
| Show database info | View document/chunk counts and storage info |
| Visual grounding | View chunk source location in document |
@ -46,12 +46,10 @@ Press `Ctrl+P` to open the command palette:
### Session Management
- Conversation history is maintained in memory for the session
- Previous Q/A pairs are used as context for follow-up questions
- Previous Q/A pairs are automatically used as context for follow-up questions via the `ask` tool
- Citations are tracked per response and can be inspected
- Document filter restricts all searches to selected documents
- 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
- Clearing chat resets session state
## Web Application
@ -63,7 +61,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
- Memory panel: set initial context before first message, view session context after
- Session state view for inspecting accumulated Q&A history, citations, and documents
### Quick Start

25
docs/architecture.md
View file

@ -24,7 +24,7 @@ flowchart TB
subgraph Agents["Agent Layer"]
QA[QA Agent]
Chat[Chat Agent]
Skill[RAG Skill]
Research[Research Graph]
RLM[RLM Agent]
end
@ -98,7 +98,7 @@ flowchart LR
### Agent Layer
Four agent types for different use cases:
Three agent types and a RAG skill for different use cases:
```mermaid
flowchart TB
@ -107,12 +107,12 @@ flowchart TB
S1 --> A1[Answer]
end
subgraph Chat["Chat Agent"]
subgraph Skill["RAG Skill"]
Q2[Question] --> Tools[Tool Selection]
Tools --> S2[Search / Ask / Get]
Tools --> S2[Search / Ask / Analyze]
S2 --> A2[Answer]
A2 --> History[Session History]
History -.-> Q2
A2 --> State[RAG State]
State -.-> Q2
end
subgraph Research["Research Graph"]
@ -138,19 +138,18 @@ flowchart TB
- Expands context around results
- Generates answer with optional citations
**Chat Agent** - Multi-turn conversational RAG:
**RAG Skill** - Multi-turn conversational RAG via [haiku.skills](https://github.com/ggozad/haiku.skills):
- Composed from reusable [toolsets](tools.md) (search, documents, QA, analysis)
- Maintains session history with prior answer recall
- Background summarization for context continuity
- Session-level document filtering
- Bundles search, list_documents, get_document, ask, analyze, and research tools
- Managed `RAGState` for session state (citations, QA history, document filters)
- Integrates with any pydantic-ai agent via `SkillToolset`
- Powers both the Chat TUI and web application
**Research Graph** - Iterative research workflow:
- Proposes one question at a time, evaluates the answer, then decides whether to continue
- Session context resolves ambiguous references
- Prior answers let the planner skip redundant searches
- Synthesizes structured report or conversational answer
- Synthesizes structured report
**RLM Agent** - Complex analytical tasks via code execution:

18
docs/benchmarks.md
View file

@ -60,27 +60,9 @@ evaluations run repliqa --config /path/to/haiku.rag.yaml --db /path/to/custom.la
- `--skip-qa` - Skip QA benchmark
- `--limit N` - Limit number of test cases
- `--name NAME` - Override the evaluation name
- `--deep` - Use deep QA mode (multi-step reasoning with research graph)
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
### Deep QA Mode
The `--deep` flag enables multi-step reasoning using the research graph instead of the simple QA agent:
```bash
evaluations run repliqa --skip-db --deep
```
In deep mode:
- Questions are decomposed into sub-questions by a planning agent
- Each sub-question is answered by searching the knowledge base
- A synthesis agent combines findings into a comprehensive answer
- The graph runs for up to 2 iterations with no early exit (confidence threshold disabled)
This matches the behavior of `haiku-rag ask --deep` in the CLI. Deep mode typically produces more thorough answers but requires more LLM calls per question.
## Methodology
### Retrieval Metrics

40
docs/cli.md
View file

@ -143,31 +143,17 @@ Ask questions with citations showing source documents:
haiku-rag ask "Who is the author of haiku.rag?" --cite
```
Use deep QA for complex questions (multi-agent decomposition):
```bash
haiku-rag ask "What are the main features and architecture of haiku.rag?" --deep --cite
```
Filter to specific documents:
```bash
haiku-rag ask "What are the main findings?" --filter "uri LIKE '%paper%'"
```
Provide background context for the question:
```bash
haiku-rag ask "What are the protocols?" --context "Focus on security best practices"
haiku-rag ask "Summarize the findings" --context-file background.txt
```
The QA agent searches your documents for relevant information and provides a comprehensive answer. When available, citations use the document title; otherwise they fall back to the URI.
Flags:
- `--cite`: Include citations showing which documents were used
- `--deep`: Decompose the question into sub-questions answered in parallel before synthesizing a final answer
- `--filter` / `-f`: Restrict searches to documents matching the filter (see [Filtering Search Results](python.md#filtering-search-results))
- `--context`: Background context for the question (passed to the agent as system context)
- `--context-file`: Path to a file containing background context
## Chat
@ -178,11 +164,6 @@ haiku-rag chat
haiku-rag chat --db /path/to/database.lancedb
```
Provide initial background context for the conversation:
```bash
haiku-rag chat --initial-context "Focus on Python programming concepts"
```
!!! note
Requires the `tui` extra: `pip install haiku.rag-slim[tui]` (included in full `haiku.rag` package)
@ -192,18 +173,6 @@ The chat interface provides:
- Expandable citations with source metadata
- Session memory for context-aware follow-up questions
- Visual grounding to inspect chunk source locations
- 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:
- `--initial-context`: Initial background context for the conversation (editable until first message)
See [Applications](apps.md#chat-tui) for keyboard shortcuts and features.
@ -242,18 +211,9 @@ Filter to specific documents:
haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'"
```
Provide background context for the research:
```bash
haiku-rag research "What are the safety protocols?" --context "Industrial manufacturing context"
haiku-rag research "Analyze the methodology" --context-file research-background.txt
```
Flags:
- `--filter` / `-f`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results))
- `--context`: Background context for the research
- `--context-file`: Path to a file containing background context
Research parameters like `max_iterations` and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.

2
docs/configuration/qa-research.md
View file

@ -40,8 +40,6 @@ qa:
- **max_iterations**: Maximum search iterations (default: 2)
- **max_concurrency**: Number of concurrent search operations (default: 1)
Deep QA mode (`haiku-rag ask --deep`) uses the research graph with a single iteration for quick, focused answers.
## Research Configuration
Configure the multi-agent research workflow:

20
docs/python.md
View file

@ -429,25 +429,23 @@ See [RLM Agent](agents/rlm.md) for details on capabilities and configuration.
## Building Custom Agents
haiku.rag provides composable toolset factories that can be mixed into any pydantic-ai agent. This lets you build custom agents with exactly the capabilities you need — search, document management, Q&A, or code analysis — sharing state across tools via `ToolContext`.
haiku.rag provides a RAG skill built on [haiku.skills](https://github.com/ggozad/haiku.skills) that bundles all capabilities into a composable agent:
```python
from pydantic_ai import Agent
from haiku.rag.tools import AgentDeps, build_toolkit
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
toolkit = build_toolkit(config, features=["search", "qa"])
skill = create_skill(db_path=db_path, config=config)
toolset = SkillToolset(skills=[skill])
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions=f"You are a helpful assistant.\n{toolkit.prompt}",
toolsets=toolkit.toolsets,
instructions=toolset.system_prompt,
toolsets=[toolset],
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = toolkit.create_context()
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What are the main findings?", deps=deps)
result = await agent.run("What are the main findings?")
```
See [Toolsets](tools.md) for the full API reference and composition guide, and the [`examples/`](https://github.com/ggozad/haiku.rag/tree/main/examples) directory for runnable scripts.
See [Toolsets](tools.md) for the full API reference.

72
docs/skills/index.md Normal file
View file

@ -0,0 +1,72 @@
# Skills
haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggozad/haiku.skills) skills. Skills are self-contained units that bundle tools, instructions, and state — they can be composed into any pydantic-ai agent via `SkillToolset`.
## Available Skills
| Skill | Description |
|-------|-------------|
| [`rag`](rag.md) | Search, retrieve, and answer questions from the knowledge base |
| [`rag-rlm`](rlm.md) | Computational analysis via code execution (requires Docker) |
## Discovery
Skills are registered as Python entrypoints under `haiku.skills`. They are discovered automatically by `haiku.skills`:
```bash
haiku-skills list --use-entrypoints
# rag — Search, retrieve and analyze documents using RAG.
# rag-rlm — Analyze documents using code execution in a Docker sandbox.
```
## Usage
```python
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
from pydantic_ai import Agent
skill = create_skill(db_path=db_path, config=config)
toolset = SkillToolset(skills=[skill])
agent = Agent(
"openai:gpt-4o",
instructions=toolset.system_prompt,
toolsets=[toolset],
)
result = await agent.run("What documents do we have?")
```
## Database Path Resolution
Both skills resolve the database path in the same order:
1. `db_path` argument passed to `create_skill()`
2. `HAIKU_RAG_DB` environment variable
3. Config default (`config.storage.data_dir / "haiku.rag.lancedb"`)
## State Management
Each skill manages its own state under a dedicated namespace. State is automatically synced via the AG-UI protocol when using `AGUIAdapter`.
```python
rag_state = toolset.get_namespace("rag")
rlm_state = toolset.get_namespace("rlm")
```
See the individual skill pages for state model details.
## AG-UI Streaming
For web applications, use pydantic-ai's `AGUIAdapter` to stream tool calls, text, and state deltas:
```python
from pydantic_ai.ag_ui import AGUIAdapter
adapter = AGUIAdapter(agent=agent, run_input=run_input)
event_stream = adapter.run_stream()
sse_event_stream = adapter.encode_stream(event_stream)
```
See the [Web Application](../apps.md#web-application) for a complete implementation.

47
docs/skills/rag.md Normal file
View file

@ -0,0 +1,47 @@
# RAG Skill
The RAG skill is the primary way to use haiku.rag tools. It bundles search, Q&A, document browsing, and research into a single skill with managed state.
## `create_skill(db_path?, config?)`
```python
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=db_path, config=config)
```
| Parameter | Default | Description |
|-----------|---------|-------------|
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
| `config` | `None` | `AppConfig` instance. If None, uses `get_config()`. |
## Tools
| Tool | Purpose |
|------|---------|
| `search(query, limit?)` | Hybrid search (vector + full-text) with context expansion |
| `list_documents(limit?, offset?, filter?)` | Paginated document listing |
| `get_document(query)` | Retrieve a document by ID, title, or URI |
| `ask(question)` | Q&A with citations via the QA agent |
| `research(question)` | Deep multi-agent research producing comprehensive reports |
## State
The skill manages a `RAGState` under the `"rag"` namespace:
```python
class RAGState(BaseModel):
citations: list[Citation] = []
qa_history: list[QAHistoryEntry] = []
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = {}
documents: list[DocumentInfo] = []
reports: list[ResearchEntry] = []
```
- **citations** — Accumulated citations from `ask` calls, with sequential indexing across calls.
- **qa_history** — Questions and answers from `ask` calls. Prior Q&A is used as context for follow-up questions when embeddings are similar.
- **document_filter** — SQL WHERE clause applied to `search`, `ask`, and `research` calls. Set this to scope queries to specific documents.
- **searches** — Search results keyed by query string.
- **documents** — Documents seen via `list_documents` or `get_document` (deduplicated by ID).
- **reports** — Research reports from `research` calls.

70
docs/skills/rlm.md Normal file
View file

@ -0,0 +1,70 @@
# RLM Skill
The RLM (Reflexion Language Model) skill provides computational analysis via code execution. It writes and runs Python code in an isolated Docker sandbox to answer questions that require computation, aggregation, or data traversal.
!!! warning "Requires Docker"
The `analyze` tool executes code in a Docker sandbox. Docker must be running on the host machine. This skill is not suitable for Docker-deployed applications — use the [`rag`](rag.md) skill alone in those environments.
## `create_skill(db_path?, config?)`
```python
from haiku.rag.skills.rlm import create_skill
skill = create_skill(db_path=db_path, config=config)
```
| Parameter | Default | Description |
|-----------|---------|-------------|
| `db_path` | `None` | Path to LanceDB database. Falls back to `HAIKU_RAG_DB` env var, then config default. |
| `config` | `None` | `AppConfig` instance. If None, uses `get_config()`. |
## Tools
| Tool | Purpose |
|------|---------|
| `analyze(question, document?, filter?)` | Answer analytical questions using code execution |
**Parameters:**
- `question` — The analytical question to answer.
- `document` — Optional document ID or title to pre-load for analysis.
- `filter` — Optional SQL WHERE clause to filter documents.
## State
The skill manages an `RLMState` under the `"rlm"` namespace:
```python
class RLMState(BaseModel):
analyses: list[AnalysisEntry] = []
class AnalysisEntry(BaseModel):
question: str
answer: str
program: str | None = None
```
Each `analyze` call appends an `AnalysisEntry` with the question, answer, and executed program.
## Usage with RAG Skill
Combine both skills to give the agent full RAG + analysis capabilities:
```python
from haiku.rag.skills.rag import create_skill as create_rag_skill
from haiku.rag.skills.rlm import create_skill as create_rlm_skill
from haiku.skills.agent import SkillToolset
from pydantic_ai import Agent
rag = create_rag_skill(db_path=db_path)
rlm = create_rlm_skill(db_path=db_path)
toolset = SkillToolset(skills=[rag, rlm])
agent = Agent(
"openai:gpt-4o",
instructions=toolset.system_prompt,
toolsets=[toolset],
)
```
See the [RLM Agent](../agents/rlm.md) documentation for details on how the underlying agent works.

369
docs/tools.md
View file

@ -1,77 +1,28 @@
# Toolsets
haiku.rag provides composable `FunctionToolset` factories in `haiku.rag.tools`. Each factory creates a pydantic-ai `FunctionToolset` that can be mixed into any agent. A shared `ToolContext` lets toolsets accumulate state (search results, citations, QA history) across invocations.
haiku.rag exposes its RAG capabilities as [haiku.skills](https://github.com/ggozad/haiku.skills) skills. See the [Skills](skills/index.md) section for the primary way to use haiku.rag tools.
## ToolContext
For lower-level access, `haiku.rag.tools` provides individual `FunctionToolset` factories used internally by agents.
`ToolContext` is a namespace-based state container. Toolsets register Pydantic models under string namespaces, and any toolset sharing the same context can read or write the same state.
## Low-Level Toolsets
For advanced use cases, individual toolset factories are available in `haiku.rag.tools`. These are used internally by the QA agent and can be composed into custom agents.
### RAGDeps Protocol
All toolsets use the `RAGDeps` protocol for dependency injection:
```python
from haiku.rag.tools import ToolContext
from haiku.rag.tools import RAGDeps
context = ToolContext()
class MyDeps:
def __init__(self, client: HaikuRAG):
self.client = client
```
### Registering and accessing state
### Search Toolset
```python
from pydantic import BaseModel
class MyState(BaseModel):
count: int = 0
context.register("my_namespace", MyState())
# Get state (returns None if not registered)
state = context.get("my_namespace")
# Get with type checking (returns None if wrong type)
state = context.get("my_namespace", MyState)
# Get or create (creates default if not registered)
state = context.get_or_create("my_namespace", MyState)
```
### Serialization
The entire context can be serialized and restored:
```python
# Serialize all namespaces (keyed by namespace)
data = context.dump_namespaces()
# {"my_namespace": {"count": 0}}
# Restore a namespace from serialized data
context.load_namespace("my_namespace", MyState, data["my_namespace"])
```
For AG-UI state management, use flat snapshots:
```python
# Flat snapshot of all namespaces (for AG-UI state)
snapshot = context.build_state_snapshot()
# {"document_filter": [], "citations": [], "citation_registry": {}, "qa_history": []}
# Restore from flat snapshot (updates registered namespaces in place)
context.restore_state_snapshot(snapshot)
```
### Preparing context for toolsets
`prepare_context()` registers the required namespaces for a given set of features:
```python
from haiku.rag.tools import ToolContext, prepare_context
context = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
```
This is idempotent and registers `SessionState` (for search, QA, and analysis features) and `QASessionState` (for QA). The chat agent's `prepare_chat_context()` is a thin wrapper that defaults to chat features and sets the AG-UI state key.
## Search Toolset
`create_search_toolset()` provides hybrid search (vector + full-text) with context expansion and citation tracking.
`create_search_toolset()` provides hybrid search with context expansion.
```python
from haiku.rag.tools import create_search_toolset
@ -79,24 +30,17 @@ from haiku.rag.tools import create_search_toolset
search = create_search_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `config` | required | `AppConfig` |
| `expand_context` | `True` | Expand results with surrounding chunks |
| `base_filter` | `None` | SQL WHERE clause applied to all searches |
| `tool_name` | `"search"` | Name of the tool exposed to the agent |
| `on_results` | `None` | Callback `(list[SearchResult]) -> None` invoked with results |
**Tool: `search(query, limit?, filter?)`**
### Document Toolset
Searches the knowledge base and returns formatted results. When a `ToolContext` with `SessionState` is registered, citations get stable indices via `citation_registry`.
**State:** Search results accumulate in `SearchState.results` under the `haiku.rag.search` namespace.
## Document Toolset
`create_document_toolset()` provides document browsing, retrieval, and summarization.
`create_document_toolset()` provides document browsing and retrieval.
```python
from haiku.rag.tools import create_document_toolset
@ -104,75 +48,20 @@ from haiku.rag.tools import create_document_toolset
docs = create_document_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig (used for summarization LLM) |
| `config` | required | `AppConfig` |
| `base_filter` | `None` | SQL WHERE clause for list operations |
**Tools:**
- `list_documents(page?)` — Paginated document listing (50 per page). Returns `DocumentListResponse` with document titles, URIs, and pagination info.
- `get_document(query)` — Retrieve a document by title or URI. Uses `find_document()` which tries exact URI match, then partial URI match, then partial title match.
- `list_documents(page?)` — Paginated document listing (50 per page).
- `get_document(query)` — Retrieve a document by title or URI.
- `summarize_document(query)` — Generate an LLM summary of a document's content.
## QA Toolset
### Analysis Toolset
`create_qa_toolset()` provides question answering via the research graph, with prior answer recall and background summarization.
```python
from haiku.rag.tools import create_qa_toolset
qa = create_qa_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `base_filter` | `None` | SQL WHERE clause applied to searches |
| `tool_name` | `"ask"` | Name of the tool exposed to the agent |
| `on_ask_complete` | `None` | Callback `(QASessionState, AppConfig) -> None` invoked after each QA cycle |
**Tool: `ask(question, document_name?)`**
Runs the research graph in conversational mode and returns a `QAResult`. When a `ToolContext` is provided:
- Prior answers from `QASessionState.qa_history` are matched via embedding similarity
- The answer is appended to `qa_history`
- `on_ask_complete` callback is invoked (if provided)
- Citations get stable indices via `SessionState.citation_registry`
**State:** QA history accumulates in `QASessionState` under the `haiku.rag.qa_session` namespace.
### Using `run_qa_core()` directly
For programmatic use without an agent, `run_qa_core()` provides the same QA flow:
```python
from haiku.rag.tools.qa import run_qa_core
result = await run_qa_core(
client=client,
config=config,
question="What are the main features?",
document_name="User Guide", # optional document filter
context=context, # optional ToolContext
session_context="User is building a web app", # optional
on_qa_complete=my_callback, # optional post-QA callback
)
print(result.answer)
print(result.confidence)
for citation in result.citations:
print(f" [{citation.index}] {citation.document_title}")
```
## Analysis Toolset
`create_analysis_toolset()` provides computational analysis via the RLM agent, which writes and executes Python code in a Docker sandbox.
`create_analysis_toolset()` provides computational analysis via the RLM agent (Docker sandbox).
```python
from haiku.rag.tools import create_analysis_toolset
@ -180,218 +69,16 @@ from haiku.rag.tools import create_analysis_toolset
analysis = create_analysis_toolset(config)
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `config` | required | `AppConfig` |
| `base_filter` | `None` | SQL WHERE clause applied to searches |
| `tool_name` | `"analyze"` | Name of the tool exposed to the agent |
**Tool: `analyze(task, document_name?)`**
Executes a computational task via code execution and returns an `AnalysisResult`. Requires Docker — see [RLM Agent](agents/rlm.md) for setup.
## Tool Prompts
`build_tools_prompt()` generates system prompt guidance for your toolsets — when to use each tool, the `document_name` parameter pattern, and usage examples. It's designed to be spliced into any agent's instructions alongside your own domain-specific guidance.
```python
from haiku.rag.tools import build_tools_prompt
# Generate guidance for the toolsets you're using
tools_prompt = build_tools_prompt(["search", "qa", "documents"])
```
Combine it with your own instructions:
```python
from pydantic_ai import Agent
from haiku.rag.tools import AgentDeps, build_tools_prompt
tools_prompt = build_tools_prompt(["search", "qa"])
agent = Agent(
"anthropic:claude-sonnet-4-5-20250929",
deps_type=AgentDeps,
instructions=f"""You are a medical research assistant.
{tools_prompt}
You also have access to:
- "check_interactions" - Use when the user asks about drug interactions.""",
toolsets=[search_toolset, qa_toolset, my_custom_toolset],
)
```
Available features: `"search"`, `"qa"`, `"documents"`, `"analysis"`.
## Composing Custom Agents
### Using `build_toolkit` (recommended)
`build_toolkit()` bundles toolsets, prompt, and context creation for a given feature set:
```python
from pydantic_ai import Agent
from haiku.rag.client import HaikuRAG
from haiku.rag.tools import AgentDeps, build_toolkit
toolkit = build_toolkit(config, features=["search", "documents", "qa"])
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions=f"You are a helpful research assistant.\n{toolkit.prompt}",
toolsets=toolkit.toolsets,
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = toolkit.create_context()
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What documents do we have about climate?", deps=deps)
print(result.output)
# Access accumulated state
from haiku.rag.tools.search import SearchState, SEARCH_NAMESPACE
search_state = context.get(SEARCH_NAMESPACE, SearchState)
if search_state:
print(f"Total search results: {len(search_state.results)}")
```
**Parameters:**
| Parameter | Default | Description |
|-----------|---------|-------------|
| `config` | required | AppConfig |
| `features` | `["search", "documents"]` | Features to enable |
| `base_filter` | `None` | SQL WHERE clause applied to all toolsets |
| `expand_context` | `True` | Expand search results with surrounding chunks |
| `on_qa_complete` | `None` | Callback invoked after each QA cycle |
**`Toolkit` properties:**
- `toolsets` — list of `FunctionToolset` instances to pass to the Agent
- `prompt` — tool guidance text for the system prompt
- `features` — the feature list this toolkit was built from
- `create_context(state_key=None)` — create a prepared `ToolContext` matching these features
- `prepare(context, state_key=None)` — register namespaces on an existing `ToolContext`
### Using individual factories
For full control, create toolsets individually with `create_*_toolset()`, `build_tools_prompt()`, and `prepare_context()`:
```python
from pydantic_ai import Agent
from haiku.rag.client import HaikuRAG
from haiku.rag.tools import (
AgentDeps,
ToolContext,
build_tools_prompt,
prepare_context,
create_search_toolset,
create_qa_toolset,
create_document_toolset,
)
search = create_search_toolset(config)
qa = create_qa_toolset(config)
docs = create_document_toolset(config)
features = ["search", "documents", "qa"]
tools_prompt = build_tools_prompt(features)
agent = Agent(
"openai:gpt-4o",
deps_type=AgentDeps,
instructions=f"You are a helpful research assistant.\n{tools_prompt}",
toolsets=[search, qa, docs],
)
async with HaikuRAG("path/to/db.lancedb") as client:
context = ToolContext()
prepare_context(context, features=features)
deps = AgentDeps(client=client, tool_context=context)
result = await agent.run("What documents do we have about climate?", deps=deps)
print(result.output)
```
`AgentDeps` satisfies the `RAGDeps` protocol and implements the AG-UI state protocol (`state` getter/setter). For AG-UI streaming, set `state_key` on the `ToolContext` (via `prepare_context` or `toolkit.create_context`):
```python
context = toolkit.create_context(state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
Tool functions access `client` and `tool_context` via pydantic-ai's `RunContext.deps`, so toolsets can be created once and reused across requests.
For complete runnable examples, see [`examples/custom_agent.py`](https://github.com/ggozad/haiku.rag/tree/main/examples/custom_agent.py) (standalone) and [`examples/custom_agent_agui.py`](https://github.com/ggozad/haiku.rag/tree/main/examples/custom_agent_agui.py) (AG-UI streaming server).
All toolsets respect session-level document filters when a `SessionState` is registered in the context. This means setting `SessionState.document_filter` restricts all tools simultaneously.
## AG-UI State Management
Both `AgentDeps` and `ChatDeps` implement the AG-UI `StateHandler` protocol. `ChatDeps` extends `AgentDeps` with chat-specific config and state handling. State is emitted under a namespaced key via `state_key` on the `ToolContext` — set it once via `prepare_context()`.
**Custom agents** use `AgentDeps` + `prepare_context`:
```python
from haiku.rag.tools import AgentDeps, ToolContext, ToolContextCache, prepare_context
context = ToolContext()
prepare_context(context, features=["search", "qa"], state_key="my_app")
deps = AgentDeps(client=client, tool_context=context)
```
**Chat agent** uses `ChatDeps` + `build_chat_toolkit` (adds chat-specific defaults like background summarization):
```python
from haiku.rag.agents.chat import (
AGUI_STATE_KEY, ChatDeps, build_chat_toolkit, create_chat_agent,
)
from haiku.rag.tools import ToolContextCache
chat_toolkit = build_chat_toolkit(config)
agent = create_chat_agent(config, toolkit=chat_toolkit)
# For multi-session apps, cache ToolContext per thread
cache = ToolContextCache()
context, is_new = cache.get_or_create(thread_id)
if is_new:
chat_toolkit.prepare(context, state_key=AGUI_STATE_KEY)
deps = ChatDeps(
config=config,
client=client,
tool_context=context,
)
```
The emitted state structure:
```json
{
"haiku.rag.chat": {
"citations": [],
"qa_history": [],
"session_context": null,
"document_filter": [],
"citation_registry": {}
}
}
```
State flows bidirectionally — the frontend sends its current state on each request, and the agent emits deltas (JSON Patch) reflecting server-side updates (new citations, QA history entries, session context). The server always prefers its own `session_context` over the client's value, since background summarization may have updated it between requests. See the [Web Application](apps.md#web-application) for a complete implementation.
## Filter Helpers
`haiku.rag.tools.filters` provides utilities for building SQL filters:
**`build_document_filter(document_name)`** — Builds a LIKE filter matching against both `uri` and `title`, case-insensitive. Also matches without spaces (e.g., "TB MED 593" matches "tbmed593").
**`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic.
**`combine_filters(filter1, filter2)`** — Combines two filters with AND logic. Returns `None` if both are `None`.
**`get_session_filter(context, base_filter?)`** — Extracts `document_filter` from `SessionState` in the `ToolContext`, builds a SQL filter from it, and combines with an optional `base_filter`.
- **`build_document_filter(document_name)`** — Builds a LIKE filter matching against both `uri` and `title`, case-insensitive. Also matches without spaces (e.g., "TB MED 593" matches "tbmed593").
- **`build_multi_document_filter(document_names)`** — Combines multiple document name filters with OR logic.
- **`combine_filters(filter1, filter2)`** — Combines two filters with AND logic. Returns `None` if both are `None`.

41
evaluations/evaluations/benchmark.py
View file

@ -17,9 +17,6 @@ from rich.progress import Progress
from evaluations.config import DatasetSpec
from evaluations.datasets import DATASETS
from evaluations.evaluators import ANSWER_EQUIVALENCE_RUBRIC
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
from haiku.rag.config.models import ModelConfig
@ -42,13 +39,11 @@ def build_experiment_metadata(
test_cases: int,
config: AppConfig,
judge_config: ModelConfig,
deep: bool = False,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
return {
"dataset": dataset_key,
"test_cases": test_cases,
"deep_ask": deep,
"embedder_provider": config.embeddings.model.provider,
"embedder_model": config.embeddings.model.name,
"embedder_dim": config.embeddings.model.vector_dim,
@ -270,7 +265,6 @@ async def run_qa_benchmark(
limit: int | None = None,
name: str | None = None,
db_path: Path | None = None,
deep: bool = False,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
corpus = spec.qa_loader()
if limit is not None:
@ -303,32 +297,19 @@ async def run_qa_benchmark(
db = spec.db_path(db_path)
async with HaikuRAG(db, config=config) as rag:
if deep:
graph = build_research_graph(config=config)
qa = get_qa_agent(rag, system_prompt=spec.system_prompt)
async def answer_question(question: str) -> str:
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=config)
deps = ResearchDeps(client=rag)
report = await graph.run(state=state, deps=deps)
return report.executive_summary if report else ""
else:
qa = get_qa_agent(rag, system_prompt=spec.system_prompt)
async def answer_question(question: str) -> str:
answer, _ = await qa.answer(question)
return answer
async def answer_question(question: str) -> str:
answer, _ = await qa.answer(question)
return answer
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
if deep:
eval_name = f"{eval_name}_deep"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
deep=deep,
)
report = await evaluation_dataset.evaluate(
@ -378,7 +359,6 @@ async def evaluate_dataset(
db_path: Path | None,
vacuum_interval: int = 100,
multimodal_only: bool = False,
deep: bool = False,
) -> None:
if not skip_db:
console.print(f"Using dataset: {spec.key}", style="bold magenta")
@ -398,11 +378,8 @@ async def evaluate_dataset(
)
if not skip_qa:
mode_label = "deep QA" if deep else "QA"
console.print(f"\nRunning {mode_label} benchmarks...", style="bold yellow")
await run_qa_benchmark(
spec, config, limit=limit, name=name, db_path=db_path, deep=deep
)
console.print("\nRunning QA benchmarks...", style="bold yellow")
await run_qa_benchmark(spec, config, limit=limit, name=name, db_path=db_path)
app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.")
@ -434,11 +411,6 @@ def run(
"--multimodal-only",
help="Only evaluate queries requiring image understanding.",
),
deep: bool = typer.Option(
False,
"--deep",
help="Use deep QA mode (multi-step reasoning with research graph).",
),
) -> None:
spec = DATASETS.get(dataset.lower())
if spec is None:
@ -477,7 +449,6 @@ def run(
db_path=db,
vacuum_interval=vacuum_interval,
multimodal_only=multimodal_only,
deep=deep,
)
)

4
examples/README.md
View file

@ -16,7 +16,7 @@ See `docker/README.md` for setup instructions.
**Script:** `custom_agent.py`
Composes `search`, `qa`, and `document` toolsets into a pydantic-ai `Agent` using `AgentDeps` and `prepare_context`. Shows how to run queries and inspect accumulated state (citations, QA history).
Uses the RAG skill with `SkillToolset` to build a conversational agent.
```bash
uv run python examples/custom_agent.py /path/to/db.lancedb
@ -26,7 +26,7 @@ uv run python examples/custom_agent.py /path/to/db.lancedb
**Script:** `custom_agent_agui.py`
A Starlette app that serves an AG-UI streaming endpoint using composed toolsets, `AgentDeps`, and `ToolContextCache` for multi-session support.
A Starlette app that serves an AG-UI streaming endpoint using the RAG skill with `SkillToolset`.
```bash
DB_PATH=/path/to/db.lancedb uv run uvicorn examples.custom_agent_agui:app --reload --port 8000

71
examples/custom_agent.py
View file

@ -1,7 +1,7 @@
"""Custom agent using haiku.rag composable toolsets.
"""Custom agent using the haiku.rag RAG skill.
Demonstrates how to compose search, QA, and document toolsets into a
pydantic-ai Agent using AgentDeps and prepare_context.
Demonstrates how to use the RAG skill with haiku.skills SkillToolset
to build a conversational agent.
Requirements:
- An Ollama instance running locally (default embedder)
@ -14,61 +14,36 @@ Usage:
import asyncio
import sys
from pathlib import Path
from pydantic_ai import Agent
from haiku.rag.client import HaikuRAG
from haiku.rag.tools import (
AgentDeps,
ToolContext,
build_tools_prompt,
create_document_toolset,
create_qa_toolset,
create_search_toolset,
prepare_context,
)
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
async def main(db_path: str) -> None:
async with HaikuRAG(db_path) as client:
# Compose toolsets into an agent
config = client.config
search_toolset = create_search_toolset(config)
qa_toolset = create_qa_toolset(config)
document_toolset = create_document_toolset(config)
skill = create_skill(db_path=Path(db_path))
toolset = SkillToolset(skills=[skill])
features = ["search", "documents", "qa"]
tools_prompt = build_tools_prompt(features)
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
instructions=toolset.system_prompt,
toolsets=[toolset],
)
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
deps_type=AgentDeps,
output_type=str,
instructions=(
"You are a helpful assistant with access to a knowledge base.\n"
f"{tools_prompt}"
),
toolsets=[search_toolset, qa_toolset, document_toolset],
)
print("Custom agent ready. Ctrl+C to exit.\n")
while True:
try:
user_input = input("You: ").strip()
except (EOFError, KeyboardInterrupt):
break
# Prepare a shared ToolContext
context = ToolContext()
prepare_context(context, features=["search", "documents", "qa"])
if not user_input:
continue
deps = AgentDeps(client=client, tool_context=context)
print("Custom agent ready. Ctrl+C to exit.\n")
while True:
try:
user_input = input("You: ").strip()
except (EOFError, KeyboardInterrupt):
break
if not user_input:
continue
result = await agent.run(user_input, deps=deps)
print(f"\nAgent: {result.output}\n")
result = await agent.run(user_input)
print(f"\nAgent: {result.output}\n")
if __name__ == "__main__":

66
examples/custom_agent_agui.py
View file

@ -1,7 +1,7 @@
"""Custom agent with AG-UI streaming.
A Starlette app that composes haiku.rag toolsets into an AG-UI compatible
agent. Multi-session support via ToolContextCache.
A Starlette app that serves an AG-UI streaming endpoint using the
haiku.rag RAG skill with haiku.skills SkillToolset.
Requirements:
- An Ollama instance running locally (default embedder)
@ -13,25 +13,18 @@ Usage:
import os
import sys
from pathlib import Path
from pydantic_ai import Agent
from pydantic_ai.ag_ui import AGUIAdapter
from pydantic_ai.ui import SSE_CONTENT_TYPE
from pydantic_ai.ui.ag_ui import AGUIAdapter
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.routing import Route
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools import (
AgentDeps,
ToolContextCache,
build_tools_prompt,
create_qa_toolset,
create_search_toolset,
prepare_context,
)
from haiku.rag.skills.rag import create_skill
from haiku.skills.agent import SkillToolset
db_path = os.environ.get("DB_PATH")
if not db_path:
@ -40,39 +33,13 @@ if not db_path:
)
sys.exit(1)
AGUI_STATE_KEY = "my_app"
skill = create_skill(db_path=Path(db_path))
toolset = SkillToolset(skills=[skill])
config = AppConfig()
# ToolContextCache maintains per-thread state across requests
context_cache = ToolContextCache()
# Singleton client
_client: HaikuRAG | None = None
def get_client() -> HaikuRAG:
global _client
if _client is None:
_client = HaikuRAG(db_path=db_path)
return _client
features = ["search", "qa"]
tools_prompt = build_tools_prompt(features)
# Create the agent once at module level
agent = Agent(
"anthropic:claude-haiku-4-5-20251001",
deps_type=AgentDeps,
output_type=str,
instructions=(
f"You are a helpful assistant with access to a knowledge base.\n{tools_prompt}"
),
toolsets=[
create_search_toolset(config),
create_qa_toolset(config),
],
instructions=toolset.system_prompt,
toolsets=[toolset],
)
@ -81,19 +48,8 @@ async def stream_chat(request: Request) -> Response:
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
run_input = AGUIAdapter.build_run_input(body)
thread_id = getattr(run_input, "thread_id", None) or "default"
context, is_new = context_cache.get_or_create(thread_id)
if is_new:
prepare_context(
context,
features=["search", "qa"],
state_key=AGUI_STATE_KEY,
)
deps = AgentDeps(client=get_client(), tool_context=context)
adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
event_stream = adapter.run_stream(deps=deps)
event_stream = adapter.run_stream()
sse_event_stream = adapter.encode_stream(event_stream)
return StreamingResponse(

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

@ -1,39 +0,0 @@
from haiku.rag.agents.chat.agent import (
DEFAULT_FEATURES,
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
ChatDeps,
build_chat_toolkit,
create_chat_agent,
prepare_chat_context,
run_chat_agent,
trigger_background_summarization,
)
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.agents.chat.state import (
AGUI_STATE_KEY,
ChatSessionState,
_rebuild_models,
)
from haiku.rag.tools.qa import QAHistoryEntry
_rebuild_models(QAHistoryEntry)
__all__ = [
"AGUI_STATE_KEY",
"DEFAULT_FEATURES",
"FEATURE_ANALYSIS",
"FEATURE_DOCUMENTS",
"FEATURE_QA",
"FEATURE_SEARCH",
"build_chat_prompt",
"build_chat_toolkit",
"create_chat_agent",
"prepare_chat_context",
"run_chat_agent",
"trigger_background_summarization",
"ChatDeps",
"ChatSessionState",
]

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

@ -1,217 +0,0 @@
from dataclasses import dataclass, field
from typing import Any
from pydantic_ai import Agent
from haiku.rag.agents.chat.context import (
trigger_background_summarization as _trigger_summarization,
)
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.deps import AgentDeps
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SessionContext
from haiku.rag.tools.toolkit import (
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
Toolkit,
build_toolkit,
)
from haiku.rag.utils import get_model
DEFAULT_FEATURES = [FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA]
def _on_qa_complete(qa_session_state: QASessionState, config: AppConfig) -> None:
_trigger_summarization(qa_session_state=qa_session_state, config=config)
@dataclass
class ChatDeps(AgentDeps):
"""Dependencies for chat agent.
Extends AgentDeps with chat-specific config and state handling.
"""
config: AppConfig = field(default_factory=AppConfig)
@AgentDeps.state.setter
def state(self, value: dict[str, Any] | None) -> None:
"""Set state from AG-UI protocol with chat-specific overrides."""
if value is None:
return
state_data = self._extract_state_data(value)
# Preserve server's session_context before restore overwrites it
qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
server_session_context = (
qa_session_state.session_context if qa_session_state is not None else None
)
self.tool_context.restore_state_snapshot(state_data)
# Chat-specific overrides after generic restore
if qa_session_state is not None:
# Prefer server's session_context (background summarizer may
# have updated it since the client's last snapshot).
if server_session_context is not None:
qa_session_state.session_context = server_session_context
# Handle initial_context -> session_context for first message
if qa_session_state.session_context is None:
if "initial_context" in state_data:
initial = state_data.get("initial_context")
if initial:
qa_session_state.session_context = SessionContext(
summary=initial
)
def build_chat_toolkit(
config: AppConfig,
features: list[str] | None = None,
) -> Toolkit:
"""Build a Toolkit configured for the chat agent.
Includes the on_qa_complete callback that triggers background
session summarization.
Args:
config: Application configuration.
features: List of features to enable. Defaults to DEFAULT_FEATURES.
Returns:
A Toolkit ready for chat agent composition and context creation.
"""
if features is None:
features = DEFAULT_FEATURES
return build_toolkit(config, features=features, on_qa_complete=_on_qa_complete)
def prepare_chat_context(
context: ToolContext,
features: list[str] | None = None,
) -> None:
"""Register required namespaces in a ToolContext for chat agent use.
Idempotent — safe to call multiple times on the same context.
Args:
context: ToolContext to prepare.
features: List of enabled features. Defaults to DEFAULT_FEATURES.
"""
from haiku.rag.tools.context import prepare_context
if features is None:
features = DEFAULT_FEATURES
prepare_context(context, features=features, state_key=AGUI_STATE_KEY)
def create_chat_agent(
config: AppConfig,
features: list[str] | None = None,
preamble: str | None = None,
toolkit: Toolkit | None = None,
) -> Agent[ChatDeps, str]:
"""Create the chat agent with composed toolsets.
Args:
config: Application configuration.
features: List of features to enable. Defaults to DEFAULT_FEATURES
(search, documents, qa). Available features: "search",
"documents", "qa", "analysis".
preamble: Optional custom identity/rules section for the system prompt.
When provided, replaces the default identity prompt. Tool guidance,
feature rules, and closing are still appended by the builder.
toolkit: Optional pre-built Toolkit. When provided, its toolsets are
used directly. When omitted, a toolkit is built from config and
features.
Returns:
The configured chat agent.
Example:
async with HaikuRAG(db_path, create=True) as client:
toolkit = build_chat_toolkit(config)
context = toolkit.create_context(state_key=AGUI_STATE_KEY)
agent = create_chat_agent(config, toolkit=toolkit)
deps = ChatDeps(config=config, client=client, tool_context=context)
result = await agent.run("Search for X", deps=deps)
"""
if features is None:
features = DEFAULT_FEATURES
if toolkit is None:
toolkit = build_chat_toolkit(config, features=features)
model = get_model(config.qa.model, config)
return Agent(
model,
deps_type=ChatDeps,
output_type=str,
instructions=build_chat_prompt(features, preamble=preamble),
toolsets=toolkit.toolsets,
retries=3,
)
def trigger_background_summarization(deps: ChatDeps) -> None:
"""Trigger background session summarization if qa_history has entries.
Call this after agent.run() or agent.run_stream() completes to update
the session context summary in the background.
Args:
deps: Chat dependencies with tool_context containing QASessionState.
"""
qa_session_state = deps.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_session_state is None or not qa_session_state.qa_history:
return
_trigger_summarization(
qa_session_state=qa_session_state,
config=deps.config,
)
async def run_chat_agent(
agent: Agent[ChatDeps, str],
deps: ChatDeps,
message: str,
) -> str:
"""Run the chat agent.
Args:
agent: The chat agent.
deps: Chat dependencies.
message: User message.
Returns:
Agent response.
"""
result = await agent.run(message, deps=deps)
return result.output
__all__ = [
"build_chat_toolkit",
"create_chat_agent",
"prepare_chat_context",
"run_chat_agent",
"trigger_background_summarization",
"ChatDeps",
"AGUI_STATE_KEY",
"FEATURE_SEARCH",
"FEATURE_DOCUMENTS",
"FEATURE_QA",
"FEATURE_ANALYSIS",
"DEFAULT_FEATURES",
]

148
haiku_rag_slim/haiku/rag/agents/chat/context.py
View file

@ -1,148 +0,0 @@
import asyncio
from datetime import datetime
from typing import TYPE_CHECKING
from pydantic_ai import Agent
from haiku.rag.agents.chat.prompts import SESSION_SUMMARY_PROMPT
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.session import SessionContext
from haiku.rag.utils import get_model
if TYPE_CHECKING:
from haiku.rag.tools.qa import QAHistoryEntry, QASessionState
# Track summarization tasks to allow cancellation
_summarization_tasks: dict[int, asyncio.Task[None]] = {}
async def summarize_session(
qa_history: list["QAHistoryEntry"],
config: AppConfig,
current_context: str | None = None,
) -> str:
"""Summarize qa_history into compact context.
Args:
qa_history: List of Q&A pairs from the conversation.
config: AppConfig for model selection.
current_context: Previous session_context.summary to incorporate.
The summarizer will build upon this.
Returns:
Markdown summary of the conversation history.
"""
if not qa_history:
return ""
model = get_model(config.qa.model, config)
agent: Agent[None, str] = Agent(
model,
output_type=str,
instructions=SESSION_SUMMARY_PROMPT,
retries=2,
)
history_text = _format_qa_history(qa_history)
if current_context:
history_text = f"## Current Context\n{current_context}\n\n{history_text}"
result = await agent.run(history_text)
return result.output
async def update_session_context(
qa_history: list["QAHistoryEntry"],
config: AppConfig,
current_context: str | None = None,
) -> SessionContext:
"""Summarize qa_history and return the resulting session context.
Args:
qa_history: List of Q&A pairs from the conversation.
config: AppConfig for model selection.
current_context: Previous summary to incorporate.
Returns:
The new SessionContext with summary and timestamp.
"""
summary = await summarize_session(
qa_history, config, current_context=current_context
)
return SessionContext(
summary=summary,
last_updated=datetime.now(),
)
def _format_qa_history(qa_history: list["QAHistoryEntry"]) -> str:
"""Format qa_history for input to summarization."""
lines: list[str] = []
for i, qa in enumerate(qa_history, 1):
lines.append(f"## Q{i}: {qa.question}")
lines.append(f"**Answer** (confidence: {qa.confidence:.0%}):")
lines.append(qa.answer)
if qa.sources:
lines.append(f"**Sources:** {', '.join(qa.sources)}")
lines.append("")
return "\n".join(lines)
async def _update_context_background(
qa_session_state: "QASessionState",
config: AppConfig,
) -> None:
"""Background task to update session context after an ask."""
try:
current_summary = (
qa_session_state.session_context.summary
if qa_session_state.session_context is not None
else None
)
result = await update_session_context(
qa_history=list(qa_session_state.qa_history),
config=config,
current_context=current_summary,
)
if result.summary:
qa_session_state.session_context = result
except asyncio.CancelledError:
pass
except Exception as e: # pragma: no cover
import logging
logging.getLogger(__name__).exception(f"Background summarization failed: {e}")
def trigger_background_summarization(
qa_session_state: "QASessionState",
config: AppConfig,
) -> None:
"""Trigger background session summarization if qa_history has entries.
Args:
qa_session_state: QASessionState with qa_history to summarize.
config: AppConfig for model selection.
"""
if not qa_session_state.qa_history:
return
key = id(qa_session_state)
# Cancel any existing summarization task for this state
if key in _summarization_tasks:
_summarization_tasks[key].cancel()
# Spawn background task
task = asyncio.create_task(
_update_context_background(
qa_session_state=qa_session_state,
config=config,
)
)
_summarization_tasks[key] = task
task.add_done_callback(lambda _t, k=key: _summarization_tasks.pop(k, None))

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

@ -1,91 +0,0 @@
from haiku.rag.tools.prompts import build_tools_prompt
_PROMPT_BASE = """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. NEVER call the same tool multiple times for a single user message
3. NEVER make up information - always use tools to get facts from the knowledge base"""
_PROMPT_QA_RULES = """
4. For questions: Use the "ask" tool EXACTLY ONCE - it automatically uses prior conversation context"""
_PROMPT_SEARCH_RULES = """
5. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally"""
_PROMPT_SEARCH_OUTPUT = """
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."""
_PROMPT_CLOSING = """
Be friendly and conversational."""
_PROMPT_QA_CLOSING = (
""" When you use the "ask" tool, summarize the key findings for the user."""
)
def build_chat_prompt(
features: list[str],
preamble: str | None = None,
) -> str:
"""Build a chat system prompt from the given feature list.
Each feature adds its relevant tool guidance to the prompt.
The base identity, critical rules, and closing are always included.
Args:
features: List of feature names (e.g., ["search", "documents", "qa"]).
preamble: Optional custom identity/rules section. When provided,
replaces the default identity prompt. Tool guidance, feature
rules, and closing are still appended.
Returns:
The composed system prompt string.
"""
parts = [preamble if preamble is not None else _PROMPT_BASE]
# Add feature-specific critical rules
if "qa" in features:
parts.append(_PROMPT_QA_RULES)
if "search" in features:
parts.append(_PROMPT_SEARCH_RULES)
# Tool guidance (reusable across agents)
tools_prompt = build_tools_prompt(features)
if tools_prompt:
parts.append(tools_prompt)
# Chat-specific search output rule
if "search" in features:
parts.append(_PROMPT_SEARCH_OUTPUT)
parts.append(_PROMPT_CLOSING)
if "qa" in features:
parts.append(_PROMPT_QA_CLOSING)
return "".join(parts)
CHAT_SYSTEM_PROMPT = build_chat_prompt(["search", "documents", "qa"])
SESSION_SUMMARY_PROMPT = """You are a session summarizer. Given a conversation history of Q&A pairs (and optionally existing context), produce a structured summary that captures key information for future context.
If a "Current Context" section is provided at the start of the input, incorporate that context into your summary. This might be initial background context from the user or a previous summary - build upon it rather than discard it.
Your summary should be concise (aim for 500-1500 tokens) and include:
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
3. **Current Focus** - What topic or question thread the user is currently exploring
Rules:
- Extract only high-signal information that would help answer follow-up questions
- When building on existing context, merge new information with prior context
- Omit small talk, greetings, or low-confidence answers
- Use bullet points for clarity
- Keep technical details but compress verbose explanations
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself."""

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

@ -1,33 +0,0 @@
from typing import TYPE_CHECKING
from pydantic import BaseModel
from haiku.rag.agents.research.models import Citation
from haiku.rag.tools.session import SessionContext
if TYPE_CHECKING:
from haiku.rag.tools.qa import QAHistoryEntry
AGUI_STATE_KEY = "haiku.rag.chat"
class ChatSessionState(BaseModel):
"""State shared between frontend and agent via AG-UI."""
initial_context: str | None = None
citations: list[Citation] = []
citations_history: list[list[Citation]] = []
qa_history: list["QAHistoryEntry"] = []
session_context: SessionContext | None = None
document_filter: list[str] = []
citation_registry: dict[str, int] = {}
def _rebuild_models(qa_history_entry_cls: type) -> None:
"""Resolve ChatSessionState forward reference to QAHistoryEntry.
Must be called after QAHistoryEntry is defined, passing the class.
"""
ChatSessionState.model_rebuild(
_types_namespace={"QAHistoryEntry": qa_history_entry_cls}
)

18
haiku_rag_slim/haiku/rag/agents/qa/agent.py
View file

@ -11,15 +11,14 @@ from haiku.rag.agents.research.models import (
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig, ModelConfig
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.search import SEARCH_NAMESPACE, SearchState, create_search_toolset
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.utils import get_model
@dataclass
class _QARunDeps:
client: HaikuRAG
tool_context: ToolContext | None = None
class QuestionAnswerAgent:
@ -47,11 +46,12 @@ class QuestionAnswerAgent:
Returns:
Tuple of (answer text, list of resolved citations)
"""
context = ToolContext()
accumulated_results: list[SearchResult] = []
search_toolset = create_search_toolset(
self._config,
base_filter=filter,
tool_name="search_documents",
on_results=accumulated_results.extend,
)
# Agent created per-call: toolset varies with filter, and Agent
@ -66,15 +66,9 @@ class QuestionAnswerAgent:
retries=3,
)
deps = _QARunDeps(client=self._client, tool_context=context)
deps = _QARunDeps(client=self._client)
result = await agent.run(question, deps=deps)
output = result.output
# Get search results from context for citation resolution
search_state = context.get(SEARCH_NAMESPACE)
search_results = (
search_state.results if isinstance(search_state, SearchState) else []
)
citations = resolve_citations(output.cited_chunks, search_results)
citations = resolve_citations(output.cited_chunks, accumulated_results)
return output.answer, citations

4
haiku_rag_slim/haiku/rag/agents/research/dependencies.py
View file

@ -16,10 +16,6 @@ class ResearchContext(BaseModel):
qa_responses: list[Any] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
session_context: str | None = Field(
default=None,
description="Session context from previous Q&A summarization",
)
def add_qa_response(self, qa: "SearchAnswer") -> None:
"""Add a structured QA response."""

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

@ -1,5 +1,4 @@
import asyncio
from typing import Literal, overload
from pydantic_ai import Agent, RunContext, format_as_xml
from pydantic_ai.output import ToolOutput
@ -7,15 +6,12 @@ from pydantic_graph.beta import Graph, GraphBuilder, StepContext
from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.agents.research.models import (
Citation,
ConversationalAnswer,
IterativePlanResult,
RawSearchAnswer,
ResearchReport,
SearchAnswer,
)
from haiku.rag.agents.research.prompts import (
CONVERSATIONAL_SYNTHESIS_PROMPT,
ITERATIVE_PLAN_PROMPT,
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT,
SEARCH_PROMPT,
@ -31,9 +27,6 @@ def format_context_for_prompt(context: ResearchContext) -> str:
"""Format the research context as XML for prompts."""
context_data: dict[str, object] = {}
if context.session_context:
context_data["background"] = context.session_context
context_data["question"] = context.original_question
if context.qa_responses:
@ -169,29 +162,13 @@ async def _search_one_step_logic(
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
@overload
def build_research_graph(
config: AppConfig = ...,
output_mode: Literal["report"] = ...,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: ...
@overload
def build_research_graph(
config: AppConfig = ...,
output_mode: Literal["conversational"] = ...,
) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]: ...
def build_research_graph(
config: AppConfig = Config,
output_mode: Literal["report", "conversational"] = "report",
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport | ConversationalAnswer]:
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the iterative research graph.
Args:
config: AppConfig object (uses config.research for provider, model, and graph parameters)
output_mode: Output format - "report" for ResearchReport, "conversational" for ConversationalAnswer
Returns:
Configured research graph with iterative planning
@ -199,18 +176,14 @@ def build_research_graph(
model_config = config.research.model
search_prompt = build_prompt(SEARCH_PROMPT, config)
if output_mode == "report":
synthesis_prompt = build_prompt(
config.prompts.synthesis or SYNTHESIS_PROMPT, config
)
else:
synthesis_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config)
synthesis_prompt = build_prompt(
config.prompts.synthesis or SYNTHESIS_PROMPT, config
)
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ResearchReport if output_mode == "report" else ConversationalAnswer,
output_type=ResearchReport,
)
@g.step
@ -236,81 +209,35 @@ def build_research_graph(
confidence=0.0,
)
if output_mode == "report":
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> ResearchReport:
"""Generate final research report."""
state = ctx.state
deps = ctx.deps
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> ResearchReport:
"""Generate final research report."""
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=synthesis_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=synthesis_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
return result.output
else:
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult],
) -> ConversationalAnswer:
"""Generate conversational answer from gathered evidence."""
state = ctx.state
deps = ctx.deps
agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[assignment]
model=get_model(model_config, config),
output_type=ConversationalAnswer,
instructions=synthesis_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
f"Answer the question based on the gathered evidence.\n\n{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
# Collect unique citations from qa_responses (dedupe by chunk_id)
seen_chunks: set[str] = set()
unique_citations: list[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)
return ConversationalAnswer(
answer=result.output.answer,
citations=unique_citations,
confidence=result.output.confidence,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
return result.output
# Build graph edges: iterative loop
#

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

@ -21,7 +21,7 @@ class IterativePlanResult(BaseModel):
class Citation(BaseModel):
"""Resolved citation with full metadata for display/visual grounding.
Used by both research graph and chat agent. The optional index field
Used by research graph and chat applications. The optional index field
supports UI display ordering in chat contexts.
"""
@ -111,18 +111,6 @@ def resolve_citations(
return citations
class ConversationalAnswer(BaseModel):
"""Conversational answer for chat context."""
answer: str = Field(description="Direct answer to the question")
citations: list[Citation] = Field(
default_factory=list, description="Citations supporting the answer"
)
confidence: float = Field(
default=1.0, description="Confidence score (0-1)", ge=0.0, le=1.0
)
class ResearchReport(BaseModel):
"""Final research report structure."""

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

@ -1,7 +1,5 @@
ITERATIVE_PLAN_PROMPT = """You are the research orchestrator planning the investigation.
If a <background> section is provided, use it to understand the conversation context.
Your task:
1. Analyze the original question
2. Propose the first question to investigate
@ -23,7 +21,6 @@ The question must be standalone and self-contained:
ITERATIVE_PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator evaluating gathered evidence.
You have access to context that may include:
- <background>: Domain context for the conversation
- <prior_answers>: Previous Q&A pairs with confidence scores
Your task:
@ -115,21 +112,3 @@ Style:
- Be professional, objective, and specific.
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
Instead, state the actual information directly."""
CONVERSATIONAL_SYNTHESIS_PROMPT = """Generate a direct, conversational answer
to the question based on the gathered evidence.
Output:
- 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.
- confidence: Score from 0.0 to 1.0 indicating answer quality.
Guidelines:
- Base your answer solely on the evidence provided in the context.
- If a <background> section is provided, use it to frame your answer appropriately.
- Be thorough - include all relevant information from the evidence.
- 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."""

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

@ -18,9 +18,6 @@ from rich.progress import (
)
from rich.syntax import Syntax
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import AppConfig, Config
from haiku.rag.mcp import create_mcp_server
@ -375,7 +372,6 @@ class HaikuRAGApp: # pragma: no cover
self,
question: str,
cite: bool = False,
deep: bool = False,
filter: str | None = None,
):
"""Ask a question using the RAG system.
@ -383,7 +379,6 @@ class HaikuRAGApp: # pragma: no cover
Args:
question: The question to ask
cite: Include citations in the answer
deep: Use deep QA mode (multi-step reasoning)
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
@ -392,46 +387,15 @@ class HaikuRAGApp: # pragma: no cover
read_only=self.read_only,
before=self.before,
) as self.client:
citations = []
if deep:
graph = build_research_graph(config=self.config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context,
config=self.config,
max_iterations=1,
)
state.search_filter = filter
deps = ResearchDeps(client=self.client)
answer, citations = await self.client.ask(question, filter=filter)
report = await graph.run(state=state, deps=deps)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
if report:
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(report.executive_summary))
if report.main_findings:
self.console.print()
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)
else:
self.console.print("[yellow]No answer generated.[/yellow]")
else:
answer, citations = await self.client.ask(question, filter=filter)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
if cite and citations:
for renderable in format_citations_rich(citations):
self.console.print(renderable)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[bold green]Answer:[/bold green]")
self.console.print(Markdown(answer))
if cite and citations:
for renderable in format_citations_rich(citations):
self.console.print(renderable)
async def rlm(
self,
@ -488,13 +452,7 @@ class HaikuRAGApp: # pragma: no cover
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
graph = build_research_graph(config=self.config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=self.config)
state.search_filter = filter
deps = ResearchDeps(client=client)
report = await graph.run(state=state, deps=deps)
report = await client.research(question=question, filter=filter)
if report is None:
self.console.print("[red]Research did not produce a report.[/red]")
@ -518,8 +476,6 @@ class HaikuRAGApp: # pragma: no cover
self.console.print(f"• {finding}")
self.console.print()
# (Themes section removed)
# Conclusions
if report.conclusions:
self.console.print("[bold cyan]Conclusions:[/bold cyan]")

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

@ -6,7 +6,7 @@ def run_chat(
db_path: Path | None = None,
read_only: bool = False,
before: datetime | None = None,
initial_context: str | None = None,
model: str | None = None,
) -> None:
"""Run the chat TUI.
@ -14,25 +14,29 @@ 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.
model: Model to use for the chat.
"""
try:
from haiku.rag.chat.app import ChatApp
except ImportError as e: # pragma: no cover
except ImportError as e:
raise ImportError(
"textual is not installed. Please install it with `pip install 'haiku.rag-slim[tui]'` or use the full haiku.rag package."
) from e
from haiku.rag.config import get_config
from haiku.rag.skills.rag import create_skill
config = get_config()
if db_path is None:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
skill = create_skill(db_path=db_path, config=config)
app = ChatApp(
db_path,
skill=skill,
read_only=read_only,
before=before,
initial_context=initial_context,
model=model,
)
app.run()

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

@ -1,30 +1,17 @@
# pyright: reportPossiblyUnboundVariable=false
import asyncio
import json
import uuid
from collections.abc import AsyncIterable, Iterable
from collections.abc import Iterable
from datetime import datetime
from pathlib import Path
from typing import TYPE_CHECKING, Any
from pydantic_ai import (
Agent,
AgentStreamEvent,
FunctionToolCallEvent,
FunctionToolResultEvent,
RunContext,
)
from pydantic_ai.messages import ModelMessage
from haiku.rag.agents.chat.agent import (
ChatDeps,
build_chat_toolkit,
create_chat_agent,
trigger_background_summarization,
)
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
from haiku.rag.skills.rag import AGENT_PREAMBLE, RAGState
from haiku.skills.agent import SkillToolset
from haiku.skills.models import Skill
if TYPE_CHECKING:
from textual.app import ComposeResult
@ -34,11 +21,26 @@ try:
logfire.configure(send_to_logfire="if-token-present", console=False)
logfire.instrument_pydantic_ai()
except ImportError: # pragma: no cover
except ImportError:
pass
try:
import textual_image.widget # noqa: F401 - import early for renderer detection
from ag_ui.core import (
AssistantMessage,
BaseEvent,
EventType,
RunAgentInput,
StateDeltaEvent,
TextMessageContentEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallStartEvent,
UserMessage,
)
from jsonpatch import JsonPatch
from pydantic_ai import Agent
from pydantic_ai.ag_ui import AGUIAdapter
from textual.app import App, SystemCommand
from textual.binding import Binding
from textual.widgets import Footer, Header, Input
@ -47,12 +49,15 @@ try:
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
TEXTUAL_AVAILABLE = True
except ImportError: # pragma: no cover
except ImportError:
TEXTUAL_AVAILABLE = False
App = object # type: ignore
SystemCommand = object # type: ignore
RAG_STATE_NAMESPACE = "rag"
class ChatApp(App):
"""Textual TUI for conversational RAG."""
@ -86,23 +91,25 @@ class ChatApp(App):
def __init__(
self,
db_path: Path,
skill: Skill,
read_only: bool = False,
before: datetime | None = None,
initial_context: str | None = None,
model: str | None = None,
) -> None:
super().__init__()
self.db_path = db_path
self._skill = skill
self.read_only = read_only
self.before = before
self._initial_context = initial_context
self._context_locked = False
self._model = model or "openai:gpt-4o"
self.client: HaikuRAG | None = None
self.config = get_config()
self.agent: Agent[ChatDeps, str] | None = None
self._toolset: SkillToolset | None = None
self._agent: Agent[None, str] | None = None
self._messages: list[Any] = []
self._state: dict[str, Any] = {}
self._is_processing = False
self._tool_call_widgets: dict[str, Any] = {}
self._current_worker: Worker[None] | None = None
self._message_history: list[ModelMessage] = []
self._document_filter: list[str] = []
def compose(self) -> "ComposeResult":
@ -136,9 +143,9 @@ class ChatApp(App):
self.action_show_info,
)
yield SystemCommand(
"Memory",
"View/edit context (editable before first message)",
self.action_show_context,
"View state",
"Show the current session state",
self.action_view_state,
)
async def on_mount(self) -> None:
@ -151,17 +158,14 @@ class ChatApp(App):
)
await self.client.__aenter__()
# Create toolkit, context, and agent
self.toolkit = build_chat_toolkit(self.config)
self.tool_context = self.toolkit.create_context(state_key=AGUI_STATE_KEY)
self.agent = create_chat_agent(self.config, toolkit=self.toolkit)
self._toolset = SkillToolset(skills=[self._skill])
self._agent = Agent(
self._model,
instructions=AGENT_PREAMBLE + self._toolset.system_prompt,
toolsets=[self._toolset],
)
self._state = self._toolset.build_state_snapshot()
# Sync document filter to tool context
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state is not None:
session_state.document_filter = self._document_filter
# Focus the input field
self.query_one(Input).focus()
async def on_unmount(self) -> None:
@ -169,60 +173,25 @@ class ChatApp(App):
if self.client:
await self.client.__aexit__(None, None, None)
async def _handle_stream_event(self, event: AgentStreamEvent) -> None:
"""Handle streaming events from the agent."""
chat_history = self.query_one(ChatHistory)
if isinstance(event, FunctionToolCallEvent):
tool_name = event.part.tool_name
tool_call_id = event.part.tool_call_id or str(uuid.uuid4())
args = event.part.args_as_dict()
widget = await chat_history.add_tool_call(tool_name, args)
self._tool_call_widgets[tool_call_id] = widget
elif isinstance(event, FunctionToolResultEvent):
tool_call_id = event.tool_call_id
if tool_call_id and tool_call_id in self._tool_call_widgets:
widget = self._tool_call_widgets[tool_call_id]
chat_history.mark_tool_complete(widget)
async def _event_stream_handler(
self,
_ctx: RunContext[ChatDeps],
event_stream: AsyncIterable[AgentStreamEvent],
) -> None:
"""Handle streaming events from the agent."""
async for event in event_stream:
await self._handle_stream_event(event)
# Yield to event loop to keep UI responsive
await asyncio.sleep(0)
async def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle user input submission."""
user_message = event.value.strip()
if not user_message or self._is_processing:
return
if not self.client or not self.agent:
return
# Lock context after first message
self._context_locked = True
# Clear the input
event.input.clear()
# Add user message to history
chat_history = self.query_one(ChatHistory)
await chat_history.add_message("user", user_message)
# Clear for new query
self._tool_call_widgets.clear()
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state:
session_state.citations.clear()
self._messages.append(
UserMessage(
id=str(uuid.uuid4()),
role="user",
content=user_message,
)
)
# Run agent in a worker to keep UI responsive
self._is_processing = True
self.query_one(Input).disabled = True
self._current_worker = self.run_worker(
@ -231,64 +200,87 @@ class ChatApp(App):
async def _run_agent(self, user_message: str) -> None:
"""Run the agent in a background worker."""
if not self.client or not self.agent:
if not self._agent or not self._toolset:
return
chat_history = self.query_one(ChatHistory)
# Show thinking indicator
await chat_history.show_thinking()
run_input = RunAgentInput(
thread_id="tui",
run_id=str(uuid.uuid4()),
messages=self._messages,
state=self._state,
tools=[],
context=[],
forwarded_props={},
)
adapter = AGUIAdapter(agent=self._agent, run_input=run_input)
message = None
accumulated_text = ""
tool_args_deltas: dict[str, str] = {}
try:
# Promote initial_context to QA session context on first run
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
qa_session_state = self.tool_context.get(
QA_SESSION_NAMESPACE, QASessionState
)
if qa_session_state is not None:
if qa_session_state.session_context is None and self._initial_context:
from haiku.rag.tools.session import SessionContext
qa_session_state.session_context = SessionContext(
summary=self._initial_context
async for event in adapter.run_stream():
if not isinstance(event, BaseEvent):
continue
if event.type == EventType.TEXT_MESSAGE_START:
chat_history.hide_thinking()
message = await chat_history.add_message("assistant")
accumulated_text = ""
elif event.type == EventType.TEXT_MESSAGE_CONTENT:
assert isinstance(event, TextMessageContentEvent)
accumulated_text += event.delta
if message:
message.update_content(accumulated_text)
chat_history.scroll_end(animate=False)
elif event.type == EventType.TEXT_MESSAGE_END:
self._messages.append(
AssistantMessage(
id=str(uuid.uuid4()),
role="assistant",
content=accumulated_text,
)
)
deps = ChatDeps(
config=self.config,
client=self.client,
tool_context=self.tool_context,
)
async with self.agent.run_stream(
user_message,
deps=deps,
message_history=self._message_history,
event_stream_handler=self._event_stream_handler,
) as stream:
# Hide thinking when we start getting content
chat_history.hide_thinking()
# Create assistant message for streaming
assistant_msg = await chat_history.add_message("assistant", "")
# Stream text updates
async for text in stream.stream_text():
assistant_msg.update_content(text)
chat_history.scroll_end(animate=False)
# Yield to event loop to keep UI responsive
await asyncio.sleep(0)
# Update message history with this conversation
self._message_history = stream.all_messages()
# Add citations from ToolContext
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state and session_state.citations:
await chat_history.add_citations(session_state.citations)
# Trigger background summarization
trigger_background_summarization(deps)
# Show citations from RAG state
await self._show_citations(chat_history)
elif event.type == EventType.TOOL_CALL_START:
assert isinstance(event, ToolCallStartEvent)
chat_history.hide_thinking()
await chat_history.add_tool_call(
event.tool_call_id, event.tool_call_name
)
tool_args_deltas[event.tool_call_id] = ""
await chat_history.show_thinking("Executing tasks...")
elif event.type == EventType.TOOL_CALL_ARGS:
assert isinstance(event, ToolCallArgsEvent)
tool_args_deltas[event.tool_call_id] = (
tool_args_deltas.get(event.tool_call_id, "") + event.delta
)
try:
args = json.loads(tool_args_deltas[event.tool_call_id])
chat_history.update_tool_args(event.tool_call_id, args)
except json.JSONDecodeError:
pass
elif event.type == EventType.TOOL_CALL_END:
assert isinstance(event, ToolCallEndEvent)
chat_history.mark_tool_complete(event.tool_call_id)
elif event.type == EventType.STATE_DELTA:
assert isinstance(event, StateDeltaEvent)
patch = JsonPatch(event.delta)
self._state = patch.apply(self._state)
self._toolset.restore_state_snapshot(self._state)
elif event.type == EventType.STATE_SNAPSHOT:
self._state = getattr(event, "snapshot", self._state)
self._toolset.restore_state_snapshot(self._state)
elif event.type == EventType.RUN_FINISHED:
chat_history.hide_thinking()
elif event.type == EventType.RUN_ERROR:
chat_history.hide_thinking()
error_msg = getattr(event, "message", "Unknown error")
await chat_history.add_message("assistant", f"Error: {error_msg}")
except asyncio.CancelledError:
chat_history.hide_thinking()
@ -303,20 +295,26 @@ class ChatApp(App):
chat_input.disabled = False
chat_input.focus()
async def _show_citations(self, chat_history: "ChatHistory") -> None:
"""Show citations from the RAG state after an agent response."""
if not self._toolset:
return
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if rag_state is None:
return
citations = getattr(rag_state, "citations", [])
if citations:
# Show only new citations (since last response)
await chat_history.add_citations(citations)
async def action_clear_chat(self) -> None:
"""Clear the chat history and reset session."""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
chat_history = self.query_one(ChatHistory)
await chat_history.clear_messages()
self._message_history.clear()
self._context_locked = False
# Re-register fresh states in ToolContext
self.tool_context.register(
SESSION_NAMESPACE,
SessionState(document_filter=self._document_filter),
)
self.tool_context.register(QA_SESSION_NAMESPACE, QASessionState())
self._messages.clear()
# Reset state
if self._toolset:
self._state = self._toolset.build_state_snapshot()
def action_focus_input(self) -> None:
"""Focus the input field, or cancel if processing."""
@ -340,7 +338,6 @@ class ChatApp(App):
if not self.client:
return
# Get citation from selected widget directly
chat_history = self.query_one(ChatHistory)
selected_widgets = list(chat_history.query(CitationWidget).filter(".selected"))
if not selected_widgets:
@ -364,38 +361,19 @@ class ChatApp(App):
await self.push_screen(InfoModal(self.client, self.db_path))
async def action_show_context(self) -> None:
"""Show context modal (edit initial context or view session context)."""
from haiku.rag.chat.widgets.context_modal import ContextModal
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
def action_view_state(self) -> None:
"""Show the current session state."""
from haiku.skills.chat.app import StateScreen
session_context = None
qa_session_state = self.tool_context.get(QA_SESSION_NAMESPACE, QASessionState)
if qa_session_state and qa_session_state.session_context is not None:
session_context = qa_session_state.session_context
await self.push_screen(
ContextModal(
initial_context=self._initial_context,
session_context=session_context,
is_locked=self._context_locked,
)
)
def on_context_modal_context_updated(self, event: Any) -> None:
"""Handle context updates from modal."""
if not self._context_locked:
self._initial_context = event.context or None
self.push_screen(StateScreen(self._state))
def on_citation_widget_selected(self, event: CitationWidget.Selected) -> None:
"""Handle citation selection."""
chat_history = self.query_one(ChatHistory)
# Remove selected class from all citations
for widget in chat_history.query(CitationWidget):
widget.remove_class("selected")
# Add selected class to the widget that was focused
event.widget.add_class("selected")
async def action_show_filter(self) -> None:
@ -414,7 +392,14 @@ class ChatApp(App):
def on_document_filter_modal_filter_changed(self, event: Any) -> None:
"""Handle document filter changes from modal."""
from haiku.rag.tools.filters import build_multi_document_filter
self._document_filter = event.selected
session_state = self.tool_context.get(SESSION_NAMESPACE, SessionState)
if session_state:
session_state.document_filter = self._document_filter
if self._toolset:
rag_state = self._toolset.get_namespace(RAG_STATE_NAMESPACE)
if isinstance(rag_state, RAGState):
rag_state.document_filter = build_multi_document_filter(
self._document_filter
)
self._state = self._toolset.build_state_snapshot()

125
haiku_rag_slim/haiku/rag/chat/widgets/chat_history.py
View file

@ -1,4 +1,4 @@
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
from textual.containers import Horizontal, VerticalScroll
from textual.message import Message
@ -32,52 +32,57 @@ class ChatMessage(Static):
class ToolCallWidget(Static):
"""Styled inline display of a tool call."""
"""Displays a single tool call with status indicator."""
TOOL_LABELS = {
"search": "Searching",
"ask": "Asking",
"get_document": "Fetching",
}
def __init__(self, tool_name: str, args: dict | None = None, **kwargs) -> None:
def __init__(
self,
tool_call_id: str,
tool_name: str,
args: dict[str, Any] | None = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
self.tool_call_id = tool_call_id
self.tool_name = tool_name
self.args = args or {}
self._complete = False
self._completed = False
def compose(self) -> "ComposeResult":
label = self.TOOL_LABELS.get(self.tool_name, self.tool_name)
# Build description based on tool type
if self.tool_name == "search":
query = self.args.get("query", "...")
doc = self.args.get("document_name")
desc = f'"{query}"'
if doc:
desc += f" in {doc}"
elif self.tool_name == "ask":
question = self.args.get("question", "...")
doc = self.args.get("document_name")
desc = f'"{question}"'
if doc:
desc += f" from {doc}"
elif self.tool_name == "get_document":
query = self.args.get("query", "...")
desc = f'"{query}"'
else:
desc = str(self.args) if self.args else ""
with Horizontal(classes="tool-row"):
if self._complete:
if self._completed:
yield Static("✓", classes="tool-status")
else:
yield LoadingIndicator(classes="tool-spinner")
yield Static(label, classes="tool-badge")
yield Static(desc, classes="tool-desc")
yield Static(self.tool_name, classes="tool-badge")
desc = self._build_description()
if desc:
yield Static(desc, classes="tool-desc")
def mark_complete(self) -> None:
self._complete = True
def _build_description(self) -> str:
if self.tool_name == "execute_skill":
skill = self.args.get("skill_name", "")
request = self.args.get("request", "...")
prefix = f"{skill}: " if skill else ""
return f'{prefix}"{request}"'
elif self.tool_name == "search":
query = self.args.get("query", "...")
return f'"{query}"'
elif self.tool_name == "ask":
question = self.args.get("question", "...")
return f'"{question}"'
elif self.tool_name == "get_document":
query = self.args.get("query", "...")
return f'"{query}"'
elif self.args:
return str(self.args)
return ""
def update_args(self, args: dict[str, Any]) -> None:
self.args = args
self.refresh(recompose=True)
def mark_completed(self) -> None:
self._completed = True
self.refresh(recompose=True)
@ -102,7 +107,6 @@ class CitationWidget(Collapsible):
pages += "..."
title += f" (p.{pages})"
# Build content widgets
content = citation.content
if len(content) > 500:
content = content[:500] + "..."
@ -132,10 +136,22 @@ class CitationWidget(Collapsible):
class ThinkingWidget(Static):
"""Thinking indicator shown while agent is processing."""
def __init__(self, text: str = "Thinking...", **kwargs) -> None:
super().__init__(**kwargs)
self._text = text
def compose(self) -> "ComposeResult":
with Horizontal(classes="thinking-row"):
yield LoadingIndicator(classes="thinking-spinner")
yield Static("Thinking...", classes="thinking-text")
yield Static(self._text, classes="thinking-text", id="thinking-label")
def update_text(self, text: str) -> None:
self._text = text
try:
label = self.query_one("#thinking-label", Static)
label.update(text)
except Exception:
pass
class SourcesHeader(Static):
@ -303,6 +319,7 @@ class ChatHistory(VerticalScroll):
def __init__(self, **kwargs) -> None:
super().__init__(**kwargs)
self.messages: list[tuple[str, str]] = []
self._tool_widgets: dict[str, ToolCallWidget] = {}
async def add_message(self, role: str, content: str = "") -> ChatMessage:
"""Add a message to the chat history."""
@ -313,18 +330,30 @@ class ChatHistory(VerticalScroll):
return message_widget
async def add_tool_call(
self, tool_name: str, args: dict | None = None
self,
tool_call_id: str,
tool_name: str,
args: dict[str, Any] | None = None,
) -> ToolCallWidget:
"""Add an inline tool call indicator."""
widget = ToolCallWidget(tool_name, args)
widget = ToolCallWidget(tool_call_id, tool_name, args)
self._tool_widgets[tool_call_id] = widget
await self.mount(widget)
self.scroll_end(animate=False)
return widget
def mark_tool_complete(self, widget: ToolCallWidget) -> None:
"""Mark a tool call as complete."""
widget.mark_complete()
widget.add_class("complete")
def update_tool_args(self, tool_call_id: str, args: dict[str, Any]) -> None:
"""Update the args of a tool call widget."""
widget = self._tool_widgets.get(tool_call_id)
if widget:
widget.update_args(args)
def mark_tool_complete(self, tool_call_id: str) -> None:
"""Mark a tool call as complete by its ID."""
widget = self._tool_widgets.get(tool_call_id)
if widget:
widget.mark_completed()
widget.add_class("complete")
async def add_citations(self, citations: list[Citation]) -> None:
"""Add citations inline after a response."""
@ -336,9 +365,12 @@ class ChatHistory(VerticalScroll):
await self.mount(widget)
self.scroll_end(animate=False)
async def show_thinking(self) -> None:
async def show_thinking(self, text: str = "Thinking...") -> None:
"""Show the thinking indicator."""
await self.mount(ThinkingWidget(id="thinking"))
try:
self.query_one("#thinking", ThinkingWidget).update_text(text)
except Exception:
await self.mount(ThinkingWidget(text, id="thinking"))
self.scroll_end(animate=False)
def hide_thinking(self) -> None:
@ -351,4 +383,5 @@ class ChatHistory(VerticalScroll):
async def clear_messages(self) -> None:
"""Clear all messages from the chat history."""
self.messages.clear()
self._tool_widgets.clear()
await self.remove_children()

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

@ -1,22 +1,12 @@
from typing import TYPE_CHECKING
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.message import Message
from textual.screen import ModalScreen
from textual.widgets import Button, Markdown, Static, TextArea
if TYPE_CHECKING:
from haiku.rag.tools.session import SessionContext
from textual.widgets import Button, Markdown, Static
class ContextModal(ModalScreen): # pragma: no cover
"""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
"""
class ContextModal(ModalScreen):
"""Modal screen for viewing session Q&A history."""
BINDINGS = [
Binding("escape", "cancel", "Close", show=False),
@ -49,12 +39,6 @@ class ContextModal(ModalScreen): # pragma: no cover
color: $text-muted;
}
#context-editor {
height: 12;
min-height: 8;
max-height: 16;
}
#context-content {
height: 1fr;
max-height: 16;
@ -73,81 +57,39 @@ class ContextModal(ModalScreen): # pragma: no cover
}
"""
class ContextUpdated(Message):
"""Emitted when the context is saved."""
def __init__(self, context: str) -> None:
super().__init__()
self.context = context
def __init__(
self,
initial_context: str | None = None,
session_context: "SessionContext | None" = None,
is_locked: bool = False,
) -> None:
def __init__(self, qa_history: list | None = None) -> None:
super().__init__()
self._initial_context = initial_context
self._session_context = session_context
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_context and self._session_context.summary
return not self._is_locked and not has_session_context
self._qa_history = qa_history or []
def compose(self) -> ComposeResult:
with Vertical(id="context-container"):
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 = self._initial_context or ""
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")
yield Static("[bold]Session Context[/bold]", id="context-header")
yield Static(
"Questions and answers from this session.",
id="context-description",
)
with VerticalScroll(id="context-content"):
yield Markdown(self._get_content())
with Horizontal(id="button-row"):
yield Button("Close", id="cancel-btn", variant="primary")
def _get_session_content(self) -> str:
if not self._session_context:
return "*No session context yet. Ask a question first.*"
def _get_content(self) -> str:
if not self._qa_history:
return "*No questions asked yet.*"
ctx = self._session_context
updated = (
ctx.last_updated.strftime("%Y-%m-%d %H:%M:%S")
if ctx.last_updated
else "unknown"
)
parts = []
for entry in self._qa_history:
q = getattr(entry, "question", str(entry))
a = getattr(entry, "answer", "")
parts.append(f"**Q:** {q}\n\n**A:** {a}")
return f"**Last updated:** {updated}\n\n---\n\n{ctx.summary}"
return "\n\n---\n\n".join(parts)
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))
"""Cancel and close."""
self.app.pop_screen()

2
haiku_rag_slim/haiku/rag/chat/widgets/document_filter_modal.py
View file

@ -8,7 +8,7 @@ from textual.widgets import Button, Checkbox, Input, Static
from haiku.rag.client import HaikuRAG
class DocumentFilterModal(ModalScreen): # pragma: no cover
class DocumentFilterModal(ModalScreen):
"""Modal screen for selecting documents to filter searches."""
BINDINGS = [

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

@ -341,11 +341,6 @@ def ask( # pragma: no cover
"--cite",
help="Include citations in the response",
),
deep: bool = typer.Option(
False,
"--deep",
help="Use deep multi-agent QA for complex questions",
),
filter: str | None = typer.Option(
None,
"--filter",
@ -358,7 +353,6 @@ def ask( # pragma: no cover
app.ask(
question=question,
cite=cite,
deep=deep,
filter=filter,
)
)
@ -626,10 +620,10 @@ def chat( # pragma: no cover
"--db",
help="Path to the LanceDB database file",
),
initial_context: str | None = typer.Option(
model: str | None = typer.Option(
None,
"--initial-context",
help="Initial background context to provide to the conversation",
"--model",
help="Model to use for the chat (e.g. openai:gpt-4o)",
),
):
"""Launch the chat TUI for conversational RAG."""
@ -641,7 +635,7 @@ def chat( # pragma: no cover
db_path,
read_only=_read_only,
before=_before,
initial_context=initial_context,
model=model,
)

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

@ -31,7 +31,10 @@ from haiku.rag.store.repositories.settings import SettingsRepository
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument
from haiku.rag.agents.research.models import Citation
from haiku.rag.agents.research.models import (
Citation,
ResearchReport,
)
from haiku.rag.agents.rlm.models import RLMResult
logger = logging.getLogger(__name__)
@ -1324,6 +1327,37 @@ class HaikuRAG:
qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt)
return await qa_agent.answer(question, filter=filter)
async def research(
self,
question: str,
*,
filter: str | None = None,
max_iterations: int | None = None,
) -> "ResearchReport":
"""Run multi-agent research to investigate a question.
Args:
question: The research question to investigate.
filter: SQL WHERE clause to filter documents.
max_iterations: Override max iterations (None uses config default).
Returns:
ResearchReport with structured findings.
"""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
graph = build_research_graph(config=self._config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context, config=self._config, max_iterations=max_iterations
)
state.search_filter = filter
deps = ResearchDeps(client=self)
return await graph.run(state=state, deps=deps)
async def rlm(
self,
question: str,

2
haiku_rag_slim/haiku/rag/inspector/__init__.py
View file

@ -1,6 +1,6 @@
try:
from haiku.rag.inspector.app import run_inspector
except ImportError as e: # pragma: no cover
except ImportError as e:
raise ImportError(
"textual is not installed. Please install it with `pip install 'haiku.rag-slim[tui]'` or use the full haiku.rag package."
) from e

6
haiku_rag_slim/haiku/rag/inspector/app.py
View file

@ -21,12 +21,12 @@ try:
from haiku.rag.inspector.widgets.search_modal import SearchModal
TEXTUAL_AVAILABLE = True
except ImportError: # pragma: no cover
except ImportError:
TEXTUAL_AVAILABLE = False
App = object # type: ignore
class InspectorApp(App): # pragma: no cover
class InspectorApp(App):
"""Textual TUI for inspecting LanceDB data."""
TITLE = "haiku.rag DB Inspector"
@ -243,7 +243,7 @@ def run_inspector(
db_path: Path | None = None,
read_only: bool = False,
before: datetime | None = None,
) -> None: # pragma: no cover
) -> None:
"""Run the inspector TUI.
Args:

2
haiku_rag_slim/haiku/rag/inspector/widgets/chunk_list.py
View file

@ -10,7 +10,7 @@ from haiku.rag.store.models import Chunk
BATCH_SIZE = 50
class ChunkList(VerticalScroll): # pragma: no cover
class ChunkList(VerticalScroll):
"""Widget for displaying and browsing chunks."""
can_focus = False

2
haiku_rag_slim/haiku/rag/inspector/widgets/context_modal.py
View file

@ -13,7 +13,7 @@ if TYPE_CHECKING:
from haiku.rag.store.models import Chunk
class ContextModal(Screen): # pragma: no cover
class ContextModal(Screen):
"""Modal screen for displaying how a chunk appears to agents."""
BINDINGS = [

2
haiku_rag_slim/haiku/rag/inspector/widgets/detail_view.py
View file

@ -16,7 +16,7 @@ class ProvenanceData(Protocol):
doc_item_refs: list[str]
class DetailView(VerticalScroll): # pragma: no cover
class DetailView(VerticalScroll):
"""Widget for displaying detailed content of documents or chunks."""
can_focus = True

2
haiku_rag_slim/haiku/rag/inspector/widgets/document_list.py
View file

@ -10,7 +10,7 @@ from haiku.rag.store.models import Document
BATCH_SIZE = 50
class DocumentList(VerticalScroll): # pragma: no cover
class DocumentList(VerticalScroll):
"""Widget for displaying and browsing documents."""
can_focus = False

2
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
View file

@ -14,7 +14,7 @@ if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
class InfoModal(ModalScreen): # pragma: no cover
class InfoModal(ModalScreen):
"""Modal screen for displaying database information."""
BINDINGS = [

2
haiku_rag_slim/haiku/rag/inspector/widgets/search_modal.py
View file

@ -11,7 +11,7 @@ from haiku.rag.inspector.widgets.detail_view import DetailView
from haiku.rag.store.models import Chunk, SearchResult
class SearchModal(Screen): # pragma: no cover
class SearchModal(Screen):
"""Screen for searching chunks."""
BINDINGS = [

2
haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py
View file

@ -15,7 +15,7 @@ if TYPE_CHECKING:
from haiku.rag.store.models import Chunk
class VisualGroundingModal(Screen): # pragma: no cover
class VisualGroundingModal(Screen):
"""Modal screen for displaying visual grounding with bounding boxes."""
BINDINGS = [

87
haiku_rag_slim/haiku/rag/mcp.py
View file

@ -2,26 +2,16 @@ from pathlib import Path
from typing import Any
from fastmcp import FastMCP
from pydantic import BaseModel
from haiku.rag.agents.research.models import ResearchReport
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, Config
from haiku.rag.store.models import SearchResult
from haiku.rag.store.models import Document, SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.utils import format_citations
class DocumentResult(BaseModel):
id: str | None
content: str
uri: str | None = None
title: str | None = None
metadata: dict[str, Any] = {}
created_at: str
updated_at: str
def create_mcp_server( # pragma: no cover
def create_mcp_server(
db_path: Path, config: AppConfig = Config, read_only: bool = False
) -> FastMCP:
"""Create an MCP server with the specified database path.
@ -111,24 +101,11 @@ def create_mcp_server( # pragma: no cover
return []
@mcp.tool()
async def get_document(document_id: str) -> DocumentResult | None:
async def get_document(document_id: str) -> Document | None:
"""Get a document by its ID."""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
document = await rag.get_document_by_id(document_id)
if document is None:
return None
return DocumentResult(
id=document.id,
content=document.content,
uri=document.uri,
title=document.title,
metadata=document.metadata,
created_at=str(document.created_at),
updated_at=str(document.updated_at),
)
return await rag.get_document_by_id(document_id)
except Exception:
return None
@ -137,30 +114,24 @@ def create_mcp_server( # pragma: no cover
limit: int | None = None,
offset: int | None = None,
filter: str | None = None,
) -> list[DocumentResult]:
) -> list[DocumentInfo]:
"""List all documents with optional pagination and filtering.
Args:
limit: Maximum number of documents to return.
offset: Number of documents to skip.
filter: Optional SQL WHERE clause to filter documents.
Returns:
List of DocumentResult instances matching the criteria.
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
documents = await rag.list_documents(limit, offset, filter)
return [
DocumentResult(
DocumentInfo(
id=doc.id,
content=doc.content,
uri=doc.uri,
title=doc.title,
metadata=doc.metadata,
created_at=str(doc.created_at),
updated_at=str(doc.updated_at),
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),
)
for doc in documents
]
@ -171,42 +142,19 @@ def create_mcp_server( # pragma: no cover
async def ask_question(
question: str,
cite: bool = False,
deep: bool = False,
) -> str:
"""Ask a question using the QA agent.
Args:
question: The question to ask.
cite: Whether to include citations in the response.
deep: Use deep multi-agent QA for complex questions that require decomposition.
Returns:
The answer as a string.
"""
try:
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
if deep:
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import (
ResearchDeps,
ResearchState,
)
graph = build_research_graph(config=config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
context=context,
config=config,
max_iterations=2,
)
deps = ResearchDeps(client=rag)
result = await graph.run(state=state, deps=deps)
answer = result.executive_summary
citations = []
else:
answer, citations = await rag.ask(question)
answer, citations = await rag.ask(question)
if cite and citations:
answer += "\n\n" + format_citations(citations)
return answer
@ -229,19 +177,8 @@ def create_mcp_server( # pragma: no cover
A research report with findings, or None if an error occurred.
"""
try:
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.state import ResearchDeps, ResearchState
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
graph = build_research_graph(config=config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=config)
deps = ResearchDeps(client=rag)
result = await graph.run(state=state, deps=deps)
return result
return await rag.research(question=question)
except Exception:
return None

0
haiku_rag_slim/haiku/rag/skills/__init__.py Normal file
View file

13
haiku_rag_slim/haiku/rag/skills/rag-rlm/SKILL.md Normal file
View file

@ -0,0 +1,13 @@
---
name: rag-rlm
description: >
Computational analysis of the knowledge base via code execution in a Docker sandbox.
Use for questions requiring counting, aggregation, statistics, data traversal,
comparison across documents, or any task best answered by writing Python code.
Examples: "how many pages?", "compare table 3 across documents",
"calculate average word count", "extract all email addresses".
---
# RLM Analysis
Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in an isolated Docker sandbox.

343
haiku_rag_slim/haiku/rag/skills/rag.py Normal file
View file

@ -0,0 +1,343 @@
import os
from pathlib import Path
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai import RunContext
from haiku.rag.agents.research.models import Citation
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.models import Skill, SkillSource
from haiku.skills.parser import parse_skill_md
from haiku.skills.state import SkillRunDeps
AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
CRITICAL RULES:
1. For greetings or casual chat: respond directly WITHOUT using any tools
2. NEVER make up information - always use tools to get facts from the knowledge base
3. For questions: Use the "ask" tool - it handles search and citation automatically
4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets
5. When you use the "ask" tool, summarize the key findings and always include citations in your response
"""
class ResearchEntry(BaseModel):
question: str
title: str
executive_summary: str
class RAGState(BaseModel):
citations: list[Citation] = Field(default_factory=list)
qa_history: list[QAHistoryEntry] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
documents: list[DocumentInfo] = Field(default_factory=list)
reports: list[ResearchEntry] = Field(default_factory=list)
def create_skill(
db_path: Path | None = None,
config: Any = None,
) -> Skill:
"""Create a RAG skill for searching and analyzing documents.
Args:
db_path: Path to the LanceDB database. Resolved from:
1. This argument
2. HAIKU_RAG_DB environment variable
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
if config is None:
config = get_config()
if db_path is None:
env_db = os.environ.get("HAIKU_RAG_DB")
if env_db:
db_path = Path(env_db).expanduser()
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
path = Path(__file__).parent / "rag"
metadata, instructions = parse_skill_md(path / "SKILL.md")
async def _find_relevant_prior_qa(
state: RAGState, query: str
) -> list[QAHistoryEntry]:
from haiku.rag.embeddings import get_embedder
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
from haiku.rag.utils import cosine_similarity
if not state.qa_history:
return []
embedder = get_embedder(config)
query_embedding = await embedder.embed_query(query)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
state.qa_history[idx].question_embedding = new_embeddings[i]
matches = []
for qa in state.qa_history:
if qa.question_embedding is not None:
similarity = cosine_similarity(query_embedding, qa.question_embedding)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matches.append(qa)
return matches
async def search(
ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
) -> str:
"""Search the knowledge base using hybrid search (vector + full-text).
Returns ranked results with content and metadata.
Args:
query: The search query.
limit: Maximum number of results.
"""
from haiku.rag.client import HaikuRAG
state = (
ctx.deps.state
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
else None
)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
results = await rag.search(
query,
limit=limit,
filter=state.document_filter if state else None,
)
results = await rag.expand_context(results)
if state:
state.searches[query] = list(results)
return "\n\n---\n\n".join(
r.format_for_agent(rank=i + 1, total=len(results))
for i, r in enumerate(results)
)
async def list_documents(
ctx: RunContext[SkillRunDeps],
limit: int | None = None,
offset: int | None = None,
filter: str | None = None,
) -> list[dict[str, Any]]:
"""List documents in the knowledge base with optional pagination and filtering.
Args:
limit: Maximum number of documents to return.
offset: Number of documents to skip.
filter: Optional SQL WHERE clause to filter documents.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = await rag.list_documents(limit, offset, filter)
result = [
{
"id": doc.id,
"title": doc.title,
"uri": doc.uri,
"metadata": doc.metadata,
"created_at": str(doc.created_at),
"updated_at": str(doc.updated_at),
}
for doc in documents
]
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
for doc_dict in result:
doc_info = DocumentInfo(
id=str(doc_dict["id"]),
title=doc_dict["title"] or "Untitled",
uri=doc_dict.get("uri") or "",
created=doc_dict.get("created_at", ""),
)
if not any(d.id == doc_info.id for d in ctx.deps.state.documents):
ctx.deps.state.documents.append(doc_info)
return result
async def get_document(
ctx: RunContext[SkillRunDeps], query: str
) -> dict[str, Any] | None:
"""Retrieve a document by ID, title, or URI.
Args:
query: Document ID, title, or URI to look up.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
document = await rag.resolve_document(query)
if document is None:
return None
result = {
"id": document.id,
"content": document.content,
"title": document.title,
"uri": document.uri,
"metadata": document.metadata,
"created_at": str(document.created_at),
"updated_at": str(document.updated_at),
}
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
doc_info = DocumentInfo(
id=str(result["id"]),
title=result["title"] or "Untitled",
uri=result.get("uri") or "",
created=result.get("created_at", ""),
)
if not any(d.id == doc_info.id for d in ctx.deps.state.documents):
ctx.deps.state.documents.append(doc_info)
return result
async def ask(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Ask a question and get an answer with citations from the knowledge base.
Args:
question: The question to ask.
"""
from haiku.rag.client import HaikuRAG
from haiku.rag.utils import format_citations
state = (
ctx.deps.state
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
else None
)
ask_question = question
if state:
matches = await _find_relevant_prior_qa(state, question)
if matches:
prior_parts = []
for qa in matches:
part = f"Q: {qa.question}\nA: {qa.answer}"
if qa.citations:
part += "\n" + format_citations(qa.citations)
prior_parts.append(part)
ask_question = (
"Context from prior questions in this session:\n\n"
+ "\n\n---\n\n".join(prior_parts)
+ "\n\n---\n\nCurrent question: "
+ question
)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
answer, citations = await rag.ask(
ask_question,
filter=state.document_filter if state else None,
)
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
next_index = len(ctx.deps.state.citations) + 1
for citation in citations:
citation.index = next_index
next_index += 1
ctx.deps.state.citations.extend(citations)
ctx.deps.state.qa_history.append(
QAHistoryEntry(question=question, answer=answer, citations=citations)
)
if citations:
answer += "\n\n" + format_citations(citations)
return answer
async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Conduct deep multi-agent research on a question.
Iteratively searches, analyzes, and synthesizes information from the
knowledge base to produce a comprehensive research report.
Only use when the user explicitly requests deep research.
Args:
question: The research question to investigate.
"""
from haiku.rag.client import HaikuRAG
state = (
ctx.deps.state
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
else None
)
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
report = await rag.research(
question, filter=state.document_filter if state else None
)
if state:
state.reports.append(
ResearchEntry(
question=question,
title=report.title,
executive_summary=report.executive_summary,
)
)
state.qa_history.append(
QAHistoryEntry(question=question, answer=report.executive_summary)
)
parts = [
f"# {report.title}",
f"\n## Executive Summary\n{report.executive_summary}",
]
if report.main_findings:
parts.append("\n## Main Findings")
for finding in report.main_findings:
parts.append(f"- {finding}")
if report.conclusions:
parts.append("\n## Conclusions")
for conclusion in report.conclusions:
parts.append(f"- {conclusion}")
if report.limitations:
parts.append("\n## Limitations")
for limitation in report.limitations:
parts.append(f"- {limitation}")
if report.recommendations:
parts.append("\n## Recommendations")
for rec in report.recommendations:
parts.append(f"- {rec}")
parts.append(f"\n## Sources\n{report.sources_summary}")
return "\n".join(parts)
return Skill(
metadata=metadata,
source=SkillSource.ENTRYPOINT,
path=path,
instructions=instructions,
tools=[
search,
list_documents,
get_document,
ask,
research,
],
state_type=RAGState,
state_namespace="rag",
)

33
haiku_rag_slim/haiku/rag/skills/rag/SKILL.md Normal file
View file

@ -0,0 +1,33 @@
---
name: rag
description: Search, retrieve and analyze documents using RAG (Retrieval Augmented Generation).
---
# RAG
You are a RAG (Retrieval Augmented Generation) assistant with access to a document knowledge base.
Use your tools to search and answer questions. 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 do you have?", "show me the documents", "list available docs").
- **get_document** — Use when the user wants the full content of a specific document (e.g., "get the paper about X", "show me document Y"). Accepts a document ID, title, or URI — partial matches work.
- **search** — Use when the user wants to find relevant passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns matching chunks with metadata.
- **ask** — Use for questions about topics in the knowledge base (e.g., "what is DocLayNet?", "explain the methodology"). Returns an answer with citations. Always include the citations in your response.
- **research** — Deep multi-agent research that produces comprehensive reports. **Only use when the user explicitly requests deep research** (e.g., "do a deep research on X", "research this topic thoroughly"). Never call this tool on your own — it is slow and expensive.
## When search returns irrelevant results
If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
- Use **ask** if the question is factual
- Report that the knowledge base doesn't contain relevant information
## When the user mentions a specific document
If the user says "search in [doc]", "find in [doc]", or "answer from [doc]":
- Extract the **topic** as the `query`/`question` parameter
- Use **get_document** or **list_documents** first to identify the document, then search/ask with a filter
Examples:
- "search for embeddings in the ML paper" -> first identify "ML paper", then search for "embeddings"
- "what does the DocLayNet paper say about annotations?" -> ask with question="what are the annotation methods?"

95
haiku_rag_slim/haiku/rag/skills/rlm.py Normal file
View file

@ -0,0 +1,95 @@
import os
from pathlib import Path
from typing import Any
from pydantic import BaseModel
from pydantic_ai import RunContext
from haiku.skills.models import Skill, SkillSource
from haiku.skills.parser import parse_skill_md
from haiku.skills.state import SkillRunDeps
class AnalysisEntry(BaseModel):
question: str
answer: str
program: str | None = None
class RLMState(BaseModel):
analyses: list[AnalysisEntry] = []
def create_skill(
db_path: Path | None = None,
config: Any = None,
) -> Skill:
"""Create an RLM analysis skill for computational document analysis.
Args:
db_path: Path to the LanceDB database. Resolved from:
1. This argument
2. HAIKU_RAG_DB environment variable
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
if config is None:
config = get_config()
if db_path is None:
env_db = os.environ.get("HAIKU_RAG_DB")
if env_db:
db_path = Path(env_db).expanduser()
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
path = Path(__file__).parent / "rag-rlm"
metadata, instructions = parse_skill_md(path / "SKILL.md")
async def analyze(
ctx: RunContext[SkillRunDeps],
question: str,
document: str | None = None,
filter: str | None = None,
) -> str:
"""Answer complex analytical questions using code execution.
Use this for questions requiring computation, aggregation, or
data traversal across documents.
Args:
question: The question to answer.
document: Optional document ID or title to pre-load for analysis.
filter: Optional SQL WHERE clause to filter documents.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = [document] if document else None
result = await rag.rlm(question, documents=documents, filter=filter)
output = result.answer
if result.program:
output += f"\n\nProgram:\n{result.program}"
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RLMState):
ctx.deps.state.analyses.append(
AnalysisEntry(
question=question,
answer=result.answer,
program=result.program,
)
)
return output
return Skill(
metadata=metadata,
source=SkillSource.ENTRYPOINT,
path=path,
instructions=instructions,
tools=[analyze],
state_type=RLMState,
state_namespace="rlm",
)

4
haiku_rag_slim/haiku/rag/store/models/document.py
View file

@ -43,8 +43,8 @@ class Document(BaseModel):
uri: str | None = None
title: str | None = None
metadata: dict = {}
docling_document: bytes | None = None
docling_version: str | None = None
docling_document: bytes | None = Field(default=None, exclude=True)
docling_version: str | None = Field(default=None, exclude=True)
created_at: datetime = Field(default_factory=datetime.now)
updated_at: datetime = Field(default_factory=datetime.now)

39
haiku_rag_slim/haiku/rag/tools/__init__.py
View file

@ -1,52 +1,23 @@
from haiku.rag.tools.analysis import create_analysis_toolset
from haiku.rag.tools.context import (
RAGDeps,
ToolContext,
ToolContextCache,
prepare_context,
)
from haiku.rag.tools.deps import AgentDeps
from haiku.rag.tools.analysis import AnalysisResult, create_analysis_toolset
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.document import create_document_toolset
from haiku.rag.tools.filters import (
build_document_filter,
build_multi_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.models import AnalysisResult, QAResult
from haiku.rag.tools.prompts import build_tools_prompt
from haiku.rag.tools.qa import create_qa_toolset
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD, QAHistoryEntry
from haiku.rag.tools.search import create_search_toolset
from haiku.rag.tools.toolkit import (
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
Toolkit,
build_toolkit,
)
__all__ = [
"AgentDeps",
"AnalysisResult",
"FEATURE_ANALYSIS",
"FEATURE_DOCUMENTS",
"FEATURE_QA",
"FEATURE_SEARCH",
"QAResult",
"PRIOR_ANSWER_RELEVANCE_THRESHOLD",
"QAHistoryEntry",
"RAGDeps",
"ToolContext",
"ToolContextCache",
"Toolkit",
"build_document_filter",
"build_multi_document_filter",
"build_toolkit",
"build_tools_prompt",
"combine_filters",
"create_analysis_toolset",
"create_document_toolset",
"create_qa_toolset",
"create_search_toolset",
"get_session_filter",
"prepare_context",
]

18
haiku_rag_slim/haiku/rag/tools/analysis.py
View file

@ -1,3 +1,4 @@
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext
from haiku.rag.agents.rlm.agent import create_rlm_agent
@ -8,9 +9,17 @@ from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import (
build_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.models import AnalysisResult
class AnalysisResult(BaseModel):
"""Result from the analysis toolset (RLM execution)."""
answer: str = Field(description="The answer produced by analysis")
code_executed: bool = Field(
default=True,
description="Whether code was executed to produce this answer",
)
def create_analysis_toolset(
@ -47,12 +56,9 @@ def create_analysis_toolset(
AnalysisResult with answer and execution metadata.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(
get_session_filter(tool_context, base_filter), doc_filter
)
effective_filter = combine_filters(base_filter, doc_filter)
rlm_context = RLMContext(filter=effective_filter)

257
haiku_rag_slim/haiku/rag/tools/context.py
View file

@ -1,13 +1,8 @@
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Protocol, TypeVar, overload, runtime_checkable
from pydantic import BaseModel, PrivateAttr
from typing import TYPE_CHECKING, Protocol, runtime_checkable
if TYPE_CHECKING:
from haiku.rag.client import HaikuRAG
T = TypeVar("T", bound=BaseModel)
@runtime_checkable
class RAGDeps(Protocol):
@ -18,253 +13,3 @@ class RAGDeps(Protocol):
"""
client: "HaikuRAG"
tool_context: "ToolContext | None"
class ToolContext(BaseModel):
"""Generic state container for haiku.rag toolsets.
Toolsets register their own Pydantic model state under namespaces.
Multiple toolsets can share state by registering under the same namespace.
All registered states must be Pydantic BaseModel subclasses, making
the entire context serializable via model_dump()/model_validate().
Example:
# Define toolset-specific state
class SearchState(BaseModel):
results: list[SearchResult] = []
filter: str | None = None
SEARCH_NAMESPACE = "haiku.rag.search"
# In toolset factory
def create_search_toolset(config):
async def search(ctx: RunContext[RAGDeps], query: str):
tool_context = ctx.deps.tool_context
if tool_context:
state = tool_context.get_or_create(SEARCH_NAMESPACE, SearchState)
...
# Usage
search_tools = create_search_toolset(config)
agent = Agent(..., toolsets=[search_tools])
await agent.run("...", deps=my_deps)
# Access accumulated state
search_state = context.get(SEARCH_NAMESPACE)
for result in search_state.results:
print(f"{result.document_title}")
# Serialize entire context
ns_data = context.dump_namespaces()
"""
state_key: str | None = None
_namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict)
_client_snapshot: dict[str, Any] | None = PrivateAttr(default=None)
def register(self, namespace: str, state: BaseModel) -> None:
"""Register state for a namespace.
Args:
namespace: Unique identifier for the toolset (e.g., "haiku.rag.search")
state: A Pydantic BaseModel instance to store
Overwrites any existing state for the namespace.
"""
self._namespaces[namespace] = state
@overload
def get(self, namespace: str) -> BaseModel | None: ...
@overload
def get(self, namespace: str, state_type: type[T]) -> T | None: ...
def get(
self, namespace: str, state_type: type[T] | None = None
) -> BaseModel | T | None:
"""Get state for a namespace, or None if not registered.
When state_type is provided, returns the state only if it matches
the expected type, otherwise returns None.
"""
state = self._namespaces.get(namespace)
if state_type is not None:
return state if isinstance(state, state_type) else None
return state
def get_or_create(self, namespace: str, state_type: type[T]) -> T:
"""Get state for a namespace, creating it if not registered.
Args:
namespace: The namespace to get or create state for.
state_type: A Pydantic BaseModel subclass to instantiate if needed.
Returns:
The state for the namespace.
"""
if namespace not in self._namespaces:
self._namespaces[namespace] = state_type()
return self._namespaces[namespace] # type: ignore[return-value]
def clear_namespace(self, namespace: str) -> None:
"""Clear state for a specific namespace."""
if namespace in self._namespaces:
del self._namespaces[namespace]
def clear_all(self) -> None:
"""Clear all namespaces."""
self._namespaces.clear()
@property
def namespaces(self) -> list[str]:
"""List all registered namespaces."""
return list(self._namespaces.keys())
@property
def client_snapshot(self) -> dict[str, Any] | None:
"""Snapshot captured after the last restore_state_snapshot call.
Represents what the client has, before any server-side overrides.
Tools use this as the baseline for delta computation so that
server-side changes (e.g. background summarization) are included.
"""
return self._client_snapshot
def dump_namespaces(self) -> dict[str, dict[str, Any]]:
"""Serialize all namespace states to a dictionary.
Returns:
Dict mapping namespace -> serialized state dict.
"""
return {ns: state.model_dump() for ns, state in self._namespaces.items()}
def build_state_snapshot(self) -> dict[str, Any]:
"""Build a flat snapshot of all namespace states for AG-UI.
Merges model_dump(mode="json") from every registered namespace
into a single flat dict.
Returns:
Combined dict of all namespace fields.
"""
snapshot: dict[str, Any] = {}
for state in self._namespaces.values():
snapshot.update(state.model_dump(mode="json"))
return snapshot
def restore_state_snapshot(self, data: dict[str, Any]) -> None:
"""Restore namespace states from a flat snapshot dict.
For each registered namespace, finds matching fields in *data*,
validates them via the namespace model, and updates the state
in place. Fields not present in *data* are left unchanged.
After restoring, captures a snapshot as ``client_snapshot`` so
tools can compute deltas against what the client actually has.
Args:
data: Flat dict as produced by build_state_snapshot().
"""
for state in self._namespaces.values():
model_fields = state.model_fields
matching = {k: v for k, v in data.items() if k in model_fields}
if matching:
# Fill in current values for fields not in data
current = state.model_dump()
current.update(matching)
updated = state.model_validate(current)
for field_name in matching:
setattr(state, field_name, getattr(updated, field_name))
self._client_snapshot = self.build_state_snapshot()
def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T:
"""Deserialize and register state for a namespace.
Args:
namespace: The namespace to register the state under.
state_type: The Pydantic model class to deserialize into.
data: The serialized state data.
Returns:
The deserialized and registered state.
"""
state = state_type.model_validate(data)
self._namespaces[namespace] = state
return state
def prepare_context(
context: ToolContext,
features: list[str] | None = None,
state_key: str | None = None,
) -> None:
"""Register required namespaces in a ToolContext based on feature flags.
Idempotent — safe to call multiple times on the same context.
Args:
context: ToolContext to prepare.
features: List of enabled features. Defaults to ["search", "documents"].
state_key: Optional AG-UI state key to set on the context.
"""
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
if features is None:
features = ["search", "documents"]
if any(f in features for f in ("search", "qa", "analysis")):
context.get_or_create(SESSION_NAMESPACE, SessionState)
if "qa" in features:
context.get_or_create(QA_SESSION_NAMESPACE, QASessionState)
if state_key is not None:
context.state_key = state_key
class ToolContextCache:
"""In-memory cache for ToolContext instances, keyed by external session/thread ID."""
def __init__(self, ttl: timedelta = timedelta(hours=1)) -> None:
self._cache: dict[str, ToolContext] = {}
self._timestamps: dict[str, datetime] = {}
self._ttl = ttl
def get_or_create(self, key: str) -> tuple[ToolContext, bool]:
"""Get an existing context or create a new one.
Returns:
Tuple of (context, is_new) where is_new is True if a new context was created.
"""
self._cleanup()
if key in self._cache:
self._timestamps[key] = datetime.now()
return self._cache[key], False
context = ToolContext()
self._cache[key] = context
self._timestamps[key] = datetime.now()
return context, True
def remove(self, key: str) -> None:
"""Remove a specific key from the cache."""
self._cache.pop(key, None)
self._timestamps.pop(key, None)
def clear(self) -> None:
"""Clear all entries."""
self._cache.clear()
self._timestamps.clear()
def _cleanup(self) -> None:
"""Remove entries older than TTL."""
now = datetime.now()
expired = [
key for key, ts in self._timestamps.items() if (now - ts) >= self._ttl
]
for key in expired:
self._cache.pop(key, None)
self._timestamps.pop(key, None)

42
haiku_rag_slim/haiku/rag/tools/deps.py
View file

@ -1,42 +0,0 @@
from dataclasses import dataclass
from typing import Any
from haiku.rag.client import HaikuRAG
from haiku.rag.tools.context import ToolContext
@dataclass
class AgentDeps:
"""Generic dependencies for agents using haiku.rag toolsets.
Implements RAGDeps protocol and AG-UI state protocol.
"""
client: HaikuRAG
tool_context: ToolContext
@property
def state(self) -> dict[str, Any]:
"""Get current state for AG-UI protocol."""
snapshot = self.tool_context.build_state_snapshot()
state_key = self.tool_context.state_key
if state_key:
return {state_key: snapshot}
return snapshot
@state.setter
def state(self, value: dict[str, Any] | None) -> None:
"""Set state from AG-UI protocol."""
if value is None:
return
data = self._extract_state_data(value)
self.tool_context.restore_state_snapshot(data)
def _extract_state_data(self, value: dict[str, Any]) -> dict[str, Any]:
"""Extract flat state dict, unwrapping state_key if present."""
state_key = self.tool_context.state_key
if state_key and state_key in value:
nested = value[state_key]
if isinstance(nested, dict):
return nested
return value

10
haiku_rag_slim/haiku/rag/tools/document.py
View file

@ -4,7 +4,6 @@ from pydantic_ai import Agent, FunctionToolset, RunContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import get_session_filter
from haiku.rag.utils import get_model
DOCUMENT_SUMMARY_PROMPT = """Generate a summary of the document content provided below.
@ -24,6 +23,7 @@ Document content:
class DocumentInfo(BaseModel):
"""Document info for list_documents response."""
id: str | None = None
title: str
uri: str
created: str
@ -91,22 +91,20 @@ def create_document_toolset(
Paginated list of documents with metadata.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
page_size = 50
offset = (page - 1) * page_size
effective_filter = get_session_filter(tool_context, base_filter)
docs = await client.list_documents(
limit=page_size, offset=offset, filter=effective_filter
limit=page_size, offset=offset, filter=base_filter
)
total = await client.count_documents(filter=effective_filter)
total = await client.count_documents(filter=base_filter)
total_pages = (total + page_size - 1) // page_size if total > 0 else 1
return DocumentListResponse(
documents=[
DocumentInfo(
id=doc.id,
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),

35
haiku_rag_slim/haiku/rag/tools/filters.py
View file

@ -1,9 +1,3 @@
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from haiku.rag.tools.context import ToolContext
def build_document_filter(document_name: str) -> str:
"""Build SQL filter for document name matching.
@ -31,35 +25,6 @@ def build_multi_document_filter(document_names: list[str]) -> str | None:
return " OR ".join(f"({f})" for f in filters)
def get_session_filter(
context: "ToolContext | None",
base_filter: str | None = None,
) -> str | None:
"""Build effective filter from session state document filter and base filter.
Checks the ToolContext for a registered SessionState. If it has a
document_filter, builds a SQL filter from it and combines with base_filter.
Args:
context: Optional ToolContext that may contain a SessionState.
base_filter: Optional base SQL WHERE clause to combine with.
Returns:
Combined filter string, or None if no filters apply.
"""
if context is None:
return base_filter
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
session_state = context.get(SESSION_NAMESPACE, SessionState)
if session_state is None or not session_state.document_filter:
return base_filter
session_filter = build_multi_document_filter(session_state.document_filter)
return combine_filters(base_filter, session_filter)
def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
"""Combine two SQL filters with AND logic.

37
haiku_rag_slim/haiku/rag/tools/models.py
View file

@ -1,37 +0,0 @@
from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
class QAResult(BaseModel):
"""Result from the QA toolset."""
question: str = Field(description="The question that was answered")
answer: str = Field(description="The answer to the question")
confidence: float = Field(
default=1.0,
description="Confidence score for this answer (0-1)",
ge=0.0,
le=1.0,
)
citations: list[Citation] = Field(
default_factory=list,
description="Citations supporting the answer",
)
@property
def sources(self) -> list[str]:
"""Source names for display."""
return list(
dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
)
class AnalysisResult(BaseModel):
"""Result from the analysis toolset (RLM execution)."""
answer: str = Field(description="The answer produced by analysis")
code_executed: bool = Field(
default=True,
description="Whether code was executed to produce this answer",
)

71
haiku_rag_slim/haiku/rag/tools/prompts.py
View file

@ -1,71 +0,0 @@
_TOOL_HEADER = """
How to decide which tool to use:"""
_TOOL_DOCUMENTS = """
- "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")."""
_TOOL_QA = """
- "ask" - Use for questions about topics in the knowledge base. Searches across documents and returns answers with citations. Prior answers are recalled to avoid redundant work."""
_TOOL_SEARCH = """
- "search" - Use when the user explicitly asks to search, find, or explore documents. Handles multi-query expansion internally and returns matching passages with surrounding context."""
_TOOL_ANALYSIS = """
- "analyze" - Use when the user asks for computation, data analysis, or quantitative tasks that require code execution (e.g., "calculate the average", "compare the numbers", "plot the data"). Runs Python code in a sandbox to produce results."""
_DOCUMENT_NAME_HEADER = """
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`"""
_DOCUMENT_NAME_SEARCH_EXAMPLES = """
- 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" """
_DOCUMENT_NAME_QA_EXAMPLES = """
- 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" """
_FEATURE_TOOLS: dict[str, str] = {
"documents": _TOOL_DOCUMENTS,
"qa": _TOOL_QA,
"search": _TOOL_SEARCH,
"analysis": _TOOL_ANALYSIS,
}
def build_tools_prompt(features: list[str]) -> str:
"""Build tool guidance for the given features.
Returns prompt text describing when and how to use each tool.
Designed to be spliced into a custom agent's system prompt.
Args:
features: List of feature names (e.g., ["search", "documents", "qa"]).
Returns:
Tool guidance prompt text.
"""
parts: list[str] = []
tool_sections = [_FEATURE_TOOLS[f] for f in features if f in _FEATURE_TOOLS]
if tool_sections:
parts.append(_TOOL_HEADER)
parts.extend(tool_sections)
if "search" in features or "qa" in features:
parts.append(_DOCUMENT_NAME_HEADER)
if "search" in features:
parts.append(_DOCUMENT_NAME_SEARCH_EXAMPLES)
if "qa" in features:
parts.append(_DOCUMENT_NAME_QA_EXAMPLES)
return "".join(parts)

256
haiku_rag_slim/haiku/rag/tools/qa.py
View file

@ -1,50 +1,17 @@
import math
from collections.abc import Callable
from pydantic import BaseModel, Field
from pydantic_ai import FunctionToolset, RunContext, ToolReturn
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import build_research_graph
from haiku.rag.agents.research.models import Citation, SearchAnswer
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.tools.context import RAGDeps, ToolContext
from haiku.rag.tools.filters import (
build_document_filter,
combine_filters,
get_session_filter,
)
from haiku.rag.tools.models import QAResult
from haiku.rag.tools.session import (
SESSION_NAMESPACE,
SessionContext,
SessionState,
compute_combined_state_delta,
)
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = sum(a * b for a, b in zip(vec1, vec2))
norm1 = math.sqrt(sum(a * a for a in vec1))
norm2 = math.sqrt(sum(b * b for b in vec2))
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
class QAHistoryEntry(BaseModel):
"""A Q&A pair with optional cached embedding for similarity matching."""
question: str
answer: str
confidence: float = 0.9
citations: list[Citation] = []
citations: list[Citation] = Field(default_factory=list)
question_embedding: list[float] | None = Field(default=None, exclude=True)
@property
@ -63,224 +30,3 @@ class QAHistoryEntry(BaseModel):
cited_chunks=[c.chunk_id for c in self.citations],
citations=self.citations,
)
class QASessionState(BaseModel):
"""Extended session state for QA with embedding cache."""
qa_history: list[QAHistoryEntry] = []
session_context: SessionContext | None = None
QA_SESSION_NAMESPACE = "haiku.rag.qa_session"
MAX_QA_HISTORY = 50
async def run_qa_core(
client: HaikuRAG,
config: AppConfig,
question: str,
document_name: str | None = None,
*,
context: ToolContext | None = None,
base_filter: str | None = None,
session_context: str | None = None,
prior_answers: list[SearchAnswer] | None = None,
on_qa_complete: Callable[[QASessionState, AppConfig], None] | None = None,
) -> QAResult:
"""Run the QA flow and return a QAResult.
This is the core QA implementation shared by toolsets and client APIs.
It updates session state and QA history when context is provided.
"""
session_state: SessionState | None = None
qa_session_state: QASessionState | None = None
if context is not None:
session_state = context.get(SESSION_NAMESPACE, SessionState)
qa_session_state = context.get(QA_SESSION_NAMESPACE, QASessionState)
doc_filter = build_document_filter(document_name) if document_name else None
effective_filter = combine_filters(
get_session_filter(context, base_filter), doc_filter
)
effective_session_context = session_context
if qa_session_state is not None and qa_session_state.session_context is not None:
effective_session_context = qa_session_state.session_context.summary
effective_prior_answers = prior_answers or []
if qa_session_state is not None and qa_session_state.qa_history:
embedder = get_embedder(config)
question_embedding = await embedder.embed_query(question)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(qa_session_state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
qa_session_state.qa_history[idx].question_embedding = new_embeddings[i]
matched_answers = []
for qa in qa_session_state.qa_history:
if qa.question_embedding is not None:
similarity = _cosine_similarity(
question_embedding, qa.question_embedding
)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matched_answers.append(qa.to_search_answer())
if matched_answers:
effective_prior_answers = matched_answers
graph = build_research_graph(config=config, output_mode="conversational")
research_context = ResearchContext(
original_question=question,
session_context=effective_session_context,
qa_responses=effective_prior_answers,
)
research_state = ResearchState(
context=research_context,
max_iterations=1,
search_filter=effective_filter,
max_concurrency=config.research.max_concurrency,
)
deps = ResearchDeps(client=client)
result = await graph.run(state=research_state, deps=deps)
# Build citations with stable indices from session state
citations = []
for i, c in enumerate(result.citations):
if session_state is not None:
index = session_state.get_or_assign_index(c.chunk_id)
else:
index = i + 1
citations.append(
Citation(
index=index,
document_id=c.document_id,
chunk_id=c.chunk_id,
document_uri=c.document_uri,
document_title=c.document_title,
page_numbers=c.page_numbers,
headings=c.headings,
content=c.content,
)
)
qa_result = QAResult(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citations,
)
if session_state is not None:
session_state.citations = citations
session_state.citations_history.append(citations)
if qa_session_state is not None:
qa_session_state.qa_history.append(
QAHistoryEntry(
question=question,
answer=result.answer,
confidence=result.confidence,
citations=citations,
)
)
# Enforce FIFO limit
if len(qa_session_state.qa_history) > MAX_QA_HISTORY:
qa_session_state.qa_history = qa_session_state.qa_history[-MAX_QA_HISTORY:]
if on_qa_complete is not None:
on_qa_complete(qa_session_state, config)
return qa_result
def create_qa_toolset(
config: AppConfig,
base_filter: str | None = None,
tool_name: str = "ask",
on_ask_complete: Callable[[QASessionState, AppConfig], None] | None = None,
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with Q&A capabilities using research graph.
Args:
config: Application configuration.
base_filter: Optional base SQL WHERE clause applied to searches.
tool_name: Name for the ask tool. Defaults to "ask".
on_ask_complete: Optional callback invoked after each QA cycle with
the updated QASessionState and config. Use this to trigger
background summarization or other post-processing.
Returns:
FunctionToolset with an ask tool.
"""
async def ask(
ctx: RunContext[RAGDeps],
question: str,
document_name: str | None = None,
) -> ToolReturn | QAResult:
"""Answer a question using the knowledge base.
Uses a research graph for searching and synthesizing answers.
Args:
question: The question to answer.
document_name: Optional document name/title to search within.
Returns:
QAResult with answer, confidence, and citations.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
state_key: str | None = None
client_snapshot: dict | None = None
if tool_context is not None:
state_key = tool_context.state_key
if tool_context.namespaces:
client_snapshot = (
tool_context.client_snapshot or tool_context.build_state_snapshot()
)
qa_result = await run_qa_core(
client=client,
config=config,
question=question,
document_name=document_name,
context=tool_context,
base_filter=base_filter,
on_qa_complete=on_ask_complete,
)
if client_snapshot is not None and tool_context is not None:
new_snapshot = tool_context.build_state_snapshot()
state_event = compute_combined_state_delta(
client_snapshot, new_snapshot, state_key=state_key
)
if state_event is not None:
answer_text = qa_result.answer
if qa_result.citations:
citation_refs = " ".join(
f"[{c.index}]" for c in qa_result.citations
)
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
return ToolReturn(return_value=answer_text, metadata=[state_event])
return qa_result
toolset: FunctionToolset[RAGDeps] = FunctionToolset()
toolset.add_function(ask, name=tool_name)
return toolset

111
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -1,23 +1,11 @@
from pydantic import BaseModel
from pydantic_ai import FunctionToolset, RunContext, ToolReturn
from collections.abc import Callable
from pydantic_ai import FunctionToolset, RunContext
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import RAGDeps
from haiku.rag.tools.filters import combine_filters, get_session_filter
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState, compute_state_delta
SEARCH_NAMESPACE = "haiku.rag.search"
class SearchState(BaseModel):
"""State for search toolset.
Accumulates search results across tool invocations.
"""
results: list[SearchResult] = []
from haiku.rag.tools.filters import combine_filters
def create_search_toolset(
@ -25,6 +13,7 @@ def create_search_toolset(
expand_context: bool = True,
base_filter: str | None = None,
tool_name: str = "search",
on_results: Callable[[list[SearchResult]], None] | None = None,
) -> FunctionToolset[RAGDeps]:
"""Create a toolset with search capabilities.
@ -35,6 +24,8 @@ def create_search_toolset(
base_filter: Optional base SQL WHERE clause applied to all searches.
Combined with any filter passed to the search tool.
tool_name: Name for the search tool. Defaults to "search".
on_results: Optional callback invoked with search results after each search.
Useful for accumulating results externally (e.g., for citation resolution).
Returns:
FunctionToolset with a search tool.
@ -45,7 +36,7 @@ def create_search_toolset(
query: str,
limit: int | None = None,
filter: str | None = None,
) -> ToolReturn | str:
) -> str:
"""Search the knowledge base for relevant documents.
Args:
@ -57,26 +48,8 @@ def create_search_toolset(
Formatted search results with content and metadata.
"""
client = ctx.deps.client
tool_context = ctx.deps.tool_context
search_state: SearchState | None = None
if tool_context is not None:
search_state = tool_context.get_or_create(SEARCH_NAMESPACE, SearchState)
session_state: SessionState | None = None
old_session_state: SessionState | None = None
state_key: str | None = None
if tool_context is not None:
session_state = tool_context.get(SESSION_NAMESPACE, SessionState)
state_key = tool_context.state_key
if session_state is not None:
old_session_state = session_state.model_copy(deep=True)
# Combine all filters: base_filter AND session_filter AND tool filter
effective_filter = combine_filters(
get_session_filter(tool_context, base_filter), filter
)
effective_filter = combine_filters(base_filter, filter)
effective_limit = limit or config.search.limit
results = await client.search(
query, limit=effective_limit, filter=effective_filter
@ -85,68 +58,18 @@ def create_search_toolset(
if expand_context:
results = await client.expand_context(results)
if search_state is not None:
search_state.results.extend(results)
results_list = list(results)
if not results:
if on_results:
on_results(results_list)
if not results_list:
return "No results found."
if session_state is not None:
citations = []
for r in results:
chunk_id = r.chunk_id or ""
if chunk_id:
index = session_state.get_or_assign_index(chunk_id)
else: # pragma: no cover
index = len(session_state.citation_registry) + 1
citations.append(
Citation(
index=index,
document_id=r.document_id or "",
chunk_id=chunk_id,
document_uri=r.document_uri or "",
document_title=r.document_title,
page_numbers=r.page_numbers or [],
headings=r.headings,
content=r.content,
)
)
session_state.citations = citations
session_state.citations_history.append(citations)
result_lines = []
for c in citations:
title = c.document_title or c.document_uri or "Unknown"
snippet = c.content[:300].replace("\n", " ").strip()
if len(c.content) > 300:
snippet += "..."
line = f"[{c.index}] **{title}**"
if c.page_numbers: # pragma: no cover
line += f" (pages {', '.join(map(str, c.page_numbers))})"
line += f"\n {snippet}"
result_lines.append(line)
formatted = f"Found {len(results)} results:\n\n" + "\n\n".join(result_lines)
if old_session_state is not None:
state_event = compute_state_delta(
old_session_state,
session_state,
state_key=state_key,
)
if state_event is not None:
return ToolReturn(
return_value=formatted,
metadata=[state_event],
)
return formatted # pragma: no cover
# Format results without citation indexing (standalone use)
total = len(results)
total = len(results_list)
formatted = [
r.format_for_agent(rank=i + 1, total=total) for i, r in enumerate(results)
r.format_for_agent(rank=i + 1, total=total)
for i, r in enumerate(results_list)
]
return "\n\n".join(formatted)

94
haiku_rag_slim/haiku/rag/tools/session.py
View file

@ -1,94 +0,0 @@
from datetime import datetime
from typing import Any
import jsonpatch
from ag_ui.core import EventType, StateDeltaEvent
from pydantic import BaseModel
from haiku.rag.agents.research.models import Citation
SESSION_NAMESPACE = "haiku.rag.session"
class SessionContext(BaseModel):
"""Compressed summary of conversation history for research graph."""
summary: str = ""
last_updated: datetime | None = None
class SessionState(BaseModel):
"""Session-level state for AG-UI integration.
This state is shared across toolsets and enables:
- Dynamic document filtering
- Stable citation indices across tool calls
- AG-UI state synchronization
"""
document_filter: list[str] = []
citation_registry: dict[str, int] = {}
citations: list[Citation] = []
citations_history: list[list[Citation]] = []
def get_or_assign_index(self, chunk_id: str) -> int:
"""Get or assign a stable citation index for a chunk_id.
Citation indices persist across tool calls within a session.
The first chunk gets index 1, subsequent new chunks get incrementing indices.
Same chunk_id always returns the same index.
"""
if chunk_id in self.citation_registry:
return self.citation_registry[chunk_id]
new_index = len(self.citation_registry) + 1
self.citation_registry[chunk_id] = new_index
return new_index
def compute_state_delta(
old_state: SessionState,
new_state: SessionState,
state_key: str | None = None,
) -> StateDeltaEvent | None:
"""Compute state delta between old and new session state.
Returns a StateDeltaEvent if there are changes, None otherwise.
"""
return compute_combined_state_delta(
old_state.model_dump(mode="json"),
new_state.model_dump(mode="json"),
state_key=state_key,
)
def compute_combined_state_delta(
old_snapshot: dict[str, Any],
new_snapshot: dict[str, Any],
state_key: str | None = None,
) -> StateDeltaEvent | None:
"""Compute state delta between old and new combined state snapshots.
This function computes delta for the combined chat state that includes
both SessionState and QASessionState fields.
Args:
old_snapshot: Previous state dict (e.g., from ChatDeps.state format).
new_snapshot: New state dict.
state_key: Optional namespace key for the state (e.g., "haiku.rag.chat").
Returns:
StateDeltaEvent if there are changes, None otherwise.
"""
wrapped_old = {state_key: old_snapshot} if state_key else old_snapshot
wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot
patch = jsonpatch.make_patch(wrapped_old, wrapped_new)
if not patch.patch:
return None
return StateDeltaEvent(
type=EventType.STATE_DELTA,
delta=patch.patch,
)

108
haiku_rag_slim/haiku/rag/tools/toolkit.py
View file

@ -1,108 +0,0 @@
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
from pydantic_ai import FunctionToolset
from haiku.rag.config.models import AppConfig
from haiku.rag.tools.context import ToolContext, prepare_context
from haiku.rag.tools.prompts import build_tools_prompt
FEATURE_SEARCH = "search"
FEATURE_DOCUMENTS = "documents"
FEATURE_QA = "qa"
FEATURE_ANALYSIS = "analysis"
@dataclass(frozen=True)
class Toolkit:
"""Bundled toolsets, prompt, and context factory for haiku.rag agents.
Created via build_toolkit(). Provides everything needed to compose
an agent with haiku.rag toolsets and create matching ToolContexts.
"""
toolsets: list[FunctionToolset[Any]] = field(default_factory=list)
prompt: str = ""
features: list[str] = field(default_factory=list)
def create_context(self, state_key: str | None = None) -> ToolContext:
"""Create a ToolContext with namespaces matching this toolkit's features.
Args:
state_key: Optional AG-UI state key to set on the context.
Returns:
A prepared ToolContext.
"""
context = ToolContext()
prepare_context(context, features=self.features, state_key=state_key)
return context
def prepare(self, context: ToolContext, state_key: str | None = None) -> None:
"""Register namespaces on an existing ToolContext for this toolkit's features.
Idempotent — safe to call multiple times on the same context.
Args:
context: ToolContext to prepare.
state_key: Optional AG-UI state key to set on the context.
"""
prepare_context(context, features=self.features, state_key=state_key)
def build_toolkit(
config: AppConfig,
features: list[str] | None = None,
base_filter: str | None = None,
expand_context: bool = True,
on_qa_complete: Callable | None = None,
) -> Toolkit:
"""Build a Toolkit with toolsets, prompt, and context factory for the given features.
Args:
config: Application configuration.
features: List of features to enable. Defaults to ["search", "documents"].
base_filter: Optional base SQL WHERE clause applied to all toolset factories.
expand_context: Whether to expand search results with surrounding context.
on_qa_complete: Optional callback invoked after each QA cycle.
Returns:
A Toolkit ready for agent composition.
"""
if features is None:
features = [FEATURE_SEARCH, FEATURE_DOCUMENTS]
toolsets: list[FunctionToolset[Any]] = []
if FEATURE_SEARCH in features:
from haiku.rag.tools.search import create_search_toolset
toolsets.append(
create_search_toolset(
config, expand_context=expand_context, base_filter=base_filter
)
)
if FEATURE_DOCUMENTS in features:
from haiku.rag.tools.document import create_document_toolset
toolsets.append(create_document_toolset(config, base_filter=base_filter))
if FEATURE_QA in features:
from haiku.rag.tools.qa import create_qa_toolset
toolsets.append(
create_qa_toolset(
config, base_filter=base_filter, on_ask_complete=on_qa_complete
)
)
if FEATURE_ANALYSIS in features:
from haiku.rag.tools.analysis import create_analysis_toolset
toolsets.append(create_analysis_toolset(config, base_filter=base_filter))
prompt = build_tools_prompt(features)
return Toolkit(toolsets=toolsets, prompt=prompt, features=features)

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

@ -1,3 +1,4 @@
import math
import sys
from datetime import UTC, datetime
from importlib import metadata
@ -14,6 +15,16 @@ if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig, ModelConfig
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = sum(a * b for a, b in zip(vec1, vec2))
norm1 = math.sqrt(sum(a * a for a in vec1))
norm2 = math.sqrt(sum(b * b for b in vec2))
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
def parse_datetime(s: str) -> datetime:
"""Parse a datetime string into a datetime object.

5
haiku_rag_slim/pyproject.toml
View file

@ -24,6 +24,7 @@ classifiers = [
dependencies = [
"cachetools>=5.5.0",
"docling-core==2.65.1",
"haiku.skills>=0.4.2",
"httpx>=0.28.1",
"jsonpatch>=1.33",
"lancedb==0.29.2",
@ -57,6 +58,10 @@ mistral = ["pydantic-ai-slim[mistral]"]
bedrock = ["pydantic-ai-slim[bedrock]"]
vertexai = ["pydantic-ai-slim[vertexai]"]
[project.entry-points."haiku.skills"]
rag = "haiku.rag.skills.rag:create_skill"
rag-rlm = "haiku.rag.skills.rlm:create_skill"
[project.scripts]
haiku-rag = "haiku.rag.cli:cli"

4
mkdocs.yml
View file

@ -74,6 +74,10 @@ nav:
- Agents:
- agents/index.md
- RLM Agent: agents/rlm.md
- Skills:
- skills/index.md
- RAG: skills/rag.md
- RLM: skills/rlm.md
- Toolsets: tools.md
- Applications: apps.md
- Server: server.md

4
pyproject.toml
View file

@ -131,6 +131,10 @@ filterwarnings = ["error", "ignore::UserWarning", "ignore::DeprecationWarning"]
[tool.coverage.run]
source = ["haiku_rag_slim"]
omit = [
"haiku_rag_slim/haiku/rag/chat/*",
"haiku_rag_slim/haiku/rag/inspector/*",
]
[tool.coverage.report]
show_missing = true

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

File diff suppressed because it is too large Load diff

336
tests/agents/chat/test_context.py
View file

@ -1,336 +0,0 @@
from datetime import datetime
from pathlib import Path
import pytest
from haiku.rag.agents.research.models import Citation
from haiku.rag.config import Config
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.rag.tools.session import SessionContext
@pytest.fixture(scope="module")
def vcr_cassette_dir():
return str(Path(__file__).parent.parent.parent / "cassettes" / "test_chat_context")
class TestSessionContext:
"""Tests for SessionContext model."""
def test_session_context_creation_empty(self):
"""Test SessionContext can be created with defaults."""
ctx = SessionContext()
assert ctx.summary == ""
assert ctx.last_updated is None
def test_session_context_creation_with_values(self):
"""Test SessionContext can be created with provided values."""
now = datetime.now()
ctx = SessionContext(
summary="User discussed authentication patterns.",
last_updated=now,
)
assert ctx.summary == "User discussed authentication patterns."
assert ctx.last_updated == now
def test_session_context_serialization_roundtrip(self):
"""Test SessionContext serializes and deserializes correctly."""
now = datetime.now()
original = SessionContext(
summary="Test summary with facts.",
last_updated=now,
)
# Serialize to dict
data = original.model_dump()
# Deserialize back
restored = SessionContext(**data)
assert restored.summary == original.summary
assert restored.last_updated == original.last_updated
class TestSummarizeSession:
"""Tests for summarize_session function."""
@pytest.mark.asyncio
async def test_summarize_session_empty_history(self):
"""Test summarize_session with empty qa_history returns empty string."""
from haiku.rag.agents.chat.context import summarize_session
result = await summarize_session(qa_history=[], config=Config)
assert result == ""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_session_single_entry(
self, allow_model_requests, temp_db_path
):
"""Test summarize_session with a single qa entry."""
from haiku.rag.agents.chat.context import summarize_session
qa_history = [
QAHistoryEntry(
question="What is the authentication method?",
answer="The API uses JWT tokens for authentication.",
confidence=0.95,
citations=[
Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-1",
document_uri="auth-guide.md",
document_title="Auth Guide",
content="JWT token details...",
)
],
)
]
result = await summarize_session(qa_history=qa_history, config=Config)
# Should produce a non-empty summary
assert len(result) > 0
# Summary should mention authentication or JWT
assert "authentication" in result.lower() or "jwt" in result.lower()
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_session_multiple_entries(
self, allow_model_requests, temp_db_path
):
"""Test summarize_session with multiple qa entries produces consolidated summary."""
from haiku.rag.agents.chat.context import summarize_session
qa_history = [
QAHistoryEntry(
question="What is the authentication method?",
answer="The API uses JWT tokens for authentication.",
confidence=0.95,
citations=[
Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-1",
document_uri="auth-guide.md",
document_title="Auth Guide",
content="JWT token details...",
)
],
),
QAHistoryEntry(
question="What is the rate limit?",
answer="Rate limiting is set to 100 requests per minute.",
confidence=0.9,
citations=[
Citation(
index=1,
document_id="doc-2",
chunk_id="chunk-2",
document_uri="api-reference.md",
document_title="API Reference",
content="Rate limit config...",
)
],
),
QAHistoryEntry(
question="How do I refresh tokens?",
answer="Use the /refresh endpoint with your refresh token.",
confidence=0.85,
citations=[
Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-3",
document_uri="auth-guide.md",
document_title="Auth Guide",
content="Token refresh...",
)
],
),
]
result = await summarize_session(qa_history=qa_history, config=Config)
# Should produce a non-empty summary
assert len(result) > 0
# Summary should contain structured sections
result_lower = result.lower()
assert "key facts" in result_lower or "established" in result_lower
assert "documents" in result_lower or "sources" in result_lower
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_summarize_session_with_current_context(self, allow_model_requests):
"""Test summarize_session incorporates current_context into the summary."""
from haiku.rag.agents.chat.context import summarize_session
qa_history = [
QAHistoryEntry(
question="What's the rate limit?",
answer="100 requests per minute.",
confidence=0.9,
)
]
# Provide current_context (e.g., previous summary)
current_context = "Focus on Python APIs. User is building a web application."
result = await summarize_session(
qa_history=qa_history,
config=Config,
current_context=current_context,
)
# Summary should be non-empty and ideally incorporate context about Python/web
assert len(result) > 0
# The context about "Python" or "web application" should influence the summary
result_lower = result.lower()
assert (
"rate" in result_lower or "limit" in result_lower or "100" in result_lower
)
class TestUpdateSessionContext:
"""Tests for update_session_context function."""
@pytest.mark.asyncio
@pytest.mark.vcr()
async def test_update_session_context_returns_context(
self, allow_model_requests, temp_db_path
):
"""Test update_session_context returns a populated SessionContext."""
from haiku.rag.agents.chat.context import update_session_context
qa_history = [
QAHistoryEntry(
question="What is the authentication method?",
answer="The API uses JWT tokens.",
confidence=0.95,
)
]
result = await update_session_context(
qa_history=qa_history,
config=Config,
)
assert result.summary != ""
assert result.last_updated is not None
@pytest.mark.asyncio
async def test_update_session_context_with_empty_history(self):
"""Test update_session_context with empty history returns empty summary."""
from haiku.rag.agents.chat.context import update_session_context
result = await update_session_context(
qa_history=[],
config=Config,
)
assert result.summary == ""
class TestTriggerBackgroundSummarization:
"""Tests for trigger_background_summarization."""
def test_trigger_with_empty_qa_history(self):
"""trigger_background_summarization returns early with empty qa_history."""
from haiku.rag.agents.chat.context import (
_summarization_tasks,
trigger_background_summarization,
)
from haiku.rag.tools.qa import QASessionState
tasks_before = len(_summarization_tasks)
qa_session_state = QASessionState()
assert len(qa_session_state.qa_history) == 0
trigger_background_summarization(qa_session_state, config=Config)
# No new task should have been created
assert len(_summarization_tasks) == tasks_before
@pytest.mark.asyncio
async def test_trigger_cancels_existing_task(self):
"""Second trigger cancels the previous background task."""
import asyncio
from unittest.mock import patch
from haiku.rag.agents.chat.context import (
_summarization_tasks,
trigger_background_summarization,
)
from haiku.rag.tools.qa import QAHistoryEntry, QASessionState
_summarization_tasks.clear()
qa_session_state = QASessionState(
qa_history=[QAHistoryEntry(question="Q1", answer="A1", confidence=0.9)]
)
# Patch _update_context_background to be a slow coroutine
async def slow_background(*args, **kwargs):
await asyncio.sleep(10)
with patch(
"haiku.rag.agents.chat.context._update_context_background",
new=slow_background,
):
# First trigger creates a task
trigger_background_summarization(qa_session_state, config=Config)
key = id(qa_session_state)
assert key in _summarization_tasks
first_task = _summarization_tasks[key]
# Second trigger should cancel the first
trigger_background_summarization(qa_session_state, config=Config)
await asyncio.sleep(0) # Let cancellation propagate
assert first_task.cancelled() or first_task.done()
# Cleanup
if key in _summarization_tasks:
_summarization_tasks[key].cancel()
try:
await _summarization_tasks[key]
except asyncio.CancelledError:
pass
_summarization_tasks.clear()
class TestUpdateSessionContextPassesCurrentContext:
"""Tests for update_session_context current_context forwarding."""
@pytest.mark.asyncio
async def test_update_session_context_passes_current_context(self):
"""Test update_session_context passes current_context to summarizer."""
from unittest.mock import patch
from haiku.rag.agents.chat.context import update_session_context
qa_history = [
QAHistoryEntry(
question="What is JWT?",
answer="JSON Web Token for authentication.",
confidence=0.95,
)
]
captured_current_context = []
async def mock_summarize(qa_history, config, current_context=None):
captured_current_context.append(current_context)
return "Mocked summary"
with patch(
"haiku.rag.agents.chat.context.summarize_session",
new=mock_summarize,
):
await update_session_context(
qa_history=qa_history,
config=Config,
current_context="Previous session summary",
)
assert len(captured_current_context) == 1
assert captured_current_context[0] == "Previous session summary"

189
tests/agents/chat/test_features.py
View file

@ -1,189 +0,0 @@
from pydantic_ai import FunctionToolset
from haiku.rag.agents.chat.agent import (
DEFAULT_FEATURES,
FEATURE_ANALYSIS,
FEATURE_DOCUMENTS,
FEATURE_QA,
FEATURE_SEARCH,
ChatDeps,
create_chat_agent,
prepare_chat_context,
)
from haiku.rag.agents.chat.prompts import build_chat_prompt
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.tools.context import ToolContext
from haiku.rag.tools.qa import QA_SESSION_NAMESPACE, QASessionState
from haiku.rag.tools.session import SESSION_NAMESPACE, SessionState
def _count_function_toolsets(agent) -> int:
"""Count FunctionToolset instances in an agent (excludes internal toolsets)."""
return sum(1 for t in agent.toolsets if type(t) is FunctionToolset)
# =============================================================================
# Feature Selection Tests
# =============================================================================
def test_default_features(temp_db_path):
"""Default features create search + document + qa toolsets and register both states."""
context = ToolContext()
prepare_chat_context(context)
agent = create_chat_agent(Config)
# Should have 3 toolsets (search, document, qa)
assert _count_function_toolsets(agent) == 3
# Both SessionState and QASessionState should be registered
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is not None
def test_search_only(temp_db_path):
"""features=["search"] creates only search toolset, no QASessionState."""
context = ToolContext()
prepare_chat_context(context, features=[FEATURE_SEARCH])
agent = create_chat_agent(Config, features=[FEATURE_SEARCH])
assert _count_function_toolsets(agent) == 1
# SessionState always registered, but QASessionState should NOT be
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
def test_search_and_documents(temp_db_path):
"""features=["search", "documents"] creates both toolsets, no QASessionState."""
context = ToolContext()
prepare_chat_context(context, features=[FEATURE_SEARCH, FEATURE_DOCUMENTS])
agent = create_chat_agent(Config, features=[FEATURE_SEARCH, FEATURE_DOCUMENTS])
assert _count_function_toolsets(agent) == 2
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
def test_all_features(temp_db_path):
"""All four features create four toolsets."""
context = ToolContext()
prepare_chat_context(
context,
features=[FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS],
)
agent = create_chat_agent(
Config,
features=[FEATURE_SEARCH, FEATURE_DOCUMENTS, FEATURE_QA, FEATURE_ANALYSIS],
)
assert _count_function_toolsets(agent) == 4
assert context.get(SESSION_NAMESPACE, SessionState) is not None
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is not None
def test_no_qa_skips_qa_session_state(temp_db_path):
"""Without QA feature, QASessionState is not registered."""
context = ToolContext()
prepare_chat_context(context, features=[FEATURE_SEARCH, FEATURE_DOCUMENTS])
assert context.get(QA_SESSION_NAMESPACE, QASessionState) is None
def test_chat_deps_state_without_qa(temp_db_path):
"""ChatDeps.state getter omits qa_history/session_context when QASessionState absent."""
from haiku.rag.agents.chat.state import AGUI_STATE_KEY
client = HaikuRAG(temp_db_path, create=True)
context = ToolContext()
prepare_chat_context(context, features=[FEATURE_SEARCH])
deps = ChatDeps(config=Config, client=client, tool_context=context)
state = deps.state
# State is wrapped under the AGUI state key
assert AGUI_STATE_KEY in state
inner = state[AGUI_STATE_KEY]
# SessionState fields should be present
assert "document_filter" in inner
assert "citation_registry" in inner
assert "citations" in inner
# QA fields should NOT be present
assert "qa_history" not in inner
assert "session_context" not in inner
client.close()
# =============================================================================
# Prompt Composition Tests
# =============================================================================
def test_build_chat_prompt_default():
"""Default features produce prompt mentioning all standard tools."""
prompt = build_chat_prompt(DEFAULT_FEATURES)
assert "list_documents" in prompt
assert "get_document" in prompt
assert "summarize_document" in prompt
assert "ask" in prompt
assert "search" in prompt
assert "analyze" not in prompt
def test_build_chat_prompt_search_only():
"""Search-only prompt doesn't mention ask or document tools."""
prompt = build_chat_prompt([FEATURE_SEARCH])
assert "search" in prompt
assert '"ask"' not in prompt
assert '"list_documents"' not in prompt
assert '"get_document"' not in prompt
assert '"summarize_document"' not in prompt
def test_build_chat_prompt_includes_analysis():
"""Analysis feature adds analyze guidance to prompt."""
prompt = build_chat_prompt(
[FEATURE_SEARCH, FEATURE_QA, FEATURE_DOCUMENTS, FEATURE_ANALYSIS]
)
assert "analyze" in prompt
assert "search" in prompt
assert "ask" in prompt
def test_build_chat_prompt_with_preamble():
"""Custom preamble replaces the default identity section."""
custom = "You are a custom assistant."
prompt = build_chat_prompt(DEFAULT_FEATURES, preamble=custom)
assert prompt.startswith(custom)
# Tool guidance should still be appended
assert "search" in prompt
assert "ask" in prompt
# Default identity should NOT be present
assert "haiku.rag" not in prompt
def test_build_chat_prompt_without_preamble_uses_default():
"""Without preamble, the default identity section is used."""
prompt = build_chat_prompt(DEFAULT_FEATURES)
assert "haiku.rag" in prompt
def test_create_chat_agent_with_preamble():
"""create_chat_agent passes preamble through to build_chat_prompt."""
custom = "You are a domain expert."
agent = create_chat_agent(Config, preamble=custom)
assert agent is not None
# _instructions is the internal list of instruction strings/callables
assert any(
custom in instr for instr in agent._instructions if isinstance(instr, str)
)

141
tests/agents/chat/test_state.py
View file

@ -1,141 +0,0 @@
from haiku.rag.agents.chat.state import ChatSessionState
from haiku.rag.agents.research.models import Citation
from haiku.rag.tools.session import SessionContext, SessionState
def test_max_qa_history_constant():
"""Test MAX_QA_HISTORY constant value."""
from haiku.rag.tools.qa import MAX_QA_HISTORY
assert MAX_QA_HISTORY == 50
def test_citation_registry_index_assignment():
"""Test get_or_assign_index basic index assignment behavior.
Verifies:
- First chunk gets index 1
- Second unique chunk gets index 2
- Same chunk_id always returns same index
"""
session_state = SessionState()
# First chunk gets index 1
index1 = session_state.get_or_assign_index("chunk-abc")
assert index1 == 1
# Second unique chunk gets index 2
index2 = session_state.get_or_assign_index("chunk-def")
assert index2 == 2
# Same chunk_id returns same index (not incremented)
index1_again = session_state.get_or_assign_index("chunk-abc")
assert index1_again == 1
def test_citation_registry_stability():
"""Test citation indices are stable across multiple calls in any order."""
session_state = SessionState()
# First round assigns indices 1, 2, 3
idx_a = session_state.get_or_assign_index("chunk-a")
idx_b = session_state.get_or_assign_index("chunk-b")
idx_c = session_state.get_or_assign_index("chunk-c")
# Second round - existing chunks keep their indices regardless of order
assert session_state.get_or_assign_index("chunk-b") == idx_b
assert session_state.get_or_assign_index("chunk-a") == idx_a
assert session_state.get_or_assign_index("chunk-c") == idx_c
# New chunk gets next index
idx_d = session_state.get_or_assign_index("chunk-d")
assert idx_d == 4
def test_citation_registry_serialization_roundtrip():
"""Test citation_registry serializes and deserializes correctly for AG-UI state."""
# Create state and assign indices
original = ChatSessionState()
original.citation_registry = {"chunk-a": 1, "chunk-b": 2}
# Serialize
state_dict = original.model_dump()
assert "citation_registry" in state_dict
assert state_dict["citation_registry"] == {"chunk-a": 1, "chunk-b": 2}
# Deserialize (simulating AG-UI state restoration)
restored = ChatSessionState.model_validate(state_dict)
assert restored.citation_registry == {"chunk-a": 1, "chunk-b": 2}
def test_chat_session_state_initial_context_default_none():
"""Initial context should default to None."""
state = ChatSessionState()
assert state.initial_context is None
def test_chat_session_state_initial_context_preserved():
"""Explicit initial_context should be preserved."""
state = ChatSessionState(initial_context="Background info about the project")
assert state.initial_context == "Background info about the project"
def test_chat_session_state_initial_context_serialization():
"""initial_context should serialize and deserialize correctly."""
state = ChatSessionState(
initial_context="User is working on authentication",
)
state_dict = state.model_dump()
assert state_dict["initial_context"] == "User is working on authentication"
restored = ChatSessionState.model_validate(state_dict)
assert restored.initial_context == "User is working on authentication"
def test_chat_session_state_model_dump_json_serializes_datetime():
"""model_dump(mode='json') should serialize datetime to ISO string.
Agent tools use model_dump(mode='json') when creating StateSnapshotEvent
to ensure datetime fields are JSON-serializable for external clients
persisting AG-UI state to database JSON columns.
"""
from datetime import datetime
session_state = ChatSessionState(
session_context=SessionContext(
summary="Test summary",
last_updated=datetime(2025, 1, 27, 12, 0, 0),
),
)
# This is how agent.py creates snapshots for StateSnapshotEvent
snapshot = session_state.model_dump(mode="json")
# datetime should be serialized as ISO string, not datetime object
assert isinstance(snapshot["session_context"]["last_updated"], str)
assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00"
def test_chat_session_state_citations_history_default():
"""citations_history defaults to empty list."""
state = ChatSessionState()
assert state.citations_history == []
def test_chat_session_state_citations_history_roundtrip():
"""citations_history serializes and deserializes correctly."""
citation = Citation(
index=1,
document_id="d1",
chunk_id="c1",
document_uri="test://doc",
document_title="Doc",
page_numbers=[],
headings=None,
content="content",
)
state = ChatSessionState(citations_history=[[citation]])
data = state.model_dump(mode="json")
restored = ChatSessionState.model_validate(data)
assert len(restored.citations_history) == 1
assert restored.citations_history[0][0].chunk_id == "c1"

2
tests/agents/research/test_models.py
View file

@ -25,7 +25,7 @@ class TestCitation:
assert citation.index is None
def test_citation_with_index(self):
"""Test Citation can be created with index (chat agent use case)."""
"""Test Citation can be created with index (chat use case)."""
citation = Citation(
index=1,
document_id="doc-1",

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

@ -68,64 +68,15 @@ def test_iterative_plan_result_model():
assert continue_result.next_question == "What are the specific requirements?"
# =============================================================================
# Conversational Graph Tests
# =============================================================================
def test_build_research_graph_conversational_mode_returns_graph():
"""Test build_research_graph with output_mode='conversational' returns a valid Graph instance."""
def test_build_research_graph_returns_graph():
"""Test build_research_graph returns a valid Graph instance."""
from pydantic_graph.beta import Graph
graph = build_research_graph(output_mode="conversational")
graph = build_research_graph()
assert graph is not None
assert isinstance(graph, Graph)
def test_build_research_graph_report_mode_returns_graph():
"""Test build_research_graph with output_mode='report' returns a valid Graph instance."""
from pydantic_graph.beta import Graph
graph = build_research_graph(output_mode="report")
assert graph is not None
assert isinstance(graph, Graph)
def test_conversational_answer_model():
"""Test ConversationalAnswer model can be created with all fields."""
from haiku.rag.agents.research.models import Citation, ConversationalAnswer
citation = Citation(
index=1,
document_id="doc-1",
chunk_id="chunk-1",
document_uri="test.md",
document_title="Test Doc",
content="Test content",
)
answer = ConversationalAnswer(
answer="The answer is 42.",
citations=[citation],
confidence=0.95,
)
assert answer.answer == "The answer is 42."
assert len(answer.citations) == 1
assert answer.confidence == 0.95
def test_conversational_answer_default_values():
"""Test ConversationalAnswer uses correct default values."""
from haiku.rag.agents.research.models import ConversationalAnswer
answer = ConversationalAnswer(answer="Just the answer.")
assert answer.answer == "Just the answer."
assert answer.citations == []
assert answer.confidence == 1.0
def test_format_context_for_prompt_basic():
"""Test format_context_for_prompt with basic context."""
from haiku.rag.agents.research.dependencies import ResearchContext
@ -138,22 +89,6 @@ def test_format_context_for_prompt_basic():
assert "What is X?" in result
def test_format_context_for_prompt_with_session_context():
"""Test format_context_for_prompt includes session_context as background."""
from haiku.rag.agents.research.dependencies import ResearchContext
from haiku.rag.agents.research.graph import format_context_for_prompt
context = ResearchContext(
original_question="What is Y?",
session_context="Previous discussion about topic Z.",
)
result = format_context_for_prompt(context)
assert "<background>" in result
assert "Previous discussion" in result
assert "What is Y?" in result
def test_format_context_for_prompt_with_prior_answers():
"""Test format_context_for_prompt includes prior_answers."""
from haiku.rag.agents.research.dependencies import ResearchContext

1311
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

1668
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

2945
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

187
tests/cassettes/test_chat_agent/test_chat_agent_get_document_not_found.yaml
View file

@ -1,187 +0,0 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '5211'
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: Get me 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:
page:
default: 1
description: 'Page number (default: 1, 50 documents per page)'
type: integer
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:
- '539'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: I’m sorry, but that document isn’t available in the knowledge base. If there’s another topic or document
you’d like help with, just let me know!
reasoning: User asking for nonexistent document. Need to respond that none exists. No tool usage.
role: assistant
created: 1769793913
id: chatcmpl-124
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 60
prompt_tokens: 1025
total_tokens: 1085
status:
code: 200
message: OK
version: 1

475
tests/cassettes/test_chat_agent/test_chat_agent_get_document_tool.yaml
View file

File diff suppressed because one or more lines are too long

3051
tests/cassettes/test_chat_agent/test_chat_agent_multi_turn_with_context.yaml
View file

File diff suppressed because one or more lines are too long

883
tests/cassettes/test_chat_agent/test_chat_agent_search_tool.yaml
View file

File diff suppressed because one or more lines are too long

1044
tests/cassettes/test_chat_agent/test_chat_agent_search_tool_with_filter.yaml
View file

File diff suppressed because one or more lines are too long

880
tests/cassettes/test_chat_agent/test_chat_agent_search_with_session_filter.yaml
View file

File diff suppressed because one or more lines are too long

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

File diff suppressed because one or more lines are too long

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

File diff suppressed because one or more lines are too long

524
tests/cassettes/test_chat_agent/test_list_documents_pagination.yaml
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
View file

File diff suppressed because one or more lines are too long

488
tests/cassettes/test_chat_agent/test_run_chat_agent.yaml
View file

File diff suppressed because one or more lines are too long

563
tests/cassettes/test_chat_agent/test_summarize_document_found.yaml
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
View file

@ -1,196 +0,0 @@
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

94
tests/cassettes/test_chat_context/TestSummarizeSession.test_summarize_session_multiple_entries.yaml
View file

@ -1,94 +0,0 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '1493'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
Your summary should be concise (aim for 500-1500 tokens) and include:
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
3. **Current Focus** - What topic or question thread the user is currently exploring
Rules:
- Extract only high-signal information that would help answer follow-up questions
- Omit small talk, greetings, or low-confidence answers
- Use bullet points for clarity
- Keep technical details but compress verbose explanations
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
role: system
- content: |
## Q1: What is the authentication method?
**Answer** (confidence: 95%):
The API uses JWT tokens for authentication.
**Sources:** Auth Guide
## Q2: What is the rate limit?
**Answer** (confidence: 90%):
Rate limiting is set to 100 requests per minute.
**Sources:** API Reference
## Q3: How do I refresh tokens?
**Answer** (confidence: 85%):
Use the /refresh endpoint with your refresh token.
**Sources:** Auth Guide
role: user
model: gpt-oss
reasoning_effort: low
stream: false
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '865'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: |-
## Key Facts Established
- **Authentication method**: JWT tokens.
- **Rate limit**: 100 requests per minute.
- **Token refresh**: use `/refresh` endpoint with a refresh token.
## Documents Referenced
- **Auth Guide** – contains details on JWT usage, token issuance, and refresh mechanism.
- **API Reference** – includes rate limiting policy and endpoint descriptions.
## Current Focus
The user is currently gathering foundational API usage details, specifically authentication methods, rate limits, and token refresh procedures.
reasoning: We need to summarize.
role: assistant
created: 1769007514
id: chatcmpl-760
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 121
prompt_tokens: 365
total_tokens: 486
status:
code: 200
message: OK
version: 1

81
tests/cassettes/test_chat_context/TestSummarizeSession.test_summarize_session_single_entry.yaml
View file

@ -1,81 +0,0 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '1207'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
Your summary should be concise (aim for 500-1500 tokens) and include:
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
3. **Current Focus** - What topic or question thread the user is currently exploring
Rules:
- Extract only high-signal information that would help answer follow-up questions
- Omit small talk, greetings, or low-confidence answers
- Use bullet points for clarity
- Keep technical details but compress verbose explanations
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
role: system
- content: |
## Q1: What is the authentication method?
**Answer** (confidence: 95%):
The API uses JWT tokens for authentication.
**Sources:** Auth Guide
role: user
model: gpt-oss
reasoning_effort: low
stream: false
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '575'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: |-
**Key Facts Established**
- The API uses **JWT tokens** for authentication.
**Documents Referenced**
- **Auth Guide** – Provides details on JWT usage for this API.
**Current Focus**
- Understanding the authentication method employed by the API.
reasoning: We need to summarize.
role: assistant
created: 1769007512
id: chatcmpl-739
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 65
prompt_tokens: 292
total_tokens: 357
status:
code: 200
message: OK
version: 1

89
tests/cassettes/test_chat_context/TestSummarizeSession.test_summarize_session_with_current_context.yaml
View file

@ -1,89 +0,0 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '1581'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a session summarizer. Given a conversation history of Q&A pairs (and optionally existing context), produce a structured summary that captures key information for future context.
If a "Current Context" section is provided at the start of the input, incorporate that context into your summary. This might be initial background context from the user or a previous summary - build upon it rather than discard it.
Your summary should be concise (aim for 500-1500 tokens) and include:
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
3. **Current Focus** - What topic or question thread the user is currently exploring
Rules:
- Extract only high-signal information that would help answer follow-up questions
- When building on existing context, merge new information with prior context
- Omit small talk, greetings, or low-confidence answers
- Use bullet points for clarity
- Keep technical details but compress verbose explanations
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
role: system
- content: |
## Current Context
Focus on Python APIs. User is building a web application.
## Q1: What's the rate limit?
**Answer** (confidence: 90%):
100 requests per minute.
role: user
model: gpt-oss
reasoning_effort: low
stream: false
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '682'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: |-
## Summary
- **Key Facts Established**
- The user is building a web application and is focused on Python APIs.
- The relevant rate limit is **100 requests per minute** (confidence 90%).
- **Documents Referenced**
- None cited in this exchange.
- **Current Focus**
- Understanding and managing API rate limits for the Python-based web application.
reasoning: We need summary.
role: assistant
created: 1769164539
id: chatcmpl-369
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 91
prompt_tokens: 362
total_tokens: 453
status:
code: 200
message: OK
version: 1

80
tests/cassettes/test_chat_context/TestUpdateSessionContext.test_update_session_context_returns_context.yaml
View file

@ -1,80 +0,0 @@
interactions:
- request:
headers:
accept:
- application/json
accept-encoding:
- gzip, deflate, zstd
connection:
- keep-alive
content-length:
- '1163'
content-type:
- application/json
host:
- localhost:11434
method: POST
parsed_body:
messages:
- content: |-
You are a session summarizer. Given a conversation history of Q&A pairs, produce a structured summary that captures key information for future context.
Your summary should be concise (aim for 500-1500 tokens) and include:
1. **Key Facts Established** - Specific facts, data, or conclusions learned during the conversation
2. **Documents Referenced** - Documents or sources that were cited, with brief notes on what they contain
3. **Current Focus** - What topic or question thread the user is currently exploring
Rules:
- Extract only high-signal information that would help answer follow-up questions
- Omit small talk, greetings, or low-confidence answers
- Use bullet points for clarity
- Keep technical details but compress verbose explanations
- Preserve document names/titles when mentioned in sources
Output the summary directly in markdown format. Do not include meta-commentary about the summary itself.
role: system
- content: |
## Q1: What is the authentication method?
**Answer** (confidence: 95%):
The API uses JWT tokens.
role: user
model: gpt-oss
reasoning_effort: low
stream: false
uri: http://localhost:11434/v1/chat/completions
response:
headers:
content-length:
- '559'
content-type:
- application/json
parsed_body:
choices:
- finish_reason: stop
index: 0
message:
content: |-
## Key Facts Established
- The API authentication method is **JWT tokens** (high confidence 95%).
## Documents Referenced
- *None provided*.
## Current Focus
- The user is exploring details related to **API authentication mechanisms**.
reasoning: We need summary.
role: assistant
created: 1769007530
id: chatcmpl-975
model: gpt-oss
object: chat.completion
system_fingerprint: fp_ollama
usage:
completion_tokens: 64
prompt_tokens: 284
total_tokens: 348
status:
code: 200
message: OK
version: 1

840
tests/cassettes/test_qa_tools/TestAskTool.test_ask_with_tool_context_returns_tool_return.yaml
View file

File diff suppressed because one or more lines are too long

Some files were not shown because too many files have changed in this diff Show more