53bbf52697
The session store rebuilt the rag namespace from four known keys, so the evidence record never survived a turn, let alone a reload from localStorage. Compaction then ran with an empty ledger: earlier evidence was replaced by receipts retaining nothing, while the citations already in citation_index kept the UI looking correct. The UI still names the fields it reads, but everything else in the namespace passes through untouched, and seeding the namespace no longer replaces sibling namespaces.
487 lines
12 KiB
TypeScript
487 lines
12 KiB
TypeScript
"use client";
|
|
|
|
import {
|
|
CopilotChatMessageView,
|
|
CopilotChatView,
|
|
CopilotKitProvider,
|
|
defineToolCallRenderer,
|
|
UseAgentUpdate,
|
|
useAgent,
|
|
useCopilotKit,
|
|
} from "@copilotkit/react-core/v2";
|
|
import {
|
|
createContext,
|
|
useCallback,
|
|
useContext,
|
|
useEffect,
|
|
useMemo,
|
|
useState,
|
|
} from "react";
|
|
import { FilterIcon } from "../lib/icons";
|
|
import type { RAGState } from "../lib/sessionStorage";
|
|
import {
|
|
createSession,
|
|
getActiveSessionId,
|
|
getLatestCitations,
|
|
getSession,
|
|
normalizeRAGState,
|
|
updateSessionMessages,
|
|
} from "../lib/sessionStorage";
|
|
import CitationBlock from "./CitationBlock";
|
|
import DbInfo from "./DbInfo";
|
|
import DocumentFilter from "./DocumentFilter";
|
|
import SessionManager from "./SessionManager";
|
|
|
|
// Must match RAGCapability.state_namespace.
|
|
const AGUI_STATE_KEY = "rag";
|
|
|
|
// AG-UI state is namespaced under AGUI_STATE_KEY
|
|
interface AgentState {
|
|
[AGUI_STATE_KEY]?: RAGState;
|
|
}
|
|
|
|
// biome-ignore lint/suspicious/noExplicitAny: CopilotKit message objects vary at runtime
|
|
function serializeMessages(messages: any[]): any[] {
|
|
return JSON.parse(JSON.stringify(messages));
|
|
}
|
|
|
|
function SpinnerIcon() {
|
|
return (
|
|
<svg
|
|
width="16"
|
|
height="16"
|
|
viewBox="0 0 24 24"
|
|
fill="none"
|
|
stroke="currentColor"
|
|
strokeWidth="2"
|
|
strokeLinecap="round"
|
|
strokeLinejoin="round"
|
|
className="tool-spinner"
|
|
>
|
|
<path d="M21 12a9 9 0 1 1-6.219-8.56" />
|
|
</svg>
|
|
);
|
|
}
|
|
|
|
function CheckIcon() {
|
|
return (
|
|
<svg
|
|
width="16"
|
|
height="16"
|
|
viewBox="0 0 24 24"
|
|
fill="none"
|
|
stroke="currentColor"
|
|
strokeWidth="2.5"
|
|
strokeLinecap="round"
|
|
strokeLinejoin="round"
|
|
>
|
|
<polyline points="20 6 9 17 4 12" />
|
|
</svg>
|
|
);
|
|
}
|
|
|
|
function SearchIcon() {
|
|
return (
|
|
<svg
|
|
width="14"
|
|
height="14"
|
|
viewBox="0 0 24 24"
|
|
fill="none"
|
|
stroke="currentColor"
|
|
strokeWidth="2"
|
|
strokeLinecap="round"
|
|
strokeLinejoin="round"
|
|
>
|
|
<circle cx="11" cy="11" r="8" />
|
|
<path d="m21 21-4.3-4.3" />
|
|
</svg>
|
|
);
|
|
}
|
|
|
|
function MessageIcon() {
|
|
return (
|
|
<svg
|
|
width="14"
|
|
height="14"
|
|
viewBox="0 0 24 24"
|
|
fill="none"
|
|
stroke="currentColor"
|
|
strokeWidth="2"
|
|
strokeLinecap="round"
|
|
strokeLinejoin="round"
|
|
>
|
|
<path d="M7.9 20A9 9 0 1 0 4 16.1L2 22Z" />
|
|
</svg>
|
|
);
|
|
}
|
|
|
|
function ToolCallIndicator({
|
|
toolName,
|
|
status,
|
|
args,
|
|
}: {
|
|
toolName: string;
|
|
status: string;
|
|
args: Record<string, unknown>;
|
|
}) {
|
|
const isComplete = status === "complete";
|
|
|
|
const getToolIcon = () => {
|
|
switch (toolName) {
|
|
case "rag_search":
|
|
return <SearchIcon />;
|
|
case "rag_cite":
|
|
return <MessageIcon />;
|
|
default:
|
|
return <SearchIcon />;
|
|
}
|
|
};
|
|
|
|
const getToolLabel = () => {
|
|
switch (toolName) {
|
|
case "rag_search":
|
|
return "Search";
|
|
case "rag_cite":
|
|
return "Cite";
|
|
default:
|
|
return toolName;
|
|
}
|
|
};
|
|
|
|
const getDescription = () => {
|
|
switch (toolName) {
|
|
case "rag_search": {
|
|
const query = args.query as string;
|
|
return <span className="tool-query">{query}</span>;
|
|
}
|
|
case "rag_cite":
|
|
return <span className="tool-query">Registering citations</span>;
|
|
default:
|
|
return <span>Processing...</span>;
|
|
}
|
|
};
|
|
|
|
return (
|
|
<div className={`tool-call-card ${isComplete ? "complete" : "loading"}`}>
|
|
<div className="tool-status-icon">
|
|
{isComplete ? <CheckIcon /> : <SpinnerIcon />}
|
|
</div>
|
|
<div className="tool-content">
|
|
<div className="tool-header">
|
|
<span className="tool-badge">
|
|
{getToolIcon()}
|
|
{getToolLabel()}
|
|
</span>
|
|
<span className="tool-status-text">
|
|
{isComplete ? "Done" : "Working..."}
|
|
</span>
|
|
</div>
|
|
<div className="tool-description">{getDescription()}</div>
|
|
</div>
|
|
</div>
|
|
);
|
|
}
|
|
|
|
// Context for sharing chat state with the message view
|
|
const ChatStateContext = createContext<RAGState | null>(null);
|
|
|
|
// Wildcard tool call renderer for all server-side tools
|
|
const toolCallRenderers = [
|
|
defineToolCallRenderer({
|
|
name: "*",
|
|
render: ({ name, args, result }) => (
|
|
<ToolCallIndicator
|
|
toolName={name}
|
|
status={result !== undefined ? "complete" : "loading"}
|
|
args={(args ?? {}) as Record<string, unknown>}
|
|
/>
|
|
),
|
|
}),
|
|
];
|
|
|
|
// Custom message view that injects CitationBlocks after assistant responses.
|
|
// Uses CopilotChatMessageView's children render prop to post-process the
|
|
// rendered message elements and inject citations at the right positions.
|
|
function MessageViewWithCitations({
|
|
messages = [],
|
|
isRunning = false,
|
|
}: {
|
|
// biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union
|
|
messages?: any[];
|
|
isRunning?: boolean;
|
|
}) {
|
|
const ragState = useContext(ChatStateContext);
|
|
const latestCitations = ragState ? getLatestCitations(ragState) : [];
|
|
|
|
const cursor = isRunning ? (
|
|
<div key="cursor" className="streaming-cursor">
|
|
<span className="dot" />
|
|
<span className="dot" />
|
|
<span className="dot" />
|
|
</div>
|
|
) : null;
|
|
|
|
// CopilotChatMessageView renders one element per user/assistant message.
|
|
// Inject CitationBlocks after assistant responses that
|
|
// followed tool calls.
|
|
return (
|
|
<CopilotChatMessageView messages={messages} isRunning={isRunning}>
|
|
{({ messageElements }) => {
|
|
const result: React.ReactNode[] = [];
|
|
let elemIdx = 0;
|
|
let seenToolCalls = false;
|
|
|
|
for (const msg of messages) {
|
|
if (msg.role === "user") {
|
|
seenToolCalls = false;
|
|
}
|
|
|
|
if (
|
|
msg.role === "assistant" &&
|
|
Array.isArray(msg.toolCalls) &&
|
|
msg.toolCalls.length > 0
|
|
) {
|
|
seenToolCalls = true;
|
|
}
|
|
|
|
if (msg.role !== "user" && msg.role !== "assistant") continue;
|
|
|
|
if (elemIdx < messageElements.length) {
|
|
result.push(messageElements[elemIdx]);
|
|
elemIdx++;
|
|
}
|
|
|
|
// After an assistant text response that followed tool calls,
|
|
// show citations from the latest turn
|
|
if (msg.role === "assistant" && msg.content && seenToolCalls) {
|
|
if (latestCitations.length > 0) {
|
|
result.push(
|
|
<CitationBlock
|
|
key={`citations-${msg.id}`}
|
|
citations={latestCitations}
|
|
/>,
|
|
);
|
|
}
|
|
seenToolCalls = false;
|
|
}
|
|
}
|
|
|
|
while (elemIdx < messageElements.length) {
|
|
result.push(messageElements[elemIdx]);
|
|
elemIdx++;
|
|
}
|
|
|
|
return (
|
|
<>
|
|
{result}
|
|
{cursor}
|
|
</>
|
|
);
|
|
}}
|
|
</CopilotChatMessageView>
|
|
);
|
|
}
|
|
MessageViewWithCitations.Cursor = CopilotChatMessageView.Cursor;
|
|
|
|
function ChatContentInner({
|
|
sessionId,
|
|
onSessionChange,
|
|
}: {
|
|
sessionId: string;
|
|
onSessionChange: (id: string) => void;
|
|
}) {
|
|
const [filterOpen, setFilterOpen] = useState(false);
|
|
// Track selected document names locally (frontend-only)
|
|
const [selectedDocuments, setSelectedDocuments] = useState<string[]>([]);
|
|
|
|
const { agent } = useAgent({
|
|
agentId: "chat_agent",
|
|
updates: [
|
|
UseAgentUpdate.OnMessagesChanged,
|
|
UseAgentUpdate.OnStateChanged,
|
|
UseAgentUpdate.OnRunStatusChanged,
|
|
],
|
|
});
|
|
const { copilotkit: ck } = useCopilotKit();
|
|
|
|
// Set threadId (CopilotChat normally does this in its connect effect)
|
|
useEffect(() => {
|
|
agent.threadId = sessionId;
|
|
}, [agent, sessionId]);
|
|
|
|
const ragState = normalizeRAGState(
|
|
(agent.state as AgentState)?.[AGUI_STATE_KEY],
|
|
);
|
|
|
|
// Restore session from localStorage when agent reference changes.
|
|
// useAgent returns a provisional agent initially, then the real agent
|
|
// after runtime connects — re-run restore each time so messages stick.
|
|
useEffect(() => {
|
|
if (agent.messages.length > 0) return;
|
|
const session = getSession(sessionId);
|
|
// Seed state for the capability; the backend replaces it after each run.
|
|
agent.setState({
|
|
...agent.state,
|
|
[AGUI_STATE_KEY]: normalizeRAGState(session?.ragState),
|
|
});
|
|
if (session && session.messages.length > 0) {
|
|
// biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union
|
|
agent.setMessages(session.messages as any[]);
|
|
}
|
|
}, [agent, sessionId]);
|
|
|
|
// Persist messages and state to localStorage.
|
|
// Read ragState from agent.state at effect time (not render time) so that
|
|
// restore and persist effects in the same commit see consistent state.
|
|
// biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes
|
|
useEffect(() => {
|
|
if (sessionId && agent.messages.length > 0) {
|
|
const currentRagState = normalizeRAGState(
|
|
(agent.state as AgentState)?.[AGUI_STATE_KEY],
|
|
);
|
|
updateSessionMessages(
|
|
sessionId,
|
|
serializeMessages(agent.messages),
|
|
currentRagState,
|
|
);
|
|
}
|
|
}, [JSON.stringify(agent.messages), ragState, sessionId]);
|
|
|
|
// Deduplicate messages by id to avoid React duplicate key warnings.
|
|
// biome-ignore lint/correctness/useExhaustiveDependencies: stable identity via agent ref
|
|
const messages = useMemo(() => {
|
|
const seen = new Map<string, number>();
|
|
const msgs = agent.messages;
|
|
for (let i = 0; i < msgs.length; i++) {
|
|
const id = msgs[i].id;
|
|
if (id) seen.set(id, i);
|
|
}
|
|
return msgs.filter((msg, i) => !msg.id || seen.get(msg.id) === i);
|
|
}, [JSON.stringify(agent.messages)]);
|
|
|
|
const onSubmitMessage = useCallback(
|
|
async (text: string) => {
|
|
agent.addMessage({
|
|
id: crypto.randomUUID(),
|
|
role: "user",
|
|
content: text,
|
|
});
|
|
try {
|
|
await ck.runAgent({ agent });
|
|
} catch (error) {
|
|
console.error("runAgent failed", error);
|
|
}
|
|
},
|
|
[agent, ck],
|
|
);
|
|
|
|
const onStop = useCallback(() => {
|
|
try {
|
|
ck.stopAgent({ agent });
|
|
} catch {
|
|
agent.abortRun();
|
|
}
|
|
}, [agent, ck]);
|
|
|
|
const handleFilterApply = (selected: string[]) => {
|
|
setSelectedDocuments(selected);
|
|
// Convert selected document names to SQL filter for the backend
|
|
const filter =
|
|
selected.length > 0
|
|
? selected
|
|
.map(
|
|
(name) =>
|
|
`(title LIKE '%${name.replace(/'/g, "''")}%' OR uri LIKE '%${name.replace(/'/g, "''")}%')`,
|
|
)
|
|
.join(" OR ")
|
|
: null;
|
|
agent.setState({
|
|
...agent.state,
|
|
[AGUI_STATE_KEY]: {
|
|
...ragState,
|
|
document_filter: filter,
|
|
},
|
|
});
|
|
};
|
|
|
|
return (
|
|
<ChatStateContext.Provider value={ragState}>
|
|
<div className="chat-wrapper">
|
|
<div className="chat-container">
|
|
<div className="chat-header">
|
|
<SessionManager
|
|
activeSessionId={sessionId}
|
|
onSessionChange={onSessionChange}
|
|
/>
|
|
<button
|
|
type="button"
|
|
className={`header-btn ${selectedDocuments.length > 0 ? "has-content" : ""}`}
|
|
onClick={() => setFilterOpen(true)}
|
|
title={
|
|
selectedDocuments.length > 0
|
|
? `Filtering: ${selectedDocuments.length} document(s)`
|
|
: "Filter documents"
|
|
}
|
|
>
|
|
<FilterIcon />
|
|
{selectedDocuments.length > 0
|
|
? `Filter (${selectedDocuments.length})`
|
|
: "Filter"}
|
|
</button>
|
|
</div>
|
|
<div className="chat-content">
|
|
<CopilotChatView
|
|
messageView={MessageViewWithCitations}
|
|
messages={messages}
|
|
isRunning={agent.isRunning}
|
|
onSubmitMessage={onSubmitMessage}
|
|
onStop={onStop}
|
|
>
|
|
{({ scrollView, input }) => (
|
|
<div className="chat-layout">
|
|
<div className="chat-scroll-area">{scrollView}</div>
|
|
<div className="chat-input-area">{input}</div>
|
|
</div>
|
|
)}
|
|
</CopilotChatView>
|
|
</div>
|
|
<DbInfo />
|
|
</div>
|
|
</div>
|
|
<DocumentFilter
|
|
isOpen={filterOpen}
|
|
onClose={() => setFilterOpen(false)}
|
|
selected={selectedDocuments}
|
|
onApply={handleFilterApply}
|
|
/>
|
|
</ChatStateContext.Provider>
|
|
);
|
|
}
|
|
|
|
export default function Chat() {
|
|
const [activeSessionId, setActiveSessionId] = useState<string | null>(null);
|
|
|
|
useEffect(() => {
|
|
let id = getActiveSessionId();
|
|
if (!id) {
|
|
id = createSession().id;
|
|
}
|
|
setActiveSessionId(id);
|
|
}, []);
|
|
|
|
if (!activeSessionId) return null;
|
|
|
|
return (
|
|
<CopilotKitProvider
|
|
key={activeSessionId}
|
|
runtimeUrl="/api/copilotkit"
|
|
useSingleEndpoint
|
|
renderToolCalls={toolCallRenderers}
|
|
>
|
|
<ChatContentInner
|
|
sessionId={activeSessionId}
|
|
onSessionChange={setActiveSessionId}
|
|
/>
|
|
</CopilotKitProvider>
|
|
);
|
|
}
|