langgraph-api 0.0.1__py3-none-any.whl

langgraph_api/js/server_sent_events.py

@@ -0,0 +1,126 @@
+"""Adapted from https://github.com/florimondmanca/httpx-sse"""
+
+from collections.abc import AsyncIterator, Iterator
+from contextlib import asynccontextmanager
+from typing import Any, TypedDict
+
+import httpx
+
+
+class ServerSentEvent(TypedDict):
+    event: str | None
+    data: str | None
+    id: str | None
+    retry: int | None
+
+
+class SSEDecoder:
+    def __init__(self) -> None:
+        self._event = ""
+        self._data: list[str] = []
+        self._last_event_id = ""
+        self._retry: int | None = None
+
+    def decode(self, line: str) -> ServerSentEvent | None:
+        # See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation  # noqa: E501
+        if not line:
+            if (
+                not self._event
+                and not self._data
+                and not self._last_event_id
+                and self._retry is None
+            ):
+                return None
+
+            sse = {
+                "event": self._event,
+                "data": "\n".join(self._data),
+                "id": self._last_event_id,
+                "retry": self._retry,
+            }
+
+            # NOTE: as per the SSE spec, do not reset last_event_id.
+            self._event = ""
+            self._data = []
+            self._retry = None
+
+            return sse
+
+        if line.startswith(":"):
+            return None
+
+        fieldname, _, value = line.partition(":")
+
+        if value.startswith(" "):
+            value = value[1:]
+
+        if fieldname == "event":
+            self._event = value
+        elif fieldname == "data":
+            self._data.append(value)
+        elif fieldname == "id":
+            if "\0" in value:
+                pass
+            else:
+                self._last_event_id = value
+        elif fieldname == "retry":
+            try:
+                self._retry = int(value)
+            except (TypeError, ValueError):
+                pass
+        else:
+            pass  # Field is ignored.
+
+        return None
+
+
+class EventSource:
+    def __init__(self, response: httpx.Response) -> None:
+        self._response = response
+
+    def _check_content_type(self) -> None:
+        """Check that the response content type is 'text/event-stream'."""
+        self._response.raise_for_status()
+        content_type = self._response.headers.get("content-type", "").partition(";")[0]
+        if "text/event-stream" not in content_type:
+            raise AssertionError(
+                "Expected response header Content-Type to contain 'text/event-stream', "
+                f"got {content_type!r}"
+            )
+
+    @property
+    def response(self) -> httpx.Response:
+        return self._response
+
+    def iter_sse(self) -> Iterator[ServerSentEvent]:
+        self._check_content_type()
+        decoder = SSEDecoder()
+        for line in self._response.iter_lines():
+            line = line.rstrip("\n")
+            sse = decoder.decode(line)
+            if sse is not None:
+                yield sse
+
+    async def aiter_sse(self) -> AsyncIterator[ServerSentEvent]:
+        self._check_content_type()
+        decoder = SSEDecoder()
+        async for line in self._response.aiter_lines():
+            line = line.rstrip("\n")
+            sse = decoder.decode(line)
+            if sse is not None:
+                yield sse
+
+
+@asynccontextmanager
+async def aconnect_sse(
+    client: httpx.AsyncClient,
+    method: str,
+    url: str,
+    **kwargs: Any,
+) -> AsyncIterator[EventSource]:
+    headers = kwargs.pop("headers", {})
+    headers["Accept"] = "text/event-stream"
+    headers["Cache-Control"] = "no-store"
+
+    async with client.stream(method, url, headers=headers, **kwargs) as response:
+        yield EventSource(response)

langgraph_api/js/src/graph.mts

