90 lines
2.6 KiB
TypeScript
90 lines
2.6 KiB
TypeScript
|
/* eslint-disable @typescript-eslint/no-explicit-any */
|
||
|
|
import {
|
||
|
|
StateGraph,
|
||
|
|
MessagesAnnotation,
|
||
|
|
START,
|
||
|
|
Annotation,
|
||
|
|
} from "@langchain/langgraph";
|
||
|
|
import { SystemMessage } from "@langchain/core/messages";
|
||
|
|
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 "./ui";
|
||
|
|
import { z, ZodTypeAny } from "zod";
|
||
|
|
|
||
|
|
// const llm = new ChatOllama({ model: "deepseek-r1" });
|
||
|
|
const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
|
||
|
|
|
||
|
|
interface ToolCall {
|
||
|
|
name: string;
|
||
|
|
args: Record<string, any>;
|
||
|
|
id?: string;
|
||
|
|
type?: "tool_call";
|
||
|
|
}
|
||
|
|
|
||
|
|
function findToolCall<Name extends string>(name: Name) {
|
||
|
|
return <Args extends ZodTypeAny>(
|
||
|
|
x: ToolCall
|
||
|
|
): x is { name: Name; args: z.infer<Args> } => x.name === name;
|
||
|
|
}
|
||
|
|
|
||
|
|
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 result = ui.interrupt("react-component", {
|
||
|
|
// instruction: "Hello world",
|
||
|
|
// });
|
||
|
|
|
||
|
|
// // throw new Error("Random error");
|
||
|
|
// // stream custom events
|
||
|
|
// for (let count = 0; count < 10; count++) config.writer?.({ count });
|
||
|
|
|
||
|
|
// How do I properly assign
|
||
|
|
const stockbrokerSchema = z.object({ company: z.string() });
|
||
|
|
const message = await llm
|
||
|
|
.bindTools([
|
||
|
|
{
|
||
|
|
name: "stockbroker",
|
||
|
|
description: "A tool to get the stock price of a company",
|
||
|
|
schema: stockbrokerSchema,
|
||
|
|
},
|
||
|
|
])
|
||
|
|
.invoke([
|
||
|
|
new SystemMessage(
|
||
|
|
"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("stockbroker")<typeof stockbrokerSchema>
|
||
|
|
);
|
||
|
|
|
||
|
|
if (stockbrokerToolCall) {
|
||
|
|
const instruction = `The stock price of ${
|
||
|
|
stockbrokerToolCall.args.company
|
||
|
|
} is ${Math.random() * 100}`;
|
||
|
|
|
||
|
|
ui.write("react-component", { instruction, logo: "hey" });
|
||
|
|
}
|
||
|
|
|
||
|
|
return { messages: message, ui: ui.collect, timestamp: Date.now() };
|
||
|
|
})
|
||
|
|
.addEdge(START, "agent");
|
||
|
|
|
||
|
|
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
|