diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 19c80c7c..c6d5446a 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -9,7 +9,7 @@ repos: - id: debug-statements - repo: https://github.com/astral-sh/ruff-pre-commit # Ruff version. - rev: v0.11.4 + rev: v0.14.3 hooks: # Run the linter. - id: ruff @@ -17,6 +17,6 @@ repos: - id: ruff-format - repo: https://github.com/RobertCraigie/pyright-python - rev: v1.1.399 + rev: v1.1.407 hooks: - id: pyright diff --git a/haiku_rag_slim/haiku/rag/a2a/__init__.py b/haiku_rag_slim/haiku/rag/a2a/__init__.py index 892056eb..7c3a00db 100644 --- a/haiku_rag_slim/haiku/rag/a2a/__init__.py +++ b/haiku_rag_slim/haiku/rag/a2a/__init__.py @@ -6,7 +6,7 @@ import logfire from pydantic_ai import Agent, RunContext from haiku.rag.config import Config -from haiku.rag.graph.common import get_model +from haiku.rag.qa.deep.common import get_model from .context import load_message_history, save_message_history from .models import AgentDependencies, SearchResult @@ -138,7 +138,11 @@ def create_a2a_app( if security_schemes or security: # Monkey-patch the agent card endpoint to include security async def _agent_card_endpoint_with_security(request): - from fasta2a.schema import AgentCapabilities, AgentCard, agent_card_ta + from fasta2a.schema import ( # type: ignore + AgentCapabilities, + AgentCard, + agent_card_ta, + ) from starlette.responses import Response if app._agent_card_json_schema is None: diff --git a/haiku_rag_slim/haiku/rag/graph/models.py b/haiku_rag_slim/haiku/rag/graph/models.py deleted file mode 100644 index 475b6e43..00000000 --- a/haiku_rag_slim/haiku/rag/graph/models.py +++ /dev/null @@ -1,24 +0,0 @@ -from pydantic import BaseModel, Field - - -class ResearchPlan(BaseModel): - main_question: str - sub_questions: list[str] - - -class SearchAnswer(BaseModel): - query: str = Field(description="The search query that was performed") - answer: str = Field(description="The answer generated based on the context") - context: list[str] = Field( - description=( - "Only the minimal set of relevant snippets (verbatim) that directly " - "support the answer" - ) - ) - sources: list[str] = Field( - description=( - "Document titles (if available) or URIs corresponding to the" - " snippets actually used in the answer (one per snippet; omit if none)" - ), - default_factory=list, - ) diff --git a/haiku_rag_slim/haiku/rag/graph/prompts.py b/haiku_rag_slim/haiku/rag/graph/prompts.py deleted file mode 100644 index 97f42aa8..00000000 --- a/haiku_rag_slim/haiku/rag/graph/prompts.py +++ /dev/null @@ -1,45 +0,0 @@ -PLAN_PROMPT = """You are the research orchestrator for a focused, iterative -workflow. - -Responsibilities: -1. Understand and decompose the main question -2. Propose a minimal, high‑leverage plan -3. Coordinate specialized agents to gather evidence -4. Iterate based on gaps and new findings - -Plan requirements: -- Produce at most 3 sub_questions that together cover the main question. -- Each sub_question must be a standalone, self‑contained query that can run - without extra context. Include concrete entities, scope, timeframe, and any - qualifiers. Avoid ambiguous pronouns (it/they/this/that). -- Prioritize the highest‑value aspects first; avoid redundancy and overlap. -- Prefer questions that are likely answerable from the current knowledge base; - if coverage is uncertain, make scopes narrower and specific. -- Order sub_questions by execution priority (most valuable first).""" - -SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist. - -Tasks: -1. Search the knowledge base for relevant evidence. -2. Analyze retrieved snippets. -3. Provide an answer strictly grounded in that evidence. - -Tool usage: -- Always call search_and_answer before drafting any answer. -- The tool returns snippets with verbatim `text`, a relevance `score`, and the - originating document identifier (document title if available, otherwise URI). -- You may call the tool multiple times to refine or broaden context, but do not - exceed 3 total calls. Favor precision over volume. -- Use scores to prioritize evidence, but include only the minimal subset of - snippet texts (verbatim) in SearchAnswer.context (typically 1‑4). -- Set SearchAnswer.sources to the corresponding document identifiers for the - snippets you used (title if available, otherwise URI; one per snippet; same - order as context). Context must be text‑only. -- If no relevant information is found, clearly say so and return an empty - context list and sources list. - -Answering rules: -- Be direct and specific; avoid meta commentary about the process. -- Do not include any claims not supported by the provided snippets. -- Prefer concise phrasing; avoid copying long passages. -- When evidence is partial, state the limits explicitly in the answer.""" diff --git a/haiku_rag_slim/haiku/rag/graph/common.py b/haiku_rag_slim/haiku/rag/qa/deep/common.py similarity index 78% rename from haiku_rag_slim/haiku/rag/graph/common.py rename to haiku_rag_slim/haiku/rag/qa/deep/common.py index 31877206..71815ba5 100644 --- a/haiku_rag_slim/haiku/rag/graph/common.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/common.py @@ -5,6 +5,7 @@ from pydantic_ai.providers.ollama import OllamaProvider from pydantic_ai.providers.openai import OpenAIProvider from haiku.rag.config import Config +from haiku.rag.qa.deep.models import SearchAnswer class HasEmitLog(Protocol): @@ -31,3 +32,10 @@ def get_model(provider: str, model: str) -> Any: def log(deps: HasEmitLog, state: Any, message: str) -> None: deps.emit_log(message, state) + + +def collect_answers_reducer( + acc: list[SearchAnswer], item: SearchAnswer | None +) -> list[SearchAnswer]: + """Reducer function to collect search answers, filtering out None values.""" + return acc + [item] if item else acc diff --git a/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py b/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py index c6017a61..eaf7957d 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/dependencies.py @@ -2,7 +2,7 @@ from pydantic import BaseModel, Field from rich.console import Console from haiku.rag.client import HaikuRAG -from haiku.rag.graph.models import SearchAnswer +from haiku.rag.qa.deep.models import SearchAnswer class DeepQAContext(BaseModel): diff --git a/haiku_rag_slim/haiku/rag/qa/deep/graph.py b/haiku_rag_slim/haiku/rag/qa/deep/graph.py index 28008f45..fb0a8760 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/graph.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/graph.py @@ -5,13 +5,18 @@ from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput from pydantic_graph.beta import GraphBuilder, StepContext -from haiku.rag.graph.common import get_model, log -from haiku.rag.graph.models import ResearchPlan, SearchAnswer -from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT +from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log from haiku.rag.qa.deep.dependencies import DeepQADependencies -from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation +from haiku.rag.qa.deep.models import ( + DeepQAAnswer, + DeepQAEvaluation, + ResearchPlan, + SearchAnswer, +) from haiku.rag.qa.deep.prompts import ( DECISION_PROMPT, + PLAN_PROMPT, + SEARCH_AGENT_PROMPT, SYNTHESIS_PROMPT, SYNTHESIS_PROMPT_WITH_CITATIONS, ) @@ -52,7 +57,7 @@ def build_deep_qa_graph(provider: str, model: str): return "\n\n".join(chunk.content for chunk, _ in expanded) prompt = ( - "Plan a focused approach for answering the main question.\n\n" + "Plan a focused approach for the main question.\n\n" f"Main question: {state.context.original_question}" ) @@ -62,9 +67,7 @@ def build_deep_qa_graph(provider: str, model: str): console=deps.console, ) plan_result = await plan_agent.run(prompt, deps=agent_deps) - state.context.sub_questions = list(plan_result.output.sub_questions)[ - : state.max_sub_questions - ] + state.context.sub_questions = list(plan_result.output.sub_questions) log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]") log( @@ -139,9 +142,21 @@ def build_deep_qa_graph(provider: str, model: str): return answer + @g.step + async def get_batch( + ctx: StepContext[DeepQAState, DeepQADeps, None | bool], + ) -> list[str] | None: + """Get next batch of questions from state.""" + state = ctx.state + take = max(1, state.max_concurrency) + batch: list[str] = [] + while state.context.sub_questions and len(batch) < take: + batch.append(state.context.sub_questions.pop(0)) + return batch if batch else None + @g.step async def decide( - ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer | None]], + ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]], ) -> bool: state = ctx.state deps = ctx.deps @@ -222,18 +237,6 @@ def build_deep_qa_graph(provider: str, model: str): return should_continue - @g.step - async def get_batch( - ctx: StepContext[DeepQAState, DeepQADeps, None | bool], - ) -> list[str] | None: - """Get next batch of questions from state.""" - state = ctx.state - take = max(1, state.max_concurrency) - batch: list[str] = [] - while state.context.sub_questions and len(batch) < take: - batch.append(state.context.sub_questions.pop(0)) - return batch if batch else None - @g.step async def synthesize( ctx: StepContext[DeepQAState, DeepQADeps, None | bool], @@ -287,13 +290,8 @@ def build_deep_qa_graph(provider: str, model: str): return result.output # Build the graph structure - def collect_reducer( - acc: list[SearchAnswer | None], item: SearchAnswer | None - ) -> list[SearchAnswer | None]: - return acc + [item] if item else acc - collect_answers = g.join( - collect_reducer, + collect_answers_reducer, initial_factory=lambda: [], ) diff --git a/haiku_rag_slim/haiku/rag/qa/deep/models.py b/haiku_rag_slim/haiku/rag/qa/deep/models.py index e33dfc95..2a7bad79 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/models.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/models.py @@ -1,6 +1,29 @@ from pydantic import BaseModel, Field +class ResearchPlan(BaseModel): + main_question: str + sub_questions: list[str] + + +class SearchAnswer(BaseModel): + query: str = Field(description="The search query that was performed") + answer: str = Field(description="The answer generated based on the context") + context: list[str] = Field( + description=( + "Only the minimal set of relevant snippets (verbatim) that directly " + "support the answer" + ) + ) + sources: list[str] = Field( + description=( + "Document titles (if available) or URIs corresponding to the" + " snippets actually used in the answer (one per snippet; omit if none)" + ), + default_factory=list, + ) + + class DeepQAEvaluation(BaseModel): is_sufficient: bool = Field( description="Whether we have sufficient information to answer the question" diff --git a/haiku_rag_slim/haiku/rag/qa/deep/prompts.py b/haiku_rag_slim/haiku/rag/qa/deep/prompts.py index 2b0d16b7..0bd399b1 100644 --- a/haiku_rag_slim/haiku/rag/qa/deep/prompts.py +++ b/haiku_rag_slim/haiku/rag/qa/deep/prompts.py @@ -1,3 +1,49 @@ +PLAN_PROMPT = """You are the research orchestrator for a focused, iterative +workflow. + +Responsibilities: +1. Understand and decompose the main question +2. Propose a minimal, high‑leverage plan +3. Coordinate specialized agents to gather evidence +4. Iterate based on gaps and new findings + +Plan requirements: +- Produce at most 3 sub_questions that together cover the main question. +- Each sub_question must be a standalone, self‑contained query that can run + without extra context. Include concrete entities, scope, timeframe, and any + qualifiers. Avoid ambiguous pronouns (it/they/this/that). +- Prioritize the highest‑value aspects first; avoid redundancy and overlap. +- Prefer questions that are likely answerable from the current knowledge base; + if coverage is uncertain, make scopes narrower and specific. +- Order sub_questions by execution priority (most valuable first).""" + +SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist. + +Tasks: +1. Search the knowledge base for relevant evidence. +2. Analyze retrieved snippets. +3. Provide an answer strictly grounded in that evidence. + +Tool usage: +- Always call search_and_answer before drafting any answer. +- The tool returns snippets with verbatim `text`, a relevance `score`, and the + originating document identifier (document title if available, otherwise URI). +- You may call the tool multiple times to refine or broaden context, but do not + exceed 3 total calls. Favor precision over volume. +- Use scores to prioritize evidence, but include only the minimal subset of + snippet texts (verbatim) in SearchAnswer.context (typically 1‑4). +- Set SearchAnswer.sources to the corresponding document identifiers for the + snippets you used (title if available, otherwise URI; one per snippet; same + order as context). Context must be text‑only. +- If no relevant information is found, clearly say so and return an empty + context list and sources list. + +Answering rules: +- Be direct and specific; avoid meta commentary about the process. +- Do not include any claims not supported by the provided snippets. +- Prefer concise phrasing; avoid copying long passages. +- When evidence is partial, state the limits explicitly in the answer.""" + SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers. Task: diff --git a/haiku_rag_slim/haiku/rag/research/__init__.py b/haiku_rag_slim/haiku/rag/research/__init__.py index e49dd051..a289a077 100644 --- a/haiku_rag_slim/haiku/rag/research/__init__.py +++ b/haiku_rag_slim/haiku/rag/research/__init__.py @@ -1,3 +1,3 @@ -from haiku.rag.graph.models import SearchAnswer +from haiku.rag.qa.deep.models import SearchAnswer from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.research.models import EvaluationResult, ResearchReport diff --git a/haiku_rag_slim/haiku/rag/research/dependencies.py b/haiku_rag_slim/haiku/rag/research/dependencies.py index 9606c658..47e2b406 100644 --- a/haiku_rag_slim/haiku/rag/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/research/dependencies.py @@ -4,7 +4,7 @@ from pydantic import BaseModel, Field from rich.console import Console from haiku.rag.client import HaikuRAG -from haiku.rag.graph.models import SearchAnswer +from haiku.rag.qa.deep.models import SearchAnswer from haiku.rag.research.models import ( GapRecord, InsightAnalysis, diff --git a/haiku_rag_slim/haiku/rag/research/graph.py b/haiku_rag_slim/haiku/rag/research/graph.py index 4b54aafd..bc2d46f4 100644 --- a/haiku_rag_slim/haiku/rag/research/graph.py +++ b/haiku_rag_slim/haiku/rag/research/graph.py @@ -5,9 +5,9 @@ from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput from pydantic_graph.beta import GraphBuilder, StepContext -from haiku.rag.graph.common import get_model, log -from haiku.rag.graph.models import ResearchPlan, SearchAnswer -from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT +from haiku.rag.qa.deep.common import collect_answers_reducer, get_model, log +from haiku.rag.qa.deep.models import ResearchPlan, SearchAnswer +from haiku.rag.qa.deep.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT from haiku.rag.research.common import ( format_analysis_for_prompt, format_context_for_prompt, @@ -60,7 +60,7 @@ def build_research_graph(provider: str, model: str): return "\n\n".join(chunk.content for chunk, _ in expanded) prompt = ( - "Plan a focused research approach for the main question.\n\n" + "Plan a focused approach for the main question.\n\n" f"Main question: {state.context.original_question}" ) @@ -73,7 +73,7 @@ def build_research_graph(provider: str, model: str): 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]✅ Research Plan Created:[/bold green]") + log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]") log( deps, state, @@ -147,9 +147,21 @@ def build_research_graph(provider: str, model: str): return answer + @g.step + async def get_batch( + ctx: StepContext[ResearchState, ResearchDeps, None | bool], + ) -> list[str] | None: + """Get next batch of questions from state.""" + state = ctx.state + take = max(1, state.max_concurrency) + batch: list[str] = [] + while state.context.sub_questions and len(batch) < take: + batch.append(state.context.sub_questions.pop(0)) + return batch if batch else None + @g.step async def analyze_insights( - ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer | None]], + ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer]], ) -> None: state = ctx.state deps = ctx.deps @@ -297,18 +309,6 @@ def build_research_graph(provider: str, model: str): return should_continue - @g.step - async def get_batch( - ctx: StepContext[ResearchState, ResearchDeps, None | bool], - ) -> list[str] | None: - """Get next batch of questions from state.""" - state = ctx.state - take = max(1, state.max_concurrency) - batch: list[str] = [] - while state.context.sub_questions and len(batch) < take: - batch.append(state.context.sub_questions.pop(0)) - return batch if batch else None - @g.step async def synthesize( ctx: StepContext[ResearchState, ResearchDeps, None | bool], @@ -348,13 +348,8 @@ def build_research_graph(provider: str, model: str): return result.output # Build the graph structure - def collect_reducer( - acc: list[SearchAnswer | None], item: SearchAnswer | None - ) -> list[SearchAnswer | None]: - return acc + [item] if item else acc - collect_answers = g.join( - collect_reducer, + collect_answers_reducer, initial_factory=lambda: [], ) diff --git a/tests/test_app.py b/tests/test_app.py index 73657202..87c484a5 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -383,11 +383,9 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): from haiku.rag.qa.deep.models import DeepQAAnswer mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"]) - mock_result = MagicMock() - mock_result.output = mock_output mock_graph = AsyncMock() - mock_graph.run.return_value = mock_result + mock_graph.run.return_value = mock_output mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client @@ -415,11 +413,9 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): mock_output = DeepQAAnswer( answer="Deep QA answer with citations [test.md]", sources=["test.md"] ) - mock_result = MagicMock() - mock_result.output = mock_output mock_graph = AsyncMock() - mock_graph.run.return_value = mock_result + mock_graph.run.return_value = mock_output mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client @@ -445,11 +441,9 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): from haiku.rag.qa.deep.models import DeepQAAnswer mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"]) - mock_result = MagicMock() - mock_result.output = mock_output mock_graph = AsyncMock() - mock_graph.run.return_value = mock_result + mock_graph.run.return_value = mock_output mock_client = AsyncMock() mock_client.__aenter__.return_value = mock_client diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py index 40971386..6a6bcf87 100644 --- a/tests/test_deep_qa.py +++ b/tests/test_deep_qa.py @@ -2,9 +2,9 @@ import pytest from pydantic_ai.models.test import TestModel from haiku.rag.client import HaikuRAG -from haiku.rag.graph.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.models import SearchAnswer from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState @@ -16,7 +16,7 @@ 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.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.qa.deep.common.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") @@ -50,7 +50,7 @@ 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.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.qa.deep.common.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") diff --git a/tests/test_mcp.py b/tests/test_mcp.py index e6a303c6..052a409f 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -257,7 +257,7 @@ async def test_mcp_ask_question_deep(): mock_graph = AsyncMock() mock_result = AsyncMock() - mock_result.output.answer = "Deep answer" + mock_result.answer = "Deep answer" mock_graph.run = AsyncMock(return_value=mock_result) mock_graph_builder.return_value = mock_graph @@ -299,9 +299,7 @@ async def test_mcp_research_question(): mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) mock_graph = AsyncMock() - mock_result = AsyncMock() - mock_result.output = mock_report - mock_graph.run = AsyncMock(return_value=mock_result) + mock_graph.run = AsyncMock(return_value=mock_report) mock_graph_builder.return_value = mock_graph tools = await mcp.get_tools() diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 93504d7f..69a684cd 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -17,7 +17,7 @@ async def test_graph_end_to_end_with_test_model(monkeypatch, temp_db_path): def test_model_factory(provider, model): return TestModel() - monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory) + monkeypatch.setattr("haiku.rag.qa.deep.common.get_model", test_model_factory) monkeypatch.setattr("haiku.rag.research.graph.get_model", test_model_factory) graph = build_research_graph(provider="test", model="test")