@@ -0,0 +1,88 @@
+import { Worker } from "node:worker_threads";
+import { register } from "node:module";
+import * as fs from "node:fs/promises";
+import type { CompiledGraph, Graph } from "@langchain/langgraph";
+import * as path from "node:path";
+import type { JSONSchema7 } from "json-schema";
+
+// enforce API @langchain/langgraph precedence
+register("./hooks.mjs", import.meta.url);
+
+export interface GraphSchema {
+  state: JSONSchema7 | undefined;
+  input: JSONSchema7 | undefined;
+  output: JSONSchema7 | undefined;
+  config: JSONSchema7 | undefined;
+}
+
+export interface GraphSpec {
+  sourceFile: string;
+  exportSymbol: string;
+}
+
+export async function resolveGraph(
+  spec: string,
+  options?: { onlyFilePresence?: boolean }
+) {
+  const [userFile, exportSymbol] = spec.split(":", 2);
+
+  // validate file exists
+  await fs.stat(userFile);
+  if (options?.onlyFilePresence) {
+    return { sourceFile: userFile, exportSymbol, resolved: undefined };
+  }
+
+  const sourceFile = path.resolve(process.cwd(), userFile);
+  type GraphLike = CompiledGraph<string> | Graph<string>;
+
+  type GraphUnknown =
+    | GraphLike
+    | Promise<GraphLike>
+    | (() => GraphLike | Promise<GraphLike>)
+    | undefined;
+
+  const isGraph = (graph: GraphLike): graph is Graph<string> => {
+    if (typeof graph !== "object" || graph == null) return false;
+    return "compile" in graph && typeof graph.compile === "function";
+  };
+
+  const graph: GraphUnknown = await import(sourceFile).then(
+    (module) => module[exportSymbol || "default"]
+  );
+
+  // obtain the graph, and if not compiled, compile it
+  const resolved: CompiledGraph<string> = await (async () => {
+    if (!graph) throw new Error("Failed to load graph: graph is nullush");
+    const graphLike = typeof graph === "function" ? await graph() : await graph;
+
+    if (isGraph(graphLike)) return graphLike.compile();
+    return graphLike;
+  })();
+
+  return { sourceFile, exportSymbol, resolved };
+}
+
+export async function runGraphSchemaWorker(spec: GraphSpec) {
+  const SCHEMA_RESOLVE_TIMEOUT_MS = 30_000;
+
+  return await new Promise<Record<string, GraphSchema>>((resolve, reject) => {
+    const worker = new Worker(
+      new URL("./parser/parser.worker.mjs", import.meta.url).pathname
+    );
+
+    // Set a timeout to reject if the worker takes too long
+    const timeoutId = setTimeout(() => {
+      worker.terminate();
+      reject(new Error("Schema extract worker timed out"));
+    }, SCHEMA_RESOLVE_TIMEOUT_MS);
+
+    worker.on("message", (result) => {
+      worker.terminate();
+      clearTimeout(timeoutId);
+      resolve(result);
+    });
+
+    worker.on("error", reject);
+    worker.postMessage(spec);
+  });
+}

langgraph_api/js/src/hooks.mjs

@@ -0,0 +1,12 @@
+// This hook is to ensure that @langchain/langgraph package
+// found in /api folder has precendence compared to user-provided package
+// found in /deps. Does not attempt to semver check for too old packages.
+export async function resolve(specifier, context, nextResolve) {
+  const parentURL = new URL("./graph.mts", import.meta.url).toString();
+
+  if (specifier.startsWith("@langchain/langgraph")) {
+    return nextResolve(specifier, { ...context, parentURL });
+  }
+
+  return nextResolve(specifier, context);
+}

langgraph_api/js/src/parser/parser.mts

