feat(sdk-js): strongly-typed messages, state/update-type, stream mode (#3360)

StateType, UpdateType is set on the Client rather than on `RunsClient.stream` because of lack of partial type arguments application. This PR also describes the message serialization format emitted by LangGraph Server (which may change ie. converting `type` to `role`). Avoiding direct import of `@langchain/core` for the core LangGraph SDK client, thus these types were copied from `@langchain/core` (a script is used to aid with keeping track with core)
2026-08-28 10:49:56 +02:00 · 2025-02-11 08:56:02 -08:00
parent 1e8f097656 69311a4135
commit e9809ae9c1
6 changed files with 420 additions and 65 deletions
@@ -27,7 +27,7 @@ import {
  createParser,
 } from "./utils/eventsource-parser/index.js";
 import { IterableReadableStream } from "./utils/stream.js";
-import {
+import type {
  RunsCreatePayload,
  RunsStreamPayload,
  RunsWaitPayload,
@@ -38,6 +38,7 @@ import {
 import { mergeSignals } from "./utils/signals.js";
 import { getEnvironmentVariable } from "./utils/env.js";
 import { _getFetchImplementation } from "./singletons/fetch.js";
+import type { TypedAsyncGenerator, StreamMode } from "./types.stream.js";
 /**
 * Get the API key from the environment.
 * Precedence:
@@ -164,8 +165,7 @@ class BaseClient {
      signal?: AbortSignal;
    },
  ): Promise<T> {
-    const response = await this.asyncCaller.call(
-      _getFetchImplementation(),
+    const response = await this.asyncCaller.fetch(
      ...this.prepareFetchOptions(path, options),
    );
    if (response.status === 202 || response.status === 204) {
@@ -457,15 +457,20 @@ export class AssistantsClient extends BaseClient {
  }
 }

-export class ThreadsClient extends BaseClient {
+export class ThreadsClient<
+  TStateType = DefaultValues,
+  TUpdateType = TStateType,
+> extends BaseClient {
  /**
   * Get a thread by ID.
   *
   * @param threadId ID of the thread.
   * @returns The thread.
   */
-  async get(threadId: string): Promise<Thread> {
-    return this.fetch<Thread>(`/threads/${threadId}`);
+  async get<ValuesType = TStateType>(
+    threadId: string,
+  ): Promise<Thread<ValuesType>> {
+    return this.fetch<Thread<ValuesType>>(`/threads/${threadId}`);
  }

  /**
@@ -481,8 +486,8 @@ export class ThreadsClient extends BaseClient {
    metadata?: Metadata;
    threadId?: string;
    ifExists?: OnConflictBehavior;
-  }): Promise<Thread> {
-    return this.fetch<Thread>(`/threads`, {
+  }): Promise<Thread<TStateType>> {
+    return this.fetch<Thread<TStateType>>(`/threads`, {
      method: "POST",
      json: {
        metadata: payload?.metadata,
@@ -497,8 +502,8 @@ export class ThreadsClient extends BaseClient {
   * @param threadId ID of the thread to be copied
   * @returns Newly copied thread
   */
-  async copy(threadId: string): Promise<Thread> {
-    return this.fetch<Thread>(`/threads/${threadId}/copy`, {
+  async copy(threadId: string): Promise<Thread<TStateType>> {
+    return this.fetch<Thread<TStateType>>(`/threads/${threadId}/copy`, {
      method: "POST",
    });
  }
@@ -542,7 +547,7 @@ export class ThreadsClient extends BaseClient {
   * @param query Query options
   * @returns List of threads
   */
-  async search(query?: {
+  async search<ValuesType = TStateType>(query?: {
    /**
     * Metadata to filter threads by.
     */
@@ -561,8 +566,8 @@ export class ThreadsClient extends BaseClient {
     * Must be one of 'idle', 'busy', 'interrupted' or 'error'.
     */
    status?: ThreadStatus;
-  }): Promise<Thread[]> {
-    return this.fetch<Thread[]>("/threads/search", {
+  }): Promise<Thread<ValuesType>[]> {
+    return this.fetch<Thread<ValuesType>[]>("/threads/search", {
      method: "POST",
      json: {
        metadata: query?.metadata ?? undefined,
@@ -579,7 +584,7 @@ export class ThreadsClient extends BaseClient {
   * @param threadId ID of the thread.
   * @returns Thread state.
   */
-  async getState<ValuesType = DefaultValues>(
+  async getState<ValuesType = TStateType>(
    threadId: string,
    checkpoint?: Checkpoint | string,
    options?: { subgraphs?: boolean },
@@ -613,7 +618,7 @@ export class ThreadsClient extends BaseClient {
   * @param threadId The ID of the thread.
   * @returns
   */
-  async updateState<ValuesType = DefaultValues>(
+  async updateState<ValuesType = TUpdateType>(
    threadId: string,
    options: {
      values: ValuesType;
@@ -672,7 +677,7 @@ export class ThreadsClient extends BaseClient {
   * @param options Additional options.
   * @returns List of thread states.
   */
-  async getHistory<ValuesType = DefaultValues>(
+  async getHistory<ValuesType = TStateType>(
    threadId: string,
    options?: {
      limit?: number;
@@ -696,24 +701,43 @@ export class ThreadsClient extends BaseClient {
  }
 }

-export class RunsClient extends BaseClient {
-  stream(
+export class RunsClient<
+  TStateType = DefaultValues,
+  TUpdateType = TStateType,
+  TCustomEventType = unknown,
+> extends BaseClient {
+  stream<
+    TStreamMode extends StreamMode | StreamMode[] = StreamMode,
+    TSubgraphs extends boolean = false,
+  >(
    threadId: null,
    assistantId: string,
-    payload?: Omit<RunsStreamPayload, "multitaskStrategy" | "onCompletion">,
-  ): AsyncGenerator<{
-    event: StreamEvent;
-    data: any;
-  }>;
+    payload?: Omit<
+      RunsStreamPayload<TStreamMode, TSubgraphs>,
+      "multitaskStrategy" | "onCompletion"
+    >,
+  ): TypedAsyncGenerator<
+    TStreamMode,
+    TSubgraphs,
+    TStateType,
+    TUpdateType,
+    TCustomEventType
+  >;

-  stream(
+  stream<
+    TStreamMode extends StreamMode | StreamMode[] = StreamMode,
+    TSubgraphs extends boolean = false,
+  >(
    threadId: string,
    assistantId: string,
-    payload?: RunsStreamPayload,
-  ): AsyncGenerator<{
-    event: StreamEvent;
-    data: any;
-  }>;
+    payload?: RunsStreamPayload<TStreamMode, TSubgraphs>,
+  ): TypedAsyncGenerator<
+    TStreamMode,
+    TSubgraphs,
+    TStateType,
+    TUpdateType,
+    TCustomEventType
+  >;

  /**
   * Create a run and stream the results.
@@ -722,14 +746,20 @@ export class RunsClient extends BaseClient {
   * @param assistantId Assistant ID to use for this run.
   * @param payload Payload for creating a run.
   */
-  async *stream(
+  async *stream<
+    TStreamMode extends StreamMode | StreamMode[] = StreamMode,
+    TSubgraphs extends boolean = false,
+  >(
    threadId: string | null,
    assistantId: string,
-    payload?: RunsStreamPayload,
-  ): AsyncGenerator<{
-    event: StreamEvent;
-    data: any;
-  }> {
+    payload?: RunsStreamPayload<TStreamMode, TSubgraphs>,
+  ): TypedAsyncGenerator<
+    TStreamMode,
+    TSubgraphs,
+    TStateType,
+    TUpdateType,
+    TCustomEventType
+  > {
    const json: Record<string, any> = {
      input: payload?.input,
      command: payload?.command,
@@ -753,8 +783,7 @@ export class RunsClient extends BaseClient {

    const endpoint =
      threadId == null ? `/runs/stream` : `/threads/${threadId}/runs/stream`;
-    const response = await this.asyncCaller.call(
-      _getFetchImplementation(),
+    const response = await this.asyncCaller.fetch(
      ...this.prepareFetchOptions(endpoint, {
        method: "POST",
        json,
@@ -767,7 +796,7 @@ export class RunsClient extends BaseClient {
    let onEndEvent: () => void;
    const textDecoder = new TextDecoder();

-    const stream: ReadableStream<{ event: string; data: any }> = (
+    const stream: ReadableStream<{ event: any; data: any }> = (
      response.body || new ReadableStream({ start: (ctrl) => ctrl.close() })
    ).pipeThrough(
      new TransformStream({
@@ -1046,8 +1075,7 @@ export class RunsClient extends BaseClient {
        ? { signal: options }
        : options;

-    const response = await this.asyncCaller.call(
-      _getFetchImplementation(),
+    const response = await this.asyncCaller.fetch(
      ...this.prepareFetchOptions(`/threads/${threadId}/runs/${runId}/stream`, {
        method: "GET",
        timeoutMs: null,
@@ -1284,7 +1312,11 @@ export class StoreClient extends BaseClient {
  }
 }

-export class Client {
+export class Client<
+  TStateType = DefaultValues,
+  TUpdateType = TStateType,
+  TCustomEventType = unknown,
+> {
  /**
   * The client for interacting with assistants.
   */
@@ -1293,12 +1325,12 @@ export class Client {
  /**
   * The client for interacting with threads.
   */
-  public threads: ThreadsClient;
+  public threads: ThreadsClient<TStateType, TUpdateType>;

  /**
   * The client for interacting with runs.
   */
-  public runs: RunsClient;
+  public runs: RunsClient<TStateType, TUpdateType, TCustomEventType>;

  /**
   * The client for interacting with cron runs.
@@ -26,3 +26,25 @@ export type {
 export { overrideFetchImplementation } from "./singletons/fetch.js";

 export type { OnConflictBehavior, Command } from "./types.js";
+export type { StreamMode } from "./types.stream.js";
+export type {
+  ValuesStreamEvent,
+  MessagesTupleStreamEvent,
+  MetadataStreamEvent,
+  UpdatesStreamEvent,
+  CustomStreamEvent,
+  MessagesStreamEvent,
+  DebugStreamEvent,
+  EventsStreamEvent,
+  ErrorStreamEvent,
+  FeedbackStreamEvent,
+} from "./types.stream.js";
+export type {
+  Message,
+  HumanMessage,
+  AIMessage,
+  ToolMessage,
+  SystemMessage,
+  FunctionMessage,
+  RemoveMessage,
+} from "./types.messages.js";
@@ -0,0 +1,108 @@
+type ImageDetail = "auto" | "low" | "high";
+type MessageContentImageUrl = {
+  type: "image_url";
+  image_url: string | { url: string; detail?: ImageDetail | undefined };
+};
+
+type MessageContentText = { type: "text"; text: string };
+type MessageContentComplex = MessageContentText | MessageContentImageUrl;
+type MessageContent = string | MessageContentComplex[];
+
+type MessageAdditionalKwargs = {
+  [x: string]: unknown;
+
+  function_call?: { arguments: string; name: string } | undefined;
+  tool_calls?:
+    | {
+        id: string;
+        function: { arguments: string; name: string };
+        type: "function";
+        index?: number | undefined;
+      }[]
+    | undefined;
+};
+
+export type HumanMessage = {
+  type: "human";
+  id?: string | undefined;
+  content: MessageContent;
+};
+
+export type AIMessage = {
+  type: "ai";
+  id?: string | undefined;
+  content: MessageContent;
+  tool_calls?:
+    | {
+        name: string;
+        args: { [x: string]: { [x: string]: any } };
+        id?: string | undefined;
+        type?: "tool_call" | undefined;
+      }[]
+    | undefined;
+  invalid_tool_calls?:
+    | {
+        name?: string | undefined;
+        args?: string | undefined;
+        id?: string | undefined;
+        error?: string | undefined;
+        type?: "invalid_tool_call" | undefined;
+      }[]
+    | undefined;
+  usage_metadata?:
+    | {
+        input_tokens: number;
+        output_tokens: number;
+        total_tokens: number;
+        input_token_details?:
+          | {
+              audio?: number | undefined;
+              cache_read?: number | undefined;
+              cache_creation?: number | undefined;
+            }
+          | undefined;
+        output_token_details?:
+          | { audio?: number | undefined; reasoning?: number | undefined }
+          | undefined;
+      }
+    | undefined;
+  additional_kwargs?: MessageAdditionalKwargs | undefined;
+  response_metadata?: Record<string, unknown> | undefined;
+};
+
+export type ToolMessage = {
+  type: "tool";
+  name?: string | undefined;
+  id?: string | undefined;
+  content: MessageContent;
+  status?: "error" | "success" | undefined;
+  tool_call_id: string;
+  additional_kwargs?: MessageAdditionalKwargs | undefined;
+  response_metadata?: Record<string, unknown> | undefined;
+};
+
+export type SystemMessage = {
+  type: "system";
+  id?: string | undefined;
+  content: MessageContent;
+};
+
+export type FunctionMessage = {
+  type: "function";
+  id?: string | undefined;
+  content: MessageContent;
+};
+
+export type RemoveMessage = {
+  type: "remove";
+  id: string;
+  content: MessageContent;
+};
+
+export type Message =
+  | HumanMessage
+  | AIMessage
+  | ToolMessage
+  | SystemMessage
+  | FunctionMessage
+  | RemoveMessage;
@@ -0,0 +1,204 @@
+import type { Message } from "./types.messages.js";
+
+/**
+ * Stream modes
+ * - "values": Stream only the state values.
+ * - "messages": Stream complete messages.
+ * - "messages-tuple": Stream (message chunk, metadata) tuples.
+ * - "updates": Stream updates to the state.
+ * - "events": Stream events occurring during execution.
+ * - "debug": Stream detailed debug information.
+ * - "custom": Stream custom events.
+ */
+export type StreamMode =
+  | "values"
+  | "messages"
+  | "updates"
+  | "events"
+  | "debug"
+  | "custom"
+  | "messages-tuple";
+
+type MessageTupleMetadata = {
+  tags: string[];
+  [key: string]: unknown;
+};
+
+type AsSubgraph<TEvent extends { event: string; data: unknown }> = {
+  event: TEvent["event"] | `${TEvent["event"]}|${string}`;
+  data: TEvent["data"];
+};
+
+/**
+ * Stream event with values after completion of each step.
+ */
+export type ValuesStreamEvent<StateType> = { event: "values"; data: StateType };
+
+/** @internal */
+export type SubgraphValuesStreamEvent<StateType> = AsSubgraph<
+  ValuesStreamEvent<StateType>
+>;
+
+/**
+ * Stream event with message chunks coming from LLM invocations inside nodes.
+ */
+export type MessagesTupleStreamEvent = {
+  event: "messages";
+  // TODO: add types for message and config, which do not depend on LangChain
+  // while making sure it's easy to keep them in sync.
+  data: [message: Message, config: MessageTupleMetadata];
+};
+
+/** @internal */
+export type SubgraphMessagesTupleStreamEvent =
+  AsSubgraph<MessagesTupleStreamEvent>;
+
+/**
+ * Metadata stream event with information about the run and thread
+ */
+export type MetadataStreamEvent = {
+  event: "metadata";
+  data: { run_id: string; thread_id: string };
+};
+
+/**
+ * Stream event with error information.
+ */
+export type ErrorStreamEvent = {
+  event: "error";
+  data: { error: string; message: string };
+};
+
+/** @internal */
+export type SubgraphErrorStreamEvent = AsSubgraph<ErrorStreamEvent>;
+
+/**
+ * Stream event with updates to the state after each step.
+ * The streamed outputs include the name of the node that
+ * produced the update as well as the update.
+ */
+export type UpdatesStreamEvent<UpdateType> = {
+  event: "updates";
+  data: { [node: string]: UpdateType };
+};
+
+/** @internal */
+export type SubgraphUpdatesStreamEvent<UpdateType> = AsSubgraph<
+  UpdatesStreamEvent<UpdateType>
+>;
+
+/**
+ * Streaming custom data from inside the nodes.
+ */
+export type CustomStreamEvent<T> = { event: "custom"; data: T };
+
+/** @internal */
+export type SubgraphCustomStreamEvent<T> = AsSubgraph<CustomStreamEvent<T>>;
+
+type MessagesMetadataStreamEvent = {
+  event: "messages/metadata";
+  data: { [messageId: string]: { metadata: unknown } };
+};
+type MessagesCompleteStreamEvent = {
+  event: "messages/complete";
+  data: Message[];
+};
+type MessagesPartialStreamEvent = {
+  event: "messages/partial";
+  data: Message[];
+};
+
+/**
+ * Message stream event specific to LangGraph Server.
+ * @deprecated Use `streamMode: "messages-tuple"` instead.
+ */
+export type MessagesStreamEvent =
+  | MessagesMetadataStreamEvent
+  | MessagesCompleteStreamEvent
+  | MessagesPartialStreamEvent;
+
+/** @internal */
+export type SubgraphMessagesStreamEvent =
+  | AsSubgraph<MessagesMetadataStreamEvent>
+  | AsSubgraph<MessagesCompleteStreamEvent>
+  | AsSubgraph<MessagesPartialStreamEvent>;
+
+/**
+ * Stream event with detailed debug information.
+ */
+export type DebugStreamEvent = { event: "debug"; data: unknown };
+
+/** @internal */
+export type SubgraphDebugStreamEvent = AsSubgraph<DebugStreamEvent>;
+
+/**
+ * Stream event with events occurring during execution.
+ */
+export type EventsStreamEvent = { event: "events"; data: unknown };
+
+/** @internal */
+export type SubgraphEventsStreamEvent = AsSubgraph<EventsStreamEvent>;
+
+/**
+ * Stream event with a feedback key to signed URL map. Set `feedbackKeys` in
+ * the `RunsStreamPayload` to receive this event.
+ */
+export type FeedbackStreamEvent = {
+  event: "feedback";
+  data: { [feedbackKey: string]: string };
+};
+
+type GetStreamModeMap<
+  TStreamMode extends StreamMode | StreamMode[],
+  TStateType = unknown,
+  TUpdateType = TStateType,
+  TCustomType = unknown,
+> =
+  | {
+      values: ValuesStreamEvent<TStateType>;
+      updates: UpdatesStreamEvent<TUpdateType>;
+      custom: CustomStreamEvent<TCustomType>;
+      debug: DebugStreamEvent;
+      messages: MessagesStreamEvent;
+      "messages-tuple": MessagesTupleStreamEvent;
+      events: EventsStreamEvent;
+    }[TStreamMode extends StreamMode[] ? TStreamMode[number] : TStreamMode]
+  | ErrorStreamEvent
+  | MetadataStreamEvent
+  | FeedbackStreamEvent;
+
+type GetSubgraphsStreamModeMap<
+  TStreamMode extends StreamMode | StreamMode[],
+  TStateType = unknown,
+  TUpdateType = TStateType,
+  TCustomType = unknown,
+> =
+  | {
+      values: SubgraphValuesStreamEvent<TStateType>;
+      updates: SubgraphUpdatesStreamEvent<TUpdateType>;
+      custom: SubgraphCustomStreamEvent<TCustomType>;
+      debug: SubgraphDebugStreamEvent;
+      messages: SubgraphMessagesStreamEvent;
+      "messages-tuple": SubgraphMessagesTupleStreamEvent;
+      events: SubgraphEventsStreamEvent;
+    }[TStreamMode extends StreamMode[] ? TStreamMode[number] : TStreamMode]
+  | SubgraphErrorStreamEvent
+  | MetadataStreamEvent
+  | FeedbackStreamEvent;
+
+export type TypedAsyncGenerator<
+  TStreamMode extends StreamMode | StreamMode[] = [],
+  TSubgraphs extends boolean = false,
+  TStateType = unknown,
+  TUpdateType = TStateType,
+  TCustomType = unknown,
+> = AsyncGenerator<
+  TSubgraphs extends true
+    ? GetSubgraphsStreamModeMap<
+        TStreamMode,
+        TStateType,
+        TUpdateType,
+        TCustomType
+      >
+    : GetStreamModeMap<TStreamMode, TStateType, TUpdateType, TCustomType>
+>;
@@ -1,23 +1,6 @@
 import { Checkpoint, Config, Metadata } from "./schema.js";
+import { StreamMode } from "./types.stream.js";

-/**
- * Stream modes
- * - "values": Stream only the state values.
- * - "messages": Stream complete messages.
- * - "messages-tuple": Stream (message chunk, metadata) tuples.
- * - "updates": Stream updates to the state.
- * - "events": Stream events occurring during execution.
- * - "debug": Stream detailed debug information.
- * - "custom": Stream custom events.
- */
-export type StreamMode =
-  | "values"
-  | "messages"
-  | "updates"
-  | "events"
-  | "debug"
-  | "custom"
-  | "messages-tuple";
 export type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
 export type OnConflictBehavior = "raise" | "do_nothing";
 export type OnCompletionBehavior = "complete" | "continue";
@@ -31,6 +14,7 @@ export type StreamEvent =
  | "messages/partial"
  | "messages/metadata"
  | "messages/complete"
+  | "messages"
  | (string & {});

 export interface Send {
@@ -148,16 +132,19 @@ interface RunsInvokePayload {
  command?: Command;
 }

-export interface RunsStreamPayload extends RunsInvokePayload {
+export interface RunsStreamPayload<
+  TStreamMode extends StreamMode | StreamMode[] = [],
+  TSubgraphs extends boolean = false,
+> extends RunsInvokePayload {
  /**
   * One of `"values"`, `"messages"`, `"messages-tuple"`, `"updates"`, `"events"`, `"debug"`, `"custom"`.
   */
-  streamMode?: StreamMode | Array<StreamMode>;
+  streamMode?: TStreamMode;

  /**
   * Stream output from subgraphs. By default, streams only the top graph.
   */
-  streamSubgraphs?: boolean;
+  streamSubgraphs?: TSubgraphs;

  /**
   * Pass one or more feedbackKeys if you want to request short-lived signed URLs
@@ -1,5 +1,6 @@
 import pRetry from "p-retry";
 import PQueueMod from "p-queue";
+import { _getFetchImplementation } from "../singletons/fetch.js";

 const STATUS_NO_RETRY = [
  400, // Bad Request
@@ -210,7 +211,8 @@ export class AsyncCaller {
  }

  fetch(...args: Parameters<typeof fetch>): ReturnType<typeof fetch> {
-    const fetchFn = this.customFetch ?? fetch;
+    const fetchFn =
+      this.customFetch ?? (_getFetchImplementation() as typeof fetch);
    return this.call(() =>
      fetchFn(...args).then((res) => (res.ok ? res : Promise.reject(res))),
    );