mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-01-02 22:32:27 +08:00
upgrade to ai-sdk 5
This commit is contained in:
dayuan.jiang
@@ -5,13 +5,13 @@ import { useRef, useEffect, useState, useCallback } from "react";
|
||||
import Image from "next/image";
|
||||
import { ScrollArea } from "@/components/ui/scroll-area";
|
||||
import ExamplePanel from "./chat-example-panel";
|
||||
import { Message } from "ai";
|
||||
import { UIMessage } from "ai";
|
||||
import { convertToLegalXml, replaceNodes } from "@/lib/utils";
|
||||
|
||||
import { useDiagram } from "@/contexts/diagram-context";
|
||||
|
||||
interface ChatMessageDisplayProps {
|
||||
messages: Message[];
|
||||
messages: UIMessage[];
|
||||
error?: Error | null;
|
||||
setInput: (input: string) => void;
|
||||
setFiles: (files: File[]) => void;
|
||||
@@ -30,15 +30,18 @@ export function ChatMessageDisplay({
|
||||
const [expandedTools, setExpandedTools] = useState<Record<string, boolean>>(
|
||||
{}
|
||||
);
|
||||
const handleDisplayChart = useCallback((xml: string) => {
|
||||
const currentXml = xml || "";
|
||||
const convertedXml = convertToLegalXml(currentXml);
|
||||
if (convertedXml !== previousXML.current) {
|
||||
previousXML.current = convertedXml;
|
||||
const replacedXML = replaceNodes(chartXML, convertedXml);
|
||||
onDisplayChart(replacedXML);
|
||||
}
|
||||
}, [chartXML, onDisplayChart]);
|
||||
const handleDisplayChart = useCallback(
|
||||
(xml: string) => {
|
||||
const currentXml = xml || "";
|
||||
const convertedXml = convertToLegalXml(currentXml);
|
||||
if (convertedXml !== previousXML.current) {
|
||||
previousXML.current = convertedXml;
|
||||
const replacedXML = replaceNodes(chartXML, convertedXml);
|
||||
onDisplayChart(replacedXML);
|
||||
}
|
||||
},
|
||||
[chartXML, onDisplayChart]
|
||||
);
|
||||
|
||||
useEffect(() => {
|
||||
if (messagesEndRef.current) {
|
||||
@@ -50,27 +53,36 @@ export function ChatMessageDisplay({
|
||||
useEffect(() => {
|
||||
messages.forEach((message) => {
|
||||
if (message.parts) {
|
||||
message.parts.forEach((part) => {
|
||||
if (part.type === "tool-invocation") {
|
||||
const { toolCallId, state, args, toolName } = part.toolInvocation;
|
||||
|
||||
message.parts.forEach((part: any) => {
|
||||
if (part.type?.startsWith("tool-")) {
|
||||
const { toolCallId, state } = part;
|
||||
|
||||
// Auto-collapse args when diagrams are generated
|
||||
if (state === "result") {
|
||||
if (state === "output-available") {
|
||||
setExpandedTools((prev) => ({
|
||||
...prev,
|
||||
[toolCallId]: false,
|
||||
}));
|
||||
}
|
||||
|
||||
|
||||
// Handle diagram updates for display_diagram tool
|
||||
if (toolName === "display_diagram" && args?.xml) {
|
||||
// For partial calls, always update to show streaming
|
||||
if (state === "partial-call") {
|
||||
handleDisplayChart(args.xml);
|
||||
if (
|
||||
part.type === "tool-display_diagram" &&
|
||||
part.input?.xml
|
||||
) {
|
||||
// For streaming input, always update to show streaming
|
||||
if (
|
||||
state === "input-streaming" ||
|
||||
state === "input-available"
|
||||
) {
|
||||
handleDisplayChart(part.input.xml);
|
||||
}
|
||||
// For completed calls, only update if not processed yet
|
||||
else if (state === "result" && !processedToolCalls.current.has(toolCallId)) {
|
||||
handleDisplayChart(args.xml);
|
||||
else if (
|
||||
state === "output-available" &&
|
||||
!processedToolCalls.current.has(toolCallId)
|
||||
) {
|
||||
handleDisplayChart(part.input.xml);
|
||||
processedToolCalls.current.add(toolCallId);
|
||||
}
|
||||
}
|
||||
@@ -80,9 +92,9 @@ export function ChatMessageDisplay({
|
||||
});
|
||||
}, [messages, handleDisplayChart]);
|
||||
|
||||
const renderToolInvocation = (toolInvocation: any) => {
|
||||
const callId = toolInvocation.toolCallId;
|
||||
const { toolName, args, state } = toolInvocation;
|
||||
const renderToolPart = (part: any) => {
|
||||
const callId = part.toolCallId;
|
||||
const { state, input } = part;
|
||||
const isExpanded = expandedTools[callId] ?? true;
|
||||
|
||||
const toggleExpanded = () => {
|
||||
@@ -100,7 +112,7 @@ export function ChatMessageDisplay({
|
||||
<div className="flex flex-col gap-2">
|
||||
<div className="flex items-center justify-between">
|
||||
<div className="text-xs">Tool: display_diagram</div>
|
||||
{args && Object.keys(args).length > 0 && (
|
||||
{input && Object.keys(input).length > 0 && (
|
||||
<button
|
||||
onClick={toggleExpanded}
|
||||
className="text-xs text-gray-500 hover:text-gray-700"
|
||||
@@ -109,20 +121,24 @@ export function ChatMessageDisplay({
|
||||
</button>
|
||||
)}
|
||||
</div>
|
||||
{args && isExpanded && (
|
||||
{input && isExpanded && (
|
||||
<div className="mt-1 font-mono text-xs overflow-hidden">
|
||||
{typeof args === "object" &&
|
||||
Object.keys(args).length > 0 &&
|
||||
`Args: ${JSON.stringify(args, null, 2)}`}
|
||||
{typeof input === "object" &&
|
||||
Object.keys(input).length > 0 &&
|
||||
`Input: ${JSON.stringify(input, null, 2)}`}
|
||||
</div>
|
||||
)}
|
||||
<div className="mt-2 text-sm">
|
||||
{state === "partial-call" ? (
|
||||
{state === "input-streaming" ? (
|
||||
<div className="h-4 w-4 border-2 border-primary border-t-transparent rounded-full animate-spin" />
|
||||
) : state === "result" ? (
|
||||
) : state === "output-available" ? (
|
||||
<div className="text-green-600">
|
||||
Diagram generated
|
||||
</div>
|
||||
) : state === "output-error" ? (
|
||||
<div className="text-red-600">
|
||||
Error generating diagram
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
</div>
|
||||
@@ -149,66 +165,35 @@ export function ChatMessageDisplay({
|
||||
: "bg-muted text-muted-foreground"
|
||||
}`}
|
||||
>
|
||||
{message.parts
|
||||
? message.parts.map((part, index) => {
|
||||
switch (part.type) {
|
||||
case "text":
|
||||
return (
|
||||
<div key={index}>
|
||||
{part.text}
|
||||
</div>
|
||||
);
|
||||
case "tool-invocation":
|
||||
return renderToolInvocation(
|
||||
part.toolInvocation
|
||||
);
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
})
|
||||
: message.content}
|
||||
{message.parts?.map((part: any, index: number) => {
|
||||
switch (part.type) {
|
||||
case "text":
|
||||
return (
|
||||
<div key={index}>{part.text}</div>
|
||||
);
|
||||
case "file":
|
||||
return (
|
||||
<div key={index} className="mt-2">
|
||||
<Image
|
||||
src={part.url}
|
||||
width={200}
|
||||
height={200}
|
||||
alt={`file-${index}`}
|
||||
className="rounded-md border"
|
||||
style={{
|
||||
objectFit: "contain",
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
default:
|
||||
if (part.type?.startsWith("tool-")) {
|
||||
return renderToolPart(part);
|
||||
}
|
||||
return null;
|
||||
}
|
||||
})}
|
||||
</div>
|
||||
|
||||
{message?.experimental_attachments
|
||||
?.filter((attachment) =>
|
||||
attachment?.contentType?.startsWith("image/")
|
||||
)
|
||||
.map((attachment, index) => (
|
||||
<div
|
||||
key={`${message.id}-${index}`}
|
||||
className={`mt-2 ${
|
||||
message.role === "user"
|
||||
? "text-right"
|
||||
: "text-left"
|
||||
}`}
|
||||
>
|
||||
<div className="inline-block">
|
||||
<Image
|
||||
src={attachment.url}
|
||||
width={200}
|
||||
height={200}
|
||||
alt={
|
||||
attachment.name ??
|
||||
`attachment-${index}`
|
||||
}
|
||||
className="rounded-md border"
|
||||
style={{
|
||||
objectFit: "contain",
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
))}
|
||||
|
||||
{(message as any).function_call && (
|
||||
<div className="mt-2 text-left">
|
||||
<div className="text-xs text-gray-500">
|
||||
Using tool:{" "}
|
||||
{(message as any).function_call.name}
|
||||
...
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
))
|
||||
)}
|
||||
|
||||
@@ -12,6 +12,7 @@ import {
|
||||
CardTitle,
|
||||
} from "@/components/ui/card";
|
||||
import { useChat } from "@ai-sdk/react";
|
||||
import { DefaultChatTransport } from "ai";
|
||||
import { ChatInput } from "@/components/chat-input";
|
||||
import { ChatMessageDisplay } from "./chat-message-display";
|
||||
import { useDiagram } from "@/contexts/diagram-context";
|
||||
@@ -48,30 +49,33 @@ export default function ChatPanel() {
|
||||
return dt.files;
|
||||
};
|
||||
|
||||
// Add state for input management
|
||||
const [input, setInput] = useState("");
|
||||
|
||||
// Remove the currentXmlRef and related useEffect
|
||||
const {
|
||||
messages,
|
||||
input,
|
||||
handleInputChange,
|
||||
handleSubmit,
|
||||
status,
|
||||
error,
|
||||
setInput,
|
||||
setMessages,
|
||||
} = useChat({
|
||||
maxSteps: 5,
|
||||
async onToolCall({ toolCall }) {
|
||||
if (toolCall.toolName === "display_diagram") {
|
||||
const { xml } = toolCall.args as { xml: string };
|
||||
// do nothing because we will handle this streamingly in the ChatMessageDisplay component
|
||||
// onDisplayChart(replaceNodes(chartXML, xml));
|
||||
return "Successfully displayed the flowchart.";
|
||||
}
|
||||
},
|
||||
onError: (error) => {
|
||||
console.error("Chat error:", error);
|
||||
},
|
||||
});
|
||||
const { messages, sendMessage, addToolResult, status, error, setMessages } =
|
||||
useChat({
|
||||
transport: new DefaultChatTransport({
|
||||
api: "/api/chat",
|
||||
}),
|
||||
async onToolCall({ toolCall }) {
|
||||
if (toolCall.toolName === "display_diagram") {
|
||||
const { xml } = toolCall.input as { xml: string };
|
||||
// do nothing because we will handle this streamingly in the ChatMessageDisplay component
|
||||
// onDisplayChart(replaceNodes(chartXML, xml));
|
||||
|
||||
// Use addToolResult instead of returning a value
|
||||
addToolResult({
|
||||
tool: "display_diagram",
|
||||
toolCallId: toolCall.toolCallId,
|
||||
output: "Successfully displayed the flowchart.",
|
||||
});
|
||||
}
|
||||
},
|
||||
onError: (error) => {
|
||||
console.error("Chat error:", error);
|
||||
},
|
||||
});
|
||||
const messagesEndRef = useRef<HTMLDivElement>(null);
|
||||
// Scroll to bottom when messages change
|
||||
useEffect(() => {
|
||||
@@ -80,21 +84,47 @@ export default function ChatPanel() {
|
||||
}
|
||||
}, [messages]);
|
||||
|
||||
console.log(JSON.stringify(messages, null, 2));
|
||||
|
||||
const onFormSubmit = async (e: React.FormEvent<HTMLFormElement>) => {
|
||||
e.preventDefault();
|
||||
if (input.trim() && status !== "streaming") {
|
||||
try {
|
||||
// Fetch chart data before setting input
|
||||
// Fetch chart data before sending message
|
||||
const chartXml = await onFetchChart();
|
||||
handleSubmit(e, {
|
||||
data: {
|
||||
xml: chartXml,
|
||||
},
|
||||
experimental_attachments:
|
||||
files.length > 0 ? createFileList(files) : undefined,
|
||||
});
|
||||
|
||||
// Clear files after submission
|
||||
// Create message parts
|
||||
const parts: any[] = [{ type: "text", text: input }];
|
||||
|
||||
// Add file parts if files exist
|
||||
if (files.length > 0) {
|
||||
for (const file of files) {
|
||||
const reader = new FileReader();
|
||||
const dataUrl = await new Promise<string>((resolve) => {
|
||||
reader.onload = () =>
|
||||
resolve(reader.result as string);
|
||||
reader.readAsDataURL(file);
|
||||
});
|
||||
|
||||
parts.push({
|
||||
type: "file",
|
||||
url: dataUrl,
|
||||
mediaType: file.type,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
sendMessage(
|
||||
{ parts },
|
||||
{
|
||||
body: {
|
||||
xml: chartXml,
|
||||
},
|
||||
}
|
||||
);
|
||||
|
||||
// Clear input and files after submission
|
||||
setInput("");
|
||||
setFiles([]);
|
||||
} catch (error) {
|
||||
console.error("Error fetching chart data:", error);
|
||||
@@ -102,6 +132,13 @@ export default function ChatPanel() {
|
||||
}
|
||||
};
|
||||
|
||||
// Handle input change
|
||||
const handleInputChange = (
|
||||
e: React.ChangeEvent<HTMLInputElement | HTMLTextAreaElement>
|
||||
) => {
|
||||
setInput(e.target.value);
|
||||
};
|
||||
|
||||
// Helper function to handle file changes
|
||||
const handleFileChange = (newFiles: File[]) => {
|
||||
setFiles(newFiles);
|
||||
|
||||
Reference in New Issue
Block a user