feat: update types

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";
@@ -48,7 +48,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,
     timestamp: Date.now(),
   };
 }