Remove unused conversational output mode from research graph

This commit is contained in:
Yiorgis Gozadinos 2026-02-20 14:37:06 +02:00
parent e6310fc484
commit eb9436eb2a
No known key found for this signature in database
11 changed files with 53 additions and 304 deletions

View file

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

View file

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

View file

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

View file

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

View file

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

View file

@ -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 <background> 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."""

View file

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

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

View file

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

View file

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

View file

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