Citations

app/backend/agent.py

@@ -13,10 +13,24 @@ if TYPE_CHECKING:
     from haiku.rag.graph.agui.emitter import AGUIEmitter
 
 
+class CitationInfo(BaseModel):
+    """Citation info for frontend display."""
+
+    index: int
+    document_id: str
+    chunk_id: str
+    document_uri: str
+    document_title: str | None = None
+    page_numbers: list[int] = []
+    headings: list[str] | None = None
+    content: str
+
+
 class ChatSessionState(BaseModel):
     """State shared between frontend and agent via AG-UI."""
 
     session_id: str = ""
+    citations: list[CitationInfo] = []
 
 
 @dataclass
@@ -35,13 +49,12 @@ You have access to a knowledge base of documents. Use your tools to search and a
 
 CRITICAL RULES:
 1. For greetings or casual chat: respond directly WITHOUT using any tools
-2. For substantive questions requiring information: use the search or ask tools
+2. For questions: ALWAYS use the "ask" tool - it provides answers with proper citations
 3. NEVER make up information - always use tools to get facts from the knowledge base
-4. When citing sources, reference the chunk IDs from search results
 
 How to decide which tool to use:
-- "search" - When you need to find relevant documents or explore what's in the knowledge base
+- "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" - When you have a specific question that needs a direct answer with citations
+- "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."""
 
@@ -111,13 +124,29 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
 
         answer, citations = await ctx.deps.client.ask(question, filter=document_filter)
 
-        # Format answer with citations
+        # Emit citations via AG-UI state for frontend rendering
-        if citations:
+        if citations and ctx.deps.agui_emitter:
-            citation_list = "\n".join(
+            citation_infos = [
-                f"  [{i + 1}] {c.document_uri or c.document_title or 'Unknown'} (chunk: {c.chunk_id})"
+                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)
+            ]
+            ctx.deps.agui_emitter.update_state(
+                ChatSessionState(citations=citation_infos)
             )
-            return f"{answer}\n\nSources:\n{citation_list}"
+
+        # 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}"
 
         return answer
 

app/frontend/components/Chat.tsx

@@ -1,12 +1,29 @@
 "use client";
 
-import { CopilotKit, useCoAgent } from "@copilotkit/react-core";
+import {
+	CopilotKit,
+	useCoAgent,
+	useCoAgentStateRender,
+} from "@copilotkit/react-core";
 import { CopilotChat } from "@copilotkit/react-ui";
 import "@copilotkit/react-ui/styles.css";
+import CitationBlock from "./CitationBlock";
 import DbInfo from "./DbInfo";
 
+interface Citation {
+	index: number;
+	document_id: string;
+	chunk_id: string;
+	document_uri: string;
+	document_title: string | null;
+	page_numbers: number[];
+	headings: string[] | null;
+	content: string;
+}
+
 interface ChatSessionState {
 	session_id: string;
+	citations: Citation[];
 }
 
 function ChatContent() {
@@ -14,6 +31,17 @@ function ChatContent() {
 		name: "chat_agent",
 		initialState: {
 			session_id: "",
+			citations: [],
+		},
+	});
+
+	useCoAgentStateRender<ChatSessionState>({
+		name: "chat_agent",
+		render: ({ state }) => {
+			if (state.citations && state.citations.length > 0) {
+				return <CitationBlock citations={state.citations} />;
+			}
+			return null;
 		},
 	});
 
