From 84f25f0627a46c6cd471de8d6293d9a4a10b3137 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 3 Dec 2025 14:36:25 +0200
Subject: [PATCH] Visual grounding in ag-ui-example

---
 examples/ag-ui-research/backend/main.py       |  38 ++++
 .../frontend/components/Agent.tsx             |   1 +
 .../frontend/components/StateDisplay.tsx      | 196 +++++++++++++++++-
 .../haiku/rag/graph/common/prompts.py         |   7 +-
 4 files changed, 238 insertions(+), 4 deletions(-)

diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py
index ccf62f15..bae6a94c 100644
--- a/examples/ag-ui-research/backend/main.py
+++ b/examples/ag-ui-research/backend/main.py
@@ -106,6 +106,8 @@ async def stream_research_agent(request: Request) -> StreamingResponse:
                 # Forward emitter events to stream
                 async def forward_events():
                     async for event in emitter:
+                        # Log events for debugging
+                        logger.info(f"AG-UI Event: {event}")
                         # Filter out ACTIVITY_SNAPSHOT - not supported by CopilotKit
                         if event.get("type") == "ACTIVITY_SNAPSHOT":
                             continue
@@ -161,10 +163,46 @@ async def health_check(_: Request) -> JSONResponse:
     )
 
 
+async def visualize_chunk(request: Request) -> JSONResponse:
+    """Return visual grounding images for a chunk as base64."""
+    import base64
+    from io import BytesIO
+
+    chunk_id = request.path_params["chunk_id"]
+    client = get_client(db_path)
+
+    # Get the chunk
+    chunk = await client.chunk_repository.get_by_id(chunk_id)
+    if not chunk:
+        return JSONResponse({"error": "Chunk not found"}, status_code=404)
+
+    # Get visualization images
+    images = await client.visualize_chunk(chunk)
+    if not images:
+        return JSONResponse({"images": [], "message": "No visual grounding available"})
+
+    # Convert PIL images to base64
+    base64_images = []
+    for img in images:
+        buffer = BytesIO()
+        img.save(buffer, format="PNG")
+        buffer.seek(0)
+        base64_images.append(base64.b64encode(buffer.read()).decode("utf-8"))
+
+    return JSONResponse(
+        {
+            "images": base64_images,
+            "chunk_id": chunk_id,
+            "document_uri": chunk.document_uri,
+        }
+    )
+
+
 # Create Starlette app
 app = Starlette(
     routes=[
         Route("/v1/research/stream", stream_research_agent, methods=["POST"]),
+        Route("/api/visualize/{chunk_id}", visualize_chunk, methods=["GET"]),
         Route("/health", health_check, methods=["GET"]),
     ],
     middleware=[
diff --git a/examples/ag-ui-research/frontend/components/Agent.tsx b/examples/ag-ui-research/frontend/components/Agent.tsx
index ac7aca85..87813d70 100644
--- a/examples/ag-ui-research/frontend/components/Agent.tsx
+++ b/examples/ag-ui-research/frontend/components/Agent.tsx
@@ -92,6 +92,7 @@ function AgentContent() {
 		},
 	});
 
