diff --git a/src/haiku/rag/graph/__init__.py b/src/haiku/rag/graph/__init__.py new file mode 100644 index 00000000..18013005 --- /dev/null +++ b/src/haiku/rag/graph/__init__.py @@ -0,0 +1 @@ +from haiku.rag.graph.models import ResearchPlan, SearchAnswer diff --git a/src/haiku/rag/graph/base.py b/src/haiku/rag/graph/base.py new file mode 100644 index 00000000..1d78fa7a --- /dev/null +++ b/src/haiku/rag/graph/base.py @@ -0,0 +1,31 @@ +from typing import Protocol, runtime_checkable + +from pydantic import BaseModel, Field +from rich.console import Console + +from haiku.rag.client import HaikuRAG +from haiku.rag.graph.models import SearchAnswer + + +@runtime_checkable +class GraphContext(Protocol): + """Protocol for graph context objects.""" + + original_question: str + sub_questions: list[str] + qa_responses: list[SearchAnswer] + + def add_qa_response(self, qa: SearchAnswer) -> None: ... + + +class BaseGraphDeps(BaseModel): + """Base dependencies for graph nodes.""" + + model_config = {"arbitrary_types_allowed": True} + + client: HaikuRAG = Field(description="RAG client for document operations") + console: Console | None = None + + def emit_log(self, message: str) -> None: + if self.console: + self.console.print(message) diff --git a/src/haiku/rag/graph/common.py b/src/haiku/rag/graph/common.py new file mode 100644 index 00000000..22009ef8 --- /dev/null +++ b/src/haiku/rag/graph/common.py @@ -0,0 +1,32 @@ +from typing import TYPE_CHECKING, Any + +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.ollama import OllamaProvider +from pydantic_ai.providers.openai import OpenAIProvider + +from haiku.rag.config import Config + +if TYPE_CHECKING: # pragma: no cover + from haiku.rag.research.state import ResearchDeps, ResearchState + + +def get_model(provider: str, model: str) -> Any: + if provider == "ollama": + return OpenAIChatModel( + model_name=model, + provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), + ) + elif provider == "vllm": + return OpenAIChatModel( + model_name=model, + provider=OpenAIProvider( + base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1", + api_key="none", + ), + ) + else: + return f"{provider}:{model}" + + +def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None: + deps.emit_log(msg, state) diff --git a/src/haiku/rag/graph/models.py b/src/haiku/rag/graph/models.py new file mode 100644 index 00000000..475b6e43 --- /dev/null +++ b/src/haiku/rag/graph/models.py @@ -0,0 +1,24 @@ +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/src/haiku/rag/graph/nodes/__init__.py b/src/haiku/rag/graph/nodes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/haiku/rag/research/nodes/analysis.py b/src/haiku/rag/graph/nodes/analysis.py similarity index 96% rename from src/haiku/rag/research/nodes/analysis.py rename to src/haiku/rag/graph/nodes/analysis.py index 18057089..1d29d935 100644 --- a/src/haiku/rag/research/nodes/analysis.py +++ b/src/haiku/rag/graph/nodes/analysis.py @@ -3,15 +3,13 @@ from dataclasses import dataclass from pydantic_ai import Agent from pydantic_graph import BaseNode, GraphRunContext +from haiku.rag.graph.common import get_model, log from haiku.rag.research.common import ( format_analysis_for_prompt, format_context_for_prompt, - get_model, - log, ) from haiku.rag.research.dependencies import ResearchDependencies from haiku.rag.research.models import EvaluationResult, InsightAnalysis, ResearchReport -from haiku.rag.research.nodes.synthesize import SynthesizeNode from haiku.rag.research.prompts import DECISION_AGENT_PROMPT, INSIGHT_AGENT_PROMPT from haiku.rag.research.state import ResearchDeps, ResearchState @@ -89,6 +87,8 @@ class AnalyzeInsightsNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]) for question in analysis.new_questions: log(deps, state, f" • {question}") + from haiku.rag.graph.nodes.analysis import DecisionNode + return DecisionNode(self.provider, self.model) @@ -169,7 +169,8 @@ class DecisionNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]" log(deps, state, f" Sufficient: {status}") - from haiku.rag.research.nodes.search import SearchDispatchNode + from haiku.rag.graph.nodes.search import SearchDispatchNode + from haiku.rag.graph.nodes.synthesize import SynthesizeNode if ( output.is_sufficient diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/graph/nodes/plan.py similarity index 89% rename from src/haiku/rag/research/nodes/plan.py rename to src/haiku/rag/graph/nodes/plan.py index 63612a55..5f2cade9 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/src/haiku/rag/graph/nodes/plan.py @@ -3,11 +3,11 @@ from dataclasses import dataclass from pydantic_ai import Agent, RunContext from pydantic_graph import BaseNode, GraphRunContext -from haiku.rag.research.common import get_model, log +from haiku.rag.graph.common import get_model, log +from haiku.rag.graph.models import ResearchPlan +from haiku.rag.graph.prompts import PLAN_PROMPT from haiku.rag.research.dependencies import ResearchDependencies -from haiku.rag.research.models import ResearchPlan, ResearchReport -from haiku.rag.research.nodes.search import SearchDispatchNode -from haiku.rag.research.prompts import PLAN_PROMPT +from haiku.rag.research.models import ResearchReport from haiku.rag.research.state import ResearchDeps, ResearchState @@ -67,4 +67,6 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): for i, sq in enumerate(state.context.sub_questions, 1): log(deps, state, f" {i}. {sq}") + from haiku.rag.graph.nodes.search import SearchDispatchNode + return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/graph/nodes/search.py similarity index 92% rename from src/haiku/rag/research/nodes/search.py rename to src/haiku/rag/graph/nodes/search.py index 664a01cc..b7e12bd3 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/graph/nodes/search.py @@ -7,10 +7,11 @@ from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput from pydantic_graph import BaseNode, GraphRunContext -from haiku.rag.research.common import get_model, log +from haiku.rag.graph.common import get_model, log +from haiku.rag.graph.models import SearchAnswer +from haiku.rag.graph.prompts import SEARCH_AGENT_PROMPT from haiku.rag.research.dependencies import ResearchDependencies -from haiku.rag.research.models import ResearchReport, SearchAnswer -from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT +from haiku.rag.research.models import ResearchReport from haiku.rag.research.state import ResearchDeps, ResearchState @@ -25,7 +26,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): state = ctx.state deps = ctx.deps if not state.context.sub_questions: - from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode + from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode return AnalyzeInsightsNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/synthesize.py b/src/haiku/rag/graph/nodes/synthesize.py similarity index 90% rename from src/haiku/rag/research/nodes/synthesize.py rename to src/haiku/rag/graph/nodes/synthesize.py index 4f8eee13..309a933d 100644 --- a/src/haiku/rag/research/nodes/synthesize.py +++ b/src/haiku/rag/graph/nodes/synthesize.py @@ -3,10 +3,9 @@ from dataclasses import dataclass from pydantic_ai import Agent from pydantic_graph import BaseNode, End, GraphRunContext -from haiku.rag.research.common import format_context_for_prompt, get_model, log -from haiku.rag.research.dependencies import ( - ResearchDependencies, -) +from haiku.rag.graph.common import get_model, log +from haiku.rag.research.common import format_context_for_prompt +from haiku.rag.research.dependencies import ResearchDependencies from haiku.rag.research.models import ResearchReport from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT from haiku.rag.research.state import ResearchDeps, ResearchState diff --git a/src/haiku/rag/graph/prompts.py b/src/haiku/rag/graph/prompts.py new file mode 100644 index 00000000..97f42aa8 --- /dev/null +++ b/src/haiku/rag/graph/prompts.py @@ -0,0 +1,45 @@ +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/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index b034748a..e49dd051 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -1,28 +1,3 @@ +from haiku.rag.graph.models import SearchAnswer from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies -from haiku.rag.research.graph import ( - PlanNode, - ResearchDeps, - ResearchState, - build_research_graph, -) -from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer -from haiku.rag.research.stream import ( - ResearchStateSnapshot, - ResearchStreamEvent, - stream_research_graph, -) - -__all__ = [ - "ResearchDependencies", - "ResearchContext", - "SearchAnswer", - "EvaluationResult", - "ResearchReport", - "ResearchDeps", - "ResearchState", - "PlanNode", - "build_research_graph", - "stream_research_graph", - "ResearchStreamEvent", - "ResearchStateSnapshot", -] +from haiku.rag.research.models import EvaluationResult, ResearchReport diff --git a/src/haiku/rag/research/common.py b/src/haiku/rag/research/common.py index fddddb99..bd6e349d 100644 --- a/src/haiku/rag/research/common.py +++ b/src/haiku/rag/research/common.py @@ -1,39 +1,8 @@ -from typing import TYPE_CHECKING, Any - from pydantic_ai import format_as_xml -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_ai.providers.openai import OpenAIProvider -from haiku.rag.config import Config from haiku.rag.research.dependencies import ResearchContext from haiku.rag.research.models import InsightAnalysis -if TYPE_CHECKING: # pragma: no cover - from haiku.rag.research.state import ResearchDeps, ResearchState - - -def get_model(provider: str, model: str) -> Any: - if provider == "ollama": - return OpenAIChatModel( - model_name=model, - provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), - ) - elif provider == "vllm": - return OpenAIChatModel( - model_name=model, - provider=OpenAIProvider( - base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1", - api_key="none", - ), - ) - else: - return f"{provider}:{model}" - - -def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None: - deps.emit_log(msg, state) - def format_context_for_prompt(context: ResearchContext) -> str: """Format the research context as XML for inclusion in prompts.""" diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index 2c698fa9..9606c658 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/src/haiku/rag/research/dependencies.py @@ -4,11 +4,11 @@ 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.research.models import ( GapRecord, InsightAnalysis, InsightRecord, - SearchAnswer, ) from haiku.rag.research.stream import ResearchStream diff --git a/src/haiku/rag/research/graph.py b/src/haiku/rag/research/graph.py index 7f70895f..422ad1fd 100644 --- a/src/haiku/rag/research/graph.py +++ b/src/haiku/rag/research/graph.py @@ -1,23 +1,12 @@ from pydantic_graph import Graph +from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode +from haiku.rag.graph.nodes.plan import PlanNode +from haiku.rag.graph.nodes.search import SearchDispatchNode +from haiku.rag.graph.nodes.synthesize import SynthesizeNode from haiku.rag.research.models import ResearchReport -from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode, DecisionNode -from haiku.rag.research.nodes.plan import PlanNode -from haiku.rag.research.nodes.search import SearchDispatchNode -from haiku.rag.research.nodes.synthesize import SynthesizeNode from haiku.rag.research.state import ResearchDeps, ResearchState -__all__ = [ - "PlanNode", - "SearchDispatchNode", - "AnalyzeInsightsNode", - "DecisionNode", - "SynthesizeNode", - "ResearchState", - "ResearchDeps", - "build_research_graph", -] - def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]: return Graph( diff --git a/src/haiku/rag/research/models.py b/src/haiku/rag/research/models.py index eedcb835..aae2f32a 100644 --- a/src/haiku/rag/research/models.py +++ b/src/haiku/rag/research/models.py @@ -131,31 +131,6 @@ class InsightAnalysis(BaseModel): ) -class ResearchPlan(BaseModel): - main_question: str - sub_questions: list[str] - - -class SearchAnswer(BaseModel): - """Structured output for the SearchSpecialist agent.""" - - 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 EvaluationResult(BaseModel): """Result of analysis and evaluation.""" diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index 4f32e5fe..5b90696e 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -1,49 +1,3 @@ -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.""" - INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the research workflow. diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 3f3aef78..e3fdd153 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -2,15 +2,15 @@ from typing import Any, cast import pytest +from haiku.rag.graph.models import SearchAnswer +from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode +from haiku.rag.graph.nodes.plan import PlanNode +from haiku.rag.graph.nodes.search import SearchDispatchNode +from haiku.rag.graph.nodes.synthesize import SynthesizeNode from haiku.rag.research.dependencies import ResearchContext from haiku.rag.research.graph import ( - AnalyzeInsightsNode, - DecisionNode, - PlanNode, ResearchDeps, ResearchState, - SearchDispatchNode, - SynthesizeNode, build_research_graph, ) from haiku.rag.research.models import ( @@ -21,7 +21,6 @@ from haiku.rag.research.models import ( InsightRecord, InsightStatus, ResearchReport, - SearchAnswer, ) from haiku.rag.research.stream import stream_research_graph