Merge pull request #1879 from langchain-ai/dqbd/sdk-js-subgraphs

feat(sdk-js): add support for subgraphs, match schema

libs/sdk-js/.gitignore

@@ -1,4 +1,6 @@
 /docs
+/*.tgz
+/*.tar
 ## GENERATED create-entrypoints.js
 /client.cjs
 /client.js

libs/sdk-js/.prettierrc

@@ -0,0 +1 @@
+{}

libs/sdk-js/src/client.ts

@@ -10,6 +10,8 @@ import {
   ThreadState,
   Cron,
   AssistantVersion,
+  Subgraphs,
+  Checkpoint,
 } from "./schema.js";
 import { AsyncCaller, AsyncCallerParams } from "./utils/async_caller.js";
 import {
@@ -221,8 +223,13 @@ export class AssistantsClient extends BaseClient {
    * @param assistantId The ID of the assistant.
    * @returns Serialized graph
    */
-  async getGraph(assistantId: string): Promise<AssistantGraph> {
-    return this.fetch<AssistantGraph>(`/assistants/${assistantId}/graph`);
+  async getGraph(
+    assistantId: string,
+    options?: { xray?: boolean },
+  ): Promise<AssistantGraph> {
+    return this.fetch<AssistantGraph>(`/assistants/${assistantId}/graph`, {
+      params: { xray: options?.xray },
+    });
   }
 
   /**
@@ -234,6 +241,31 @@ export class AssistantsClient extends BaseClient {
     return this.fetch<GraphSchema>(`/assistants/${assistantId}/schemas`);
   }
 
+  /**
+   * Get the schemas of an assistant by ID.
+   *
+   * @param assistantId The ID of the assistant to get the schema of.
+   * @param options Additional options for getting subgraphs, such as namespace or recursion extraction.
+   * @returns The subgraphs of the assistant.
+   */
+  async getSubgraphs(
+    assistantId: string,
+    options?: {
+      namespace?: string;
+      recurse?: boolean;
+    },
+  ): Promise<Subgraphs> {
+    if (options?.namespace) {
+      return this.fetch<Subgraphs>(
+        `/assistants/${assistantId}/subgraphs/${options.namespace}`,
+        { params: { recurse: options?.recurse } },
+      );
+    }
+    return this.fetch<Subgraphs>(`/assistants/${assistantId}/subgraphs`, {
+      params: { recurse: options?.recurse },
+    });
+  }
+
   /**
    * Create a new assistant.
    * @param payload Payload for creating an assistant.
@@ -479,13 +511,30 @@ export class ThreadsClient extends BaseClient {
    */
   async getState<ValuesType = DefaultValues>(
     threadId: string,
-    checkpointId?: string,
+    checkpoint?: Checkpoint | string,
+    options?: { subgraphs?: boolean },
   ): Promise<ThreadState<ValuesType>> {
-    return this.fetch<ThreadState<ValuesType>>(
-      checkpointId != null
-        ? `/threads/${threadId}/state/${checkpointId}`
-        : `/threads/${threadId}/state`,
-    );
+    if (checkpoint != null) {
+      if (typeof checkpoint !== "string") {
+        return this.fetch<ThreadState<ValuesType>>(
+          `/threads/${threadId}/state/checkpoint`,
+          {
+            method: "POST",
+            json: { checkpoint, subgraphs: options?.subgraphs },
+          },
+        );
+      }
+
+      // deprecated
+      return this.fetch<ThreadState<ValuesType>>(
+        `/threads/${threadId}/state/${checkpoint}`,
+        { params: { subgraphs: options?.subgraphs } },
+      );
+    }
+
+    return this.fetch<ThreadState<ValuesType>>(`/threads/${threadId}/state`, {
+      params: { subgraphs: options?.subgraphs },
+    });
   }
 
   /**
@@ -496,7 +545,12 @@ export class ThreadsClient extends BaseClient {
    */
   async updateState<ValuesType = DefaultValues>(
     threadId: string,
-    options: { values: ValuesType; checkpointId?: string; asNode?: string },
+    options: {
+      values: ValuesType;
+      checkpoint?: Checkpoint;
+      checkpointId?: string;
+      asNode?: string;
+    },
   ): Promise<Pick<Config, "configurable">> {
     return this.fetch<Pick<Config, "configurable">>(
       `/threads/${threadId}/state`,
@@ -505,6 +559,7 @@ export class ThreadsClient extends BaseClient {
         json: {
           values: options.values,
           checkpoint_id: options.checkpointId,
+          checkpoint: options.checkpoint,
           as_node: options?.asNode,
         },
       },
@@ -608,6 +663,7 @@ export class RunsClient extends BaseClient {
       config: payload?.config,
       metadata: payload?.metadata,
       stream_mode: payload?.streamMode,
+      stream_subgraphs: payload?.streamSubgraphs,
       feedback_keys: payload?.feedbackKeys,
       assistant_id: assistantId,
       interrupt_before: payload?.interruptBefore,

libs/sdk-js/src/schema.ts

@@ -75,6 +75,8 @@ export interface GraphSchema {
   config_schema?: JSONSchema7;
 }
 
+export type Subgraphs = Record<string, GraphSchema>;
+
 export type Metadata = Optional<Record<string, unknown>>;
 
 export interface AssistantBase {

libs/sdk-js/src/types.ts

@@ -98,6 +98,11 @@ export interface RunsStreamPayload extends RunsInvokePayload {
    */
   streamMode?: StreamMode | Array<StreamMode>;
 
+  /**
+   * Stream output from subgraphs. By default, streams only the top graph.
+   */
+  streamSubgraphs?: boolean;
+
   /**
    * Pass one or more feedbackKeys if you want to request short-lived signed URLs
    * for submitting feedback to LangSmith with this key for this run.

libs/sdk-py/langgraph_sdk/client.py

@@ -304,11 +304,14 @@ class AssistantsClient:
         """  # noqa: E501
         return await self.http.get(f"/assistants/{assistant_id}")
 
-    async def get_graph(self, assistant_id: str) -> dict[str, list[dict[str, Any]]]:
+    async def get_graph(
+        self, assistant_id: str, *, xray: bool = False
+    ) -> dict[str, list[dict[str, Any]]]:
         """Get the graph of an assistant by ID.
 
         Args:
             assistant_id: The ID of the assistant to get the graph of.
+            xray: Include graph representation of subgraphs.
 
         Returns:
             Graph: The graph information for the assistant in JSON format.
@@ -338,7 +341,9 @@ class AssistantsClient:
 
 
         """  # noqa: E501
-        return await self.http.get(f"/assistants/{assistant_id}/graph")
+        return await self.http.get(
+            f"/assistants/{assistant_id}/graph", params={"xray": xray}
+        )
 
     async def get_schemas(self, assistant_id: str) -> GraphSchema:
         """Get the schemas of an assistant by ID.
@@ -2068,11 +2073,14 @@ class SyncAssistantsClient:
         """  # noqa: E501
         return self.http.get(f"/assistants/{assistant_id}")
 
-    def get_graph(self, assistant_id: str) -> dict[str, list[dict[str, Any]]]:
+    def get_graph(
+        self, assistant_id: str, *, xray: bool = False
+    ) -> dict[str, list[dict[str, Any]]]:
         """Get the graph of an assistant by ID.
 
         Args:
             assistant_id: The ID of the assistant to get the graph of.
+            xray: Include graph representation of subgraphs.
 
         Returns:
             Graph: The graph information for the assistant in JSON format.
@@ -2102,7 +2110,7 @@ class SyncAssistantsClient:
 
 
         """  # noqa: E501
-        return self.http.get(f"/assistants/{assistant_id}/graph")
+        return self.http.get(f"/assistants/{assistant_id}/graph", params={"xray": xray})
 
     def get_schemas(self, assistant_id: str) -> GraphSchema:
         """Get the schemas of an assistant by ID.