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

Build a "conversational" research graph

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-09 15:51:09 +02:00
parent df71e7a893
commit 26bc02ef14
No known key found for this signature in database
5 changed files with 450 additions and 186 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

158
app/backend/agent.py
View file

@ -27,11 +27,19 @@ class CitationInfo(BaseModel):
class QAResponse(BaseModel):
"""A Q&A pair from conversation history."""
"""A Q&A pair from conversation history with citations."""
question: str
answer: str
sources: list[str] = []
confidence: float = 0.9
citations: list[CitationInfo] = []
@property
def sources(self) -> list[str]:
"""Source names for display."""
return list(
dict.fromkeys(c.document_title or c.document_uri for c in self.citations)
)
class ChatSessionState(BaseModel):
@ -71,6 +79,16 @@ class ChatDeps:
session_state: ChatSessionState | None = None
def build_document_filter(document_name: str) -> str:
"""Build SQL filter for document name matching."""
escaped = document_name.replace("'", "''")
no_spaces = escaped.replace(" ", "")
return (
f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
)
CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system.
You have access to a knowledge base of documents. Use your tools to search and answer questions.
@ -141,15 +159,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
context = format_conversation_context(ctx.deps.session_state.qa_history)
# Build filter from document_name
doc_filter = None
if document_name:
escaped = document_name.replace("'", "''")
# Also try without spaces for matching "TB MED 593" to "tbmed593"
no_spaces = escaped.replace(" ", "")
doc_filter = (
f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
)
doc_filter = build_document_filter(document_name) if document_name else None
# Use search agent for query expansion and deduplication
search_agent = SearchAgent(ctx.deps.client, ctx.deps.config)
@ -220,11 +230,17 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
"""Answer a specific question using the knowledge base.
Use this for direct questions that need a focused answer with citations.
Uses a research graph for planning, searching, and synthesis.
Args:
question: The question to answer
document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual")
"""
from haiku.rag.graph.research.dependencies import ResearchContext
from haiku.rag.graph.research.graph import build_conversational_graph
from haiku.rag.graph.research.models import Citation, SearchAnswer
from haiku.rag.graph.research.state import ResearchDeps, ResearchState
if ctx.deps.agui_emitter:
msg = f"Answering: {question}"
if document_name:
@ -232,64 +248,81 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
ctx.deps.agui_emitter.log(msg)
# Build filter from document_name
doc_filter = None
if document_name:
escaped = document_name.replace("'", "''")
# Also try without spaces for matching "TB MED 593" to "tbmed593"
no_spaces = escaped.replace(" ", "")
doc_filter = (
f"LOWER(uri) LIKE LOWER('%{escaped}%') OR LOWER(title) LIKE LOWER('%{escaped}%') "
f"OR LOWER(uri) LIKE LOWER('%{no_spaces}%') OR LOWER(title) LIKE LOWER('%{no_spaces}%')"
)
doc_filter = build_document_filter(document_name) if document_name else None
# Build context-aware system prompt if we have history
system_prompt = None
# Convert existing qa_history to SearchAnswers for context seeding
existing_qa: list[SearchAnswer] = []
if ctx.deps.session_state and ctx.deps.session_state.qa_history:
from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT
for qa in ctx.deps.session_state.qa_history:
citations = [
Citation(
document_id=c.document_id,
chunk_id=c.chunk_id,
document_uri=c.document_uri,
document_title=c.document_title,
page_numbers=c.page_numbers,
headings=c.headings,
content=c.content,
)
for c in qa.citations
]
existing_qa.append(
SearchAnswer(
query=qa.question,
answer=qa.answer,
confidence=qa.confidence,
cited_chunks=[c.chunk_id for c in qa.citations],
citations=citations,
)
)
context_xml = format_conversation_context(ctx.deps.session_state.qa_history)
system_prompt = (
f"{QA_SYSTEM_PROMPT}\n\n"
f"{context_xml}\n\n"
"Use this conversation context to provide informed answers. "
"Reference previous answers when relevant."
)
# Build and run the conversational research graph
graph = build_conversational_graph(config=ctx.deps.config)
answer, citations = await ctx.deps.client.ask(
question, system_prompt=system_prompt, filter=doc_filter
context = ResearchContext(
original_question=question,
qa_responses=existing_qa,
)
state = ResearchState(
context=context,
max_iterations=1,
confidence_threshold=0.0,
search_filter=doc_filter,
max_concurrency=ctx.deps.config.research.max_concurrency,
)
# Don't pass agui_emitter to research graph - its state model differs from ChatSessionState
# The ask tool handles final state emission with citations
deps = ResearchDeps(
client=ctx.deps.client,
)
# Accumulate Q&A in session state
if ctx.deps.session_state is not None:
sources = (
[c.document_title or c.document_uri for c in citations]
if citations
else []
result = await graph.run(state=state, deps=deps)
# Build citation infos for frontend and history
citation_infos = [
CitationInfo(
index=i + 1,
document_id=c.document_id,
chunk_id=c.chunk_id,
document_uri=c.document_uri,
document_title=c.document_title,
page_numbers=c.page_numbers,
headings=c.headings,
content=c.content,
)
for i, c in enumerate(result.citations)
]
# Accumulate Q&A in session state with full citation metadata
if ctx.deps.session_state is not None:
qa_response = QAResponse(
question=question,
answer=answer,
sources=list(dict.fromkeys(sources)), # dedupe preserving order
answer=result.answer,
confidence=result.confidence,
citations=citation_infos,
)
ctx.deps.session_state.qa_history.append(qa_response)
# Build citation infos for frontend
citation_infos = []
if citations:
citation_infos = [
CitationInfo(
index=i + 1,
document_id=c.document_id,
chunk_id=c.chunk_id,
document_uri=c.document_uri,
document_title=c.document_title,
page_numbers=c.page_numbers,
headings=c.headings,
content=c.content,
)
for i, c in enumerate(citations)
]
# Emit updated state with citations AND accumulated qa_history
if ctx.deps.agui_emitter:
ctx.deps.agui_emitter.update_state(
@ -308,12 +341,13 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
)
)
# Format answer with citation references
if citations:
citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citations)))
return f"{answer}\n\nSources: {citation_refs}"
# Format answer with citation references and confidence
answer_text = result.answer
if citation_infos:
citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citation_infos)))
answer_text = f"{answer_text}\n\nSources: {citation_refs}"
return answer
return answer_text
@agent.tool
async def get_document(

3
app/frontend/components/Chat.tsx
View file

@ -24,7 +24,8 @@ interface Citation {
interface QAResponse {
question: string;
answer: string;
sources: string[];
confidence: number;
citations: Citation[];
}
interface ChatSessionState {

425
haiku_rag_slim/haiku/rag/graph/research/graph.py
View file

@ -18,6 +18,8 @@ from haiku.rag.graph.agui.emitter import (
)
from haiku.rag.graph.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.graph.research.models import (
Citation,
ConversationalAnswer,
EvaluationResult,
RawSearchAnswer,
ResearchPlan,
@ -25,8 +27,10 @@ from haiku.rag.graph.research.models import (
SearchAnswer,
)
from haiku.rag.graph.research.prompts import (
CONVERSATIONAL_SYNTHESIS_PROMPT,
DECISION_PROMPT,
PLAN_PROMPT,
PLAN_PROMPT_WITH_CONTEXT,
SEARCH_PROMPT,
SYNTHESIS_PROMPT,
)
@ -60,6 +64,135 @@ def format_context_for_prompt(context: ResearchContext) -> str:
return format_as_xml(context_data, root_tag="research_context")
# =============================================================================
# Shared step logic helpers
# =============================================================================
async def _plan_step_logic(
state: ResearchState,
deps: ResearchDeps,
config: AppConfig,
plan_prompt: str,
) -> None:
"""Shared logic for the plan step."""
model_config = config.research.model
# Use context-aware prompt if we have existing qa_responses
has_context = bool(state.context.qa_responses)
effective_plan_prompt = (
build_prompt(PLAN_PROMPT_WITH_CONTEXT, config) if has_context else plan_prompt
)
plan_agent = Agent(
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=effective_plan_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
@plan_agent.tool
async def gather_context(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
return "\n\n".join(r.content for r in results)
# Build prompt with existing context if available
if has_context:
context_xml = format_context_for_prompt(state.context)
prompt = (
f"Review existing context and plan additional research if needed.\n\n"
f"{context_xml}\n\n"
f"Main question: {state.context.original_question}"
)
else:
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
async def _search_one_step_logic(
state: ResearchState,
deps: ResearchDeps,
config: AppConfig,
search_prompt: str,
sub_q: str,
) -> SearchAnswer:
"""Shared logic for the search_one step."""
model_config = config.research.model
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
agent = Agent(
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=search_prompt,
retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
"""Search the knowledge base for relevant documents."""
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
ctx2.deps.search_results = results
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found for: {query}"
return "\n\n".join(parts)
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
if raw_answer:
answer = SearchAnswer.from_raw(raw_answer, agent_deps.search_results)
state.context.add_qa_response(answer)
return answer
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
def _get_batch_logic(state: ResearchState) -> list[str] | None:
"""Shared logic for the get_batch step."""
if not state.context.sub_questions:
return None
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
# =============================================================================
# Research graph (full version with decide loop)
# =============================================================================
def build_research_graph(
config: AppConfig = Config,
include_plan: bool = True,
@ -107,39 +240,7 @@ def build_research_graph(
)
try:
plan_agent = Agent(
model=get_model(model_config, config),
output_type=ResearchPlan,
instructions=plan_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
search_filter = state.search_filter
@plan_agent.tool
async def gather_context(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
return "\n\n".join(r.content for r in results)
_ = gather_context
prompt = (
"Plan a focused approach for the main question.\n\n"
f"Main question: {state.context.original_question}"
)
agent_deps = ResearchDependencies(client=deps.client, context=state.context)
plan_result = await plan_agent.run(prompt, deps=agent_deps)
state.context.sub_questions = list(plan_result.output.sub_questions)
await _plan_step_logic(state, deps, config, plan_prompt)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
@ -168,92 +269,47 @@ def build_research_graph(
if deps.agui_emitter:
deps.agui_emitter.start_step(step_name)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
try:
if deps.semaphore is None:
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
async with deps.semaphore:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Searching: {sub_q}",
"query": sub_q,
},
)
agent = Agent(
model=get_model(model_config, config),
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
instructions=search_prompt,
retries=3,
deps_type=ResearchDependencies,
answer = await _search_one_step_logic(
state, deps, config, search_prompt, sub_q
)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Found answer with {answer.confidence:.0%} confidence",
"query": sub_q,
"confidence": answer.confidence,
},
)
search_filter = state.search_filter
@agent.tool
async def search_and_answer(
ctx2: RunContext[ResearchDependencies],
query: str,
limit: int | None = None,
) -> str:
"""Search the knowledge base for relevant documents."""
results = await ctx2.deps.client.search(
query, limit=limit, filter=search_filter
)
results = await ctx2.deps.client.expand_context(results)
ctx2.deps.search_results = results
parts = [r.format_for_agent() for r in results]
if not parts:
return f"No relevant information found for: {query}"
return "\n\n".join(parts)
_ = search_and_answer
agent_deps = ResearchDependencies(
client=deps.client, context=state.context
return answer
except Exception as e:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
try:
result = await agent.run(sub_q, deps=agent_deps)
raw_answer = result.output
if raw_answer:
answer = SearchAnswer.from_raw(
raw_answer, agent_deps.search_results
)
state.context.add_qa_response(answer)
if deps.agui_emitter:
deps.agui_emitter.update_state(state)
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Found answer with {answer.confidence:.0%} confidence",
"query": sub_q,
"confidence": answer.confidence,
},
)
return answer
return SearchAnswer(query=sub_q, answer="", confidence=0.0)
except Exception as e:
if deps.agui_emitter:
deps.agui_emitter.update_activity(
"searching",
{
"stepName": "search_one",
"message": f"Search failed: {e}",
"query": sub_q,
"error": str(e),
},
)
return SearchAnswer(
query=sub_q,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
return SearchAnswer(
query=sub_q,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
finally:
if deps.agui_emitter:
deps.agui_emitter.finish_step(step_name)
@ -263,14 +319,7 @@ def build_research_graph(
ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
state = ctx.state
if not state.context.sub_questions:
return None
batch = list(state.context.sub_questions)
state.context.sub_questions.clear()
return batch
return _get_batch_logic(ctx.state)
@g.step
async def decide(
@ -548,3 +597,133 @@ def build_research_graph(
)
return g.build()
# =============================================================================
# Conversational graph (simplified, single iteration)
# =============================================================================
def build_conversational_graph(
config: AppConfig = Config,
) -> Graph[ResearchState, ResearchDeps, None, ConversationalAnswer]:
"""Build a simplified research graph for conversational chat.
This graph is optimized for single-iteration Q&A:
- Context-aware planning (generates fewer sub-questions when context exists)
- Single search iteration (no decide loop)
- Conversational output (direct answer, not formal report)
Args:
config: AppConfig object
Returns:
Graph that outputs ConversationalAnswer
"""
# Build prompts
plan_prompt = build_prompt(
PLAN_PROMPT
+ "\n\nUse the gather_context tool once on the main question before planning.",
config,
)
search_prompt = build_prompt(SEARCH_PROMPT, config)
conversational_prompt = build_prompt(CONVERSATIONAL_SYNTHESIS_PROMPT, config)
g = GraphBuilder(
state_type=ResearchState,
deps_type=ResearchDeps,
output_type=ConversationalAnswer,
)
@g.step
async def plan(ctx: StepContext[ResearchState, ResearchDeps, None]) -> None:
"""Create research plan with sub-questions."""
await _plan_step_logic(ctx.state, ctx.deps, config, plan_prompt)
@g.step
async def search_one(
ctx: StepContext[ResearchState, ResearchDeps, str],
) -> SearchAnswer:
"""Answer a single sub-question using the knowledge base."""
try:
return await _search_one_step_logic(
ctx.state, ctx.deps, config, search_prompt, ctx.inputs
)
except Exception as e:
return SearchAnswer(
query=ctx.inputs,
answer=f"Search failed: {str(e)}",
confidence=0.0,
)
@g.step
async def get_batch(
ctx: StepContext[ResearchState, ResearchDeps, None],
) -> list[str] | None:
"""Get all remaining questions for this iteration."""
return _get_batch_logic(ctx.state)
@g.step
async def synthesize(
ctx: StepContext[ResearchState, ResearchDeps, list[SearchAnswer] | None],
) -> ConversationalAnswer:
"""Generate conversational answer from gathered evidence."""
state = ctx.state
deps = ctx.deps
agent = Agent(
model=get_model(config.research.model, config),
output_type=ConversationalAnswer,
instructions=conversational_prompt,
retries=3,
output_retries=3,
deps_type=ResearchDependencies,
)
context_xml = format_context_for_prompt(state.context)
prompt = (
f"Answer the following question based on the gathered evidence.\n\n"
f"{context_xml}\n\n"
f"Question: {state.context.original_question}"
)
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,
)
# Build the graph structure (simplified: plan → search → synthesize)
collect_answers = g.join(
reduce_list_append,
initial_factory=list[SearchAnswer],
)
g.add(
g.edge_from(g.start_node).to(plan),
g.edge_from(plan).to(get_batch),
g.edge_from(get_batch).to(
g.decision()
.branch(g.match(list).label("Has questions").map().to(search_one))
.branch(g.match(type(None)).label("No questions").to(synthesize))
),
g.edge_from(search_one).to(collect_answers),
g.edge_from(collect_answers).to(synthesize),
g.edge_from(synthesize).to(g.end_node),
)
return g.build()

12
haiku_rag_slim/haiku/rag/graph/research/models.py
View file

@ -125,6 +125,18 @@ class EvaluationResult(BaseModel):
)
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."""

38
haiku_rag_slim/haiku/rag/graph/research/prompts.py
View file

@ -20,6 +20,27 @@ Plan requirements:
Use the gather_context tool once on the main question before planning."""
PLAN_PROMPT_WITH_CONTEXT = """You are the research orchestrator for a focused workflow.
You have access to PREVIOUS CONVERSATION CONTEXT in the qa_responses section below.
Review this context first - if it already answers the question, generate minimal
or no sub-questions. Only create sub-questions to fill gaps in the existing context.
Responsibilities:
1. Review existing qa_responses to understand what's already known
2. Identify gaps that need additional research
3. Propose minimal sub-questions only for missing information
Plan requirements:
- If existing context fully answers the question, return a SINGLE sub-question
to verify or slightly expand the answer.
- Only create new sub-questions for genuine gaps in the existing knowledge.
- sub_questions must be a list of plain strings (max 3).
- Each sub_question must be standalone and self-contained.
- Prioritize the highest-value gaps first.
Use the gather_context tool once on the main question before planning."""
SEARCH_PROMPT = """You are a search and question-answering specialist.
Process:
@ -113,3 +134,20 @@ 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 collected evidence in qa_responses.
- 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."""