Build a "conversational" research graph
This commit is contained in:
parent
df71e7a893
commit
26bc02ef14
5 changed files with 450 additions and 186 deletions
|
|
@ -27,11 +27,19 @@ class CitationInfo(BaseModel):
|
||||||
|
|
||||||
|
|
||||||
class QAResponse(BaseModel):
|
class QAResponse(BaseModel):
|
||||||
"""A Q&A pair from conversation history."""
|
"""A Q&A pair from conversation history with citations."""
|
||||||
|
|
||||||
question: str
|
question: str
|
||||||
answer: 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):
|
class ChatSessionState(BaseModel):
|
||||||
|
|
@ -71,6 +79,16 @@ class ChatDeps:
|
||||||
session_state: ChatSessionState | None = None
|
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.
|
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.
|
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)
|
context = format_conversation_context(ctx.deps.session_state.qa_history)
|
||||||
|
|
||||||
# Build filter from document_name
|
# Build filter from document_name
|
||||||
doc_filter = None
|
doc_filter = build_document_filter(document_name) if document_name else 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}%')"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Use search agent for query expansion and deduplication
|
# Use search agent for query expansion and deduplication
|
||||||
search_agent = SearchAgent(ctx.deps.client, ctx.deps.config)
|
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.
|
"""Answer a specific question using the knowledge base.
|
||||||
|
|
||||||
Use this for direct questions that need a focused answer with citations.
|
Use this for direct questions that need a focused answer with citations.
|
||||||
|
Uses a research graph for planning, searching, and synthesis.
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
question: The question to answer
|
question: The question to answer
|
||||||
document_name: Optional document name/title to search within (e.g., "tbmed593", "army manual")
|
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:
|
if ctx.deps.agui_emitter:
|
||||||
msg = f"Answering: {question}"
|
msg = f"Answering: {question}"
|
||||||
if document_name:
|
if document_name:
|
||||||
|
|
@ -232,64 +248,81 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
ctx.deps.agui_emitter.log(msg)
|
ctx.deps.agui_emitter.log(msg)
|
||||||
|
|
||||||
# Build filter from document_name
|
# Build filter from document_name
|
||||||
doc_filter = None
|
doc_filter = build_document_filter(document_name) if document_name else 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}%')"
|
|
||||||
)
|
|
||||||
|
|
||||||
# Build context-aware system prompt if we have history
|
# Convert existing qa_history to SearchAnswers for context seeding
|
||||||
system_prompt = None
|
existing_qa: list[SearchAnswer] = []
|
||||||
if ctx.deps.session_state and ctx.deps.session_state.qa_history:
|
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)
|
# Build and run the conversational research graph
|
||||||
system_prompt = (
|
graph = build_conversational_graph(config=ctx.deps.config)
|
||||||
f"{QA_SYSTEM_PROMPT}\n\n"
|
|
||||||
f"{context_xml}\n\n"
|
|
||||||
"Use this conversation context to provide informed answers. "
|
|
||||||
"Reference previous answers when relevant."
|
|
||||||
)
|
|
||||||
|
|
||||||
answer, citations = await ctx.deps.client.ask(
|
context = ResearchContext(
|
||||||
question, system_prompt=system_prompt, filter=doc_filter
|
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
|
result = await graph.run(state=state, deps=deps)
|
||||||
if ctx.deps.session_state is not None:
|
|
||||||
sources = (
|
# Build citation infos for frontend and history
|
||||||
[c.document_title or c.document_uri for c in citations]
|
citation_infos = [
|
||||||
if citations
|
CitationInfo(
|
||||||
else []
|
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(
|
qa_response = QAResponse(
|
||||||
question=question,
|
question=question,
|
||||||
answer=answer,
|
answer=result.answer,
|
||||||
sources=list(dict.fromkeys(sources)), # dedupe preserving order
|
confidence=result.confidence,
|
||||||
|
citations=citation_infos,
|
||||||
)
|
)
|
||||||
ctx.deps.session_state.qa_history.append(qa_response)
|
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
|
# Emit updated state with citations AND accumulated qa_history
|
||||||
if ctx.deps.agui_emitter:
|
if ctx.deps.agui_emitter:
|
||||||
ctx.deps.agui_emitter.update_state(
|
ctx.deps.agui_emitter.update_state(
|
||||||
|
|
@ -308,12 +341,13 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
# Format answer with citation references
|
# Format answer with citation references and confidence
|
||||||
if citations:
|
answer_text = result.answer
|
||||||
citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citations)))
|
if citation_infos:
|
||||||
return f"{answer}\n\nSources: {citation_refs}"
|
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
|
@agent.tool
|
||||||
async def get_document(
|
async def get_document(
|
||||||
|
|
|
||||||
|
|
@ -24,7 +24,8 @@ interface Citation {
|
||||||
interface QAResponse {
|
interface QAResponse {
|
||||||
question: string;
|
question: string;
|
||||||
answer: string;
|
answer: string;
|
||||||
sources: string[];
|
confidence: number;
|
||||||
|
citations: Citation[];
|
||||||
}
|
}
|
||||||
|
|
||||||
interface ChatSessionState {
|
interface ChatSessionState {
|
||||||
|
|
|
||||||
|
|
@ -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.dependencies import ResearchContext, ResearchDependencies
|
||||||
from haiku.rag.graph.research.models import (
|
from haiku.rag.graph.research.models import (
|
||||||
|
Citation,
|
||||||
|
ConversationalAnswer,
|
||||||
EvaluationResult,
|
EvaluationResult,
|
||||||
RawSearchAnswer,
|
RawSearchAnswer,
|
||||||
ResearchPlan,
|
ResearchPlan,
|
||||||
|
|
@ -25,8 +27,10 @@ from haiku.rag.graph.research.models import (
|
||||||
SearchAnswer,
|
SearchAnswer,
|
||||||
)
|
)
|
||||||
from haiku.rag.graph.research.prompts import (
|
from haiku.rag.graph.research.prompts import (
|
||||||
|
CONVERSATIONAL_SYNTHESIS_PROMPT,
|
||||||
DECISION_PROMPT,
|
DECISION_PROMPT,
|
||||||
PLAN_PROMPT,
|
PLAN_PROMPT,
|
||||||
|
PLAN_PROMPT_WITH_CONTEXT,
|
||||||
SEARCH_PROMPT,
|
SEARCH_PROMPT,
|
||||||
SYNTHESIS_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")
|
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(
|
def build_research_graph(
|
||||||
config: AppConfig = Config,
|
config: AppConfig = Config,
|
||||||
include_plan: bool = True,
|
include_plan: bool = True,
|
||||||
|
|
@ -107,39 +240,7 @@ def build_research_graph(
|
||||||
)
|
)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
plan_agent = Agent(
|
await _plan_step_logic(state, deps, config, plan_prompt)
|
||||||
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)
|
|
||||||
|
|
||||||
if deps.agui_emitter:
|
if deps.agui_emitter:
|
||||||
deps.agui_emitter.update_state(state)
|
deps.agui_emitter.update_state(state)
|
||||||
|
|
@ -168,92 +269,47 @@ def build_research_graph(
|
||||||
|
|
||||||
if deps.agui_emitter:
|
if deps.agui_emitter:
|
||||||
deps.agui_emitter.start_step(step_name)
|
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:
|
try:
|
||||||
if deps.semaphore is None:
|
answer = await _search_one_step_logic(
|
||||||
deps.semaphore = asyncio.Semaphore(state.max_concurrency)
|
state, deps, config, search_prompt, sub_q
|
||||||
|
)
|
||||||
async with deps.semaphore:
|
if deps.agui_emitter:
|
||||||
if deps.agui_emitter:
|
deps.agui_emitter.update_state(state)
|
||||||
deps.agui_emitter.update_activity(
|
deps.agui_emitter.update_activity(
|
||||||
"searching",
|
"searching",
|
||||||
{
|
{
|
||||||
"stepName": "search_one",
|
"stepName": "search_one",
|
||||||
"message": f"Searching: {sub_q}",
|
"message": f"Found answer with {answer.confidence:.0%} confidence",
|
||||||
"query": sub_q,
|
"query": sub_q,
|
||||||
},
|
"confidence": answer.confidence,
|
||||||
)
|
},
|
||||||
|
|
||||||
agent = Agent(
|
|
||||||
model=get_model(model_config, config),
|
|
||||||
output_type=ToolOutput(RawSearchAnswer, max_retries=3),
|
|
||||||
instructions=search_prompt,
|
|
||||||
retries=3,
|
|
||||||
deps_type=ResearchDependencies,
|
|
||||||
)
|
)
|
||||||
|
return answer
|
||||||
search_filter = state.search_filter
|
except Exception as e:
|
||||||
|
if deps.agui_emitter:
|
||||||
@agent.tool
|
deps.agui_emitter.update_activity(
|
||||||
async def search_and_answer(
|
"searching",
|
||||||
ctx2: RunContext[ResearchDependencies],
|
{
|
||||||
query: str,
|
"stepName": "search_one",
|
||||||
limit: int | None = None,
|
"message": f"Search failed: {e}",
|
||||||
) -> str:
|
"query": sub_q,
|
||||||
"""Search the knowledge base for relevant documents."""
|
"error": str(e),
|
||||||
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 SearchAnswer(
|
||||||
try:
|
query=sub_q,
|
||||||
result = await agent.run(sub_q, deps=agent_deps)
|
answer=f"Search failed: {str(e)}",
|
||||||
raw_answer = result.output
|
confidence=0.0,
|
||||||
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,
|
|
||||||
)
|
|
||||||
finally:
|
finally:
|
||||||
if deps.agui_emitter:
|
if deps.agui_emitter:
|
||||||
deps.agui_emitter.finish_step(step_name)
|
deps.agui_emitter.finish_step(step_name)
|
||||||
|
|
@ -263,14 +319,7 @@ def build_research_graph(
|
||||||
ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
|
ctx: StepContext[ResearchState, ResearchDeps, None | bool | str],
|
||||||
) -> list[str] | None:
|
) -> list[str] | None:
|
||||||
"""Get all remaining questions for this iteration."""
|
"""Get all remaining questions for this iteration."""
|
||||||
state = ctx.state
|
return _get_batch_logic(ctx.state)
|
||||||
|
|
||||||
if not state.context.sub_questions:
|
|
||||||
return None
|
|
||||||
|
|
||||||
batch = list(state.context.sub_questions)
|
|
||||||
state.context.sub_questions.clear()
|
|
||||||
return batch
|
|
||||||
|
|
||||||
@g.step
|
@g.step
|
||||||
async def decide(
|
async def decide(
|
||||||
|
|
@ -548,3 +597,133 @@ def build_research_graph(
|
||||||
)
|
)
|
||||||
|
|
||||||
return g.build()
|
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()
|
||||||
|
|
|
||||||
|
|
@ -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):
|
class ResearchReport(BaseModel):
|
||||||
"""Final research report structure."""
|
"""Final research report structure."""
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -20,6 +20,27 @@ Plan requirements:
|
||||||
|
|
||||||
Use the gather_context tool once on the main question before planning."""
|
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.
|
SEARCH_PROMPT = """You are a search and question-answering specialist.
|
||||||
|
|
||||||
Process:
|
Process:
|
||||||
|
|
@ -113,3 +134,20 @@ Style:
|
||||||
- Be professional, objective, and specific.
|
- Be professional, objective, and specific.
|
||||||
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
|
- NEVER use meta-commentary like "This report covers..." or "The findings show...".
|
||||||
Instead, state the actual information directly."""
|
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."""
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue