kexsh/agent-chat-ui
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

split up

Browse Source
This commit is contained in:
bracesproul
2025-02-27 15:41:47 -08:00
parent c7b61071a1
commit 3f4aad48e6
14 changed files with 254 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

164
agent/agent.tsx
View File

@@ -1,100 +1,94 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
import {
StateGraph,
MessagesAnnotation,
START,
Annotation,
} from "@langchain/langgraph";
import { SystemMessage } from "@langchain/core/messages";
import { StateGraph, START, END } from "@langchain/langgraph";
import { ChatGoogleGenerativeAI } from "@langchain/google-genai";
import { z } from "zod";
import { GenerativeUIAnnotation, GenerativeUIState } from "./types";
import { stockbrokerGraph } from "./stockbroker";
import { ChatOpenAI } from "@langchain/openai";
import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
import { uiMessageReducer } from "@langchain/langgraph-sdk/react-ui/types";
import type ComponentMap from "./uis/index";
import { z, ZodTypeAny } from "zod";
const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
async function router(
state: GenerativeUIState,
): Promise<Partial<GenerativeUIState>> {
const routerDescription = `The route to take based on the user's input.
- stockbroker: can fetch the price of a ticker, purchase/sell a ticker, or get the user's portfolio
- weather: can fetch the current weather conditions for a location
- generalInput: handles all other cases where the above tools don't apply
`;
const routerSchema = z.object({
route: z
.enum(["stockbroker", "weather", "generalInput"])
.describe(routerDescription),
});
const routerTool = {
name: "router",
description: "A tool to route the user's query to the appropriate tool.",
schema: routerSchema,
};
const getStockPriceSchema = z.object({
ticker: z.string().describe("The ticker symbol of the company"),
});
const getPortfolioSchema = z.object({
get_portfolio: z.boolean().describe("Should be true."),
});
const llm = new ChatGoogleGenerativeAI({
model: "gemini-2.0-flash",
temperature: 0,
}).bindTools([routerTool], { tool_choice: "router" });
const STOCKBROKER_TOOLS = [
{
name: "get_stock_price",
description: "A tool to get the stock price of a company",
schema: getStockPriceSchema,
},
{
name: "get_portfolio",
description:
"A tool to get the user's portfolio details. Only call this tool if the user requests their portfolio details.",
schema: getPortfolioSchema,
},
];
const prompt = `You're a highly helpful AI assistant, tasked with routing the user's query to the appropriate tool.
You should analyze the user's input, and choose the appropriate tool to use.`;
interface ToolCall {
name: string;
args: Record<string, any>;
id?: string;
type?: "tool_call";
}
const recentHumanMessage = state.messages
.reverse()
.find((m) => m.getType() === "human");
function findToolCall<Name extends string>(name: Name) {
return <Args extends ZodTypeAny>(
x: ToolCall,
): x is { name: Name; args: z.infer<Args> } => x.name === name;
}
if (!recentHumanMessage) {
throw new Error("No human message found in state");
}
const builder = new StateGraph(
Annotation.Root({
messages: MessagesAnnotation.spec["messages"],
ui: Annotation({ default: () => [], reducer: uiMessageReducer }),
timestamp: Annotation<number>,
}),
)
.addNode("agent", async (state, config) => {
const ui = typedUi<typeof ComponentMap>(config);
const message = await llm
.bindTools(STOCKBROKER_TOOLS)
.invoke([
new SystemMessage(
"You are a stockbroker agent that uses tools to get the stock price of a company",
),
...state.messages,
const response = await llm.invoke([
{ role: "system", content: prompt },
recentHumanMessage,
]);
const stockbrokerToolCall = message.tool_calls?.find(
findToolCall("get_stock_price")<typeof getStockPriceSchema>,
);
const portfolioToolCall = message.tool_calls?.find(
findToolCall("get_portfolio")<typeof getStockPriceSchema>,
);
if (stockbrokerToolCall) {
const instruction = `The stock price of ${
stockbrokerToolCall.args.ticker
} is ${Math.random() * 100}`;
ui.write("stock-price", { instruction, logo: "hey" });
const toolCall = response.tool_calls?.[0]?.args as
| z.infer<typeof routerSchema>
| undefined;
if (!toolCall) {
throw new Error("No tool call found in response");
}
if (portfolioToolCall) {
ui.write("portfolio-view", {});
}
return {
next: toolCall.route,
};
}
return { messages: message, ui: ui.collect, timestamp: Date.now() };
function handleRoute(
state: GenerativeUIState,
): "stockbroker" | "weather" | "generalInput" {
return state.next;
}
async function handleGeneralInput(state: GenerativeUIState) {
const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
const response = await llm.invoke(state.messages);
return {
messages: [response],
};
}
const builder = new StateGraph(GenerativeUIAnnotation)
.addNode("router", router)
.addNode("stockbroker", stockbrokerGraph)
.addNode("weather", () => {
throw new Error("Weather not implemented");
})
.addEdge(START, "agent");
.addNode("generalInput", handleGeneralInput)
.addConditionalEdges("router", handleRoute, [
"stockbroker",
"weather",
"generalInput",
])
.addEdge(START, "router")
.addEdge("stockbroker", END)
.addEdge("weather", END)
.addEdge("generalInput", END);
export const graph = builder.compile();
// event handler of evetns ˇtypes)
// event handler for specific node -> handle node
// TODO:
// - Send run ID & additional metadata for the client to properly use messages (maybe we even have a config)
// - Store that run ID in messages
graph.name = "Generative UI Agent";

14
agent/find-tool-call.ts Normal file
View File

@@ -0,0 +1,14 @@
import { z, ZodTypeAny } from "zod";
interface ToolCall {
name: string;
args: Record<string, any>;
id?: string;
type?: "tool_call";
}
export function findToolCall<Name extends string>(name: Name) {
return <Args extends ZodTypeAny>(
x: ToolCall,
): x is { name: Name; args: z.infer<Args> } => x.name === name;
}

11
agent/stockbroker/index.tsx Normal file
View File

@@ -0,0 +1,11 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
import { StateGraph, START } from "@langchain/langgraph";
import { StockbrokerAnnotation } from "./types";
import { callTools } from "./nodes/tools";
const builder = new StateGraph(StockbrokerAnnotation)
.addNode("agent", callTools)
.addEdge(START, "agent");
export const stockbrokerGraph = builder.compile();
stockbrokerGraph.name = "Stockbroker";

84
agent/stockbroker/nodes/tools.tsx Normal file
View File

@@ -0,0 +1,84 @@
import { StockbrokerState } from "../types";
import { ToolMessage } from "@langchain/core/messages";
import { ChatOpenAI } from "@langchain/openai";
import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
import type ComponentMap from "../../uis/index";
import { z } from "zod";
import { LangGraphRunnableConfig } from "@langchain/langgraph";
import { findToolCall } from "../../find-tool-call";
const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
const getStockPriceSchema = z.object({
ticker: z.string().describe("The ticker symbol of the company"),
});
const getPortfolioSchema = z.object({
get_portfolio: z.boolean().describe("Should be true."),
});
const STOCKBROKER_TOOLS = [
{
name: "stock-price",
description: "A tool to get the stock price of a company",
schema: getStockPriceSchema,
},
{
name: "portfolio",
description:
"A tool to get the user's portfolio details. Only call this tool if the user requests their portfolio details.",
schema: getPortfolioSchema,
},
];
export async function callTools(
state: StockbrokerState,
config: LangGraphRunnableConfig,
): Promise<Partial<StockbrokerState>> {
const ui = typedUi<typeof ComponentMap>(config);
const message = await llm.bindTools(STOCKBROKER_TOOLS).invoke([
{
role: "system",
content:
"You are a stockbroker agent that uses tools to get the stock price of a company",
},
...state.messages,
]);
const stockbrokerToolCall = message.tool_calls?.find(
findToolCall("stock-price")<typeof getStockPriceSchema>,
);
const portfolioToolCall = message.tool_calls?.find(
findToolCall("portfolio")<typeof getStockPriceSchema>,
);
if (stockbrokerToolCall) {
const instruction = `The stock price of ${
stockbrokerToolCall.args.ticker
} is ${Math.random() * 100}`;
ui.write("stock-price", { instruction, logo: "hey" });
}
if (portfolioToolCall) {
ui.write("portfolio", {});
}
const toolMessages =
message.tool_calls?.map((tc) => {
return new ToolMessage({
name: tc.name,
tool_call_id: tc.id ?? "",
content: "Successfully handled tool call",
});
}) || [];
console.log("Returning", [message, ...toolMessages]);
return {
messages: [message, ...toolMessages],
// TODO: Fix the ui return type.
ui: ui.collect as any[],
timestamp: Date.now(),
};
}

11
agent/stockbroker/types.ts Normal file
View File

@@ -0,0 +1,11 @@
import { Annotation } from "@langchain/langgraph";
import { GenerativeUIAnnotation } from "../types";
export const StockbrokerAnnotation = Annotation.Root({
messages: GenerativeUIAnnotation.spec.messages,
ui: GenerativeUIAnnotation.spec.ui,
timestamp: GenerativeUIAnnotation.spec.timestamp,
next: Annotation<"stockbroker" | "weather">(),
});
export type StockbrokerState = typeof StockbrokerAnnotation.State;

11
agent/types.ts Normal file
View File

@@ -0,0 +1,11 @@
import { MessagesAnnotation, Annotation } from "@langchain/langgraph";
import { uiMessageReducer } from "@langchain/langgraph-sdk/react-ui/types";
export const GenerativeUIAnnotation = Annotation.Root({
messages: MessagesAnnotation.spec["messages"],
ui: Annotation({ default: () => [], reducer: uiMessageReducer }),
timestamp: Annotation<number>,
next: Annotation<"stockbroker" | "weather" | "generalInput">(),
});
export type GenerativeUIState = typeof GenerativeUIAnnotation.State;

2
agent/uis/index.tsx
View File

@@ -3,6 +3,6 @@ import PortfolioView from "./portfolio-view";
const ComponentMap = {
"stock-price": StockPrice,
"portfolio-view": PortfolioView,
portfolio: PortfolioView,
} as const;
export default ComponentMap;

2
langgraph.json
View File

@@ -3,7 +3,7 @@
"agent": "./agent/agent.tsx:graph"
},
"ui": {
"agent": "./agent/ui.tsx"
"agent": "./agent/uis/index.tsx"
},
"_INTERNAL_docker_tag": "20",
"env": ".env"

1
package.json
View File

@@ -15,6 +15,7 @@
"@assistant-ui/react": "^0.8.0",
"@assistant-ui/react-markdown": "^0.8.0",
"@langchain/core": "^0.3.41",
"@langchain/google-genai": "^0.1.10",
"@langchain/langgraph": "^0.2.49",
"@langchain/langgraph-api": "*",
"@langchain/langgraph-cli": "*",

29
pnpm-lock.yaml generated
View File

@@ -21,6 +21,9 @@ importers:
"@langchain/core":
specifier: ^0.3.41
version: 0.3.41(openai@4.85.4(zod@3.24.2))
"@langchain/google-genai":
specifier: ^0.1.10
version: 0.1.10(@langchain/core@0.3.41(openai@4.85.4(zod@3.24.2)))(zod@3.24.2)
"@langchain/langgraph":
specifier: ^0.2.49
version: 0.2.49(@langchain/core@0.3.41(openai@4.85.4(zod@3.24.2)))(react@19.0.0)
@@ -879,6 +882,13 @@ packages:
integrity: sha512-MDWhGtE+eHw5JW7lq4qhc5yRLS11ERl1c7Z6Xd0a58DozHES6EnNNwUWbMiG4J9Cgj053Bhk8zvlhFYKVhULwg==,
}
"@google/generative-ai@0.21.0":
resolution:
{
integrity: sha512-7XhUbtnlkSEZK15kN3t+tzIMxsbKm/dSkKBFalj+20NvPKe1kBY7mR2P7vuijEn+f06z5+A8bVGKO0v39cr6Wg==,
}
engines: { node: ">=18.0.0" }
"@hono/node-server@1.13.8":
resolution:
{
@@ -986,6 +996,15 @@ packages:
}
engines: { node: ">=18" }
"@langchain/google-genai@0.1.10":
resolution:
{
integrity: sha512-+0xFWvauNDNp8Nvhy5F5g8RbB5g4WWQSIxoPI4IQIUICBBT/kS/Omf1VJI6Loc0IH93m9ZSwYxRVCRu3qx51TQ==,
}
engines: { node: ">=18" }
peerDependencies:
"@langchain/core": ">=0.3.17 <0.4.0"
"@langchain/langgraph-api@http://localhost:3123/17/@langchain/langgraph-api":
resolution: { tarball: http://localhost:3123/17/@langchain/langgraph-api }
version: 0.0.10
@@ -5313,6 +5332,8 @@ snapshots:
"@floating-ui/utils@0.2.9": {}
"@google/generative-ai@0.21.0": {}
"@hono/node-server@1.13.8(hono@4.7.2)":
dependencies:
hono: 4.7.2
@@ -5382,6 +5403,14 @@ snapshots:
transitivePeerDependencies:
- openai
"@langchain/google-genai@0.1.10(@langchain/core@0.3.41(openai@4.85.4(zod@3.24.2)))(zod@3.24.2)":
dependencies:
"@google/generative-ai": 0.21.0
"@langchain/core": 0.3.41(openai@4.85.4(zod@3.24.2))
zod-to-json-schema: 3.24.3(zod@3.24.2)
transitivePeerDependencies:
- zod
"@langchain/langgraph-api@http://localhost:3123/17/@langchain/langgraph-api(@langchain/core@0.3.41(openai@4.85.4(zod@3.24.2)))(@langchain/langgraph-checkpoint@0.0.15(@langchain/core@0.3.41(openai@4.85.4(zod@3.24.2))))(@langchain/langgraph@0.2.49(@langchain/core@0.3.41(openai@4.85.4(zod@3.24.2)))(react@19.0.0))(openai@4.85.4(zod@3.24.2))(typescript@5.7.3)":
dependencies:
"@babel/code-frame": 7.26.2

2
src/App.tsx
View File

@@ -3,7 +3,7 @@ import { Thread } from "@/components/assistant-ui/thread";
function App() {
return (
<div className="h-full">
<div className="h-screen">
<Thread />
</div>
);

13
src/components/assistant-ui/thread.tsx
View File

@@ -218,6 +218,7 @@ function CustomComponent({
}) {
const meta = thread.getMessagesMetadata(message, idx);
const seenState = meta?.firstSeenState;
console.log("seenState", meta);
const customComponent = seenState?.values.ui
.slice()
.reverse()
@@ -226,9 +227,13 @@ function CustomComponent({
additional_kwargs.run_id === seenState.metadata?.run_id,
);
if (!customComponent) {
console.log("no custom component", message, meta);
return null;
}
return (
<div key={message.id}>
<pre>{JSON.stringify(message, null, 2)}</pre>
{customComponent && (
<LoadExternalComponent
assistantId="agent"
@@ -245,10 +250,12 @@ const AssistantMessage: FC = () => {
const assistantMsgs = useMessage((m) => {
const langchainMessage = getExternalStoreMessages<Message>(m);
return langchainMessage;
})?.[0];
});
const assistantMsg = assistantMsgs[0];
let threadMsgIdx: number | undefined = undefined;
const threadMsg = thread.messages.find((m, idx) => {
if (m.id === assistantMsgs?.id) {
if (m.id === assistantMsg?.id) {
threadMsgIdx = idx;
return true;
}

7
src/providers/Runtime.tsx
View File

@@ -21,14 +21,7 @@ export function RuntimeProvider({
const input = message.content[0].text;
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({

5
src/providers/Stream.tsx
View File

@@ -6,7 +6,6 @@ import type {
RemoveUIMessage,
} from "@langchain/langgraph-sdk/react-ui/types";
// Define the type for the context value
type StreamContextType = ReturnType<
typeof useStream<
{ messages: Message[]; ui: UIMessage[] },
@@ -18,10 +17,8 @@ type StreamContextType = ReturnType<
>
>;
// 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,
}) => {
@@ -37,8 +34,6 @@ export const StreamProvider: React.FC<{ children: ReactNode }> = ({
assistantId: "agent",
});
console.log("StreamProvider", streamValue);
return (
<StreamContext.Provider value={streamValue}>
{children}