Merge pull request #198 from ggozad/feat/interactive-research

Interactive research through AG-UI in CLI & web example
This commit is contained in:
Yiorgis Gozadinos 2025-12-18 11:49:27 +02:00 committed by GitHub
commit f6618e6037
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
22 changed files with 1479 additions and 208 deletions

View file

@ -3,6 +3,19 @@
### Added ### Added
- **Interactive Research Mode**: Human-in-the-loop research using graph-based decision nodes
- `haiku-rag research --interactive` starts conversational CLI chat
- Natural language interpretation for user commands (search, modify questions, synthesize)
- Chat with assistant before starting research, and during decision points
- Review collected answers and pending questions at each decision point
- Add, remove, or modify sub-questions through natural conversation
- New `human_decide` graph node emits AG-UI tool calls (`TOOL_CALL_START/ARGS/END`) for frontend integration
- New `emit_tool_call_start()`, `emit_tool_call_args()`, `emit_tool_call_end()` AG-UI event helpers
- New `AGUIEmitter.emit()` method for direct event emission
- **AG-UI Research Example**: Human-in-the-loop research with client-side tool calling
- Frontend handles `human_decision` tool calls via AG-UI `TOOL_CALL_*` events
- Tool results sent directly to backend `/v1/research/stream` endpoint
- Backend queues decisions and continues the research graph
- **HotpotQA Evaluation**: Added HotpotQA dataset adapter for multi-hop QA benchmarks - **HotpotQA Evaluation**: Added HotpotQA dataset adapter for multi-hop QA benchmarks
- Extracts unique documents from validation set context paragraphs - Extracts unique documents from validation set context paragraphs
- Uses MAP for retrieval evaluation (multiple supporting documents per question) - Uses MAP for retrieval evaluation (multiple supporting documents per question)

View file

@ -196,3 +196,35 @@ report = await graph.run(state=state, deps=deps)
``` ```
The filter applies to all search operations in the graph. See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax. The filter applies to all search operations in the graph. See [Filtering Search Results](python.md#filtering-search-results) for available filter columns and syntax.
### Interactive Research Mode
Interactive mode provides human-in-the-loop control over the research process through a conversational interface.
**CLI usage:**
```bash
# Start interactive research mode
haiku-rag research --interactive
# Start with a specific question
haiku-rag research --interactive "How does X work?"
# With document filter
haiku-rag research --interactive --filter "uri LIKE '%report%'"
```
In interactive mode, you can:
- Chat with the assistant before starting research
- Review the generated sub-questions after planning
- Add, remove, or modify questions through natural conversation
- Execute searches and review collected answers
- Continue researching or synthesize when ready
For a web-based interactive experience, see the [AG-UI Research Example](https://github.com/ggozad/haiku.rag/tree/main/examples/ag-ui-research). The example demonstrates AG-UI client-side tool calling:
- Frontend handles `human_decision` tool calls via AG-UI `TOOL_CALL_*` events
- Decision UI rendered inline in the chat at each decision point
- Question editing (add/remove) and action buttons (Search, Generate Report)
- Tool results sent directly to the backend endpoint which queues decisions and continues the graph

View file

@ -183,10 +183,24 @@ Filter to specific documents:
haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'" haiku-rag research "What are the key findings?" --filter "uri LIKE '%paper%'"
``` ```
Interactive mode with human-in-the-loop:
```bash
# Start interactive research mode
haiku-rag research --interactive
# Start with a specific question
haiku-rag research --interactive "How does haiku.rag work?"
# With document filter
haiku-rag research --interactive --filter "uri LIKE '%docs%'"
```
Flags: Flags:
- `--verbose`: Show planning, searching previews, evaluation summary, and stop reason - `--verbose`: Show planning, searching previews, evaluation summary, and stop reason
- `--filter`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results)) - `--filter`: SQL WHERE clause to filter documents (see [Filtering Search Results](python.md#filtering-search-results))
- `--interactive` / `-i`: Start interactive research mode with human-in-the-loop decision points
Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section. Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration/index.md) under the `research` section.

View file

@ -6,7 +6,7 @@ import logfire
from pydantic_ai import Agent, RunContext from pydantic_ai import Agent, RunContext
from haiku.rag.config import AppConfig, Config from haiku.rag.config import AppConfig, Config
from haiku.rag.graph.common import get_model from haiku.rag.utils import get_model
from .context import load_message_history, save_message_history from .context import load_message_history, save_message_history
from .models import A2AConfig, AgentDependencies, SearchResult from .models import A2AConfig, AgentDependencies, SearchResult

View file

@ -1,13 +1,13 @@
# Interactive Research Assistant # Interactive Research Assistant
Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time. Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time with human-in-the-loop control.
[Watch demo video](https://vimeo.com/1128874386) [Watch demo video](https://vimeo.com/1128874386)
## Features ## Features
- **Multi-iteration research graph**: Automated question decomposition and search - **Human-in-the-loop research**: Review and modify questions at decision points, then continue searching or generate report
- **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered - **Multi-iteration research graph**: Automated question decomposition and parallel search
- **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol - **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol
- **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources - **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources
@ -25,9 +25,7 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
**Option A: Create a new database** **Option A: Create a new database**
```bash ```bash
mkdir -p data haiku-rag init --db data/haiku_rag.lancedb
haiku-rag add "Your documents here" --db data/haiku_rag.lancedb
# Or add from files
haiku-rag add-src document.pdf --db data/haiku_rag.lancedb haiku-rag add-src document.pdf --db data/haiku_rag.lancedb
``` ```
@ -63,27 +61,29 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
DB_PATH=/path/to/your/existing/haiku_rag.lancedb # If using an existing db. DB_PATH=/path/to/your/existing/haiku_rag.lancedb # If using an existing db.
``` ```
1. **Start the application** 4. **Start the application**
```bash ```bash
docker compose up --build docker compose up --build
``` ```
2. **Access the interface** 5. **Access the interface**
- Frontend: http://localhost:3000 - Frontend: http://localhost:3000
- Backend health: http://localhost:8000/health - Backend health: http://localhost:8000/health
## How It Works ## How It Works
1. **Ask a question**: Type your research question in the chat 1. **Ask a question**: Type your research question in the chat
2. **Plan phase**: The research graph automatically: 2. **Plan phase**: The research graph decomposes your question into targeted sub-questions
- Decomposes your question into targeted sub-questions 3. **Decision point**: Review the proposed questions in the right panel
- Gathers initial context about the topic - Add new questions using the input field
3. **Research iterations**: The graph autonomously: - Remove questions you don't need
- Searches the knowledge base for each sub-question in parallel - Click **Search** to execute searches for pending questions
- Assesses confidence in gathered information - Click **Generate Report** to skip to synthesis (when you have enough answers)
- Generates new follow-up questions if needed 4. **Research iterations**: After each search cycle, you return to a decision point where you can:
- Iterates until confidence threshold is met or max iterations reached - Review collected answers
4. **Synthesis**: Generates a comprehensive research report with: - Add follow-up questions based on findings
- Continue searching or generate the final report
5. **Synthesis**: Generates a comprehensive research report with:
- Executive summary - Executive summary
- Main findings with supporting evidence - Main findings with supporting evidence
- Conclusions and recommendations - Conclusions and recommendations
@ -93,38 +93,42 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
### Agent + Graph Pattern ### Agent + Graph Pattern
This example demonstrates the **agent+graph** architecture pattern: This example demonstrates the **agent+graph** architecture with AG-UI client-side tool calls:
1. **Conversational Agent** (`agent.py`): 1. **Conversational Agent** (`agent.py`):
- Pydantic AI agent handles user conversations - Pydantic AI agent handles user conversations
- Decides when to invoke the research tool based on user intent - Decides when to invoke the research tool based on user intent
- Responds directly to greetings/casual chat without tools - Responds directly to greetings/casual chat without tools
- Formats research results for the user
2. **Research Graph** (haiku.rag): 2. **Interactive Research Graph** (haiku.rag):
- Multi-step research workflow invoked by the agent's tool - Multi-step research workflow invoked by the agent's tool
- Autonomous execution with plan → search → analyze → decide → synthesize flow - At decision points, emits AG-UI `TOOL_CALL_START/ARGS/END` events for `human_decision`
- Emits AG-UI events for real-time progress tracking - Waits for tool result via async queue before continuing
3. **Shared Event Stream**: 3. **Client-Side Tool Handling** (AG-UI pattern):
- Frontend listens for `human_decision` tool calls via AG-UI events
- Renders decision UI inline in chat when tool call is received
- User decision sent directly to backend `/v1/research/stream` endpoint
- Backend extracts tool result from messages and routes to waiting graph via async queue
4. **Shared Event Stream**:
- `AGUIEmitter` is shared between agent and graph - `AGUIEmitter` is shared between agent and graph
- Events from both flow through a single stream to the frontend - Events from both flow through a single stream to the frontend
- Custom streaming endpoint (`main.py`) uses anyio memory streams for proper async handling - `STATE_DELTA` events sync research state to frontend in real-time
### Components ### Components
- **Backend** (Python): - **Backend** (Python):
- Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base - Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
- `agent.py`: Pydantic AI agent with `run_research` tool - `agent.py`: Pydantic AI agent with `run_research` tool, manages `ActiveResearch` registry
- `main.py`: Custom AG-UI streaming endpoint with anyio memory object streams - `main.py`: Custom AG-UI streaming endpoint, extracts tool results from messages
- Real-time event forwarding from emitter to SSE stream - Real-time event forwarding from emitter to SSE stream
- Filters out `ACTIVITY_SNAPSHOT` events (not yet supported by CopilotKit)
- **Frontend** (Next.js/React): - **Frontend** (Next.js/React):
- CopilotKit for AG-UI protocol integration - AG-UI protocol integration for real-time streaming
- Handles `human_decision` tool calls with inline decision UI
- Split-pane UI: chat on left, live research state on right - Split-pane UI: chat on left, live research state on right
- Real-time state synchronization via Server-Sent Events (SSE) - Tool results sent directly to backend endpoint
- `StateDisplay` component with collapsible sections for questions and report
## Configuration ## Configuration

View file

@ -1,6 +1,7 @@
"""Research assistant agent with graph integration.""" """Research assistant agent with graph integration."""
from dataclasses import dataclass import asyncio
from dataclasses import dataclass, field
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@ -9,10 +10,10 @@ from pydantic_ai import Agent, RunContext
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState
from haiku.rag.utils import get_model
if TYPE_CHECKING: if TYPE_CHECKING:
from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.agui.emitter import AGUIEmitter
@ -27,6 +28,20 @@ Config = (
) )
@dataclass
class ActiveResearch:
"""Tracks state for active research awaiting human decision."""
queue: asyncio.Queue[HumanDecision]
sub_questions: list[str] = field(default_factory=list)
qa_responses: list[dict] = field(default_factory=list)
original_question: str = ""
# Global registry of active research by thread_id
_active_research: dict[str, ActiveResearch] = {}
@dataclass @dataclass
class AgentDeps: class AgentDeps:
"""Dependencies for research agent.""" """Dependencies for research agent."""
@ -34,6 +49,8 @@ class AgentDeps:
client: HaikuRAG client: HaikuRAG
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
search_filter: str | None = None search_filter: str | None = None
thread_id: str | None = None
research_result: "ResearchReport | None" = None
model = get_model(Config.research.model, Config) model = get_model(Config.research.model, Config)
@ -50,10 +67,10 @@ CRITICAL RULES:
4. NEVER answer substantive questions from your own knowledge - always use the tool 4. NEVER answer substantive questions from your own knowledge - always use the tool
How to decide: How to decide:
- "Hi" / "Hello" / "How are you?" Respond directly, NO tools - "Hi" / "Hello" / "How are you?" -> Respond directly, NO tools
- "What can you do?" Respond directly, NO tools - "What can you do?" -> Respond directly, NO tools
- "How does X work in the codebase?" Use run_research tool - "How does X work in the codebase?" -> Use run_research tool
- "Tell me about Y" Use run_research tool - "Tell me about Y" -> Use run_research tool
When you use run_research, the graph will decompose questions, search the knowledge base, When you use run_research, the graph will decompose questions, search the knowledge base,
and generate a comprehensive report. and generate a comprehensive report.
@ -70,23 +87,41 @@ async def run_research(ctx: RunContext[AgentDeps], question: str) -> str:
DO NOT use for greetings or casual conversation. DO NOT use for greetings or casual conversation.
""" """
if ctx.deps.agui_emitter: if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.log(f"🔍 Starting research on: {question}") ctx.deps.agui_emitter.log(f"Starting research on: {question}")
graph = build_research_graph(Config) # Create queue for human decisions
queue: asyncio.Queue[HumanDecision] = asyncio.Queue()
# Build interactive graph
graph = build_research_graph(Config, interactive=True)
context = ResearchContext(original_question=question) context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config) state = ResearchState.from_config(context=context, config=Config)
state.search_filter = ctx.deps.search_filter state.search_filter = ctx.deps.search_filter
# Register active research for decision endpoint to find
thread_id = ctx.deps.thread_id
if thread_id:
_active_research[thread_id] = ActiveResearch(
queue=queue,
sub_questions=[],
qa_responses=[],
original_question=question,
)
graph_deps = ResearchDeps( graph_deps = ResearchDeps(
client=ctx.deps.client, client=ctx.deps.client,
agui_emitter=ctx.deps.agui_emitter, agui_emitter=ctx.deps.agui_emitter,
human_input_queue=queue,
interactive=True,
) )
try: try:
result = await graph.run(state=state, deps=graph_deps) result = await graph.run(state=state, deps=graph_deps)
if ctx.deps.agui_emitter: if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.log("✅ Research complete!") ctx.deps.agui_emitter.log("Research complete!")
# Store result for main.py to emit RUN_FINISHED after agent completes
ctx.deps.research_result = result
return f"""Research completed successfully! return f"""Research completed successfully!
@ -108,5 +143,9 @@ The full research report with all citations has been provided to the user.
except Exception as e: except Exception as e:
if ctx.deps.agui_emitter: if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.log(f"Research error: {str(e)}") ctx.deps.agui_emitter.log(f"Research error: {str(e)}")
return f"I encountered an error while researching: {str(e)}" return f"I encountered an error while researching: {str(e)}"
finally:
# Cleanup
if thread_id and thread_id in _active_research:
del _active_research[thread_id]

