609 lines
15 KiB
TypeScript
609 lines
15 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 state_namespace from haiku.rag.skills.rag
|
|
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 FileIcon() {
|
|
return (
|
|
<svg
|
|
width="14"
|
|
height="14"
|
|
viewBox="0 0 24 24"
|
|
fill="none"
|
|
stroke="currentColor"
|
|
strokeWidth="2"
|
|
strokeLinecap="round"
|
|
strokeLinejoin="round"
|
|
>
|
|
<path d="M15 2H6a2 2 0 0 0-2 2v16a2 2 0 0 0 2 2h12a2 2 0 0 0 2-2V7Z" />
|
|
<path d="M14 2v4a2 2 0 0 0 2 2h4" />
|
|
</svg>
|
|
);
|
|
}
|
|
|
|
function ToolCallIndicator({
|
|
toolName,
|
|
status,
|
|
args,
|
|
}: {
|
|
toolName: string;
|
|
status: string;
|
|
args: Record<string, unknown>;
|
|
}) {
|
|
const isComplete = status === "complete";
|
|
|
|
const getToolIcon = () => {
|
|
switch (toolName) {
|
|
case "search":
|
|
return <SearchIcon />;
|
|
case "get_document":
|
|
return <FileIcon />;
|
|
case "execute_skill":
|
|
case "execute_code":
|
|
case "cite":
|
|
return <MessageIcon />;
|
|
default:
|
|
return <SearchIcon />;
|
|
}
|
|
};
|
|
|
|
const getToolLabel = () => {
|
|
switch (toolName) {
|
|
case "search":
|
|
return "Search";
|
|
case "get_document":
|
|
return "Document";
|
|
case "execute_skill":
|
|
return "Skill";
|
|
case "execute_code":
|
|
return "Code";
|
|
case "cite":
|
|
return "Cite";
|
|
case "list_documents":
|
|
return "Documents";
|
|
default:
|
|
return toolName;
|
|
}
|
|
};
|
|
|
|
const getDescription = () => {
|
|
switch (toolName) {
|
|
case "execute_skill": {
|
|
const skill = args.skill_name as string | undefined;
|
|
const request = args.request as string | undefined;
|
|
return (
|
|
<span className="tool-query">
|
|
{skill ? `${skill}: ` : ""}
|
|
{request ?? "Processing..."}
|
|
</span>
|
|
);
|
|
}
|
|
case "search": {
|
|
const query = args.query as string;
|
|
return <span className="tool-query">{query}</span>;
|
|
}
|
|
case "get_document":
|
|
return <span className="tool-query">{args.query as string}</span>;
|
|
case "execute_code": {
|
|
const code = args.code as string | undefined;
|
|
return (
|
|
<span className="tool-query">
|
|
{code ? code.slice(0, 80) : "Running code..."}
|
|
</span>
|
|
);
|
|
}
|
|
case "cite":
|
|
return <span className="tool-query">Registering citations</span>;
|
|
case "list_documents":
|
|
return <span className="tool-query">Listing documents</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>
|
|
);
|
|
}
|
|
|
|
// Render an activity message from a skill sub-agent tool call/result
|
|
function ActivityIndicator({
|
|
message,
|
|
isComplete,
|
|
}: {
|
|
// biome-ignore lint/suspicious/noExplicitAny: AG-UI activity message shape
|
|
message: any;
|
|
isComplete: boolean;
|
|
}) {
|
|
const content = message.content ?? {};
|
|
const toolName = content.tool_name ?? "tool";
|
|
|
|
let args: Record<string, unknown> = {};
|
|
if (content.args) {
|
|
try {
|
|
args =
|
|
typeof content.args === "string"
|
|
? JSON.parse(content.args)
|
|
: content.args;
|
|
} catch {
|
|
// ignore parse errors
|
|
}
|
|
}
|
|
|
|
return (
|
|
<ToolCallIndicator
|
|
toolName={toolName}
|
|
status={isComplete ? "complete" : "loading"}
|
|
args={args}
|
|
/>
|
|
);
|
|
}
|
|
|
|
// 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) : [];
|
|
|
|
// Collect completed tool_call_ids from skill_tool_result activity messages
|
|
const completedToolCallIds = useMemo(() => {
|
|
const ids = new Set<string>();
|
|
for (const msg of messages) {
|
|
if (
|
|
msg.role === "activity" &&
|
|
msg.activityType === "skill_tool_result" &&
|
|
msg.content?.tool_call_id
|
|
) {
|
|
ids.add(msg.content.tool_call_id);
|
|
}
|
|
}
|
|
return ids;
|
|
}, [messages]);
|
|
|
|
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.
|
|
// We interleave activity indicators (skill sub-agent tool calls) and
|
|
// optionally 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;
|
|
}
|
|
|
|
// Activity messages are not rendered by CopilotKit —
|
|
// render them ourselves without consuming messageElements
|
|
if (msg.role === "activity") {
|
|
if (msg.activityType === "skill_tool_call") {
|
|
const toolCallId = msg.content?.tool_call_id;
|
|
result.push(
|
|
<ActivityIndicator
|
|
key={`activity-${msg.id}`}
|
|
message={msg}
|
|
isComplete={
|
|
toolCallId ? completedToolCallIds.has(toolCallId) : false
|
|
}
|
|
/>,
|
|
);
|
|
}
|
|
continue;
|
|
}
|
|
|
|
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 the namespaced AG-UI state so the backend's first STATE_DELTA
|
|
// (e.g. add /rag/searches/...) has a namespace object to patch into.
|
|
agent.setState({
|
|
[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, keeping the last occurrence.
|
|
// haiku.skills 0.10.0+ sends activity snapshots with the same id
|
|
// and replace=true — CopilotKit doesn't deduplicate these, so we
|
|
// must do it 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>
|
|
);
|
|
}
|