e1dd8517f9
Both optional capabilities read what earlier questions retrieved and cited from the capability's state, so a host that carries only the message history hands every run an empty record. Compaction then replaced the earlier evidence with receipts and retained nothing, and the loss was invisible: the citations the host already displayed were still there. It now refuses when it finds evidence from an earlier question and no record of what that question cited. `state_carried` reaches the optional capabilities through discovery, so the refusal distinguishes a host that never carries state from a question that simply cited nothing. The documentation taught the pattern that breaks: the compose example is now stateful and the requirement is stated where each capability is introduced. The app's browser storage was doing exactly this, keeping only the fields the UI reads. It now persists the whole namespace map, so the citation policy's violations survive a reload as well as the evidence record.
482 lines
12 KiB
TypeScript
482 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 {
|
|
AGUI_STATE_KEY,
|
|
agentStateOf,
|
|
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";
|
|
|
|
// AG-UI state is namespaced under AGUI_STATE_KEY (see sessionStorage).
|
|
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 capabilities; the backend replaces it after each run.
|
|
// The whole namespace map goes back, not just the fields this UI reads:
|
|
// compaction and the citation policy read what earlier questions recorded.
|
|
agent.setState({ ...agent.state, ...agentStateOf(session ?? undefined) });
|
|
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) {
|
|
updateSessionMessages(
|
|
sessionId,
|
|
serializeMessages(agent.messages),
|
|
(agent.state ?? {}) as Record<string, unknown>,
|
|
);
|
|
}
|
|
}, [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>
|
|
);
|
|
}
|