@@ -41,6 +69,12 @@ function ChatContent() {
 				.chat-content {
 					flex: 1;
 					min-height: 0;
+					display: flex;
+					flex-direction: column;
+				}
+				.chat-content > * {
+					flex: 1;
+					min-height: 0;
 				}
 			`}</style>
 			<div className="chat-wrapper">

app/frontend/components/CitationBlock.tsx

@@ -0,0 +1,155 @@
+"use client";
+
+import { useState } from "react";
+
+interface Citation {
+	index: number;
+	document_id: string;
+	chunk_id: string;
+	document_uri: string;
+	document_title: string | null;
+	page_numbers: number[];
+	headings: string[] | null;
+	content: string;
+}
+
+interface CitationBlockProps {
+	citations: Citation[];
+}
+
+function CitationItem({ citation }: { citation: Citation }) {
+	const [expanded, setExpanded] = useState(false);
+
+	const title = citation.document_title || citation.document_uri || "Unknown";
+	const pageInfo =
+		citation.page_numbers.length > 0
+			? `p. ${citation.page_numbers.join(", ")}`
+			: null;
+
+	return (
+		<div className="citation-item">
+			<button
+				type="button"
+				className="citation-header"
+				onClick={() => setExpanded(!expanded)}
+			>
+				<span className="citation-index">[{citation.index}]</span>
+				<span className="citation-title">{title}</span>
+				{pageInfo && <span className="citation-page">{pageInfo}</span>}
+				<span className={`citation-chevron ${expanded ? "expanded" : ""}`}>
+					{expanded ? "▼" : "▶"}
+				</span>
+			</button>
+			{expanded && (
+				<div className="citation-content">
+					{citation.headings && citation.headings.length > 0 && (
+						<div className="citation-headings">
+							{citation.headings.join(" › ")}
+						</div>
+					)}
+					<div className="citation-text">{citation.content}</div>
+				</div>
+			)}
+		</div>
+	);
+}
+
+export default function CitationBlock({ citations }: CitationBlockProps) {
+	if (!citations || citations.length === 0) {
+		return null;
+	}
+
+	return (
+		<>
+			<style>{`
+				.citation-block {
+					margin-top: 0.75rem;
+					border: 1px solid #e2e8f0;
+					border-radius: 8px;
+					overflow: hidden;
+					font-size: 0.875rem;
+				}
+				.citation-block-header {
+					background: #f8fafc;
+					padding: 0.5rem 0.75rem;
+					font-weight: 500;
+					color: #475569;
+					border-bottom: 1px solid #e2e8f0;
+				}
+				.citation-item {
+					border-bottom: 1px solid #f1f5f9;
+				}
+				.citation-item:last-child {
+					border-bottom: none;
+				}
+				.citation-header {
+					display: flex;
+					align-items: center;
+					gap: 0.5rem;
+					width: 100%;
+					padding: 0.5rem 0.75rem;
+					background: none;
+					border: none;
+					cursor: pointer;
+					text-align: left;
+					transition: background 0.15s;
+				}
+				.citation-header:hover {
+					background: #f8fafc;
+				}
+				.citation-index {
+					font-weight: 600;
+					color: #3b82f6;
+					flex-shrink: 0;
+				}
+				.citation-title {
+					flex: 1;
+					color: #334155;
+					overflow: hidden;
+					text-overflow: ellipsis;
+					white-space: nowrap;
+				}
+				.citation-page {
+					color: #94a3b8;
+					font-size: 0.75rem;
+					flex-shrink: 0;
+				}
+				.citation-chevron {
+					color: #94a3b8;
+					font-size: 0.625rem;
+					flex-shrink: 0;
+					transition: transform 0.15s;
+				}
+				.citation-chevron.expanded {
+					transform: rotate(0deg);
+				}
+				.citation-content {
+					padding: 0.75rem;
+					background: #fafafa;
+					border-top: 1px solid #f1f5f9;
+				}
+				.citation-headings {
+					font-size: 0.75rem;
+					color: #64748b;
+					margin-bottom: 0.5rem;
+					font-style: italic;
+				}
+				.citation-text {
+					color: #475569;
+					line-height: 1.5;
+					white-space: pre-wrap;
+					max-height: 200px;
+					overflow-y: auto;
+				}
+			`}</style>
+			<div className="citation-block">
+				<div className="citation-block-header">
+					Sources ({citations.length})
+				</div>
+				{citations.map((citation) => (
+					<CitationItem key={citation.chunk_id} citation={citation} />
+				))}
+			</div>
+		</>
+	);
+}