From 16470a75c077e3c1bf52823d7cb42e4ab9f74886 Mon Sep 17 00:00:00 2001
From: bracesproul <braceasproul@gmail.com>
Date: Wed, 5 Mar 2025 12:09:57 -0800
Subject: [PATCH] fix some types

---
 agent/stockbroker/nodes/tools.tsx  |  2 +-
 agent/trip-planner/nodes/tools.tsx |  2 +-
 agent/trip-planner/types.ts        |  2 +-
 agent/types.ts                     | 11 +++++++++--
 4 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/agent/stockbroker/nodes/tools.tsx b/agent/stockbroker/nodes/tools.tsx
index d31a7dc..d434b06 100644
--- a/agent/stockbroker/nodes/tools.tsx
+++ b/agent/stockbroker/nodes/tools.tsx
@@ -65,7 +65,7 @@ export async function callTools(
 
   return {
     messages: [message],
-    ui: ui.collect,
+    ui: ui.collect as StockbrokerUpdate["ui"],
     timestamp: Date.now(),
   };
 }
diff --git a/agent/trip-planner/nodes/tools.tsx b/agent/trip-planner/nodes/tools.tsx
index e32e2b4..be4e8d3 100644
--- a/agent/trip-planner/nodes/tools.tsx
+++ b/agent/trip-planner/nodes/tools.tsx
@@ -111,7 +111,7 @@ export async function callTools(
 
   return {
     messages: [response],
-    ui: ui.collect,
+    ui: ui.collect as TripPlannerUpdate["ui"],
     timestamp: Date.now(),
   };
 }
diff --git a/agent/trip-planner/types.ts b/agent/trip-planner/types.ts
index 2fc66a4..391aa0d 100644
--- a/agent/trip-planner/types.ts
+++ b/agent/trip-planner/types.ts
@@ -16,4 +16,4 @@ export const TripPlannerAnnotation = Annotation.Root({
 });
 
 export type TripPlannerState = typeof TripPlannerAnnotation.State;
-export type TripPlannerUpdate = typeof TripPlannerAnnotation.Update;
\ No newline at end of file
+export type TripPlannerUpdate = typeof TripPlannerAnnotation.Update;
diff --git a/agent/types.ts b/agent/types.ts
index a5355bb..4f109dc 100644
--- a/agent/types.ts
+++ b/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/types";
 
 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">(),
 });