Visual grounding in ag-ui-example
This commit is contained in:
Yiorgis Gozadinos
parent
6e31ea38a8
commit
84f25f0627
4 changed files with 238 additions and 4 deletions
|
|
@ -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=[
|
||||
|
|
|
|||
|
|
@ -92,6 +92,7 @@ function AgentContent() {
|
|||
},
|
||||
});
|
||||
|
||||
|
||||
return (
|
||||
<>
|
||||
<style>{`
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
Loading…
Reference in a new issue