split up

agent/stockbroker/index.tsx

@@ -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";

agent/stockbroker/nodes/tools.tsx

@@ -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(),
+  };
+}

agent/stockbroker/types.ts

@@ -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;