feat: drop assistant ui, use custom chat ui

This commit is contained in:
bracesproul
2025-03-03 12:31:27 -08:00
parent a75c710990
commit 3f3f50d5c5
20 changed files with 4553 additions and 2477 deletions

View File

@@ -1,78 +0,0 @@
import { ReactNode, useEffect } from "react";
import {
useExternalStoreRuntime,
AppendMessage,
AssistantRuntimeProvider,
} from "@assistant-ui/react";
import { HumanMessage, Message, ToolMessage } from "@langchain/langgraph-sdk";
import { useStreamContext } from "./Stream";
import { convertLangChainMessages } from "./convert-messages";
function ensureToolCallsHaveResponses(messages: Message[]): Message[] {
const newMessages: ToolMessage[] = [];
messages.forEach((message, index) => {
if (message.type !== "ai" || message.tool_calls?.length === 0) {
// If it's not an AI message, or it doesn't have tool calls, we can ignore.
return;
}
// If it has tool calls, ensure the message which follows this is a tool message
const followingMessage = messages[index + 1];
if (followingMessage && followingMessage.type === "tool") {
// Following message is a tool message, so we can ignore.
return;
}
// Since the following message is not a tool message, we must create a new tool message
newMessages.push(
...(message.tool_calls?.map((tc) => ({
type: "tool" as const,
tool_call_id: tc.id ?? "",
id: tc.id ?? "",
name: tc.name,
content: "Successfully handled tool call.",
})) ?? [])
);
});
return newMessages;
}
export function RuntimeProvider({
children,
}: Readonly<{
children: ReactNode;
}>) {
const stream = useStreamContext();
const onNew = async (message: AppendMessage) => {
if (message.content[0]?.type !== "text")
throw new Error("Only text messages are supported");
const input = message.content[0].text;
const humanMessage: HumanMessage = { type: "human", content: input };
const newMessages = [
...ensureToolCallsHaveResponses(stream.messages),
humanMessage,
];
console.log("Sending new messages", newMessages);
stream.submit({ messages: newMessages }, { streamMode: ["values"] });
};
useEffect(() => {
console.log("useEffect - stream.messages", stream.messages);
}, [stream.messages]);
const runtime = useExternalStoreRuntime({
isRunning: stream.isLoading,
messages: stream.messages,
convertMessage: convertLangChainMessages,
onNew,
});
return (
<AssistantRuntimeProvider runtime={runtime}>
{children}
</AssistantRuntimeProvider>
);
}

View File

@@ -13,7 +13,7 @@ const useTypedStream = useStream<
messages?: Message[] | Message | string;
ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
};
CustomType: UIMessage | RemoveUIMessage;
CustomUpdateType: UIMessage | RemoveUIMessage;
}
>;

View File

@@ -1,122 +0,0 @@
import { ThreadMessageLike, ToolCallContentPart } from "@assistant-ui/react";
import { Message, AIMessage, ToolMessage } from "@langchain/langgraph-sdk";
export const getMessageType = (message: Record<string, any>): string => {
if (Array.isArray(message.id)) {
const lastItem = message.id[message.id.length - 1];
if (lastItem.startsWith("HumanMessage")) {
return "human";
} else if (lastItem.startsWith("AIMessage")) {
return "ai";
} else if (lastItem.startsWith("ToolMessage")) {
return "tool";
} else if (
lastItem.startsWith("BaseMessage") ||
lastItem.startsWith("SystemMessage")
) {
return "system";
}
}
if ("getType" in message && typeof message.getType === "function") {
return message.getType();
} else if ("_getType" in message && typeof message._getType === "function") {
return message._getType();
} else if ("type" in message) {
return message.type as string;
} else {
console.error(message);
throw new Error("Unsupported message type");
}
};
function getMessageContentOrThrow(message: unknown): string {
if (typeof message !== "object" || message === null) {
return "";
}
const castMsg = message as Record<string, any>;
if (
typeof castMsg?.content !== "string" &&
(!Array.isArray(castMsg.content) || castMsg.content[0]?.type !== "text") &&
(!castMsg.kwargs ||
!castMsg.kwargs?.content ||
typeof castMsg.kwargs?.content !== "string")
) {
console.error(castMsg);
throw new Error("Only text messages are supported");
}
let content = "";
if (Array.isArray(castMsg.content) && castMsg.content[0]?.type === "text") {
content = castMsg.content[0].text;
} else if (typeof castMsg.content === "string") {
content = castMsg.content;
} else if (
castMsg?.kwargs &&
castMsg.kwargs?.content &&
typeof castMsg.kwargs?.content === "string"
) {
content = castMsg.kwargs.content;
}
return content;
}
export function convertLangChainMessages(message: Message): ThreadMessageLike {
const content = getMessageContentOrThrow(message);
switch (getMessageType(message)) {
case "system":
return {
role: "system",
id: message.id,
content: [{ type: "text", text: content }],
};
case "human":
return {
role: "user",
id: message.id,
content: [{ type: "text", text: content }],
};
case "ai":
const aiMsg = message as AIMessage;
const toolCallsContent: ToolCallContentPart[] = aiMsg.tool_calls?.length
? aiMsg.tool_calls.map((tc) => ({
type: "tool-call" as const,
toolCallId: tc.id ?? "",
toolName: tc.name,
args: tc.args,
argsText: JSON.stringify(tc.args),
}))
: [];
return {
role: "assistant",
id: message.id,
content: [
...toolCallsContent,
{
type: "text",
text: content,
},
],
};
case "tool":
const toolMsg = message as ToolMessage;
return {
role: "assistant",
content: [
{
type: "tool-call",
toolName: toolMsg.name ?? "ToolCall",
toolCallId: toolMsg.tool_call_id,
result: content,
},
],
};
default:
console.error(message);
throw new Error(`Unsupported message type: ${getMessageType(message)}`);
}
}