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

Remove custom AG-UI infrastructure in favor of pydantic-ai native support

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-09 18:58:48 +02:00
parent baefbe9416
commit bb61e8fab6
No known key found for this signature in database
52 changed files with 129 additions and 19717 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
CHANGELOG.md
View file

@ -1,6 +1,18 @@
# Changelog
## [Unreleased]
### Removed
- **BREAKING: Custom AG-UI Infrastructure**: Removed custom AG-UI event handling in favor of pydantic-ai's native AG-UI support
- Deleted `haiku.rag.graph.agui` module (`AGUIEmitter`, `AGUIConsoleRenderer`, `stream_graph()`, `create_agui_server()`)
- Removed `--agui` flag from `serve` command
- Removed `--verbose` flags from `ask` and `research` commands
- Removed `--interactive` flag from `research` command
- Removed `AGUIConfig` from configuration
- Deleted `cli_chat.py` interactive chat module
- Research graph now uses `graph.run()` directly instead of `stream_graph()`
- For AG-UI streaming, use pydantic-ai's native `AGUIAdapter` with `ToolReturn` and `StateSnapshotEvent` (see `app/backend/` for example)
## [0.25.0] - 2026-01-12
### Fixed

15
examples/ag-ui-research/.env.example
View file

@ -1,15 +0,0 @@
# Path to the LanceDB database on your HOST machine
# This will be mounted into the Docker container at /app/data/haiku.rag.lancedb
# Must be an absolute path to an existing database created with haiku-rag
DB_PATH=/absolute/path/to/your/haiku.rag.lancedb
# Ollama API base URL (if using Ollama for local models)
# If running Ollama on your host machine, use your machine's IP address
# that the Docker container can reach (not localhost)
OLLAMA_BASE_URL=http://host.docker.internal:11434
# API keys (set as needed for your QA provider)
# OPENAI_API_KEY=your-key-here
# ANTHROPIC_API_KEY=your-key-here
# VOYAGE_API_KEY=your-key-here
# CO_API_KEY=your-key-here

3
examples/ag-ui-research/.gitignore vendored
View file

@ -1,3 +0,0 @@
haiku.rag.yaml
.env
data/

154
examples/ag-ui-research/README.md
View file

