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

Merge pull request #136 from ggozad/feat/ag-ui

Browse source 
AG-UI support in research / deep ask graphs.
This commit is contained in:
Yiorgis Gozadinos 2025-11-13 13:30:41 +02:00 committed by GitHub
commit cb1001dc29
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
77 changed files with 16448 additions and 16051 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

46
CHANGELOG.md
View file

@ -1,15 +1,59 @@
# Changelog
## [Unreleased]
### Added
- **AG-UI Protocol Support**: Full AG-UI (Agent-UI) protocol implementation for graph execution with event streaming
- New `AGUIEmitter` class for emitting AG-UI events from graphs
- Support for all AG-UI event types: lifecycle events (`RUN_STARTED`, `RUN_FINISHED`, `RUN_ERROR`), step events (`STEP_STARTED`, `STEP_FINISHED`), state updates (`STATE_SNAPSHOT`, `STATE_DELTA`), activity narration (`ACTIVITY_SNAPSHOT`), and text messages (`TEXT_MESSAGE_CHUNK`)
- `AGUIConsoleRenderer` for rendering AG-UI event streams to terminal with Rich formatting
- `stream_graph()` utility function for executing graphs with AG-UI event emission
- State diff computation for efficient state synchronization
- **Delta State Updates**: AG-UI emitter now supports incremental state updates via JSON Patch operations (`STATE_DELTA` events) to reduce bandwidth, configurable via `use_deltas` parameter (enabled by default)
- **AG-UI Server**: Starlette-based HTTP server for serving graphs via AG-UI protocol
- Server-Sent Events (SSE) streaming endpoint at `/v1/agent/stream`
- Health check endpoint at `/health`
- Full CORS support configurable via `agui` config section
- `create_agui_server()` function for programmatic server creation
- **Deep QA AG-UI Support**: Deep QA graph now fully supports AG-UI event streaming
- Integration with `AGUIEmitter` for progress tracking
- Step-by-step execution visibility via AG-UI events
- **CLI AG-UI Flag**: New `--agui` flag for `serve` command to start AG-UI server
- **Graph Module**: New unified `haiku.rag.graph` module containing all graph-related functionality
- **Common Graph Nodes**: New factory functions (`create_plan_node`, `create_search_node`) in `haiku.rag.graph.common.nodes` for reusable graph components
- **AG-UI Research Example**: New full-stack example (`examples/ag-ui-research`) demonstrating agent+graph architecture with CopilotKit frontend
- Pydantic AI agent with research tool that invokes the research graph
- Custom AG-UI streaming endpoint with anyio memory streams
- React/Next.js frontend with split-pane UI showing live research state
- Real-time progress tracking of questions, answers, insights, and gaps
- Docker Compose setup for easy local development
### Changed
- **BREAKING**: Major refactoring of graph-related code into unified `haiku.rag.graph` module structure:
- `haiku.rag.research` → `haiku.rag.graph.research`
- `haiku.rag.qa.deep` → `haiku.rag.graph.deep_qa`
- `haiku.rag.agui` → `haiku.rag.graph.agui`
- `haiku.rag.graph_common` → `haiku.rag.graph.common`
- **BREAKING**: Research and Deep QA graphs now use AG-UI event protocol instead of direct console logging
- Removed `console` and `stream` parameters from graph dependencies
- All progress updates now emit through `AGUIEmitter`
- **BREAKING**: `ResearchState` converted from dataclass to Pydantic `BaseModel` for JSON serialization and AG-UI compatibility
- Research and Deep QA graphs now emit detailed execution events for better observability
- CLI research command now uses AG-UI event rendering for `--verbose` output
- Improved graph execution visibility with step-by-step progress tracking
- Updated all documentation to reflect new import paths and AG-UI usage
- Updated examples (ag-ui-research, a2a-server) to use new import paths
### Fixed
- **Document Creation**: Optimized `create_document` to skip unnecessary DoclingDocument conversion when chunks are pre-provided
- **FileReader**: Error messages now include both original exception details and file path for easier debugging
### Removed
- **BREAKING**: Removed legacy `ResearchStream` and `ResearchStreamEvent` classes (replaced by AG-UI event protocol)
## [0.15.0] - 2025-11-07
### Added

22
README.md
View file

