diff --git a/app/backend/agent.py b/app/backend/agent.py index 3a232518..c653d688 100644 --- a/app/backend/agent.py +++ b/app/backend/agent.py @@ -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( diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 3a2f3494..2558573e 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -24,7 +24,8 @@ interface Citation { interface QAResponse { question: string; answer: string; - sources: string[]; + confidence: number; + citations: Citation[]; } interface ChatSessionState { diff --git a/haiku_rag_slim/haiku/rag/graph/research/graph.py b/haiku_rag_slim/haiku/rag/graph/research/graph.py index 91e3536b..a59c7fd0 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/research/graph.py @@ -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() diff --git a/haiku_rag_slim/haiku/rag/graph/research/models.py b/haiku_rag_slim/haiku/rag/graph/research/models.py index 922132f4..9cd37168 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/models.py +++ b/haiku_rag_slim/haiku/rag/graph/research/models.py @@ -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.""" diff --git a/haiku_rag_slim/haiku/rag/graph/research/prompts.py b/haiku_rag_slim/haiku/rag/graph/research/prompts.py index 6bfa8ed4..b472e18e 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/research/prompts.py @@ -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."""