@ -1,154 +0,0 @@
# 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 with human-in-the-loop control.
[Watch demo video](https://vimeo.com/1128874386)
## Features
- **Human-in-the-loop research**: Review and modify questions at decision points, then continue searching or generate report
- **Multi-iteration research graph**: Automated question decomposition and parallel search
- **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
## Quick Start
### Prerequisites
- Docker and Docker Compose
- A haiku.rag database with indexed documents
- Ollama (or configure another LLM provider)
### Setup
1. **Prepare your knowledge base**
**Option A: Create a new database**
```bash
haiku-rag init --db data/haiku_rag.lancedb
haiku-rag add-src document.pdf --db data/haiku_rag.lancedb
```
**Option B: Use an existing database**
Set the `DB_PATH` environment variable to point to your existing haiku.rag database:
```bash
# In .env file
DB_PATH=/path/to/your/existing/haiku_rag.lancedb
```
Or export it before running docker compose:
```bash
export DB_PATH=/path/to/your/existing/haiku_rag.lancedb
docker compose up --build
```
The database will be mounted as read-write, so the research assistant can access all documents in your existing knowledge base.
2. **Configure haiku.rag**
```bash
cp haiku.rag.yaml.example haiku.rag.yaml
# Edit haiku.rag.yaml to customize provider/model
```
See [haiku.rag configuration](https://ggozad.github.io/haiku.rag/configuration/) for details.
3. **Set API keys** (if using non-Ollama providers)
```bash
cp .env.example .env
# Edit .env to set your API keys
OPENAI_API_KEY=your-key-here
ANTHROPIC_API_KEY=your-key-here
DB_PATH=/path/to/your/existing/haiku_rag.lancedb # If using an existing db.
```
4. **Pull the base image**
```bash
docker pull ghcr.io/ggozad/haiku.rag-slim:latest
```
5. **Start the application**
```bash
docker compose up --build
```
6. **Access the interface**
- Frontend: http://localhost:3000
- Backend health: http://localhost:8000/health
## How It Works
1. **Ask a question**: Type your research question in the chat
2. **Plan phase**: The research graph decomposes your question into targeted sub-questions
3. **Decision point**: Review the proposed questions in the right panel
- Add new questions using the input field
- Remove questions you don't need
- Click **Search** to execute searches for pending questions
- Click **Generate Report** to skip to synthesis (when you have enough answers)
4. **Research iterations**: After each search cycle, you return to a decision point where you can:
- Review collected answers
- Add follow-up questions based on findings
- Continue searching or generate the final report
5. **Synthesis**: Generates a comprehensive research report with:
- Executive summary
- Main findings with supporting evidence
- Conclusions and recommendations
- Source citations
## Architecture
### Agent + Graph Pattern
This example demonstrates the **agent+graph** architecture with AG-UI client-side tool calls:
1. **Conversational Agent** (`agent.py`):
- Pydantic AI agent handles user conversations
- Decides when to invoke the research tool based on user intent
- Responds directly to greetings/casual chat without tools
2. **Interactive Research Graph** (haiku.rag):
- Multi-step research workflow invoked by the agent's tool
- At decision points, emits AG-UI `TOOL_CALL_START/ARGS/END` events for `human_decision`
- Waits for tool result via async queue before continuing
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
- Events from both flow through a single stream to the frontend
- `STATE_DELTA` events sync research state to frontend in real-time
### Components
- **Backend** (Python):
- Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
- `agent.py`: Pydantic AI agent with `run_research` tool, manages `ActiveResearch` registry
- `main.py`: Custom AG-UI streaming endpoint, extracts tool results from messages
- Real-time event forwarding from emitter to SSE stream
- **Frontend** (Next.js/React):
- 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
- Tool results sent directly to backend endpoint
## Configuration
Configuration is done through `haiku.rag.yaml` (see `haiku.rag.yaml.example`):
- `research.provider`: LLM provider (default: `ollama`)
- `research.model`: Model name (default: `gpt-oss:latest`)
- `research.max_iterations`: Maximum research iterations (default: `3`)
- `research.confidence_threshold`: Confidence threshold for completion (default: `0.8`)
- `research.max_concurrency`: Parallel sub-question processing (default: `1`)
- `providers.ollama.base_url`: Ollama endpoint (default: `http://host.docker.internal:11434`)
Environment variables (see `.env.example`):
- `DB_PATH`: Path to haiku.rag database (default: `haiku_rag.lancedb`)
- `OPENAI_API_KEY`, `ANTHROPIC_API_KEY`: API keys for cloud providers
For full configuration options, see [haiku.rag configuration docs](https://ggozad.github.io/haiku.rag/configuration/).

12
examples/ag-ui-research/backend/.gitignore vendored
View file

@ -1,12 +0,0 @@
__pycache__/
*.py[cod]
*$py.class
*.so
.Python
env/
venv/
.env
.venv
*.egg-info/
dist/
build/

16
examples/ag-ui-research/backend/Dockerfile
View file

@ -1,16 +0,0 @@
FROM ghcr.io/ggozad/haiku.rag-slim:latest
WORKDIR /app
# Copy backend application files
COPY main.py agent.py ./
# Install additional dependencies for the example
RUN pip install --no-cache-dir \
starlette>=0.50.0 \
uvicorn[standard]>=0.40.0 \
python-dotenv>=1.2.1
EXPOSE 8000
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000", "--reload"]

25
examples/ag-ui-research/backend/README.md
View file

@ -1,25 +0,0 @@
# Haiku.rag Research Assistant Backend
Starlette backend for the haiku.rag interactive research assistant, using the research graph with AG-UI protocol support.
## Setup
```bash
uv sync
uv run python main.py
```
The server starts on `http://localhost:8000` and uses [haiku.rag configuration](https://ggozad.github.io/haiku.rag/configuration/).
## Architecture
The backend uses `create_agui_server()` from `haiku.rag.graph.agui.server` which provides:
- **Research graph execution**: Multi-iteration research workflow
- **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates
- **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations
## Endpoints
- `GET /health` - Health check with configuration info
- `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol)

151
examples/ag-ui-research/backend/agent.py
View file

@ -1,151 +0,0 @@
"""Research assistant agent with graph integration."""
import asyncio
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING
from pydantic_ai import Agent, RunContext
from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import HumanDecision, ResearchDeps, ResearchState
from haiku.rag.utils import get_model
if TYPE_CHECKING:
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.research.models import ResearchReport
# Load config
config_path = Path("/app/haiku.rag.yaml")
Config = (
AppConfig.model_validate(load_yaml_config(config_path))
if config_path.exists()
else AppConfig()
)
@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
class AgentDeps:
"""Dependencies for research agent."""
client: HaikuRAG
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
search_filter: str | None = None
thread_id: str | None = None
research_result: "ResearchReport | None" = None
model = get_model(Config.research.model, Config)
agent = Agent(
model,
deps_type=AgentDeps,
system_prompt="""You are an advanced research assistant powered by haiku.rag.
CRITICAL RULES:
1. For greetings (hi, hello, hey, etc) or casual chat: respond directly WITHOUT using any tools
2. For questions about yourself or the system: respond directly WITHOUT using any tools
3. For substantive questions requiring information: ALWAYS use the run_research tool
4. NEVER answer substantive questions from your own knowledge - always use the tool
How to decide:
- "Hi" / "Hello" / "How are you?" -> Respond directly, NO tools
- "What can you do?" -> Respond directly, NO tools
- "How does X work in the codebase?" -> 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,
and generate a comprehensive report.
Be friendly and conversational in all responses.""",
)
@agent.tool
async def run_research(ctx: RunContext[AgentDeps], question: str) -> str:
"""Execute research graph on a substantive question.
Use for questions requiring knowledge base search.
DO NOT use for greetings or casual conversation.
"""
if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.log(f"Starting research on: {question}")
# 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)
state = ResearchState.from_config(context=context, config=Config)
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(
client=ctx.deps.client,
agui_emitter=ctx.deps.agui_emitter,
human_input_queue=queue,
interactive=True,
)
try:
result = await graph.run(state=state, deps=graph_deps)
if ctx.deps.agui_emitter:
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!
Question: {question}
Executive Summary: {result.executive_summary}
Main Findings:
{chr(10).join(f"- {finding}" for finding in result.main_findings[:3])}
Conclusions:
{chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])}
Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"}
Iterations completed: {state.iterations}
The full research report with all citations has been provided to the user.
"""
except Exception as e:
if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.log(f"Research error: {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]

343
examples/ag-ui-research/backend/main.py
View file

@ -1,343 +0,0 @@
import json
import logging
import os
from pathlib import Path
from agent import AgentDeps, _active_research, agent
from anyio import create_memory_object_stream, create_task_group
from anyio.streams.memory import MemoryObjectSendStream
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, StreamingResponse
from starlette.routing import Route
from haiku.rag.client import HaikuRAG
from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.agui.server import RunAgentInput, format_sse_event
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import ResearchReport
from haiku.rag.graph.research.state import HumanDecision, ResearchState
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
# Load config from mounted haiku.rag.yaml
config_path = Path("/app/haiku.rag.yaml")
if config_path.exists():
yaml_data = load_yaml_config(config_path)
Config = AppConfig.model_validate(yaml_data)
else:
# Fallback to default config
Config = AppConfig()
# Get DB path from environment
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
db_path = Path(db_path_str)
if not db_path.exists():
logger.error(f"Database not found at {db_path}")
logger.error("Run: haiku-rag add <path-to-documents>")
raise RuntimeError(f"Database not found: {db_path}")
logger.info(f"Initializing research assistant with database: {db_path}")
logger.info(
f"Research Provider: {Config.research.model.provider}, Model: {Config.research.model.name}"
)
# Store client reference for proper lifecycle management
_client_cache: dict[str, HaikuRAG] = {}
def get_client(effective_db_path: Path) -> HaikuRAG:
"""Get or create cached client."""
path_key = str(effective_db_path)
if path_key not in _client_cache:
_client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=Config)
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:
"""Agent streaming endpoint with research graph integration."""
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)
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 = ""
if input_data.messages:
user_message = input_data.messages[-1].get("content", "")
send_stream, receive_stream = create_memory_object_stream[str]()
async def run_agent_with_streaming(
send_stream: MemoryObjectSendStream[str],
) -> None:
"""Execute agent and forward emitter events to memory stream."""
async with send_stream:
try:
# Create shared emitter (use_deltas=True for CopilotKit compatibility)
emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter(
thread_id=input_data.thread_id,
run_id=input_data.run_id,
use_deltas=True,
)
# Get client
effective_db_path = input_data.config.get("db_path") or db_path
if isinstance(effective_db_path, str):
effective_db_path = Path(effective_db_path)
client = get_client(effective_db_path)
# Build search filter from document IDs (empty list = search all)
document_ids = input_data.state.get("documentFilter") or []
search_filter = None
if document_ids:
ids_str = ", ".join(f"'{id}'" for id in document_ids)
search_filter = f"id IN ({ids_str})"
# Create agent dependencies with shared emitter
agent_deps = AgentDeps(
client=client,
agui_emitter=emitter,
search_filter=search_filter,
thread_id=thread_id,
)
# Start run with empty initial state
emitter.start_run(
initial_state=ResearchState.from_config(
context=ResearchContext(original_question=""),
config=Config,
)
)
# Forward emitter events to stream
async def forward_events():
async for event in emitter:
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
# As CopilotKit does not handle ACTIVITY_SNAPSHOT events
if event_type == "ACTIVITY_SNAPSHOT":
activity_type = event.get("activityType", "")
content = event.get("content", {})
message = content.get("message", "")
# Emit STATE_DELTA to patch activity info into state
delta_event = {
"type": "STATE_DELTA",
"delta": [
{
"op": "add",
"path": "/current_activity",
"value": activity_type,
},
{
"op": "add",
"path": "/current_activity_message",
"value": message,
},
],
}
await send_stream.send(format_sse_event(delta_event))
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))
# Run agent and event forwarding concurrently
async with create_task_group() as tg:
tg.start_soon(forward_events)
result = await agent.run(user_message, deps=agent_deps)
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()
except Exception as e:
logger.exception("Error executing agent")
try:
await send_stream.send(
format_sse_event({"type": "error", "error": str(e)})
)
except Exception:
pass
async def event_generator():
"""Generate SSE events from memory stream."""
async with create_task_group() as tg:
tg.start_soon(run_agent_with_streaming, send_stream)
async with receive_stream:
async for event_str in receive_stream:
yield event_str
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
async def health_check(_: Request) -> JSONResponse:
"""Health check endpoint with configuration info."""
return JSONResponse(
{
"status": "healthy",
"agent_model": str(agent.model),
"research_provider": Config.research.model.provider,
"research_model": Config.research.model.name,
"db_path": str(db_path),
"db_exists": db_path.exists(),
}
)
async def list_documents(_: Request) -> JSONResponse:
"""List all documents in the database."""
client = get_client(db_path)
docs = await client.document_repository.list_all()
return JSONResponse(
{
"documents": [
{"id": doc.id, "title": doc.title, "uri": doc.uri} for doc in docs
]
}
)
async def visualize_chunk(request: Request) -> JSONResponse:
"""Return visual grounding images for a chunk as base64."""
import base64
from io import BytesIO
chunk_id = request.path_params["chunk_id"]
client = get_client(db_path)
# Get the chunk
chunk = await client.chunk_repository.get_by_id(chunk_id)
if not chunk:
return JSONResponse({"error": "Chunk not found"}, status_code=404)
# Get visualization images
images = await client.visualize_chunk(chunk)
if not images:
return JSONResponse({"images": [], "message": "No visual grounding available"})
# Convert PIL images to base64
base64_images = []
for img in images:
buffer = BytesIO()
img.save(buffer, format="PNG")
buffer.seek(0)
base64_images.append(base64.b64encode(buffer.read()).decode("utf-8"))
return JSONResponse(
{
"images": base64_images,
"chunk_id": chunk_id,
"document_uri": chunk.document_uri,
}
)
# Create Starlette app
app = Starlette(
routes=[
Route("/v1/research/stream", stream_research_agent, methods=["POST"]),
Route("/api/documents", list_documents, methods=["GET"]),
Route("/api/visualize/{chunk_id}", visualize_chunk, methods=["GET"]),
Route("/health", health_check, methods=["GET"]),
],
middleware=[
Middleware(
CORSMiddleware,
allow_origins=["http://localhost:3000", "http://frontend:3000"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
],
)
if __name__ == "__main__":
import uvicorn
uvicorn.run(
"main:app",
host="0.0.0.0",
port=8000,
reload=True,
)

26
examples/ag-ui-research/backend/pyproject.toml
View file

@ -1,26 +0,0 @@
[project]
name = "haiku-rag-research-assistant"
version = "0.1.0"
description = "Haiku.rag research assistant with AG-UI protocol support"
readme = "README.md"
requires-python = ">=3.13"
dependencies = [
"starlette>=0.50.0",
"uvicorn[standard]>=0.40.0",
"pydantic-ai-slim[ag-ui,openai]>=1.39.0",
"python-dotenv>=1.2.1",
"haiku.rag-slim[agui]>=0.25.0",
]
[dependency-groups]
dev = ["pyright>=1.1.407", "ruff>=0.14.10"]
[tool.hatch.metadata]
allow-direct-references = true
[tool.hatch.build.targets.wheel]
packages = ["."]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

2172
examples/ag-ui-research/backend/uv.lock
View file

File diff suppressed because it is too large Load diff

76
examples/ag-ui-research/docker-compose.yml
View file

@ -1,76 +0,0 @@
services:
docling-serve:
image: quay.io/docling-project/docling-serve:latest
container_name: ag-ui-docling-serve
ports:
- "5001:5001"
environment:
- DOCLING_SERVE_ENABLE_UI=1
networks:
- ag-ui-network
restart: unless-stopped
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:5001/health"]
interval: 30s
timeout: 10s
retries: 3
start_period: 40s
start_interval: 5s
backend:
build:
context: ./backend
dockerfile: Dockerfile
ports:
- "8000:8000"
environment:
- DB_PATH=/app/data/haiku.rag.lancedb
# API keys (set these in your shell or .env file)
- OPENAI_API_KEY=${OPENAI_API_KEY}
- ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
# Ollama connection (use value from .env)
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL}
# Prevent Python bytecode caching for development
- PYTHONDONTWRITEBYTECODE=1
# Use local haiku_rag_slim for development
- PYTHONPATH=/app/haiku_rag_slim
volumes:
- ${DB_PATH}:/app/data/haiku.rag.lancedb
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
# Mount local haiku_rag_slim for development
- ../../haiku_rag_slim:/app/haiku_rag_slim:ro
networks:
- ag-ui-network
extra_hosts:
- "host.docker.internal:host-gateway"
depends_on:
docling-serve:
condition: service_healthy
restart: unless-stopped
healthcheck:
test: ["CMD", "python", "-c", "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')"]
interval: 30s
timeout: 10s
retries: 3
frontend:
build:
context: ./frontend
dockerfile: Dockerfile
ports:
- "3000:3000"
environment:
- BACKEND_URL=http://backend:8000
volumes:
- ./frontend:/app
- /app/node_modules
- /app/.next
depends_on:
- backend
networks:
- ag-ui-network
restart: unless-stopped
networks:
ag-ui-network:
driver: bridge

7
examples/ag-ui-research/frontend/.dockerignore
View file

@ -1,7 +0,0 @@
node_modules
.next
.git
.gitignore
README.md
npm-debug.log
.env*.local

39
examples/ag-ui-research/frontend/.gitignore vendored
View file

@ -1,39 +0,0 @@
# dependencies
/node_modules
/.pnp
.pnp.*
.yarn/*
!.yarn/patches
!.yarn/plugins
!.yarn/releases
!.yarn/versions
# testing
/coverage
# next.js
/.next/
/out/
# production
/build
# misc
.DS_Store
*.pem
# debug
npm-debug.log*
yarn-debug.log*
yarn-error.log*
# env files
.env*.local
.env
# vercel
.vercel
# typescript
*.tsbuildinfo
next-env.d.ts

12
examples/ag-ui-research/frontend/Dockerfile
View file

@ -1,12 +0,0 @@
# Development Dockerfile for Next.js frontend
FROM node:22-alpine
WORKDIR /app
COPY package.json package-lock.json* ./
RUN npm install --legacy-peer-deps
COPY . .
EXPOSE 3000
# Run in development mode with hot reload
CMD ["npm", "run", "dev"]

33
examples/ag-ui-research/frontend/app/api/copilotkit/route.ts
View file

@ -1,33 +0,0 @@
import { HttpAgent } from "@ag-ui/client";
import {
CopilotRuntime,
copilotRuntimeNextJSAppRouterEndpoint,
ExperimentalEmptyAdapter,
} from "@copilotkit/runtime";
import type { NextRequest } from "next/server";
// Connect CopilotKit to PydanticAI via HttpAgent
// The HttpAgent creates a bridge between the Next.js frontend and the Python backend
// It communicates with the server created by agent.to_ag_ui()
const runtime = new CopilotRuntime({
agents: {
// "research_agent" maps to the agent name used in useCoAgent() on the frontend
research_agent: new HttpAgent({
url: `${process.env.BACKEND_URL || "http://backend:8000"}/v1/research/stream`,
}),
},
});
// Service adapter for multi-agent support (empty since we only have one agent)
const serviceAdapter = new ExperimentalEmptyAdapter();
// Next.js API route handler that proxies requests between frontend and backend
export async function POST(request: NextRequest) {
const { handleRequest } = copilotRuntimeNextJSAppRouterEndpoint({
runtime,
serviceAdapter,
endpoint: "/api/copilotkit",
});
return handleRequest(request);
}

24
examples/ag-ui-research/frontend/app/globals.css
View file

@ -1,24 +0,0 @@
* {
box-sizing: border-box;
padding: 0;
margin: 0;
}
html,
body {
max-width: 100vw;
overflow-x: hidden;
font-family:
system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto,
"Helvetica Neue", Arial, sans-serif;
}
body {
background: linear-gradient(to bottom, #f8f9fa, #e9ecef);
min-height: 100vh;
}
a {
color: inherit;
text-decoration: none;
}

20
examples/ag-ui-research/frontend/app/layout.tsx
View file

@ -1,20 +0,0 @@
import type { Metadata } from "next";
import "./globals.css";
export const metadata: Metadata = {
title: "Haiku.rag Research Assistant",
description:
"Interactive research powered by Haiku.rag, Pydantic AI, and AG-UI",
};
export default function RootLayout({
children,
}: Readonly<{
children: React.ReactNode;
}>) {
return (
<html lang="en">
<body>{children}</body>
</html>
);
}

9
examples/ag-ui-research/frontend/app/page.tsx
View file

@ -1,9 +0,0 @@
import Agent from "@/components/Agent";
export default function Home() {
return (
<main>
<Agent />
</main>
);
}

34
examples/ag-ui-research/frontend/biome.json
View file

@ -1,34 +0,0 @@
{
"$schema": "https://biomejs.dev/schemas/2.2.6/schema.json",
"vcs": {
"enabled": false,
"clientKind": "git",
"useIgnoreFile": false
},
"files": {
"ignoreUnknown": false
},
"formatter": {
"enabled": true,
"indentStyle": "tab"
},
"linter": {
"enabled": true,
"rules": {
"recommended": true
}
},
"javascript": {
"formatter": {
"quoteStyle": "double"
}
},
"assist": {
"enabled": true,
"actions": {
"source": {
"organizeImports": "on"
}
}
}
}

449
examples/ag-ui-research/frontend/components/Agent.tsx
View file

@ -1,449 +0,0 @@
"use client";
import {
CopilotKit,
useCoAgent,
useCopilotAction,
useCopilotContext,
} from "@copilotkit/react-core";
import { CopilotChat } from "@copilotkit/react-ui";
import "@copilotkit/react-ui/styles.css";
import { useState } from "react";
import DocumentSelector from "./DocumentSelector";
import StateDisplay from "./StateDisplay";
interface SearchAnswer {
query: string;
answer: string;
confidence: number;
cited_chunks: string[];
citations: {
document_id: string;
chunk_id: string;
document_uri: string;
document_title?: string;
page_numbers: number[];
headings?: string[];
content: string;
}[];
}
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;
sub_questions: string[];
qa_responses: SearchAnswer[];
}
type DecisionAction = "search" | "synthesize" | "modify_questions";
interface DecisionResult {
action: DecisionAction;
questions?: string[];
}
function DecisionUI({
args,
onResolve,
}: {
args: DecisionArgs;
onResolve: (result: DecisionResult) => void | Promise<void>;
}) {
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>
);
}
const BACKEND_URL =
process.env.NEXT_PUBLIC_BACKEND_URL || "http://localhost:8000";
function AgentContent() {
const { state, setState, running } = useCoAgent<ResearchState>({
name: "research_agent",
initialState: {
context: {
original_question: "",
sub_questions: [],
qa_responses: [],
},
iterations: 0,
max_iterations: 3,
confidence_threshold: 0.8,
max_concurrency: 1,
last_eval: null,
documentFilter: [],
},
});
const { threadId } = useCopilotContext();
const handleDocumentFilterChange = (ids: string[]) => {
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 (
<>
<style>{`
.chat-container {
width: 50%;
height: 100vh;
border-right: 1px solid #e2e8f0;
display: flex;
flex-direction: column;
}
.chat-container > * {
flex: 1;
min-height: 0;
}
`}</style>
<div style={{ display: "flex", height: "100vh" }}>
<div className="chat-container">
<CopilotChat
labels={{
title: "Research Assistant",
initial:
"Hello! I can help you conduct deep research on complex questions using the haiku.rag knowledge base. Ask me anything!",
}}
/>
</div>
<div
style={{
width: "50%",
height: "100vh",
overflow: "auto",
background: "#f7fafc",
}}
>
<div style={{ padding: "2rem" }}>
<header style={{ marginBottom: "2rem" }}>
<h1
style={{
fontSize: "2rem",
fontWeight: "bold",
marginBottom: "0.5rem",
color: "#1a202c",
}}
>
Research State
</h1>
<p
style={{
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
Live updates from the research agent
</p>
</header>
{!running && (
<div style={{ marginBottom: "1rem" }}>
<DocumentSelector
selected={state.documentFilter || []}
onChange={handleDocumentFilterChange}
/>
</div>
)}
<StateDisplay state={state} />
</div>
</div>
</div>
</>
);
}
export default function Agent() {
return (
<CopilotKit runtimeUrl="/api/copilotkit" agent="research_agent">
<AgentContent />
</CopilotKit>
);
}

359
examples/ag-ui-research/frontend/components/DocumentSelector.tsx
View file

@ -1,359 +0,0 @@
"use client";
import { useCallback, useEffect, useMemo, useState } from "react";
interface Document {
id: string;
title: string;
uri: string;
}
interface DocumentSelectorProps {
selected: string[];
onChange: (ids: string[]) => void;
disabled: boolean;
}
export default function DocumentSelector({
selected,
onChange,
disabled,
}: DocumentSelectorProps) {
const [documents, setDocuments] = useState<Document[]>([]);
const [loading, setLoading] = useState(true);
const [error, setError] = useState<string | null>(null);
const [expanded, setExpanded] = useState(true);
const [searchQuery, setSearchQuery] = useState("");
useEffect(() => {
const fetchDocuments = async () => {
try {
const response = await fetch(
`${process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"}/api/documents`,
);
if (!response.ok) {
throw new Error("Failed to fetch documents");
}
const data = await response.json();
const docs = data.documents || [];
setDocuments(docs);
// Select all documents by default if none are selected
if (selected.length === 0 && docs.length > 0) {
onChange(docs.map((d: Document) => d.id));
}
setError(null);
} catch (err) {
setError(err instanceof Error ? err.message : "Unknown error");
} finally {
setLoading(false);
}
};
fetchDocuments();
}, []);
const filteredDocuments = useMemo(() => {
if (!searchQuery.trim()) return documents;
const query = searchQuery.toLowerCase();
return documents.filter(
(doc) =>
(doc.title || "").toLowerCase().includes(query) ||
(doc.uri || "").toLowerCase().includes(query),
);
}, [documents, searchQuery]);
const handleToggle = useCallback(
(id: string) => {
if (disabled) return;
if (selected.includes(id)) {
onChange(selected.filter((s) => s !== id));
} else {
onChange([...selected, id]);
}
},
[selected, onChange, disabled],
);
const handleSelectAll = useCallback(() => {
if (disabled) return;
const filteredIds = filteredDocuments.map((d) => d.id);
const allFilteredSelected = filteredIds.every((id) =>
selected.includes(id),
);
if (allFilteredSelected) {
// Deselect all filtered documents
onChange(selected.filter((id) => !filteredIds.includes(id)));
} else {
// Select all filtered documents (add to existing selection)
const newSelection = [...new Set([...selected, ...filteredIds])];
onChange(newSelection);
}
}, [selected, filteredDocuments, onChange, disabled]);
const selectedCount = selected.length;
const totalCount = documents.length;
const filterActive = selectedCount > 0 && selectedCount < totalCount;
return (
<div
style={{
background: "white",
borderRadius: "8px",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
overflow: "hidden",
opacity: disabled ? 0.6 : 1,
transition: "opacity 0.2s",
}}
>
<button
type="button"
onClick={() => setExpanded(!expanded)}
style={{
width: "100%",
display: "flex",
justifyContent: "space-between",
alignItems: "center",
padding: "0.75rem",
background: filterActive ? "#ebf8ff" : "#edf2f7",
border: filterActive ? "1px solid #90cdf4" : "1px solid #e2e8f0",
borderRadius: expanded ? "8px 8px 0 0" : "8px",
cursor: "pointer",
fontSize: "0.875rem",
fontWeight: "600",
color: filterActive ? "#2b6cb0" : "#2d3748",
}}
>
<span>
Document Filter
{filterActive && ` (${selectedCount}/${totalCount})`}
{!filterActive && selectedCount === 0 && " (All)"}
</span>
<span>{expanded ? "▼" : "▶"}</span>
</button>
{expanded && (
<div
style={{
padding: "0.75rem",
background: "#f7fafc",
borderLeft: "1px solid #e2e8f0",
borderRight: "1px solid #e2e8f0",
borderBottom: "1px solid #e2e8f0",
borderRadius: "0 0 8px 8px",
}}
>
{loading && (
<div
style={{
padding: "1rem",
textAlign: "center",
color: "#718096",
fontSize: "0.875rem",
}}
>
Loading documents...
</div>
)}
{error && (
<div
style={{
padding: "0.75rem",
background: "#fed7d7",
color: "#c53030",
borderRadius: "4px",
fontSize: "0.875rem",
}}
>
{error}
</div>
)}
{!loading && !error && documents.length === 0 && (
<div
style={{
padding: "1rem",
textAlign: "center",
color: "#718096",
fontSize: "0.875rem",
}}
>
No documents in database
</div>
)}
{!loading && !error && documents.length > 0 && (
<>
{/* Search Input */}
<div style={{ marginBottom: "0.5rem" }}>
<input
type="text"
placeholder="Search by title or URI..."
value={searchQuery}
onChange={(e) => setSearchQuery(e.target.value)}
disabled={disabled}
style={{
width: "100%",
padding: "0.5rem 0.75rem",
fontSize: "0.875rem",
border: "1px solid #e2e8f0",
borderRadius: "4px",
background: disabled ? "#f7fafc" : "white",
color: disabled ? "#a0aec0" : "#2d3748",
outline: "none",
}}
/>
</div>
{/* Select All / Clear */}
<div
style={{
marginBottom: "0.5rem",
paddingBottom: "0.5rem",
borderBottom: "1px solid #e2e8f0",
display: "flex",
alignItems: "center",
justifyContent: "space-between",
}}
>
<div>
<button
type="button"
onClick={handleSelectAll}
disabled={disabled}
style={{
padding: "0.375rem 0.75rem",
fontSize: "0.75rem",
background: disabled ? "#e2e8f0" : "#4299e1",
color: disabled ? "#a0aec0" : "white",
border: "none",
borderRadius: "4px",
cursor: disabled ? "not-allowed" : "pointer",
}}
>
{filteredDocuments.every((d) => selected.includes(d.id))
? "Clear Visible"
: "Select Visible"}
</button>
<span
style={{
marginLeft: "0.75rem",
fontSize: "0.75rem",
color: "#718096",
}}
>
{selectedCount} of {totalCount} selected
</span>
</div>
{searchQuery && (
<span
style={{
fontSize: "0.75rem",
color: "#718096",
}}
>
Showing {filteredDocuments.length} of {totalCount}
</span>
)}
</div>
{/* Document List */}
<div
style={{
maxHeight: "200px",
overflowY: "auto",
display: "flex",
flexDirection: "column",
gap: "0.25rem",
}}
>
{filteredDocuments.map((doc) => {
const isSelected = selected.includes(doc.id);
return (
<label
key={doc.id}
style={{
display: "flex",
alignItems: "center",
gap: "0.5rem",
padding: "0.5rem",
background: isSelected ? "#ebf8ff" : "white",
border: isSelected
? "1px solid #90cdf4"
: "1px solid #e2e8f0",
borderRadius: "4px",
cursor: disabled ? "not-allowed" : "pointer",
transition: "all 0.15s",
}}
>
<input
type="checkbox"
checked={isSelected}
onChange={() => handleToggle(doc.id)}
disabled={disabled}
style={{
width: "1rem",
height: "1rem",
cursor: disabled ? "not-allowed" : "pointer",
}}
/>
<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
style={{
fontSize: doc.title ? "0.7rem" : "0.875rem",
fontWeight: doc.title
? "400"
: isSelected
? "600"
: "400",
color: doc.title ? "#718096" : "#2d3748",
whiteSpace: "nowrap",
overflow: "hidden",
textOverflow: "ellipsis",
}}
>
{doc.uri}
</div>
</div>
</label>
);
})}
</div>
</>
)}
{disabled && (
<div
style={{
marginTop: "0.5rem",
padding: "0.5rem",
background: "#fef5e7",
border: "1px solid #f6ad55",
borderRadius: "4px",
fontSize: "0.75rem",
color: "#744210",
textAlign: "center",
}}
>
Filter locked during research
</div>
)}
</div>
)}
</div>
);
}

