langgraph-api 0.0.1__py3-none-any.whl

langgraph_api/graph.py

@@ -0,0 +1,314 @@
+import asyncio
+import glob
+import importlib.util
+import inspect
+import json
+import os
+import sys
+from collections.abc import Callable
+from itertools import filterfalse
+from random import choice
+from typing import NamedTuple
+from uuid import UUID, uuid5
+
+import structlog
+from langchain_core.runnables.config import run_in_executor
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import Graph
+from langgraph.pregel import Pregel
+from langgraph.store.base import BaseStore
+from starlette.exceptions import HTTPException
+
+from langgraph_api.js.remote import RemotePregel
+from langgraph_api.schema import Config
+
+logger = structlog.stdlib.get_logger(__name__)
+
+GraphFactoryFromConfig = Callable[[Config], Pregel | Graph]
+GraphFactory = Callable[[], Pregel | Graph]
+GraphValue = Pregel | GraphFactory
+
+
+GRAPHS: dict[str, Pregel | GraphFactoryFromConfig | GraphFactory] = {}
+NAMESPACE_GRAPH = UUID("6ba7b821-9dad-11d1-80b4-00c04fd430c8")
+FACTORY_ACCEPTS_CONFIG: dict[str, bool] = {}
+
+
+async def register_graph(graph_id: str, graph: GraphValue, config: dict | None) -> None:
+    """Register a graph."""
+    from langgraph_storage.database import connect
+    from langgraph_storage.ops import Assistants
+
+    await logger.ainfo(f"Registering graph with id '{graph_id}'", graph_id=graph_id)
+    GRAPHS[graph_id] = graph
+    if callable(graph):
+        FACTORY_ACCEPTS_CONFIG[graph_id] = len(inspect.signature(graph).parameters) > 0
+    async with connect() as conn:
+        await Assistants.put(
+            conn,
+            str(uuid5(NAMESPACE_GRAPH, graph_id)),
+            graph_id=graph_id,
+            metadata={"created_by": "system"},
+            config=config or {},
+            if_exists="do_nothing",
+            name=graph_id,
+        )
+
+
+def get_graph(
+    graph_id: str,
+    config: Config,
+    *,
+    checkpointer: BaseCheckpointSaver | None = None,
+    store: BaseStore | None = None,
+) -> Pregel:
+    """Return the runnable."""
+    assert_graph_exists(graph_id)
+    value = GRAPHS[graph_id]
+    if graph_id in FACTORY_ACCEPTS_CONFIG:
+        value = value(config) if FACTORY_ACCEPTS_CONFIG[graph_id] else value()
+    if isinstance(value, Graph):
+        value = value.compile()
+    if not isinstance(value, Pregel) and not isinstance(value, RemotePregel):
+        raise HTTPException(
+            status_code=424,
+            detail=f"Graph '{graph_id}' is not valid. Review graph registration.",
+        )
+    if isinstance(value, RemotePregel):
+        value.checkpointer = checkpointer
+        value.name = graph_id
+        return value
+
+    update = {
+        "checkpointer": checkpointer,
+        "store": store,
+    }
+    if value.name == "LangGraph":
+        update["name"] = graph_id
+    return value.copy(update=update)
+
+
+def graph_exists(graph_id: str) -> bool:
+    """Return whether a graph exists."""
+    return graph_id in GRAPHS
+
+
+def assert_graph_exists(graph_id: str) -> None:
+    """Assert that a graph exists."""
+    if not graph_exists(graph_id):
+        raise HTTPException(status_code=404, detail=f"Graph '{graph_id}' not found")
+
+
+def get_assistant_id(assistant_id: str) -> str:
+    """Check if assistant_id is a valid graph_id. If so, retrieve the
+    assistant_id from the graph_id. Otherwise, return the assistant_id
+    as is.
+
+    This method is used where the API allows passing both assistant_id
+    and graph_id interchangeably.
+    """
+    if assistant_id in GRAPHS:
+        assistant_id = str(uuid5(NAMESPACE_GRAPH, assistant_id))
+    return assistant_id
+
+
+class GraphSpec(NamedTuple):
+    """A graph specification."""
+
+    id: str
+    path: str | None = None
+    module: str | None = None
+    variable: str | None = None
+    config: dict | None = None
+
+
+js_bg_tasks: set[asyncio.Task] = set()
+
+
+def is_js_spec(spec: GraphSpec) -> bool:
+    return spec.path is not None and os.path.splitext(spec.path)[1] in (
+        ".ts",
+        ".mts",
+        ".cts",
+        ".js",
+        ".mjs",
+        ".cjs",
+    )
+
+
+def _load_graph_config_from_env() -> dict | None:
+    """Return graph config from env."""
+    config_str = os.getenv("LANGGRAPH_CONFIG")
+    if not config_str:
+        return None
+
+    config_per_id = json.loads(config_str)
+    if not isinstance(config_per_id, dict):
+        raise ValueError("LANGGRAPH_CONFIG must be a JSON object")
+
+    return config_per_id
+
+
+async def collect_graphs_from_env(register: bool = False) -> None:
+    """Return graphs from env."""
+
+    paths_str = os.getenv("LANGSERVE_GRAPHS")
+    config_per_graph = _load_graph_config_from_env() or {}
+
+    if paths_str:
+        specs = [
+            GraphSpec(
+                key,
+                module=value.split(":")[0],
+                variable=value.split(":")[1],
+                config=config_per_graph.get(key),
+            )
+            if "/" not in value
+            else GraphSpec(
+                key,
+                path=value.split(":")[0],
+                variable=value.split(":")[1],
+                config=config_per_graph.get(key),
+            )
+            for key, value in json.loads(paths_str).items()
+        ]
+    else:
+        specs = [
+            GraphSpec(
+                id=graph_path.split("/")[-1].replace(".py", ""),
+                path=graph_path,
+                config=config_per_graph.get(
+                    graph_path.split("/")[-1].replace(".py", "")
+                ),
+            )
+            for graph_path in glob.glob("/graphs/*.py")
+        ]
+
+    js_specs = list(filter(is_js_spec, specs))
+    py_specs = list(filterfalse(is_js_spec, specs))
+
+    if js_specs:
+        import sys
+
+        from langgraph_api.js.remote import (
+            RemotePregel,
+            run_js_process,
+            run_remote_checkpointer,
+            run_remote_store,
+            wait_until_js_ready,
+        )
+
+        js_bg_tasks.add(
+            asyncio.create_task(
+                run_remote_checkpointer(),
+                name="remote-checkpointer",
+            )
+        )
+        js_bg_tasks.add(
+            asyncio.create_task(
+                run_remote_store(),
+                name="remote-store",
+            )
+        )
+        js_bg_tasks.add(
+            asyncio.create_task(
+                run_js_process(paths_str, watch="--reload" in sys.argv[1:]),
+                name="remote-graphs",
+            )
+        )
+        for task in js_bg_tasks:
+            task.add_done_callback(_handle_exception)
+
+        await wait_until_js_ready()
+
+        for spec in js_specs:
+            graph = await RemotePregel.load(graph_id=spec.id)
+            if register:
+                await register_graph(spec.id, graph, spec.config)
+
+    for spec in py_specs:
+        graph = await run_in_executor(None, _graph_from_spec, spec)
+        if register:
+            await register_graph(spec.id, graph, spec.config)
+
+
+def _handle_exception(task: asyncio.Task) -> None:
+    try:
+        task.result()
+    except asyncio.CancelledError:
+        pass
+    finally:
+        # if the task died either with exception or not, we should exit
+        sys.exit(1)
+
+
+async def stop_remote_graphs() -> None:
+    logger.info("Cancelling remote graphs")
+    for task in js_bg_tasks:
+        task.cancel()
+
+
+def verify_graphs() -> None:
+    asyncio.run(collect_graphs_from_env())
+
+
+def _graph_from_spec(spec: GraphSpec) -> GraphValue:
+    """Return a graph from a spec."""
+    # import the graph module
+    if spec.module:
+        module = importlib.import_module(spec.module)
+    elif spec.path:
+        try:
+            modname = "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(24))
+            modspec = importlib.util.spec_from_file_location(modname, spec.path)
+            if modspec is None:
+                raise ValueError(f"Could not find python file for graph: {spec}")
+            module = importlib.util.module_from_spec(modspec)
+            sys.modules[modname] = module
+            modspec.loader.exec_module(module)
+        except ImportError as e:
+            e.add_note(f"Could not import python module for graph: {spec}")
+            raise e
+    else:
+        raise ValueError("Graph specification must have a path or module")
+
+    if spec.variable:
+        graph: GraphValue = getattr(module, spec.variable)
+        if callable(graph):
+            sig = inspect.signature(graph)
+            if not sig.parameters:
+                if not isinstance(graph(), (Graph, Pregel)):  # noqa: UP038
+                    raise ValueError(
+                        f"Graph factory function '{spec.variable}' in module '{spec.path}' must return a Graph or Pregel"
+                    )
+            elif len(sig.parameters) != 1:
+                raise ValueError(
+                    f"Graph factory function '{spec.variable}' in module '{spec.path}' must take exactly one argument, a RunnableConfig"
+                )
+        elif isinstance(graph, Graph):
+            graph = graph.compile()
+        elif isinstance(graph, Pregel):
+            pass
+        else:
+            raise ValueError(
+                f"Variable '{spec.variable}' in module '{spec.path}' is not a Graph or Graph factory function"
+            )
+    else:
+        # find the graph in the module
+        # - first look for a compiled graph (Pregel)
+        # - if not found, look for a Graph and compile it
+        for _, member in inspect.getmembers(module):
+            if isinstance(member, Pregel):
+                graph = member
+                break
+        else:
+            for _, member in inspect.getmembers(module):
+                if isinstance(member, Graph):
+                    graph = member.compile()
+                    break
+            else:
+                raise ValueError(
+                    f"Could not find a Graph in module at path: {spec.path}"
+                )
+
+    return graph

