From 4855051936313e44e36ab20c4743b356119ee68b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 19 Feb 2026 13:41:19 +0200 Subject: [PATCH] Add research() to client and add haiku.skills dependency, remove --deep flag and simplify app --- evaluations/evaluations/benchmark.py | 41 ++-------- haiku_rag_slim/haiku/rag/app.py | 60 ++------------ haiku_rag_slim/haiku/rag/cli.py | 6 -- haiku_rag_slim/haiku/rag/client.py | 61 +++++++++++++- haiku_rag_slim/haiku/rag/mcp.py | 38 +-------- haiku_rag_slim/pyproject.toml | 1 + tests/test_client_research.py | 118 +++++++++++++++++++++++++++ uv.lock | 19 +++++ 8 files changed, 214 insertions(+), 130 deletions(-) create mode 100644 tests/test_client_research.py diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 247f2df5..2dcae976 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -17,9 +17,6 @@ from rich.progress import Progress from evaluations.config import DatasetSpec from evaluations.datasets import DATASETS from evaluations.evaluators import ANSWER_EQUIVALENCE_RUBRIC -from haiku.rag.agents.research.dependencies import ResearchContext -from haiku.rag.agents.research.graph import build_research_graph -from haiku.rag.agents.research.state import ResearchDeps, ResearchState from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, find_config_file, load_yaml_config from haiku.rag.config.models import ModelConfig @@ -42,13 +39,11 @@ def build_experiment_metadata( test_cases: int, config: AppConfig, judge_config: ModelConfig, - deep: bool = False, ) -> dict[str, Any]: """Build experiment metadata for Logfire tracking.""" return { "dataset": dataset_key, "test_cases": test_cases, - "deep_ask": deep, "embedder_provider": config.embeddings.model.provider, "embedder_model": config.embeddings.model.name, "embedder_dim": config.embeddings.model.vector_dim, @@ -270,7 +265,6 @@ async def run_qa_benchmark( limit: int | None = None, name: str | None = None, db_path: Path | None = None, - deep: bool = False, ) -> ReportCaseFailure[str, str, dict[str, str]] | None: corpus = spec.qa_loader() if limit is not None: @@ -303,32 +297,19 @@ async def run_qa_benchmark( db = spec.db_path(db_path) async with HaikuRAG(db, config=config) as rag: - if deep: - graph = build_research_graph(config=config) + qa = get_qa_agent(rag, system_prompt=spec.system_prompt) - async def answer_question(question: str) -> str: - context = ResearchContext(original_question=question) - state = ResearchState.from_config(context=context, config=config) - deps = ResearchDeps(client=rag) - report = await graph.run(state=state, deps=deps) - return report.executive_summary if report else "" - else: - qa = get_qa_agent(rag, system_prompt=spec.system_prompt) - - async def answer_question(question: str) -> str: - answer, _ = await qa.answer(question) - return answer + async def answer_question(question: str) -> str: + answer, _ = await qa.answer(question) + return answer eval_name = name if name is not None else f"{spec.key}_qa_evaluation" - if deep: - eval_name = f"{eval_name}_deep" experiment_metadata = build_experiment_metadata( dataset_key=spec.key, test_cases=len(cases), config=config, judge_config=judge_config, - deep=deep, ) report = await evaluation_dataset.evaluate( @@ -378,7 +359,6 @@ async def evaluate_dataset( db_path: Path | None, vacuum_interval: int = 100, multimodal_only: bool = False, - deep: bool = False, ) -> None: if not skip_db: console.print(f"Using dataset: {spec.key}", style="bold magenta") @@ -398,11 +378,8 @@ async def evaluate_dataset( ) if not skip_qa: - mode_label = "deep QA" if deep else "QA" - console.print(f"\nRunning {mode_label} benchmarks...", style="bold yellow") - await run_qa_benchmark( - spec, config, limit=limit, name=name, db_path=db_path, deep=deep - ) + console.print("\nRunning QA benchmarks...", style="bold yellow") + await run_qa_benchmark(spec, config, limit=limit, name=name, db_path=db_path) app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.") @@ -434,11 +411,6 @@ def run( "--multimodal-only", help="Only evaluate queries requiring image understanding.", ), - deep: bool = typer.Option( - False, - "--deep", - help="Use deep QA mode (multi-step reasoning with research graph).", - ), ) -> None: spec = DATASETS.get(dataset.lower()) if spec is None: @@ -477,7 +449,6 @@ def run( db_path=db, vacuum_interval=vacuum_interval, multimodal_only=multimodal_only, - deep=deep, ) ) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 7fa5618a..dbf00936 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -18,9 +18,6 @@ from rich.progress import ( ) from rich.syntax import Syntax -from haiku.rag.agents.research.dependencies import ResearchContext -from haiku.rag.agents.research.graph import build_research_graph -from haiku.rag.agents.research.state import ResearchDeps, ResearchState from haiku.rag.client import HaikuRAG, RebuildMode from haiku.rag.config import AppConfig, Config from haiku.rag.mcp import create_mcp_server @@ -375,7 +372,6 @@ class HaikuRAGApp: # pragma: no cover self, question: str, cite: bool = False, - deep: bool = False, filter: str | None = None, ): """Ask a question using the RAG system. @@ -383,7 +379,6 @@ class HaikuRAGApp: # pragma: no cover Args: question: The question to ask cite: Include citations in the answer - deep: Use deep QA mode (multi-step reasoning) filter: SQL WHERE clause to filter documents """ async with HaikuRAG( @@ -392,46 +387,15 @@ class HaikuRAGApp: # pragma: no cover read_only=self.read_only, before=self.before, ) as self.client: - citations = [] - if deep: - graph = build_research_graph(config=self.config) - context = ResearchContext(original_question=question) - state = ResearchState.from_config( - context=context, - config=self.config, - max_iterations=1, - ) - state.search_filter = filter - deps = ResearchDeps(client=self.client) + answer, citations = await self.client.ask(question, filter=filter) - report = await graph.run(state=state, deps=deps) - - self.console.print(f"[bold blue]Question:[/bold blue] {question}") - self.console.print() - if report: - self.console.print("[bold green]Answer:[/bold green]") - self.console.print(Markdown(report.executive_summary)) - if report.main_findings: - self.console.print() - self.console.print("[bold cyan]Key Findings:[/bold cyan]") - for finding in report.main_findings: - self.console.print(f"• {finding}") - if report.sources_summary: - self.console.print() - self.console.print("[bold cyan]Sources:[/bold cyan]") - self.console.print(report.sources_summary) - else: - self.console.print("[yellow]No answer generated.[/yellow]") - else: - answer, citations = await self.client.ask(question, filter=filter) - - self.console.print(f"[bold blue]Question:[/bold blue] {question}") - self.console.print() - self.console.print("[bold green]Answer:[/bold green]") - self.console.print(Markdown(answer)) - if cite and citations: - for renderable in format_citations_rich(citations): - self.console.print(renderable) + self.console.print(f"[bold blue]Question:[/bold blue] {question}") + self.console.print() + self.console.print("[bold green]Answer:[/bold green]") + self.console.print(Markdown(answer)) + if cite and citations: + for renderable in format_citations_rich(citations): + self.console.print(renderable) async def rlm( self, @@ -488,13 +452,7 @@ class HaikuRAGApp: # pragma: no cover 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) - state.search_filter = filter - deps = ResearchDeps(client=client) - - report = await graph.run(state=state, deps=deps) + report = await client.research(question=question, filter=filter) if report is None: self.console.print("[red]Research did not produce a report.[/red]") diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 79e2144e..f4053122 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -341,11 +341,6 @@ def ask( # pragma: no cover "--cite", help="Include citations in the response", ), - deep: bool = typer.Option( - False, - "--deep", - help="Use deep multi-agent QA for complex questions", - ), filter: str | None = typer.Option( None, "--filter", @@ -358,7 +353,6 @@ def ask( # pragma: no cover app.ask( question=question, cite=cite, - deep=deep, filter=filter, ) ) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index b1159876..73d98d97 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -9,7 +9,7 @@ from dataclasses import dataclass from datetime import datetime from enum import Enum from pathlib import Path -from typing import TYPE_CHECKING, overload +from typing import TYPE_CHECKING, Literal, overload from urllib.parse import urlparse import httpx @@ -31,7 +31,11 @@ from haiku.rag.store.repositories.settings import SettingsRepository if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument - from haiku.rag.agents.research.models import Citation + from haiku.rag.agents.research.models import ( + Citation, + ConversationalAnswer, + ResearchReport, + ) from haiku.rag.agents.rlm.models import RLMResult logger = logging.getLogger(__name__) @@ -1324,6 +1328,59 @@ class HaikuRAG: qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt) return await qa_agent.answer(question, filter=filter) + @overload + async def research( + self, + question: str, + *, + output_mode: Literal["report"] = ..., + filter: str | None = ..., + max_iterations: int | None = ..., + ) -> "ResearchReport": ... + + @overload + async def research( + self, + question: str, + *, + output_mode: Literal["conversational"], + filter: str | None = ..., + max_iterations: int | None = ..., + ) -> "ConversationalAnswer": ... + + async def research( + self, + question: str, + *, + output_mode: Literal["report", "conversational"] = "report", + filter: str | None = None, + max_iterations: int | None = None, + ) -> "ResearchReport | ConversationalAnswer": + """Run multi-agent research to investigate a question. + + Args: + question: The research question to investigate. + output_mode: "report" for ResearchReport, "conversational" for ConversationalAnswer. + filter: SQL WHERE clause to filter documents. + max_iterations: Override max iterations (None uses config default). + + Returns: + ResearchReport or ConversationalAnswer based on output_mode. + """ + from haiku.rag.agents.research.dependencies import ResearchContext + from haiku.rag.agents.research.graph import build_research_graph + from haiku.rag.agents.research.state import ResearchDeps, ResearchState + + graph = build_research_graph(config=self._config, output_mode=output_mode) + context = ResearchContext(original_question=question) + state = ResearchState.from_config( + context=context, config=self._config, max_iterations=max_iterations + ) + state.search_filter = filter + deps = ResearchDeps(client=self) + + return await graph.run(state=state, deps=deps) + async def rlm( self, question: str, diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 974d9439..06d93eeb 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -171,42 +171,19 @@ def create_mcp_server( # pragma: no cover async def ask_question( question: str, cite: bool = False, - deep: bool = False, ) -> str: """Ask a question using the QA agent. Args: question: The question to ask. cite: Whether to include citations in the response. - deep: Use deep multi-agent QA for complex questions that require decomposition. Returns: The answer as a string. """ try: async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - if deep: - from haiku.rag.agents.research.dependencies import ResearchContext - from haiku.rag.agents.research.graph import build_research_graph - from haiku.rag.agents.research.state import ( - ResearchDeps, - ResearchState, - ) - - graph = build_research_graph(config=config) - context = ResearchContext(original_question=question) - state = ResearchState.from_config( - context=context, - config=config, - max_iterations=2, - ) - deps = ResearchDeps(client=rag) - - result = await graph.run(state=state, deps=deps) - answer = result.executive_summary - citations = [] - else: - answer, citations = await rag.ask(question) + answer, citations = await rag.ask(question) if cite and citations: answer += "\n\n" + format_citations(citations) return answer @@ -229,19 +206,8 @@ def create_mcp_server( # pragma: no cover A research report with findings, or None if an error occurred. """ try: - from haiku.rag.agents.research.dependencies import ResearchContext - from haiku.rag.agents.research.graph import build_research_graph - from haiku.rag.agents.research.state import ResearchDeps, ResearchState - async with HaikuRAG(db_path, config=config, read_only=read_only) as rag: - 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) - - return result + return await rag.research(question=question) except Exception: return None diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 5ba9eac3..84f34710 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -24,6 +24,7 @@ classifiers = [ dependencies = [ "cachetools>=5.5.0", "docling-core==2.65.1", + "haiku.skills>=0.3.0", "httpx>=0.28.1", "jsonpatch>=1.33", "lancedb==0.29.2", diff --git a/tests/test_client_research.py b/tests/test_client_research.py new file mode 100644 index 00000000..f798afa4 --- /dev/null +++ b/tests/test_client_research.py @@ -0,0 +1,118 @@ +from pathlib import Path +from unittest.mock import AsyncMock, patch + +import pytest + +from haiku.rag.agents.research.models import ConversationalAnswer, ResearchReport +from haiku.rag.client import HaikuRAG + + +@pytest.fixture(scope="module") +def vcr_cassette_dir(): + return str(Path(__file__).parent / "cassettes" / "test_client_research") + + +async def test_client_research_report(temp_db_path): + """Test client.research() delegates to research graph in report mode.""" + mock_report = ResearchReport( + title="Test Report", + executive_summary="Summary", + main_findings=["Finding 1"], + conclusions=["Conclusion 1"], + sources_summary="Sources", + ) + + with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build: + mock_graph = AsyncMock() + mock_graph.run = AsyncMock(return_value=mock_report) + mock_build.return_value = mock_graph + + async with HaikuRAG(temp_db_path, create=True) as client: + result = await client.research(question="What is X?") + + assert result is mock_report + mock_build.assert_called_once() + # Verify output_mode passed correctly + _, kwargs = mock_build.call_args + assert kwargs["output_mode"] == "report" + + # Verify graph.run was called with correct state/deps + mock_graph.run.assert_called_once() + call_kwargs = mock_graph.run.call_args[1] + assert call_kwargs["state"].context.original_question == "What is X?" + assert isinstance(call_kwargs["deps"].client, HaikuRAG) + + +async def test_client_research_conversational(temp_db_path): + """Test client.research() with conversational output mode.""" + mock_answer = ConversationalAnswer( + answer="The answer is 42.", + confidence=0.95, + ) + + with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build: + mock_graph = AsyncMock() + mock_graph.run = AsyncMock(return_value=mock_answer) + mock_build.return_value = mock_graph + + async with HaikuRAG(temp_db_path, create=True) as client: + result = await client.research( + question="What is X?", + output_mode="conversational", + ) + + assert result is mock_answer + _, kwargs = mock_build.call_args + assert kwargs["output_mode"] == "conversational" + + +async def test_client_research_passes_filter(temp_db_path): + """Test client.research() passes filter to state.""" + mock_report = ResearchReport( + title="Test", + executive_summary="Summary", + main_findings=[], + conclusions=[], + sources_summary="", + ) + + with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build: + mock_graph = AsyncMock() + mock_graph.run = AsyncMock(return_value=mock_report) + mock_build.return_value = mock_graph + + async with HaikuRAG(temp_db_path, create=True) as client: + await client.research( + question="What is X?", + filter="uri LIKE '%test%'", + ) + + call_kwargs = mock_graph.run.call_args[1] + assert call_kwargs["state"].search_filter == "uri LIKE '%test%'" + + +async def test_client_research_uses_config(temp_db_path): + """Test client.research() passes config to graph builder and state.""" + mock_report = ResearchReport( + title="Test", + executive_summary="Summary", + main_findings=[], + conclusions=[], + sources_summary="", + ) + + with patch("haiku.rag.agents.research.graph.build_research_graph") as mock_build: + mock_graph = AsyncMock() + mock_graph.run = AsyncMock(return_value=mock_report) + mock_build.return_value = mock_graph + + async with HaikuRAG(temp_db_path, create=True) as client: + await client.research(question="What is X?") + + _, kwargs = mock_build.call_args + assert kwargs["config"] is client._config + + call_kwargs = mock_graph.run.call_args[1] + state = call_kwargs["state"] + assert state.max_iterations == client._config.research.max_iterations + assert state.max_concurrency == client._config.research.max_concurrency diff --git a/uv.lock b/uv.lock index 229f13b1..17c4041f 100644 --- a/uv.lock +++ b/uv.lock @@ -1449,6 +1449,7 @@ source = { editable = "haiku_rag_slim" } dependencies = [ { name = "cachetools" }, { name = "docling-core" }, + { name = "haiku-skills" }, { name = "httpx" }, { name = "jsonpatch" }, { name = "lancedb" }, @@ -1512,6 +1513,7 @@ requires-dist = [ { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.20.1" }, { name = "docling", marker = "extra == 'docling'", specifier = "==2.73.1" }, { name = "docling-core", specifier = "==2.65.1" }, + { name = "haiku-skills", specifier = ">=0.3.0" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "jsonpatch", specifier = ">=1.33" }, { name = "lancedb", specifier = "==0.29.2" }, @@ -1540,6 +1542,20 @@ requires-dist = [ ] provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"] +[[package]] +name = "haiku-skills" +version = "0.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pydantic" }, + { name = "pydantic-ai-slim", extra = ["mcp"] }, + { name = "pyyaml" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/d6/11dbef98d7b04f5aacc7da3354fbe7bbb9ab5aa5948517d135f68f457d7d/haiku_skills-0.3.0.tar.gz", hash = "sha256:191414a840653ba938aa8ce1804f10cc489426d44fdcc7f261feb2e32e1be041", size = 158600, upload-time = "2026-02-19T10:34:28.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/a2/279333965841a2ecda2e5cc695306998f9770ac5b92b2d6ad8c50206b162/haiku_skills-0.3.0-py3-none-any.whl", hash = "sha256:ac5ddaea07d920ffec368ea7cf8e88963e727feece01571fd40dfc31da4b3ccd", size = 20359, upload-time = "2026-02-19T10:34:27.235Z" }, +] + [[package]] name = "hf-xet" version = "1.2.0" @@ -3649,6 +3665,9 @@ groq = [ logfire = [ { name = "logfire", extra = ["httpx"] }, ] +mcp = [ + { name = "mcp" }, +] mistral = [ { name = "mistralai" }, ]