diff --git a/README.md b/README.md index 6c285fb6..eb0a429b 100644 --- a/README.md +++ b/README.md @@ -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}") diff --git a/docs/agents.md b/docs/agents.md index b9dc8065..e9fb688a 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -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( diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 7eea7e6c..6aa2c178 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -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 ) diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 11e059c7..edefec59 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -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") diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 4c654e7a..801c6dcb 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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): diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 8a4573fa..3c18587c 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/qa/deep/graph.py b/haiku_rag_slim/haiku/rag/qa/deep/graph.py index 4be28cd4..9ded283c 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/graph.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/graph.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/qa/deep/state.py b/haiku_rag_slim/haiku/rag/qa/deep/state.py index 46750242..0e07098e 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/state.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/state.py @@ -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, + ) diff --git a/haiku_rag_slim/haiku/rag/research/graph.py b/haiku_rag_slim/haiku/rag/research/graph.py index 117bf741..a756476e 100644 --- a/haiku_rag_slim/haiku/rag/research/graph.py +++ b/haiku_rag_slim/haiku/rag/research/graph.py @@ -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, diff --git a/haiku_rag_slim/haiku/rag/research/state.py b/haiku_rag_slim/haiku/rag/research/state.py index 2c748103..bfc48cbb 100644 --- a/haiku_rag_slim/haiku/rag/research/state.py +++ b/haiku_rag_slim/haiku/rag/research/state.py @@ -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, + ) diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py index 5dab9475..c220fcda 100644 --- a/tests/test_deep_qa.py +++ b/tests/test_deep_qa.py @@ -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), diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 052a409f..de2956d3 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -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 diff --git a/tests/test_research_graph.py b/tests/test_research_graph.py index d4c41afc..0a7c0b05 100644 --- a/tests/test_research_graph.py +++ b/tests/test_research_graph.py @@ -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( diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 8c4239f7..5bdf40a6 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -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?"),