@ -86,12 +86,12 @@ To customize settings, create a `haiku.rag.yaml` config file (see [Configuration
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.research import (
from haiku.rag.graph.agui import stream_graph
from haiku.rag.graph.research import (
ResearchContext,
ResearchDeps,
ResearchState,
build_research_graph,
stream_research_graph,
)
async with HaikuRAG("database.lancedb") as client:
@ -126,15 +126,17 @@ async with HaikuRAG("database.lancedb") as client:
report = await graph.run(state=state, deps=deps)
print(report.title)
# Streaming progress (log/report/error events)
async for event in stream_research_graph(graph, state, deps):
if event.type == "log":
iteration = event.state.iterations if event.state else state.iterations
print(f"[{iteration}] {event.message}")
elif event.type == "report":
# Streaming progress (AG-UI events)
async for event in stream_graph(graph, state, deps):
if event["type"] == "STEP_STARTED":
print(f"Starting step: {event['stepName']}")
elif event["type"] == "ACTIVITY_SNAPSHOT":
print(f" {event['content']}")
elif event["type"] == "RUN_FINISHED":
print("\nResearch complete!\n")
print(event.report.title)
print(event.report.executive_summary)
result = event["result"]
print(result["title"])
print(result["executive_summary"])
```
## MCP Server

42
docs/agents.md
View file

@ -96,9 +96,9 @@ Python usage:
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
async with HaikuRAG(path_to_db) as client:
# Use global config (recommended)
@ -205,9 +205,9 @@ Python usage (blocking result):
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
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
async with HaikuRAG(path_to_db) as client:
# Use global config (recommended)
@ -250,15 +250,15 @@ deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
```
Python usage (streamed events):
Python usage (streamed AG-UI events):
```python
from haiku.rag.graph.agui import stream_graph
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
from haiku.rag.research.stream import stream_research_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
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(config=Config)
@ -267,16 +267,14 @@ async with HaikuRAG(path_to_db) as client:
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
async for event in stream_research_graph(
graph,
state,
deps,
):
if event.type == "log":
iteration = event.state.iterations if event.state else state.iterations
print(f"[{iteration}] {event.message}")
elif event.type == "report":
async for event in stream_graph(graph, state, deps):
if event["type"] == "STEP_STARTED":
print(f"Starting step: {event['stepName']}")
elif event["type"] == "ACTIVITY_SNAPSHOT":
print(f" {event['content']}")
elif event["type"] == "RUN_FINISHED":
print("\nResearch complete!\n")
print(event.report.title)
print(event.report.executive_summary)
result = event["result"]
print(result["title"])
print(result["executive_summary"])
```

37
docs/cli.md
View file

@ -146,20 +146,23 @@ When available, citations use the document title; otherwise they fall back to th
Run the multi-step research graph:
```bash
haiku-rag research "How does haiku.rag organize and query documents?" \
--max-iterations 2 \
--confidence-threshold 0.8 \
--max-concurrency 3 \
--verbose
haiku-rag research "How does haiku.rag organize and query documents?"
```
With verbose output to see progress:
```bash
haiku-rag research "How does haiku.rag organize and query documents?" --verbose
```
Flags:
- `--max-iterations, -n`: maximum search/evaluate cycles (default: 3)
- `--confidence-threshold`: stop once evaluation confidence meets/exceeds this (default: 0.8)
- `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3)
- `--verbose`: show planning, searching previews, evaluation summary, and stop reason
- `--verbose`: Show planning, searching previews, evaluation summary, and stop reason
When `--verbose` is set the CLI also consumes the internal research stream, printing every `log` event as agents progress through planning, search, evaluation, and synthesis. If you build your own integration, call `stream_research_graph` to access the same `log`, `report`, and `error` events and render them however you like while the graph is running.
Research parameters like `max_iterations`, `confidence_threshold`, and `max_concurrency` are configured in your [configuration file](configuration.md) under the `research` section.
When `--verbose` is set, the CLI consumes the research graph's AG-UI event stream, displaying step events and activity snapshots as agents progress through planning, search, evaluation, and synthesis. Without `--verbose`, only the final research report is displayed.
If you build your own integration, import `stream_graph` from `haiku.rag.graph.agui` to access AG-UI events (`STEP_STARTED`, `ACTIVITY_SNAPSHOT`, `STATE_SNAPSHOT`, `RUN_FINISHED`, etc.) and render them however you like while the graph is running.
## Server
@ -174,10 +177,18 @@ haiku-rag serve --mcp --stdio
# File monitoring only
haiku-rag serve --monitor
# Both services
haiku-rag serve --monitor --mcp
# AG-UI server only
haiku-rag serve --agui
# Custom port
# Multiple services
haiku-rag serve --monitor --mcp
haiku-rag serve --monitor --agui
haiku-rag serve --mcp --agui
# All services
haiku-rag serve --monitor --mcp --agui
# Custom MCP port
haiku-rag serve --mcp --mcp-port 9000
```

57
docs/configuration.md
View file

@ -83,6 +83,17 @@ qa:
research:
provider: "" # Empty to use qa settings
model: ""
max_iterations: 3
confidence_threshold: 0.8
max_concurrency: 1
agui:
host: "0.0.0.0"
port: 8000
cors_origins: ["*"]
cors_credentials: true
cors_methods: ["GET", "POST", "OPTIONS"]
cors_headers: ["*"]
processing:
chunk_size: 256
@ -460,6 +471,52 @@ providers:
**Note:** vLLM reranking uses the `/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded. Consult the specific model's documentation for proper vLLM serving configuration.
## Research Configuration
Configure the multi-agent research workflow:
```yaml
research:
provider: "" # Empty to use qa settings
model: "" # Empty to use qa model
max_iterations: 3 # Maximum search/evaluate cycles
confidence_threshold: 0.8 # Stop when confidence meets/exceeds this
max_concurrency: 1 # Sub-questions searched in parallel per iteration
```
- **provider/model**: LLM provider and model for research. Leave empty to use the same settings as `qa`.
- **max_iterations**: Maximum number of search/evaluate cycles before stopping (default: 3)
- **confidence_threshold**: Stop research when evaluation confidence score meets or exceeds this threshold (default: 0.8)
- **max_concurrency**: Number of sub-questions to search in parallel during each iteration (default: 1)
The research workflow plans sub-questions, searches in parallel batches, evaluates findings, and iterates until reaching the confidence threshold or max iterations.
## AG-UI Server Configuration
Configure the AG-UI HTTP server for streaming graph execution events:
```yaml
agui:
host: "0.0.0.0"
port: 8000
cors_origins: ["*"]
cors_credentials: true
cors_methods: ["GET", "POST", "OPTIONS"]
cors_headers: ["*"]
```
Start the AG-UI server with:
```bash
haiku-rag serve --agui
```
The server exposes:
- `GET /health` - Health check endpoint
- `POST /v1/agent/stream` - Research graph streaming endpoint (Server-Sent Events)
See [Server Mode](server.md) for more details.
## Other Settings
### Database and Storage

123
docs/server.md
View file

@ -1,10 +1,10 @@
# Server Mode
The server provides automatic file monitoring and MCP functionality.
The server provides automatic file monitoring, MCP functionality, and AG-UI graph streaming.
## Starting the Server
The `serve` command requires at least one service flag. You can enable file monitoring, MCP server, or both:
The `serve` command requires at least one service flag. You can enable file monitoring, MCP server, AG-UI server, or any combination:
### MCP Server Only
@ -93,3 +93,122 @@ The server can parse 40+ file formats including:
- And more...
URLs are also supported for web content.
## AG-UI Server
The AG-UI server provides HTTP streaming of both research and deep ask graph execution using Server-Sent Events (SSE).
### Starting the AG-UI Server
```bash
haiku-rag serve --agui
```
This starts an HTTP server (default: http://0.0.0.0:8000) that exposes:
- `GET /health` - Health check endpoint
- `POST /v1/research/stream` - Research graph streaming endpoint
- `POST /v1/deep-ask/stream` - Deep ask graph streaming endpoint
### Configuration
Configure the AG-UI server in your `haiku.rag.yaml`:
```yaml
agui:
host: "0.0.0.0"
port: 8000
cors_origins: ["*"]
cors_credentials: true
```
See [Configuration](configuration.md#ag-ui-server-configuration) for all available options.
### Using the Streaming Endpoints
Both endpoints accept POST requests with the same AG-UI RunAgentInput format and stream AG-UI events.
**Request format:**
```json
{
"threadId": "optional-thread-id",
"runId": "optional-run-id",
"state": {
"question": "What are the key features of haiku.rag?"
},
"messages": [],
"config": {}
}
```
**Research endpoint example:**
```bash
curl -X POST http://localhost:8000/v1/research/stream \
-H "Content-Type: application/json" \
-d '{
"state": {
"question": "What are the key features of haiku.rag?"
}
}' \
--no-buffer
```
**Deep ask endpoint example:**
```bash
curl -X POST http://localhost:8000/v1/deep-ask/stream \
-H "Content-Type: application/json" \
-d '{
"state": {
"question": "How does haiku.rag handle document chunking?",
"use_citations": true
}
}' \
--no-buffer
```
The `--no-buffer` flag ensures curl displays events as they arrive instead of buffering them.
**Note:** The `state` object can include:
- `question`: The question to answer (required)
- `use_citations`: Enable citations in deep ask responses (optional, deep ask only)
**Response:** Server-Sent Events stream with AG-UI protocol events:
- `RUN_STARTED` - Graph execution started
- `STATE_SNAPSHOT` - Current state snapshot
- `STEP_STARTED` - Node execution started
- `STEP_FINISHED` - Node execution completed
- `ACTIVITY_SNAPSHOT` - Progress update
- `RUN_FINISHED` - Graph execution completed with result
- `RUN_ERROR` - Error during execution
Example event output:
```
data: {"type":"RUN_STARTED","threadId":"abc123","runId":"xyz789"}
data: {"type":"STATE_SNAPSHOT","snapshot":{"context":{"original_question":"What are the key features of haiku.rag?"},"iterations":0}}
data: {"type":"STEP_STARTED","stepName":"plan"}
data: {"type":"ACTIVITY_SNAPSHOT","messageId":"msg-1","activityType":"planning","content":"Creating research plan"}
data: {"type":"STEP_FINISHED","stepName":"plan"}
data: {"type":"RUN_FINISHED","threadId":"abc123","runId":"xyz789","result":{"title":"Research Report","executive_summary":"..."}}
```
The endpoint follows the [AG-UI protocol](https://docs.ag-ui.com/concepts/events) for event streaming.
### Running Multiple Services
You can run any combination of services:
```bash
# File monitoring + AG-UI
haiku-rag serve --monitor --agui
# MCP + AG-UI
haiku-rag serve --mcp --agui
# All services
haiku-rag serve --monitor --mcp --agui
```

2
examples/a2a-server/haiku_rag_a2a/a2a/__init__.py
View file

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

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

@ -3,6 +3,11 @@
# 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

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

@ -1,16 +1,16 @@
# Interactive Research Assistant
Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic AI](https://ai.pydantic.dev/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time.
Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/), [Pydantic Graph](https://ai.pydantic.dev/graph/), and [AG-UI](https://docs.ag-ui.com/). Ask complex questions and watch the research process unfold in real-time.
[Watch demo video](https://vimeo.com/1128874386)
## Features
- **Multi-step research workflow**: Question decomposition, search, analysis, and synthesis
- **Human-in-the-loop**: Approve or revise research plans before execution
- **Live state synchronization**: Real-time updates of research progress between backend and frontend
- **Context expansion**: Automatically expands top search results for better context
- **Rich reporting**: Generates structured reports with findings, conclusions, and citations
- **Multi-iteration research graph**: Automated question decomposition, search, insight extraction, and gap analysis
- **Intelligent evaluation**: Confidence-based decision making with automatic iteration until sufficient information is gathered
- **Live state synchronization**: Real-time delta updates of research progress via AG-UI protocol
- **Insight & gap tracking**: Structured insights with provenance and automatic gap identification
- **Rich reporting**: Generates comprehensive research reports with findings, conclusions, and sources
## Quick Start
@ -57,32 +57,67 @@ Research assistant powered by [haiku.rag](https://ggozad.github.io/haiku.rag/),
## How It Works
1. **Ask a question**: Type your research question in the chat
2. **Review the plan**: The agent decomposes your question into 3 sub-questions
3. **Approve or revise**: Choose to approve the plan or request changes
4. **Watch it work**: The agent automatically:
- Searches the knowledge base for each sub-question
- Extracts key insights from search results
- Evaluates overall confidence in findings
5. **Get your report**: Receive a structured research report with citations
2. **Plan phase**: The research graph automatically:
- Decomposes your question into targeted sub-questions
- Gathers initial context about the topic
3. **Research iterations**: The graph autonomously:
- Searches the knowledge base for each sub-question in parallel
- Extracts structured insights with source provenance
- Identifies information gaps and assesses confidence
- Generates new follow-up questions for gaps
- Iterates until confidence threshold is met or max iterations reached
4. **Synthesis**: Generates a comprehensive research report with:
- Executive summary
- Main findings with supporting evidence
- Conclusions and recommendations
- Source citations
## Architecture
- **Backend** (Python): Pydantic AI agent with haiku.rag integration
- Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
- `agent.py`: Research agent with tool definitions
- `main.py`: Starlette app serving AG-UI protocol
### Agent + Graph Pattern
- **Frontend** (Next.js): CopilotKit/AG-UI interface
- Real-time state synchronization with backend
- Interactive approval workflow
- Collapsible research plan and insights display
This example demonstrates the **agent+graph** architecture pattern:
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
- Formats research results for the user
2. **Research Graph** (haiku.rag):
- Multi-step research workflow invoked by the agent's tool
- Autonomous execution with plan → search → analyze → decide → synthesize flow
- Emits AG-UI events for real-time progress tracking
3. **Shared Event Stream**:
- `AGUIEmitter` is shared between agent and graph
- Events from both flow through a single stream to the frontend
- Custom streaming endpoint (`main.py`) uses anyio memory streams for proper async handling
### Components
- **Backend** (Python):
- Uses published `ghcr.io/ggozad/haiku.rag:latest` Docker image as base
- `agent.py`: Pydantic AI agent with `run_research` tool
- `main.py`: Custom AG-UI streaming endpoint with anyio memory object streams
- Real-time event forwarding from emitter to SSE stream
- Filters out `ACTIVITY_SNAPSHOT` events (not yet supported by CopilotKit)
- **Frontend** (Next.js/React):
- CopilotKit for AG-UI protocol integration
- Split-pane UI: chat on left, live research state on right
- Real-time state synchronization via Server-Sent Events (SSE)
- `StateDisplay` component with collapsible sections for questions, insights, and gaps
## Configuration
Configuration is done through `haiku.rag.yaml` (see `haiku.rag.yaml.example`):
- `qa.provider`: LLM provider (default: `ollama`)
- `qa.model`: Model name (default: `gpt-oss:latest`)
- `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`):

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

@ -3,15 +3,13 @@ FROM ghcr.io/ggozad/haiku.rag:latest
WORKDIR /app
# Copy backend application files
COPY agent.py main.py ./
COPY pyproject.toml ./
COPY main.py agent.py ./
# Install backend dependencies
# Note: haiku-rag is already installed in the base image
# Install additional dependencies for the example
# Note: haiku-rag-slim is already installed in the base image
RUN pip install --no-cache-dir \
starlette>=0.45.2 \
uvicorn[standard]>=0.34.2 \
pydantic-ai-slim[ag-ui,openai]>=1.1.0 \
python-dotenv>=1.0.1
EXPOSE 8000

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

@ -1,6 +1,6 @@
# Haiku.rag Research Assistant Backend
FastAPI backend for the haiku.rag interactive research assistant, using Pydantic AI with AG-UI protocol support.
Starlette backend for the haiku.rag interactive research assistant, using the research graph with AG-UI protocol support.
## Setup
@ -11,7 +11,17 @@ 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 with insight/gap tracking
- **AG-UI protocol**: Server-Sent Events (SSE) streaming for real-time state updates
- **Delta state updates**: Efficient incremental state synchronization using JSON Patch operations
- **Both research and deep_qa endpoints**: `/agent/research` and `/agent/deep_qa`
## Endpoints
- `GET /health` - Health check
- `POST /agent` - AG-UI protocol endpoint
- `GET /health` - Health check with configuration info
- `POST /agent/research/stream` - Research graph streaming endpoint (AG-UI protocol)
- `POST /agent/deep_qa/stream` - Deep QA graph streaming endpoint (AG-UI protocol)

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

@ -1,432 +1,111 @@
import json
from dataclasses import dataclass
"""Research assistant agent with graph integration."""
from dataclasses import dataclass
from pathlib import Path
from typing import TYPE_CHECKING
from ag_ui.core import EventType, StateSnapshotEvent
from pydantic import BaseModel
from pydantic_ai import Agent, RunContext
from pydantic_ai.ag_ui import StateDeps
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.graph_common import get_model
from haiku.rag.config import load_yaml_config
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.common import get_model
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_research_graph
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
if TYPE_CHECKING:
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.research.models import ResearchReport
class ResearchState(BaseModel):
"""Shared state between research agent and frontend."""
question: str = ""
phase: str = "idle"
status: str = ""
plan: list[dict] = []
current_question_index: int = 0
insights: list[dict] = []
document_registry: dict[str, dict] = {}
current_document: dict | None = None
confidence: float = 0.0
final_report: dict | None = None
# 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 ResearchDeps(StateDeps[ResearchState]):
"""Dependencies for the research agent with HaikuRAG client."""
class AgentDeps:
"""Dependencies for research agent."""
client: HaikuRAG
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
def _as_state_snapshot(ctx: RunContext[ResearchDeps]) -> StateSnapshotEvent:
return StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=ctx.deps.state)
model = get_model(Config.research.provider, Config.research.model)
def create_agent(
qa_provider: str = Config.qa.provider, qa_model: str = Config.qa.model
) -> Agent[ResearchDeps, str]:
"""Create and configure the research agent.
Args:
qa_provider: QA provider for the agent (default: from Config.qa.provider)
qa_model: Model name to use (default: from Config.qa.model)
"""
print(f"[AGENT SETUP] Creating agent with provider={qa_provider}, model={qa_model}")
agent = Agent(
model=get_model(qa_provider, qa_model),
deps_type=ResearchDeps,
instructions="""You are a research co-pilot powered by haiku.rag.
Your workflow MUST follow these exact steps in order:
1. Call propose_research_plan with the user's question
2. After propose_research_plan completes, IMMEDIATELY call approve_research_plan (with no arguments)
3. WAIT for approve_research_plan to return:
- If it returns "APPROVED", proceed to step 4
- If it returns "REVISE", ask the user "How would you like me to revise the research plan?" and wait for their response
- Once you receive their revision feedback, revise the plan and go back to step 1
4. Once approved, process questions ONE AT A TIME:
- Call search_question(question_id=0) and WAIT for it to complete
- Then call extract_insights_from_results(question_id=0) and WAIT for it to complete
- Then call search_question(question_id=1) and WAIT for it to complete
- Then call extract_insights_from_results(question_id=1) and WAIT for it to complete
- Then call search_question(question_id=2) and WAIT for it to complete
- Then call extract_insights_from_results(question_id=2) and WAIT for it to complete
5. After all questions are processed, call evaluate_research_confidence
6. Ask user if they want to finalize or continue researching
7. When user approves, call synthesize_final_report
agent = Agent(
model,
deps_type=AgentDeps,
system_prompt="""You are an advanced research assistant powered by haiku.rag.
CRITICAL RULES:
- MANDATORY: Call approve_research_plan immediately after propose_research_plan - NO EXCEPTIONS
- If approve_research_plan returns "REVISE", ask the user for revision feedback naturally in chat
- Call ONE tool at a time - wait for each tool to return before calling the next
- NEVER call extract_insights_from_results until search_question has completed and returned results
- DO NOT explain what you're about to do - just call the tool
- The state updates will show the user what's happening - you don't need to narrate
- Process all 3 questions automatically without asking for approval between them
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
Document Viewing:
- When user asks to "show document X", call get_full_document with the document_uri
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
Remember: Call tools ONE AT A TIME in sequence. Each tool must complete before calling the next.
""",
When you use run_research, the graph will decompose questions, search the knowledge base,
extract insights, 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}")
graph = build_research_graph(Config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config)
graph_deps = ResearchDeps(
client=ctx.deps.client,
agui_emitter=ctx.deps.agui_emitter,
)
@agent.tool
async def propose_research_plan(
ctx: RunContext[ResearchDeps], question: str
) -> StateSnapshotEvent:
"""Propose a research plan by decomposing the question into sub-questions."""
ctx.deps.state.question = question
ctx.deps.state.phase = "planning"
ctx.deps.state.status = "Decomposing question into sub-questions..."
try:
result = await graph.run(state=state, deps=graph_deps)
decompose_prompt = f"""Break down this research question into exactly 3 specific sub-questions that would help answer it comprehensively.
if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.log("✅ Research complete!")
Research Question: {question}
return f"""Research completed successfully!
Return ONLY a JSON array of sub-questions, like: ["Question 1?", "Question 2?", "Question 3?"]"""
Question: {question}
response = await ctx.deps.client.ask(decompose_prompt)
Executive Summary: {result.executive_summary}
try:
sub_questions = json.loads(response)
except json.JSONDecodeError:
sub_questions = [
q.strip().lstrip("0123456789.-) ")
for q in response.split("\n")
if q.strip()
][:3]
Main Findings:
{chr(10).join(f"- {finding}" for finding in result.main_findings[:3])}
plan = [
{"id": i, "question": q, "status": "pending"}
for i, q in enumerate(sub_questions)
]
Conclusions:
{chr(10).join(f"- {conclusion}" for conclusion in result.conclusions[:2])}
ctx.deps.state.plan = plan
ctx.deps.state.current_question_index = 0
ctx.deps.state.status = f"Proposed plan with {len(plan)} sub-questions"
Total insights gathered: {len(state.context.insights)}
Confidence: {f"{state.last_eval.confidence_score:.0%}" if state.last_eval else "N/A"}
Iterations completed: {state.iterations}
return _as_state_snapshot(ctx)
The full research report with all citations has been provided to the user.
"""
@agent.tool
async def search_question(
ctx: RunContext[ResearchDeps],
question_id: int,
search_type: str = "hybrid",
) -> StateSnapshotEvent:
"""Execute search for a specific sub-question."""
plan = ctx.deps.state.plan
if question_id >= len(plan):
raise ValueError(f"Question ID {question_id} not found in plan")
question = plan[question_id]["question"]
ctx.deps.state.phase = "searching"
ctx.deps.state.current_question_index = question_id
ctx.deps.state.status = f"Searching: {question}"
plan[question_id]["status"] = "searching"
search_results = await ctx.deps.client.search(
question, limit=5, search_type=search_type
)
expanded_map = {}
if search_results:
expanded_results = await ctx.deps.client.expand_context(
search_results[:3], radius=2
)
expanded_map = {
chunk.id: (chunk, score) for chunk, score in expanded_results
}
results = []
for chunk, score in search_results:
doc_uri = chunk.document_uri or "unknown"
doc_title = chunk.document_title or chunk.document_uri or "Unknown"
if doc_uri not in ctx.deps.state.document_registry:
ctx.deps.state.document_registry[doc_uri] = {
"title": doc_title,
"chunks_referenced": [],
}
if (
chunk.id
not in ctx.deps.state.document_registry[doc_uri]["chunks_referenced"]
):
ctx.deps.state.document_registry[doc_uri]["chunks_referenced"].append(
chunk.id
)
expanded_chunk, _ = (
expanded_map[chunk.id] if chunk.id in expanded_map else (chunk, score)
)
result_data = {
"chunk": expanded_chunk.content[:500],
"chunk_id": chunk.id,
"document_uri": doc_uri,
"document_title": doc_title,
"chunk_position": chunk.order,
"full_chunk_content": expanded_chunk.content,
"score": round(score, 3),
"expanded": chunk.id in expanded_map,
}
results.append(result_data)
plan[question_id]["search_results"] = {
"type": search_type,
"results": results,
}
plan[question_id]["status"] = "searched"
ctx.deps.state.status = f"Found {len(results)} results"
return _as_state_snapshot(ctx)
@agent.tool
async def extract_insights_from_results(
ctx: RunContext[ResearchDeps],
question_id: int,
) -> StateSnapshotEvent:
"""Extract key insights from search results for a specific question."""
plan = ctx.deps.state.plan
if question_id >= len(plan):
raise ValueError(f"Question ID {question_id} not found in plan")
question_item = plan[question_id]
if "search_results" not in question_item:
raise ValueError(
f"No search results found for question ID {question_id}. "
f"You must call search_question(question_id={question_id}) first."
)
search_results = question_item["search_results"]
ctx.deps.state.phase = "analyzing"
ctx.deps.state.status = "Extracting insights from results..."
context_parts = [
f"[Result {idx}] [Source: {r['document_title']}] {r['full_chunk_content']}"
for idx, r in enumerate(search_results["results"])
]
context = "\n\n".join(context_parts)
class InsightResult(BaseModel):
summary: str
confidence: float
result_indices: list[int]
class InsightsList(BaseModel):
insights: list[InsightResult]
question_text = question_item["question"]
extract_prompt = f"""Analyze these search results and extract 1-3 key insights that help answer the question: "{question_text}"
Search Results:
{context}
For each insight, reference which result numbers (0, 1, 2, etc.) support it."""
insight_agent: Agent[None, InsightsList] = Agent(
ctx.model,
output_type=InsightsList,
retries=3,
)
result = await insight_agent.run(extract_prompt)
raw_insights = [
{
"summary": insight.summary,
"confidence": insight.confidence,
"result_indices": insight.result_indices,
}
for insight in result.output.insights
]
new_insights = []
for insight in raw_insights:
source_refs = []
for idx in insight.get("result_indices", []):
if 0 <= idx < len(search_results["results"]):
result = search_results["results"][idx]
source_refs.append(
{
"chunk_id": result["chunk_id"],
"document_uri": result["document_uri"],
"document_title": result["document_title"],
"chunk_position": result["chunk_position"],
}
)
new_insights.append(
{
"summary": insight["summary"],
"confidence": insight.get("confidence", 0.7),
"source_refs": source_refs,
}
)
ctx.deps.state.insights.extend(new_insights)
plan[question_id]["status"] = "done"
ctx.deps.state.status = f"Extracted {len(new_insights)} insights"
return _as_state_snapshot(ctx)
@agent.tool
async def evaluate_research_confidence(
ctx: RunContext[ResearchDeps],
) -> StateSnapshotEvent:
"""Evaluate overall confidence in the research findings."""
insights = ctx.deps.state.insights
if not insights:
raise ValueError("No insights collected yet")
ctx.deps.state.phase = "evaluating"
ctx.deps.state.status = "Evaluating research confidence..."
confidences = [i.get("confidence", 0.5) for i in insights]
overall_confidence = sum(confidences) / len(confidences) if confidences else 0
eval_prompt = f"""Evaluate if these insights provide a confident answer to: "{ctx.deps.state.question}"
Insights collected:
{chr(10).join([f"- {i['summary']}" for i in insights])}
Assess:
1. Do we have enough information to answer the question?
2. What gaps remain?
3. Overall confidence (0.0-1.0)
Return JSON: {{"confidence": 0.0-1.0, "gaps": ["gap1", "gap2"], "recommendation": "continue" or "finalize"}}"""
response = await ctx.deps.client.ask(eval_prompt)
try:
evaluation = json.loads(response)
overall_confidence = evaluation.get("confidence", overall_confidence)
except json.JSONDecodeError:
pass
ctx.deps.state.confidence = overall_confidence
ctx.deps.state.status = f"Confidence: {overall_confidence:.0%}"
return _as_state_snapshot(ctx)
@agent.tool
async def synthesize_final_report(
ctx: RunContext[ResearchDeps],
) -> StateSnapshotEvent:
"""Generate final research report with citations."""
insights = ctx.deps.state.insights
if not insights:
raise ValueError("No insights to synthesize")
ctx.deps.state.phase = "synthesizing"
ctx.deps.state.status = "Generating final report..."
insights_summary = []
for i in insights:
source_titles = [ref["document_title"] for ref in i.get("source_refs", [])]
unique_sources = list(dict.fromkeys(source_titles))
insights_summary.append(
f"- {i['summary']} (sources: {', '.join(unique_sources[:2])})"
)
report_prompt = f"""Generate a comprehensive research report answering: "{ctx.deps.state.question}"
Based on these insights:
{chr(10).join(insights_summary)}
Create a structured report with:
- Executive Summary (2-3 sentences)
- Main Findings (bullet points)
- Conclusions
- Sources (list the document titles mentioned above)
Return JSON with format:
{{
"title": "...",
"summary": "...",
"findings": ["finding1", "finding2", ...],
"conclusions": ["conclusion1", ...],
"sources": ["source1", "source2", ...]
}}"""
response = await ctx.deps.client.ask(report_prompt)
try:
report = json.loads(response)
except json.JSONDecodeError:
report = {
"title": ctx.deps.state.question,
"summary": response[:300],
"findings": [i["summary"] for i in insights],
"conclusions": ["See findings above"],
"sources": [],
}
citations = [
{
"document_uri": doc_uri,
"document_title": doc_info["title"],
"chunk_ids": doc_info["chunks_referenced"],
}
for doc_uri, doc_info in ctx.deps.state.document_registry.items()
]
report["citations"] = citations
ctx.deps.state.final_report = report
ctx.deps.state.phase = "done"
ctx.deps.state.status = "Research complete"
return _as_state_snapshot(ctx)
@agent.tool
async def get_full_document(
ctx: RunContext[ResearchDeps],
document_uri: str,
) -> StateSnapshotEvent:
"""Retrieve and display the full content of a document by its URI."""
ctx.deps.state.status = f"Retrieving document: {document_uri}"
document = await ctx.deps.client.get_document_by_uri(document_uri)
if document is None:
ctx.deps.state.status = f"Document not found: {document_uri}"
ctx.deps.state.current_document = {
"uri": document_uri,
"title": "Not Found",
"content": f"Document with URI '{document_uri}' was not found.",
"total_chunks": 0,
}
else:
all_chunks = await ctx.deps.client.search(
query="", limit=1000, search_type="fts"
)
chunks_for_doc = [
c for c, _ in all_chunks if c.document_uri == document_uri
]
ctx.deps.state.current_document = {
"uri": document.uri or document_uri,
"title": document.title or "Untitled",
"content": document.content,
"total_chunks": len(chunks_for_doc),
"metadata": document.metadata,
}
ctx.deps.state.status = f"Retrieved: {document.title or document_uri}"
return _as_state_snapshot(ctx)
return agent
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)}"

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

@ -1,87 +1,171 @@
import logging
import os
from contextlib import asynccontextmanager
from pathlib import Path
from agent import ResearchDeps, ResearchState, create_agent
from agent import AgentDeps, 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.responses import JSONResponse
from starlette.routing import Mount, Route
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 Config
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 ResearchState
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
client: HaikuRAG | None = None
ag_ui_app = None
# 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()
@asynccontextmanager
async def lifespan(app):
global client
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
db_path = Path(db_path_str)
# 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}")
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 HaikuRAG client with database: {db_path}")
client = HaikuRAG(db_path)
logger.info("Research assistant backend ready")
logger.info(f"QA Provider: {Config.qa.provider}, Model: {Config.qa.model}")
logger.info(f"Initializing research assistant with database: {db_path}")
logger.info(
f"Research Provider: {Config.research.provider}, Model: {Config.research.model}"
)
yield
if client:
logger.info("Closing HaikuRAG client")
client.close()
# Store client reference for proper lifecycle management
_client_cache: dict[str, HaikuRAG] = {}
agent = create_agent()
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]
async def health(request):
db_path_str = os.getenv("DB_PATH", "haiku_rag.lancedb")
async def stream_research_agent(request: Request) -> StreamingResponse:
"""Agent streaming endpoint with research graph integration."""
body = await request.json()
input_data = RunAgentInput(**body)
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
emitter: AGUIEmitter[ResearchState, ResearchReport] = AGUIEmitter(
thread_id=input_data.thread_id,
run_id=input_data.run_id,
use_deltas=False,
)
# 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)
# Create agent dependencies with shared emitter
agent_deps = AgentDeps(client=client, agui_emitter=emitter)
# 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:
# Filter out ACTIVITY_SNAPSHOT - not supported by CopilotKit
if event.get("type") == "ACTIVITY_SNAPSHOT":
continue
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)
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),
"qa_provider": Config.qa.provider,
"qa_model": Config.qa.model,
"ollama_base_url": Config.providers.ollama.base_url,
"db_path": db_path_str,
"db_exists": Path(db_path_str).exists(),
"research_provider": Config.research.provider,
"research_model": Config.research.model,
"db_path": str(db_path),
"db_exists": db_path.exists(),
}
)
def get_ag_ui_app():
global ag_ui_app
if ag_ui_app is None and client is not None:
research_deps = ResearchDeps(client=client, state=ResearchState())
logger.info("Creating AG-UI app")
ag_ui_app = agent.to_ag_ui(deps=research_deps)
return ag_ui_app
async def agent_endpoint(scope, receive, send):
app = get_ag_ui_app()
if app is None:
response = JSONResponse({"error": "Client not initialized"}, status_code=503)
await response(scope, receive, send)
return
await app(scope, receive, send)
# Create Starlette app
app = Starlette(
routes=[
Route("/health", health),
Mount("/agent", agent_endpoint),
Route("/v1/research/stream", stream_research_agent, methods=["POST"]),
Route("/health", health_check, methods=["GET"]),
],
middleware=[
Middleware(
@ -92,17 +176,11 @@ app = Starlette(
allow_headers=["*"],
)
],
lifespan=lifespan,
)
if __name__ == "__main__":
import uvicorn
print("Starting haiku.rag research assistant backend...")
print(f"Agent model: {agent.model}")
print(f"QA provider: {Config.qa.provider}")
print(f"QA model: {Config.qa.model}")
uvicorn.run(
"main:app",
host="0.0.0.0",

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

@ -9,7 +9,7 @@ dependencies = [
"uvicorn[standard]>=0.34.2",
"pydantic-ai-slim[ag-ui,openai]>=1.1.0",
"python-dotenv>=1.0.1",
"haiku-rag>=0.12.1",
"haiku-rag-slim @ file:///Users/ggozad/dev/open-source/haiku.rag-agui/haiku_rag_slim",
]
[dependency-groups]
@ -18,6 +18,9 @@ dev = [
"ruff>=0.13.0",
]
[tool.hatch.metadata]
allow-direct-references = true
[tool.hatch.build.targets.wheel]
packages = ["."]

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

File diff suppressed because it is too large Load diff

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

@ -10,11 +10,16 @@ services:
# 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
volumes:
- ./backend:/app
- /app/.venv
- ${DB_PATH}:/app/data/haiku.rag.lancedb
- ./haiku.rag.yaml:/app/haiku.rag.yaml:ro
- ./backend/main.py:/app/main.py
- ./backend/agent.py:/app/agent.py
- ../../haiku_rag_slim/haiku:/app/.venv/lib/python3.13/site-packages/haiku
networks:
- ag-ui-network
extra_hosts:

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

@ -13,7 +13,7 @@ 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"}/agent`,
url: `${process.env.BACKEND_URL || "http://backend:8000"}/v1/research/stream`,
}),
},
});

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