langgraph_api/http.py

@@ -0,0 +1,168 @@
+import asyncio
+
+import httpx
+from tenacity import (
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
+
+from langgraph_api.serde import json_dumpb
+
+
+class JsonHttpClient:
+    """HTTPX client for JSON requests."""
+
+    def __init__(self, client: httpx.AsyncClient) -> None:
+        """Initialize the auth client."""
+        self.client = client
+
+    async def post(
+        self,
+        path: str,
+        /,
+        *,
+        params: dict[str, str] | None = None,
+        headers: dict[str, str] | None = None,
+        json: dict | None = None,
+        content: bytes | None = None,
+        connect_timeout: float | None = None,
+        request_timeout: float | None = None,
+        total_timeout: float | None = None,
+        raise_error: bool = True,
+    ) -> None:
+        request_headers = {"Content-Type": "application/json"}
+        if headers:
+            request_headers.update(headers)
+
+        if json and content:
+            raise ValueError("Cannot specify both 'json' and 'content'")
+
+        try:
+            res = await asyncio.wait_for(
+                self.client.post(
+                    path,
+                    content=json_dumpb(json) if json else content,
+                    headers=request_headers,
+                    timeout=httpx.Timeout(
+                        request_timeout or connect_timeout,
+                        connect=connect_timeout,
+                        read=request_timeout,
+                    )
+                    if connect_timeout or request_timeout
+                    else None,
+                    params=params,
+                ),
+                # httpx timeout controls are additive for each operation
+                # (connect, read, write), so we need an asyncio timeout instead
+                total_timeout,
+            )
+            # Raise for retriable errors
+            if raise_error:
+                res.raise_for_status()
+        finally:
+            # We don't need the response body, so we close the response
+            try:
+                await res.aclose()
+            except UnboundLocalError:
+                pass
+
+
+_http_client: JsonHttpClient
+
+
+async def start_http_client() -> None:
+    global _http_client
+    _http_client = JsonHttpClient(
+        client=httpx.AsyncClient(
+            transport=httpx.AsyncHTTPTransport(
+                retries=2,  # this applies only to ConnectError, ConnectTimeout
+                limits=httpx.Limits(
+                    max_keepalive_connections=10, keepalive_expiry=60.0
+                ),
+            ),
+        ),
+    )
+
+
+async def stop_http_client() -> None:
+    global _http_client
+    await _http_client.client.aclose()
+    del _http_client
+
+
+def get_http_client() -> JsonHttpClient:
+    return _http_client
+
+
+def is_retriable_error(exception: Exception) -> bool:
+    if isinstance(exception, httpx.HTTPError):
+        return (
+            getattr(exception, "response", None) is not None
+            and exception.response.status_code >= 500
+        )
+    return False
+
+
+retry_http = retry(
+    reraise=True,
+    retry=retry_if_exception(is_retriable_error),
+    wait=wait_exponential_jitter(),
+    stop=stop_after_attempt(3),
+)
+
+
+@retry_http
+async def http_request(
+    method: str,
+    path: str,
+    /,
+    *,
+    params: dict[str, str] | None = None,
+    headers: dict[str, str] | None = None,
+    body: bytes | str | None = None,
+    json: dict | None = None,
+    connect_timeout: float | None = 5,
+    request_timeout: float | None = 30,
+    raise_error: bool = True,
+) -> httpx.Response:
+    """Make an HTTP request with retries.
+
+    Args:
+        method: HTTP method
+        path: URL path
+        params: Query parameters
+        headers: Request headers
+        body: Raw request body (bytes or str)
+        json: JSON body (mutually exclusive with body)
+        connect_timeout: Connection timeout in seconds
+        request_timeout: Request timeout in seconds
+        raise_error: Whether to raise for HTTP errors
+
+    Returns:
+        httpx.Response object
+    """
+    if not path.startswith(("http://", "https://", "/")):
+        raise ValueError("path must start with / or http")
+
+    client = get_http_client()
+
+    content = None
+    if body is not None:
+        content = body
+    elif json is not None:
+        content = json_dumpb(json)
+
+    if method.upper() == "POST":
+        return await client.post(
+            path,
+            params=params,
+            headers=headers,
+            content=content,
+            connect_timeout=connect_timeout,
+            request_timeout=request_timeout,
+            raise_error=raise_error,
+        )
+    else:
+        raise ValueError(f"Unsupported HTTP method: {method}")

langgraph_api/http_logger.py

@@ -0,0 +1,89 @@
+import asyncio
+import logging
+
+import structlog
+from starlette.requests import ClientDisconnect
+from starlette.types import Message, Receive, Scope, Send
+
+from langgraph_api.asyncio import create_task
+from langgraph_api.logging import LOG_JSON
+
+asgi = structlog.stdlib.get_logger("asgi")
+
+PATHS_IGNORE = {"/ok", "/metrics"}
+
+
+class AccessLoggerMiddleware:
+    def __init__(
+        self,
+        app,
+        logger: structlog.stdlib.BoundLogger,
+    ) -> None:
+        self.app = app
+        self.logger = logger
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http" or (LOG_JSON and scope.get("path") in PATHS_IGNORE):
+            return await self.app(scope, receive, send)  # pragma: no cover
+
+        loop = asyncio.get_event_loop()
+        info = {"response": {}}
+
+        if self.logger.isEnabledFor(logging.DEBUG):
+
+            async def inner_receive() -> Message:
+                message = await receive()
+                asgi.debug(f"ASGI receive {message['type']}", **message)
+                return message
+
+            async def inner_send(message: Message) -> None:
+                if message["type"] == "http.response.start":
+                    info["response"] = message
+                await send(message)
+                asgi.debug(f"ASGI send {message['type']}", **message)
+        else:
+            inner_receive = receive
+
+            async def inner_send(message) -> None:
+                if message["type"] == "http.response.start":
+                    info["response"] = message
+                await send(message)
+
+        try:
+            info["start_time"] = loop.time()
+            await self.app(scope, inner_receive, inner_send)
+        except ClientDisconnect as exc:
+            info["response"]["status"] = 499
+            raise exc
+        except Exception as exc:
+            info["response"]["status"] = 500
+            raise exc
+        finally:
+            info["end_time"] = loop.time()
+            latency = int((info["end_time"] - info["start_time"]) * 1_000)
+            create_task(
+                self.logger.ainfo(
+                    f"{scope.get('method')} {scope.get('path')} {info['response'].get('status')} {latency}ms",
+                    method=scope.get("method"),
+                    path=scope.get("path"),
+                    status=info["response"].get("status"),
+                    latency_ms=latency,
+                    route=scope.get("route"),
+                    path_params=scope.get("path_params"),
+                    query_string=scope.get("query_string").decode(),
+                    proto=scope.get("http_version"),
+                    req_header=_headers_to_dict(scope.get("headers")),
+                    res_header=_headers_to_dict(info["response"].get("headers")),
+                )
+            )
+
+
+HEADERS_IGNORE = {"authorization", "cookie", "set-cookie", "x-api-key"}
+
+
+def _headers_to_dict(headers: list[tuple[bytes, bytes]] | None) -> dict[str, str]:
+    if headers is None:
+        return {}
+    return {
+        k.decode(): v.decode() for k, v in headers if k.lower() not in HEADERS_IGNORE
+    }

langgraph_api/js/.gitignore

@@ -0,0 +1,2 @@
+node_modules
+client.schemas.json

langgraph_api/js/build.mts

@@ -0,0 +1,49 @@
+import { z } from "zod";
+import * as fs from "node:fs/promises";
+import * as path from "node:path";
+import {
+  GraphSchema,
+  resolveGraph,
+  runGraphSchemaWorker,
+} from "./src/graph.mts";
+
+const __dirname = new URL(".", import.meta.url).pathname;
+
+async function main() {
+  const specs = z
+    .record(z.string())
+    .parse(JSON.parse(process.env.LANGSERVE_GRAPHS));
+
+  const GRAPH_SCHEMAS: Record<string, Record<string, GraphSchema>> = {};
+
+  try {
+    await Promise.all(
+      Object.entries(specs).map(async ([graphId, rawSpec]) => {
+        console.info(`[${graphId}]: Checking for source file existence`);
+        const { resolved, ...spec } = await resolveGraph(rawSpec, {
+          onlyFilePresence: true,
+        });
+
+        try {
+          console.info(`[${graphId}]: Extracting schema`);
+          GRAPH_SCHEMAS[graphId] = await runGraphSchemaWorker(spec);
+        } catch (error) {
+          console.error(`[${graphId}]: Error extracting schema: ${error}`);
+        }
+      })
+    );
+
+    await fs.writeFile(
+      path.resolve(__dirname, "client.schemas.json"),
+      JSON.stringify(GRAPH_SCHEMAS),
+      { encoding: "utf-8" }
+    );
+  } catch (error) {
+    console.error(`Error resolving graphs: ${error}`);
+    process.exit(1);
+  }
+
+  console.info("All graphs resolved");
+}
+
+main();