+
 	return (
 		<>
 			<style>{`
diff --git a/examples/ag-ui-research/frontend/components/StateDisplay.tsx b/examples/ag-ui-research/frontend/components/StateDisplay.tsx
index a98a1442..df08e64a 100644
--- a/examples/ag-ui-research/frontend/components/StateDisplay.tsx
+++ b/examples/ag-ui-research/frontend/components/StateDisplay.tsx
@@ -1,7 +1,15 @@
 "use client";
 
 import { Markdown } from "@copilotkit/react-ui";
-import { useState } from "react";
+import { useState, useCallback } from "react";
+
+interface VisualGroundingState {
+	isOpen: boolean;
+	chunkId: string | null;
+	images: string[];
+	loading: boolean;
+	error: string | null;
+}
 
 interface InsightRecord {
 	id: string;
@@ -98,6 +106,58 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 		Record<string, boolean>
 	>({});
 
+	const [visualGrounding, setVisualGrounding] = useState<VisualGroundingState>({
+		isOpen: false,
+		chunkId: null,
+		images: [],
+		loading: false,
+		error: null,
+	});
+
+	const fetchVisualGrounding = useCallback(async (chunkId: string) => {
+		setVisualGrounding({
+			isOpen: true,
+			chunkId,
+			images: [],
+			loading: true,
+			error: null,
+		});
+
+		try {
+			const response = await fetch(
+				`${process.env.NEXT_PUBLIC_API_URL || "http://localhost:8000"}/api/visualize/${chunkId}`
+			);
+			const data = await response.json();
+
+			if (!response.ok) {
+				throw new Error(data.error || "Failed to fetch visual grounding");
+			}
+
+			setVisualGrounding((prev) => ({
+				...prev,
+				images: data.images || [],
+				loading: false,
+				error: data.images?.length === 0 ? data.message : null,
+			}));
+		} catch (err) {
+			setVisualGrounding((prev) => ({
+				...prev,
+				loading: false,
+				error: err instanceof Error ? err.message : "Unknown error",
+			}));
+		}
+	}, []);
+
+	const closeVisualGrounding = useCallback(() => {
+		setVisualGrounding({
+			isOpen: false,
+			chunkId: null,
+			images: [],
+			loading: false,
+			error: null,
+		});
+	}, []);
+
 	const toggleSection = (section: string) => {
 		setExpandedSections((prev) => ({
 			...prev,
@@ -605,6 +665,22 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 																		{citation.content.slice(0, 200)}
 																		{citation.content.length > 200 && "…"}
 																	</div>
+																	<button
+																		type="button"
+																		onClick={() => fetchVisualGrounding(citation.chunk_id)}
+																		style={{
+																			marginTop: "0.5rem",
+																			padding: "0.25rem 0.5rem",
+																			fontSize: "0.7rem",
+																			background: "#4299e1",
+																			color: "white",
+																			border: "none",
+																			borderRadius: "4px",
+																			cursor: "pointer",
+																		}}
+																	>
+																		📍 View in Document
+																	</button>
 																</div>
 															))}
 														</div>
@@ -1130,6 +1206,124 @@ export default function StateDisplay({ state }: StateDisplayProps) {
 					)}
 				</div>
 			)}
+
+			{/* Visual Grounding Modal */}
+			{visualGrounding.isOpen && (
+				<div
+					style={{
+						position: "fixed",
+						top: 0,
+						left: 0,
+						right: 0,
+						bottom: 0,
+						background: "rgba(0, 0, 0, 0.75)",
+						display: "flex",
+						alignItems: "center",
+						justifyContent: "center",
+						zIndex: 1000,
+					}}
+					onClick={closeVisualGrounding}
+				>
+					<div
+						style={{
+							background: "white",
+							borderRadius: "8px",
+							padding: "1.5rem",
+							maxWidth: "90vw",
+							maxHeight: "90vh",
+							overflow: "auto",
+							position: "relative",
+						}}
+						onClick={(e) => e.stopPropagation()}
+					>
+						<button
+							type="button"
+							onClick={closeVisualGrounding}
+							style={{
+								position: "absolute",
+								top: "0.5rem",
+								right: "0.5rem",
+								background: "#e53e3e",
+								color: "white",
+								border: "none",
+								borderRadius: "50%",
+								width: "2rem",
+								height: "2rem",
+								cursor: "pointer",
+								fontSize: "1rem",
+							}}
+						>
+							✕
+						</button>
+						<h3
+							style={{
+								margin: "0 0 1rem 0",
+								fontSize: "1.125rem",
+								color: "#2d3748",
+							}}
+						>
+							Visual Grounding
+						</h3>
+						{visualGrounding.loading && (
+							<div
+								style={{
+									padding: "2rem",
+									textAlign: "center",
+									color: "#718096",
+								}}
+							>
+								Loading...
+							</div>
+						)}
+						{visualGrounding.error && (
+							<div
+								style={{
+									padding: "1rem",
+									background: "#fed7d7",
+									color: "#c53030",
+									borderRadius: "4px",
+								}}
+							>
+								{visualGrounding.error}
+							</div>
+						)}
+						{!visualGrounding.loading &&
+							!visualGrounding.error &&
+							visualGrounding.images.length > 0 && (
+								<div
+									style={{
+										display: "flex",
+										flexDirection: "column",
+										gap: "1rem",
+									}}
+								>
+									{visualGrounding.images.map((img, idx) => (
+										<div key={idx}>
+											<div
+												style={{
+													fontSize: "0.75rem",
+													color: "#718096",
+													marginBottom: "0.5rem",
+												}}
+											>
+												Page {idx + 1} of {visualGrounding.images.length}
+											</div>
+											<img
+												src={`data:image/png;base64,${img}`}
+												alt={`Page ${idx + 1}`}
+												style={{
+													maxWidth: "100%",
+													border: "1px solid #e2e8f0",
+													borderRadius: "4px",
+												}}
+											/>
+										</div>
+									))}
+								</div>
+							)}
+					</div>
+				</div>
+			)}
 		</div>
 	);
 }
diff --git a/haiku_rag_slim/haiku/rag/graph/common/prompts.py b/haiku_rag_slim/haiku/rag/graph/common/prompts.py
index 8073ba4a..2a3a460c 100644
--- a/haiku_rag_slim/haiku/rag/graph/common/prompts.py
+++ b/haiku_rag_slim/haiku/rag/graph/common/prompts.py
@@ -27,10 +27,11 @@ Process:
 4. Provide a concise answer based strictly on the retrieved content.
 
 The search tool returns results like:
-[chunk_abc123] (score: 0.85) Content text here...
-[chunk_def456] (score: 0.72) More content...
+[9bde5847-44c9-400a-8997-0e6b65babf92] (score: 0.85) Content text here...
+[d5a63c82-cb40-439f-9b2e-de7d177829b7] (score: 0.72) More content...
 
-In your response, include the chunk IDs you used in cited_chunks.
+IMPORTANT: In cited_chunks, use the EXACT, COMPLETE chunk ID (the full UUID).
+Do NOT truncate or shorten chunk IDs.
 
 Guidelines:
 - Base answers strictly on retrieved content - do not use external knowledge.