View file

@ -1,8 +1,9 @@
import json
import logging import logging
import os import os
from pathlib import Path from pathlib import Path
from agent import AgentDeps, agent from agent import AgentDeps, _active_research, agent
from anyio import create_memory_object_stream, create_task_group from anyio import create_memory_object_stream, create_task_group
from anyio.streams.memory import MemoryObjectSendStream from anyio.streams.memory import MemoryObjectSendStream
from starlette.applications import Starlette from starlette.applications import Starlette
@ -19,7 +20,7 @@ from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import ResearchReport from haiku.rag.graph.research.models import ResearchReport
from haiku.rag.graph.research.state import ResearchState from haiku.rag.graph.research.state import HumanDecision, ResearchState
logging.basicConfig( logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
@ -62,11 +63,54 @@ def get_client(effective_db_path: Path) -> HaikuRAG:
return _client_cache[path_key] return _client_cache[path_key]
def extract_tool_result(messages: list[dict]) -> dict | None:
"""Extract human_decision tool result from messages if present."""
for msg in reversed(messages):
# Check for tool result message (CopilotKit sends role="tool")
if msg.get("role") == "tool":
content = msg.get("content")
# Content may be a string (JSON) or dict
if isinstance(content, str):
try:
content = json.loads(content)
except json.JSONDecodeError:
continue
if isinstance(content, dict) and "action" in content:
return content
return None
async def stream_research_agent(request: Request) -> StreamingResponse: async def stream_research_agent(request: Request) -> StreamingResponse:
"""Agent streaming endpoint with research graph integration.""" """Agent streaming endpoint with research graph integration."""
body = await request.json() body = await request.json()
logger.info(f"Received request body keys: {list(body.keys())}")
if "tools" in body:
logger.info(f"Frontend tools received: {body['tools']}")
input_data = RunAgentInput(**body) input_data = RunAgentInput(**body)
thread_id = input_data.thread_id
active_research = _active_research.get(thread_id) if thread_id else None
# Check if this is a tool result for active research
if active_research and input_data.messages:
tool_result = extract_tool_result(input_data.messages)
if tool_result:
logger.info(f"Received tool result: {tool_result}")
action = tool_result.get("action", "search")
questions = tool_result.get("questions")
decision = HumanDecision(
action=action,
questions=questions,
)
await active_research.queue.put(decision)
# Return acknowledgment - the original stream will continue
return StreamingResponse(
iter([format_sse_event({"type": "TOOL_RESULT_RECEIVED"})]),
media_type="text/event-stream",
)
user_message = "" user_message = ""
if input_data.messages: if input_data.messages:
user_message = input_data.messages[-1].get("content", "") user_message = input_data.messages[-1].get("content", "")
@ -79,11 +123,11 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
"""Execute agent and forward emitter events to memory stream.""" """Execute agent and forward emitter events to memory stream."""
async with send_stream: async with send_stream:
try: try:
# Create shared emitter # Create shared emitter (use_deltas=True for CopilotKit compatibility)
emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter( emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter(
thread_id=input_data.thread_id, thread_id=input_data.thread_id,
run_id=input_data.run_id, run_id=input_data.run_id,
use_deltas=False, use_deltas=True,
) )
# Get client # Get client
@ -104,6 +148,7 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
client=client, client=client,
agui_emitter=emitter, agui_emitter=emitter,
search_filter=search_filter, search_filter=search_filter,
thread_id=thread_id,
) )
# Start run with empty initial state # Start run with empty initial state
@ -117,9 +162,12 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
# Forward emitter events to stream # Forward emitter events to stream
async def forward_events(): async def forward_events():
async for event in emitter: async for event in emitter:
# Log events for debugging
logger.info(f"AG-UI Event: {event}")
event_type = event.get("type") event_type = event.get("type")
logger.info(f"AG-UI event: {event_type}")
# Log tool call events for debugging
if event_type and event_type.startswith("TOOL_CALL"):
logger.info(f"Tool call event: {event}")
# Convert ACTIVITY_SNAPSHOT to STATE_DELTA for CopilotKit # Convert ACTIVITY_SNAPSHOT to STATE_DELTA for CopilotKit
# As CopilotKit does not handle ACTIVITY_SNAPSHOT events # As CopilotKit does not handle ACTIVITY_SNAPSHOT events
@ -129,7 +177,6 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
message = content.get("message", "") message = content.get("message", "")
# Emit STATE_DELTA to patch activity info into state # Emit STATE_DELTA to patch activity info into state
# Use "add" op which creates or replaces the value
delta_event = { delta_event = {
"type": "STATE_DELTA", "type": "STATE_DELTA",
"delta": [ "delta": [
@ -148,6 +195,21 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
await send_stream.send(format_sse_event(delta_event)) await send_stream.send(format_sse_event(delta_event))
continue continue
# Sync state to ActiveResearch when human_decision tool call
if event_type == "TOOL_CALL_ARGS" and thread_id:
delta = event.get("delta", "{}")
args = (
json.loads(delta) if isinstance(delta, str) else delta
)
active = _active_research.get(thread_id)
if active:
active.sub_questions = list(
args.get("sub_questions", [])
)
active.qa_responses = list(args.get("qa_responses", []))
if "original_question" in args:
active.original_question = args["original_question"]
await send_stream.send(format_sse_event(event)) await send_stream.send(format_sse_event(event))
# Run agent and event forwarding concurrently # Run agent and event forwarding concurrently
@ -156,6 +218,9 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
result = await agent.run(user_message, deps=agent_deps) result = await agent.run(user_message, deps=agent_deps)
emitter.log(result.output) emitter.log(result.output)
# Emit RUN_FINISHED with research result if available
if agent_deps.research_result:
emitter.finish_run(agent_deps.research_result)
await emitter.close() await emitter.close()
except Exception as e: except Exception as e:

View file

@ -32,9 +32,13 @@ services:
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL} - OLLAMA_BASE_URL=${OLLAMA_BASE_URL}
# Prevent Python bytecode caching for development # Prevent Python bytecode caching for development
- PYTHONDONTWRITEBYTECODE=1 - PYTHONDONTWRITEBYTECODE=1
# Use local haiku_rag_slim for development
- PYTHONPATH=/app/haiku_rag_slim
volumes: volumes:
- ${DB_PATH}:/app/data/haiku.rag.lancedb - ${DB_PATH}:/app/data/haiku.rag.lancedb
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro - ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
# Mount local haiku_rag_slim for development
- ../../haiku_rag_slim:/app/haiku_rag_slim:ro
networks: networks:
- ag-ui-network - ag-ui-network
extra_hosts: extra_hosts:

View file

@ -4,7 +4,7 @@ FROM node:22-alpine
WORKDIR /app WORKDIR /app
COPY package.json package-lock.json* ./ COPY package.json package-lock.json* ./
RUN npm ci RUN npm install --legacy-peer-deps
COPY . . COPY . .
EXPOSE 3000 EXPOSE 3000

View file

@ -1,64 +1,282 @@
"use client"; "use client";
import { CopilotKit, useCoAgent } from "@copilotkit/react-core"; import {
CopilotKit,
useCoAgent,
useCopilotAction,
useCopilotContext,
} from "@copilotkit/react-core";
import { CopilotChat } from "@copilotkit/react-ui"; import { CopilotChat } from "@copilotkit/react-ui";
import "@copilotkit/react-ui/styles.css"; import "@copilotkit/react-ui/styles.css";
import { useState } from "react";
import DocumentSelector from "./DocumentSelector"; import DocumentSelector from "./DocumentSelector";
import StateDisplay from "./StateDisplay"; import StateDisplay from "./StateDisplay";
interface Citation {
document_id: string;
chunk_id: string;
document_uri: string;
document_title?: string;
page_numbers: number[];
headings?: string[];
content: string;
}
interface SearchAnswer { interface SearchAnswer {
query: string; query: string;
answer: string; answer: string;
confidence: number; confidence: number;
cited_chunks: string[]; cited_chunks: string[];
citations: Citation[]; citations: {
document_id: string;
chunk_id: string;
document_uri: string;
document_title?: string;
page_numbers: number[];
headings?: string[];
content: string;
}[];
} }
interface ResearchContext { interface ResearchState {
context: {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
};
iterations: number;
max_iterations: number;
confidence_threshold: number;
max_concurrency: number;
last_eval: {
new_questions: string[];
confidence_score: number;
is_sufficient: boolean;
reasoning: string;
} | null;
result?: {
title: string;
executive_summary: string;
main_findings: string[];
conclusions: string[];
limitations: string[];
recommendations: string[];
sources_summary: string;
};
current_activity?: string;
current_activity_message?: string;
documentFilter?: string[];
}
interface DecisionArgs {
original_question: string; original_question: string;
sub_questions: string[]; sub_questions: string[];
qa_responses: SearchAnswer[]; qa_responses: SearchAnswer[];
} }
interface EvaluationResult { type DecisionAction = "search" | "synthesize" | "modify_questions";
new_questions: string[];
confidence_score: number; interface DecisionResult {
is_sufficient: boolean; action: DecisionAction;
reasoning: string; questions?: string[];
} }
interface ResearchReport { function DecisionUI({
title: string; args,
executive_summary: string; onResolve,
main_findings: string[]; }: {
conclusions: string[]; args: DecisionArgs;
limitations: string[]; onResolve: (result: DecisionResult) => void | Promise<void>;
recommendations: string[]; }) {
sources_summary: string; const [editableQuestions, setEditableQuestions] = useState<string[]>(
args.sub_questions || [],
);
const [newQuestion, setNewQuestion] = useState("");
const [submitting, setSubmitting] = useState(false);
const qaCount = args.qa_responses?.length || 0;
const hasQuestions = editableQuestions.length > 0;
const canSearch = hasQuestions && !submitting;
const canSynthesize = qaCount > 0 && !submitting;
const questionsModified =
editableQuestions.length !== args.sub_questions.length ||
editableQuestions.some((q, i) => q !== args.sub_questions[i]);
const handleSubmit = (action: DecisionAction, questions?: string[]) => {
setSubmitting(true);
onResolve({ action, questions });
};
const handleSearch = () => {
handleSubmit(
questionsModified ? "modify_questions" : "search",
editableQuestions,
);
};
const handleSynthesize = () => {
handleSubmit("synthesize");
};
const handleRemoveQuestion = (index: number) => {
if (submitting) return;
setEditableQuestions(editableQuestions.filter((_, i) => i !== index));
};
const handleAddQuestion = () => {
if (submitting || !newQuestion.trim()) return;
setEditableQuestions([...editableQuestions, newQuestion.trim()]);
setNewQuestion("");
};
if (submitting) {
return null;
}
return (
<div
style={{
marginBottom: "1rem",
background: "#f0f9ff",
border: "2px solid #0ea5e9",
borderRadius: "8px",
padding: "1rem",
}}
>
<div
style={{
fontWeight: "bold",
color: "#0369a1",
marginBottom: "0.75rem",
fontSize: "1rem",
}}
>
Research Decision Point
</div>
<div
style={{
fontSize: "0.85rem",
color: "#64748b",
marginBottom: "0.75rem",
}}
>
{qaCount} answers collected
</div>
<div style={{ marginBottom: "0.75rem" }}>
<div
style={{
fontSize: "0.8rem",
color: "#475569",
marginBottom: "0.5rem",
}}
>
Pending Questions ({editableQuestions.length}):
</div>
{editableQuestions.map((q, idx) => (
<div
key={`question-${idx}`}
style={{
display: "flex",
alignItems: "center",
gap: "0.5rem",
padding: "0.375rem 0.5rem",
background: "white",
borderRadius: "4px",
marginBottom: "0.25rem",
fontSize: "0.85rem",
}}
>
<span style={{ flex: 1 }}>{q}</span>
<button
type="button"
onClick={() => handleRemoveQuestion(idx)}
style={{
background: "#ef4444",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.25rem 0.5rem",
cursor: "pointer",
fontSize: "0.75rem",
}}
>
Remove
</button>
</div>
))}
</div>
<div style={{ display: "flex", gap: "0.5rem", marginBottom: "1rem" }}>
<input
type="text"
value={newQuestion}
onChange={(e) => setNewQuestion(e.target.value)}
placeholder="Add a new question..."
style={{
flex: 1,
padding: "0.5rem",
border: "1px solid #cbd5e1",
borderRadius: "4px",
fontSize: "0.85rem",
}}
onKeyDown={(e) => {
if (e.key === "Enter") handleAddQuestion();
}}
/>
<button
type="button"
onClick={handleAddQuestion}
style={{
background: "#22c55e",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.5rem 1rem",
cursor: "pointer",
fontSize: "0.85rem",
}}
>
Add
</button>
</div>
<div style={{ display: "flex", gap: "0.5rem" }}>
<button
type="button"
onClick={handleSearch}
disabled={!canSearch}
style={{
flex: 1,
background: canSearch ? "#0ea5e9" : "#94a3b8",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.75rem",
cursor: canSearch ? "pointer" : "not-allowed",
fontWeight: "bold",
fontSize: "0.9rem",
}}
>
Search ({editableQuestions.length})
</button>
<button
type="button"
onClick={handleSynthesize}
disabled={!canSynthesize}
style={{
flex: 1,
background: canSynthesize ? "#8b5cf6" : "#94a3b8",
color: "white",
border: "none",
borderRadius: "4px",
padding: "0.75rem",
cursor: canSynthesize ? "pointer" : "not-allowed",
fontWeight: "bold",
fontSize: "0.9rem",
}}
>
Generate Report
</button>
</div>
</div>
);
} }
interface ResearchState { const BACKEND_URL =
context: ResearchContext; process.env.NEXT_PUBLIC_BACKEND_URL || "http://localhost:8000";
iterations: number;
max_iterations: number;
confidence_threshold: number;
max_concurrency: number;
last_eval: EvaluationResult | null;
result?: ResearchReport;
current_activity?: string;
current_activity_message?: string;
documentFilter?: string[];
}
function AgentContent() { function AgentContent() {
const { state, setState, running } = useCoAgent<ResearchState>({ const { state, setState, running } = useCoAgent<ResearchState>({
@ -78,10 +296,76 @@ function AgentContent() {
}, },
}); });
const { threadId } = useCopilotContext();
const handleDocumentFilterChange = (ids: string[]) => { const handleDocumentFilterChange = (ids: string[]) => {
setState({ ...state, documentFilter: ids }); setState({ ...state, documentFilter: ids });
}; };
const sendToolResult = async (result: DecisionResult) => {
if (!threadId) {
console.error("No threadId available to send tool result");
return;
}
try {
const response = await fetch(`${BACKEND_URL}/v1/research/stream`, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
threadId,
messages: [
{
id: crypto.randomUUID(),
role: "tool",
content: JSON.stringify(result),
},
],
}),
});
if (!response.ok) {
console.error("Failed to send tool result:", response.status);
}
} catch (error) {
console.error("Error sending tool result:", error);
}
};
useCopilotAction({
name: "human_decision",
description: "Pause for human decision on research direction",
parameters: [
{
name: "original_question",
type: "string",
description: "The original research question",
},
{
name: "sub_questions",
type: "string[]",
description: "Pending sub-questions to search",
},
{
name: "qa_responses",
type: "object[]",
description: "Answers collected so far",
},
],
renderAndWaitForResponse: ({ args, status }) => {
if (status === "complete") {
return null;
}
return (
<DecisionUI
args={args as unknown as DecisionArgs}
onResolve={sendToolResult}
/>
);
},
});
return ( return (
<> <>
<style>{` <style>{`
@ -98,7 +382,6 @@ function AgentContent() {
} }
`}</style> `}</style>
<div style={{ display: "flex", height: "100vh" }}> <div style={{ display: "flex", height: "100vh" }}>
{/* Chat on the left */}
<div className="chat-container"> <div className="chat-container">
<CopilotChat <CopilotChat
labels={{ labels={{
@ -109,7 +392,6 @@ function AgentContent() {
/> />
</div> </div>
{/* State display on the right */}
<div <div
style={{ style={{
width: "50%", width: "50%",
@ -141,13 +423,14 @@ function AgentContent() {
</p> </p>
</header> </header>
<div style={{ marginBottom: "1rem" }}> {!running && (
<DocumentSelector <div style={{ marginBottom: "1rem" }}>
selected={state.documentFilter || []} <DocumentSelector
onChange={handleDocumentFilterChange} selected={state.documentFilter || []}
disabled={running} onChange={handleDocumentFilterChange}
/> />
</div> </div>
)}
<StateDisplay state={state} /> <StateDisplay state={state} />
</div> </div>

View file

@ -298,22 +298,29 @@ export default function DocumentSelector({
}} }}
/> />
<div style={{ flex: 1, minWidth: 0 }}> <div style={{ flex: 1, minWidth: 0 }}>
{doc.title && (
<div
style={{
fontSize: "0.875rem",
fontWeight: isSelected ? "600" : "400",
color: "#2d3748",
whiteSpace: "nowrap",
overflow: "hidden",
textOverflow: "ellipsis",
}}
>
{doc.title}
</div>
)}
<div <div
style={{ style={{
fontSize: "0.875rem", fontSize: doc.title ? "0.7rem" : "0.875rem",
fontWeight: isSelected ? "600" : "400", fontWeight: doc.title
color: "#2d3748", ? "400"
whiteSpace: "nowrap", : isSelected
overflow: "hidden", ? "600"
textOverflow: "ellipsis", : "400",
}} color: doc.title ? "#718096" : "#2d3748",
>
{doc.title || "Untitled"}
</div>
<div
style={{
fontSize: "0.7rem",
color: "#718096",
whiteSpace: "nowrap", whiteSpace: "nowrap",
overflow: "hidden", overflow: "hidden",
textOverflow: "ellipsis", textOverflow: "ellipsis",

View file

@ -162,7 +162,7 @@ export default function StateDisplay({ state }: StateDisplayProps) {
}} }}
> >
{/* Question */} {/* Question */}
{state.context.original_question && ( {state.context?.original_question && (
<div <div
style={{ style={{
background: "white", background: "white",
@ -192,8 +192,8 @@ export default function StateDisplay({ state }: StateDisplayProps) {
</div> </div>
)} )}
{/* Research Progress - only show when research has started */} {/* Research Progress - only show when research is in progress (not when complete) */}
{(state.iterations > 0 || (state.current_activity && !state.result)) && ( {(state.iterations > 0 || state.current_activity) && !state.result && (
<div <div
style={{ style={{
background: "white", background: "white",
@ -202,8 +202,8 @@ export default function StateDisplay({ state }: StateDisplayProps) {
boxShadow: "0 1px 3px rgba(0,0,0,0.1)", boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
}} }}
> >
{/* Current Activity - hide when complete */} {/* Current Activity */}
{state.current_activity && !state.result && ( {state.current_activity && (
<div <div
style={{ style={{
padding: "0.75rem", padding: "0.75rem",
@ -363,9 +363,8 @@ export default function StateDisplay({ state }: StateDisplayProps) {
</div> </div>
)} )}
{/* Sub-Questions and QA Responses */} {/* Answers */}
{(state.context.sub_questions.length > 0 || {state.context?.qa_responses && state.context.qa_responses.length > 0 && (
state.context.qa_responses.length > 0) && (
<div <div
style={{ style={{
background: "white", background: "white",
@ -392,10 +391,7 @@ export default function StateDisplay({ state }: StateDisplayProps) {
color: "#2d3748", color: "#2d3748",
}} }}
> >
<span> <span>Answers ({state.context.qa_responses.length})</span>
Sub-Questions ({state.context.sub_questions.length}) Answers (
{state.context.qa_responses.length})
</span>
<span>{expandedSections.questions ? "▼" : "▶"}</span> <span>{expandedSections.questions ? "▼" : "▶"}</span>
</button> </button>
{expandedSections.questions && ( {expandedSections.questions && (
@ -408,38 +404,6 @@ export default function StateDisplay({ state }: StateDisplayProps) {
borderRadius: "0 0 4px 4px", borderRadius: "0 0 4px 4px",
}} }}
> >
{/* Show pending sub_questions */}
{state.context.sub_questions.map((question, idx) => (
<div
key={`pending-${idx}`}
style={{
marginBottom: "0.5rem",
background: "white",
borderRadius: "4px",
border: "1px solid #e2e8f0",
padding: "0.75rem",
display: "flex",
gap: "0.75rem",
alignItems: "center",
}}
>
<div
style={{
fontSize: "1.25rem",
color: "#a0aec0",
flexShrink: 0,
}}
>
</div>
<div
style={{ flex: 1, fontSize: "0.875rem", color: "#4a5568" }}
>
<Markdown content={question} />
</div>
</div>
))}
{/* Show all qa_responses (each has query + answer) */} {/* Show all qa_responses (each has query + answer) */}
{state.context.qa_responses.map((qaResponse, idx) => { {state.context.qa_responses.map((qaResponse, idx) => {
const questionId = `q-${idx}`; const questionId = `q-${idx}`;

View file

@ -12,9 +12,9 @@
}, },
"dependencies": { "dependencies": {
"@ag-ui/client": "^0.0.42", "@ag-ui/client": "^0.0.42",
"@copilotkit/react-core": "^1.10.6", "@copilotkit/react-core": "^1.50.0",
"@copilotkit/react-ui": "^1.10.6", "@copilotkit/react-ui": "^1.50.0",
"@copilotkit/runtime": "^1.10.6", "@copilotkit/runtime": "^1.50.0",
"next": "15.5.5", "next": "15.5.5",
"react": "^19.0.0", "react": "^19.0.0",
"react-dom": "^19.0.0" "react-dom": "^19.0.0"

View file

@ -321,7 +321,8 @@ def ask(
@cli.command("research", help="Run multi-agent research and output a concise report") @cli.command("research", help="Run multi-agent research and output a concise report")
def research( def research(
question: str = typer.Argument( question: str = typer.Argument(
help="The research question to investigate", None,
help="The research question to investigate (required unless --interactive)",
), ),
db: Path | None = typer.Option( db: Path | None = typer.Option(
None, None,
@ -339,9 +340,34 @@ def research(
"-f", "-f",
help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")", help="SQL WHERE clause to filter documents (e.g., \"uri LIKE '%arxiv%'\")",
), ),
interactive: bool = typer.Option(
False,
"--interactive",
"-i",
help="Start interactive research mode with human-in-the-loop",
),
): ):
app = create_app(db) app = create_app(db)
asyncio.run(app.research(question=question, verbose=verbose, filter=filter))
if interactive:
from haiku.rag.cli_chat import interactive_research
from haiku.rag.client import HaikuRAG
client = HaikuRAG(db_path=app.db_path, config=app.config)
try:
interactive_research(
client=client,
config=app.config,
search_filter=filter,
question=question,
)
finally:
client.close()
else:
if question is None:
typer.echo("Error: Question is required unless using --interactive mode")
raise typer.Exit(1)
asyncio.run(app.research(question=question, verbose=verbose, filter=filter))
@cli.command("settings", help="Display current configuration settings") @cli.command("settings", help="Display current configuration settings")

View file

@ -0,0 +1,489 @@
"""Interactive CLI chat loop for research graph with human-in-the-loop."""
import asyncio
import json
from pydantic_ai import Agent
from rich.console import Console
from rich.markdown import Markdown
from rich.panel import Panel
from rich.prompt import Prompt
from haiku.rag.client import HaikuRAG
from haiku.rag.config import get_config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.models import ResearchReport
from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState
from haiku.rag.utils import get_model
INITIAL_CHAT_PROMPT = """You are a research assistant. The user hasn't started a research task yet.
You can:
1. Chat with the user - greet them, answer questions about what you can do
2. Detect when they want to research something
## Actions:
- "chat": User is chatting, greeting, or asking questions (set message with your response)
- "research": User wants to research a topic (extract the research question into research_question)
## Guidelines:
- If the user provides a clear research question or topic, set action="research" and extract the question
- If the user is just chatting or asking what you can do, set action="chat" and respond helpfully
- Be friendly and explain you can help them research topics by searching a knowledge base
Examples:
- "hi" action="chat", message="Hello! I'm a research assistant. I can help you research topics by searching through documents and synthesizing findings. What would you like to explore?"
- "what can you do?" action="chat", message="I help you conduct research! Give me a question or topic, and I'll break it into sub-questions, search for answers, and synthesize a report. What are you curious about?"
- "tell me about Python's memory management" action="research", research_question="How does Python's memory management work?"
- "I want to understand how RAG systems work" action="research", research_question="How do RAG (Retrieval-Augmented Generation) systems work?"
"""
RESEARCH_ASSISTANT_PROMPT = """You are a research assistant helping the user conduct research on a topic.
You are at a decision point in the research workflow. You can:
1. Chat with the user - answer questions, discuss the research, make suggestions
2. Take workflow actions when the user requests them
## Workflow Actions (set in the action field):
- "search": Search the pending questions (user says: "go", "search", "yes", "continue", "looks good")
- "synthesize": Generate final report (user says: "done", "finish", "synthesize", "generate report")
- "add_questions": Add NEW research questions to the existing list
- "modify_questions": REPLACE all pending questions with a new list (use when user wants to remove, keep only certain questions, or change the questions)
- "chat": Have a conversation without modifying questions
## IMPORTANT - Modifying Questions:
- "use only the first question" action="modify_questions", questions=[first question from the list]
- "drop questions 2 and 3" action="modify_questions", questions=[remaining questions]
- "keep only questions about X" action="modify_questions", questions=[filtered list]
- "remove the duplicate" action="modify_questions", questions=[deduplicated list]
- When user wants to reduce/filter/keep-only, use "modify_questions" NOT "chat"
## Guidelines:
- If the user wants to modify the question list in ANY way (remove, keep only, filter), use "modify_questions"
- For "modify_questions", include ALL questions that should remain in the questions field
- You can combine "chat" with a message to explain what you're doing
- If just chatting without changes, set action="chat" and provide helpful response in message
"""
async def initial_chat(
user_message: str,
config: AppConfig,
) -> HumanDecision:
"""Handle initial conversation before research starts.
Args:
user_message: The user's message
config: Application configuration
Returns:
HumanDecision with chat response or research question
"""
agent: Agent[None, HumanDecision] = Agent(
model=get_model(config.research.model, config),
output_type=HumanDecision,
instructions=INITIAL_CHAT_PROMPT,
retries=2,
)
result = await agent.run(user_message)
return result.output
async def interpret_user_decision(
user_message: str,
sub_questions: list[str],
qa_responses: list[dict],
config: AppConfig,
) -> HumanDecision:
"""Interpret a natural language user message into a HumanDecision.
Args:
user_message: The user's natural language input
sub_questions: Current sub-questions pending search
qa_responses: Answers already collected
config: Application configuration
Returns:
HumanDecision with the interpreted action, questions, and/or message
"""
agent: Agent[None, HumanDecision] = Agent(
model=get_model(config.research.model, config),
output_type=HumanDecision,
instructions=RESEARCH_ASSISTANT_PROMPT,
retries=2,
)
# Build context with full research state
answers_summary = ""
if qa_responses:
answers_parts = []
for qa in qa_responses:
conf = f"{qa['confidence']:.0%}" if qa.get("confidence") else "N/A"
answers_parts.append(
f"Q: {qa['query']}\nA: {qa['answer'][:300]}... (confidence: {conf})"
)
answers_summary = "\n\n".join(answers_parts)
context = f"""Current research state:
- Answers collected: {len(qa_responses)}
- Pending questions to search: {len(sub_questions)}
Pending questions:
{chr(10).join(f"- {q}" for q in sub_questions) if sub_questions else "(none)"}
{f"Collected answers:{chr(10)}{answers_summary}" if answers_summary else ""}
User message: {user_message}"""
result = await agent.run(context)
return result.output
async def run_interactive_research(
question: str,
client: HaikuRAG,
config: AppConfig | None = None,
search_filter: str | None = None,
) -> ResearchReport:
"""Run interactive research with human-in-the-loop decision points.
Args:
question: The research question
client: HaikuRAG client for document operations
config: Application configuration (uses global config if None)
search_filter: Optional SQL WHERE clause to filter documents
Returns:
ResearchReport with the final synthesis
"""
config = config or get_config()
console = Console()
# Build interactive graph
graph = build_research_graph(config=config, include_plan=True, interactive=True)
# Create async queue for human input
human_input_queue: asyncio.Queue[HumanDecision] = asyncio.Queue()
# Create emitter
emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter()
# Create deps with queue
deps = ResearchDeps(
client=client,
agui_emitter=emitter,
human_input_queue=human_input_queue,
interactive=True,
)
# Create initial state
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=config)
state.search_filter = search_filter
# Start the run
emitter.start_run(state)
# Run graph in background task
async def run_graph() -> ResearchReport:
try:
result = await graph.run(state=state, deps=deps)
emitter.finish_run(result)
return result
except Exception as e:
emitter.error(e)
raise
graph_task = asyncio.create_task(run_graph())
# Process events and handle human decision points
try:
async for event in emitter:
event_type = event.get("type")
if event_type == "STEP_STARTED":
step_name = event.get("stepName", "")
if step_name == "plan":
console.print("[dim]Planning research...[/dim]")
elif step_name.startswith("search:"):
query = step_name.replace("search: ", "")
console.print(f"[dim]Searching: {query}[/dim]")
elif step_name == "synthesize":
console.print("[dim]Synthesizing report...[/dim]")
elif event_type == "STATE_SNAPSHOT" or event_type == "STATE_DELTA":
# State updated, could show progress
pass
elif event_type == "TOOL_CALL_START":
tool_name = event.get("toolCallName")
if tool_name == "human_decision":
# Will get args in next event
pass
elif event_type == "TOOL_CALL_ARGS":
delta = event.get("delta", "{}")
args = json.loads(delta) if isinstance(delta, str) else delta
original_question = args.get("original_question", "")
sub_questions = list(args.get("sub_questions", []))
qa_responses = args.get("qa_responses", [])
iterations = args.get("iterations", 0)
# Loop for modifications until user wants to proceed
while True:
# Show research state
console.print()
console.print(
Panel(
f"[bold]{original_question}[/bold]",
title="Research Question",
border_style="blue",
)
)
# Show collected answers
if qa_responses:
answers_text = []
for i, qa in enumerate(qa_responses, 1):
conf = (
f"{qa['confidence']:.0%}"
if qa.get("confidence")
else "N/A"
)
answer_preview = (
qa["answer"][:200] + "..."
if len(qa["answer"]) > 200
else qa["answer"]
)
answers_text.append(
f"[cyan]{i}. {qa['query']}[/cyan]\n"
f" [dim]Confidence: {conf} | Citations: {qa.get('citations_count', 0)}[/dim]\n"
f" {answer_preview}"
)
console.print(
Panel(
"\n\n".join(answers_text),
title=f"Answers Collected ({len(qa_responses)})",
border_style="green",
)
)
# Show pending questions
if sub_questions:
console.print(
Panel(
"\n".join(
f"{i + 1}. {q}" for i, q in enumerate(sub_questions)
),
title="Pending Questions to Search",
border_style="cyan",
)
)
else:
console.print("[dim]No pending questions.[/dim]")
if iterations > 0:
console.print(f"[dim]Iteration: {iterations}[/dim]")
# Prompt user with context-aware hints
console.print()
hints = []
if sub_questions:
hints.append("search questions")
hints.append("modify questions")
if qa_responses:
hints.append("generate report")
hint_text = f" [dim]({', '.join(hints)})[/dim]" if hints else ""
user_input = Prompt.ask(
f"[bold]What would you like to do?[/bold]{hint_text}"
)
# Chat with research assistant
console.print("[dim]Thinking...[/dim]")
decision = await interpret_user_decision(
user_message=user_input,
sub_questions=sub_questions,
qa_responses=qa_responses,
config=config,
)
# Handle modifications and chat locally, continue loop
if decision.action == "chat":
if decision.message:
console.print(
f"\n[bold cyan]Assistant:[/bold cyan] {decision.message}"
)
continue
elif decision.action == "add_questions" and decision.questions:
sub_questions.extend(decision.questions)
console.print(
f"[green]Added {len(decision.questions)} question(s)[/green]"
)
continue
elif decision.action == "modify_questions" and decision.questions:
sub_questions = list(decision.questions)
console.print(
f"[green]Replaced with {len(decision.questions)} question(s)[/green]"
)
continue
# User wants to proceed - send final decision
action_display = {
"search": "Searching questions",
"synthesize": "Generating report",
}
console.print(
f"[dim]→ {action_display.get(decision.action, decision.action)}[/dim]"
)
# Include any accumulated question changes
if decision.action == "search":
decision = HumanDecision(
action="modify_questions", questions=sub_questions
)
await human_input_queue.put(decision)
break
elif event_type == "TEXT_MESSAGE_CHUNK":
# Log message from graph
message = event.get("delta", "")
if message:
console.print(f"[dim]{message}[/dim]")
elif event_type == "RUN_FINISHED":
break
elif event_type == "RUN_ERROR":
error_msg = event.get("message", "Unknown error")
console.print(f"[red]Error: {error_msg}[/red]")
break
# Wait for graph to complete
report = await graph_task
return report
except Exception as e:
graph_task.cancel()
raise e
finally:
await emitter.close()
async def run_chat_loop(
client: HaikuRAG,
config: AppConfig | None = None,
search_filter: str | None = None,
question: str | None = None,
) -> None:
"""Run an interactive chat loop for research.
Args:
client: HaikuRAG client for document operations
config: Application configuration (uses global config if None)
search_filter: Optional SQL WHERE clause to filter documents
question: Optional initial research question (skips initial chat if provided)
"""
config = config or get_config()
console = Console()
console.print(
Panel(
"[bold cyan]Interactive Research Mode[/bold cyan]\n\n"
"Chat with me or tell me what you'd like to research.\n"
"Type [green]exit[/green] or [green]quit[/green] to end the session.",
title="haiku.rag Research Assistant",
border_style="cyan",
)
)
while True:
try:
# Use provided question or get one through conversation
if question:
research_question = question
console.print(f"[dim]Starting research: {research_question}[/dim]")
question = None # Clear so subsequent loops go through chat
else:
# Initial conversation loop - chat until user wants to research
research_question = None
while research_question is None:
user_input = Prompt.ask("\n[bold blue]You[/bold blue]")
if not user_input.strip():
continue
if user_input.lower().strip() in ("exit", "quit", "q"):
console.print("[dim]Goodbye![/dim]")
return
console.print("[dim]Thinking...[/dim]")
decision = await initial_chat(user_input, config)
if decision.action == "research" and decision.research_question:
research_question = decision.research_question
console.print(
f"[dim]Starting research: {research_question}[/dim]"
)
elif decision.action == "chat" and decision.message:
console.print(
f"\n[bold cyan]Assistant:[/bold cyan] {decision.message}"
)
else:
# Fallback - treat as research question
research_question = user_input
console.print()
report = await run_interactive_research(
question=research_question,
client=client,
config=config,
search_filter=search_filter,
)
# Display final report
console.print()
console.print(
Panel(
Markdown(f"## {report.title}\n\n{report.executive_summary}"),
title="Research Report",
border_style="green",
)
)
if report.main_findings:
findings = "\n".join(f"- {f}" for f in report.main_findings[:5])
console.print(Markdown(f"**Key Findings:**\n{findings}"))
if report.conclusions:
conclusions = "\n".join(f"- {c}" for c in report.conclusions[:3])
console.print(Markdown(f"**Conclusions:**\n{conclusions}"))
console.print(Markdown(f"**Sources:** {report.sources_summary}"))
except KeyboardInterrupt:
console.print("\n[dim]Goodbye![/dim]")
return
except Exception as e:
console.print(f"[red]Error: {e}[/red]")
def interactive_research(
client: HaikuRAG,
config: AppConfig | None = None,
search_filter: str | None = None,
question: str | None = None,
) -> None:
"""Entry point for interactive research mode.
Args:
client: HaikuRAG client for document operations
config: Application configuration (uses global config if None)
search_filter: Optional SQL WHERE clause to filter documents
question: Optional initial research question (skips initial chat if provided)
"""
asyncio.run(run_chat_loop(client, config, search_filter, question))

View file

@ -80,8 +80,8 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
self._thread_id = self._generate_thread_id(state_json) self._thread_id = self._generate_thread_id(state_json)
# RunStarted (state snapshot follows immediately with full state) # RunStarted (state snapshot follows immediately with full state)
self._emit(emit_run_started(self._thread_id, self._run_id)) self.emit(emit_run_started(self._thread_id, self._run_id))
self._emit(emit_state_snapshot(initial_state)) self.emit(emit_state_snapshot(initial_state))
# Store a deep copy to detect future changes # Store a deep copy to detect future changes
self._last_state = initial_state.model_copy(deep=True) self._last_state = initial_state.model_copy(deep=True)
@ -92,12 +92,12 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
step_name: Name of the step being started step_name: Name of the step being started
""" """
self._current_step = step_name self._current_step = step_name
self._emit(emit_step_started(step_name)) self.emit(emit_step_started(step_name))
def finish_step(self) -> None: def finish_step(self) -> None:
"""Emit StepFinished event for the current step.""" """Emit StepFinished event for the current step."""
if self._current_step: if self._current_step:
self._emit(emit_step_finished(self._current_step)) self.emit(emit_step_finished(self._current_step))
self._current_step = None self._current_step = None
def log(self, message: str, role: str = "assistant") -> None: def log(self, message: str, role: str = "assistant") -> None:
@ -107,7 +107,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
message: The message content message: The message content
role: The role of the sender (default: assistant) role: The role of the sender (default: assistant)
""" """
self._emit(emit_text_message(message, role)) self.emit(emit_text_message(message, role))
def update_state(self, new_state: StateT) -> None: def update_state(self, new_state: StateT) -> None:
"""Emit StateDelta or StateSnapshot for state change. """Emit StateDelta or StateSnapshot for state change.
@ -117,10 +117,10 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
""" """
if self._use_deltas and self._last_state is not None: if self._use_deltas and self._last_state is not None:
# Emit delta for incremental updates # Emit delta for incremental updates
self._emit(emit_state_delta(self._last_state, new_state)) self.emit(emit_state_delta(self._last_state, new_state))
else: else:
# Emit full snapshot for initial state or when deltas disabled # Emit full snapshot for initial state or when deltas disabled
self._emit(emit_state_snapshot(new_state)) self.emit(emit_state_snapshot(new_state))
# Store a deep copy to detect future changes # Store a deep copy to detect future changes
self._last_state = new_state.model_copy(deep=True) self._last_state = new_state.model_copy(deep=True)
@ -139,7 +139,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
""" """
if message_id is None: if message_id is None:
message_id = str(uuid4()) message_id = str(uuid4())
self._emit(emit_activity(message_id, activity_type, content)) self.emit(emit_activity(message_id, activity_type, content))
def finish_run(self, result: ResultT) -> None: def finish_run(self, result: ResultT) -> None:
"""Emit RunFinished event. """Emit RunFinished event.
@ -147,7 +147,7 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
Args: Args:
result: The final result from the graph result: The final result from the graph
""" """
self._emit(emit_run_finished(self._thread_id, self._run_id, result)) self.emit(emit_run_finished(self._thread_id, self._run_id, result))
def error(self, error: Exception, code: str | None = None) -> None: def error(self, error: Exception, code: str | None = None) -> None:
"""Emit RunError event. """Emit RunError event.
@ -156,9 +156,9 @@ class AGUIEmitter[StateT: BaseModel, ResultT]:
error: The exception that occurred error: The exception that occurred
code: Optional error code code: Optional error code
""" """
self._emit(emit_run_error(str(error), code)) self.emit(emit_run_error(str(error), code))
def _emit(self, event: AGUIEvent) -> None: def emit(self, event: AGUIEvent) -> None:
"""Put event in queue. """Put event in queue.
Args: Args:

View file

@ -252,3 +252,62 @@ def emit_activity_delta(
"activityType": activity_type, "activityType": activity_type,
"patch": patch, "patch": patch,
} }
def emit_tool_call_start(
tool_call_id: str,
tool_name: str,
parent_message_id: str | None = None,
) -> dict[str, Any]:
"""Create a ToolCallStart event.
Args:
tool_call_id: Unique identifier for this tool call
tool_name: Name of the tool being called
parent_message_id: Optional parent message ID
Returns:
ToolCallStart event dict
"""
event: dict[str, Any] = {
"type": "TOOL_CALL_START",
"toolCallId": tool_call_id,
"toolCallName": tool_name,
}
if parent_message_id:
event["parentMessageId"] = parent_message_id
return event
def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> dict[str, Any]:
"""Create a ToolCallArgs event.
Args:
tool_call_id: Identifier for the tool call
args: Tool arguments
Returns:
ToolCallArgs event dict
"""
import json
return {
"type": "TOOL_CALL_ARGS",
"toolCallId": tool_call_id,
"delta": json.dumps(args),
}
def emit_tool_call_end(tool_call_id: str) -> dict[str, Any]:
"""Create a ToolCallEnd event.
Args:
tool_call_id: Identifier for the tool call being completed
Returns:
ToolCallEnd event dict
"""
return {
"type": "TOOL_CALL_END",
"toolCallId": tool_call_id,
}

View file

@ -166,7 +166,10 @@ def create_agui_server( # pragma: no cover
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState from haiku.rag.graph.research.state import (
ResearchDeps,
ResearchState,
)
# Store client reference for proper lifecycle management # Store client reference for proper lifecycle management
_client_cache: dict[str, HaikuRAG] = {} _client_cache: dict[str, HaikuRAG] = {}

View file

@ -1,4 +1,6 @@
import asyncio import asyncio
from typing import Literal
from uuid import uuid4
from pydantic_ai import Agent, RunContext, format_as_xml from pydantic_ai import Agent, RunContext, format_as_xml
from pydantic_ai.output import ToolOutput from pydantic_ai.output import ToolOutput
@ -7,6 +9,13 @@ from pydantic_graph.beta.join import reduce_list_append
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.events import (
emit_text_message_end,
emit_text_message_start,
emit_tool_call_args,
emit_tool_call_end,
emit_tool_call_start,
)
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import ( from haiku.rag.graph.research.models import (
EvaluationResult, EvaluationResult,
@ -54,12 +63,14 @@ def format_context_for_prompt(context: ResearchContext) -> str:
def build_research_graph( def build_research_graph(
config: AppConfig = Config, config: AppConfig = Config,
include_plan: bool = True, include_plan: bool = True,
interactive: bool = False,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: ) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the Research graph. """Build the Research graph.
Args: Args:
config: AppConfig object (uses config.research for provider, model, and graph parameters) config: AppConfig object (uses config.research for provider, model, and graph parameters)
include_plan: Whether to include the planning step (False for execute-only mode) include_plan: Whether to include the planning step (False for execute-only mode)
interactive: Whether to include human decision nodes for HIL
Returns: Returns:
Configured Research graph Configured Research graph
@ -240,7 +251,7 @@ def build_research_graph(
@g.step @g.step
async def get_batch( async def get_batch(
ctx: StepContext[ResearchState, ResearchDeps, None | bool], ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
) -> list[str] | None: ) -> list[str] | None:
"""Get all remaining questions for this iteration.""" """Get all remaining questions for this iteration."""
state = ctx.state state = ctx.state
@ -302,9 +313,16 @@ def build_research_graph(
state.last_eval = output state.last_eval = output
state.iterations += 1 state.iterations += 1
# Get already-answered questions to avoid duplicates
answered_queries = {qa.query.lower() for qa in state.context.qa_responses}
for new_q in output.new_questions: for new_q in output.new_questions:
if new_q not in state.context.sub_questions: # Skip if already in pending or already answered
state.context.sub_questions.append(new_q) if new_q in state.context.sub_questions:
continue
if new_q.lower() in answered_queries:
continue
state.context.sub_questions.append(new_q)
if deps.agui_emitter: if deps.agui_emitter:
deps.agui_emitter.update_state(state) deps.agui_emitter.update_state(state)
@ -329,9 +347,82 @@ def build_research_graph(
if deps.agui_emitter: if deps.agui_emitter:
deps.agui_emitter.finish_step() deps.agui_emitter.finish_step()
@g.step
async def human_decide(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer] | None | bool],
) -> Literal["search", "synthesize"]:
"""Wait for human decision on whether to continue searching or synthesize."""
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("human_decide")
deps.agui_emitter.update_state(state)
try:
# Emit tool call for human input wrapped in a message context
# This makes the tool call appear as if emitted by the LLM
message_id = str(uuid4())
tool_call_id = str(uuid4())
if deps.agui_emitter:
# Start an assistant message to contain the tool call
deps.agui_emitter.emit(emit_text_message_start(message_id))
# Emit tool call with parent message reference
deps.agui_emitter.emit(
emit_tool_call_start(tool_call_id, "human_decision", message_id)
)
# Include full state for display
qa_responses = [
{
"query": qa.query,
"answer": qa.answer,
"confidence": qa.confidence,
"citations_count": len(qa.citations),
}
for qa in state.context.qa_responses
]
deps.agui_emitter.emit(
emit_tool_call_args(
tool_call_id,
{
"original_question": state.context.original_question,
"sub_questions": list(state.context.sub_questions),
"qa_responses": qa_responses,
"iterations": state.iterations,
},
)
)
deps.agui_emitter.emit(emit_tool_call_end(tool_call_id))
# End the message after tool call
deps.agui_emitter.emit(emit_text_message_end(message_id))
# Wait for human input
if deps.human_input_queue is None:
raise RuntimeError("human_input_queue is required for interactive mode")
decision = await deps.human_input_queue.get()
# Process decision
if decision.action == "modify_questions" and decision.questions:
state.context.sub_questions = list(decision.questions)
elif decision.action == "add_questions" and decision.questions:
state.context.sub_questions.extend(decision.questions)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
if decision.action in ("search", "modify_questions", "add_questions"):
return "search"
else:
return "synthesize"
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step @g.step
async def synthesize( async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool], ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
) -> ResearchReport: ) -> ResearchReport:
"""Generate final research report.""" """Generate final research report."""
state = ctx.state state = ctx.state
@ -375,39 +466,76 @@ def build_research_graph(
initial_factory=list[SearchAnswer], initial_factory=list[SearchAnswer],
) )
if include_plan: if interactive:
# Interactive mode: human decides after plan and after evaluation
if include_plan:
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(human_decide),
)
else:
g.add(g.edge_from(g.start_node).to(human_decide))
g.add( g.add(
g.edge_from(g.start_node).to(plan), g.edge_from(human_decide).to(
g.edge_from(plan).to(get_batch), g.decision()
.branch(
g.match(str, matches=lambda x: x == "search")
.label("Search")
.to(get_batch)
)
.branch(
g.match(str, matches=lambda x: x == "synthesize")
.label("Synthesize")
.to(synthesize)
)
),
g.edge_from(get_batch).to(
g.decision()
.branch(g.match(list).label("Has questions").map().to(search_one))
.branch(g.match(type(None)).label("No questions").to(human_decide))
),
g.edge_from(search_one).to(collect_answers),
# After search, evaluate to suggest new questions, then human decides
g.edge_from(collect_answers).to(decide),
g.edge_from(decide).to(human_decide),
g.edge_from(synthesize).to(g.end_node),
) )
else: else:
g.add(g.edge_from(g.start_node).to(get_batch)) # Non-interactive mode: automatic decision based on confidence/iterations
if include_plan:
g.add( g.add(
g.edge_from(get_batch).to( g.edge_from(g.start_node).to(plan),
g.decision() g.edge_from(plan).to(get_batch),
.branch(g.match(list).label("Has questions").map().to(search_one))
.branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(decide),
)
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x)
.label("Continue research")
.to(get_batch)
) )
.branch( else:
g.match(bool, matches=lambda x: not x) g.add(g.edge_from(g.start_node).to(get_batch))
.label("Done researching")
.to(synthesize) g.add(
) g.edge_from(get_batch).to(
), g.decision()
g.edge_from(synthesize).to(g.end_node), .branch(g.match(list).label("Has questions").map().to(search_one))
) .branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(decide),
)
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x)
.label("Continue research")
.to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done researching")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build() return g.build()

View file

@ -1,6 +1,6 @@
import asyncio import asyncio
from dataclasses import dataclass from dataclasses import dataclass
from typing import TYPE_CHECKING from typing import TYPE_CHECKING, Literal
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@ -13,6 +13,17 @@ if TYPE_CHECKING:
from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.agui.emitter import AGUIEmitter
class HumanDecision(BaseModel):
"""Human decision input for interactive research."""
action: Literal[
"search", "synthesize", "modify_questions", "add_questions", "chat", "research"
]
questions: list[str] | None = None
message: str | None = None
research_question: str | None = None
@dataclass @dataclass
class ResearchDeps: class ResearchDeps:
"""Dependencies for research graph execution.""" """Dependencies for research graph execution."""
@ -20,6 +31,8 @@ class ResearchDeps:
client: HaikuRAG client: HaikuRAG
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
semaphore: asyncio.Semaphore | None = None semaphore: asyncio.Semaphore | None = None
human_input_queue: asyncio.Queue[HumanDecision] | None = None
interactive: bool = False
def emit_log(self, message: str, state: "ResearchState | None" = None) -> None: def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
"""Emit a log message through AG-UI events.""" """Emit a log message through AG-UI events."""

View file

@ -11,6 +11,9 @@ from haiku.rag.graph.agui.events import (
emit_step_finished, emit_step_finished,
emit_step_started, emit_step_started,
emit_text_message, emit_text_message,
emit_tool_call_args,
emit_tool_call_end,
emit_tool_call_start,
) )
@ -135,6 +138,46 @@ def test_emit_activity():
assert event["content"] == {"message": "Working on task"} assert event["content"] == {"message": "Working on task"}
def test_emit_tool_call_start():
"""Test TOOL_CALL_START event creation."""
event = emit_tool_call_start("call-1", "search_documents")
assert event["type"] == "TOOL_CALL_START"
assert event["toolCallId"] == "call-1"
assert event["toolCallName"] == "search_documents"
assert "parentMessageId" not in event
def test_emit_tool_call_start_with_parent():
"""Test TOOL_CALL_START event with parent message ID."""
event = emit_tool_call_start("call-1", "search", parent_message_id="msg-1")
assert event["type"] == "TOOL_CALL_START"
assert event["toolCallId"] == "call-1"
assert event["toolCallName"] == "search"
assert event["parentMessageId"] == "msg-1"
def test_emit_tool_call_args():
"""Test TOOL_CALL_ARGS event creation."""
import json
args = {"query": "test query", "limit": 10}
event = emit_tool_call_args("call-1", args)
assert event["type"] == "TOOL_CALL_ARGS"
assert event["toolCallId"] == "call-1"
assert event["delta"] == json.dumps(args)
def test_emit_tool_call_end():
"""Test TOOL_CALL_END event creation."""
event = emit_tool_call_end("call-1")
assert event["type"] == "TOOL_CALL_END"
assert event["toolCallId"] == "call-1"
def test_event_structure_consistency(): def test_event_structure_consistency():
"""Test that all events have consistent structure.""" """Test that all events have consistent structure."""
events = [ events = [
@ -146,6 +189,9 @@ def test_event_structure_consistency():
emit_text_message("text"), emit_text_message("text"),
emit_state_snapshot(TestState(value=1)), emit_state_snapshot(TestState(value=1)),
emit_activity("m1", "type", {"content": "value"}), emit_activity("m1", "type", {"content": "value"}),
emit_tool_call_start("c1", "tool"),
emit_tool_call_args("c1", {"arg": "value"}),
emit_tool_call_end("c1"),
] ]
for event in events: for event in events:

View file

@ -1,3 +1,5 @@
import asyncio
import pytest import pytest
from pydantic_ai.models.test import TestModel from pydantic_ai.models.test import TestModel
@ -5,7 +7,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.graph.agui.stream import stream_graph from haiku.rag.graph.agui.stream import stream_graph
from haiku.rag.graph.research.dependencies import ResearchContext from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState
@pytest.mark.asyncio @pytest.mark.asyncio
@ -61,3 +63,83 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
assert "STEP_STARTED" in event_types assert "STEP_STARTED" in event_types
client.close() client.close()
@pytest.mark.asyncio
async def test_interactive_graph_with_human_decision(monkeypatch, temp_db_path):
"""Test interactive research graph pauses and resumes with human decisions."""
# Mock get_model to return TestModel
def test_model_factory(_provider, _model, _config=None):
return TestModel()
monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
# Build interactive graph
graph = build_research_graph(interactive=True)
state = ResearchState(
context=ResearchContext(original_question="What is haiku.rag?"),
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=2,
)
# Create human input queue
human_input_queue: asyncio.Queue[HumanDecision] = asyncio.Queue()
client = HaikuRAG(temp_db_path, create=True)
deps = ResearchDeps(
client=client,
human_input_queue=human_input_queue,
interactive=True,
)
events = []
tool_call_received = asyncio.Event()
result = None
async def run_graph():
nonlocal result
async for event in stream_graph(graph, state, deps):
events.append(event)
if event["type"] == "TOOL_CALL_START":
tool_name = event.get("toolCallName")
if tool_name == "human_decision":
tool_call_received.set()
elif event["type"] == "RUN_FINISHED":
result = event["result"]
elif event["type"] == "RUN_ERROR":
pytest.fail(f"Graph execution failed: {event['message']}")
async def send_decisions():
# Wait for first tool call (after planning)
await asyncio.wait_for(tool_call_received.wait(), timeout=30)
tool_call_received.clear()
# Send search decision
await human_input_queue.put(HumanDecision(action="search"))
# Wait for second tool call (after search cycle)
await asyncio.wait_for(tool_call_received.wait(), timeout=30)
# Send synthesize decision
await human_input_queue.put(HumanDecision(action="synthesize"))
# Run graph and decision sender concurrently
await asyncio.gather(run_graph(), send_decisions())
# Verify result
assert result is not None, (
f"No result. Events collected: {[e['type'] for e in events]}"
)
assert isinstance(result, dict)
assert "title" in result
# Verify human_decision tool calls were emitted
event_types = [e["type"] for e in events]
assert "TOOL_CALL_START" in event_types
assert "TOOL_CALL_END" in event_types
client.close()