968
examples/ag-ui-research/frontend/components/StateDisplay.tsx
View file

@ -1,968 +0,0 @@
"use client";
import { Markdown } from "@copilotkit/react-ui";
import { useCallback, useState } from "react";
interface VisualGroundingState {
isOpen: boolean;
chunkId: string | null;
images: string[];
loading: boolean;
error: string | null;
}
interface Citation {
document_id: string;
chunk_id: string;
document_uri: string;
document_title?: string;
page_numbers: number[];
headings?: string[];
content: string;
}
interface SearchAnswer {
query: string;
answer: string;
confidence: number;
cited_chunks: string[];
citations: Citation[];
}
interface ResearchContext {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
}
interface EvaluationResult {
new_questions: string[];
confidence_score: number;
is_sufficient: boolean;
reasoning: string;
}
interface ResearchReport {
title: string;
executive_summary: string;
main_findings: string[];
conclusions: string[];
limitations: string[];
recommendations: string[];
sources_summary: string;
}
interface ResearchState {
context: ResearchContext;
iterations: number;
max_iterations: number;
confidence_threshold: number;
max_concurrency: number;
last_eval: EvaluationResult | null;
result?: ResearchReport;
current_activity?: string;
current_activity_message?: string;
}
interface StateDisplayProps {
state: ResearchState;
}
export default function StateDisplay({ state }: StateDisplayProps) {
const [expandedSections, setExpandedSections] = useState<
Record<string, boolean>
>({
questions: true,
report: true,
});
const [expandedQuestions, setExpandedQuestions] = useState<
Record<string, boolean>
>({});
const [visualGrounding, setVisualGrounding] = useState<VisualGroundingState>({
isOpen: false,
chunkId: null,
images: [],
loading: false,
error: null,
});
const fetchVisualGrounding = useCallback(async (chunkId: string) => {
setVisualGrounding({
isOpen: true,
chunkId,
images: [],
loading: true,
error: null,
});
try {
const response = await fetch(
`${process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"}/api/visualize/${chunkId}`,
);
const data = await response.json();
if (!response.ok) {
throw new Error(data.error || "Failed to fetch visual grounding");
}
setVisualGrounding((prev) => ({
...prev,
images: data.images || [],
loading: false,
error: data.images?.length === 0 ? data.message : null,
}));
} catch (err) {
setVisualGrounding((prev) => ({
...prev,
loading: false,
error: err instanceof Error ? err.message : "Unknown error",
}));
}
}, []);
const closeVisualGrounding = useCallback(() => {
setVisualGrounding({
isOpen: false,
chunkId: null,
images: [],
loading: false,
error: null,
});
}, []);
const toggleSection = (section: string) => {
setExpandedSections((prev) => ({
...prev,
[section]: !prev[section],
}));
};
const toggleQuestion = (questionId: string) => {
setExpandedQuestions((prev) => ({
...prev,
[questionId]: !prev[questionId],
}));
};
// Calculate research progress based on iterations
const researchProgress =
state.max_iterations > 0
? (state.iterations / state.max_iterations) * 100
: 0;
const confidence = state.last_eval?.confidence_score || 0;
return (
<div
style={{
display: "flex",
flexDirection: "column",
gap: "1rem",
}}
>
{/* Question */}
{state.context?.original_question && (
<div
style={{
background: "white",
borderRadius: "8px",
padding: "1.5rem",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
}}
>
<div
style={{
fontSize: "0.875rem",
color: "#718096",
marginBottom: "0.25rem",
}}
>
Question
</div>
<div
style={{
fontSize: "1.125rem",
fontWeight: "bold",
color: "#2d3748",
}}
>
{state.context.original_question}
</div>
</div>
)}
{/* Research Progress - only show when research is in progress (not when complete) */}
{(state.iterations > 0 || state.current_activity) && !state.result && (
<div
style={{
background: "white",
borderRadius: "8px",
padding: "1.5rem",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
}}
>
{/* Current Activity */}
{state.current_activity && (
<div
style={{
padding: "0.75rem",
background: "#ebf8ff",
borderRadius: "6px",
border: "1px solid #90cdf4",
marginBottom: state.iterations > 0 ? "1rem" : 0,
}}
>
<div
style={{ display: "flex", alignItems: "center", gap: "0.5rem" }}
>
<span>⏳</span>
<span style={{ fontWeight: "600", color: "#2b6cb0" }}>
{state.current_activity
.replace(/_/g, " ")
.replace(/\b\w/g, (c) => c.toUpperCase())}
</span>
</div>
{state.current_activity_message && (
<div
style={{
fontSize: "0.875rem",
color: "#4299e1",
marginTop: "0.25rem",
marginLeft: "1.5rem",
}}
>
{state.current_activity_message}
</div>
)}
</div>
)}
{/* Iteration Progress Bar */}
{state.iterations > 0 && (
<div>
<div
style={{
display: "flex",
justifyContent: "space-between",
alignItems: "center",
marginBottom: "0.5rem",
}}
>
<span
style={{
fontSize: "0.75rem",
color: "#718096",
}}
>
Iterations
</span>
<span
style={{
fontSize: "0.75rem",
fontWeight: "600",
color: "#2d3748",
}}
>
{state.iterations}/{state.max_iterations}
</span>
</div>
<div
style={{
height: "0.5rem",
background: "#e2e8f0",
borderRadius: "4px",
overflow: "hidden",
}}
>
<div
style={{
width: `${researchProgress}%`,
height: "100%",
background: state.result ? "#48bb78" : "#4299e1",
transition: "width 0.3s ease",
}}
/>
</div>
</div>
)}
</div>
)}
{/* Confidence Meter */}
{confidence > 0 && (
<div
style={{
background: "white",
borderRadius: "8px",
padding: "1.5rem",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
}}
>
<div
style={{
fontSize: "0.875rem",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Confidence
</div>
<div style={{ display: "flex", alignItems: "center", gap: "1rem" }}>
<div
style={{
flex: 1,
height: "1rem",
background: "#e2e8f0",
borderRadius: "4px",
overflow: "hidden",
}}
>
<div
style={{
width: `${confidence * 100}%`,
height: "100%",
background:
confidence > 0.8
? "#48bb78"
: confidence > 0.5
? "#ed8936"
: "#f56565",
transition: "width 0.3s ease",
}}
/>
</div>
<div
style={{
fontSize: "1.5rem",
fontWeight: "bold",
color:
confidence > 0.8
? "#48bb78"
: confidence > 0.5
? "#ed8936"
: "#f56565",
}}
>
{(confidence * 100).toFixed(0)}%
</div>
</div>
{state.last_eval?.reasoning && (
<div
style={{
marginTop: "0.75rem",
fontSize: "0.875rem",
color: "#4a5568",
padding: "0.75rem",
background: "#f7fafc",
borderRadius: "4px",
}}
>
<Markdown content={state.last_eval.reasoning} />
</div>
)}
</div>
)}
{/* Answers */}
{state.context?.qa_responses && state.context.qa_responses.length > 0 && (
<div
style={{
background: "white",
borderRadius: "8px",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
overflow: "hidden",
}}
>
<button
type="button"
onClick={() => toggleSection("questions")}
style={{
width: "100%",
display: "flex",
justifyContent: "space-between",
alignItems: "center",
padding: "0.75rem",
background: "#edf2f7",
border: "1px solid #e2e8f0",
borderRadius: "4px",
cursor: "pointer",
fontSize: "1rem",
fontWeight: "600",
color: "#2d3748",
}}
>
<span>Answers ({state.context.qa_responses.length})</span>
<span>{expandedSections.questions ? "▼" : "▶"}</span>
</button>
{expandedSections.questions && (
<div
style={{
padding: "1rem",
background: "#f7fafc",
border: "1px solid #e2e8f0",
borderTop: "none",
borderRadius: "0 0 4px 4px",
}}
>
{/* Show all qa_responses (each has query + answer) */}
{state.context.qa_responses.map((qaResponse, idx) => {
const questionId = `q-${idx}`;
return (
<div
key={questionId}
style={{
marginBottom: "0.5rem",
background: "white",
borderRadius: "4px",
border: "1px solid #e2e8f0",
overflow: "hidden",
}}
>
<button
type="button"
onClick={() => toggleQuestion(questionId)}
style={{
width: "100%",
display: "flex",
gap: "0.75rem",
padding: "0.75rem",
background: "white",
border: "none",
cursor: "pointer",
textAlign: "left",
alignItems: "center",
}}
>
<div
style={{
fontSize: "1.25rem",
color: "#48bb78",
flexShrink: 0,
}}
>
✓
</div>
<div style={{ flex: 1 }}>
<div
style={{
fontSize: "0.875rem",
color: "#4a5568",
}}
>
<Markdown content={qaResponse.query} />
</div>
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginTop: "0.25rem",
}}
>
Confidence: {(qaResponse.confidence * 100).toFixed(0)}
%
</div>
</div>
<span
style={{
fontSize: "0.875rem",
color: "#718096",
}}
>
{expandedQuestions[questionId] ? "▼" : "▶"}
</span>
</button>
{/* QA Response nested inside question */}
{expandedQuestions[questionId] && (
<div
style={{
padding: "1rem",
background: "#f7fafc",
borderTop: "1px solid #e2e8f0",
}}
>
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginBottom: "0.5rem",
fontWeight: "600",
}}
>
Answer
</div>
<div
style={{
padding: "0.75rem",
background: "white",
borderRadius: "4px",
border: "1px solid #e2e8f0",
}}
>
<div
style={{
fontSize: "0.875rem",
color: "#2d3748",
lineHeight: "1.5",
}}
>
<Markdown content={qaResponse.answer} />
</div>
</div>
{/* Citations with visual grounding info */}
{qaResponse.citations &&
qaResponse.citations.length > 0 && (
<div style={{ marginTop: "1rem" }}>
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginBottom: "0.5rem",
fontWeight: "600",
}}
>
Citations ({qaResponse.citations.length})
</div>
{qaResponse.citations.map((citation, citIdx) => (
<div
key={citIdx}
style={{
padding: "0.75rem",
background: "white",
borderRadius: "4px",
border: "1px solid #e2e8f0",
marginBottom: "0.5rem",
}}
>
<div
style={{
display: "flex",
justifyContent: "space-between",
alignItems: "flex-start",
marginBottom: "0.5rem",
}}
>
<div
style={{
fontSize: "0.75rem",
fontWeight: "600",
color: "#2d3748",
}}
>
{citation.document_title ||
citation.document_uri}
</div>
{citation.page_numbers &&
citation.page_numbers.length > 0 && (
<div
style={{
fontSize: "0.7rem",
color: "#718096",
background: "#edf2f7",
padding: "0.125rem 0.375rem",
borderRadius: "4px",
}}
>
{citation.page_numbers.length === 1
? `p. ${citation.page_numbers[0]}`
: `pp. ${citation.page_numbers[0]}-${citation.page_numbers[citation.page_numbers.length - 1]}`}
</div>
)}
</div>
{citation.headings &&
citation.headings.length > 0 && (
<div
style={{
fontSize: "0.7rem",
color: "#718096",
marginBottom: "0.375rem",
}}
>
{citation.headings.join(" › ")}
</div>
)}
<div
style={{
fontSize: "0.8rem",
color: "#4a5568",
lineHeight: "1.4",
maxHeight: "4.5rem",
overflow: "hidden",
textOverflow: "ellipsis",
}}
>
{citation.content.slice(0, 200)}
{citation.content.length > 200 && "…"}
</div>
<button
type="button"
onClick={() =>
fetchVisualGrounding(citation.chunk_id)
}
style={{
marginTop: "0.5rem",
padding: "0.25rem 0.5rem",
fontSize: "0.7rem",
background: "#4299e1",
color: "white",
border: "none",
borderRadius: "4px",
cursor: "pointer",
}}
>
📍 View in Document
</button>
</div>
))}
</div>
)}
</div>
)}
</div>
);
})}
</div>
)}
</div>
)}
{/* Final Report */}
{state.result && (
<div
style={{
background: "white",
borderRadius: "8px",
boxShadow: "0 1px 3px rgba(0,0,0,0.1)",
overflow: "hidden",
}}
>
<button
type="button"
onClick={() => toggleSection("report")}
style={{
width: "100%",
display: "flex",
justifyContent: "space-between",
alignItems: "center",
padding: "0.75rem",
background: "#edf2f7",
border: "1px solid #e2e8f0",
borderRadius: "4px",
cursor: "pointer",
fontSize: "1rem",
fontWeight: "600",
color: "#2d3748",
}}
>
<span>Final Report</span>
<span>{expandedSections.report ? "▼" : "▶"}</span>
</button>
{expandedSections.report && (
<div
style={{
padding: "1.5rem",
background: "white",
border: "1px solid #e2e8f0",
borderTop: "none",
borderRadius: "0 0 4px 4px",
}}
>
<h3
style={{
fontSize: "1.25rem",
fontWeight: "600",
marginBottom: "1rem",
color: "#2d3748",
}}
>
{state.result.title}
</h3>
<div style={{ marginBottom: "1.5rem" }}>
<h4
style={{
fontSize: "0.875rem",
fontWeight: "600",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Executive Summary
</h4>
<div
style={{
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
<Markdown content={state.result.executive_summary} />
</div>
</div>
<div style={{ marginBottom: "1.5rem" }}>
<h4
style={{
fontSize: "0.875rem",
fontWeight: "600",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Main Findings
</h4>
<ul
style={{
paddingLeft: "1.5rem",
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
{state.result.main_findings.map((finding, idx) => (
<li
key={`finding-${idx}-${finding.substring(0, 30)}`}
style={{ marginBottom: "0.5rem" }}
>
<Markdown content={finding} />
</li>
))}
</ul>
</div>
<div style={{ marginBottom: "1.5rem" }}>
<h4
style={{
fontSize: "0.875rem",
fontWeight: "600",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Conclusions
</h4>
<ul
style={{
paddingLeft: "1.5rem",
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
{state.result.conclusions.map((conclusion, idx) => (
<li
key={`conclusion-${idx}-${conclusion.substring(0, 30)}`}
style={{ marginBottom: "0.5rem" }}
>
<Markdown content={conclusion} />
</li>
))}
</ul>
</div>
{state.result.recommendations.length > 0 && (
<div style={{ marginBottom: "1.5rem" }}>
<h4
style={{
fontSize: "0.875rem",
fontWeight: "600",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Recommendations
</h4>
<ul
style={{
paddingLeft: "1.5rem",
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
{state.result.recommendations.map((rec, idx) => (
<li
key={`rec-${idx}-${rec.substring(0, 30)}`}
style={{ marginBottom: "0.5rem" }}
>
<Markdown content={rec} />
</li>
))}
</ul>
</div>
)}
{state.result.limitations.length > 0 && (
<div style={{ marginBottom: "1.5rem" }}>
<h4
style={{
fontSize: "0.875rem",
fontWeight: "600",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Limitations
</h4>
<ul
style={{
paddingLeft: "1.5rem",
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
{state.result.limitations.map((lim, idx) => (
<li
key={`lim-${idx}-${lim.substring(0, 30)}`}
style={{ marginBottom: "0.5rem" }}
>
<Markdown content={lim} />
</li>
))}
</ul>
</div>
)}
<div>
<h4
style={{
fontSize: "0.875rem",
fontWeight: "600",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Sources
</h4>
<div
style={{
fontSize: "0.875rem",
color: "#4a5568",
lineHeight: "1.6",
}}
>
<Markdown content={state.result.sources_summary} />
</div>
</div>
</div>
)}
</div>
)}
{/* Visual Grounding Modal */}
{visualGrounding.isOpen && (
<div
style={{
position: "fixed",
top: 0,
left: 0,
right: 0,
bottom: 0,
background: "rgba(0, 0, 0, 0.75)",
display: "flex",
alignItems: "center",
justifyContent: "center",
zIndex: 1000,
}}
onClick={closeVisualGrounding}
>
<div
style={{
background: "white",
borderRadius: "8px",
padding: "1.5rem",
maxWidth: "90vw",
maxHeight: "90vh",
overflow: "auto",
position: "relative",
}}
onClick={(e) => e.stopPropagation()}
>
<button
type="button"
onClick={closeVisualGrounding}
style={{
position: "absolute",
top: "0.5rem",
right: "0.5rem",
background: "#e53e3e",
color: "white",
border: "none",
borderRadius: "50%",
width: "2rem",
height: "2rem",
cursor: "pointer",
fontSize: "1rem",
}}
>
✕
</button>
<h3
style={{
margin: "0 0 1rem 0",
fontSize: "1.125rem",
color: "#2d3748",
}}
>
Visual Grounding
</h3>
{visualGrounding.loading && (
<div
style={{
padding: "2rem",
textAlign: "center",
color: "#718096",
}}
>
Loading...
</div>
)}
{visualGrounding.error && (
<div
style={{
padding: "1rem",
background: "#fed7d7",
color: "#c53030",
borderRadius: "4px",
}}
>
{visualGrounding.error}
</div>
)}
{!visualGrounding.loading &&
!visualGrounding.error &&
visualGrounding.images.length > 0 && (
<div
style={{
display: "flex",
flexDirection: "column",
gap: "1rem",
}}
>
{visualGrounding.images.map((img, idx) => (
<div key={idx}>
<div
style={{
fontSize: "0.75rem",
color: "#718096",
marginBottom: "0.5rem",
}}
>
Page {idx + 1} of {visualGrounding.images.length}
</div>
<img
src={`data:image/png;base64,${img}`}
alt={`Page ${idx + 1}`}
style={{
maxWidth: "100%",
border: "1px solid #e2e8f0",
borderRadius: "4px",
}}
/>
</div>
))}
</div>
)}
</div>
</div>
)}
</div>
);
}

7
examples/ag-ui-research/frontend/next.config.ts
View file

@ -1,7 +0,0 @@
import type { NextConfig } from "next";
const nextConfig: NextConfig = {
output: "standalone",
};
export default nextConfig;

11689
examples/ag-ui-research/frontend/package-lock.json generated
View file

File diff suppressed because it is too large Load diff

29
examples/ag-ui-research/frontend/package.json
View file

@ -1,29 +0,0 @@
{
"name": "ag-ui-frontend",
"version": "0.1.0",
"private": true,
"scripts": {
"dev": "next dev",
"build": "next build",
"start": "next start",
"lint": "next lint",
"check": "biome check app components",
"format": "biome check --write app components"
},
"dependencies": {
"@ag-ui/client": "^0.0.42",
"@copilotkit/react-core": "^1.50.0",
"@copilotkit/react-ui": "^1.50.0",
"@copilotkit/runtime": "^1.50.0",
"next": "15.5.5",
"react": "^19.0.0",
"react-dom": "^19.0.0"
},
"devDependencies": {
"@biomejs/biome": "2.2.6",
"@types/node": "^22",
"@types/react": "^19",
"@types/react-dom": "^19",
"typescript": "^5"
}
}

0
examples/ag-ui-research/frontend/public/.gitkeep
View file

27
examples/ag-ui-research/frontend/tsconfig.json
View file

@ -1,27 +0,0 @@
{
"compilerOptions": {
"target": "ES2017",
"lib": ["dom", "dom.iterable", "esnext"],
"allowJs": true,
"skipLibCheck": true,
"strict": true,
"noEmit": true,
"esModuleInterop": true,
"module": "esnext",
"moduleResolution": "bundler",
"resolveJsonModule": true,
"isolatedModules": true,
"jsx": "preserve",
"incremental": true,
"plugins": [
{
"name": "next"
}
],
"paths": {
"@/*": ["./*"]
}
},
"include": ["next-env.d.ts", "**/*.ts", "**/*.tsx", ".next/types/**/*.ts"],
"exclude": ["node_modules"]
}

36
examples/ag-ui-research/haiku.rag.yaml.example
View file

@ -1,36 +0,0 @@
# haiku.rag configuration for ag-ui-research example
# Copy to haiku.rag.yaml and customize
# Document processing with docling-serve
processing:
converter: docling-serve
chunker: docling-serve
chunk_size: 256
chunker_type: hybrid
providers:
docling_serve:
base_url: http://docling-serve:5001
api_key: ""
ollama:
base_url: http://host.docker.internal:11434
research:
model:
provider: ollama
name: gpt-oss:latest
max_iterations: 3
confidence_threshold: 0.8
max_concurrency: 1
# For OpenAI:
# research:
# model:
# provider: openai
# name: gpt-4o-mini
# For Anthropic:
# research:
# model:
# provider: anthropic
# name: claude-3-5-haiku-20241022

6
examples/ag-ui-research/package-lock.json generated
View file

@ -1,6 +0,0 @@
{
"name": "haiku-ag-ui",
"lockfileVersion": 3,
"requires": true,
"packages": {}
}

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

@ -20,7 +20,6 @@ from rich.progress import (
from haiku.rag.client import HaikuRAG, RebuildMode
from haiku.rag.config import AppConfig, Config
from haiku.rag.graph.agui import AGUIConsoleRenderer, stream_graph
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
@ -409,7 +408,6 @@ class HaikuRAGApp:
question: str,
cite: bool = False,
deep: bool = False,
verbose: bool = False,
filter: str | None = None,
):
"""Ask a question using the RAG system.
@ -418,7 +416,6 @@ class HaikuRAGApp:
question: The question to ask
cite: Include citations in the answer
deep: Use deep QA mode (multi-step reasoning)
verbose: Show verbose output
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
@ -430,8 +427,6 @@ class HaikuRAGApp:
try:
citations = []
if deep:
from haiku.rag.graph.research.models import ResearchReport
graph = build_research_graph(config=self.config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(
@ -443,18 +438,7 @@ class HaikuRAGApp:
state.search_filter = filter
deps = ResearchDeps(client=self.client)
if verbose:
renderer = AGUIConsoleRenderer(self.console)
result_dict = await renderer.render(
stream_graph(graph, state, deps)
)
report = (
ResearchReport.model_validate(result_dict)
if result_dict
else None
)
else:
report = await graph.run(state=state, deps=deps)
report = await graph.run(state=state, deps=deps)
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
@ -485,14 +469,11 @@ class HaikuRAGApp:
except Exception as e:
self.console.print(f"[red]Error: {e}[/red]")
async def research(
self, question: str, verbose: bool = False, filter: str | None = None
):
async def research(self, question: str, filter: str | None = None):
"""Run research via the pydantic-graph pipeline.
Args:
question: The research question
verbose: Show AG-UI event stream during execution
filter: SQL WHERE clause to filter documents
"""
async with HaikuRAG(
@ -512,28 +493,12 @@ class HaikuRAGApp:
state.search_filter = filter
deps = ResearchDeps(client=client)
if verbose:
# Use AG-UI renderer to process and display events
renderer = AGUIConsoleRenderer(self.console)
report_dict = await renderer.render(
stream_graph(graph, state, deps)
)
else:
# Run without rendering events, just get the result
report = await graph.run(state=state, deps=deps)
report_dict = (
report.model_dump() if hasattr(report, "model_dump") else report
)
report = await graph.run(state=state, deps=deps)
if report_dict is None:
if report is None:
self.console.print("[red]Research did not produce a report.[/red]")
return
# Convert dict to ResearchReport model
from haiku.rag.graph.research.models import ResearchReport
report = ResearchReport.model_validate(report_dict)
# Display the report
self.console.print("[bold green]Research Report[/bold green]")
self.console.rule()
@ -819,7 +784,6 @@ class HaikuRAGApp:
enable_mcp: bool = True,
mcp_transport: str | None = None,
mcp_port: int = 8001,
enable_agui: bool = False,
):
"""Start the server with selected services."""
async with HaikuRAG(
@ -859,30 +823,6 @@ class HaikuRAGApp:
mcp_task = asyncio.create_task(run_mcp())
tasks.append(mcp_task)
# Start AG-UI server if enabled
if enable_agui:
async def run_agui():
import uvicorn
from haiku.rag.graph.agui import create_agui_server
logger.info(
f"Starting AG-UI server on {self.config.agui.host}:{self.config.agui.port}"
)
app = create_agui_server(self.config, db_path=self.db_path)
config = uvicorn.Config(
app=app,
host=self.config.agui.host,
port=self.config.agui.port,
log_level="info",
)
server = uvicorn.Server(config)
await server.serve()
agui_task = asyncio.create_task(run_agui())
tasks.append(agui_task)
if not tasks:
logger.warning("No services enabled")
return

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

@ -332,11 +332,6 @@ def ask(
"--deep",
help="Use deep multi-agent QA for complex questions",
),
verbose: bool = typer.Option(
False,
"--verbose",
help="Show verbose progress output (only with --deep)",
),
filter: str | None = typer.Option(
None,
"--filter",
@ -345,63 +340,26 @@ def ask(
),
):
app = create_app(db)
asyncio.run(
app.ask(question=question, cite=cite, deep=deep, verbose=verbose, filter=filter)
)
asyncio.run(app.ask(question=question, cite=cite, deep=deep, filter=filter))
@cli.command("research", help="Run multi-agent research and output a concise report")
def research(
question: str = typer.Argument(
None,
help="The research question to investigate (required unless --interactive)",
),
question: str = typer.Argument(..., help="The research question to investigate"),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
verbose: bool = typer.Option(
False,
"--verbose",
help="Show planning, searching previews, evaluation summary, and stop reason",
),
filter: str | None = typer.Option(
None,
"--filter",
"-f",
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)
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, read_only=_read_only, before=_before
)
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))
asyncio.run(app.research(question=question, filter=filter))
@cli.command("settings", help="Display current configuration settings")
@ -584,7 +542,7 @@ def inspect(
@cli.command(
"serve",
help="Start haiku.rag server. Use --monitor, --mcp, and/or --agui to enable services.",
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
)
def serve(
db: Path | None = typer.Option(
@ -612,17 +570,12 @@ def serve(
"--mcp-port",
help="Port to bind MCP server to (ignored with --stdio)",
),
agui: bool = typer.Option(
False,
"--agui",
help="Enable AG-UI HTTP server for graph streaming",
),
) -> None:
"""Start the server with selected services."""
# Require at least one service flag
if not (monitor or mcp or agui):
if not (monitor or mcp):
typer.echo(
"Error: At least one service flag (--monitor, --mcp, or --agui) must be specified"
"Error: At least one service flag (--monitor or --mcp) must be specified"
)
raise typer.Exit(1)
@ -640,7 +593,6 @@ def serve(
enable_mcp=mcp,
mcp_transport=transport,
mcp_port=mcp_port,
enable_agui=agui,
)
)

489
haiku_rag_slim/haiku/rag/cli_chat.py
View file

@ -1,489 +0,0 @@
"""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))

4
haiku_rag_slim/haiku/rag/config/__init__.py
View file

@ -1,12 +1,9 @@
import os
from haiku.rag.config.loader import (
find_config_file,
generate_default_config,
load_yaml_config,
)
from haiku.rag.config.models import (
AGUIConfig,
AppConfig,
ConversionOptions,
EmbeddingModelConfig,
@ -26,7 +23,6 @@ from haiku.rag.config.models import (
__all__ = [
"Config",
"AGUIConfig",
"AppConfig",
"ConversionOptions",
"EmbeddingModelConfig",

10
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -171,15 +171,6 @@ class ProvidersConfig(BaseModel):
docling_serve: DoclingServeConfig = Field(default_factory=DoclingServeConfig)
class AGUIConfig(BaseModel):
host: str = "0.0.0.0"
port: int = 8000
cors_origins: list[str] = ["*"]
cors_credentials: bool = True
cors_methods: list[str] = ["GET", "POST", "OPTIONS"]
cors_headers: list[str] = ["*"]
class PromptsConfig(BaseModel):
domain_preamble: str = ""
qa: str | None = None
@ -204,5 +195,4 @@ class AppConfig(BaseModel):
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
search: SearchConfig = Field(default_factory=SearchConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
agui: AGUIConfig = Field(default_factory=AGUIConfig)
prompts: PromptsConfig = Field(default_factory=PromptsConfig)

10
haiku_rag_slim/haiku/rag/graph/__init__.py
View file

@ -1,15 +1,5 @@
from haiku.rag.graph.agui import (
AGUIConsoleRenderer,
AGUIEmitter,
create_agui_server,
stream_graph,
)
from haiku.rag.graph.research.graph import build_research_graph
__all__ = [
"AGUIConsoleRenderer",
"AGUIEmitter",
"build_research_graph",
"create_agui_server",
"stream_graph",
]

59
haiku_rag_slim/haiku/rag/graph/agui/__init__.py
View file

@ -1,59 +0,0 @@
"""Generic AG-UI protocol support for haiku.rag graphs."""
from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
from haiku.rag.graph.agui.emitter import (
AGUIEmitter,
AGUIEvent,
emit_activity,
emit_activity_delta,
emit_run_error,
emit_run_finished,
emit_run_started,
emit_state_delta,
emit_state_snapshot,
emit_step_finished,
emit_step_started,
emit_text_message,
emit_text_message_content,
emit_text_message_end,
emit_text_message_start,
emit_tool_call_args,
emit_tool_call_end,
emit_tool_call_start,
)
from haiku.rag.graph.agui.server import (
RunAgentInput,
create_agui_app,
create_agui_server,
format_sse_event,
)
from haiku.rag.graph.agui.state import compute_state_delta
from haiku.rag.graph.agui.stream import stream_graph
__all__ = [
"AGUIConsoleRenderer",
"AGUIEmitter",
"AGUIEvent",
"RunAgentInput",
"compute_state_delta",
"create_agui_app",
"create_agui_server",
"emit_activity",
"emit_activity_delta",
"emit_run_error",
"emit_run_finished",
"emit_run_started",
"emit_state_delta",
"emit_state_snapshot",
"emit_step_finished",
"emit_step_started",
"emit_text_message",
"emit_text_message_content",
"emit_text_message_end",
"emit_text_message_start",
"emit_tool_call_args",
"emit_tool_call_end",
"emit_tool_call_start",
"format_sse_event",
"stream_graph",
]

135
haiku_rag_slim/haiku/rag/graph/agui/cli_renderer.py
View file

@ -1,135 +0,0 @@
"""Generic CLI renderer for AG-UI events with Rich console output."""
from collections.abc import AsyncIterator
from typing import Any
from rich.console import Console
from haiku.rag.graph.agui.emitter import AGUIEvent
class AGUIConsoleRenderer:
"""Renders AG-UI events to Rich console with formatted output.
Generic renderer that processes AG-UI protocol events and renders them
with Rich formatting. Works with any graph that emits AG-UI events.
"""
def __init__(self, console: Console | None = None):
"""Initialize the renderer.
Args:
console: Optional Rich console instance (creates new one if not provided)
"""
self.console = console or Console()
async def render(self, events: AsyncIterator[AGUIEvent]) -> Any | None:
"""Process events and render to console, return final result.
Args:
events: Async iterator of AG-UI events
Returns:
The final result from RunFinished event, or None
"""
result = None
async for event in events:
event_type = event.get("type")
if event_type == "RUN_STARTED":
self._render_run_started(event)
elif event_type == "RUN_FINISHED":
result = event.get("result")
self._render_run_finished()
elif event_type == "RUN_ERROR":
self._render_error(event)
elif event_type == "STEP_STARTED":
self._render_step_started(event)
elif event_type == "STEP_FINISHED":
self._render_step_finished(event)
elif event_type == "TEXT_MESSAGE_CHUNK":
self._render_text_message(event)
elif event_type == "TEXT_MESSAGE_START":
pass # Start of streaming message, no output needed
elif event_type == "TEXT_MESSAGE_CONTENT":
self._render_text_content(event)
elif event_type == "TEXT_MESSAGE_END":
pass # End of streaming message, no output needed
elif event_type == "STATE_SNAPSHOT":
self._render_state_snapshot(event)
elif event_type == "STATE_DELTA":
self._render_state_delta(event)
elif event_type == "ACTIVITY_SNAPSHOT":
self._render_activity(event)
elif event_type == "ACTIVITY_DELTA":
pass # Activity deltas don't need separate rendering
return result
def _render_run_started(self, event: AGUIEvent) -> None:
"""Render run start event."""
run_id = event.get("runId", "")
if run_id:
# Show shortened run ID (first 8 chars like our UUIDs)
short_id = run_id[:8] if len(run_id) > 8 else run_id
self.console.print(f"[bold green][RUN_STARTED][/bold green] Run {short_id}")
def _render_run_finished(self) -> None:
"""Render run completion."""
self.console.print("[bold green][RUN_FINISHED][/bold green] Completed")
def _render_error(self, event: AGUIEvent) -> None:
"""Render error event."""
message = event.get("message", "Unknown error")
self.console.print(f"[bold red][RUN_ERROR][/bold red] {message}")
def _render_step_started(self, event: AGUIEvent) -> None:
"""Render step start event."""
step_name = event.get("stepName", "")
if step_name:
display_name = step_name.replace("_", " ").title()
self.console.print(
f"\n[bold cyan][STEP_STARTED][/bold cyan] {display_name}"
)
def _render_step_finished(self, event: AGUIEvent) -> None:
"""Render step finish event."""
step_name = event.get("stepName", "")
if step_name:
display_name = step_name.replace("_", " ").title()
self.console.print(f"[cyan][STEP_FINISHED][/cyan] {display_name}")
def _render_text_message(self, event: AGUIEvent) -> None:
"""Render complete text message."""
delta = event.get("delta", "")
self.console.print(f"[magenta][TEXT_MESSAGE][/magenta] {delta}")
def _render_text_content(self, event: AGUIEvent) -> None:
"""Render streaming text content delta."""
delta = event.get("delta", "")
self.console.print(delta, end="")
def _render_activity(self, event: AGUIEvent) -> None:
"""Render activity update."""
content = event.get("content", "")
if content:
self.console.print(f"[yellow][ACTIVITY][/yellow] {content}")
def _render_state_snapshot(self, event: AGUIEvent) -> None:
"""Render full state snapshot."""
snapshot = event.get("snapshot")
if not snapshot:
return
self.console.print("[blue][STATE_SNAPSHOT][/blue]")
self.console.print(snapshot, style="dim")
def _render_state_delta(self, event: AGUIEvent) -> None:
"""Render state delta operations."""
delta = event.get("delta", [])
if not delta:
return
self.console.print("[blue][STATE_DELTA][/blue]")
self.console.print(delta, style="dim")

385
haiku_rag_slim/haiku/rag/graph/agui/emitter.py
View file

@ -1,385 +0,0 @@
"""Generic AG-UI event emitter for any graph execution."""
import asyncio
import hashlib
import json
from collections.abc import AsyncIterator
from typing import Any
from uuid import uuid4
from ag_ui.core import (
ActivitySnapshotEvent,
BaseEvent,
RunErrorEvent,
RunFinishedEvent,
RunStartedEvent,
StateDeltaEvent,
StateSnapshotEvent,
StepFinishedEvent,
StepStartedEvent,
TextMessageChunkEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallEndEvent,
ToolCallStartEvent,
)
from pydantic import BaseModel
from haiku.rag.graph.agui.state import compute_state_delta
AGUIEvent = dict[str, Any]
def _serialize_event(event: BaseEvent) -> AGUIEvent:
"""Serialize an ag_ui event to a dict with camelCase keys."""
return event.model_dump(mode="json", by_alias=True, exclude_none=True)
class AGUIEmitter[StateT: BaseModel, ResultT]:
"""Generic queue-backed AG-UI event emitter for any graph.
Manages the lifecycle of AG-UI events including:
- Run lifecycle (start, finish, error)
- Step lifecycle (start, finish)
- Text messages
- State synchronization (snapshots and deltas)
- Activity updates
Type parameters:
StateT: The Pydantic BaseModel type for graph state
ResultT: The result type returned by the graph
"""
def __init__(
self,
thread_id: str | None = None,
run_id: str | None = None,
use_deltas: bool = True,
):
"""Initialize the emitter.
Args:
thread_id: Optional thread ID (generated from input hash if not provided)
run_id: Optional run ID (random UUID if not provided)
use_deltas: Whether to emit state deltas instead of full snapshots (default: True)
"""
self._queue: asyncio.Queue[AGUIEvent | None] = asyncio.Queue()
self._closed = False
self._thread_id = thread_id or str(uuid4())
self._run_id = run_id or str(uuid4())
self._last_state: StateT | None = None
self._active_steps: set[str] = set()
self._use_deltas = use_deltas
@property
def thread_id(self) -> str:
"""Get the thread ID for this emitter."""
return self._thread_id
@property
def run_id(self) -> str:
"""Get the run ID for this emitter."""
return self._run_id
def start_run(self, initial_state: StateT) -> None:
"""Emit RunStarted and initial StateSnapshot.
Args:
initial_state: The initial state of the graph
"""
# If thread_id wasn't provided, generate from state hash
if not self._thread_id or self._thread_id == str(uuid4()):
state_json = initial_state.model_dump_json()
self._thread_id = self._generate_thread_id(state_json)
# RunStarted (state snapshot follows immediately with full state)
self.emit(
_serialize_event(
RunStartedEvent(thread_id=self._thread_id, run_id=self._run_id)
)
)
self.emit(
_serialize_event(StateSnapshotEvent(snapshot=initial_state.model_dump()))
)
# Store a deep copy to detect future changes
self._last_state = initial_state.model_copy(deep=True)
def start_step(self, step_name: str) -> None:
"""Emit StepStarted event.
Args:
step_name: Name of the step being started
"""
self._active_steps.add(step_name)
self.emit(_serialize_event(StepStartedEvent(step_name=step_name)))
def finish_step(self, step_name: str) -> None:
"""Emit StepFinished event for the specified step.
Args:
step_name: Name of the step being finished
"""
self._active_steps.discard(step_name)
self.emit(_serialize_event(StepFinishedEvent(step_name=step_name)))
def log(self, message: str, role: str = "assistant") -> None:
"""Emit a text message event.
Args:
message: The message content
role: The role of the sender (default: assistant)
"""
message_id = str(uuid4())
self.emit(
_serialize_event(
TextMessageChunkEvent(
message_id=message_id,
role=role, # type: ignore[arg-type]
delta=message,
)
)
)
def update_state(self, new_state: StateT) -> None:
"""Emit StateDelta or StateSnapshot for state change.
Args:
new_state: The updated state
"""
if self._use_deltas and self._last_state is not None:
# Emit delta for incremental updates
delta = compute_state_delta(self._last_state, new_state)
self.emit(_serialize_event(StateDeltaEvent(delta=delta)))
else:
# Emit full snapshot for initial state or when deltas disabled
self.emit(
_serialize_event(StateSnapshotEvent(snapshot=new_state.model_dump()))
)
# Store a deep copy to detect future changes
self._last_state = new_state.model_copy(deep=True)
def update_activity(
self,
activity_type: str,
content: dict[str, Any],
message_id: str | None = None,
) -> None:
"""Emit ActivitySnapshot event.
Args:
activity_type: Type of activity (e.g., "planning", "searching")
content: Structured payload representing the activity state
message_id: Optional message ID to associate activity with (auto-generated if None)
"""
if message_id is None:
message_id = str(uuid4())
self.emit(
_serialize_event(
ActivitySnapshotEvent(
message_id=message_id,
activity_type=activity_type,
content=content,
)
)
)
def finish_run(self, result: ResultT) -> None:
"""Emit RunFinished event.
Args:
result: The final result from the graph
"""
# Convert result to dict if it's a Pydantic model
result_data: Any = result
if hasattr(result, "model_dump"):
result_data = result.model_dump() # type: ignore[union-attr]
self.emit(
_serialize_event(
RunFinishedEvent(
thread_id=self._thread_id, run_id=self._run_id, result=result_data
)
)
)
def error(self, error: Exception, code: str | None = None) -> None:
"""Emit RunError event.
Args:
error: The exception that occurred
code: Optional error code
"""
self.emit(_serialize_event(RunErrorEvent(message=str(error), code=code)))
def emit(self, event: AGUIEvent) -> None:
"""Put event in queue.
Args:
event: The event to emit
"""
if not self._closed:
self._queue.put_nowait(event)
async def close(self) -> None:
"""Close the emitter and stop event iteration."""
if self._closed:
return
self._closed = True
await self._queue.put(None)
def __aiter__(self) -> AsyncIterator[AGUIEvent]:
"""Enable async iteration over events."""
return self._iter_events()
async def _iter_events(self) -> AsyncIterator[AGUIEvent]:
"""Iterate over events from the queue."""
while True:
event = await self._queue.get()
if event is None:
break
yield event
@staticmethod
def _generate_thread_id(input_data: str) -> str:
"""Generate a deterministic thread ID from input data.
Args:
input_data: The input data (e.g., question, prompt)
Returns:
A stable thread ID based on input hash
"""
# Use hash of input for deterministic thread ID
hash_obj = hashlib.sha256(input_data.encode("utf-8"))
return hash_obj.hexdigest()[:16]
def emit_text_message_start(message_id: str, role: str = "assistant") -> AGUIEvent:
"""Create a TextMessageStart event."""
return _serialize_event(
TextMessageStartEvent(message_id=message_id, role=role) # type: ignore[arg-type]
)
def emit_text_message_content(message_id: str, delta: str) -> AGUIEvent:
"""Create a TextMessageContent event."""
return _serialize_event(TextMessageContentEvent(message_id=message_id, delta=delta))
def emit_text_message_end(message_id: str) -> AGUIEvent:
"""Create a TextMessageEnd event."""
return _serialize_event(TextMessageEndEvent(message_id=message_id))
def emit_tool_call_start(
tool_call_id: str,
tool_name: str,
parent_message_id: str | None = None,
) -> AGUIEvent:
"""Create a ToolCallStart event."""
return _serialize_event(
ToolCallStartEvent(
tool_call_id=tool_call_id,
tool_call_name=tool_name,
parent_message_id=parent_message_id,
)
)
def emit_tool_call_args(tool_call_id: str, args: dict[str, Any]) -> AGUIEvent:
"""Create a ToolCallArgs event."""
return _serialize_event(
ToolCallArgsEvent(tool_call_id=tool_call_id, delta=json.dumps(args))
)
def emit_tool_call_end(tool_call_id: str) -> AGUIEvent:
"""Create a ToolCallEnd event."""
return _serialize_event(ToolCallEndEvent(tool_call_id=tool_call_id))
def emit_run_started(thread_id: str, run_id: str) -> AGUIEvent:
"""Create a RunStarted event."""
return _serialize_event(RunStartedEvent(thread_id=thread_id, run_id=run_id))
def emit_run_finished(thread_id: str, run_id: str, result: Any) -> AGUIEvent:
"""Create a RunFinished event."""
# Convert result to dict if it's a Pydantic model
if hasattr(result, "model_dump"):
result = result.model_dump()
return _serialize_event(
RunFinishedEvent(thread_id=thread_id, run_id=run_id, result=result)
)
def emit_run_error(message: str, code: str | None = None) -> AGUIEvent:
"""Create a RunError event."""
return _serialize_event(RunErrorEvent(message=message, code=code))
def emit_step_started(step_name: str) -> AGUIEvent:
"""Create a StepStarted event."""
return _serialize_event(StepStartedEvent(step_name=step_name))
def emit_step_finished(step_name: str) -> AGUIEvent:
"""Create a StepFinished event."""
return _serialize_event(StepFinishedEvent(step_name=step_name))
def emit_text_message(content: str, role: str = "assistant") -> AGUIEvent:
"""Create a TextMessageChunk event (convenience wrapper)."""
message_id = str(uuid4())
return _serialize_event(
TextMessageChunkEvent(
message_id=message_id,
role=role, # type: ignore[arg-type]
delta=content,
)
)
def emit_state_snapshot(state: BaseModel) -> AGUIEvent:
"""Create a StateSnapshot event."""
return _serialize_event(StateSnapshotEvent(snapshot=state.model_dump()))
def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> AGUIEvent:
"""Create a StateDelta event with JSON Patch operations."""
delta = compute_state_delta(old_state, new_state)
return _serialize_event(StateDeltaEvent(delta=delta))
def emit_activity(
message_id: str,
activity_type: str,
content: dict[str, Any],
) -> AGUIEvent:
"""Create an ActivitySnapshot event."""
return _serialize_event(
ActivitySnapshotEvent(
message_id=message_id,
activity_type=activity_type,
content=content,
)
)
def emit_activity_delta(
message_id: str,
activity_type: str,
patch: list[dict[str, Any]],
) -> AGUIEvent:
"""Create an ActivityDelta event with JSON Patch operations."""
from ag_ui.core import ActivityDeltaEvent
return _serialize_event(
ActivityDeltaEvent(
message_id=message_id,
activity_type=activity_type,
patch=patch,
)
)

268
haiku_rag_slim/haiku/rag/graph/agui/server.py
View file

@ -1,268 +0,0 @@
"""AG-UI HTTP server implementation for graph execution."""
import json
from collections.abc import AsyncIterator, Callable
from pathlib import Path
from typing import TYPE_CHECKING, Any, Protocol
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
from pydantic import BaseModel, Field
from pydantic_graph.beta import Graph
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, StreamingResponse
from starlette.routing import Route
from haiku.rag.config.models import AGUIConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
from haiku.rag.graph.agui.stream import stream_graph
class GraphDeps(Protocol):
"""Protocol for graph dependencies that support AG-UI emission."""
agui_emitter: AGUIEmitter[Any, Any] | None
class RunAgentInput(BaseModel):
"""AG-UI protocol run agent input.
See: https://docs.ag-ui.com/concepts/agents#runagentinput
"""
thread_id: str | None = Field(None, alias="threadId")
run_id: str | None = Field(None, alias="runId")
state: dict[str, Any] = Field(default_factory=dict)
messages: list[dict[str, Any]] = Field(default_factory=list)
config: dict[str, Any] = Field(default_factory=dict)
def create_agui_app(
graph_factory: Callable[[], Graph],
state_factory: Callable[[dict[str, Any]], BaseModel],
deps_factory: Callable[[dict[str, Any]], GraphDeps],
config: AGUIConfig,
) -> Starlette:
"""Create Starlette app with AG-UI endpoint.
Args:
graph_factory: Factory function to create graph instance
state_factory: Factory to create initial state from input
deps_factory: Factory to create graph dependencies
config: AG-UI server configuration
Returns:
Starlette application with AG-UI endpoints
"""
async def event_stream(
input_data: RunAgentInput,
) -> AsyncIterator[str]:
"""Generate SSE event stream from graph execution.
Yields:
Server-Sent Events formatted strings
"""
# Create graph, state, and dependencies
graph = graph_factory()
# Create initial state from input
initial_state = state_factory(input_data.state)
# Create dependencies (may use config from input)
deps = deps_factory(input_data.config)
# Execute graph and stream events
async for event in stream_graph(graph, initial_state, deps):
# Format as SSE event
event_data = format_sse_event(event)
yield event_data
async def stream_agent(request: Request) -> StreamingResponse:
"""AG-UI agent stream endpoint.
Accepts AG-UI RunAgentInput and streams events via SSE.
"""
# Parse request body
body = await request.json()
input_data = RunAgentInput(**body)
# Return SSE stream
return StreamingResponse(
event_stream(input_data),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no", # Disable buffering in nginx
},
)
async def health_check(_: Request) -> JSONResponse:
"""Health check endpoint."""
return JSONResponse({"status": "healthy"})
# Define routes
routes = [
Route("/v1/agent/stream", stream_agent, methods=["POST"]),
Route("/health", health_check, methods=["GET"]),
]
# Configure CORS middleware
middleware = [
Middleware(
CORSMiddleware,
allow_origins=config.cors_origins,
allow_credentials=config.cors_credentials,
allow_methods=config.cors_methods,
allow_headers=config.cors_headers,
)
]
# Create Starlette app
app = Starlette(
routes=routes,
middleware=middleware,
debug=False,
)
return app
def format_sse_event(event: AGUIEvent) -> str:
"""Format AG-UI event as Server-Sent Event.
Args:
event: AG-UI event dictionary
Returns:
SSE formatted string with event data
"""
# Convert event to JSON
event_json = json.dumps(event, ensure_ascii=False)
# Format as SSE
# Each event is: data: <json>\n\n
return f"data: {event_json}\n\n"
def create_agui_server( # pragma: no cover
config: "AppConfig", db_path: Path | None = None
) -> Starlette:
"""Create AG-UI server with research endpoint.
Args:
config: Application config with research settings
db_path: Optional database path override
Returns:
Starlette app with research endpoint
"""
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import (
ResearchDeps,
ResearchState,
)
# Store client reference for proper lifecycle management
_client_cache: dict[str, HaikuRAG] = {}
def get_client(effective_db_path: Path) -> HaikuRAG:
"""Get or create cached client."""
path_key = str(effective_db_path)
if path_key not in _client_cache:
_client_cache[path_key] = HaikuRAG(db_path=effective_db_path, config=config)
return _client_cache[path_key]
# Research graph factories
def research_graph_factory() -> Graph:
return build_research_graph(config)
def research_state_factory(input_state: dict[str, Any]) -> ResearchState:
question = input_state.get("question", "")
if not question:
messages = input_state.get("messages", [])
if messages:
question = messages[0].get("content", "")
context = ResearchContext(original_question=question)
max_iterations = input_state.get("max_iterations")
confidence_threshold = input_state.get("confidence_threshold")
return ResearchState.from_config(
context=context,
config=config,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
)
def research_deps_factory(input_config: dict[str, Any]) -> ResearchDeps:
effective_db_path = (
db_path
or input_config.get("db_path")
or config.storage.data_dir / "haiku.rag.lancedb"
)
return ResearchDeps(client=get_client(effective_db_path))
# Create event stream function
async def research_event_stream(
input_data: RunAgentInput,
) -> AsyncIterator[str]:
"""Generate SSE event stream from research graph execution."""
graph = research_graph_factory()
initial_state = research_state_factory(input_data.state)
deps = research_deps_factory(input_data.config)
async for event in stream_graph(graph, initial_state, deps):
event_data = format_sse_event(event)
yield event_data
# Endpoint handlers
async def stream_research(request: Request) -> StreamingResponse:
"""Research graph streaming endpoint."""
body = await request.json()
input_data = RunAgentInput(**body)
return StreamingResponse(
research_event_stream(input_data),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
async def health_check(_: Request) -> JSONResponse:
"""Health check endpoint."""
return JSONResponse({"status": "healthy"})
# Define routes
routes = [
Route("/v1/research/stream", stream_research, methods=["POST"]),
Route("/health", health_check, methods=["GET"]),
]
# Configure CORS middleware
middleware = [
Middleware(
CORSMiddleware,
allow_origins=config.agui.cors_origins,
allow_credentials=config.agui.cors_credentials,
allow_methods=config.agui.cors_methods,
allow_headers=config.agui.cors_headers,
)
]
# Create Starlette app
app = Starlette(
routes=routes,
middleware=middleware,
debug=False,
)
return app

34
haiku_rag_slim/haiku/rag/graph/agui/state.py
View file

@ -1,34 +0,0 @@
"""Generic AG-UI state utilities for any Pydantic BaseModel."""
from typing import Any
from pydantic import BaseModel
def compute_state_delta(
old_state: BaseModel, new_state: BaseModel
) -> list[dict[str, Any]]:
"""Compute JSON Patch (RFC 6902) operations from old state to new state.
Args:
old_state: Previous state (any Pydantic BaseModel)
new_state: Current state (same type as old_state)
Returns:
List of JSON Patch operations
"""
operations: list[dict[str, Any]] = []
# Convert states to dicts for comparison
old_dict = old_state.model_dump()
new_dict = new_state.model_dump()
# Compare each field and generate patches
for key, new_value in new_dict.items():
old_value = old_dict.get(key)
if old_value != new_value:
# Simple replace operation
operations.append({"op": "replace", "path": f"/{key}", "value": new_value})
return operations

85
haiku_rag_slim/haiku/rag/graph/agui/stream.py
View file

@ -1,85 +0,0 @@
"""Generic graph streaming with AG-UI events."""
import asyncio
from collections.abc import AsyncIterator
from contextlib import suppress
from typing import Protocol, TypeVar
from pydantic import BaseModel
from pydantic_graph.beta import Graph
from haiku.rag.graph.agui.emitter import AGUIEmitter, AGUIEvent
StateT = TypeVar("StateT", bound=BaseModel)
ResultT = TypeVar("ResultT")
class GraphDeps[StateT: BaseModel, ResultT](Protocol):
"""Protocol for graph dependencies that support AG-UI emission."""
agui_emitter: AGUIEmitter[StateT, ResultT] | None
async def stream_graph[StateT: BaseModel, DepsT: GraphDeps, ResultT](
graph: Graph[StateT, DepsT, None, ResultT],
state: StateT,
deps: DepsT,
use_deltas: bool = True,
) -> AsyncIterator[AGUIEvent]:
"""Run a graph and yield AG-UI events as they occur.
This is a generic streaming function that works with any pydantic-graph
that follows the AG-UI pattern:
- State must be a Pydantic BaseModel
- Deps must have an optional agui_emitter attribute
- Graph must be a pydantic-graph Graph instance
Args:
graph: The pydantic-graph Graph to execute
state: Initial state (Pydantic BaseModel)
deps: Graph dependencies with agui_emitter support
use_deltas: Whether to emit state deltas instead of full snapshots (default: True)
Yields:
AG-UI event dictionaries
Raises:
TypeError: If deps doesn't support agui_emitter
RuntimeError: If graph doesn't produce a result
"""
if not hasattr(deps, "agui_emitter"):
raise TypeError("deps must have an 'agui_emitter' attribute")
# Create AG-UI emitter
emitter: AGUIEmitter[StateT, ResultT] = AGUIEmitter(use_deltas=use_deltas)
deps.agui_emitter = emitter # type: ignore[assignment]
async def _execute() -> None:
try:
# Start the run with initial state
emitter.start_run(initial_state=state)
# Execute the graph
result = await graph.run(state=state, deps=deps)
if result is None:
raise RuntimeError("Graph did not produce a result")
# Finish the run with the result
emitter.finish_run(result)
except Exception as exc:
# Emit error event
emitter.error(exc)
finally:
await emitter.close()
runner = asyncio.create_task(_execute())
try:
async for event in emitter:
yield event
finally:
if not runner.done():
runner.cancel()
with suppress(asyncio.CancelledError):
await runner

418
haiku_rag_slim/haiku/rag/graph/research/graph.py
View file

@ -1,6 +1,4 @@
import asyncio
from typing import Literal
from uuid import uuid4
from pydantic_ai import Agent, RunContext, format_as_xml
from pydantic_ai.output import ToolOutput
@ -9,13 +7,6 @@ from pydantic_graph.beta.join import reduce_list_append
from haiku.rag.config import Config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter 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.models import (
Citation,
@ -196,14 +187,12 @@ def _get_batch_logic(state: ResearchState) -> list[str] | None:
def build_research_graph(
config: AppConfig = Config,
include_plan: bool = True,
interactive: bool = False,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the Research graph.
Args:
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)
interactive: Whether to include human decision nodes for HIL
Returns:
Configured Research graph
@ -230,89 +219,23 @@ def build_research_graph(
@g.step
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
"""Create research plan with sub-questions."""
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity(
"planning", {"stepName": "plan", "message": "Creating research plan"}
)
try:
await _plan_step_logic(state, deps, config, plan_prompt)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning",
{
"stepName": "plan",
"message": f"Created plan with {count} sub-questions",
"sub_questions": list(state.context.sub_questions),
},
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step("plan")
await _plan_step_logic(ctx.state, ctx.deps, config, plan_prompt)
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
"""Answer a single sub-question using the knowledge base."""
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
step_name = f"search: {sub_q}"
if deps.agui_emitter:
deps.agui_emitter.start_step(step_name)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
try:
answer = await _search_one_step_logic(
state, deps, config, search_prompt, sub_q
return await _search_one_step_logic(
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Found answer with {answer.confidence:.0%} confidence",
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
except Exception as e:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
return SearchAnswer(
query=sub_q,
query=ctx.inputs,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step(step_name)
@g.step
async def get_batch(
@ -329,154 +252,58 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("decide")
deps.agui_emitter.update_activity(
"evaluating", {"message": "Evaluating research sufficiency"}
agent = Agent(
model=get_model(model_config, config),
output_type=EvaluationResult,
instructions=decision_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
]
if state.last_eval is not None:
prev = state.last_eval
prompt_parts.append(
"<previous_evaluation>"
f"<confidence>{prev.confidence_score:.2f}</confidence>"
f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
f"<reasoning>{prev.reasoning}</reasoning>"
"</previous_evaluation>"
)
prompt = "\n\n".join(part for part in prompt_parts if part)
try:
agent = Agent(
model=get_model(model_config, config),
output_type=EvaluationResult,
instructions=decision_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
decision_result = await agent.run(prompt, deps=agent_deps)
output = decision_result.output
context_xml = format_context_for_prompt(state.context)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
]
if state.last_eval is not None:
prev = state.last_eval
prompt_parts.append(
"<previous_evaluation>"
f"<confidence>{prev.confidence_score:.2f}</confidence>"
f"<is_sufficient>{str(prev.is_sufficient).lower()}</is_sufficient>"
f"<reasoning>{prev.reasoning}</reasoning>"
"</previous_evaluation>"
)
prompt = "\n\n".join(part for part in prompt_parts if part)
state.last_eval = output
state.iterations += 1
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
decision_result = await agent.run(prompt, deps=agent_deps)
output = decision_result.output
# Get already-answered questions to avoid duplicates
answered_queries = {qa.query.lower() for qa in state.context.qa_responses}
state.last_eval = output
state.iterations += 1
for new_q in output.new_questions:
# Skip if already in pending or already answered
if new_q in state.context.sub_questions:
continue
if new_q.lower() in answered_queries:
continue
state.context.sub_questions.append(new_q)
# Get already-answered questions to avoid duplicates
answered_queries = {qa.query.lower() for qa in state.context.qa_responses}
should_continue = (
not output.is_sufficient
or output.confidence_score < state.confidence_threshold
) and state.iterations < state.max_iterations
for new_q in output.new_questions:
# Skip if already in pending or already answered
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:
deps.agui_emitter.update_state(state)
sufficient = "Yes" if output.is_sufficient else "No"
deps.agui_emitter.update_activity(
"evaluating",
{
"stepName": "decide",
"message": f"Confidence: {output.confidence_score:.0%}, Sufficient: {sufficient}",
"confidence": output.confidence_score,
"is_sufficient": output.is_sufficient,
},
)
should_continue = (
not output.is_sufficient
or output.confidence_score < state.confidence_threshold
) and state.iterations < state.max_iterations
return should_continue
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step("decide")
@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("human_decide")
return should_continue
@g.step
async def synthesize(
@ -486,37 +313,27 @@ def build_research_graph(
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("synthesize")
deps.agui_emitter.update_activity(
"synthesizing", {"message": "Generating final research report"}
)
agent = Agent(
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=synthesis_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
try:
agent = Agent(
model=get_model(model_config, config),
output_type=ResearchReport,
instructions=synthesis_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
return result.output
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step("synthesize")
context_xml = format_context_for_prompt(state.context)
prompt = (
"Generate a comprehensive research report based on all gathered information.\n\n"
f"{context_xml}\n\n"
"Create a detailed report that synthesizes all findings into a coherent response."
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
return result.output
# Build the graph structure
collect_answers = g.join(
@ -524,77 +341,40 @@ def build_research_graph(
initial_factory=list[SearchAnswer],
)
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))
if include_plan:
g.add(
g.edge_from(human_decide).to(
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),
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
else:
# Non-interactive mode: automatic decision based on confidence/iterations
if include_plan:
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
g.add(g.edge_from(g.start_node).to(get_batch))
g.add(
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(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)
)
else:
g.add(g.edge_from(g.start_node).to(get_batch))
g.add(
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(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),
)
.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()

26
haiku_rag_slim/haiku/rag/graph/research/state.py
View file

@ -1,27 +1,15 @@
import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal
from typing import TYPE_CHECKING
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport
from haiku.rag.graph.research.models import EvaluationResult
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
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
@ -29,17 +17,7 @@ class ResearchDeps:
"""Dependencies for research graph execution."""
client: HaikuRAG
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | 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:
"""Emit a log message through AG-UI events."""
if self.agui_emitter:
self.agui_emitter.log(message)
if state:
self.agui_emitter.update_state(state)
class ResearchState(BaseModel):

1
tests/graph/agui/__init__.py
View file

@ -1 +0,0 @@
"""Tests for AG-UI implementation."""

204
tests/graph/agui/test_cli_renderer.py
View file

@ -1,204 +0,0 @@
"""Tests for AGUIConsoleRenderer."""
import pytest
from pydantic import BaseModel
from rich.console import Console
from haiku.rag.graph.agui.cli_renderer import AGUIConsoleRenderer
from haiku.rag.graph.agui.emitter import (
emit_activity,
emit_run_error,
emit_run_finished,
emit_run_started,
emit_state_delta,
emit_state_snapshot,
emit_step_finished,
emit_step_started,
emit_text_message,
)
class SimpleState(BaseModel):
"""Simple state for testing."""
value: int
async def async_gen(items):
"""Helper to create async generator from list."""
for item in items:
yield item
@pytest.mark.asyncio
async def test_renderer_basic_flow():
"""Test basic event rendering flow."""
console = Console(file=None, force_terminal=False) # Don't actually print
renderer = AGUIConsoleRenderer(console)
events = [
emit_run_started("t1", "r1-test-run-id"),
emit_state_snapshot(SimpleState(value=1)),
emit_step_started("plan"),
emit_step_finished("plan"),
emit_activity("m1", "planning", {"message": "Planning research"}),
emit_run_finished("t1", "r1", {"status": "complete"}),
]
result = await renderer.render(async_gen(events))
assert result == {"status": "complete"}
@pytest.mark.asyncio
async def test_renderer_multiple_snapshots():
"""Test that renderer handles multiple snapshots without errors."""
renderer = AGUIConsoleRenderer()
events = [
emit_state_snapshot(SimpleState(value=1)),
emit_state_snapshot(SimpleState(value=2)),
]
# Should render both snapshots without error
result = await renderer.render(async_gen(events))
assert result is None # No run finished event
@pytest.mark.asyncio
async def test_renderer_handles_all_event_types():
"""Test that renderer handles all event types without errors."""
renderer = AGUIConsoleRenderer()
events = [
emit_run_started("t1", "r1"),
emit_state_snapshot(SimpleState(value=1)),
emit_step_started("step1"),
emit_step_finished("step1"),
emit_text_message("message"),
emit_activity("m1", "work", {"message": "Working"}),
emit_run_error("error occurred"),
emit_run_finished("t1", "r1", {"result": "done"}),
]
# Should not raise any exceptions
result = await renderer.render(async_gen(events))
assert result == {"result": "done"}
@pytest.mark.asyncio
async def test_renderer_state_snapshots():
"""Test that state snapshots are rendered."""
renderer = AGUIConsoleRenderer()
events = [
emit_state_snapshot(SimpleState(value=1)),
emit_state_snapshot(SimpleState(value=2)),
emit_run_finished("t1", "r1", {"done": True}),
]
result = await renderer.render(async_gen(events))
assert result == {"done": True}
@pytest.mark.asyncio
async def test_renderer_no_result():
"""Test renderer when no RUN_FINISHED event."""
renderer = AGUIConsoleRenderer()
events = [
emit_run_started("t1", "r1"),
emit_step_started("step1"),
]
result = await renderer.render(async_gen(events))
assert result is None
@pytest.mark.asyncio
async def test_renderer_snapshot_then_delta():
"""Test that renderer handles snapshot followed by deltas."""
renderer = AGUIConsoleRenderer()
state1 = SimpleState(value=1)
state2 = SimpleState(value=2)
state3 = SimpleState(value=3)
events = [
emit_state_snapshot(state1), # Initial snapshot
emit_state_delta(state1, state2), # Delta to value=2
emit_state_delta(state2, state3), # Delta to value=3
emit_run_finished("t1", "r1", {"complete": True}),
]
result = await renderer.render(async_gen(events))
assert result == {"complete": True}
@pytest.mark.asyncio
async def test_renderer_with_empty_state():
"""Test renderer handles empty state gracefully."""
renderer = AGUIConsoleRenderer()
# Create a state with no changes
events = [
emit_step_started("step1"), # No state snapshot
emit_run_finished("t1", "r1", {"result": "ok"}),
]
result = await renderer.render(async_gen(events))
assert result == {"result": "ok"}
@pytest.mark.asyncio
async def test_renderer_state_delta():
"""Test that renderer renders state deltas."""
renderer = AGUIConsoleRenderer()
state1 = SimpleState(value=1)
state2 = SimpleState(value=2)
events = [
emit_state_snapshot(state1), # Initial state
emit_state_delta(state1, state2), # Delta update
emit_run_finished("t1", "r1", None),
]
result = await renderer.render(async_gen(events))
assert result is None
@pytest.mark.asyncio
async def test_renderer_state_delta_without_initial():
"""Test that renderer handles delta without initial state gracefully."""
renderer = AGUIConsoleRenderer()
state1 = SimpleState(value=1)
state2 = SimpleState(value=2)
# Send delta without initial snapshot - should still render it
events = [
emit_state_delta(state1, state2),
emit_run_finished("t1", "r1", {"ok": True}),
]
result = await renderer.render(async_gen(events))
assert result == {"ok": True}
@pytest.mark.asyncio
async def test_renderer_lifecycle_events():
"""Test that lifecycle events are rendered properly."""
renderer = AGUIConsoleRenderer()
events = [
emit_run_started("thread-123", "run-abc-def-ghi"), # Long run ID
emit_step_started("search_one"),
emit_step_finished("search_one"),
emit_step_started("analyze_insights"),
emit_step_finished("analyze_insights"),
emit_run_finished("thread-123", "run-abc-def-ghi", {"complete": True}),
]
result = await renderer.render(async_gen(events))
assert result == {"complete": True}

317
tests/graph/agui/test_emitter.py
View file

@ -1,317 +0,0 @@
"""Tests for AGUIEmitter."""
import asyncio
import pytest
from pydantic import BaseModel
from haiku.rag.graph.agui.emitter import AGUIEmitter
class TestState(BaseModel):
"""Test state model."""
value: int
text: str
class TestResult(BaseModel):
"""Test result model."""
status: str
@pytest.mark.asyncio
async def test_emitter_lifecycle():
"""Test emitter lifecycle events."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
initial_state = TestState(value=1, text="initial")
emitter.start_run(initial_state)
events = []
async for event in emitter:
events.append(event)
if event["type"] == "RUN_STARTED":
# Close after getting run started
result = TestResult(status="complete")
emitter.finish_run(result)
await emitter.close()
assert len(events) >= 3 # RUN_STARTED, STATE_SNAPSHOT, RUN_FINISHED
assert events[0]["type"] == "RUN_STARTED"
assert events[-1]["type"] == "RUN_FINISHED"
assert events[-1]["result"] == {"status": "complete"}
@pytest.mark.asyncio
async def test_emitter_step_events():
"""Test step lifecycle events."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
emitter.start_step("test_step")
emitter.finish_step("test_step")
await emitter.close()
events = []
async for event in emitter:
events.append(event)
step_events = [e for e in events if e["type"] in ("STEP_STARTED", "STEP_FINISHED")]
assert len(step_events) == 2
assert step_events[0]["type"] == "STEP_STARTED"
assert step_events[0]["stepName"] == "test_step"
assert step_events[1]["type"] == "STEP_FINISHED"
assert step_events[1]["stepName"] == "test_step"
@pytest.mark.asyncio
async def test_emitter_state_updates():
"""Test state update events with snapshots."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter(use_deltas=False)
state1 = TestState(value=1, text="first")
state2 = TestState(value=2, text="second")
emitter.update_state(state1)
emitter.update_state(state2)
await emitter.close()
events = []
async for event in emitter:
events.append(event)
state_events = [e for e in events if e["type"] == "STATE_SNAPSHOT"]
assert len(state_events) == 2
assert state_events[0]["snapshot"] == {"value": 1, "text": "first"}
assert state_events[1]["snapshot"] == {"value": 2, "text": "second"}
@pytest.mark.asyncio
async def test_emitter_state_deltas():
"""Test state update events with deltas."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter(use_deltas=True)
state1 = TestState(value=1, text="first")
state2 = TestState(value=2, text="second")
emitter.update_state(state1)
emitter.update_state(state2)
await emitter.close()
events = []
async for event in emitter:
events.append(event)
# First update should be a snapshot (no previous state)
snapshot_events = [e for e in events if e["type"] == "STATE_SNAPSHOT"]
assert len(snapshot_events) == 1
assert snapshot_events[0]["snapshot"] == {"value": 1, "text": "first"}
# Second update should be a delta
delta_events = [e for e in events if e["type"] == "STATE_DELTA"]
assert len(delta_events) == 1
# Delta should contain replace operations for changed fields
delta = delta_events[0]["delta"]
assert isinstance(delta, list)
assert len(delta) == 2 # Two fields changed
assert any(op["path"] == "/value" and op["value"] == 2 for op in delta)
assert any(op["path"] == "/text" and op["value"] == "second" for op in delta)
@pytest.mark.asyncio
async def test_emitter_activity_events():
"""Test activity events."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
# Activity without a step
emitter.update_activity("processing", {"message": "Processing data"})
# Activity within a step (stepName explicitly included in content)
emitter.start_step("analyze")
emitter.update_activity(
"done", {"stepName": "analyze", "message": "Completed"}, message_id="msg-1"
)
emitter.finish_step("analyze")
await emitter.close()
events = []
async for event in emitter:
events.append(event)
activity_events = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"]
assert len(activity_events) == 2
assert activity_events[0]["activityType"] == "processing"
assert activity_events[0]["content"]["message"] == "Processing data"
assert "stepName" not in activity_events[0]["content"] # No step context
assert activity_events[1]["messageId"] == "msg-1"
assert activity_events[1]["activityType"] == "done"
assert activity_events[1]["content"]["message"] == "Completed"
assert activity_events[1]["content"]["stepName"] == "analyze" # Has step context
@pytest.mark.asyncio
async def test_emitter_text_messages():
"""Test text message events."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
emitter.log("Test message", role="assistant")
emitter.log("Another message", role="user")
await emitter.close()
events = []
async for event in emitter:
events.append(event)
text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"]
assert len(text_events) == 2
assert text_events[0]["delta"] == "Test message"
assert text_events[0]["role"] == "assistant"
assert text_events[1]["delta"] == "Another message"
assert text_events[1]["role"] == "user"
@pytest.mark.asyncio
async def test_emitter_error():
"""Test error event emission."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
error = ValueError("Test error")
emitter.error(error, code="TEST_ERROR")
await emitter.close()
events = []
async for event in emitter:
events.append(event)
error_events = [e for e in events if e["type"] == "RUN_ERROR"]
assert len(error_events) == 1
assert error_events[0]["message"] == "Test error"
assert error_events[0]["code"] == "TEST_ERROR"
@pytest.mark.asyncio
async def test_emitter_thread_and_run_ids():
"""Test thread and run ID management."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter(
thread_id="thread-1", run_id="run-1"
)
assert emitter.thread_id == "thread-1"
assert emitter.run_id == "run-1"
initial_state = TestState(value=1, text="test")
emitter.start_run(initial_state)
await emitter.close()
events = []
async for event in emitter:
events.append(event)
run_started = [e for e in events if e["type"] == "RUN_STARTED"][0]
assert run_started["threadId"] == "thread-1"
assert run_started["runId"] == "run-1"
@pytest.mark.asyncio
async def test_emitter_generates_thread_id():
"""Test that thread ID is generated from state hash when not provided."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
initial_state = TestState(value=42, text="test")
emitter.start_run(initial_state)
# Thread ID should be generated deterministically from state
assert emitter.thread_id is not None
assert len(emitter.thread_id) > 0
await emitter.close()
async for _ in emitter:
pass
@pytest.mark.asyncio
async def test_emitter_closes_properly():
"""Test that emitter closes and stops iteration."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
emitter.log("Message 1")
await emitter.close()
# Attempting to iterate after close should work and stop
events = []
async for event in emitter:
events.append(event)
# Should have received the message and then stopped
assert len(events) == 1
@pytest.mark.asyncio
async def test_emitter_concurrent_emission():
"""Test that multiple events can be emitted concurrently."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
async def emit_many():
for i in range(10):
emitter.log(f"Message {i}")
await asyncio.sleep(0.001) # Simulate some work
await emitter.close()
# Start emission in background
emit_task = asyncio.create_task(emit_many())
# Collect events
events = []
async for event in emitter:
events.append(event)
await emit_task
# Should have all 10 messages
text_events = [e for e in events if e["type"] == "TEXT_MESSAGE_CHUNK"]
assert len(text_events) == 10
@pytest.mark.asyncio
async def test_emitter_concurrent_steps():
"""Test that multiple steps can run concurrently and finish independently."""
emitter: AGUIEmitter[TestState, TestResult] = AGUIEmitter()
async def run_step(step_name: str, delay: float):
emitter.start_step(step_name)
await asyncio.sleep(delay)
emitter.finish_step(step_name)
# Start collecting events in background
events: list[dict] = []
async def collect():
async for event in emitter:
events.append(event)
collector = asyncio.create_task(collect())
# Run three steps concurrently with different durations
# Step A finishes last, Step B first, Step C middle
await asyncio.gather(
run_step("step_a", 0.03),
run_step("step_b", 0.01),
run_step("step_c", 0.02),
)
await emitter.close()
await collector
started = [e for e in events if e["type"] == "STEP_STARTED"]
finished = [e for e in events if e["type"] == "STEP_FINISHED"]
# All three steps should have started
started_names = {e["stepName"] for e in started}
assert started_names == {"step_a", "step_b", "step_c"}
# All three steps should have finished
finished_names = {e["stepName"] for e in finished}
assert finished_names == {"step_a", "step_b", "step_c"}

227
tests/graph/agui/test_server.py
View file

@ -1,227 +0,0 @@
"""Tests for AG-UI server."""
import pytest
from pydantic import BaseModel
from starlette.testclient import TestClient
from haiku.rag.config.models import AGUIConfig
from haiku.rag.graph.agui.server import RunAgentInput, create_agui_app, format_sse_event
class SimpleState(BaseModel):
"""Simple state for testing."""
question: str
class SimpleResult(BaseModel):
"""Simple result for testing."""
answer: str
class MockGraph:
"""Mock graph that returns immediately."""
async def run(self, state, deps): # type: ignore[no-untyped-def]
"""Return a simple result."""
return SimpleResult(answer=f"Answer to: {state.question}")
def test_run_agent_input_parsing():
"""Test RunAgentInput model parsing."""
data = {
"threadId": "thread-1",
"runId": "run-1",
"state": {"question": "What is AI?"},
"messages": [],
"config": {},
}
input_data = RunAgentInput(**data)
assert input_data.thread_id == "thread-1"
assert input_data.run_id == "run-1"
assert input_data.state == {"question": "What is AI?"}
def test_format_sse_event():
"""Test SSE event formatting."""
event = {"type": "TEST_EVENT", "data": "test"}
sse = format_sse_event(event)
assert sse.startswith("data: ")
assert sse.endswith("\n\n")
assert '{"type": "TEST_EVENT"' in sse
def test_create_agui_app_basic():
"""Test basic app creation."""
config = AGUIConfig(
host="localhost",
port=8000,
cors_origins=["http://localhost"],
)
def graph_factory():
return MockGraph()
def state_factory(input_state):
return SimpleState(question=input_state.get("question", ""))
def deps_factory(input_config):
from dataclasses import dataclass
@dataclass
class SimpleDeps:
agui_emitter: None = None
return SimpleDeps()
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
# Should return a Starlette app
assert app is not None
assert hasattr(app, "routes")
def test_server_health_endpoint():
"""Test health check endpoint."""
config = AGUIConfig()
def graph_factory():
return MockGraph()
def state_factory(input_state):
return SimpleState(question="")
def deps_factory(input_config):
from dataclasses import dataclass
@dataclass
class SimpleDeps:
agui_emitter: None = None
return SimpleDeps()
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
client = TestClient(app)
response = client.get("/health")
assert response.status_code == 200
assert response.json() == {"status": "healthy"}
@pytest.mark.asyncio
async def test_server_stream_endpoint():
"""Test AG-UI streaming endpoint."""
config = AGUIConfig()
def graph_factory():
return MockGraph()
def state_factory(input_state):
question = input_state.get("question", "")
return SimpleState(question=question)
def deps_factory(input_config):
from dataclasses import dataclass
@dataclass
class SimpleDeps:
agui_emitter: None = None
return SimpleDeps()
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
client = TestClient(app)
request_data = {
"threadId": "test-1",
"runId": "run-1",
"state": {"question": "What is pydantic-graph?"},
"messages": [],
"config": {},
}
response = client.post("/v1/agent/stream", json=request_data)
assert response.status_code == 200
assert response.headers["content-type"] == "text/event-stream; charset=utf-8"
# Read the streamed events
events = []
for line in response.iter_lines():
if line.startswith("data: "):
import json
event_data = line[6:] # Remove "data: " prefix
event = json.loads(event_data)
events.append(event)
# Should have received multiple events
assert len(events) > 0
# Should have RUN_STARTED and RUN_FINISHED
event_types = [e["type"] for e in events]
assert "RUN_STARTED" in event_types
assert "RUN_FINISHED" in event_types
def test_server_cors_headers():
"""Test CORS middleware is configured."""
config = AGUIConfig(
cors_origins=["http://example.com"],
cors_credentials=True,
)
def graph_factory():
return MockGraph()
def state_factory(input_state):
return SimpleState(question="")
def deps_factory(input_config):
from dataclasses import dataclass
@dataclass
class SimpleDeps:
agui_emitter: None = None
return SimpleDeps()
app = create_agui_app(
graph_factory=graph_factory, # type: ignore[arg-type]
state_factory=state_factory,
deps_factory=deps_factory, # type: ignore[arg-type]
config=config,
)
client = TestClient(app)
# GET request with Origin header should get CORS headers
response = client.get("/health", headers={"Origin": "http://example.com"})
assert response.status_code == 200
# CORS middleware should add access-control headers
assert (
"access-control-allow-origin" in response.headers or response.status_code == 200
)

199
tests/graph/agui/test_stream.py
View file

@ -1,199 +0,0 @@
"""Tests for stream_graph function."""
from dataclasses import dataclass
import pytest
from pydantic import BaseModel
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.agui.stream import stream_graph
class TestState(BaseModel):
"""Test state model."""
value: int
@dataclass
class TestDeps:
"""Test dependencies."""
agui_emitter: AGUIEmitter | None = None
class MockGraph:
"""Mock graph for testing."""
def __init__(self, result: str | dict[str, str | int] = "done"):
self.result = result
self.run_called = False
async def run(self, state, deps): # type: ignore[no-untyped-def]
"""Mock run method."""
self.run_called = True
# Emit some events through the emitter
if deps.agui_emitter:
deps.agui_emitter.start_step("mock_step")
deps.agui_emitter.update_activity("working", {"message": "Doing work"})
deps.agui_emitter.finish_step("mock_step")
return self.result
@pytest.mark.asyncio
async def test_stream_graph_basic():
"""Test basic graph streaming."""
graph = MockGraph(result="success")
state = TestState(value=1)
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have collected events
assert len(events) > 0
# Should have RUN_STARTED and RUN_FINISHED
event_types = [e["type"] for e in events]
assert "RUN_STARTED" in event_types
assert "RUN_FINISHED" in event_types
# Graph should have been executed
assert graph.run_called
@pytest.mark.asyncio
async def test_stream_graph_emits_initial_state():
"""Test that initial state is emitted."""
graph = MockGraph()
state = TestState(value=42)
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have initial state snapshot
state_snapshots = [e for e in events if e["type"] == "STATE_SNAPSHOT"]
assert len(state_snapshots) > 0
# First snapshot should be initial state
assert state_snapshots[0]["snapshot"] == {"value": 42}
@pytest.mark.asyncio
async def test_stream_graph_emits_step_events():
"""Test that step events from graph are emitted."""
graph = MockGraph()
state = TestState(value=1)
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have step events from MockGraph
step_started = [e for e in events if e["type"] == "STEP_STARTED"]
assert len(step_started) > 0
assert step_started[0]["stepName"] == "mock_step"
@pytest.mark.asyncio
async def test_stream_graph_emits_activity():
"""Test that activity events from graph are emitted."""
graph = MockGraph()
state = TestState(value=1)
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have activity events from MockGraph
activities = [e for e in events if e["type"] == "ACTIVITY_SNAPSHOT"]
assert len(activities) > 0
assert activities[0]["content"] == {"message": "Doing work"}
@pytest.mark.asyncio
async def test_stream_graph_handles_error():
"""Test that graph errors are captured and emitted."""
class ErrorGraph:
async def run(self, state, deps):
raise ValueError("Test error")
graph = ErrorGraph()
state = TestState(value=1)
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Should have error event
errors = [e for e in events if e["type"] == "RUN_ERROR"]
assert len(errors) > 0
assert "Test error" in errors[0]["message"]
@pytest.mark.asyncio
async def test_stream_graph_closes_emitter():
"""Test that emitter is properly closed."""
class NeverReturnsGraph:
async def run(self, state, deps):
# Don't return anything
pass
graph = NeverReturnsGraph()
state = TestState(value=1)
deps = TestDeps()
events = []
try:
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
except RuntimeError:
# Expected - graph didn't return a result
pass
# Should have error event about no result
errors = [e for e in events if e["type"] == "RUN_ERROR"]
assert len(errors) > 0
@pytest.mark.asyncio
async def test_stream_graph_without_emitter_support():
"""Test error when deps doesn't support agui_emitter."""
@dataclass
class BadDeps:
pass
graph = MockGraph()
state = TestState(value=1)
deps = BadDeps()
with pytest.raises(TypeError, match="agui_emitter"):
async for _ in stream_graph(graph, state, deps): # type: ignore[arg-type]
pass
@pytest.mark.asyncio
async def test_stream_graph_result_in_finish_event():
"""Test that graph result is included in RUN_FINISHED event."""
graph = MockGraph(result={"status": "complete", "count": 42})
state = TestState(value=1)
deps = TestDeps()
events = []
async for event in stream_graph(graph, state, deps): # type: ignore[arg-type]
events.append(event)
# Find RUN_FINISHED event
finished = [e for e in events if e["type"] == "RUN_FINISHED"]
assert len(finished) == 1
assert finished[0]["result"] == {"status": "complete", "count": 42}

25
tests/graph/test_research_graph.py
View file

@ -1,9 +1,9 @@
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.agui.stream import stream_graph
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 ResearchDeps, ResearchState
@ -27,24 +27,11 @@ async def test_graph_end_to_end(allow_model_requests, temp_db_path, qa_corpus):
deps = ResearchDeps(client=client)
events = []
result = None
async for event in stream_graph(graph, state, deps):
events.append(event)
if event["type"] == "RUN_FINISHED":
result = event["result"]
elif event["type"] == "RUN_ERROR":
pytest.fail(f"Graph execution failed: {event['message']}")
result = await graph.run(state=state, deps=deps)
assert result is not None, (
f"No result. Events collected: {[e['type'] for e in events]}"
)
assert isinstance(result, dict)
assert "title" in result
assert "executive_summary" in result
event_types = [e["type"] for e in events]
assert "RUN_STARTED" in event_types
assert "RUN_FINISHED" in event_types
assert result is not None
assert isinstance(result, ResearchReport)
assert result.title
assert result.executive_summary
client.close()

47
tests/test_app.py
View file

@ -338,24 +338,6 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch):
assert mock_print.call_count >= 1
@pytest.mark.asyncio
async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with verbose (should be ignored for non-deep)."""
mock_answer = "Test answer"
mock_citations = []
mock_client = AsyncMock()
mock_client.ask.return_value = (mock_answer, mock_citations)
mock_client.__aenter__.return_value = mock_client
mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print)
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
await app.ask("test question", verbose=True)
mock_client.ask.assert_called_once_with("test question", filter=None)
@pytest.mark.asyncio
async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep mode uses research graph."""
@ -427,35 +409,6 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
assert call_kwargs["state"].context.original_question == "test question"
@pytest.mark.asyncio
async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep mode and verbose output."""
import haiku.rag.app as app_module
mock_output = {"executive_summary": "Deep research answer"}
mock_renderer = AsyncMock()
mock_renderer.render.return_value = mock_output
mock_graph = AsyncMock()
mock_client = AsyncMock()
mock_print = MagicMock()
monkeypatch.setattr(app.console, "print", mock_print)
monkeypatch.setattr(app_module, "build_research_graph", lambda **kwargs: mock_graph)
with patch("haiku.rag.app.HaikuRAG") as mock_rag_class:
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_client)
mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None)
with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer):
await app.ask("test question", deep=True, verbose=True)
# With verbose, it should use AGUIConsoleRenderer.render, not graph.run
mock_renderer.render.assert_called_once()
mock_graph.run.assert_not_called()
@pytest.mark.asyncio
async def test_history_all_tables(tmp_path, monkeypatch):
"""Test history command shows version history for all tables."""

22
tests/test_cli.py
View file

@ -282,7 +282,6 @@ def test_ask():
question="What is Python?",
cite=False,
deep=False,
verbose=False,
filter=None,
)
@ -300,7 +299,6 @@ def test_ask_with_cite():
question="What is Python?",
cite=True,
deep=False,
verbose=False,
filter=None,
)
@ -318,7 +316,6 @@ def test_ask_with_deep():
question="What is Python?",
cite=False,
deep=True,
verbose=False,
filter=None,
)
@ -336,25 +333,6 @@ def test_ask_with_deep_and_cite():
question="What is Python?",
cite=True,
deep=True,
verbose=False,
filter=None,
)
def test_ask_with_deep_and_verbose():
with patch("haiku.rag.cli.HaikuRAGApp") as mock_app:
mock_app_instance = MagicMock()
mock_app_instance.ask = AsyncMock()
mock_app.return_value = mock_app_instance
result = runner.invoke(cli, ["ask", "What is Python?", "--deep", "--verbose"])
assert result.exit_code == 0
mock_app_instance.ask.assert_called_once_with(
question="What is Python?",
cite=False,
deep=True,
verbose=True,
filter=None,
)