@ -1,208 +1,94 @@
"use client";
import {
CopilotKit,
useCoAgent,
useCoAgentStateRender,
useCopilotAction,
} from "@copilotkit/react-core";
import { CopilotKit, useCoAgent } from "@copilotkit/react-core";
import { CopilotChat } from "@copilotkit/react-ui";
import "@copilotkit/react-ui/styles.css";
import StateDisplay from "./StateDisplay";
interface SourceRef {
chunk_id: string;
document_uri: string;
document_title: string;
chunk_position: number;
interface InsightRecord {
id: string;
summary: string;
status: string;
notes?: string;
supporting_sources: string[];
originating_questions: string[];
}
interface GapRecord {
id: string;
description: string;
severity: string;
blocking: boolean;
resolved: boolean;
notes?: string;
supporting_sources: string[];
resolved_by: string[];
}
interface SearchAnswer {
query: string;
answer: string;
confidence: number;
context: string[];
sources: string[];
}
interface ResearchContext {
original_question: string;
sub_questions: string[];
qa_responses: SearchAnswer[];
insights: InsightRecord[];
gaps: GapRecord[];
}
interface EvaluationResult {
confidence: number;
reasoning: string;
should_continue: boolean;
gaps_identified: string[];
follow_up_questions: string[];
}
interface ResearchReport {
question: string;
summary: string;
findings: string[];
conclusions: string[];
insights_used: string[];
methodology: string;
}
interface ResearchState {
question: string;
phase: string;
status: string;
plan: Array<{
id: number;
question: string;
status: string;
search_results?: {
type: string;
results: Array<{
chunk: string;
chunk_id: string;
document_uri: string;
document_title: string;
chunk_position: number;
full_chunk_content: string;
score: number;
expanded: boolean;
}>;
};
}>;
current_question_index: number;
insights: Array<{
summary: string;
confidence: number;
source_refs: SourceRef[];
}>;
document_registry: Record<
string,
{
title: string;
chunks_referenced: string[];
}
>;
current_document: {
uri: string;
title: string;
content: string;
total_chunks: number;
metadata?: Record<string, unknown>;
} | null;
confidence: number;
final_report: {
title: string;
summary: string;
findings: string[];
conclusions: string[];
sources: string[];
citations: Array<{
document_uri: string;
document_title: string;
chunk_ids: string[];
}>;
context: ResearchContext;
iterations: number;
max_iterations: number;
confidence_threshold: number;
max_concurrency: number;
last_eval: EvaluationResult | null;
last_analysis: {
insights_extracted: InsightRecord[];
gaps_identified: GapRecord[];
} | null;
result?: ResearchReport;
}
function AgentContent() {
const { state } = useCoAgent<ResearchState>({
name: "research_agent",
initialState: {
question: "",
phase: "idle",
status: "",
plan: [],
current_question_index: 0,
insights: [],
document_registry: {},
current_document: null,
confidence: 0.0,
final_report: null,
},
});
useCopilotAction({
name: "approve_research_plan",
description:
"Request user approval for the research plan. Returns 'APPROVED' if approved or 'REVISE' if user wants to revise.",
parameters: [],
renderAndWaitForResponse: ({ respond, status }) => (
<div
style={{
padding: "1.5rem",
background: "white",
borderRadius: "8px",
border: "2px solid #4299e1",
marginBottom: "1rem",
boxShadow: "0 2px 8px rgba(0,0,0,0.1)",
}}
>
<h3
style={{
fontSize: "1.25rem",
fontWeight: "bold",
marginBottom: "1rem",
color: "#2d3748",
}}
>
Research Plan Approval
</h3>
<p
style={{
fontSize: "0.875rem",
color: "#4a5568",
marginBottom: "1rem",
}}
>
Please review the research plan in the right pane.
</p>
<div
style={{
display: "flex",
gap: "1rem",
}}
className={status !== "executing" ? "hidden" : ""}
>
<button
type="button"
onClick={() => respond?.("REVISE")}
disabled={status !== "executing"}
style={{
flex: 1,
padding: "0.75rem",
background: "white",
border: "2px solid #e2e8f0",
borderRadius: "6px",
fontSize: "0.875rem",
fontWeight: "600",
cursor: status === "executing" ? "pointer" : "not-allowed",
opacity: status === "executing" ? 1 : 0.5,
}}
>
Revise Plan
</button>
<button
type="button"
onClick={() => respond?.("APPROVED")}
disabled={status !== "executing"}
style={{
flex: 1,
padding: "0.75rem",
background: "#4299e1",
color: "white",
border: "none",
borderRadius: "6px",
fontSize: "0.875rem",
fontWeight: "600",
cursor: status === "executing" ? "pointer" : "not-allowed",
opacity: status === "executing" ? 1 : 0.5,
}}
>
Approve & Start Research
</button>
</div>
</div>
),
});
useCoAgentStateRender<ResearchState>({
name: "research_agent",
render: ({ state: newState }) => {
const phaseMessages: Record<string, string> = {
planning: "Planning research...",
searching: "Searching...",
analyzing: "Extracting insights...",
evaluating: `Evaluating confidence: ${(newState.confidence * 100).toFixed(0)}%`,
synthesizing: "Generating final report...",
done: "Research complete!",
};
const phaseMessage =
phaseMessages[newState.phase] || newState.status || "Ready";
return (
<div
style={{
padding: "1rem",
background: "#e6f7ff",
borderRadius: "4px",
marginBottom: "0.5rem",
border: "1px solid #91d5ff",
}}
>
<strong>Research Update:</strong> {phaseMessage}
</div>
);
context: {
original_question: "",
sub_questions: [],
qa_responses: [],
insights: [],
gaps: [],
},
iterations: 0,
max_iterations: 3,
confidence_threshold: 0.8,
max_concurrency: 1,
last_eval: null,
last_analysis: null,
},
});

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

File diff suppressed because it is too large Load diff

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

File diff suppressed because it is too large Load diff

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

@ -1,20 +1,23 @@
# haiku.rag configuration for ag-ui-research example
# Copy to haiku.rag.yaml and customize
qa:
research:
provider: ollama
model: gpt-oss:latest
max_iterations: 3
confidence_threshold: 0.8
max_concurrency: 1
providers:
ollama:
base_url: http://host.docker.internal:11434
# For OpenAI:
# qa:
# research:
# provider: openai
# model: gpt-4o-mini
# For Anthropic:
# qa:
# research:
# provider: anthropic
# model: claude-3-5-haiku-20241022

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

@ -10,12 +10,12 @@ from rich.progress import Progress
from haiku.rag.client import HaikuRAG
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
from haiku.rag.mcp import create_mcp_server
from haiku.rag.monitor import FileWatcher
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
from haiku.rag.research.stream import stream_research_graph
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
@ -216,23 +216,31 @@ class HaikuRAGApp:
async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client:
try:
if deep:
from rich.console import Console
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
graph = build_deep_qa_graph(config=self.config)
context = DeepQAContext(
original_question=question, use_citations=cite
)
state = DeepQAState.from_config(context=context, config=self.config)
deps = DeepQADeps(
client=self.client, console=Console() if verbose else None
)
deps = DeepQADeps(client=self.client)
result = await graph.run(state=state, deps=deps)
answer = result.answer
if verbose:
# Use AG-UI renderer to process and display events
from haiku.rag.graph.agui import AGUIConsoleRenderer
renderer = AGUIConsoleRenderer(self.console)
result_dict = await renderer.render(
stream_graph(graph, state, deps)
)
# Result should be a dict with 'answer' key
answer = result_dict.get("answer", "") if result_dict else ""
else:
# Run without rendering events, just get the result
result = await graph.run(state=state, deps=deps)
answer = result.answer
else:
answer = await self.client.ask(question, cite=cite)
@ -243,46 +251,46 @@ class HaikuRAGApp:
except Exception as e:
self.console.print(f"[red]Error: {e}[/red]")
async def research(
self,
question: str,
verbose: bool = False,
):
async def research(self, question: str, verbose: bool = False):
"""Run research via the pydantic-graph pipeline.
Args:
question: The research question
verbose: Show verbose output
verbose: Show AG-UI event stream during execution
"""
async with HaikuRAG(db_path=self.db_path, config=self.config) as client:
try:
if verbose:
self.console.print("[bold cyan]Starting research[/bold cyan]")
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
self.console.print("[bold cyan]Starting research[/bold cyan]")
self.console.print(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
graph = build_research_graph(config=self.config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=self.config)
deps = ResearchDeps(
client=client, console=self.console if verbose else None
)
deps = ResearchDeps(client=client)
report = None
async for event in stream_research_graph(graph, state, deps):
if event.type == "report":
report = event.report
break
if event.type == "error":
self.console.print(
f"[red]Error during research: {event.message}[/red]"
)
return
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
)
if report is None:
if report_dict 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()
@ -455,6 +463,7 @@ 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(self.db_path, config=self.config) as client:
@ -482,6 +491,30 @@ 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

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

@ -302,7 +302,7 @@ def research(
verbose: bool = typer.Option(
False,
"--verbose",
help="Show verbose progress output",
help="Show planning, searching previews, evaluation summary, and stop reason",
),
):
app = create_app(db)
@ -401,7 +401,7 @@ def download_models_cmd():
@cli.command(
"serve",
help="Start haiku.rag server. Use --monitor and/or --mcp to enable services.",
help="Start haiku.rag server. Use --monitor, --mcp, and/or --agui to enable services.",
)
def serve(
db: Path | None = typer.Option(
@ -429,12 +429,17 @@ 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):
if not (monitor or mcp or agui):
typer.echo(
"Error: At least one service flag (--monitor or --mcp) must be specified"
"Error: At least one service flag (--monitor, --mcp, or --agui) must be specified"
)
raise typer.Exit(1)
@ -452,6 +457,7 @@ def serve(
enable_mcp=mcp,
mcp_transport=transport,
mcp_port=mcp_port,
enable_agui=agui,
)
)

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

@ -6,6 +6,7 @@ from haiku.rag.config.loader import (
load_yaml_config,
)
from haiku.rag.config.models import (
AGUIConfig,
AppConfig,
EmbeddingsConfig,
LanceDBConfig,
@ -22,6 +23,7 @@ from haiku.rag.config.models import (
__all__ = [
"Config",
"AGUIConfig",
"AppConfig",
"StorageConfig",
"MonitorConfig",

8
haiku_rag_slim/haiku/rag/config/loader.py
View file

@ -84,4 +84,12 @@ def generate_default_config() -> dict:
"research_base_url": "",
},
},
"agui": {
"host": "0.0.0.0",
"port": 8000,
"cors_origins": ["*"],
"cors_credentials": True,
"cors_methods": ["GET", "POST", "OPTIONS"],
"cors_headers": ["*"],
},
}

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

@ -77,6 +77,15 @@ class ProvidersConfig(BaseModel):
vllm: VLLMConfig = Field(default_factory=VLLMConfig)
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 AppConfig(BaseModel):
environment: str = "production"
storage: StorageConfig = Field(default_factory=StorageConfig)
@ -88,3 +97,4 @@ class AppConfig(BaseModel):
research: ResearchConfig = Field(default_factory=ResearchConfig)
processing: ProcessingConfig = Field(default_factory=ProcessingConfig)
providers: ProvidersConfig = Field(default_factory=ProvidersConfig)
agui: AGUIConfig = Field(default_factory=AGUIConfig)

26
haiku_rag_slim/haiku/rag/graph/__init__.py Normal file
View file

@ -0,0 +1,26 @@
"""Graph module for haiku.rag.
This module contains all graph-related functionality including:
- AG-UI protocol for graph streaming
- Common graph utilities and models
- Research graph implementation
- Deep QA graph implementation
"""
from haiku.rag.graph.agui import (
AGUIConsoleRenderer,
AGUIEmitter,
create_agui_server,
stream_graph,
)
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.research.graph import build_research_graph
__all__ = [
"AGUIConsoleRenderer",
"AGUIEmitter",
"build_deep_qa_graph",
"build_research_graph",
"create_agui_server",
"stream_graph",
]

53
haiku_rag_slim/haiku/rag/graph/agui/__init__.py Normal file
View file

@ -0,0 +1,53 @@
"""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
from haiku.rag.graph.agui.events import (
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,
)
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",
"format_sse_event",
"stream_graph",
]

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

@ -0,0 +1,135 @@
"""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.events 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")

197
haiku_rag_slim/haiku/rag/graph/agui/emitter.py Normal file
View file

@ -0,0 +1,197 @@
"""Generic AG-UI event emitter for any graph execution."""
import asyncio
import hashlib
from collections.abc import AsyncIterator
from uuid import uuid4
from pydantic import BaseModel
from haiku.rag.graph.agui.events import (
AGUIEvent,
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 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._current_step: str | None = None
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(emit_run_started(self._thread_id, self._run_id))
self._emit(emit_state_snapshot(initial_state))
# 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._current_step = step_name
self._emit(emit_step_started(step_name))
def finish_step(self) -> None:
"""Emit StepFinished event for the current step."""
if self._current_step:
self._emit(emit_step_finished(self._current_step))
self._current_step = None
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)
"""
self._emit(emit_text_message(message, role))
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
self._emit(emit_state_delta(self._last_state, new_state))
else:
# Emit full snapshot for initial state or when deltas disabled
self._emit(emit_state_snapshot(new_state))
# 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: str, message_id: str | None = None
) -> None:
"""Emit ActivitySnapshot event.
Args:
activity_type: Type of activity (e.g., "planning", "searching")
content: Description of the activity
message_id: Optional message ID to associate activity with (auto-generated if None)
"""
if message_id is None:
message_id = str(uuid4())
self._emit(emit_activity(message_id, activity_type, content))
def finish_run(self, result: ResultT) -> None:
"""Emit RunFinished event.
Args:
result: The final result from the graph
"""
self._emit(emit_run_finished(self._thread_id, self._run_id, result))
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(emit_run_error(str(error), 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]

254
haiku_rag_slim/haiku/rag/graph/agui/events.py Normal file
View file

@ -0,0 +1,254 @@
"""Generic AG-UI event creation utilities for any graph."""
from typing import Any
from uuid import uuid4
from pydantic import BaseModel
from haiku.rag.graph.agui.state import compute_state_delta
# Type aliases for AG-UI events (actual types from ag_ui.core will be used at runtime)
AGUIEvent = dict[str, Any]
def emit_run_started(
thread_id: str, run_id: str, input_data: str | None = None
) -> dict[str, Any]:
"""Create a RunStarted event.
Args:
thread_id: Unique identifier for the conversation thread
run_id: Unique identifier for this run
input_data: Optional input that started the run
Returns:
RunStarted event dict
"""
event: dict[str, Any] = {
"type": "RUN_STARTED",
"threadId": thread_id,
"runId": run_id,
}
if input_data:
event["input"] = input_data
return event
def emit_run_finished(thread_id: str, run_id: str, result: Any) -> dict[str, Any]:
"""Create a RunFinished event.
Args:
thread_id: Unique identifier for the conversation thread
run_id: Unique identifier for this run
result: The final result of the run
Returns:
RunFinished event dict
"""
# Convert result to dict if it's a Pydantic model
if hasattr(result, "model_dump"):
result = result.model_dump()
return {
"type": "RUN_FINISHED",
"threadId": thread_id,
"runId": run_id,
"result": result,
}
def emit_run_error(message: str, code: str | None = None) -> dict[str, Any]:
"""Create a RunError event.
Args:
message: Error message
code: Optional error code
Returns:
RunError event dict
"""
event: dict[str, Any] = {
"type": "RUN_ERROR",
"message": message,
}
if code:
event["code"] = code
return event
def emit_step_started(step_name: str) -> dict[str, Any]:
"""Create a StepStarted event.
Args:
step_name: Name of the step being started
Returns:
StepStarted event dict
"""
return {
"type": "STEP_STARTED",
"stepName": step_name,
}
def emit_step_finished(step_name: str) -> dict[str, Any]:
"""Create a StepFinished event.
Args:
step_name: Name of the step that finished
Returns:
StepFinished event dict
"""
return {
"type": "STEP_FINISHED",
"stepName": step_name,
}
def emit_text_message(content: str, role: str = "assistant") -> dict[str, Any]:
"""Create a TextMessageChunk event (convenience wrapper).
This creates a complete text message in one event.
Args:
content: The message content
role: The role of the sender (default: assistant)
Returns:
TextMessageChunk event dict
"""
message_id = str(uuid4())
return {
"type": "TEXT_MESSAGE_CHUNK",
"messageId": message_id,
"role": role,
"delta": content,
}
def emit_text_message_start(message_id: str, role: str = "assistant") -> dict[str, Any]:
"""Create a TextMessageStart event.
Args:
message_id: Unique identifier for this message
role: The role of the sender
Returns:
TextMessageStart event dict
"""
return {
"type": "TEXT_MESSAGE_START",
"messageId": message_id,
"role": role,
}
def emit_text_message_content(message_id: str, delta: str) -> dict[str, Any]:
"""Create a TextMessageContent event.
Args:
message_id: Identifier for the message being streamed
delta: Content chunk to append
Returns:
TextMessageContent event dict
"""
return {
"type": "TEXT_MESSAGE_CONTENT",
"messageId": message_id,
"delta": delta,
}
def emit_text_message_end(message_id: str) -> dict[str, Any]:
"""Create a TextMessageEnd event.
Args:
message_id: Identifier for the message being completed
Returns:
TextMessageEnd event dict
"""
return {
"type": "TEXT_MESSAGE_END",
"messageId": message_id,
}
def emit_state_snapshot(state: BaseModel) -> dict[str, Any]:
"""Create a StateSnapshot event.
Args:
state: The complete state to snapshot (any Pydantic BaseModel)
Returns:
StateSnapshot event dict
"""
return {
"type": "STATE_SNAPSHOT",
"snapshot": state.model_dump(),
}
def emit_state_delta(old_state: BaseModel, new_state: BaseModel) -> dict[str, Any]:
"""Create a StateDelta event with JSON Patch operations.
Args:
old_state: Previous state (any Pydantic BaseModel)
new_state: Current state (same type as old_state)
Returns:
StateDelta event dict
"""
delta = compute_state_delta(old_state, new_state)
return {
"type": "STATE_DELTA",
"delta": delta,
}
def emit_activity(
message_id: str,
activity_type: str,
content: str,
) -> dict[str, Any]:
"""Create an ActivitySnapshot event.
Args:
message_id: Message ID to associate activity with (required)
activity_type: Type of activity (e.g., "planning", "searching")
content: Description of the activity
Returns:
ActivitySnapshot event dict
"""
return {
"type": "ACTIVITY_SNAPSHOT",
"messageId": message_id,
"activityType": activity_type,
"content": content,
}
def emit_activity_delta(
message_id: str,
activity_type: str,
patch: list[dict[str, Any]],
) -> dict[str, Any]:
"""Create an ActivityDelta event with JSON Patch operations.
Args:
message_id: Message ID of the activity being updated
activity_type: Type of activity being updated
patch: JSON Patch operations to apply
Returns:
ActivityDelta event dict
"""
return {
"type": "ACTIVITY_DELTA",
"messageId": message_id,
"activityType": activity_type,
"patch": patch,
}

310
haiku_rag_slim/haiku/rag/graph/agui/server.py Normal file
View file

@ -0,0 +1,310 @@
"""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
from haiku.rag.graph.agui.events import 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(config: "AppConfig", db_path: Path | None = None) -> Starlette:
"""Create AG-UI server with both research and deep ask endpoints.
Args:
config: Application config with research and qa settings
db_path: Optional database path override
Returns:
Starlette app with research and deep ask endpoints
"""
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
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)
return ResearchState.from_config(context=context, config=config)
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))
# Deep ask graph factories
def deep_ask_graph_factory() -> Graph:
return build_deep_qa_graph(config)
def deep_ask_state_factory(input_state: dict[str, Any]) -> DeepQAState:
question = input_state.get("question", "")
if not question:
messages = input_state.get("messages", [])
if messages:
question = messages[0].get("content", "")
use_citations = input_state.get("use_citations", False)
context = DeepQAContext(original_question=question, use_citations=use_citations)
return DeepQAState.from_config(context=context, config=config)
def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps:
effective_db_path = (
db_path
or input_config.get("db_path")
or config.storage.data_dir / "haiku.rag.lancedb"
)
return DeepQADeps(client=get_client(effective_db_path))
# Create event stream functions for each graph type
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
async def deep_ask_event_stream(
input_data: RunAgentInput,
) -> AsyncIterator[str]:
"""Generate SSE event stream from deep ask graph execution."""
graph = deep_ask_graph_factory()
initial_state = deep_ask_state_factory(input_data.state)
deps = deep_ask_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 stream_deep_ask(request: Request) -> StreamingResponse:
"""Deep ask graph streaming endpoint."""
body = await request.json()
input_data = RunAgentInput(**body)
return StreamingResponse(
deep_ask_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("/v1/deep-ask/stream", stream_deep_ask, 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 Normal file
View file

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

86
haiku_rag_slim/haiku/rag/graph/agui/stream.py Normal file
View file

@ -0,0 +1,86 @@
"""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
from haiku.rag.graph.agui.events import 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

5
haiku_rag_slim/haiku/rag/graph/common/__init__.py Normal file
View file

@ -0,0 +1,5 @@
"""Common utilities for graph implementations."""
from haiku.rag.graph.common.utils import get_model
__all__ = ["get_model"]

0
haiku_rag_slim/haiku/rag/graph_common/models.py → haiku_rag_slim/haiku/rag/graph/common/models.py
View file

265
haiku_rag_slim/haiku/rag/graph/common/nodes.py Normal file
View file

@ -0,0 +1,265 @@
"""Common node implementations for graph workflows."""
import asyncio
from collections.abc import Awaitable, Callable
from typing import Any, Protocol
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import StepContext
from haiku.rag.client import HaikuRAG
from haiku.rag.graph.agui.emitter import AGUIEmitter
from haiku.rag.graph.common import get_model
from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
class GraphContext(Protocol):
"""Protocol for graph context objects."""
original_question: str
sub_questions: list[str]
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a QA response to context."""
...
class GraphState(Protocol):
"""Protocol for graph state objects."""
context: GraphContext
max_concurrency: int
class GraphDeps(Protocol):
"""Protocol for graph dependencies."""
client: HaikuRAG
agui_emitter: AGUIEmitter[Any, Any] | None
semaphore: asyncio.Semaphore | None
class GraphAgentDeps(Protocol):
"""Protocol for agent dependencies."""
client: HaikuRAG
context: GraphContext
def create_plan_node[AgentDepsT: GraphAgentDeps](
provider: str,
model: str,
deps_type: type[AgentDepsT],
activity_message: str = "Creating plan",
output_retries: int | None = None,
) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]:
"""Create a plan node for any graph.
Args:
provider: Model provider (e.g., 'openai', 'anthropic')
model: Model name
deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
activity_message: Message to show during planning activity
output_retries: Number of output retries for the agent (optional)
Returns:
Async function that can be used as a graph step
"""
async def plan(ctx: StepContext[Any, Any, None], /) -> None:
state: GraphState = ctx.state # type: ignore[assignment]
deps: GraphDeps = ctx.deps # type: ignore[assignment]
if deps.agui_emitter:
deps.agui_emitter.start_step("plan")
deps.agui_emitter.update_activity("planning", activity_message)
try:
# Build agent configuration
agent_config = {
"model": get_model(provider, model),
"output_type": ResearchPlan,
"instructions": (
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
"retries": 3,
"deps_type": deps_type,
}
if output_retries is not None:
agent_config["output_retries"] = output_retries
plan_agent = Agent(**agent_config)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[AgentDepsT], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
# Tool is registered via decorator above
_ = gather_context
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
# Create agent dependencies
agent_deps = deps_type(client=deps.client, context=state.context) # type: ignore[call-arg]
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
# State now contains the plan - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
count = len(state.context.sub_questions)
deps.agui_emitter.update_activity(
"planning", f"Created plan with {count} sub-questions"
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
return plan
def create_search_node[AgentDepsT: GraphAgentDeps](
provider: str,
model: str,
deps_type: type[AgentDepsT],
with_step_wrapper: bool = True,
success_message_format: str = "Answered: {sub_q}",
handle_exceptions: bool = False,
) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]:
"""Create a search_one node for any graph.
Args:
provider: Model provider
model: Model name
deps_type: Type of dependencies for the agent
with_step_wrapper: Whether to wrap with agui_emitter start/finish step
success_message_format: Format string for success activity message
handle_exceptions: Whether to handle exceptions with fallback answer
Returns:
Async function that can be used as a graph step
"""
async def search_one(ctx: StepContext[Any, Any, str], /) -> SearchAnswer:
state: GraphState = ctx.state # type: ignore[assignment]
deps: GraphDeps = ctx.deps # type: ignore[assignment]
sub_q = ctx.inputs
# Create unique step name from question text
step_name = f"search: {sub_q}"
if deps.agui_emitter and with_step_wrapper:
deps.agui_emitter.start_step(step_name)
try:
# Create semaphore if not already provided
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
# Use semaphore to control concurrency
async with deps.semaphore:
return await _do_search(
state,
deps,
sub_q,
provider,
model,
deps_type,
success_message_format,
handle_exceptions,
)
finally:
if deps.agui_emitter and with_step_wrapper:
deps.agui_emitter.finish_step()
return search_one
async def _do_search[AgentDepsT: GraphAgentDeps](
state: GraphState,
deps: GraphDeps,
sub_q: str,
provider: str,
model: str,
deps_type: type[AgentDepsT],
success_message_format: str,
handle_exceptions: bool,
) -> SearchAnswer:
"""Internal search implementation."""
if deps.agui_emitter:
deps.agui_emitter.update_activity("searching", f"Searching: {sub_q}")
agent = Agent(
model=get_model(provider, model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=deps_type,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[AgentDepsT], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_title or chunk.document_uri or ""),
}
for chunk, score in expanded
]
if not entries:
return f"No relevant information found in the knowledge base for: {query}"
return format_as_xml(entries, root_tag="snippets")
# Tool is registered via decorator above
_ = search_and_answer
agent_deps = deps_type(client=deps.client, context=state.context) # type: ignore[call-arg]
try:
result = await agent.run(sub_q, deps=agent_deps)
answer = result.output
if answer:
state.context.add_qa_response(answer)
# State updated with new answer - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
# Format the success message
if "{confidence" in success_message_format:
message = success_message_format.format(
sub_q=sub_q, confidence=answer.confidence
)
else:
message = success_message_format.format(sub_q=sub_q)
deps.agui_emitter.update_activity("searching", message)
return answer
except Exception as e:
if handle_exceptions:
# Narrate the error
if deps.agui_emitter:
deps.agui_emitter.update_activity("searching", f"Search failed: {e}")
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",
confidence=0.0,
)
return failure_answer
else:
raise

