Merge branch 'main' of https://github.com/langchain-ai/agent-ui-demo into brace/hookup-accommodations

agent/stockbroker/nodes/tools.tsx

@@ -1,4 +1,4 @@
-import { StockbrokerState } from "../types";
+import { StockbrokerState, StockbrokerUpdate } from "../types";
 import { ChatOpenAI } from "@langchain/openai";
 import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
 import type ComponentMap from "../../uis/index";
@@ -32,7 +32,7 @@ const STOCKBROKER_TOOLS = [
 export async function callTools(
   state: StockbrokerState,
   config: LangGraphRunnableConfig,
-): Promise<Partial<StockbrokerState>> {
+): Promise<StockbrokerUpdate> {
   const ui = typedUi<typeof ComponentMap>(config);
 
   const message = await llm.bindTools(STOCKBROKER_TOOLS).invoke([
@@ -65,8 +65,7 @@ export async function callTools(
 
   return {
     messages: [message],
-    // TODO: Fix the ui return type.
-    ui: ui.collect as any[],
+    ui: ui.collect as StockbrokerUpdate["ui"],
     timestamp: Date.now(),
   };
 }

agent/stockbroker/types.ts

@@ -8,3 +8,4 @@ export const StockbrokerAnnotation = Annotation.Root({
 });
 
 export type StockbrokerState = typeof StockbrokerAnnotation.State;
+export type StockbrokerUpdate = typeof StockbrokerAnnotation.Update;

agent/trip-planner/nodes/classify.ts

@@ -1,11 +1,11 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { TripPlannerState } from "../types";
+import { TripPlannerState, TripPlannerUpdate } from "../types";
 import { z } from "zod";
 import { formatMessages } from "agent/utils/format-messages";
 
 export async function classify(
   state: TripPlannerState,
-): Promise<Partial<TripPlannerState>> {
+): Promise<TripPlannerUpdate> {
   if (!state.tripDetails) {
     // Can not classify if tripDetails are undefined
     return {};

agent/trip-planner/nodes/extraction.tsx

@@ -1,5 +1,5 @@
 import { ChatOpenAI } from "@langchain/openai";
-import { TripDetails, TripPlannerState } from "../types";
+import { TripDetails, TripPlannerState, TripPlannerUpdate } from "../types";
 import { z } from "zod";
 import { formatMessages } from "agent/utils/format-messages";
 
@@ -43,7 +43,7 @@ function calculateDates(
 
 export async function extraction(
   state: TripPlannerState,
-): Promise<Partial<TripPlannerState>> {
+): Promise<TripPlannerUpdate> {
   const schema = z.object({
     location: z
       .string()

agent/trip-planner/nodes/tools.tsx

@@ -1,4 +1,4 @@
-import { TripPlannerState } from "../types";
+import { TripPlannerState, TripPlannerUpdate } from "../types";
 import { ChatOpenAI } from "@langchain/openai";
 import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
 import type ComponentMap from "../../uis/index";
@@ -49,7 +49,7 @@ const schema = z.object({
 export async function callTools(
   state: TripPlannerState,
   config: LangGraphRunnableConfig,
-): Promise<Partial<TripPlannerState>> {
+): Promise<TripPlannerUpdate> {
   if (!state.tripDetails) {
     throw new Error("No trip details found");
   }
@@ -111,8 +111,7 @@ export async function callTools(
 
   return {
     messages: [response],
-    // TODO: Fix the ui return type.
-    ui: ui.collect as any[],
+    ui: ui.collect as TripPlannerUpdate["ui"],
     timestamp: Date.now(),
   };
 }

agent/trip-planner/types.ts

@@ -16,3 +16,4 @@ export const TripPlannerAnnotation = Annotation.Root({
 });
 
 export type TripPlannerState = typeof TripPlannerAnnotation.State;
+export type TripPlannerUpdate = typeof TripPlannerAnnotation.Update;

agent/types.ts

@@ -1,9 +1,16 @@
 import { MessagesAnnotation, Annotation } from "@langchain/langgraph";
-import { uiMessageReducer } from "@langchain/langgraph-sdk/react-ui/types";
+import {
+  RemoveUIMessage,
+  UIMessage,
+  uiMessageReducer,
+} from "@langchain/langgraph-sdk/react-ui/server";
 
 export const GenerativeUIAnnotation = Annotation.Root({
   messages: MessagesAnnotation.spec["messages"],
-  ui: Annotation({ default: () => [], reducer: uiMessageReducer }),
+  ui: Annotation<
+    UIMessage[],
+    UIMessage | RemoveUIMessage | (UIMessage | RemoveUIMessage)[]
+  >({ default: () => [], reducer: uiMessageReducer }),
   timestamp: Annotation<number>,
   next: Annotation<"stockbroker" | "tripPlanner" | "generalInput">(),
 });

agent/uis/stockbroker/stock-price/index.tsx

@@ -1,12 +1,20 @@
 import "./index.css";
-import { useStream } from "@langchain/langgraph-sdk/react";
-import type { AIMessage, Message } from "@langchain/langgraph-sdk";
+import {
+  useStreamContext,
+  type UIMessage,
+} from "@langchain/langgraph-sdk/react-ui";
+import type { AIMessage, Message, ToolMessage } from "@langchain/langgraph-sdk";
 import { useState, useEffect, useCallback } from "react";
 
 export default function StockPrice(props: {
   instruction: string;
   logo: string;
 }) {
+  const thread = useStreamContext<
+    { messages: Message[]; ui: UIMessage[] },
+    { MetaType: { ui: UIMessage | undefined } }
+  >();
+
   const [quantity, setQuantity] = useState(1);
   const [stockData, setStockData] = useState({
     symbol: "AAPL",
@@ -51,23 +59,16 @@ export default function StockPrice(props: {
   const [limitPrice, setLimitPrice] = useState(stockData.price.toFixed(2));
   const [showOrderSuccess, setShowOrderSuccess] = useState(false);
 
-  // useStream should be able to be infered from context
-  const thread = useStream<{ messages: Message[] }>({
-    assistantId: "assistant_123",
-    apiUrl: "http://localhost:3123",
-  });
-
-  const messagesCopy = thread.messages;
-
-  const aiTool = messagesCopy
-    .slice()
-    .reverse()
-    .find(
-      (message): message is AIMessage =>
-        message.type === "ai" && !!message.tool_calls?.length,
-    );
+  const aiTool = thread.messages.findLast(
+    (message): message is AIMessage =>
+      message.type === "ai" && !!message.tool_calls?.length,
+  );
 
   const toolCallId = aiTool?.tool_calls?.[0]?.id;
+  const toolResponse = thread.messages.findLast(
+    (message): message is ToolMessage =>
+      message.type === "tool" && message.tool_call_id === toolCallId,
+  );
 
   // Simulated price history generation on component mount
   useEffect(() => {
@@ -248,6 +249,7 @@ export default function StockPrice(props: {
   const range = max - min || 1;
   const chartPath = generateChartPath();
 
+  if (toolResponse) return <div>Responded</div>;
   return (
     <div className="w-full max-w-md bg-white rounded-xl shadow-lg overflow-hidden border border-gray-200">
       <div