improved agent
This commit is contained in:
bracesproul
@@ -1,93 +1,46 @@
|
||||
import { useState, ReactNode } from "react";
|
||||
import { ReactNode } from "react";
|
||||
import {
|
||||
useExternalStoreRuntime,
|
||||
ThreadMessageLike,
|
||||
AppendMessage,
|
||||
AssistantRuntimeProvider,
|
||||
} from "@assistant-ui/react";
|
||||
import { Message } from "@langchain/langgraph-sdk";
|
||||
|
||||
function langChainRoleToAssistantRole(role: Message["type"]): "system" | "assistant" | "user" {
|
||||
if (role === "ai") return "assistant";
|
||||
if (role === "system") return "system";
|
||||
if (["human", "tool", "function"].includes(role)) return "user";
|
||||
throw new Error(`Unknown role: ${role}`);
|
||||
}
|
||||
|
||||
function langChainContentToAssistantContent(content: Message["content"]): ThreadMessageLike["content"] {
|
||||
if (!content) return [];
|
||||
|
||||
if (typeof content === "string") return content;
|
||||
|
||||
if (typeof content === "object") {
|
||||
if ("text" in content) {
|
||||
return [{
|
||||
type: "text",
|
||||
text: content.text as string,
|
||||
}]
|
||||
}
|
||||
|
||||
if ("thinking" in content) {
|
||||
return [{
|
||||
type: "reasoning",
|
||||
text: content.thinking as string,
|
||||
}]
|
||||
}
|
||||
}
|
||||
|
||||
throw new Error(`Unknown content: ${content}`);
|
||||
}
|
||||
|
||||
const convertMessage = (message: Message): ThreadMessageLike => {
|
||||
return {
|
||||
role: langChainRoleToAssistantRole(message.type),
|
||||
content: langChainContentToAssistantContent(message.content),
|
||||
};
|
||||
};
|
||||
|
||||
async function sleep(ms: number) {
|
||||
return new Promise((resolve) => setTimeout(resolve, ms));
|
||||
}
|
||||
import { HumanMessage } from "@langchain/langgraph-sdk";
|
||||
import { useStreamContext } from "./Stream";
|
||||
import { convertLangChainMessages } from "./convert-messages";
|
||||
|
||||
export function RuntimeProvider({
|
||||
children,
|
||||
}: Readonly<{
|
||||
children: ReactNode;
|
||||
}>) {
|
||||
const [isRunning, setIsRunning] = useState(false);
|
||||
const [messages, setMessages] = useState<Message[]>([]);
|
||||
|
||||
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;
|
||||
setMessages((currentConversation) => [
|
||||
...currentConversation,
|
||||
{ type: "human", content: input },
|
||||
]);
|
||||
|
||||
setIsRunning(true);
|
||||
// CALL API HERE
|
||||
// const assistantMessage = await backendApi(input);
|
||||
await sleep(2000);
|
||||
setMessages((currentConversation) => [
|
||||
...currentConversation,
|
||||
{ type: "ai", content: [{ type: "text", text: "This is an assistant message" }] },
|
||||
]);
|
||||
setIsRunning(false);
|
||||
const humanMessage: HumanMessage = { type: "human", content: input };
|
||||
// TODO: I dont think I need to do this, since we're passing stream.messages into the state hook, and it should update when we call `submit`
|
||||
// setMessages((currentConversation) => [
|
||||
// ...currentConversation,
|
||||
// humanMessage,
|
||||
// ]);
|
||||
|
||||
stream.submit({ messages: [humanMessage] });
|
||||
console.log("Sent message", humanMessage);
|
||||
};
|
||||
|
||||
|
||||
const runtime = useExternalStoreRuntime({
|
||||
isRunning,
|
||||
messages,
|
||||
convertMessage,
|
||||
isRunning: stream.isLoading,
|
||||
messages: stream.messages,
|
||||
convertMessage: convertLangChainMessages,
|
||||
onNew,
|
||||
});
|
||||
|
||||
|
||||
return (
|
||||
<AssistantRuntimeProvider runtime={runtime}>
|
||||
{children}
|
||||
</AssistantRuntimeProvider>
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
58
src/providers/Stream.tsx
Normal file
58
src/providers/Stream.tsx
Normal file
@@ -0,0 +1,58 @@
|
||||
import React, { createContext, useContext, ReactNode } from "react";
|
||||
import { useStream } from "@langchain/langgraph-sdk/react";
|
||||
import type { Message } from "@langchain/langgraph-sdk";
|
||||
import type {
|
||||
UIMessage,
|
||||
RemoveUIMessage,
|
||||
} from "@langchain/langgraph-sdk/react-ui/types";
|
||||
|
||||
// Define the type for the context value
|
||||
type StreamContextType = ReturnType<
|
||||
typeof useStream<
|
||||
{ messages: Message[]; ui: UIMessage[] },
|
||||
{
|
||||
messages?: Message[] | Message | string;
|
||||
ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
|
||||
},
|
||||
UIMessage | RemoveUIMessage
|
||||
>
|
||||
>;
|
||||
|
||||
// Create the context with a default undefined value
|
||||
const StreamContext = createContext<StreamContextType | undefined>(undefined);
|
||||
|
||||
// Create a provider component
|
||||
export const StreamProvider: React.FC<{ children: ReactNode }> = ({
|
||||
children,
|
||||
}) => {
|
||||
const streamValue = useStream<
|
||||
{ messages: Message[]; ui: UIMessage[] },
|
||||
{
|
||||
messages?: Message[] | Message | string;
|
||||
ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
|
||||
},
|
||||
UIMessage | RemoveUIMessage
|
||||
>({
|
||||
apiUrl: "http://localhost:2024",
|
||||
assistantId: "agent",
|
||||
});
|
||||
|
||||
console.log("StreamProvider", streamValue);
|
||||
|
||||
return (
|
||||
<StreamContext.Provider value={streamValue}>
|
||||
{children}
|
||||
</StreamContext.Provider>
|
||||
);
|
||||
};
|
||||
|
||||
// Create a custom hook to use the context
|
||||
export const useStreamContext = (): StreamContextType => {
|
||||
const context = useContext(StreamContext);
|
||||
if (context === undefined) {
|
||||
throw new Error("useStreamContext must be used within a StreamProvider");
|
||||
}
|
||||
return context;
|
||||
};
|
||||
|
||||
export default StreamContext;
|
||||
140
src/providers/convert-messages.ts
Normal file
140
src/providers/convert-messages.ts
Normal file
@@ -0,0 +1,140 @@
|
||||
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 }],
|
||||
// ...(message.additional_kwargs
|
||||
// ? {
|
||||
// metadata: {
|
||||
// custom: {
|
||||
// ...message.additional_kwargs,
|
||||
// },
|
||||
// },
|
||||
// }
|
||||
// : {}),
|
||||
};
|
||||
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,
|
||||
},
|
||||
],
|
||||
// ...(message.additional_kwargs
|
||||
// ? {
|
||||
// metadata: {
|
||||
// custom: {
|
||||
// ...message.additional_kwargs,
|
||||
// },
|
||||
// },
|
||||
// }
|
||||
// : {}),
|
||||
};
|
||||
case "tool":
|
||||
const toolMsg = message as ToolMessage;
|
||||
return {
|
||||
role: "user",
|
||||
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)}`);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user