@@ -0,0 +1,443 @@
+import * as ts from "typescript";
+import * as vfs from "@typescript/vfs";
+import * as path from "node:path";
+import dedent from "dedent";
+import { buildGenerator } from "../schema/types.mts";
+
+const __dirname = new URL(".", import.meta.url).pathname;
+
+const compilerOptions = {
+  noEmit: true,
+  strict: true,
+  allowUnusedLabels: true,
+};
+
+export class SubgraphExtractor {
+  protected program: ts.Program;
+  protected checker: ts.TypeChecker;
+  protected sourceFile: ts.SourceFile;
+  protected inferFile: ts.SourceFile;
+
+  protected anyPregelType: ts.Type;
+  protected anyGraphType: ts.Type;
+
+  protected strict: boolean;
+
+  constructor(
+    program: ts.Program,
+    sourceFile: ts.SourceFile,
+    inferFile: ts.SourceFile,
+    options?: { strict?: boolean }
+  ) {
+    this.program = program;
+    this.sourceFile = sourceFile;
+    this.inferFile = inferFile;
+
+    this.checker = program.getTypeChecker();
+    this.strict = options?.strict ?? false;
+
+    this.anyPregelType = this.findTypeByName("AnyPregel");
+    this.anyGraphType = this.findTypeByName("AnyGraph");
+  }
+
+  private findTypeByName = (needle: string) => {
+    let result: ts.Type | undefined = undefined;
+
+    const visit = (node: ts.Node) => {
+      if (ts.isTypeAliasDeclaration(node)) {
+        const symbol = (node as any).symbol;
+
+        if (symbol != null) {
+          const name = this.checker
+            .getFullyQualifiedName(symbol)
+            .replace(/".*"\./, "");
+          if (name === needle) result = this.checker.getTypeAtLocation(node);
+        }
+      }
+      if (result == null) ts.forEachChild(node, visit);
+    };
+
+    ts.forEachChild(this.inferFile, visit);
+    if (!result) throw new Error(`Failed to find "${needle}" type`);
+    return result;
+  };
+
+  private find = (
+    root: ts.Node,
+    predicate: (node: ts.Node) => boolean
+  ): ts.Node | undefined => {
+    let result: ts.Node | undefined = undefined;
+
+    const visit = (node: ts.Node) => {
+      if (predicate(node)) {
+        result = node;
+      } else {
+        ts.forEachChild(node, visit);
+      }
+    };
+
+    if (predicate(root)) return root;
+    ts.forEachChild(root, visit);
+    return result;
+  };
+
+  protected findSubgraphs = (
+    node: ts.Node,
+    namespace: string[] = []
+  ): {
+    node: string;
+    namespace: string[];
+    subgraph: { name: string; node: ts.Node };
+  }[] => {
+    const findAllAddNodeCalls = (
+      acc: {
+        node: string;
+        namespace: string[];
+        subgraph: { name: string; node: ts.Node };
+      }[],
+      node: ts.Node
+    ) => {
+      if (ts.isCallExpression(node)) {
+        const firstChild = node.getChildAt(0);
+
+        if (
+          ts.isPropertyAccessExpression(firstChild) &&
+          this.getText(firstChild.name) === "addNode"
+        ) {
+          let nodeName: string = "unknown";
+          let variables: { node: ts.Node; name: string }[] = [];
+
+          const [subgraphNode, callArg] = node.arguments;
+
+          if (subgraphNode && ts.isStringLiteralLike(subgraphNode)) {
+            nodeName = this.getText(subgraphNode);
+            if (
+              (nodeName.startsWith(`"`) && nodeName.endsWith(`"`)) ||
+              (nodeName.startsWith(`'`) && nodeName.endsWith(`'`))
+            ) {
+              nodeName = nodeName.slice(1, -1);
+            }
+          }
+
+          if (callArg) {
+            if (
+              ts.isFunctionLike(callArg) ||
+              ts.isCallLikeExpression(callArg)
+            ) {
+              variables = this.reduceChildren(
+                callArg,
+                this.findSubgraphIdentifiers,
+                []
+              );
+            } else if (ts.isIdentifier(callArg)) {
+              variables = this.findSubgraphIdentifiers([], callArg);
+            }
+          }
+
+          if (variables.length > 0) {
+            if (variables.length > 1) {
+              const targetName = [...namespace, nodeName].join("|");
+              const errMsg = `Multiple unique subgraph invocations found for "${targetName}"`;
+              if (this.strict) throw new Error(errMsg);
+              console.warn(errMsg);
+            }
+
+            acc.push({
+              namespace: namespace,
+              node: nodeName,
+              subgraph: variables[0],
+            });
+          }
+        }
+      }
+
+      return acc;
+    };
+
+    let subgraphs = this.reduceChildren(node, findAllAddNodeCalls, []);
+
+    // TODO: make this more strict, only traverse the flow graph only
+    // if no `addNode` calls were found
+    if (!subgraphs.length) {
+      // internal property, however relied upon by ts-ast-viewer et all
+      // so that we don't need to traverse the control flow ourselves
+      // https://github.com/microsoft/TypeScript/pull/58036
+      type InternalFlowNode = ts.Node & { flowNode?: { node: ts.Node } };
+      const candidate = this.find(
+        node,
+        (node: any) => node && "flowNode" in node && node.flowNode
+      ) as InternalFlowNode | undefined;
+
+      if (
+        candidate?.flowNode &&
+        this.isGraphOrPregelType(
+          this.checker.getTypeAtLocation(candidate.flowNode.node)
+        )
+      ) {
+        subgraphs = this.findSubgraphs(candidate.flowNode.node, namespace);
+      }
+    }
+
+    // handle recursive behaviour
+    if (subgraphs.length > 0) {
+      return [
+        ...subgraphs,
+        ...subgraphs.map(({ subgraph, node }) =>
+          this.findSubgraphs(subgraph.node, [...namespace, node])
+        ),
+      ].flat();
+    }
+
+    return subgraphs;
+  };
+
+  protected getSubgraphsVariables = (name: string) => {
+    const sourceSymbol = this.checker.getSymbolAtLocation(this.sourceFile)!;
+    const exports = this.checker.getExportsOfModule(sourceSymbol);
+
+    const targetExport = exports.find((item) => item.name === name);
+    if (!targetExport) throw new Error(`Failed to find export "${name}"`);
+    const varDecls = (targetExport.declarations ?? []).filter(
+      ts.isVariableDeclaration
+    );
+
+    return varDecls.flatMap((varDecl) => {
+      if (!varDecl.initializer) return [];
+      return this.findSubgraphs(varDecl.initializer, [name]);
+    });
+  };
+
+  public getAugmentedSourceFile = (
+    name: string
+  ): {
+    files: [filePath: string, contents: string][];
+    exports: { typeName: string; valueName: string; graphName: string }[];
+  } => {
+    const vars = this.getSubgraphsVariables(name);
+    type TypeExport = {
+      typeName: `__langgraph__${string}`;
+      valueName: string;
+      graphName: string;
+    };
+
+    const typeExports: TypeExport[] = [
+      { typeName: `__langgraph__${name}`, valueName: name, graphName: name },
+    ];
+
+    for (const { subgraph, node, namespace } of vars) {
+      typeExports.push({
+        typeName: `__langgraph__${namespace.join("_")}_${node}`,
+        valueName: subgraph.name,
+        graphName: [...namespace, node].join("|"),
+      });
+    }
+
+    const sourceFilePath = "__langgraph__source.mts";
+    const sourceContents = [
+      this.getText(this.sourceFile),
+      ...typeExports.map(
+        ({ typeName, valueName }) =>
+          `export type ${typeName} = typeof ${valueName}`
+      ),
+    ].join("\n\n");
+
+    const inferFilePath = "__langraph__infer.mts";
+    const inferContents = [
+      ...typeExports.map(
+        ({ typeName }) =>
+          `import type { ${typeName}} from "./__langgraph__source.mts"`
+      ),
+      this.inferFile.getText(this.inferFile),
+
+      ...typeExports.flatMap(({ typeName }) => {
+        return [
+          dedent`
+            type ${typeName}__reflect = Reflect<${typeName}>;
+            export type ${typeName}__state = Inspect<${typeName}__reflect["state"]>;
+            export type ${typeName}__update = Inspect<${typeName}__reflect["update"]>;
+
+            type ${typeName}__builder = BuilderReflect<${typeName}>;
+            export type ${typeName}__input = Inspect<FilterAny<${typeName}__builder["input"]>>;
+            export type ${typeName}__output = Inspect<FilterAny<${typeName}__builder["output"]>>;
+            export type ${typeName}__config = Inspect<FilterAny<${typeName}__builder["config"]>>;
+          `,
+        ];
+      }),
+    ].join("\n\n");
+
+    return {
+      files: [
+        [sourceFilePath, sourceContents],
+        [inferFilePath, inferContents],
+      ],
+      exports: typeExports,
+    };
+  };
+
+  protected findSubgraphIdentifiers = (
+    acc: { node: ts.Node; name: string }[],
+    node: ts.Node
+  ) => {
+    if (ts.isIdentifier(node)) {
+      const smb = this.checker.getSymbolAtLocation(node);
+
+      if (
+        smb?.valueDeclaration &&
+        ts.isVariableDeclaration(smb.valueDeclaration)
+      ) {
+        const target = smb.valueDeclaration;
+        const targetType = this.checker.getTypeAtLocation(target);
+
+        if (this.isGraphOrPregelType(targetType)) {
+          acc.push({ name: this.getText(target.name), node: target });
+        }
+      }
+
+      if (smb?.declarations) {
+        const target = smb.declarations.find(ts.isImportSpecifier);
+        if (target) {
+          const targetType = this.checker.getTypeAtLocation(target);
+          if (this.isGraphOrPregelType(targetType)) {
+            acc.push({ name: this.getText(target.name), node: target });
+          }
+        }
+      }
+    }
+
+    return acc;
+  };
+
+  protected isGraphOrPregelType = (type: ts.Type) => {
+    return (
+      this.checker.isTypeAssignableTo(type, this.anyPregelType) ||
+      this.checker.isTypeAssignableTo(type, this.anyGraphType)
+    );
+  };
+
+  protected getText(node: ts.Node) {
+    return node.getText(this.sourceFile);
+  }
+
+  protected reduceChildren<Acc>(
+    node: ts.Node,
+    fn: (acc: Acc, node: ts.Node) => Acc,
+    initial: Acc
+  ): Acc {
+    let acc = initial;
+    function it(node: ts.Node) {
+      acc = fn(acc, node);
+      ts.forEachChild(node, it.bind(this));
+    }
+
+    ts.forEachChild(node, it.bind(this));
+    return acc;
+  }
+
+  static extractSchemas(
+    target:
+      | string
+      | {
+          contents: string;
+          files?: [fileName: string, contents: string][];
+        },
+    name: string,
+    options?: { strict?: boolean }
+  ) {
+    const dirname =
+      typeof target === "string" ? path.dirname(target) : __dirname;
+
+    const fsMap = new Map<string, string>();
+    const system = vfs.createFSBackedSystem(fsMap, dirname, ts);
+    const host = vfs.createVirtualCompilerHost(system, compilerOptions, ts);
+
+    const targetPath =
+      typeof target === "string"
+        ? target
+        : path.resolve(dirname, "./__langgraph__target.mts");
+
+    const inferTemplatePath = path.resolve(
+      __dirname,
+      "../schema/types.template.mts"
+    );
+
+    if (typeof target !== "string") {
+      fsMap.set(targetPath, target.contents);
+      for (const [name, contents] of target.files ?? []) {
+        fsMap.set(path.resolve(dirname, name), contents);
+      }
+    }
+
+    host.compilerHost.resolveModuleNames = (moduleNames, containingFile) => {
+      const resolvedModules: (ts.ResolvedModule | undefined)[] = [];
+      for (const moduleName of moduleNames) {
+        let target = containingFile;
+        const relative = path.relative(dirname, containingFile);
+        if (
+          moduleName.startsWith("@langchain/langgraph") &&
+          relative &&
+          !relative.startsWith("..") &&
+          !path.isAbsolute(relative)
+        ) {
+          target = path.resolve(__dirname, relative);
+        }
+
+        resolvedModules.push(
+          ts.resolveModuleName(
+            moduleName,
+            target,
+            compilerOptions,
+            host.compilerHost
+          ).resolvedModule
+        );
+      }
+
+      return resolvedModules;
+    };
+
+    const research = ts.createProgram({
+      rootNames: [inferTemplatePath, targetPath],
+      options: compilerOptions,
+      host: host.compilerHost,
+    });
+
+    const extractor = new SubgraphExtractor(
+      research,
+      research.getSourceFile(targetPath)!,
+      research.getSourceFile(inferTemplatePath)!,
+      options
+    );
+
+    const { files, exports } = extractor.getAugmentedSourceFile(name);
+    for (const [name, source] of files) {
+      system.writeFile(path.resolve(dirname, name), source);
+    }
+
+    const extract = ts.createProgram({
+      rootNames: [path.resolve(dirname, "./__langraph__infer.mts")],
+      options: compilerOptions,
+      host: host.compilerHost,
+    });
+
+    const schemaGenerator = buildGenerator(extract);
+    const trySymbol = (schema: typeof schemaGenerator, symbol: string) => {
+      try {
+        return schema?.getSchemaForSymbol(symbol) ?? undefined;
+      } catch (e) {
+        console.warn(`Failed to obtain symbol "${symbol}":`, e?.message);
+      }
+      return undefined;
+    };
+
+    return Object.fromEntries(
+      exports.map(({ typeName, graphName }) => [
+        graphName,
+        {
+          input: trySymbol(schemaGenerator, `${typeName}__input`),
+          output: trySymbol(schemaGenerator, `${typeName}__output`),
+          state: trySymbol(schemaGenerator, `${typeName}__update`),
+          config: trySymbol(schemaGenerator, `${typeName}__config`),
+        },
+      ])
+    );
+  }
+}

langgraph_api/js/src/parser/parser.worker.mjs

@@ -0,0 +1,12 @@
+import { tsImport } from "tsx/esm/api";
+import { parentPort } from "node:worker_threads";
+
+parentPort?.on("message", async (payload) => {
+  const { SubgraphExtractor } = await tsImport("./parser.mts", import.meta.url);
+  const result = SubgraphExtractor.extractSchemas(
+    payload.sourceFile,
+    payload.exportSymbol,
+    { strict: false }
+  );
+  parentPort?.postMessage(result);
+});