Merge pull request #240 from ggozad/chore/improve-agui
Improve conversational agent & AG-UI integration
This commit is contained in:
Yiorgis Gozadinos
GitHub
commit
9756add372
25 changed files with 1018 additions and 105 deletions
22
.github/workflows/test.yml
vendored
22
.github/workflows/test.yml
vendored
|
|
@ -25,8 +25,28 @@ jobs:
|
|||
- name: Type check
|
||||
run: uv run pyright
|
||||
|
||||
lint-frontend:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: pnpm/action-setup@v4
|
||||
with:
|
||||
version: 10
|
||||
- name: Set up Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: "22"
|
||||
cache: "pnpm"
|
||||
cache-dependency-path: app/frontend/pnpm-lock.yaml
|
||||
- name: Install dependencies
|
||||
working-directory: app/frontend
|
||||
run: pnpm install --frozen-lockfile
|
||||
- name: Lint and format check
|
||||
working-directory: app/frontend
|
||||
run: pnpm run check
|
||||
|
||||
test:
|
||||
needs: lint
|
||||
needs: [lint, lint-frontend]
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
|
|
|
|||
|
|
@ -20,3 +20,12 @@ repos:
|
|||
rev: v1.1.407
|
||||
hooks:
|
||||
- id: pyright
|
||||
|
||||
- repo: local
|
||||
hooks:
|
||||
- id: biome
|
||||
name: biome check
|
||||
entry: bash -c 'cd app/frontend && npm run check'
|
||||
language: system
|
||||
files: ^app/frontend/
|
||||
types_or: [javascript, jsx, ts, tsx, json]
|
||||
|
|
|
|||
13
CHANGELOG.md
13
CHANGELOG.md
|
|
@ -1,6 +1,19 @@
|
|||
# Changelog
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- **Background Context Support**: Pass background context to agents via CLI or Python API
|
||||
- `haiku-rag ask --context "..." --context-file path` for Q&A with background context
|
||||
- `haiku-rag research --context "..." --context-file path` for research with background context
|
||||
- `haiku-rag chat --context "..." --context-file path` for chat sessions with persistent context
|
||||
- `ResearchContext(background_context="...")` for Python API usage
|
||||
- `ChatSessionState(background_context="...")` for chat agent sessions
|
||||
- Context is included in agent system prompts and research graph planning
|
||||
- **Frontend Background Context**: Settings panel in the chat app to configure persistent background context
|
||||
- Context is stored in localStorage and sent with each conversation
|
||||
- **Frontend Linting**: Added Biome for linting and formatting the frontend codebase
|
||||
|
||||
## [0.26.4] - 2026-01-15
|
||||
|
||||
### Added
|
||||
|
|
|
|||
|
|
@ -80,8 +80,9 @@ async def stream_chat(request: Request) -> Response:
|
|||
accept = request.headers.get("accept", SSE_CONTENT_TYPE)
|
||||
run_input = AGUIAdapter.build_run_input(body)
|
||||
|
||||
# Restore qa_history from incoming state (look under namespaced key)
|
||||
# Restore session state from incoming AG-UI state (look under namespaced key)
|
||||
initial_qa_history: list[QAResponse] = []
|
||||
background_context: str | None = None
|
||||
state = getattr(run_input, "state", None)
|
||||
if state:
|
||||
chat_state = state.get(AGUI_STATE_KEY, state)
|
||||
|
|
@ -89,6 +90,7 @@ async def stream_chat(request: Request) -> Response:
|
|||
initial_qa_history = [
|
||||
QAResponse(**qa) for qa in chat_state.get("qa_history", [])
|
||||
]
|
||||
background_context = chat_state.get("background_context")
|
||||
|
||||
# Build deps with session state
|
||||
thread_id = getattr(run_input, "thread_id", None)
|
||||
|
|
@ -98,6 +100,7 @@ async def stream_chat(request: Request) -> Response:
|
|||
session_state=ChatSessionState(
|
||||
session_id=thread_id or "",
|
||||
qa_history=initial_qa_history,
|
||||
background_context=background_context,
|
||||
),
|
||||
state_key=AGUI_STATE_KEY,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"$schema": "https://biomejs.dev/schemas/2.2.6/schema.json",
|
||||
"vcs": {
|
||||
"enabled": false,
|
||||
"enabled": true,
|
||||
"clientKind": "git",
|
||||
"useIgnoreFile": false
|
||||
"useIgnoreFile": true
|
||||
},
|
||||
"files": {
|
||||
"ignoreUnknown": false
|
||||
|
|
@ -15,7 +15,11 @@
|
|||
"linter": {
|
||||
"enabled": true,
|
||||
"rules": {
|
||||
"recommended": true
|
||||
"recommended": true,
|
||||
"a11y": {
|
||||
"noAutofocus": "off",
|
||||
"noSvgWithoutTitle": "off"
|
||||
}
|
||||
}
|
||||
},
|
||||
"javascript": {
|
||||
|
|
|
|||
|
|
@ -7,9 +7,11 @@ import {
|
|||
useCopilotAction,
|
||||
} from "@copilotkit/react-core";
|
||||
import { CopilotChat } from "@copilotkit/react-ui";
|
||||
import { useCallback, useEffect, useState } from "react";
|
||||
import "@copilotkit/react-ui/styles.css";
|
||||
import CitationBlock from "./CitationBlock";
|
||||
import DbInfo from "./DbInfo";
|
||||
import SettingsPanel, { STORAGE_KEY } from "./SettingsPanel";
|
||||
|
||||
// Must match AGUI_STATE_KEY from haiku.rag.agents.chat
|
||||
const AGUI_STATE_KEY = "haiku.rag.chat";
|
||||
|
|
@ -36,6 +38,7 @@ interface ChatSessionState {
|
|||
session_id: string;
|
||||
citations: Citation[];
|
||||
qa_history: QAResponse[];
|
||||
background_context: string | null;
|
||||
}
|
||||
|
||||
// AG-UI state is namespaced under AGUI_STATE_KEY
|
||||
|
|
@ -131,6 +134,24 @@ function FileIcon() {
|
|||
);
|
||||
}
|
||||
|
||||
function SettingsIcon() {
|
||||
return (
|
||||
<svg
|
||||
width="18"
|
||||
height="18"
|
||||
viewBox="0 0 24 24"
|
||||
fill="none"
|
||||
stroke="currentColor"
|
||||
strokeWidth="2"
|
||||
strokeLinecap="round"
|
||||
strokeLinejoin="round"
|
||||
>
|
||||
<path d="M12.22 2h-.44a2 2 0 0 0-2 2v.18a2 2 0 0 1-1 1.73l-.43.25a2 2 0 0 1-2 0l-.15-.08a2 2 0 0 0-2.73.73l-.22.38a2 2 0 0 0 .73 2.73l.15.1a2 2 0 0 1 1 1.72v.51a2 2 0 0 1-1 1.74l-.15.09a2 2 0 0 0-.73 2.73l.22.38a2 2 0 0 0 2.73.73l.15-.08a2 2 0 0 1 2 0l.43.25a2 2 0 0 1 1 1.73V20a2 2 0 0 0 2 2h.44a2 2 0 0 0 2-2v-.18a2 2 0 0 1 1-1.73l.43-.25a2 2 0 0 1 2 0l.15.08a2 2 0 0 0 2.73-.73l.22-.39a2 2 0 0 0-.73-2.73l-.15-.08a2 2 0 0 1-1-1.74v-.5a2 2 0 0 1 1-1.74l.15-.09a2 2 0 0 0 .73-2.73l-.22-.38a2 2 0 0 0-2.73-.73l-.15.08a2 2 0 0 1-2 0l-.43-.25a2 2 0 0 1-1-1.73V4a2 2 0 0 0-2-2z" />
|
||||
<circle cx="12" cy="12" r="3" />
|
||||
</svg>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolCallIndicator({
|
||||
toolName,
|
||||
status,
|
||||
|
|
@ -336,7 +357,27 @@ function ToolCallIndicator({
|
|||
);
|
||||
}
|
||||
|
||||
function ChatContent() {
|
||||
function ChatContentInner({
|
||||
backgroundContext,
|
||||
setBackgroundContext,
|
||||
}: {
|
||||
backgroundContext: string;
|
||||
setBackgroundContext: (value: string) => void;
|
||||
}) {
|
||||
const [settingsOpen, setSettingsOpen] = useState(false);
|
||||
|
||||
const handleSaveContext = useCallback(
|
||||
(value: string) => {
|
||||
setBackgroundContext(value);
|
||||
if (value) {
|
||||
localStorage.setItem(STORAGE_KEY, value);
|
||||
} else {
|
||||
localStorage.removeItem(STORAGE_KEY);
|
||||
}
|
||||
},
|
||||
[setBackgroundContext],
|
||||
);
|
||||
|
||||
useCoAgent<AgentState>({
|
||||
name: "chat_agent",
|
||||
initialState: {
|
||||
|
|
@ -344,6 +385,7 @@ function ChatContent() {
|
|||
session_id: "",
|
||||
citations: [],
|
||||
qa_history: [],
|
||||
background_context: backgroundContext || null,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
|
@ -425,6 +467,40 @@ function ChatContent() {
|
|||
display: flex;
|
||||
flex-direction: column;
|
||||
}
|
||||
.chat-header {
|
||||
display: flex;
|
||||
justify-content: flex-end;
|
||||
padding: 0.5rem 0.75rem;
|
||||
border-bottom: 1px solid #e2e8f0;
|
||||
background: #f8fafc;
|
||||
}
|
||||
.settings-btn {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 0.375rem;
|
||||
padding: 0.375rem 0.625rem;
|
||||
background: white;
|
||||
border: 1px solid #e2e8f0;
|
||||
border-radius: 6px;
|
||||
cursor: pointer;
|
||||
color: #64748b;
|
||||
font-size: 0.8125rem;
|
||||
transition: all 0.15s;
|
||||
}
|
||||
.settings-btn:hover {
|
||||
background: #f1f5f9;
|
||||
border-color: #cbd5e1;
|
||||
color: #475569;
|
||||
}
|
||||
.settings-btn.has-context {
|
||||
background: #eff6ff;
|
||||
border-color: #bfdbfe;
|
||||
color: #2563eb;
|
||||
}
|
||||
.settings-btn.has-context:hover {
|
||||
background: #dbeafe;
|
||||
border-color: #93c5fd;
|
||||
}
|
||||
.chat-content {
|
||||
flex: 1;
|
||||
min-height: 0;
|
||||
|
|
@ -438,6 +514,21 @@ function ChatContent() {
|
|||
`}</style>
|
||||
<div className="chat-wrapper">
|
||||
<div className="chat-container">
|
||||
<div className="chat-header">
|
||||
<button
|
||||
type="button"
|
||||
className={`settings-btn ${backgroundContext ? "has-context" : ""}`}
|
||||
onClick={() => setSettingsOpen(true)}
|
||||
title={
|
||||
backgroundContext
|
||||
? "Background context is set"
|
||||
: "Set background context"
|
||||
}
|
||||
>
|
||||
<SettingsIcon />
|
||||
Context
|
||||
</button>
|
||||
</div>
|
||||
<div className="chat-content">
|
||||
<CopilotChat
|
||||
labels={{
|
||||
|
|
@ -450,10 +541,38 @@ function ChatContent() {
|
|||
<DbInfo />
|
||||
</div>
|
||||
</div>
|
||||
<SettingsPanel
|
||||
isOpen={settingsOpen}
|
||||
onClose={() => setSettingsOpen(false)}
|
||||
onSave={handleSaveContext}
|
||||
currentValue={backgroundContext}
|
||||
/>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
function ChatContent() {
|
||||
const [backgroundContext, setBackgroundContext] = useState<string | null>(
|
||||
null,
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
const stored = localStorage.getItem(STORAGE_KEY);
|
||||
setBackgroundContext(stored || "");
|
||||
}, []);
|
||||
|
||||
if (backgroundContext === null) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<ChatContentInner
|
||||
backgroundContext={backgroundContext}
|
||||
setBackgroundContext={setBackgroundContext}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
export default function Chat() {
|
||||
return (
|
||||
<CopilotKit runtimeUrl="/api/copilotkit" agent="chat_agent">
|
||||
|
|
|
|||
188
app/frontend/components/SettingsPanel.tsx
Normal file
188
app/frontend/components/SettingsPanel.tsx
Normal file
|
|
@ -0,0 +1,188 @@
|
|||
"use client";
|
||||
|
||||
import { useCallback, useEffect, useId, useState } from "react";
|
||||
|
||||
const STORAGE_KEY = "haiku.rag.settings.background_context";
|
||||
|
||||
interface SettingsPanelProps {
|
||||
isOpen: boolean;
|
||||
onClose: () => void;
|
||||
onSave: (backgroundContext: string) => void;
|
||||
currentValue: string;
|
||||
}
|
||||
|
||||
export default function SettingsPanel({
|
||||
isOpen,
|
||||
onClose,
|
||||
onSave,
|
||||
currentValue,
|
||||
}: SettingsPanelProps) {
|
||||
const [value, setValue] = useState(currentValue);
|
||||
const titleId = useId();
|
||||
|
||||
useEffect(() => {
|
||||
if (isOpen) {
|
||||
setValue(currentValue);
|
||||
}
|
||||
}, [isOpen, currentValue]);
|
||||
|
||||
const handleSave = useCallback(() => {
|
||||
onSave(value);
|
||||
onClose();
|
||||
}, [value, onSave, onClose]);
|
||||
|
||||
const handleKeyDown = useCallback(
|
||||
(e: React.KeyboardEvent) => {
|
||||
if (e.key === "Escape") {
|
||||
onClose();
|
||||
}
|
||||
},
|
||||
[onClose],
|
||||
);
|
||||
|
||||
if (!isOpen) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<style>{`
|
||||
.settings-modal-overlay {
|
||||
position: fixed;
|
||||
top: 0;
|
||||
left: 0;
|
||||
right: 0;
|
||||
bottom: 0;
|
||||
background: rgba(0, 0, 0, 0.5);
|
||||
display: flex;
|
||||
align-items: center;
|
||||
justify-content: center;
|
||||
z-index: 1000;
|
||||
}
|
||||
.settings-modal {
|
||||
background: white;
|
||||
border-radius: 12px;
|
||||
padding: 1.5rem;
|
||||
width: 90%;
|
||||
max-width: 600px;
|
||||
max-height: 80vh;
|
||||
overflow: auto;
|
||||
position: relative;
|
||||
box-shadow: 0 20px 25px -5px rgba(0, 0, 0, 0.1), 0 10px 10px -5px rgba(0, 0, 0, 0.04);
|
||||
}
|
||||
.settings-modal-title {
|
||||
margin: 0 0 0.5rem 0;
|
||||
font-size: 1.25rem;
|
||||
font-weight: 600;
|
||||
color: #1e293b;
|
||||
}
|
||||
.settings-modal-description {
|
||||
margin: 0 0 1rem 0;
|
||||
font-size: 0.875rem;
|
||||
color: #64748b;
|
||||
line-height: 1.5;
|
||||
}
|
||||
.settings-textarea {
|
||||
width: 100%;
|
||||
min-height: 150px;
|
||||
padding: 0.75rem;
|
||||
border: 1px solid #e2e8f0;
|
||||
border-radius: 8px;
|
||||
font-size: 0.875rem;
|
||||
font-family: inherit;
|
||||
resize: vertical;
|
||||
color: #334155;
|
||||
line-height: 1.5;
|
||||
}
|
||||
.settings-textarea:focus {
|
||||
outline: none;
|
||||
border-color: #3b82f6;
|
||||
box-shadow: 0 0 0 3px rgba(59, 130, 246, 0.1);
|
||||
}
|
||||
.settings-textarea::placeholder {
|
||||
color: #94a3b8;
|
||||
}
|
||||
.settings-modal-actions {
|
||||
display: flex;
|
||||
justify-content: flex-end;
|
||||
gap: 0.75rem;
|
||||
margin-top: 1.25rem;
|
||||
}
|
||||
.settings-btn {
|
||||
padding: 0.5rem 1rem;
|
||||
font-size: 0.875rem;
|
||||
font-weight: 500;
|
||||
border-radius: 6px;
|
||||
cursor: pointer;
|
||||
transition: all 0.15s;
|
||||
}
|
||||
.settings-btn-cancel {
|
||||
background: white;
|
||||
color: #475569;
|
||||
border: 1px solid #e2e8f0;
|
||||
}
|
||||
.settings-btn-cancel:hover {
|
||||
background: #f8fafc;
|
||||
border-color: #cbd5e1;
|
||||
}
|
||||
.settings-btn-save {
|
||||
background: #3b82f6;
|
||||
color: white;
|
||||
border: none;
|
||||
}
|
||||
.settings-btn-save:hover {
|
||||
background: #2563eb;
|
||||
}
|
||||
`}</style>
|
||||
<div
|
||||
className="settings-modal-overlay"
|
||||
onClick={onClose}
|
||||
onKeyDown={handleKeyDown}
|
||||
role="dialog"
|
||||
aria-modal="true"
|
||||
aria-labelledby={titleId}
|
||||
>
|
||||
{/* biome-ignore lint/a11y/noStaticElementInteractions: modal content wrapper */}
|
||||
<div
|
||||
className="settings-modal"
|
||||
onClick={(e) => e.stopPropagation()}
|
||||
onKeyDown={(e) => e.stopPropagation()}
|
||||
>
|
||||
<h2 id={titleId} className="settings-modal-title">
|
||||
Background Context
|
||||
</h2>
|
||||
<p className="settings-modal-description">
|
||||
Provide background information that will be used throughout your
|
||||
conversation. This helps the assistant understand domain-specific
|
||||
context, terminology, or any relevant details about your questions.
|
||||
</p>
|
||||
<textarea
|
||||
className="settings-textarea"
|
||||
value={value}
|
||||
onChange={(e) => setValue(e.target.value)}
|
||||
placeholder="e.g., Focus on Python programming concepts and best practices..."
|
||||
autoFocus
|
||||
/>
|
||||
<div className="settings-modal-actions">
|
||||
<button
|
||||
type="button"
|
||||
className="settings-btn settings-btn-cancel"
|
||||
onClick={onClose}
|
||||
>
|
||||
Cancel
|
||||
</button>
|
||||
<button
|
||||
type="button"
|
||||
className="settings-btn settings-btn-save"
|
||||
onClick={handleSave}
|
||||
>
|
||||
Save
|
||||
</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
);
|
||||
}
|
||||
|
||||
export { STORAGE_KEY };
|
||||
|
|
@ -103,6 +103,7 @@ The `ChatSessionState` maintains:
|
|||
|
||||
- `session_id` — Unique identifier for the session
|
||||
- `qa_history` — List of previous Q/A pairs (FIFO, max 50)
|
||||
- `background_context` — Optional background context for the conversation
|
||||
- `embedding_cache` — Cached embeddings for semantic ranking
|
||||
|
||||
Q/A history is used to:
|
||||
|
|
@ -111,6 +112,19 @@ Q/A history is used to:
|
|||
2. Avoid repeating previous answers
|
||||
3. Enable semantic ranking of relevant past answers
|
||||
|
||||
### Background Context
|
||||
|
||||
You can provide background context that persists throughout the conversation:
|
||||
|
||||
```python
|
||||
session = ChatSessionState(
|
||||
background_context="Focus on Python programming concepts and best practices."
|
||||
)
|
||||
deps = ChatDeps(client=client, config=config, session_state=session)
|
||||
```
|
||||
|
||||
The context is included in the agent's system prompt and passed to the research graph when answering questions.
|
||||
|
||||
### AG-UI Integration
|
||||
|
||||
When using the chat agent with AG-UI streaming, state is emitted under a namespaced key to avoid conflicts with other agents:
|
||||
|
|
@ -140,7 +154,7 @@ The emitted state structure:
|
|||
}
|
||||
```
|
||||
|
||||
Frontend clients should extract state from under this key. See the [Conversational RAG App](apps.md#conversational-rag-app) for a complete implementation example.
|
||||
Frontend clients should extract state from under this key. See the [Web Application](apps.md#web-application) for a complete implementation example.
|
||||
|
||||
## Research Graph
|
||||
|
||||
|
|
@ -216,6 +230,18 @@ async with HaikuRAG(path_to_db) as client:
|
|||
print(report.executive_summary)
|
||||
```
|
||||
|
||||
**With background context:**
|
||||
|
||||
```python
|
||||
context = ResearchContext(
|
||||
original_question="What are the safety protocols?",
|
||||
background_context="Industrial manufacturing and workplace safety domain."
|
||||
)
|
||||
state = ResearchState.from_config(context=context, config=Config)
|
||||
```
|
||||
|
||||
The `background_context` provides domain background that helps the planning and synthesis agents understand the context of the research question.
|
||||
|
||||
**With custom config:**
|
||||
|
||||
```python
|
||||
|
|
|
|||
29
docs/cli.md
29
docs/cli.md
|
|
@ -153,6 +153,12 @@ Filter to specific documents:
|
|||
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:
|
||||
|
|
@ -160,6 +166,8 @@ 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
|
||||
|
||||
|
|
@ -170,6 +178,12 @@ haiku-rag chat
|
|||
haiku-rag chat --db /path/to/database.lancedb
|
||||
```
|
||||
|
||||
Provide background context for the conversation:
|
||||
```bash
|
||||
haiku-rag chat --context "Focus on Python programming concepts"
|
||||
haiku-rag chat --context-file domain-context.txt
|
||||
```
|
||||
|
||||
!!! note
|
||||
Requires the `tui` extra: `pip install haiku.rag-slim[tui]` (included in full `haiku.rag` package)
|
||||
|
||||
|
|
@ -179,6 +193,12 @@ The chat interface provides:
|
|||
- Expandable citations with source metadata
|
||||
- Session memory for context-aware follow-up questions
|
||||
- Visual grounding to inspect chunk source locations
|
||||
- Background context that persists across the entire conversation
|
||||
|
||||
Flags:
|
||||
|
||||
- `--context`: Background context for the conversation
|
||||
- `--context-file`: Path to a file containing background context
|
||||
|
||||
See [Applications](apps.md#chat-tui) for keyboard shortcuts and features.
|
||||
|
||||
|
|
@ -217,9 +237,18 @@ 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`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.
|
||||
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
|
|
@ -29,8 +29,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
deps_type=ChatDeps,
|
||||
output_type=str,
|
||||
instructions=CHAT_SYSTEM_PROMPT,
|
||||
retries=3,
|
||||
)
|
||||
|
||||
@agent.system_prompt
|
||||
async def add_background_context(ctx: RunContext[ChatDeps]) -> str:
|
||||
"""Add background_context to system prompt when available."""
|
||||
if ctx.deps.session_state and ctx.deps.session_state.background_context:
|
||||
return f"\nBACKGROUND CONTEXT:\n{ctx.deps.session_state.background_context}"
|
||||
return ""
|
||||
|
||||
@agent.tool
|
||||
async def search(
|
||||
ctx: RunContext[ChatDeps],
|
||||
|
|
@ -85,6 +93,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
qa_history=(
|
||||
ctx.deps.session_state.qa_history if ctx.deps.session_state else []
|
||||
),
|
||||
background_context=(
|
||||
ctx.deps.session_state.background_context
|
||||
if ctx.deps.session_state
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# Return detailed results for the agent to present
|
||||
|
|
@ -181,9 +194,16 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
# Build and run the conversational research graph
|
||||
graph = build_conversational_graph(config=ctx.deps.config)
|
||||
|
||||
background_context = (
|
||||
ctx.deps.session_state.background_context
|
||||
if ctx.deps.session_state
|
||||
else None
|
||||
)
|
||||
|
||||
context = ResearchContext(
|
||||
original_question=question,
|
||||
qa_responses=existing_qa,
|
||||
background_context=background_context,
|
||||
)
|
||||
state = ResearchState(
|
||||
context=context,
|
||||
|
|
@ -237,6 +257,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
|||
qa_history=(
|
||||
ctx.deps.session_state.qa_history if ctx.deps.session_state else []
|
||||
),
|
||||
background_context=(
|
||||
ctx.deps.session_state.background_context
|
||||
if ctx.deps.session_state
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# Format answer with citation references and confidence
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ class SearchAgent:
|
|||
deps_type=SearchDeps,
|
||||
output_type=str,
|
||||
instructions=SEARCH_SYSTEM_PROMPT,
|
||||
retries=3,
|
||||
)
|
||||
|
||||
@self._agent.tool
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import hashlib
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
|
@ -61,6 +61,7 @@ class ChatSessionState(BaseModel):
|
|||
session_id: str = ""
|
||||
citations: list[CitationInfo] = []
|
||||
qa_history: list[QAResponse] = []
|
||||
background_context: str | None = None
|
||||
|
||||
|
||||
def format_conversation_context(qa_history: list[QAResponse]) -> str:
|
||||
|
|
@ -155,7 +156,10 @@ async def rank_qa_history_by_similarity(
|
|||
|
||||
@dataclass
|
||||
class ChatDeps:
|
||||
"""Dependencies for chat agent."""
|
||||
"""Dependencies for chat agent.
|
||||
|
||||
Implements StateHandler protocol for AG-UI state management.
|
||||
"""
|
||||
|
||||
client: HaikuRAG
|
||||
config: AppConfig
|
||||
|
|
@ -163,6 +167,46 @@ class ChatDeps:
|
|||
session_state: ChatSessionState | None = None
|
||||
state_key: str | None = None
|
||||
|
||||
@property
|
||||
def state(self) -> dict[str, Any] | None:
|
||||
"""Get current state for AG-UI protocol."""
|
||||
if self.session_state is None:
|
||||
return None
|
||||
snapshot = self.session_state.model_dump()
|
||||
if self.state_key:
|
||||
return {self.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
|
||||
# Extract from namespaced key if present
|
||||
state_data: dict[str, Any] = value
|
||||
if self.state_key and self.state_key in value:
|
||||
nested = value[self.state_key]
|
||||
if isinstance(nested, dict):
|
||||
state_data = nested
|
||||
# Update session_state from incoming state
|
||||
if self.session_state is not None:
|
||||
if "qa_history" in state_data:
|
||||
self.session_state.qa_history = [
|
||||
QAResponse(**qa) if isinstance(qa, dict) else qa
|
||||
for qa in state_data.get("qa_history", [])
|
||||
]
|
||||
if "citations" in state_data:
|
||||
self.session_state.citations = [
|
||||
CitationInfo(**c) if isinstance(c, dict) else c
|
||||
for c in state_data.get("citations", [])
|
||||
]
|
||||
if "background_context" in state_data:
|
||||
self.session_state.background_context = state_data.get(
|
||||
"background_context"
|
||||
)
|
||||
if "session_id" in state_data:
|
||||
self.session_state.session_id = state_data.get("session_id", "")
|
||||
|
||||
|
||||
@dataclass
|
||||
class SearchDeps:
|
||||
|
|
|
|||
|
|
@ -19,6 +19,10 @@ class ResearchContext(BaseModel):
|
|||
qa_responses: list[Any] = Field(
|
||||
default_factory=list, description="Structured QA pairs used during research"
|
||||
)
|
||||
background_context: str | None = Field(
|
||||
default=None,
|
||||
description="Optional background context provided at session start",
|
||||
)
|
||||
|
||||
def add_qa_response(self, qa: "SearchAnswer") -> None:
|
||||
"""Add a structured QA response."""
|
||||
|
|
|
|||
|
|
@ -30,44 +30,51 @@ from haiku.rag.utils import build_prompt, get_model
|
|||
|
||||
|
||||
def format_context_for_prompt(context: ResearchContext) -> str:
|
||||
"""Format the research context as XML for inclusion in prompts."""
|
||||
context_data = {
|
||||
"original_question": context.original_question,
|
||||
"unanswered_questions": context.sub_questions,
|
||||
"qa_responses": [
|
||||
"""Format the research context as XML for planning prompts."""
|
||||
context_data: dict[str, object] = {}
|
||||
|
||||
if context.background_context:
|
||||
context_data["background"] = context.background_context
|
||||
|
||||
context_data["question"] = context.original_question
|
||||
|
||||
if context.sub_questions:
|
||||
context_data["pending_questions"] = context.sub_questions
|
||||
|
||||
if context.qa_responses:
|
||||
context_data["prior_answers"] = [
|
||||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"confidence": qa.confidence,
|
||||
"sources": [
|
||||
{
|
||||
"document_uri": c.document_uri,
|
||||
"document_title": c.document_title,
|
||||
"page_numbers": c.page_numbers,
|
||||
"headings": c.headings,
|
||||
}
|
||||
for c in qa.citations
|
||||
],
|
||||
"source": qa.citations[0].document_title or qa.citations[0].document_uri
|
||||
if qa.citations
|
||||
else None,
|
||||
}
|
||||
for qa in context.qa_responses
|
||||
],
|
||||
}
|
||||
return format_as_xml(context_data, root_tag="research_context")
|
||||
]
|
||||
|
||||
return format_as_xml(context_data, root_tag="context")
|
||||
|
||||
|
||||
def format_conversational_context_for_prompt(context: ResearchContext) -> str:
|
||||
"""Format context for conversational mode - excludes unanswered_questions."""
|
||||
context_data: dict[str, object] = {
|
||||
"question": context.original_question,
|
||||
}
|
||||
"""Format context for synthesis prompts."""
|
||||
context_data: dict[str, object] = {}
|
||||
|
||||
if context.background_context:
|
||||
context_data["background"] = context.background_context
|
||||
|
||||
context_data["question"] = context.original_question
|
||||
|
||||
# Only include conversation_history if there are qa_responses
|
||||
if context.qa_responses:
|
||||
context_data["conversation_history"] = [
|
||||
context_data["prior_answers"] = [
|
||||
{
|
||||
"question": qa.query,
|
||||
"answer": qa.answer,
|
||||
"sources": [c.document_title or c.document_uri for c in qa.citations],
|
||||
"confidence": qa.confidence,
|
||||
"source": qa.citations[0].document_title or qa.citations[0].document_uri
|
||||
if qa.citations
|
||||
else None,
|
||||
}
|
||||
for qa in context.qa_responses
|
||||
]
|
||||
|
|
@ -90,9 +97,12 @@ async def _plan_step_logic(
|
|||
model_config = config.research.model
|
||||
|
||||
# Use context-aware prompt if we have existing qa_responses
|
||||
has_context = bool(state.context.qa_responses)
|
||||
has_prior_answers = bool(state.context.qa_responses)
|
||||
has_background = bool(state.context.background_context)
|
||||
effective_plan_prompt = (
|
||||
build_prompt(PLAN_PROMPT_WITH_CONTEXT, config) if has_context else plan_prompt
|
||||
build_prompt(PLAN_PROMPT_WITH_CONTEXT, config)
|
||||
if has_prior_answers
|
||||
else plan_prompt
|
||||
)
|
||||
|
||||
plan_agent = Agent(
|
||||
|
|
@ -119,13 +129,20 @@ async def _plan_step_logic(
|
|||
return "\n\n".join(r.content for r in results)
|
||||
|
||||
# Build prompt with existing context if available
|
||||
if has_context:
|
||||
if has_prior_answers:
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = (
|
||||
f"Review existing context and plan additional research if needed.\n\n"
|
||||
f"{context_xml}\n\n"
|
||||
f"Main question: {state.context.original_question}"
|
||||
)
|
||||
elif has_background:
|
||||
context_xml = format_context_for_prompt(state.context)
|
||||
prompt = (
|
||||
f"Plan a focused approach for the main question.\n\n"
|
||||
f"{context_xml}\n\n"
|
||||
f"Main question: {state.context.original_question}"
|
||||
)
|
||||
else:
|
||||
prompt = (
|
||||
"Plan a focused approach for the main question.\n\n"
|
||||
|
|
|
|||
|
|
@ -1,10 +1,11 @@
|
|||
PLAN_PROMPT = """You are the research orchestrator for a focused, iterative workflow.
|
||||
PLAN_PROMPT = """You are the research orchestrator for a focused workflow.
|
||||
|
||||
If a <background> section is provided, use it to understand the domain context.
|
||||
|
||||
Responsibilities:
|
||||
1. Understand and decompose the main question
|
||||
2. Propose a minimal, high-leverage plan
|
||||
3. Coordinate specialized agents to gather evidence
|
||||
4. Iterate based on gaps and new findings
|
||||
|
||||
Plan requirements:
|
||||
- Produce at most 3 sub_questions that together cover the main question.
|
||||
|
|
@ -22,19 +23,22 @@ Use the gather_context tool once on the main question before planning."""
|
|||
|
||||
PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator for a focused workflow.
|
||||
|
||||
You have access to PREVIOUS CONVERSATION CONTEXT in the qa_responses section below.
|
||||
Review this context first - if it already answers the question, generate minimal
|
||||
or no sub-questions. Only create sub-questions to fill gaps in the existing context.
|
||||
You have access to context that may include:
|
||||
- <background>: Domain context for the conversation
|
||||
- <prior_answers>: Previous Q&A pairs with confidence scores
|
||||
|
||||
Review this first - if prior answers already answer the question completely,
|
||||
you may return an empty sub_questions list. Only create sub-questions to
|
||||
fill genuine gaps.
|
||||
|
||||
Responsibilities:
|
||||
1. Review existing qa_responses to understand what's already known
|
||||
1. Review prior_answers to understand what's already known
|
||||
2. Identify gaps that need additional research
|
||||
3. Propose minimal sub-questions only for missing information
|
||||
|
||||
Plan requirements:
|
||||
- If existing context fully answers the question, return a SINGLE sub-question
|
||||
to verify or slightly expand the answer.
|
||||
- Only create new sub-questions for genuine gaps in the existing knowledge.
|
||||
- If prior answers fully answer the question, return an empty sub_questions list.
|
||||
- Only create new sub-questions for genuine gaps in existing knowledge.
|
||||
- sub_questions must be a list of plain strings (max 3).
|
||||
- Each sub_question must be standalone and self-contained.
|
||||
- Prioritize the highest-value gaps first.
|
||||
|
|
@ -145,7 +149,8 @@ Output:
|
|||
- confidence: Score from 0.0 to 1.0 indicating answer quality.
|
||||
|
||||
Guidelines:
|
||||
- Base your answer solely on the collected evidence in qa_responses.
|
||||
- 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..."
|
||||
|
|
|
|||
|
|
@ -376,6 +376,7 @@ class HaikuRAGApp:
|
|||
cite: bool = False,
|
||||
deep: bool = False,
|
||||
filter: str | None = None,
|
||||
background_context: str | None = None,
|
||||
):
|
||||
"""Ask a question using the RAG system.
|
||||
|
||||
|
|
@ -384,6 +385,7 @@ class HaikuRAGApp:
|
|||
cite: Include citations in the answer
|
||||
deep: Use deep QA mode (multi-step reasoning)
|
||||
filter: SQL WHERE clause to filter documents
|
||||
background_context: Optional background context for the question
|
||||
"""
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
|
|
@ -394,7 +396,9 @@ class HaikuRAGApp:
|
|||
citations = []
|
||||
if deep:
|
||||
graph = build_research_graph(config=self.config)
|
||||
context = ResearchContext(original_question=question)
|
||||
context = ResearchContext(
|
||||
original_question=question, background_context=background_context
|
||||
)
|
||||
state = ResearchState.from_config(
|
||||
context=context,
|
||||
config=self.config,
|
||||
|
|
@ -423,7 +427,14 @@ class HaikuRAGApp:
|
|||
else:
|
||||
self.console.print("[yellow]No answer generated.[/yellow]")
|
||||
else:
|
||||
answer, citations = await self.client.ask(question, filter=filter)
|
||||
system_prompt = (
|
||||
f"BACKGROUND CONTEXT:\n{background_context}"
|
||||
if background_context
|
||||
else None
|
||||
)
|
||||
answer, citations = await self.client.ask(
|
||||
question, system_prompt=system_prompt, filter=filter
|
||||
)
|
||||
|
||||
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
|
||||
self.console.print()
|
||||
|
|
@ -433,12 +444,18 @@ class HaikuRAGApp:
|
|||
for renderable in format_citations_rich(citations):
|
||||
self.console.print(renderable)
|
||||
|
||||
async def research(self, question: str, filter: str | None = None):
|
||||
async def research(
|
||||
self,
|
||||
question: str,
|
||||
filter: str | None = None,
|
||||
background_context: str | None = None,
|
||||
):
|
||||
"""Run research via the pydantic-graph pipeline.
|
||||
|
||||
Args:
|
||||
question: The research question
|
||||
filter: SQL WHERE clause to filter documents
|
||||
background_context: Optional background context for the research
|
||||
"""
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
|
|
@ -451,7 +468,9 @@ class HaikuRAGApp:
|
|||
self.console.print()
|
||||
|
||||
graph = build_research_graph(config=self.config)
|
||||
context = ResearchContext(original_question=question)
|
||||
context = ResearchContext(
|
||||
original_question=question, background_context=background_context
|
||||
)
|
||||
state = ResearchState.from_config(context=context, config=self.config)
|
||||
state.search_filter = filter
|
||||
deps = ResearchDeps(client=client)
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ def run_chat(
|
|||
db_path: Path | None = None,
|
||||
read_only: bool = False,
|
||||
before: datetime | None = None,
|
||||
background_context: str | None = None,
|
||||
) -> None:
|
||||
"""Run the chat TUI.
|
||||
|
||||
|
|
@ -13,6 +14,7 @@ def run_chat(
|
|||
db_path: Path to the LanceDB database. If None, uses default from config.
|
||||
read_only: Whether to open the database in read-only mode.
|
||||
before: Query database as it existed before this datetime.
|
||||
background_context: Optional background context for the conversation.
|
||||
"""
|
||||
try:
|
||||
from haiku.rag.chat.app import ChatApp
|
||||
|
|
@ -27,5 +29,10 @@ def run_chat(
|
|||
if db_path is None:
|
||||
db_path = config.storage.data_dir / "haiku.rag.lancedb"
|
||||
|
||||
app = ChatApp(db_path, read_only=read_only, before=before)
|
||||
app = ChatApp(
|
||||
db_path,
|
||||
read_only=read_only,
|
||||
before=before,
|
||||
background_context=background_context,
|
||||
)
|
||||
app.run()
|
||||
|
|
|
|||
|
|
@ -85,12 +85,17 @@ class ChatApp(App): # type: ignore[misc]
|
|||
]
|
||||
|
||||
def __init__(
|
||||
self, db_path: Path, read_only: bool = False, before: datetime | None = None
|
||||
self,
|
||||
db_path: Path,
|
||||
read_only: bool = False,
|
||||
before: datetime | None = None,
|
||||
background_context: str | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.db_path = db_path
|
||||
self.read_only = read_only
|
||||
self.before = before
|
||||
self.background_context = background_context
|
||||
self.client: HaikuRAG | None = None
|
||||
self.config = get_config()
|
||||
self.agent: Agent[ChatDeps, str] | None = None
|
||||
|
|
@ -121,7 +126,10 @@ class ChatApp(App): # type: ignore[misc]
|
|||
|
||||
# Create agent and session state
|
||||
self.agent = create_chat_agent(self.config)
|
||||
self.session_state = ChatSessionState(session_id=str(uuid.uuid4()))
|
||||
self.session_state = ChatSessionState(
|
||||
session_id=str(uuid.uuid4()),
|
||||
background_context=self.background_context,
|
||||
)
|
||||
|
||||
# Focus the input field
|
||||
self.query_one(Input).focus()
|
||||
|
|
@ -266,8 +274,11 @@ class ChatApp(App): # type: ignore[misc]
|
|||
self._last_citations.clear()
|
||||
self._selected_citation_idx = None
|
||||
self._message_history.clear()
|
||||
# Reset session state for fresh conversation
|
||||
self.session_state = ChatSessionState(session_id=str(uuid.uuid4()))
|
||||
# Reset session state for fresh conversation (preserve background_context)
|
||||
self.session_state = ChatSessionState(
|
||||
session_id=str(uuid.uuid4()),
|
||||
background_context=self.background_context,
|
||||
)
|
||||
|
||||
def action_focus_input(self) -> None:
|
||||
"""Focus the input field, or cancel if processing."""
|
||||
|
|
|
|||
|
|
@ -338,9 +338,34 @@ def ask(
|
|||
"-f",
|
||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
||||
),
|
||||
context: str | None = typer.Option(
|
||||
None,
|
||||
"--context",
|
||||
help="Background context for the question",
|
||||
),
|
||||
context_file: Path | None = typer.Option(
|
||||
None,
|
||||
"--context-file",
|
||||
help="Path to a file containing background context",
|
||||
),
|
||||
):
|
||||
# Resolve initial context from flag or file
|
||||
background_context: str | None = None
|
||||
if context_file:
|
||||
background_context = context_file.read_text()
|
||||
elif context:
|
||||
background_context = context
|
||||
|
||||
app = create_app(db)
|
||||
asyncio.run(app.ask(question=question, cite=cite, deep=deep, filter=filter))
|
||||
asyncio.run(
|
||||
app.ask(
|
||||
question=question,
|
||||
cite=cite,
|
||||
deep=deep,
|
||||
filter=filter,
|
||||
background_context=background_context,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@cli.command("research", help="Run multi-agent research and output a concise report")
|
||||
|
|
@ -357,9 +382,30 @@ def research(
|
|||
"-f",
|
||||
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
|
||||
),
|
||||
context: str | None = typer.Option(
|
||||
None,
|
||||
"--context",
|
||||
help="Background context for the research",
|
||||
),
|
||||
context_file: Path | None = typer.Option(
|
||||
None,
|
||||
"--context-file",
|
||||
help="Path to a file containing background context",
|
||||
),
|
||||
):
|
||||
# Resolve initial context from flag or file
|
||||
background_context: str | None = None
|
||||
if context_file:
|
||||
background_context = context_file.read_text()
|
||||
elif context:
|
||||
background_context = context
|
||||
|
||||
app = create_app(db)
|
||||
asyncio.run(app.research(question=question, filter=filter))
|
||||
asyncio.run(
|
||||
app.research(
|
||||
question=question, filter=filter, background_context=background_context
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@cli.command("settings", help="Display current configuration settings")
|
||||
|
|
@ -547,12 +593,35 @@ def chat(
|
|||
"--db",
|
||||
help="Path to the LanceDB database file",
|
||||
),
|
||||
context: str | None = typer.Option(
|
||||
None,
|
||||
"--context",
|
||||
help="Initial context/background information for the conversation",
|
||||
),
|
||||
context_file: Path | None = typer.Option(
|
||||
None,
|
||||
"--context-file",
|
||||
help="Path to a file containing initial context",
|
||||
),
|
||||
):
|
||||
"""Launch the chat TUI for conversational RAG."""
|
||||
from haiku.rag.chat import run_chat
|
||||
|
||||
db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
|
||||
run_chat(db_path, read_only=_read_only, before=_before)
|
||||
|
||||
# Resolve initial context from flag or file
|
||||
background_context: str | None = None
|
||||
if context_file:
|
||||
background_context = context_file.read_text()
|
||||
elif context:
|
||||
background_context = context
|
||||
|
||||
run_chat(
|
||||
db_path,
|
||||
read_only=_read_only,
|
||||
before=_before,
|
||||
background_context=background_context,
|
||||
)
|
||||
|
||||
|
||||
@cli.command(
|
||||
|
|
|
|||
|
|
@ -76,6 +76,17 @@ def test_chat_session_state():
|
|||
assert state.qa_history == []
|
||||
|
||||
|
||||
def test_chat_agent_has_dynamic_system_prompt():
|
||||
"""Test that chat agent registers a dynamic system prompt for background_context."""
|
||||
agent = create_chat_agent(Config)
|
||||
# The agent should have at least one system prompt function registered
|
||||
# (the add_background_context function)
|
||||
assert len(agent._system_prompt_functions) >= 1
|
||||
# Verify it's the add_background_context function
|
||||
func_names = [r.function.__name__ for r in agent._system_prompt_functions]
|
||||
assert "add_background_context" in func_names
|
||||
|
||||
|
||||
def test_citation_info():
|
||||
"""Test CitationInfo model."""
|
||||
citation = CitationInfo(
|
||||
|
|
|
|||
|
|
@ -248,3 +248,225 @@ def test_build_document_filter_escapes_quotes():
|
|||
def test_max_qa_history_constant():
|
||||
"""Test MAX_QA_HISTORY constant value."""
|
||||
assert MAX_QA_HISTORY == 50
|
||||
|
||||
|
||||
def test_chat_session_state_background_context():
|
||||
"""Test ChatSessionState accepts background_context."""
|
||||
from haiku.rag.agents.chat.state import ChatSessionState
|
||||
|
||||
state = ChatSessionState(
|
||||
session_id="test-session",
|
||||
background_context="This is background knowledge about the topic.",
|
||||
)
|
||||
assert state.background_context == "This is background knowledge about the topic."
|
||||
|
||||
|
||||
def test_chat_session_state_background_context_defaults_to_none():
|
||||
"""Test ChatSessionState background_context defaults to None."""
|
||||
from haiku.rag.agents.chat.state import ChatSessionState
|
||||
|
||||
state = ChatSessionState(session_id="test-session")
|
||||
assert state.background_context is None
|
||||
|
||||
|
||||
def test_chat_deps_state_getter_returns_namespaced_state():
|
||||
"""Test ChatDeps.state getter returns state under namespaced key."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
session_state = ChatSessionState(
|
||||
session_id="test-123",
|
||||
qa_history=[
|
||||
QAResponse(question="Q1", answer="A1", confidence=0.9),
|
||||
],
|
||||
background_context="Background info",
|
||||
)
|
||||
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=session_state,
|
||||
state_key=AGUI_STATE_KEY,
|
||||
)
|
||||
|
||||
state = deps.state
|
||||
assert state is not None
|
||||
assert AGUI_STATE_KEY in state
|
||||
assert state[AGUI_STATE_KEY]["session_id"] == "test-123"
|
||||
assert len(state[AGUI_STATE_KEY]["qa_history"]) == 1
|
||||
assert state[AGUI_STATE_KEY]["qa_history"][0]["question"] == "Q1"
|
||||
assert state[AGUI_STATE_KEY]["background_context"] == "Background info"
|
||||
|
||||
|
||||
def test_chat_deps_state_getter_without_namespace():
|
||||
"""Test ChatDeps.state getter returns flat state when no state_key."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
session_state = ChatSessionState(session_id="test-123")
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=session_state,
|
||||
state_key=None,
|
||||
)
|
||||
|
||||
state = deps.state
|
||||
assert state is not None
|
||||
assert "session_id" in state
|
||||
assert state["session_id"] == "test-123"
|
||||
|
||||
|
||||
def test_chat_deps_state_getter_returns_none_without_session():
|
||||
"""Test ChatDeps.state getter returns None when no session_state."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import ChatDeps
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=None,
|
||||
)
|
||||
|
||||
assert deps.state is None
|
||||
|
||||
|
||||
def test_chat_deps_state_setter_updates_from_namespaced_state():
|
||||
"""Test ChatDeps.state setter updates session_state from namespaced incoming state."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
session_state = ChatSessionState(session_id="initial")
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=session_state,
|
||||
state_key=AGUI_STATE_KEY,
|
||||
)
|
||||
|
||||
# Simulate incoming AG-UI state with namespaced key
|
||||
incoming_state = {
|
||||
AGUI_STATE_KEY: {
|
||||
"session_id": "updated-123",
|
||||
"qa_history": [
|
||||
{"question": "Q1", "answer": "A1", "confidence": 0.9, "citations": []}
|
||||
],
|
||||
"citations": [],
|
||||
"background_context": "New context",
|
||||
}
|
||||
}
|
||||
|
||||
deps.state = incoming_state
|
||||
|
||||
assert deps.session_state is not None
|
||||
assert deps.session_state.session_id == "updated-123"
|
||||
assert len(deps.session_state.qa_history) == 1
|
||||
assert deps.session_state.qa_history[0].question == "Q1"
|
||||
assert deps.session_state.background_context == "New context"
|
||||
|
||||
|
||||
def test_chat_deps_state_setter_handles_none():
|
||||
"""Test ChatDeps.state setter handles None gracefully."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import ChatDeps, ChatSessionState
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
session_state = ChatSessionState(session_id="original")
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=session_state,
|
||||
)
|
||||
|
||||
# Setting None should not raise and should not change state
|
||||
deps.state = None
|
||||
|
||||
assert deps.session_state is not None
|
||||
assert deps.session_state.session_id == "original"
|
||||
|
||||
|
||||
def test_chat_deps_state_setter_without_session_state():
|
||||
"""Test ChatDeps.state setter does nothing when session_state is None."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import ChatDeps
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=None,
|
||||
)
|
||||
|
||||
# Should not raise even with valid incoming state
|
||||
deps.state = {"session_id": "test", "qa_history": [], "citations": []}
|
||||
|
||||
assert deps.session_state is None
|
||||
|
||||
|
||||
def test_chat_deps_state_setter_with_citation_dicts():
|
||||
"""Test ChatDeps.state setter converts citation dicts to CitationInfo."""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_config = MagicMock()
|
||||
|
||||
session_state = ChatSessionState(session_id="test")
|
||||
deps = ChatDeps(
|
||||
client=mock_client,
|
||||
config=mock_config,
|
||||
session_state=session_state,
|
||||
state_key=AGUI_STATE_KEY,
|
||||
)
|
||||
|
||||
incoming_state = {
|
||||
AGUI_STATE_KEY: {
|
||||
"session_id": "test",
|
||||
"qa_history": [],
|
||||
"citations": [
|
||||
{
|
||||
"index": 1,
|
||||
"document_id": "doc-1",
|
||||
"chunk_id": "chunk-1",
|
||||
"document_uri": "test.md",
|
||||
"document_title": "Test Doc",
|
||||
"page_numbers": [1, 2],
|
||||
"headings": ["Intro"],
|
||||
"content": "Test content",
|
||||
}
|
||||
],
|
||||
"background_context": None,
|
||||
}
|
||||
}
|
||||
|
||||
deps.state = incoming_state
|
||||
|
||||
assert deps.session_state is not None
|
||||
assert len(deps.session_state.citations) == 1
|
||||
citation = deps.session_state.citations[0]
|
||||
assert citation.document_id == "doc-1"
|
||||
assert citation.chunk_id == "chunk-1"
|
||||
assert citation.page_numbers == [1, 2]
|
||||
|
|
|
|||
|
|
@ -44,3 +44,62 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
|
|||
assert result.executive_summary
|
||||
|
||||
client.close()
|
||||
|
||||
|
||||
def test_research_context_background_context():
|
||||
"""Test ResearchContext accepts background_context."""
|
||||
context = ResearchContext(
|
||||
original_question="What is X?",
|
||||
background_context="Background: X is a concept in domain Y.",
|
||||
)
|
||||
assert context.background_context == "Background: X is a concept in domain Y."
|
||||
|
||||
|
||||
def test_research_context_background_context_defaults_to_none():
|
||||
"""Test ResearchContext background_context defaults to None."""
|
||||
context = ResearchContext(original_question="What is X?")
|
||||
assert context.background_context is None
|
||||
|
||||
|
||||
def test_format_context_for_prompt_includes_background():
|
||||
"""Test format_context_for_prompt includes background in output."""
|
||||
from haiku.rag.agents.research.graph import format_context_for_prompt
|
||||
|
||||
context = ResearchContext(
|
||||
original_question="What is X?",
|
||||
background_context="X is a concept in domain Y.",
|
||||
)
|
||||
result = format_context_for_prompt(context)
|
||||
assert "X is a concept in domain Y." in result
|
||||
assert "<background>" in result
|
||||
|
||||
|
||||
def test_format_context_for_prompt_excludes_background_when_none():
|
||||
"""Test format_context_for_prompt excludes background when None."""
|
||||
from haiku.rag.agents.research.graph import format_context_for_prompt
|
||||
|
||||
context = ResearchContext(original_question="What is X?")
|
||||
result = format_context_for_prompt(context)
|
||||
assert "<background>" not in result
|
||||
|
||||
|
||||
def test_format_conversational_context_for_prompt_includes_background():
|
||||
"""Test format_conversational_context_for_prompt includes background."""
|
||||
from haiku.rag.agents.research.graph import format_conversational_context_for_prompt
|
||||
|
||||
context = ResearchContext(
|
||||
original_question="What is X?",
|
||||
background_context="X is a concept in domain Y.",
|
||||
)
|
||||
result = format_conversational_context_for_prompt(context)
|
||||
assert "X is a concept in domain Y." in result
|
||||
assert "<background>" in result
|
||||
|
||||
|
||||
def test_format_conversational_context_for_prompt_excludes_background_when_none():
|
||||
"""Test format_conversational_context_for_prompt excludes background when None."""
|
||||
from haiku.rag.agents.research.graph import format_conversational_context_for_prompt
|
||||
|
||||
context = ResearchContext(original_question="What is X?")
|
||||
result = format_conversational_context_for_prompt(context)
|
||||
assert "<background>" not in result
|
||||
|
|
|
|||
|
|
@ -304,7 +304,9 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch):
|
|||
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
|
||||
await app.ask("test question")
|
||||
|
||||
mock_client.ask.assert_called_once_with("test question", filter=None)
|
||||
mock_client.ask.assert_called_once_with(
|
||||
"test question", system_prompt=None, filter=None
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -333,7 +335,9 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
|
|||
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
|
||||
await app.ask("test question", cite=True)
|
||||
|
||||
mock_client.ask.assert_called_once_with("test question", filter=None)
|
||||
mock_client.ask.assert_called_once_with(
|
||||
"test question", system_prompt=None, filter=None
|
||||
)
|
||||
# Verify print was called (once for answer, once for citations)
|
||||
assert mock_print.call_count >= 1
|
||||
|
||||
|
|
|
|||
|
|
@ -283,6 +283,7 @@ def test_ask():
|
|||
cite=False,
|
||||
deep=False,
|
||||
filter=None,
|
||||
background_context=None,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -300,6 +301,7 @@ def test_ask_with_cite():
|
|||
cite=True,
|
||||
deep=False,
|
||||
filter=None,
|
||||
background_context=None,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -317,6 +319,7 @@ def test_ask_with_deep():
|
|||
cite=False,
|
||||
deep=True,
|
||||
filter=None,
|
||||
background_context=None,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -334,6 +337,7 @@ def test_ask_with_deep_and_cite():
|
|||
cite=True,
|
||||
deep=True,
|
||||
filter=None,
|
||||
background_context=None,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue