95 lines
2.8 KiB
TypeScript
95 lines
2.8 KiB
TypeScript
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";
|
|
|
|
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 llm = new ChatGoogleGenerativeAI({
|
|
model: "gemini-2.0-flash",
|
|
temperature: 0,
|
|
}).bindTools([routerTool], { tool_choice: "router" });
|
|
|
|
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.`;
|
|
|
|
const recentHumanMessage = state.messages.findLast(
|
|
(m) => m.getType() === "human"
|
|
);
|
|
|
|
if (!recentHumanMessage) {
|
|
throw new Error("No human message found in state");
|
|
}
|
|
|
|
const response = await llm.invoke([
|
|
{ role: "system", content: prompt },
|
|
recentHumanMessage,
|
|
]);
|
|
|
|
const toolCall = response.tool_calls?.[0]?.args as
|
|
| z.infer<typeof routerSchema>
|
|
| undefined;
|
|
if (!toolCall) {
|
|
throw new Error("No tool call found in response");
|
|
}
|
|
|
|
return {
|
|
next: toolCall.route,
|
|
};
|
|
}
|
|
|
|
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");
|
|
})
|
|
.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();
|
|
graph.name = "Generative UI Agent";
|