From 14873862f885b10ff2a83946f09a918e15086e79 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 8 Jan 2026 11:57:54 +0200 Subject: [PATCH] Q/A history --- app/backend/agent.py | 85 +++++++++++++++++++++++++++++--- app/backend/main.py | 24 ++++++--- app/frontend/components/Chat.tsx | 8 +++ 3 files changed, 104 insertions(+), 13 deletions(-) diff --git a/app/backend/agent.py b/app/backend/agent.py index 486c3a97..d8f372c0 100644 --- a/app/backend/agent.py +++ b/app/backend/agent.py @@ -2,7 +2,7 @@ from dataclasses import dataclass from typing import TYPE_CHECKING from pydantic import BaseModel -from pydantic_ai import Agent, RunContext +from pydantic_ai import Agent, RunContext, format_as_xml from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig @@ -26,11 +26,38 @@ class CitationInfo(BaseModel): content: str +class QAResponse(BaseModel): + """A Q&A pair from conversation history.""" + + question: str + answer: str + sources: list[str] = [] + + class ChatSessionState(BaseModel): """State shared between frontend and agent via AG-UI.""" session_id: str = "" citations: list[CitationInfo] = [] + qa_history: list[QAResponse] = [] + + +def format_conversation_context(qa_history: list[QAResponse]) -> str: + """Format conversation history as XML for inclusion in prompts.""" + if not qa_history: + return "" + + context_data = { + "previous_qa": [ + { + "question": qa.question, + "answer": qa.answer, + "sources": qa.sources, + } + for qa in qa_history + ], + } + return format_as_xml(context_data, root_tag="conversation_context") @dataclass @@ -41,6 +68,7 @@ class ChatDeps: config: AppConfig agui_emitter: "AGUIEmitter | None" = None search_results: list[SearchResult] | None = None + session_state: ChatSessionState | None = None CHAT_SYSTEM_PROMPT = """You are a helpful research assistant powered by haiku.rag, a knowledge base system. @@ -53,7 +81,7 @@ CRITICAL RULES: 3. NEVER make up information - always use tools to get facts from the knowledge base How to decide which tool to use: -- "ask" - DEFAULT CHOICE for any question. Use this for questions like "What is X?", "How does Y work?", "Explain Z", etc. Returns answers with citations. +- "ask" - DEFAULT CHOICE for any question. Use this for questions like "What is X?", "How does Y work?", "Explain Z", etc. Returns answers with citations. The ask tool maintains conversation context, so follow-up questions benefit from previous answers. - "search" - ONLY use when explicitly exploring/browsing the knowledge base, or when the user asks to "search for" or "find" something without needing an answer. Be friendly and conversational. When you use tools, summarize the key findings for the user.""" @@ -122,10 +150,40 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: if ctx.deps.agui_emitter: ctx.deps.agui_emitter.log(f"Answering: {question}") - answer, citations = await ctx.deps.client.ask(question, filter=document_filter) + # Build context-aware system prompt if we have history + system_prompt = None + if ctx.deps.session_state and ctx.deps.session_state.qa_history: + from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT - # Emit citations via AG-UI state for frontend rendering - if citations and ctx.deps.agui_emitter: + 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." + ) + + answer, citations = await ctx.deps.client.ask( + question, system_prompt=system_prompt, filter=document_filter + ) + + # 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 [] + ) + qa_response = QAResponse( + question=question, + answer=answer, + sources=list(dict.fromkeys(sources)), # dedupe preserving order + ) + ctx.deps.session_state.qa_history.append(qa_response) + + # Build citation infos for frontend + citation_infos = [] + if citations: citation_infos = [ CitationInfo( index=i + 1, @@ -139,8 +197,23 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: ) 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( - ChatSessionState(citations=citation_infos) + ChatSessionState( + session_id=( + ctx.deps.session_state.session_id + if ctx.deps.session_state + else "" + ), + citations=citation_infos, + qa_history=( + ctx.deps.session_state.qa_history + if ctx.deps.session_state + else [] + ), + ) ) # Format answer with citation references diff --git a/app/backend/main.py b/app/backend/main.py index c9091e75..3264f4ac 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -2,7 +2,7 @@ import logging import os from pathlib import Path -from agent import ChatDeps, ChatSessionState, create_chat_agent +from agent import ChatDeps, ChatSessionState, QAResponse, create_chat_agent from anyio import create_memory_object_stream, create_task_group from anyio.streams.memory import MemoryObjectSendStream from dotenv import load_dotenv @@ -122,16 +122,26 @@ async def stream_chat(request: Request) -> StreamingResponse: effective_db_path = Path(input_data.config["db_path"]) client = get_client(effective_db_path) - # Create deps + # Parse incoming state to restore qa_history + initial_qa_history: list[QAResponse] = [] + if input_data.state and "qa_history" in input_data.state: + initial_qa_history = [ + QAResponse(**qa) + for qa in input_data.state.get("qa_history", []) + ] + + # Create initial state with restored history + initial_state = ChatSessionState( + session_id=input_data.thread_id or "", + qa_history=initial_qa_history, + ) + + # Create deps with session state deps = ChatDeps( client=client, config=Config, agui_emitter=emitter, - ) - - # Start run with empty state - initial_state = ChatSessionState( - session_id=input_data.thread_id or "", + session_state=initial_state, ) emitter.start_run(initial_state=initial_state) diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 2fabe376..2b3d37c7 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -21,9 +21,16 @@ interface Citation { content: string; } +interface QAResponse { + question: string; + answer: string; + sources: string[]; +} + interface ChatSessionState { session_id: string; citations: Citation[]; + qa_history: QAResponse[]; } function ChatContent() { @@ -32,6 +39,7 @@ function ChatContent() { initialState: { session_id: "", citations: [], + qa_history: [], }, });