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

Add research() to client and add haiku.skills dependency, remove --deep flag and simplify app

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-02-19 13:41:19 +02:00
parent 0c270ba42f
commit 4855051936
No known key found for this signature in database
8 changed files with 214 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
evaluations/evaluations/benchmark.py
View file

@ -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,
)
)

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

@ -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]")

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

@ -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,
)
)

61
haiku_rag_slim/haiku/rag/client.py
View file

@ -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,

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

@ -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

1
haiku_rag_slim/pyproject.toml
View file

@ -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",

118
tests/test_client_research.py Normal file
View file

@ -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

19
uv.lock
View file

@ -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" },
]