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

Make graph settings part of config for both deep ask and research graphs

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-11-06 13:15:00 +02:00
parent fa93226a28
commit 782fbafb3d
No known key found for this signature in database
14 changed files with 206 additions and 139 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
README.md
View file

@ -88,8 +88,8 @@ 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 (
PlanNode,
ResearchContext,
ResearchDeps,
ResearchState,
@ -115,34 +115,22 @@ async with HaikuRAG("database.lancedb") as client:
print(answer)
# Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
graph = build_research_graph()
# Graph settings (provider, model, max_iterations, etc.) come from config
graph = build_research_graph(config=Config)
question = (
"What are the main drivers and trends of global temperature "
"anomalies since 1990?"
)
state = ResearchState(
context=ResearchContext(original_question=question),
max_iterations=2,
confidence_threshold=0.8,
max_concurrency=2,
)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
# Blocking run (final result only)
result = await graph.run(
PlanNode(provider="openai", model="gpt-4o-mini"),
state=state,
deps=deps,
)
print(result.output.title)
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,
PlanNode(provider="openai", model="gpt-4o-mini"),
state,
deps,
):
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}")

107
docs/agents.md
View file

@ -78,7 +78,7 @@ Key differences from Research:
Note on parallel execution:
- The `search_one` node is mapped over all questions in a batch
- Parallelism is controlled via `max_concurrency` using asyncio.Semaphore
- Parallelism is controlled via `max_concurrency`
- All questions in an iteration are processed before evaluation
CLI usage:
@ -95,25 +95,19 @@ 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
async with HaikuRAG(path_to_db) as client:
graph = build_deep_qa_graph(
provider="openai",
model="gpt-4o-mini"
)
# Use global config (recommended)
graph = build_deep_qa_graph(config=Config)
context = DeepQAContext(
original_question="What are the main features of haiku.rag?",
use_citations=True
)
state = DeepQAState(
context=context,
max_sub_questions=3,
max_iterations=2,
max_concurrency=1
)
state = DeepQAState.from_config(context=context, config=Config)
deps = DeepQADeps(client=client)
result = await graph.run(
@ -125,6 +119,33 @@ async with HaikuRAG(path_to_db) as client:
print(result.sources)
```
Alternative usage with custom config:
```python
# Create a custom config with different settings
from haiku.rag.config.models import AppConfig, QAConfig
custom_config = AppConfig(
qa=QAConfig(
provider="openai",
model="gpt-4o-mini",
max_sub_questions=5,
max_iterations=3,
max_concurrency=2,
)
)
graph = build_deep_qa_graph(config=custom_config)
context = DeepQAContext(
original_question="What are the main features of haiku.rag?",
use_citations=True
)
state = DeepQAState.from_config(context=context, config=custom_config)
deps = DeepQADeps(client=client)
result = await graph.run(state=state, deps=deps)
```
### Research Graph
The research workflow is implemented as a typed pydantic‑graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report — with clear stop conditions and shared state.
@ -166,39 +187,34 @@ Primary models:
Note on parallel execution:
- The `search_one` node is mapped over all questions in a batch
- Parallelism is controlled via `max_concurrency` using asyncio.Semaphore
- Parallelism is controlled via `max_concurrency`
- Analysis and decision nodes process results after each batch completes
CLI usage:
```bash
haiku-rag research "How does haiku.rag organize and query documents?" \
--max-iterations 2 \
--confidence-threshold 0.8 \
--max-concurrency 3 \
--verbose
# Basic usage (uses config from file or defaults)
haiku-rag research "How does haiku.rag organize and query documents?" --verbose
# With custom config file
haiku-rag --config my-research-config.yaml research "How does haiku.rag organize and query documents?" --verbose
```
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
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(
provider="openai",
model="gpt-4o-mini"
)
# Use global config (recommended)
graph = build_research_graph(config=Config)
question = "What are the main drivers and trends of global temperature anomalies since 1990?"
state = ResearchState(
context=ResearchContext(original_question=question),
max_iterations=2,
confidence_threshold=0.8,
max_concurrency=2,
)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
result = await graph.run(
@ -211,27 +227,44 @@ async with HaikuRAG(path_to_db) as client:
print(report.executive_summary)
```
Alternative usage with custom config:
```python
from haiku.rag.config.models import AppConfig, ResearchConfig
custom_config = AppConfig(
research=ResearchConfig(
provider="openai",
model="gpt-4o-mini",
max_iterations=5,
confidence_threshold=0.85,
max_concurrency=3,
)
)
graph = build_research_graph(config=custom_config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=custom_config)
deps = ResearchDeps(client=client)
result = await graph.run(state=state, deps=deps)
```
Python usage (streamed events):
```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.research.stream import stream_research_graph
async with HaikuRAG(path_to_db) as client:
graph = build_research_graph(
provider="openai",
model="gpt-4o-mini"
)
graph = build_research_graph(config=Config)
question = "What are the main drivers and trends of global temperature anomalies since 1990?"
state = ResearchState(
context=ResearchContext(original_question=question),
max_iterations=2,
confidence_threshold=0.8,
max_concurrency=2,
)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=client)
async for event in stream_research_graph(

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

@ -204,23 +204,29 @@ class HaikuRAGApp:
deep: bool = False,
verbose: bool = False,
):
"""Ask a question using the RAG system.
Args:
question: The question to ask
cite: Include citations in the answer
deep: Use deep QA mode (multi-step reasoning)
verbose: Show verbose output
"""
async with HaikuRAG(db_path=self.db_path) as self.client:
try:
if deep:
from rich.console import Console
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
graph = build_deep_qa_graph(
provider=Config.qa.provider,
model=Config.qa.model,
)
graph = build_deep_qa_graph(config=Config)
context = DeepQAContext(
original_question=question, use_citations=cite
)
state = DeepQAState(context=context)
state = DeepQAState.from_config(context=context, config=Config)
deps = DeepQADeps(
client=self.client, console=Console() if verbose else None
)
@ -240,28 +246,26 @@ class HaikuRAGApp:
async def research(
self,
question: str,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
verbose: bool = False,
):
"""Run research via the pydantic-graph pipeline (default)."""
"""Run research via the pydantic-graph pipeline.
Args:
question: The research question
verbose: Show verbose output
"""
async with HaikuRAG(db_path=self.db_path) as client:
try:
from haiku.rag.config import Config
if verbose:
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(
provider=Config.research.provider or Config.qa.provider,
model=Config.research.model or Config.qa.model,
)
graph = build_research_graph(config=Config)
context = ResearchContext(original_question=question)
state = ResearchState(
context=context,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
)
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(
client=client, console=self.console if verbose else None
)

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

@ -290,17 +290,6 @@ def research(
question: str = typer.Argument(
help="The research question to investigate",
),
max_iterations: int = typer.Option(
3,
"--max-iterations",
"-n",
help="Maximum search/analyze iterations",
),
confidence_threshold: float = typer.Option(
0.8,
"--confidence-threshold",
help="Minimum confidence (0-1) to stop",
),
db: Path = typer.Option(
Config.storage.data_dir / "haiku.rag.lancedb",
"--db",
@ -315,14 +304,7 @@ def research(
from haiku.rag.app import HaikuRAGApp
app = HaikuRAGApp(db_path=db)
asyncio.run(
app.research(
question=question,
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
verbose=verbose,
)
)
asyncio.run(app.research(question=question, verbose=verbose))
@cli.command("settings", help="Display current configuration settings")

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

@ -37,11 +37,17 @@ class RerankingConfig(BaseModel):
class QAConfig(BaseModel):
provider: str = "ollama"
model: str = "gpt-oss"
max_sub_questions: int = 3
max_iterations: int = 2
max_concurrency: int = 1
class ResearchConfig(BaseModel):
provider: str = "ollama"
model: str = "gpt-oss"
max_iterations: int = 3
confidence_threshold: float = 0.8
max_concurrency: int = 1
class ProcessingConfig(BaseModel):

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

@ -2,11 +2,11 @@ from pathlib import Path
from typing import Any
from fastmcp import FastMCP
from haiku.rag.client import HaikuRAG
from haiku.rag.research.models import ResearchReport
from pydantic import BaseModel
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.research.models import ResearchReport
class SearchResult(BaseModel):
@ -191,20 +191,16 @@ def create_mcp_server(db_path: Path) -> FastMCP:
try:
async with HaikuRAG(db_path) as rag:
if deep:
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.config import Config
graph = build_deep_qa_graph(
provider=Config.qa.provider,
model=Config.qa.model,
)
graph = build_deep_qa_graph(config=Config)
context = DeepQAContext(
original_question=question, use_citations=cite
)
state = DeepQAState(context=context)
state = DeepQAState.from_config(context=context, config=Config)
deps = DeepQADeps(client=rag)
result = await graph.run(state=state, deps=deps)
@ -218,9 +214,6 @@ def create_mcp_server(db_path: Path) -> FastMCP:
@mcp.tool()
async def research_question(
question: str,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
max_concurrency: int = 1,
) -> ResearchReport | None:
"""Run multi-agent research to investigate a complex question.
@ -229,9 +222,6 @@ def create_mcp_server(db_path: Path) -> FastMCP:
Args:
question: The research question to investigate.
max_iterations: Maximum search/analyze iterations (default: 3).
confidence_threshold: Minimum confidence score (0-1) to stop early (default: 0.8).
max_concurrency: Maximum concurrent sub-questions to process (default: 1).
Returns:
A research report with findings, or None if an error occurred.
@ -242,16 +232,9 @@ def create_mcp_server(db_path: Path) -> FastMCP:
from haiku.rag.research.state import ResearchDeps, ResearchState
async with HaikuRAG(db_path) as rag:
graph = build_research_graph(
provider=Config.research.provider or Config.qa.provider,
model=Config.research.model or Config.qa.model,
)
state = ResearchState(
context=ResearchContext(original_question=question),
max_iterations=max_iterations,
confidence_threshold=confidence_threshold,
max_concurrency=max_concurrency,
)
graph = build_research_graph(config=Config)
context = ResearchContext(original_question=question)
state = ResearchState.from_config(context=context, config=Config)
deps = ResearchDeps(client=rag)
result = await graph.run(state=state, deps=deps)

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

@ -1,5 +1,13 @@
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
@ -11,16 +19,21 @@ from haiku.rag.qa.deep.prompts import (
SYNTHESIS_PROMPT_WITH_CITATIONS,
)
from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
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
def build_deep_qa_graph(
provider: str, model: str
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,

25
haiku_rag_slim/haiku/rag/qa/deep/state.py
View file

@ -1,9 +1,14 @@
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.qa.deep.dependencies import DeepQAContext
from rich.console import Console
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
@dataclass
@ -24,3 +29,21 @@ class DeepQAState:
max_iterations: int = 2
max_concurrency: int = 1
iterations: int = 0
@classmethod
def from_config(cls, context: DeepQAContext, config: "AppConfig") -> "DeepQAState":
"""Create a DeepQAState from an AppConfig.
Args:
context: The DeepQAContext containing the question and settings
config: The AppConfig object (uses config.qa for state parameters)
Returns:
A configured DeepQAState instance
"""
return cls(
context=context,
max_sub_questions=config.qa.max_sub_questions,
max_iterations=config.qa.max_iterations,
max_concurrency=config.qa.max_concurrency,
)

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

@ -1,5 +1,13 @@
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
@ -19,16 +27,21 @@ from haiku.rag.research.prompts import (
SYNTHESIS_AGENT_PROMPT,
)
from haiku.rag.research.state import ResearchDeps, ResearchState
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
def build_research_graph(
provider: str, model: str
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,

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

@ -1,11 +1,16 @@
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
from rich.console import Console
if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig
@dataclass
@ -31,3 +36,23 @@ class ResearchState:
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,
)

4
tests/test_deep_qa.py
View file

@ -19,7 +19,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path):
monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
graph = build_deep_qa_graph(provider="test", model="test")
graph = build_deep_qa_graph()
state = DeepQAState(
context=DeepQAContext(
@ -53,7 +53,7 @@ async def test_deep_qa_with_citations(monkeypatch, temp_db_path):
monkeypatch.setattr("haiku.rag.graph_common.utils.get_model", test_model_factory)
monkeypatch.setattr("haiku.rag.qa.deep.graph.get_model", test_model_factory)
graph = build_deep_qa_graph(provider="test", model="test")
graph = build_deep_qa_graph()
state = DeepQAState(
context=DeepQAContext(original_question="What is Python?", use_citations=True),

3
tests/test_mcp.py
View file

@ -309,9 +309,6 @@ async def test_mcp_research_question():
result = await research_tool.fn( # type: ignore[attr-defined]
question="Research question?",
max_iterations=1,
confidence_threshold=0.5,
max_concurrency=1,
)
assert result is not None

2
tests/test_research_graph.py
View file

@ -6,7 +6,7 @@ from haiku.rag.research.state import ResearchState
def test_build_graph_and_state():
graph = build_research_graph(provider="openai", model="gpt-4")
graph = build_research_graph()
assert graph is not None
state = ResearchState(

2
tests/test_research_graph_integration.py
View file

@ -20,7 +20,7 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path):
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(provider="test", model="test")
graph = build_research_graph()
state = ResearchState(
context=ResearchContext(original_question="What is haiku.rag?"),