0
haiku_rag_slim/haiku/rag/graph_common/prompts.py → haiku_rag_slim/haiku/rag/graph/common/prompts.py
View file

20
haiku_rag_slim/haiku/rag/graph_common/utils.py → haiku_rag_slim/haiku/rag/graph/common/utils.py
View file

@ -1,7 +1,5 @@
"""Common utilities for all graph implementations."""
from typing import Any, Protocol
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.ollama import OllamaProvider
from pydantic_ai.providers.openai import OpenAIProvider
@ -9,12 +7,6 @@ from pydantic_ai.providers.openai import OpenAIProvider
from haiku.rag.config import Config
class HasEmitLog(Protocol):
"""Protocol for objects that can emit log messages."""
def emit_log(self, message: str, state: Any = None) -> None: ...
def get_model(provider: str, model: str) -> OpenAIChatModel | str:
"""
Get a model instance for the specified provider and model name.
@ -50,15 +42,3 @@ def get_model(provider: str, model: str) -> OpenAIChatModel | str:
f"Unknown model provider: {provider}. "
f"Supported providers: ollama, vllm, openai, anthropic, gemini, groq, bedrock"
)
def log(deps: HasEmitLog, state: Any, message: str) -> None:
"""
Emit a log message through the dependencies.
Args:
deps: Dependencies object with emit_log method
state: Current state (passed to emit_log)
message: The message to log
"""
deps.emit_log(message, state)

1
haiku_rag_slim/haiku/rag/graph/deep_qa/__init__.py Normal file
View file

@ -0,0 +1 @@
from haiku.rag.graph.deep_qa.models import DeepQAAnswer

4
haiku_rag_slim/haiku/rag/qa/deep/dependencies.py → haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py
View file

@ -1,8 +1,7 @@
from pydantic import BaseModel, Field
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.graph_common.models import SearchAnswer
from haiku.rag.graph.common.models import SearchAnswer
class DeepQAContext(BaseModel):
@ -26,4 +25,3 @@ class DeepQADependencies(BaseModel):
client: HaikuRAG = Field(description="RAG client for document operations")
context: DeepQAContext = Field(description="Shared QA context")
console: Console | None = None

243
haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py Normal file
View file

@ -0,0 +1,243 @@
from pydantic_ai import Agent
from pydantic_ai.format_prompt import format_as_xml
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
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.common import get_model
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies
from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation
from haiku.rag.graph.deep_qa.prompts import (
DECISION_PROMPT,
SYNTHESIS_PROMPT,
SYNTHESIS_PROMPT_WITH_CITATIONS,
)
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
def build_deep_qa_graph(
config: AppConfig = Config,
) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]:
"""Build the Deep QA graph.
Args:
config: AppConfig object (uses config.qa for provider, model, and graph parameters)
Returns:
Configured Deep QA graph
"""
provider = config.qa.provider
model = config.qa.model
g = GraphBuilder(
state_type=DeepQAState,
deps_type=DeepQADeps,
output_type=DeepQAAnswer,
)
# Create and register the plan node using the factory
plan = g.step(
create_plan_node(
provider=provider,
model=model,
deps_type=DeepQADependencies, # type: ignore[arg-type]
activity_message="Planning approach",
output_retries=None, # Deep QA doesn't use output_retries
)
) # type: ignore[arg-type]
# Create and register the search_one node using the factory
search_one = g.step(
create_search_node(
provider=provider,
model=model,
deps_type=DeepQADependencies, # type: ignore[arg-type]
with_step_wrapper=False, # Deep QA doesn't wrap with agui_emitter step
success_message_format="Answered: {sub_q}",
handle_exceptions=True,
)
) # type: ignore[arg-type]
@g.step
async def get_batch(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
state = ctx.state
if not state.context.sub_questions:
return None
# Take ALL remaining questions - max_concurrency controls parallel execution within .map()
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def decide(
ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]],
) -> bool:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("decide")
deps.agui_emitter.update_activity(
"evaluating", "Evaluating information sufficiency"
)
try:
agent = Agent(
model=get_model(provider, model),
output_type=DeepQAEvaluation,
instructions=DECISION_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"gathered_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = (
"Evaluate whether we have sufficient information to answer the question.\n\n"
f"{context_xml}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
evaluation = result.output
state.iterations += 1
for new_q in evaluation.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
status = "sufficient" if evaluation.is_sufficient else "insufficient"
deps.agui_emitter.update_activity(
"evaluating",
f"Information {status} after {state.iterations} iteration(s)",
)
should_continue = (
not evaluation.is_sufficient and state.iterations < state.max_iterations
)
return should_continue
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def synthesize(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> DeepQAAnswer:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("synthesize")
deps.agui_emitter.update_activity(
"synthesizing", "Synthesizing final answer"
)
try:
prompt_template = (
SYNTHESIS_PROMPT_WITH_CITATIONS
if state.context.use_citations
else SYNTHESIS_PROMPT
)
agent = Agent(
model=get_model(provider, model),
output_type=DeepQAAnswer,
instructions=prompt_template,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"sub_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
if deps.agui_emitter:
deps.agui_emitter.update_activity("synthesizing", "Answer complete")
return result.output
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
# Build the graph structure
collect_answers = g.join(
reduce_list_append,
initial_factory=list[SearchAnswer],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
# Branch based on whether we have questions
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),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done with QA")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

0
haiku_rag_slim/haiku/rag/qa/deep/models.py → haiku_rag_slim/haiku/rag/graph/deep_qa/models.py
View file

0
haiku_rag_slim/haiku/rag/qa/deep/prompts.py → haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py
View file

33
haiku_rag_slim/haiku/rag/qa/deep/state.py → haiku_rag_slim/haiku/rag/graph/deep_qa/state.py
View file

@ -2,33 +2,40 @@ import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING
from rich.console import Console
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
@dataclass
class DeepQADeps:
client: HaikuRAG
console: Console | None = None
agui_emitter: "AGUIEmitter[DeepQAState, DeepQAAnswer] | None" = None
semaphore: asyncio.Semaphore | None = None
def emit_log(self, message: str, state: "DeepQAState | None" = None) -> None:
if self.console:
self.console.print(message)
class DeepQAState(BaseModel):
"""Deep QA state for multi-agent question answering."""
@dataclass
class DeepQAState:
context: DeepQAContext
max_sub_questions: int = 3
max_iterations: int = 2
max_concurrency: int = 1
iterations: int = 0
model_config = {"arbitrary_types_allowed": True}
context: DeepQAContext = Field(description="Shared QA context")
max_sub_questions: int = Field(
default=3, description="Maximum number of sub-questions"
)
max_iterations: int = Field(
default=2, description="Maximum number of QA iterations"
)
max_concurrency: int = Field(
default=1, description="Maximum parallel sub-question searches"
)
iterations: int = Field(default=0, description="Current iteration number")
@classmethod
def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState":

0
haiku_rag_slim/haiku/rag/graph/nodes/__init__.py
View file

3
haiku_rag_slim/haiku/rag/graph/research/__init__.py Normal file
View file

@ -0,0 +1,3 @@
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import EvaluationResult, ResearchReport

4
haiku_rag_slim/haiku/rag/research/common.py → haiku_rag_slim/haiku/rag/graph/research/common.py
View file

@ -1,7 +1,7 @@
from pydantic_ai import format_as_xml
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.models import InsightAnalysis
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.models import InsightAnalysis
def format_context_for_prompt(context: ResearchContext) -> str:

128
haiku_rag_slim/haiku/rag/research/dependencies.py → haiku_rag_slim/haiku/rag/graph/research/dependencies.py
View file

@ -1,16 +1,14 @@
from collections.abc import Iterable
from pydantic import BaseModel, Field
from rich.console import Console
from pydantic import BaseModel, Field, PrivateAttr
from haiku.rag.client import HaikuRAG
from haiku.rag.graph_common.models import SearchAnswer
from haiku.rag.research.models import (
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.research.models import (
GapRecord,
InsightAnalysis,
InsightRecord,
)
from haiku.rag.research.stream import ResearchStream
class ResearchContext(BaseModel):
@ -30,26 +28,29 @@ class ResearchContext(BaseModel):
default_factory=list, description="Identified information gaps"
)
# Private dict indexes for O(1) lookups
_insights_by_id: dict[str, InsightRecord] = PrivateAttr(default_factory=dict)
_gaps_by_id: dict[str, GapRecord] = PrivateAttr(default_factory=dict)
def model_post_init(self, __context: object) -> None:
"""Build indexes after initialization."""
self._insights_by_id = {ins.id: ins for ins in self.insights}
self._gaps_by_id = {gap.id: gap for gap in self.gaps}
def add_qa_response(self, qa: SearchAnswer) -> None:
"""Add a structured QA response (minimal context already included)."""
self.qa_responses.append(qa)
def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]:
"""Merge one or more insights into the shared context with deduplication."""
merged: list[InsightRecord] = []
for record in records:
candidate = InsightRecord.model_validate(record)
existing = next(
(ins for ins in self.insights if ins.id == candidate.id), None
)
if not existing:
existing = next(
(ins for ins in self.insights if ins.summary == candidate.summary),
None,
)
existing = self._insights_by_id.get(candidate.id)
if existing:
# Update existing insight
existing.summary = candidate.summary
existing.status = candidate.status
if candidate.notes:
@ -62,36 +63,24 @@ class ResearchContext(BaseModel):
)
merged.append(existing)
else:
candidate = candidate.model_copy(deep=True)
if candidate.id is None: # pragma: no cover - defensive
raise ValueError(
"InsightRecord.id must be populated after validation"
)
candidate_id: str = candidate.id
candidate.id = self._allocate_insight_id(candidate_id)
self.insights.append(candidate)
merged.append(candidate)
# Add new insight
new_insight = candidate.model_copy(deep=True)
self.insights.append(new_insight)
self._insights_by_id[new_insight.id] = new_insight
merged.append(new_insight)
return merged
def upsert_gaps(self, records: Iterable[GapRecord]) -> list[GapRecord]:
"""Merge one or more gap records into the shared context with deduplication."""
merged: list[GapRecord] = []
for record in records:
candidate = GapRecord.model_validate(record)
existing = next((gap for gap in self.gaps if gap.id == candidate.id), None)
if not existing:
existing = next(
(
gap
for gap in self.gaps
if gap.description == candidate.description
),
None,
)
existing = self._gaps_by_id.get(candidate.id)
if existing:
# Update existing gap
existing.description = candidate.description
existing.severity = candidate.severity
existing.blocking = candidate.blocking
@ -106,22 +95,19 @@ class ResearchContext(BaseModel):
)
merged.append(existing)
else:
candidate = candidate.model_copy(deep=True)
if candidate.id is None: # pragma: no cover - defensive
raise ValueError("GapRecord.id must be populated after validation")
candidate_id: str = candidate.id
candidate.id = self._allocate_gap_id(candidate_id)
self.gaps.append(candidate)
merged.append(candidate)
# Add new gap
new_gap = candidate.model_copy(deep=True)
self.gaps.append(new_gap)
self._gaps_by_id[new_gap.id] = new_gap
merged.append(new_gap)
return merged
def mark_gap_resolved(
self, identifier: str, resolved_by: Iterable[str] | None = None
) -> GapRecord | None:
"""Mark a gap as resolved by identifier (id or description)."""
gap = self._find_gap(identifier)
"""Mark a gap as resolved by identifier."""
gap = self._gaps_by_id.get(identifier)
if gap is None:
return None
@ -133,7 +119,6 @@ class ResearchContext(BaseModel):
def integrate_analysis(self, analysis: InsightAnalysis) -> None:
"""Apply an analysis result to the shared context."""
merged_insights: list[InsightRecord] = []
if analysis.highlights:
merged_insights = self.upsert_insights(analysis.highlights)
@ -143,9 +128,7 @@ class ResearchContext(BaseModel):
analysis.gap_assessments = merged_gaps
if analysis.resolved_gaps:
resolved_by_list = (
[ins.id for ins in merged_insights if ins.id is not None]
if merged_insights
else None
[ins.id for ins in merged_insights] if merged_insights else None
)
for resolved in analysis.resolved_gaps:
self.mark_gap_resolved(resolved, resolved_by=resolved_by_list)
@ -153,29 +136,6 @@ class ResearchContext(BaseModel):
if question not in self.sub_questions:
self.sub_questions.append(question)
def _allocate_insight_id(self, candidate_id: str) -> str:
taken: set[str] = set()
for ins in self.insights:
if ins.id is not None:
taken.add(ins.id)
return _allocate_sequential_id(candidate_id, taken)
def _allocate_gap_id(self, candidate_id: str) -> str:
taken: set[str] = set()
for gap in self.gaps:
if gap.id is not None:
taken.add(gap.id)
return _allocate_sequential_id(candidate_id, taken)
def _find_gap(self, identifier: str) -> GapRecord | None:
normalized = identifier.lower().strip()
for gap in self.gaps:
if gap.id is not None and gap.id == normalized:
return gap
if gap.description.lower().strip() == normalized:
return gap
return None
class ResearchDependencies(BaseModel):
"""Dependencies for research agents with multi-agent context."""
@ -184,32 +144,8 @@ class ResearchDependencies(BaseModel):
client: HaikuRAG = Field(description="RAG client for document operations")
context: ResearchContext = Field(description="Shared research context")
console: Console | None = None
stream: ResearchStream | None = Field(
default=None, description="Optional research event stream"
)
def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]:
"""Merge two iterables preserving order while removing duplicates."""
merged = list(existing)
seen = {item for item in existing if item}
for item in incoming:
if item and item not in seen:
merged.append(item)
seen.add(item)
return merged
def _allocate_sequential_id(candidate: str, taken: set[str]) -> str:
slug = candidate
if slug not in taken:
return slug
base = slug
counter = 2
while True:
slug = f"{base}-{counter}"
if slug not in taken:
return slug
counter += 1
return [k for k in dict.fromkeys([*existing, *incoming]) if k]

295
haiku_rag_slim/haiku/rag/graph/research/graph.py Normal file
View file

@ -0,0 +1,295 @@
from pydantic_ai import Agent
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
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.common import get_model
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
from haiku.rag.graph.research.common import (
format_analysis_for_prompt,
format_context_for_prompt,
)
from haiku.rag.graph.research.dependencies import ResearchDependencies
from haiku.rag.graph.research.models import (
EvaluationResult,
InsightAnalysis,
ResearchReport,
)
from haiku.rag.graph.research.prompts import (
DECISION_AGENT_PROMPT,
INSIGHT_AGENT_PROMPT,
SYNTHESIS_AGENT_PROMPT,
)
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
def build_research_graph(
config: AppConfig = Config,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the Research graph.
Args:
config: AppConfig object (uses config.research for provider, model, and graph parameters)
Returns:
Configured Research graph
"""
provider = config.research.provider
model = config.research.model
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ResearchReport,
)
# Create and register the plan node using the factory
plan = g.step(
create_plan_node(
provider=provider,
model=model,
deps_type=ResearchDependencies, # type: ignore[arg-type]
activity_message="Creating research plan",
output_retries=3,
)
) # type: ignore[arg-type]
# Create and register the search_one node using the factory
search_one = g.step(
create_search_node(
provider=provider,
model=model,
deps_type=ResearchDependencies, # type: ignore[arg-type]
with_step_wrapper=True,
success_message_format="Found answer with {confidence:.0%} confidence",
handle_exceptions=True,
)
) # type: ignore[arg-type]
@g.step
async def get_batch(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
state = ctx.state
if not state.context.sub_questions:
return None
# Take ALL remaining questions and process them in parallel
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def analyze_insights(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
) -> None:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("analyze_insights")
deps.agui_emitter.update_activity(
"analyzing", "Synthesizing insights and gaps"
)
try:
agent = Agent(
model=get_model(provider, model),
output_type=InsightAnalysis,
instructions=INSIGHT_AGENT_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Review the latest research context and update the shared ledger of insights, gaps,"
" and follow-up questions.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
result = await agent.run(prompt, deps=agent_deps)
analysis: InsightAnalysis = result.output
state.context.integrate_analysis(analysis)
state.last_analysis = analysis
# State updated with insights/gaps - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
highlights = len(analysis.highlights) if analysis.highlights else 0
gaps = len(analysis.gap_assessments) if analysis.gap_assessments else 0
resolved = len(analysis.resolved_gaps) if analysis.resolved_gaps else 0
parts = []
if highlights:
parts.append(f"{highlights} insights")
if gaps:
parts.append(f"{gaps} gaps")
if resolved:
parts.append(f"{resolved} resolved")
summary = ", ".join(parts) if parts else "No updates"
deps.agui_emitter.update_activity("analyzing", f"Analysis: {summary}")
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step()
@g.step
async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("decide")
deps.agui_emitter.update_activity(
"evaluating", "Evaluating research sufficiency"
)
try:
agent = Agent(
model=get_model(provider, model),
output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
analysis_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)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
)
decision_result = await agent.run(prompt, deps=agent_deps)
output = decision_result.output
state.last_eval = output
state.iterations += 1
for new_q in output.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
# State updated with evaluation - emit state update and narrate
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
sufficient = "Yes" if output.is_sufficient else "No"
deps.agui_emitter.update_activity(
"evaluating",
f"Confidence: {output.confidence_score:.0%}, Sufficient: {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()
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> ResearchReport:
state = ctx.state
deps = ctx.deps
if deps.agui_emitter:
deps.agui_emitter.start_step("synthesize")
deps.agui_emitter.update_activity(
"synthesizing", "Generating final research report"
)
try:
agent = Agent(
model=get_model(provider, model),
output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_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()
# Build the graph structure
collect_answers = g.join(
reduce_list_append,
initial_factory=list[SearchAnswer],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
# Branch based on whether we have questions
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(analyze_insights),
g.edge_from(analyze_insights).to(decide),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x)
.label("Continue research")
.to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done researching")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

96
haiku_rag_slim/haiku/rag/research/models.py → haiku_rag_slim/haiku/rag/graph/research/models.py
View file

@ -1,19 +1,12 @@
import re
import uuid
from enum import Enum
from pydantic import BaseModel, Field, model_validator
_SLUG_RE = re.compile(r"[^a-z0-9]+")
from pydantic import BaseModel, Field, field_validator
def _make_slug(text: str, prefix: str) -> str:
"""Generate a lowercase slug with the given prefix as fallback."""
base = _SLUG_RE.sub("-", text.lower()).strip("-")
if not base:
base = prefix
# Trim overly long slugs but keep enough entropy for readability
return base[:48]
def _deduplicate_list(items: list[str]) -> list[str]:
"""Remove duplicates while preserving order."""
return list(dict.fromkeys(items))
class InsightStatus(str, Enum):
@ -28,48 +21,54 @@ class GapSeverity(str, Enum):
HIGH = "high"
class InsightRecord(BaseModel):
class TrackedRecord(BaseModel):
"""Base model for tracked entities with sources and metadata."""
model_config = {"validate_assignment": True}
id: str = Field(
default_factory=lambda: str(uuid.uuid4())[:8],
description="Unique identifier for the record",
)
supporting_sources: list[str] = Field(
default_factory=list,
description="Source identifiers backing this record",
)
notes: str | None = Field(
default=None,
description="Optional elaboration or caveats",
)
@field_validator("supporting_sources", mode="before")
@classmethod
def deduplicate_sources(cls, v: list[str]) -> list[str]:
"""Ensure supporting_sources has no duplicates."""
return _deduplicate_list(v) if v else []
class InsightRecord(TrackedRecord):
"""Structured insight with provenance and lifecycle metadata."""
id: str | None = Field(
default=None,
description="Stable slug identifier for the insight (auto-generated if omitted)",
)
summary: str = Field(description="Concise description of the insight")
status: InsightStatus = Field(
default=InsightStatus.OPEN,
description="Lifecycle status for the insight",
)
supporting_sources: list[str] = Field(
default_factory=list,
description="Source identifiers backing the insight",
)
originating_questions: list[str] = Field(
default_factory=list,
description="Research sub-questions that produced this insight",
)
notes: str | None = Field(
default=None,
description="Optional elaboration or caveats for the insight",
)
@model_validator(mode="after")
def _set_defaults(self) -> "InsightRecord":
if not self.id:
self.id = _make_slug(self.summary, "insight")
self.id = self.id.lower()
self.supporting_sources = list(dict.fromkeys(self.supporting_sources))
self.originating_questions = list(dict.fromkeys(self.originating_questions))
return self
@field_validator("originating_questions", mode="before")
@classmethod
def deduplicate_questions(cls, v: list[str]) -> list[str]:
"""Ensure originating_questions has no duplicates."""
return _deduplicate_list(v) if v else []
class GapRecord(BaseModel):
class GapRecord(TrackedRecord):
"""Structured representation of an identified research gap."""
id: str | None = Field(
default=None,
description="Stable slug identifier for the gap (auto-generated if omitted)",
)
description: str = Field(description="Concrete statement of what is missing")
severity: GapSeverity = Field(
default=GapSeverity.MEDIUM,
@ -87,23 +86,12 @@ class GapRecord(BaseModel):
default_factory=list,
description="Insight IDs or notes explaining how the gap was closed",
)
supporting_sources: list[str] = Field(
default_factory=list,
description="Sources confirming the gap status (e.g., evidence of absence)",
)
notes: str | None = Field(
default=None,
description="Optional clarification about the gap or follow-up actions",
)
@model_validator(mode="after")
def _set_defaults(self) -> "GapRecord":
if not self.id:
self.id = _make_slug(self.description, "gap")
self.id = self.id.lower()
self.resolved_by = list(dict.fromkeys(self.resolved_by))
self.supporting_sources = list(dict.fromkeys(self.supporting_sources))
return self
@field_validator("resolved_by", mode="before")
@classmethod
def deduplicate_resolved_by(cls, v: list[str]) -> list[str]:
"""Ensure resolved_by has no duplicates."""
return _deduplicate_list(v) if v else []
class InsightAnalysis(BaseModel):

0
haiku_rag_slim/haiku/rag/research/prompts.py → haiku_rag_slim/haiku/rag/graph/research/prompts.py
View file

85
haiku_rag_slim/haiku/rag/graph/research/state.py Normal file
View file

@ -0,0 +1,85 @@
import asyncio
from dataclasses import dataclass
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,
InsightAnalysis,
ResearchReport,
)
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
from haiku.rag.graph.agui.emitter import AGUIEmitter
@dataclass
class ResearchDeps:
"""Dependencies for research graph execution."""
client: HaikuRAG
agui_emitter: "AGUIEmitter[ResearchState, ResearchReport] | None" = None
semaphore: asyncio.Semaphore | None = None
def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
"""Emit a log message through AG-UI events.
Args:
message: The message to log
state: Optional state to include in state update
"""
if self.agui_emitter:
self.agui_emitter.log(message)
if state:
self.agui_emitter.update_state(state)
class ResearchState(BaseModel):
"""Research graph state model.
Fully JSON-serializable Pydantic model suitable for AG-UI state synchronization.
"""
model_config = {"arbitrary_types_allowed": True}
context: ResearchContext = Field(
description="Shared research context with questions, insights, and gaps"
)
iterations: int = Field(default=0, description="Current iteration number")
max_iterations: int = Field(default=3, description="Maximum allowed iterations")
confidence_threshold: float = Field(
default=0.8, description="Confidence threshold for completion", ge=0.0, le=1.0
)
max_concurrency: int = Field(
default=1, description="Maximum concurrent search operations", ge=1
)
last_eval: EvaluationResult | None = Field(
default=None, description="Last evaluation result"
)
last_analysis: InsightAnalysis | None = Field(
default=None, description="Last insight analysis"
)
@classmethod
def from_config(
cls, context: ResearchContext, config: "AppConfig"
) -> "ResearchState":
"""Create a ResearchState from an AppConfig.
Args:
context: The ResearchContext containing the question and settings
config: The AppConfig object (uses config.research for state parameters)
Returns:
A configured ResearchState instance
"""
return cls(
context=context,
max_iterations=config.research.max_iterations,
confidence_threshold=config.research.confidence_threshold,
max_concurrency=config.research.max_concurrency,
)

5
haiku_rag_slim/haiku/rag/graph_common/__init__.py
View file

@ -1,5 +0,0 @@
"""Common utilities for graph implementations."""
from haiku.rag.graph_common.utils import get_model, log
__all__ = ["get_model", "log"]

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

@ -6,7 +6,7 @@ from pydantic import BaseModel
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, Config
from haiku.rag.research.models import ResearchReport
from haiku.rag.graph.research.models import ResearchReport
class SearchResult(BaseModel):
@ -191,9 +191,9 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
try:
async with HaikuRAG(db_path, config=config) as rag:
if deep:
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
graph = build_deep_qa_graph(config=config)
context = DeepQAContext(
@ -226,9 +226,9 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP:
A research report with findings, or None if an error occurred.
"""
try:
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchDeps, ResearchState
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
async with HaikuRAG(db_path, config=config) as rag:
graph = build_research_graph(config=config)

1
haiku_rag_slim/haiku/rag/qa/deep/__init__.py
View file

@ -1 +0,0 @@
from haiku.rag.qa.deep.models import DeepQAAnswer

363
haiku_rag_slim/haiku/rag/qa/deep/graph.py
View file

@ -1,363 +0,0 @@
from typing import Any
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
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_common import get_model, log
from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer
from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.qa.deep.dependencies import DeepQADependencies
from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
from haiku.rag.qa.deep.prompts import (
DECISION_PROMPT,
SYNTHESIS_PROMPT,
SYNTHESIS_PROMPT_WITH_CITATIONS,
)
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
def build_deep_qa_graph(
config: AppConfig = Config,
) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]:
"""Build the Deep QA graph.
Args:
config: AppConfig object (uses config.qa for provider, model, and graph parameters)
Returns:
Configured Deep QA graph
"""
provider = config.qa.provider
model = config.qa.model
g = GraphBuilder(
state_type=DeepQAState,
deps_type=DeepQADeps,
output_type=DeepQAAnswer,
)
@g.step
async def plan(ctx: StepContext[DeepQAState, DeepQADeps, None]) -> None:
state = ctx.state
deps = ctx.deps
log(deps, state, "\n[bold cyan]📋 Planning approach...[/bold cyan]")
plan_agent = Agent(
model=get_model(provider, model),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
deps_type=DeepQADependencies,
)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
log(
deps,
state,
f" [bold]Main Question:[/bold] {state.context.original_question}",
)
log(deps, state, " [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
@g.step
async def search_one(
ctx: StepContext[DeepQAState, DeepQADeps, str],
) -> SearchAnswer:
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
# Create semaphore if not already provided
if deps.semaphore is None:
import asyncio
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
# Use semaphore to control concurrency
async with deps.semaphore:
return await _do_search(state, deps, sub_q)
async def _do_search(
state: DeepQAState,
deps: DeepQADeps,
sub_q: str,
) -> SearchAnswer:
log(
deps,
state,
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
)
agent = Agent(
model=get_model(provider, model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[DeepQADependencies], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_title or chunk.document_uri or ""),
}
for chunk, score in expanded
]
if not entries:
return (
f"No relevant information found in the knowledge base for: {query}"
)
return format_as_xml(entries, root_tag="snippets")
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
answer = result.output
if answer:
state.context.add_qa_response(answer)
preview = answer.answer[:150] + (
"…" if len(answer.answer) > 150 else ""
)
log(deps, state, f" [green]✓[/green] {preview}")
return answer
except Exception as e:
log(deps, state, f"[red]Search failed:[/red] {e}")
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",
confidence=0.0,
)
return failure_answer
@g.step
async def get_batch(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
state = ctx.state
if not state.context.sub_questions:
return None
# Take ALL remaining questions - max_concurrency controls parallel execution within .map()
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def decide(
ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]],
) -> bool:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=DeepQAEvaluation,
instructions=DECISION_PROMPT,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"gathered_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = (
"Evaluate whether we have sufficient information to answer the question.\n\n"
f"{context_xml}"
)
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
result = await agent.run(prompt, deps=agent_deps)
evaluation = result.output
state.iterations += 1
log(deps, state, f" [bold]Assessment:[/bold] {evaluation.reasoning}")
status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]"
log(deps, state, f" Sufficient: {status}")
for new_q in evaluation.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if evaluation.new_questions:
log(deps, state, " [cyan]New questions:[/cyan]")
for question in evaluation.new_questions:
log(deps, state, f" • {question}")
should_continue = (
not evaluation.is_sufficient and state.iterations < state.max_iterations
)
if not should_continue:
if state.iterations >= state.max_iterations:
log(
deps,
state,
f"\n[bold yellow]⚠️ Reached max iterations ({state.max_iterations})[/bold yellow]",
)
log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]")
else:
log(
deps,
state,
f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]",
)
return should_continue
@g.step
async def synthesize(
ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
) -> DeepQAAnswer:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]",
)
prompt_template = (
SYNTHESIS_PROMPT_WITH_CITATIONS
if state.context.use_citations
else SYNTHESIS_PROMPT
)
agent = Agent(
model=get_model(provider, model),
output_type=DeepQAAnswer,
instructions=prompt_template,
retries=3,
deps_type=DeepQADependencies,
)
context_data = {
"original_question": state.context.original_question,
"sub_answers": [
{
"question": qa.query,
"answer": qa.answer,
"sources": qa.sources,
}
for qa in state.context.qa_responses
],
}
context_xml = format_as_xml(context_data, root_tag="gathered_information")
prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"
agent_deps = DeepQADependencies(
client=deps.client,
context=state.context,
console=deps.console,
)
result = await agent.run(prompt, deps=agent_deps)
log(deps, state, "[bold green]✅ Answer complete![/bold green]")
return result.output
# Build the graph structure
collect_answers = g.join(
reduce_list_append,
initial_factory=list[SearchAnswer],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
# Branch based on whether we have questions
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),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done with QA")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

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

@ -1,3 +0,0 @@
from haiku.rag.graph_common.models import SearchAnswer
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.models import EvaluationResult, ResearchReport

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

@ -1,429 +0,0 @@
from typing import Any
from pydantic_ai import Agent, RunContext
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.output import ToolOutput
from pydantic_graph.beta import Graph, GraphBuilder, StepContext
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_common import get_model, log
from haiku.rag.graph_common.models import ResearchPlan, SearchAnswer
from haiku.rag.graph_common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
from haiku.rag.research.common import (
format_analysis_for_prompt,
format_context_for_prompt,
)
from haiku.rag.research.dependencies import ResearchDependencies
from haiku.rag.research.models import (
EvaluationResult,
InsightAnalysis,
ResearchReport,
)
from haiku.rag.research.prompts import (
DECISION_AGENT_PROMPT,
INSIGHT_AGENT_PROMPT,
SYNTHESIS_AGENT_PROMPT,
)
from haiku.rag.research.state import ResearchDeps, ResearchState
def build_research_graph(
config: AppConfig = Config,
) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]:
"""Build the Research graph.
Args:
config: AppConfig object (uses config.research for provider, model, and graph parameters)
Returns:
Configured Research graph
"""
provider = config.research.provider
model = config.research.model
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ResearchReport,
)
@g.step
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
state = ctx.state
deps = ctx.deps
log(deps, state, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
plan_agent = Agent(
model=get_model(provider, model),
output_type=ResearchPlan,
instructions=(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning."
),
retries=3,
deps_type=ResearchDependencies,
)
@plan_agent.tool
async def gather_context(
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
) -> str:
results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded)
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
log(
deps,
state,
f" [bold]Main Question:[/bold] {state.context.original_question}",
)
log(deps, state, " [bold]Sub-questions:[/bold]")
for i, sq in enumerate(state.context.sub_questions, 1):
log(deps, state, f" {i}. {sq}")
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
state = ctx.state
deps = ctx.deps
sub_q = ctx.inputs
# Create semaphore if not already provided
if deps.semaphore is None:
import asyncio
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
# Use semaphore to control concurrency
async with deps.semaphore:
return await _do_search(state, deps, sub_q)
async def _do_search(
state: ResearchState,
deps: ResearchDeps,
sub_q: str,
) -> SearchAnswer:
log(
deps,
state,
f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
)
agent = Agent(
model=get_model(provider, model),
output_type=ToolOutput(SearchAnswer, max_retries=3),
instructions=SEARCH_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies], query: str, limit: int = 5
) -> str:
search_results = await ctx2.deps.client.search(query, limit=limit)
expanded = await ctx2.deps.client.expand_context(search_results)
entries: list[dict[str, Any]] = [
{
"text": chunk.content,
"score": score,
"document_uri": (chunk.document_title or chunk.document_uri or ""),
}
for chunk, score in expanded
]
if not entries:
return (
f"No relevant information found in the knowledge base for: {query}"
)
return format_as_xml(entries, root_tag="snippets")
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
try:
result = await agent.run(sub_q, deps=agent_deps)
answer = result.output
if answer:
state.context.add_qa_response(answer)
preview = answer.answer[:150] + (
"…" if len(answer.answer) > 150 else ""
)
log(deps, state, f" [green]✓[/green] {preview}")
return answer
except Exception as e:
log(deps, state, f"[red]Search failed:[/red] {e}")
failure_answer = SearchAnswer(
query=sub_q,
answer=f"Search failed after retries: {str(e)}",
confidence=0.0,
)
return failure_answer
@g.step
async def get_batch(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
state = ctx.state
if not state.context.sub_questions:
return None
# Take ALL remaining questions and process them in parallel
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
@g.step
async def analyze_insights(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]],
) -> None:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]🧭 Synthesizing new insights and gap status...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=InsightAnalysis,
instructions=INSIGHT_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
"Review the latest research context and update the shared ledger of insights, gaps,"
" and follow-up questions.\n\n"
f"{context_xml}"
)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
result = await agent.run(prompt, deps=agent_deps)
analysis: InsightAnalysis = result.output
state.context.integrate_analysis(analysis)
state.last_analysis = analysis
if analysis.commentary:
log(deps, state, f" Summary: {analysis.commentary}")
if analysis.highlights:
log(deps, state, " [bold]Updated insights:[/bold]")
for insight in analysis.highlights:
label = insight.status.value
log(
deps,
state,
f" • ({label}) {insight.summary}",
)
if analysis.gap_assessments:
log(deps, state, " [bold yellow]Gap updates:[/bold yellow]")
for gap in analysis.gap_assessments:
status = "resolved" if gap.resolved else "open"
severity = gap.severity.value
log(
deps,
state,
f" • ({severity}/{status}) {gap.description}",
)
if analysis.resolved_gaps:
log(deps, state, " [green]Resolved gaps:[/green]")
for resolved in analysis.resolved_gaps:
log(deps, state, f" • {resolved}")
if analysis.new_questions:
log(deps, state, " [cyan]Proposed follow-ups:[/cyan]")
for question in analysis.new_questions:
log(deps, state, f" • {question}")
@g.step
async def decide(ctx: StepContext[ResearchState, ResearchDeps, None]) -> bool:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📊 Evaluating research sufficiency...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=EvaluationResult,
instructions=DECISION_AGENT_PROMPT,
retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
analysis_xml = format_analysis_for_prompt(state.last_analysis)
prompt_parts = [
"Assess whether the research now answers the original question with adequate confidence.",
context_xml,
analysis_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)
agent_deps = ResearchDependencies(
client=deps.client,
context=state.context,
console=deps.console,
stream=deps.stream,
)
decision_result = await agent.run(prompt, deps=agent_deps)
output = decision_result.output
state.last_eval = output
state.iterations += 1
for new_q in output.new_questions:
if new_q not in state.context.sub_questions:
state.context.sub_questions.append(new_q)
if output.key_insights:
log(deps, state, " [bold]Key insights:[/bold]")
for insight in output.key_insights:
log(deps, state, f" • {insight}")
if output.gaps:
log(deps, state, " [bold yellow]Remaining gaps:[/bold yellow]")
for gap in output.gaps:
log(deps, state, f" • {gap}")
log(
deps,
state,
f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
)
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
log(deps, state, f" Sufficient: {status}")
should_continue = (
not output.is_sufficient
or output.confidence_score < state.confidence_threshold
) and state.iterations < state.max_iterations
if not should_continue:
log(deps, state, "\n[bold green]✅ Stopping research.[/bold green]")
return should_continue
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, None | bool],
) -> ResearchReport:
state = ctx.state
deps = ctx.deps
log(
deps,
state,
"\n[bold cyan]📝 Generating final research report...[/bold cyan]",
)
agent = Agent(
model=get_model(provider, model),
output_type=ResearchReport,
instructions=SYNTHESIS_AGENT_PROMPT,
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,
console=deps.console,
stream=deps.stream,
)
result = await agent.run(prompt, deps=agent_deps)
log(deps, state, "[bold green]✅ Research complete![/bold green]")
return result.output
# Build the graph structure
collect_answers = g.join(
reduce_list_append,
initial_factory=list[SearchAnswer],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
)
# Branch based on whether we have questions
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(analyze_insights),
g.edge_from(analyze_insights).to(decide),
)
# Branch based on decision
g.add(
g.edge_from(decide).to(
g.decision()
.branch(
g.match(bool, matches=lambda x: x)
.label("Continue research")
.to(get_batch)
)
.branch(
g.match(bool, matches=lambda x: not x)
.label("Done researching")
.to(synthesize)
)
),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

58
haiku_rag_slim/haiku/rag/research/state.py
View file

@ -1,58 +0,0 @@
import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.models import EvaluationResult, InsightAnalysis
from haiku.rag.research.stream import ResearchStream
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
@dataclass
class ResearchDeps:
client: HaikuRAG
console: Console | None = None
stream: ResearchStream | None = None
semaphore: asyncio.Semaphore | None = None
def emit_log(self, message: str, state: "ResearchState | None" = None) -> None:
if self.console:
self.console.print(message)
if self.stream:
self.stream.log(message, state)
@dataclass
class ResearchState:
context: ResearchContext
iterations: int = 0
max_iterations: int = 3
confidence_threshold: float = 0.8
max_concurrency: int = 1
last_eval: EvaluationResult | None = None
last_analysis: InsightAnalysis | None = None
@classmethod
def from_config(
cls, context: ResearchContext, config: "AppConfig"
) -> "ResearchState":
"""Create a ResearchState from an AppConfig.
Args:
context: The ResearchContext containing the question and settings
config: The AppConfig object (uses config.research for state parameters)
Returns:
A configured ResearchState instance
"""
return cls(
context=context,
max_iterations=config.research.max_iterations,
confidence_threshold=config.research.confidence_threshold,
max_concurrency=config.research.max_concurrency,
)

162
haiku_rag_slim/haiku/rag/research/stream.py
View file

@ -1,162 +0,0 @@
import asyncio
from collections.abc import AsyncIterator
from dataclasses import dataclass
from typing import TYPE_CHECKING, Literal
from haiku.rag.research.models import ResearchReport
if TYPE_CHECKING: # pragma: no cover
from haiku.rag.research.state import ResearchState
@dataclass(slots=True)
class ResearchStateSnapshot:
question: str
sub_questions: list[str]
iterations: int
max_iterations: int
confidence_threshold: float
pending_sub_questions: int
answered_questions: int
insights: list[str]
gaps: list[str]
last_confidence: float | None
last_sufficient: bool | None
@classmethod
def from_state(cls, state: "ResearchState") -> "ResearchStateSnapshot":
context = state.context
last_confidence: float | None = None
last_sufficient: bool | None = None
if state.last_eval:
last_confidence = state.last_eval.confidence_score
last_sufficient = state.last_eval.is_sufficient
return cls(
question=context.original_question,
sub_questions=list(context.sub_questions),
iterations=state.iterations,
max_iterations=state.max_iterations,
confidence_threshold=state.confidence_threshold,
pending_sub_questions=len(context.sub_questions),
answered_questions=len(context.qa_responses),
insights=[
f"{insight.status.value}:{insight.summary}"
for insight in context.insights
],
gaps=[
f"{gap.severity.value}/{'resolved' if gap.resolved else 'open'}:{gap.description}"
for gap in context.gaps
],
last_confidence=last_confidence,
last_sufficient=last_sufficient,
)
@dataclass(slots=True)
class ResearchStreamEvent:
type: Literal["log", "report", "error"]
message: str | None = None
state: ResearchStateSnapshot | None = None
report: ResearchReport | None = None
error: str | None = None
class ResearchStream:
"""Queue-backed stream for research graph events."""
def __init__(self) -> None:
self._queue: asyncio.Queue[ResearchStreamEvent | None] = asyncio.Queue()
self._closed = False
def _snapshot(self, state: "ResearchState | None") -> ResearchStateSnapshot | None:
if state is None:
return None
return ResearchStateSnapshot.from_state(state)
def log(self, message: str, state: "ResearchState | None" = None) -> None:
if self._closed:
return
event = ResearchStreamEvent(
type="log", message=message, state=self._snapshot(state)
)
self._queue.put_nowait(event)
def report(self, report: ResearchReport, state: "ResearchState") -> None:
if self._closed:
return
event = ResearchStreamEvent(
type="report",
report=report,
state=self._snapshot(state),
)
self._queue.put_nowait(event)
def error(self, error: Exception, state: "ResearchState | None" = None) -> None:
if self._closed:
return
event = ResearchStreamEvent(
type="error",
message=str(error),
error=str(error),
state=self._snapshot(state),
)
self._queue.put_nowait(event)
async def close(self) -> None:
if self._closed:
return
self._closed = True
await self._queue.put(None)
def __aiter__(self) -> AsyncIterator[ResearchStreamEvent]:
return self._iter_events()
async def _iter_events(self) -> AsyncIterator[ResearchStreamEvent]:
while True:
event = await self._queue.get()
if event is None:
break
yield event
async def stream_research_graph(
graph,
state: "ResearchState",
deps,
) -> AsyncIterator[ResearchStreamEvent]:
"""Run the research graph and yield streaming events as they occur."""
from contextlib import suppress
from haiku.rag.research.state import ResearchDeps
if not isinstance(deps, ResearchDeps):
raise TypeError("deps must be an instance of ResearchDeps")
stream = ResearchStream()
deps.stream = stream
async def _execute() -> None:
try:
report = await graph.run(state=state, deps=deps)
if report is None:
raise RuntimeError("Graph did not produce a report")
stream.report(report, state)
except Exception as exc:
stream.error(exc, state)
finally:
await stream.close()
runner = asyncio.create_task(_execute())
try:
async for event in stream:
yield event
finally:
if not runner.done():
runner.cancel()
with suppress(asyncio.CancelledError):
await runner

2
haiku_rag_slim/pyproject.toml
View file

@ -27,7 +27,7 @@ dependencies = [
"lancedb==0.25.2",
"pathspec>=0.12.1",
"pydantic>=2.12.3",
"pydantic-ai-slim[openai,fastmcp,logfire]>=1.11.1",
"pydantic-ai-slim[openai,fastmcp,logfire,ag-ui]>=1.11.1",
"python-dotenv>=1.2.1",
"pyyaml>=6.0.3",
"rich>=14.2.0",

101
haiku_rag_slim/test_agui_server.py Executable file
View file

@ -0,0 +1,101 @@
#!/usr/bin/env python3
"""Test script for AG-UI server functionality."""
import asyncio
import json
import httpx
async def test_agui_server():
"""Test the AG-UI server endpoint."""
base_url = "http://localhost:8000"
# Test health check
async with httpx.AsyncClient() as client:
print("Testing health check...")
response = await client.get(f"{base_url}/health")
print(f"Health check response: {response.status_code}")
print(f"Health check data: {response.json()}")
# Test AG-UI streaming endpoint
print("\nTesting AG-UI stream endpoint...")
request_data = {
"threadId": "test-thread-1",
"runId": "test-run-1",
"state": {"question": "What is pydantic-graph?"},
"messages": [],
"config": {},
}
print(f"Request data: {json.dumps(request_data, indent=2)}")
# Send request and stream response
async with client.stream(
"POST",
f"{base_url}/v1/agent/stream",
json=request_data,
timeout=120.0,
) as response:
print(f"Response status: {response.status_code}")
print("Streaming events...\n")
event_count = 0
async for line in response.aiter_lines():
if line.startswith("data: "):
event_data = line[6:] # Remove "data: " prefix
try:
event = json.loads(event_data)
event_type = event.get("type", "UNKNOWN")
print(f"Event {event_count}: {event_type}")
# Show specific event details
if event_type == "RUN_STARTED":
print(f" Thread ID: {event.get('threadId')}")
print(f" Run ID: {event.get('runId')}")
elif event_type == "STEP_STARTED":
print(f" Step: {event.get('stepName')}")
elif event_type == "ACTIVITY_SNAPSHOT":
print(f" Activity: {event.get('content')}")
elif event_type == "STATE_SNAPSHOT":
state = event.get("snapshot", {})
if "context" in state:
context = state["context"]
if "sub_questions" in context:
num_questions = len(context["sub_questions"])
print(f" Sub-questions: {num_questions}")
elif event_type == "RUN_FINISHED":
result = event.get("result", {})
if "title" in result:
print(f" Report Title: {result['title']}")
elif event_type == "RUN_ERROR":
print(f" Error: {event.get('message')}")
event_count += 1
except json.JSONDecodeError as e:
print(f"Failed to parse event: {e}")
print(f"Raw line: {line}")
print(f"\nTotal events received: {event_count}")
if __name__ == "__main__":
print("AG-UI Server Test")
print("=" * 50)
print("Make sure to start the server first with:")
print(" haiku-rag serve --agui --agui-port 8000")
print("=" * 50)
print()
try:
asyncio.run(test_agui_server())
except httpx.ConnectError:
print("ERROR: Could not connect to server at http://localhost:8000")
print("Make sure the AG-UI server is running.")
except KeyboardInterrupt:
print("\nTest interrupted by user")
except Exception as e:
print(f"Test failed with error: {e}")
import traceback
traceback.print_exc()

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

@ -0,0 +1 @@
"""Tests for haiku.rag.graph module."""

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

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

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

@ -0,0 +1,204 @@
"""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.events 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", "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", "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}

263
tests/graph/agui/test_emitter.py Normal file
View file

@ -0,0 +1,263 @@
"""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()
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()
emitter.update_activity("processing", "Processing data")
emitter.update_activity("done", "Completed", message_id="msg-1")
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"] == "Processing data"
assert activity_events[1]["messageId"] == "msg-1"
assert activity_events[1]["activityType"] == "done"
@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

155
tests/graph/agui/test_events.py Normal file
View file

@ -0,0 +1,155 @@
"""Tests for AG-UI event creation utilities."""
from pydantic import BaseModel
from haiku.rag.graph.agui.events import (
emit_activity,
emit_run_error,
emit_run_finished,
emit_run_started,
emit_state_snapshot,
emit_step_finished,
emit_step_started,
emit_text_message,
)
class TestState(BaseModel):
"""Test state model."""
value: int
class TestResult(BaseModel):
"""Test result model."""
status: str
def test_emit_run_started():
"""Test RUN_STARTED event creation."""
event = emit_run_started("thread-1", "run-1")
assert event["type"] == "RUN_STARTED"
assert event["threadId"] == "thread-1"
assert event["runId"] == "run-1"
assert "input" not in event
def test_emit_run_started_with_input():
"""Test RUN_STARTED event with input data."""
event = emit_run_started("thread-1", "run-1", input_data="test input")
assert event["type"] == "RUN_STARTED"
assert event["input"] == "test input"
def test_emit_run_finished():
"""Test RUN_FINISHED event creation."""
result = TestResult(status="complete")
event = emit_run_finished("thread-1", "run-1", result)
assert event["type"] == "RUN_FINISHED"
assert event["threadId"] == "thread-1"
assert event["runId"] == "run-1"
assert event["result"] == {"status": "complete"}
def test_emit_run_finished_with_dict():
"""Test RUN_FINISHED event with dict result."""
result = {"status": "complete", "count": 42}
event = emit_run_finished("thread-1", "run-1", result)
assert event["type"] == "RUN_FINISHED"
assert event["result"] == result
def test_emit_run_error():
"""Test RUN_ERROR event creation."""
event = emit_run_error("Something went wrong")
assert event["type"] == "RUN_ERROR"
assert event["message"] == "Something went wrong"
assert "code" not in event
def test_emit_run_error_with_code():
"""Test RUN_ERROR event with error code."""
event = emit_run_error("Something went wrong", code="ERR_001")
assert event["type"] == "RUN_ERROR"
assert event["message"] == "Something went wrong"
assert event["code"] == "ERR_001"
def test_emit_step_started():
"""Test STEP_STARTED event creation."""
event = emit_step_started("plan")
assert event["type"] == "STEP_STARTED"
assert event["stepName"] == "plan"
def test_emit_step_finished():
"""Test STEP_FINISHED event creation."""
event = emit_step_finished("plan")
assert event["type"] == "STEP_FINISHED"
assert event["stepName"] == "plan"
def test_emit_text_message():
"""Test TEXT_MESSAGE_CHUNK event creation."""
event = emit_text_message("Hello world")
assert event["type"] == "TEXT_MESSAGE_CHUNK"
assert event["delta"] == "Hello world"
assert event["role"] == "assistant"
assert "messageId" in event
def test_emit_text_message_with_role():
"""Test TEXT_MESSAGE_CHUNK event with custom role."""
event = emit_text_message("Hello", role="user")
assert event["type"] == "TEXT_MESSAGE_CHUNK"
assert event["role"] == "user"
def test_emit_state_snapshot():
"""Test STATE_SNAPSHOT event creation."""
state = TestState(value=42)
event = emit_state_snapshot(state)
assert event["type"] == "STATE_SNAPSHOT"
assert event["snapshot"] == {"value": 42}
def test_emit_activity():
"""Test ACTIVITY_SNAPSHOT event creation."""
event = emit_activity("msg-1", "processing", "Working on task")
assert event["type"] == "ACTIVITY_SNAPSHOT"
assert event["messageId"] == "msg-1"
assert event["activityType"] == "processing"
assert event["content"] == "Working on task"
def test_event_structure_consistency():
"""Test that all events have consistent structure."""
events = [
emit_run_started("t1", "r1"),
emit_run_finished("t1", "r1", {"result": "done"}),
emit_run_error("error"),
emit_step_started("step1"),
emit_step_finished("step1"),
emit_text_message("text"),
emit_state_snapshot(TestState(value=1)),
emit_activity("m1", "type", "content"),
]
for event in events:
assert isinstance(event, dict)
assert "type" in event
assert isinstance(event["type"], str)
assert event["type"].isupper() # Event types are uppercase

250
tests/graph/agui/test_server.py Normal file
View file

@ -0,0 +1,250 @@
"""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_run_agent_input_defaults():
"""Test RunAgentInput with defaults."""
input_data = RunAgentInput() # type: ignore[call-arg]
assert input_data.thread_id is None
assert input_data.run_id is None
assert input_data.state == {}
assert input_data.messages == []
assert input_data.config == {}
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,
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,
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,
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,
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
)
def test_agui_config_defaults():
"""Test AGUIConfig default values."""
config = AGUIConfig()
assert config.host == "0.0.0.0"
assert config.port == 8000
assert config.cors_origins == ["*"]
assert config.cors_credentials is True
assert "GET" in config.cors_methods
assert "POST" in config.cors_methods

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

@ -0,0 +1,199 @@
"""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", "Doing work")
deps.agui_emitter.finish_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):
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):
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):
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):
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"] == "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):
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):
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):
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}

20
tests/test_deep_qa.py → tests/graph/test_deep_qa.py
View file

@ -2,10 +2,10 @@ import pytest
from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG
from haiku.rag.graph_common.models import SearchAnswer
from haiku.rag.qa.deep.dependencies import DeepQAContext
from haiku.rag.qa.deep.graph import build_deep_qa_graph
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
from haiku.rag.graph.common.models import SearchAnswer
from haiku.rag.graph.deep_qa.dependencies import DeepQAContext
from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph
from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
@pytest.mark.asyncio
@ -16,8 +16,8 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
graph = build_deep_qa_graph()
@ -30,7 +30,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
# Use real client but with TestModel for LLM calls
client = HaikuRAG(temp_db_path)
deps = DeepQADeps(client=client, console=None)
deps = DeepQADeps(client=client)
result = await graph.run(state=state, deps=deps)
@ -50,8 +50,8 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory)
graph = build_deep_qa_graph()
@ -62,7 +62,7 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
# Use real client but with TestModel for LLM calls
client = HaikuRAG(temp_db_path)
deps = DeepQADeps(client=client, console=None)
deps = DeepQADeps(client=client)
result = await graph.run(state=state, deps=deps)

85
tests/graph/test_research_graph.py Normal file
View file

@ -0,0 +1,85 @@
import asyncio
import pytest
from pydantic_ai.models.test import TestModel
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.state import ResearchDeps, ResearchState
def test_build_graph_and_state():
graph = build_research_graph()
assert graph is not None
state = ResearchState(
context=ResearchContext(
original_question="What are the key features of haiku.rag?"
),
max_iterations=1,
confidence_threshold=0.8,
)
assert state.iterations == 0
assert state.context.sub_questions == []
def test_async_loop_available():
# Ensure an event loop can be created in test env
loop = asyncio.new_event_loop()
loop.close()
@pytest.mark.asyncio
async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
"""Test research graph with mocked LLM using AG-UI events."""
# Mock get_model to return TestModel which generates valid schema-compliant data
def test_model_factory(_provider, _model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph.common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.graph.research.graph.get_model", test_model_factory)
graph = build_research_graph()
state = ResearchState(
context=ResearchContext(original_question="What is haiku.rag?"),
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=2,
)
# Use real client but with TestModel for LLM calls
client = HaikuRAG(temp_db_path)
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']}")
# TestModel will generate valid structured output for each node
assert result is not None, (
f"No result. Events collected: {[e['type'] for e in events]}"
)
# Result is serialized as dict in AG-UI events
assert isinstance(result, dict)
assert "title" in result
assert isinstance(result["title"], str)
assert "executive_summary" in result
assert "main_findings" in result
# Verify AG-UI events were emitted
event_types = [e["type"] for e in events]
assert "RUN_STARTED" in event_types
assert "RUN_FINISHED" in event_types
assert "STATE_SNAPSHOT" in event_types
assert "STEP_STARTED" in event_types
client.close()

28
tests/test_app.py
View file

@ -343,7 +343,7 @@ async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA."""
from haiku.rag.qa.deep.models import DeepQAAnswer
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
@ -358,7 +358,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
with patch(
"haiku.rag.qa.deep.graph.build_deep_qa_graph", return_value=mock_graph
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
):
await app.ask("test question", deep=True)
@ -371,7 +371,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA and citations."""
from haiku.rag.qa.deep.models import DeepQAAnswer
from haiku.rag.graph.deep_qa.models import DeepQAAnswer
mock_output = DeepQAAnswer(
answer="Deep QA answer with citations [test.md]", sources=["test.md"]
@ -388,7 +388,7 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
with patch(
"haiku.rag.qa.deep.graph.build_deep_qa_graph", return_value=mock_graph
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
):
await app.ask("test question", deep=True, cite=True)
@ -401,12 +401,13 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch):
@pytest.mark.asyncio
async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
"""Test asking a question with deep QA and verbose output."""
from haiku.rag.qa.deep.models import DeepQAAnswer
mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"])
mock_output = {"answer": "Deep QA answer", "sources": ["test.md"]}
mock_renderer = AsyncMock()
mock_renderer.render.return_value = mock_output
mock_graph = AsyncMock()
mock_graph.run.return_value = mock_output
mock_client = AsyncMock()
mock_client.__aenter__.return_value = mock_client
@ -416,10 +417,13 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch):
with patch("haiku.rag.app.HaikuRAG", return_value=mock_client):
with patch(
"haiku.rag.qa.deep.graph.build_deep_qa_graph", return_value=mock_graph
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph
):
await app.ask("test question", deep=True, verbose=True)
with patch(
"haiku.rag.graph.agui.AGUIConsoleRenderer", return_value=mock_renderer
):
await app.ask("test question", deep=True, verbose=True)
mock_graph.run.assert_called_once()
call_kwargs = mock_graph.run.call_args[1]
assert call_kwargs["deps"].console is not None
# With verbose, it should use AGUIConsoleRenderer.render, not graph.run
mock_renderer.render.assert_called_once()
mock_graph.run.assert_not_called()

8
tests/test_mcp.py
View file

@ -4,8 +4,8 @@ from unittest.mock import AsyncMock, patch
import pytest
from haiku.rag.graph.research.models import ResearchReport
from haiku.rag.mcp import create_mcp_server
from haiku.rag.research.models import ResearchReport
from haiku.rag.store.models.document import Document
@ -249,7 +249,9 @@ async def test_mcp_ask_question_deep():
with (
patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
patch("haiku.rag.qa.deep.graph.build_deep_qa_graph") as mock_graph_builder,
patch(
"haiku.rag.graph.deep_qa.graph.build_deep_qa_graph"
) as mock_graph_builder,
):
mock_rag = AsyncMock()
mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag)
@ -291,7 +293,7 @@ async def test_mcp_research_question():
with (
patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class,
patch(
"haiku.rag.research.graph.build_research_graph"
"haiku.rag.graph.research.graph.build_research_graph"
) as mock_graph_builder,
):
mock_rag = AsyncMock()

26
tests/test_research_graph.py
View file

@ -1,26 +0,0 @@
import asyncio
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.state import ResearchState
def test_build_graph_and_state():
graph = build_research_graph()
assert graph is not None
state = ResearchState(
context=ResearchContext(
original_question="What are the key features of haiku.rag?"
),
max_iterations=1,
confidence_threshold=0.8,
)
assert state.iterations == 0
assert state.context.sub_questions == []
def test_async_loop_available():
# Ensure an event loop can be created in test env
loop = asyncio.new_event_loop()
loop.close()

55
tests/test_research_graph_integration.py
View file

@ -1,55 +0,0 @@
import pytest
from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG
from haiku.rag.research.dependencies import ResearchContext
from haiku.rag.research.graph import build_research_graph
from haiku.rag.research.models import ResearchReport
from haiku.rag.research.state import ResearchDeps, ResearchState
from haiku.rag.research.stream import stream_research_graph
@pytest.mark.asyncio
async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
"""Test research graph with mocked LLM using TestModel."""
# Mock get_model to return TestModel which generates valid schema-compliant data
def test_model_factory(provider, model):
return TestModel()
monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory)
graph = build_research_graph()
state = ResearchState(
context=ResearchContext(original_question="What is haiku.rag?"),
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=2,
)
# Use real client but with TestModel for LLM calls
client = HaikuRAG(temp_db_path)
deps = ResearchDeps(client=client, console=None)
collected = []
report = None
async for event in stream_research_graph(graph, state, deps):
collected.append(event)
if event.type == "report":
report = event.report
break
elif event.type == "error":
pytest.fail(f"Graph execution failed: {event.error}")
# TestModel will generate valid structured output for each node
assert report is not None, (
f"No report generated. Events collected: {[e.type for e in collected]}"
)
assert isinstance(report, ResearchReport)
assert report.title is not None
assert isinstance(report.title, str)
assert any(evt.type == "log" for evt in collected)
client.close()

20
uv.lock
View file

@ -38,6 +38,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/5f/a0/d9ef19f780f319c21ee90ecfef4431cbeeca95bec7f14071785c17b6029b/accelerate-1.10.1-py3-none-any.whl", hash = "sha256:3621cff60b9a27ce798857ece05e2b9f56fcc71631cfb31ccf71f0359c311f11", size = 374909, upload-time = "2025-08-25T13:57:04.55Z" },
]
[[package]]
name = "ag-ui-protocol"
version = "0.1.10"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pydantic" },
]
sdist = { url = "https://files.pythonhosted.org/packages/67/bb/5a5ec893eea5805fb9a3db76a9888c3429710dfb6f24bbb37568f2cf7320/ag_ui_protocol-0.1.10.tar.gz", hash = "sha256:3213991c6b2eb24bb1a8c362ee270c16705a07a4c5962267a083d0959ed894f4", size = 6945, upload-time = "2025-11-06T15:17:17.068Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/8f/78/eb55fabaab41abc53f52c0918a9a8c0f747807e5306273f51120fd695957/ag_ui_protocol-0.1.10-py3-none-any.whl", hash = "sha256:c81e6981f30aabdf97a7ee312bfd4df0cd38e718d9fc10019c7d438128b93ab5", size = 7889, upload-time = "2025-11-06T15:17:15.325Z" },
]
[[package]]
name = "aiohappyeyeballs"
version = "2.6.1"
@ -1188,7 +1200,7 @@ dependencies = [
{ name = "lancedb" },
{ name = "pathspec" },
{ name = "pydantic" },
{ name = "pydantic-ai-slim", extra = ["fastmcp", "logfire", "openai"] },
{ name = "pydantic-ai-slim", extra = ["ag-ui", "fastmcp", "logfire", "openai"] },
{ name = "python-dotenv" },
{ name = "pyyaml" },
{ name = "rich" },
@ -1247,7 +1259,7 @@ requires-dist = [
{ name = "pydantic-ai-slim", extras = ["google"], marker = "extra == 'google'" },
{ name = "pydantic-ai-slim", extras = ["groq"], marker = "extra == 'groq'" },
{ name = "pydantic-ai-slim", extras = ["mistral"], marker = "extra == 'mistral'" },
{ name = "pydantic-ai-slim", extras = ["openai", "fastmcp", "logfire"], specifier = ">=1.11.1" },
{ name = "pydantic-ai-slim", extras = ["openai", "fastmcp", "logfire", "ag-ui"], specifier = ">=1.11.1" },
{ name = "pydantic-ai-slim", extras = ["vertexai"], marker = "extra == 'vertexai'" },
{ name = "python-dotenv", specifier = ">=1.2.1" },
{ name = "pyyaml", specifier = ">=6.0.3" },
@ -3019,6 +3031,10 @@ wheels = [
]
[package.optional-dependencies]
ag-ui = [
{ name = "ag-ui-protocol" },
{ name = "starlette" },
]
anthropic = [
{ name = "anthropic" },
]