diff --git a/app/backend/main.py b/app/backend/main.py index f554251a..50876a8c 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -16,7 +16,7 @@ from starlette.routing import Route from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config from haiku.rag.config.models import AppConfig -from haiku.rag.skills.rag import create_skill +from haiku.rag.skills.rag import AGENT_PREAMBLE, create_skill from haiku.skills import SkillDeps, SkillToolset load_dotenv(find_dotenv(usecwd=True)) @@ -66,16 +66,6 @@ def get_client() -> HaikuRAG: skill = create_skill(db_path=db_path, config=Config) toolset = SkillToolset(skills=[skill]) -AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system. - -CRITICAL RULES: -1. For greetings or casual chat: respond directly WITHOUT using any tools -2. NEVER make up information - always use tools to get facts from the knowledge base -3. For questions: Use the "ask" tool - it handles search and citation automatically -4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets -5. When you use the "ask" tool, summarize the key findings and always include citations in your response -""" - agent = Agent( os.getenv("HAIKU_CHAT_MODEL", "openai:gpt-4o"), instructions=AGENT_PREAMBLE + toolset.system_prompt, diff --git a/docs/agents/index.md b/docs/agents/index.md index a733ee8c..cbc81edb 100644 --- a/docs/agents/index.md +++ b/docs/agents/index.md @@ -84,15 +84,6 @@ When prior answers are provided, the planner uses a context-aware prompt that ev - **search_one**: Answers a single question using the knowledge base (up to 3 search calls per question). Each answer is added to `ResearchContext.qa_responses` for the next planning iteration. - **synthesize**: Generates the final output from all gathered evidence. -**Output modes:** - -The graph supports two output modes via `build_research_graph(output_mode=...)`: - -| Mode | Output type | Used by | -|------|-------------|---------| -| `"report"` | `ResearchReport` (title, executive summary, findings, conclusions, recommendations) | CLI `haiku-rag research`, Python API | -| `"conversational"` | `ConversationalAnswer` (answer, citations, confidence) | Chat agent's `ask` tool | - **Iterative flow:** - Each iteration: planner evaluates context → proposes one question → search answers it → loop back @@ -160,40 +151,6 @@ async with HaikuRAG(path_to_db) as client: report = await graph.run(state=state, deps=deps) ``` -**Conversational mode with prior answers:** - -```python -from haiku.rag.config import Config -from haiku.rag.agents.research.dependencies import ResearchContext -from haiku.rag.agents.research.graph import build_research_graph -from haiku.rag.agents.research.models import SearchAnswer -from haiku.rag.agents.research.state import ResearchDeps, ResearchState - -# Conversational mode returns ConversationalAnswer instead of ResearchReport -graph = build_research_graph(config=Config, output_mode="conversational") - -# Pass session context and prior answers from conversation history -context = ResearchContext( - original_question="How does it handle authentication?", - session_context="User is building a Python web app with FastAPI.", - qa_responses=[ - SearchAnswer( - query="What authentication methods are supported?", - answer="JWT and OAuth2 are supported.", - confidence=0.95, - cited_chunks=["chunk-1"], - ) - ], -) -state = ResearchState.from_config(context=context, config=Config) -deps = ResearchDeps(client=client) - -result = await graph.run(state=state, deps=deps) -print(result.answer) # Direct conversational answer -print(result.confidence) # 0.0-1.0 -print(result.citations) # Deduplicated citations from all searches -``` - ### Filtering Documents Restrict searches to specific documents via the `search_filter` parameter: diff --git a/docs/architecture.md b/docs/architecture.md index 00417156..aa3b1f1d 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -150,7 +150,7 @@ flowchart TB - Proposes one question at a time, evaluates the answer, then decides whether to continue - Session context resolves ambiguous references - Prior answers let the planner skip redundant searches -- Synthesizes structured report or conversational answer +- Synthesizes structured report **RLM Agent** - Complex analytical tasks via code execution: diff --git a/haiku_rag_slim/haiku/rag/agents/research/graph.py b/haiku_rag_slim/haiku/rag/agents/research/graph.py index 18b95a9f..3734dfde 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/graph.py +++ b/haiku_rag_slim/haiku/rag/agents/research/graph.py @@ -1,5 +1,4 @@ import asyncio -from typing import Literal, overload from pydantic_ai import Agent, RunContext, format_as_xml from pydantic_ai.output import ToolOutput @@ -7,15 +6,12 @@ from pydantic_graph.beta import Graph, GraphBuilder, StepContext from haiku.rag.agents.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.agents.research.models import ( - Citation, - ConversationalAnswer, IterativePlanResult, RawSearchAnswer, ResearchReport, SearchAnswer, ) from haiku.rag.agents.research.prompts import ( - CONVERSATIONAL_SYNTHESIS_PROMPT, ITERATIVE_PLAN_PROMPT, ITERATIVE_PLAN_PROMPT_WITH_CONTEXT, SEARCH_PROMPT, @@ -169,29 +165,13 @@ async def _search_one_step_logic( return SearchAnswer(query=sub_q, answer="", confidence=0.0) -@overload -def build_research_graph( - config: AppConfig = ..., - output_mode: Literal["report"] = ..., -) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: ... - - -@overload -def build_research_graph( - config: AppConfig = ..., - output_mode: Literal["conversational"] = ..., -) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]: ... - - def build_research_graph( config: AppConfig = Config, - output_mode: Literal["report", "conversational"] = "report", -) -> Graph[ResearchState, ResearchDeps, None, ResearchReport | ConversationalAnswer]: +) -> Graph[ResearchState, ResearchDeps, None, ResearchReport]: """Build the iterative research graph. Args: config: AppConfig object (uses config.research for provider, model, and graph parameters) - output_mode: Output format - "report" for ResearchReport, "conversational" for ConversationalAnswer Returns: Configured research graph with iterative planning @@ -199,18 +179,14 @@ def build_research_graph( model_config = config.research.model search_prompt = build_prompt(SEARCH_PROMPT, config) - - if output_mode == "report": - synthesis_prompt = build_prompt( - config.prompts.synthesis or SYNTHESIS_PROMPT, config - ) - else: - synthesis_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config) + synthesis_prompt = build_prompt( + config.prompts.synthesis or SYNTHESIS_PROMPT, config + ) g = GraphBuilder( state_type=ResearchState, deps_type=ResearchDeps, - output_type=ResearchReport if output_mode == "report" else ConversationalAnswer, + output_type=ResearchReport, ) @g.step @@ -236,81 +212,35 @@ def build_research_graph( confidence=0.0, ) - if output_mode == "report": + @g.step + async def synthesize( + ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult], + ) -> ResearchReport: + """Generate final research report.""" + state = ctx.state + deps = ctx.deps - @g.step - async def synthesize( - ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult], - ) -> ResearchReport: - """Generate final research report.""" - state = ctx.state - deps = ctx.deps + agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment] + model=get_model(model_config, config), + output_type=ResearchReport, + instructions=synthesis_prompt, + retries=3, + output_retries=3, + deps_type=ResearchDependencies, + ) - agent: Agent[ResearchDependencies, ResearchReport] = Agent( # type: ignore[assignment] - model=get_model(model_config, config), - output_type=ResearchReport, - instructions=synthesis_prompt, - retries=3, - output_retries=3, - deps_type=ResearchDependencies, - ) - - context_xml = format_context_for_prompt(state.context) - prompt = ( - "Generate a comprehensive research report based on all gathered information.\n\n" - f"{context_xml}\n\n" - "Create a detailed report that synthesizes all findings into a coherent response." - ) - agent_deps = ResearchDependencies( - client=deps.client, - context=state.context, - ) - result = await agent.run(prompt, deps=agent_deps) - return result.output - - else: - - @g.step - async def synthesize( - ctx: StepContext[ResearchState, ResearchDeps, IterativePlanResult], - ) -> ConversationalAnswer: - """Generate conversational answer from gathered evidence.""" - state = ctx.state - deps = ctx.deps - - agent: Agent[ResearchDependencies, ConversationalAnswer] = Agent( # type: ignore[assignment] - model=get_model(model_config, config), - output_type=ConversationalAnswer, - instructions=synthesis_prompt, - retries=3, - output_retries=3, - deps_type=ResearchDependencies, - ) - - context_xml = format_context_for_prompt(state.context) - prompt = ( - f"Answer the question based on the gathered evidence.\n\n{context_xml}" - ) - agent_deps = ResearchDependencies( - client=deps.client, - context=state.context, - ) - result = await agent.run(prompt, deps=agent_deps) - - # Collect unique citations from qa_responses (dedupe by chunk_id) - seen_chunks: set[str] = set() - unique_citations: list[Citation] = [] - for qa in state.context.qa_responses: - for c in qa.citations: - if c.chunk_id not in seen_chunks: - seen_chunks.add(c.chunk_id) - unique_citations.append(c) - - return ConversationalAnswer( - answer=result.output.answer, - citations=unique_citations, - confidence=result.output.confidence, - ) + context_xml = format_context_for_prompt(state.context) + prompt = ( + "Generate a comprehensive research report based on all gathered information.\n\n" + f"{context_xml}\n\n" + "Create a detailed report that synthesizes all findings into a coherent response." + ) + agent_deps = ResearchDependencies( + client=deps.client, + context=state.context, + ) + result = await agent.run(prompt, deps=agent_deps) + return result.output # Build graph edges: iterative loop # diff --git a/haiku_rag_slim/haiku/rag/agents/research/models.py b/haiku_rag_slim/haiku/rag/agents/research/models.py index bb9e0d93..0c845e07 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/models.py +++ b/haiku_rag_slim/haiku/rag/agents/research/models.py @@ -111,18 +111,6 @@ def resolve_citations( return citations -class ConversationalAnswer(BaseModel): - """Conversational answer for chat context.""" - - answer: str = Field(description="Direct answer to the question") - citations: list[Citation] = Field( - default_factory=list, description="Citations supporting the answer" - ) - confidence: float = Field( - default=1.0, description="Confidence score (0-1)", ge=0.0, le=1.0 - ) - - class ResearchReport(BaseModel): """Final research report structure.""" diff --git a/haiku_rag_slim/haiku/rag/agents/research/prompts.py b/haiku_rag_slim/haiku/rag/agents/research/prompts.py index bb21bc55..fae407e2 100644 --- a/haiku_rag_slim/haiku/rag/agents/research/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/research/prompts.py @@ -115,21 +115,3 @@ Style: - Be professional, objective, and specific. - NEVER use meta-commentary like "This report covers..." or "The findings show...". Instead, state the actual information directly.""" - -CONVERSATIONAL_SYNTHESIS_PROMPT = """Generate a direct, conversational answer -to the question based on the gathered evidence. - -Output: -- answer: Direct, comprehensive answer with a natural, helpful tone. - Write the actual answer, not a description of what you found. - Use as many sentences as needed to fully address the question. -- confidence: Score from 0.0 to 1.0 indicating answer quality. - -Guidelines: -- Base your answer solely on the evidence provided in the context. -- If a section is provided, use it to frame your answer appropriately. -- Be thorough - include all relevant information from the evidence. -- Use formatting (bullet points, numbered lists) when it improves clarity. -- Do NOT use meta-commentary like "Based on the research..." or "The evidence shows..." - Instead, directly state the information. -- If the evidence is incomplete, acknowledge limitations briefly.""" diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index a5af5f7e..969b33b4 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, Any from haiku.rag.client import HaikuRAG from haiku.rag.config import get_config -from haiku.rag.skills.rag import RAGState +from haiku.rag.skills.rag import AGENT_PREAMBLE, RAGState from haiku.skills.agent import SkillToolset from haiku.skills.models import Skill @@ -55,16 +55,6 @@ except ImportError: # pragma: no cover RAG_STATE_NAMESPACE = "rag" -AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system. - -CRITICAL RULES: -1. For greetings or casual chat: respond directly WITHOUT using any tools -2. NEVER make up information - always use tools to get facts from the knowledge base -3. For questions: Use the "ask" tool - it handles search and citation automatically -4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets -5. When you use the "ask" tool, summarize the key findings and always include citations in your response -""" - class ChatApp(App): """Textual TUI for conversational RAG.""" diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 73d98d97..b18f30d0 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, Literal, overload +from typing import TYPE_CHECKING, overload from urllib.parse import urlparse import httpx @@ -33,7 +33,6 @@ if TYPE_CHECKING: from haiku.rag.agents.research.models import ( Citation, - ConversationalAnswer, ResearchReport, ) from haiku.rag.agents.rlm.models import RLMResult @@ -1328,50 +1327,28 @@ 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": + ) -> "ResearchReport": """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. + ResearchReport with structured findings. """ 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) + graph = build_research_graph(config=self._config) context = ResearchContext(original_question=question) state = ResearchState.from_config( context=context, config=self._config, max_iterations=max_iterations diff --git a/haiku_rag_slim/haiku/rag/skills/rag.py b/haiku_rag_slim/haiku/rag/skills/rag.py index 099a2d8e..ffed170f 100644 --- a/haiku_rag_slim/haiku/rag/skills/rag.py +++ b/haiku_rag_slim/haiku/rag/skills/rag.py @@ -13,6 +13,16 @@ from haiku.skills.models import Skill, SkillSource from haiku.skills.parser import parse_skill_md from haiku.skills.state import SkillRunDeps +AGENT_PREAMBLE = """You are a helpful research assistant powered by haiku.rag, a knowledge base system. + +CRITICAL RULES: +1. For greetings or casual chat: respond directly WITHOUT using any tools +2. NEVER make up information - always use tools to get facts from the knowledge base +3. For questions: Use the "ask" tool - it handles search and citation automatically +4. For searches: Use the "search" tool - copy the ENTIRE tool response to your output INCLUDING content snippets +5. When you use the "ask" tool, summarize the key findings and always include citations in your response +""" + class ResearchEntry(BaseModel): question: str diff --git a/tests/agents/research/test_research_graph.py b/tests/agents/research/test_research_graph.py index c08e9693..dc7750ba 100644 --- a/tests/agents/research/test_research_graph.py +++ b/tests/agents/research/test_research_graph.py @@ -68,64 +68,15 @@ def test_iterative_plan_result_model(): assert continue_result.next_question == "What are the specific requirements?" -# ============================================================================= -# Conversational Graph Tests -# ============================================================================= - - -def test_build_research_graph_conversational_mode_returns_graph(): - """Test build_research_graph with output_mode='conversational' returns a valid Graph instance.""" +def test_build_research_graph_returns_graph(): + """Test build_research_graph returns a valid Graph instance.""" from pydantic_graph.beta import Graph - graph = build_research_graph(output_mode="conversational") + graph = build_research_graph() assert graph is not None assert isinstance(graph, Graph) -def test_build_research_graph_report_mode_returns_graph(): - """Test build_research_graph with output_mode='report' returns a valid Graph instance.""" - from pydantic_graph.beta import Graph - - graph = build_research_graph(output_mode="report") - assert graph is not None - assert isinstance(graph, Graph) - - -def test_conversational_answer_model(): - """Test ConversationalAnswer model can be created with all fields.""" - from haiku.rag.agents.research.models import Citation, ConversationalAnswer - - citation = Citation( - index=1, - document_id="doc-1", - chunk_id="chunk-1", - document_uri="test.md", - document_title="Test Doc", - content="Test content", - ) - - answer = ConversationalAnswer( - answer="The answer is 42.", - citations=[citation], - confidence=0.95, - ) - - assert answer.answer == "The answer is 42." - assert len(answer.citations) == 1 - assert answer.confidence == 0.95 - - -def test_conversational_answer_default_values(): - """Test ConversationalAnswer uses correct default values.""" - from haiku.rag.agents.research.models import ConversationalAnswer - - answer = ConversationalAnswer(answer="Just the answer.") - - assert answer.answer == "Just the answer." - assert answer.citations == [] - assert answer.confidence == 1.0 - - def test_format_context_for_prompt_basic(): """Test format_context_for_prompt with basic context.""" from haiku.rag.agents.research.dependencies import ResearchContext diff --git a/tests/test_client_research.py b/tests/test_client_research.py index f798afa4..95656d18 100644 --- a/tests/test_client_research.py +++ b/tests/test_client_research.py @@ -3,7 +3,7 @@ from unittest.mock import AsyncMock, patch import pytest -from haiku.rag.agents.research.models import ConversationalAnswer, ResearchReport +from haiku.rag.agents.research.models import ResearchReport from haiku.rag.client import HaikuRAG @@ -32,9 +32,6 @@ async def test_client_research_report(temp_db_path): 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() @@ -43,29 +40,6 @@ async def test_client_research_report(temp_db_path): 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(