langgraph-api 0.2.26__py3-none-any.whl → 0.2.28__py3-none-any.whl

langgraph_api/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.26"
+__version__ = "0.2.28"

langgraph_api/api/assistants.py

@@ -9,6 +9,7 @@ from starlette.exceptions import HTTPException
 from starlette.responses import Response
 from starlette.routing import BaseRoute
 
+from langgraph_api import store as api_store
 from langgraph_api.graph import get_assistant_id, get_graph
 from langgraph_api.js.base import BaseRemotePregel
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
@@ -25,7 +26,6 @@ from langgraph_runtime.checkpoint import Checkpointer
 from langgraph_runtime.database import connect
 from langgraph_runtime.ops import Assistants
 from langgraph_runtime.retry import retry_db
-from langgraph_runtime.store import Store
 
 logger = structlog.stdlib.get_logger(__name__)
 
@@ -194,7 +194,7 @@ async def get_assistant_graph(
             assistant["graph_id"],
             config,
             checkpointer=Checkpointer(conn),
-            store=Store(),
+            store=(await api_store.get_store()),
         ) as graph:
             xray: bool | int = False
             xray_query = request.query_params.get("xray")
@@ -240,7 +240,7 @@ async def get_assistant_subgraphs(
             assistant["graph_id"],
             config,
             checkpointer=Checkpointer(conn),
-            store=Store(),
+            store=(await api_store.get_store()),
         ) as graph:
             namespace = request.path_params.get("namespace")
 
@@ -286,7 +286,7 @@ async def get_assistant_schemas(
             assistant["graph_id"],
             config,
             checkpointer=Checkpointer(conn),
-            store=Store(),
+            store=(await api_store.get_store()),
         ) as graph:
             if isinstance(graph, BaseRemotePregel):
                 schemas = await graph.fetch_state_schema()

langgraph_api/api/store.py

@@ -6,6 +6,7 @@ from starlette.routing import BaseRoute
 
 from langgraph_api.auth.custom import handle_event as _handle_event
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
+from langgraph_api.store import get_store
 from langgraph_api.utils import get_auth_ctx
 from langgraph_api.validation import (
     StoreDeleteRequest,
@@ -14,7 +15,6 @@ from langgraph_api.validation import (
     StoreSearchRequest,
 )
 from langgraph_runtime.retry import retry_db
-from langgraph_runtime.store import Store
 
 
 def _validate_namespace(namespace: tuple[str, ...]) -> Response | None:
@@ -57,7 +57,9 @@ async def put_item(request: ApiRequest):
         "value": payload["value"],
     }
     await handle_event("put", handler_payload)
-    await Store().aput(namespace, handler_payload["key"], handler_payload["value"])
+    await (await get_store()).aput(
+        namespace, handler_payload["key"], handler_payload["value"]
+    )
     return Response(status_code=204)
 
 
@@ -75,7 +77,7 @@ async def get_item(request: ApiRequest):
         "key": key,
     }
     await handle_event("get", handler_payload)
-    result = await Store().aget(namespace, key)
+    result = await (await get_store()).aget(namespace, key)
     return ApiResponse(result.dict() if result is not None else None)
 
 
@@ -91,7 +93,9 @@ async def delete_item(request: ApiRequest):
         "key": payload["key"],
     }
     await handle_event("delete", handler_payload)
-    await Store().adelete(handler_payload["namespace"], handler_payload["key"])
+    await (await get_store()).adelete(
+        handler_payload["namespace"], handler_payload["key"]
+    )
     return Response(status_code=204)
 
 
@@ -114,7 +118,7 @@ async def search_items(request: ApiRequest):
         "query": query,
     }
     await handle_event("search", handler_payload)
-    items = await Store().asearch(
+    items = await (await get_store()).asearch(
         handler_payload["namespace"],
         filter=handler_payload["filter"],
         limit=handler_payload["limit"],
@@ -145,7 +149,7 @@ async def list_namespaces(request: ApiRequest):
         "offset": offset,
     }
     await handle_event("list_namespaces", handler_payload)
-    result = await Store().alist_namespaces(
+    result = await (await get_store()).alist_namespaces(
         prefix=handler_payload["namespace"],
         suffix=handler_payload["suffix"],
         max_depth=handler_payload["max_depth"],

langgraph_api/asgi_transport.py

@@ -0,0 +1,171 @@
+"""ASGI transport that lets you schedule to the main loop.
+
+Adapted from: https://github.com/encode/httpx/blob/6c7af967734bafd011164f2a1653abc87905a62b/httpx/_transports/asgi.py#L1
+"""
+
+from __future__ import annotations
+
+import typing
+
+from httpx import ASGITransport as ASGITransportBase
+from httpx import AsyncByteStream, Request, Response
+
+if typing.TYPE_CHECKING:  # pragma: no cover
+    import asyncio
+
+    import trio
+
+    Event = asyncio.Event | trio.Event
+
+__all__ = ["ASGITransport"]
+
+
+def is_running_trio() -> bool:
+    try:
+        # sniffio is a dependency of trio.
+
+        # See https://github.com/python-trio/trio/issues/2802
+        import sniffio
+
+        if sniffio.current_async_library() == "trio":
+            return True
+    except ImportError:  # pragma: nocover
+        pass
+
+    return False
+
+
+def create_event() -> Event:
+    if is_running_trio():
+        import trio
+
+        return trio.Event()
+
+    import asyncio
+
+    return asyncio.Event()
+
+
+class ASGIResponseStream(AsyncByteStream):
+    def __init__(self, body: list[bytes]) -> None:
+        self._body = body
+
+    async def __aiter__(self) -> typing.AsyncIterator[bytes]:
+        yield b"".join(self._body)
+
+
+class ASGITransport(ASGITransportBase):
+    """
+    A custom AsyncTransport that handles sending requests directly to an ASGI app.
+
+    ```python
+    transport = httpx.ASGITransport(
+        app=app,
+        root_path="/submount",
+        client=("1.2.3.4", 123)
+    )
+    client = httpx.AsyncClient(transport=transport)
+    ```
+
+    Arguments:
+
+    * `app` - The ASGI application.
+    * `raise_app_exceptions` - Boolean indicating if exceptions in the application
+       should be raised. Default to `True`. Can be set to `False` for use cases
+       such as testing the content of a client 500 response.
+    * `root_path` - The root path on which the ASGI application should be mounted.
+    * `client` - A two-tuple indicating the client IP and port of incoming requests.
+    ```
+    """
+
+    async def handle_async_request(
+        self,
+        request: Request,
+    ) -> Response:
+        from langgraph_api.asyncio import call_soon_in_main_loop
+
+        assert isinstance(request.stream, AsyncByteStream)
+
+        # ASGI scope.
+        scope = {
+            "type": "http",
+            "asgi": {"version": "3.0"},
+            "http_version": "1.1",
+            "method": request.method,
+            "headers": [(k.lower(), v) for (k, v) in request.headers.raw],
+            "scheme": request.url.scheme,
+            "path": request.url.path,
+            "raw_path": request.url.raw_path.split(b"?")[0],
+            "query_string": request.url.query,
+            "server": (request.url.host, request.url.port),
+            "client": self.client,
+            "root_path": self.root_path,
+        }
+
+        # Request.
+        request_body_chunks = request.stream.__aiter__()
+        request_complete = False
+
+        # Response.
+        status_code = None
+        response_headers = None
+        body_parts = []
+        response_started = False
+        response_complete = create_event()
+
+        # ASGI callables.
+
+        async def receive() -> dict[str, typing.Any]:
+            nonlocal request_complete
+
+            if request_complete:
+                await response_complete.wait()
+                return {"type": "http.disconnect"}
+
+            try:
+                body = await request_body_chunks.__anext__()
+            except StopAsyncIteration:
+                request_complete = True
+                return {"type": "http.request", "body": b"", "more_body": False}
+            return {"type": "http.request", "body": body, "more_body": True}
+
+        async def send(message: typing.MutableMapping[str, typing.Any]) -> None:
+            nonlocal status_code, response_headers, response_started
+
+            if message["type"] == "http.response.start":
+                assert not response_started
+
+                status_code = message["status"]
+                response_headers = message.get("headers", [])
+                response_started = True
+
+            elif message["type"] == "http.response.body":
+                assert not response_complete.is_set()
+                body = message.get("body", b"")
+                more_body = message.get("more_body", False)
+
+                if body and request.method != "HEAD":
+                    body_parts.append(body)
+
+                if not more_body:
+                    response_complete.set()
+
+        try:
+            await call_soon_in_main_loop(self.app(scope, receive, send))
+        except Exception:  # noqa: PIE-786
+            if self.raise_app_exceptions:
+                raise
+
+            response_complete.set()
+            if status_code is None:
+                status_code = 500
+            if response_headers is None:
+                response_headers = {}
+
+        assert response_complete.is_set()
+        assert status_code is not None
+        assert response_headers is not None
+
+        stream = ASGIResponseStream(body_parts)
+
+        return Response(status_code, headers=response_headers, stream=stream)

langgraph_api/asyncio.py

@@ -6,6 +6,7 @@ from functools import partial
 from typing import Any, Generic, TypeVar
 
 import structlog
+from langgraph.utils.future import chain_future
 
 T = TypeVar("T")
 
@@ -19,6 +20,12 @@ def set_event_loop(loop: asyncio.AbstractEventLoop) -> None:
     _MAIN_LOOP = loop
 
 
+def get_event_loop() -> asyncio.AbstractEventLoop:
+    if _MAIN_LOOP is None:
+        raise RuntimeError("No event loop set")
+    return _MAIN_LOOP
+
+
 async def sleep_if_not_done(delay: float, done: asyncio.Event) -> None:
     try:
         await asyncio.wait_for(done.wait(), delay)
@@ -118,6 +125,16 @@ def run_coroutine_threadsafe(
     return future
 
 
+def call_soon_in_main_loop(coro: Coroutine[Any, Any, T]) -> asyncio.Future[T]:
+    """Run a coroutine in the main event loop."""
+    if _MAIN_LOOP is None:
+        raise RuntimeError("No event loop set")
+    main_loop_fut = asyncio.ensure_future(coro, loop=_MAIN_LOOP)
+    this_loop_fut = asyncio.get_running_loop().create_future()
+    _MAIN_LOOP.call_soon_threadsafe(chain_future, main_loop_fut, this_loop_fut)
+    return this_loop_fut
+
+
 class SimpleTaskGroup(AbstractAsyncContextManager["SimpleTaskGroup"]):
     """An async task group that can be configured to wait and/or cancel tasks on exit.
 

langgraph_api/config.py

@@ -109,6 +109,7 @@ class TTLConfig(TypedDict, total=False):
 
 
 class StoreConfig(TypedDict, total=False):
+    path: str
     index: IndexConfig
     ttl: TTLConfig
 

langgraph_api/graph.py

@@ -3,7 +3,6 @@ import functools
 import glob
 import importlib.util
 import inspect
-import json
 import os
 import sys
 import warnings
@@ -14,6 +13,7 @@ from random import choice
 from typing import TYPE_CHECKING, Any, NamedTuple
 from uuid import UUID, uuid5
 
+import orjson
 import structlog
 from langchain_core.runnables.config import run_in_executor, var_child_runnable_config
 from langgraph.checkpoint.base import BaseCheckpointSaver
@@ -60,6 +60,12 @@ async def register_graph(
     if callable(graph):
         FACTORY_ACCEPTS_CONFIG[graph_id] = len(inspect.signature(graph).parameters) > 0
     async with connect() as conn:
+        graph_name = getattr(graph, "name", None) if isinstance(graph, Pregel) else None
+        assistant_name = (
+            graph_name
+            if graph_name is not None and graph_name != "LangGraph"
+            else graph_id
+        )
         await Assistants.put(
             conn,
             str(uuid5(NAMESPACE_GRAPH, graph_id)),
@@ -67,7 +73,7 @@ async def register_graph(
             metadata={"created_by": "system"},
             config=config or {},
             if_exists="do_nothing",
-            name=graph_id,
+            name=assistant_name,
             description=description,
         )
 
@@ -200,10 +206,19 @@ def _load_graph_config_from_env() -> dict | None:
     config_str = os.getenv("LANGGRAPH_CONFIG")
     if not config_str:
         return None
+    try:
+        config_per_id = orjson.loads(config_str)
+    except orjson.JSONDecodeError as e:
+        raise ValueError(
+            "Provided environment variable LANGGRAPH_CONFIG must be a valid JSON object"
+            f"\nFound: {config_str}"
+        ) from e
 
-    config_per_id = json.loads(config_str)
     if not isinstance(config_per_id, dict):
-        raise ValueError("LANGGRAPH_CONFIG must be a JSON object")
+        raise ValueError(
+            "Provided environment variable LANGGRAPH_CONFIG must be a JSON object"
+            f"\nFound: {config_str}"
+        )
 
     return config_per_id
 
@@ -218,7 +233,15 @@ async def collect_graphs_from_env(register: bool = False) -> None:
         specs = []
         # graphs-config can be either a mapping from graph id to path where the graph
         # is defined or graph id to a dictionary containing information about the graph.
-        graphs_config = json.loads(paths_str)
+        try:
+            graphs_config = orjson.loads(paths_str)
+        except orjson.JSONDecodeError as e:
+            raise ValueError(
+                "LANGSERVE_GRAPHS must be a valid JSON object."
+                f"\nFound: {paths_str}"
+                "\n The LANGSERVE_GRAPHS environment variable is typically set"
+                'from the "graphs" field in your configuration (langgraph.json) file.'
+            ) from e
 
         for key, value in graphs_config.items():
             if isinstance(value, dict) and "path" in value:

langgraph_api/js/remote.py

@@ -39,6 +39,7 @@ from starlette.exceptions import HTTPException
 from starlette.requests import HTTPConnection, Request
 from starlette.routing import Route
 
+from langgraph_api import store as api_store
 from langgraph_api.auth.custom import DotDict, ProxyUser
 from langgraph_api.config import LANGGRAPH_AUTH_TYPE
 from langgraph_api.js.base import BaseRemotePregel
@@ -70,6 +71,12 @@ _client = httpx.AsyncClient(
 )
 
 
+def _snapshot_defaults():
+    if not hasattr(StateSnapshot, "interrupts"):
+        return {}
+    return {"interrupts": tuple()}
+
+
 def default_command(obj):
     if isinstance(obj, Send):
         return {"node": obj.node, "args": obj.arg}
@@ -251,7 +258,7 @@ class RemotePregel(BaseRemotePregel):
             item.get("parentConfig"),
             _convert_tasks(item.get("tasks", [])),
             # TODO: add handling of interrupts when multiple resumes land in JS
-            tuple(),
+            **_snapshot_defaults(),
         )
 
     async def aget_state(
@@ -473,10 +480,8 @@ def _get_passthrough_checkpointer(conn: AsyncConnectionProto):
     return checkpointer
 
 
-def _get_passthrough_store():
-    from langgraph_runtime.store import Store
-
-    return Store()
+async def _get_passthrough_store():
+    return await api_store.get_store()
 
 
 # Setup a HTTP server on top of CHECKPOINTER_SOCKET unix socket
@@ -574,7 +579,7 @@ async def run_remote_checkpointer():
             else:
                 raise ValueError(f"Unknown operation type: {op}")
 
-        store = _get_passthrough_store()
+        store = await _get_passthrough_store()
         results = await store.abatch(processed_operations)
 
         # Handle potentially undefined or non-dict results
@@ -613,7 +618,7 @@ async def run_remote_checkpointer():
 
         namespaces = namespaces_str.split(".")
 
-        store = _get_passthrough_store()
+        store = await _get_passthrough_store()
         result = await store.aget(namespaces, key)
 
         return result
@@ -626,7 +631,7 @@ async def run_remote_checkpointer():
         value = payload["value"]
         index = payload.get("index")
 
-        store = _get_passthrough_store()
+        store = await _get_passthrough_store()
         await store.aput(namespace, key, value, index=index)
 
         return {"success": True}
@@ -639,7 +644,7 @@ async def run_remote_checkpointer():
         offset = payload.get("offset", 0)
         query = payload.get("query")
 
-        store = _get_passthrough_store()
+        store = await _get_passthrough_store()
         result = await store.asearch(
             namespace_prefix, filter=filter, limit=limit, offset=offset, query=query
         )
@@ -652,7 +657,7 @@ async def run_remote_checkpointer():
         namespace = tuple(payload["namespace"])
         key = payload["key"]
 
-        store = _get_passthrough_store()
+        store = await _get_passthrough_store()
         await store.adelete(namespace, key)
 
         return {"success": True}
@@ -665,7 +670,7 @@ async def run_remote_checkpointer():
         limit = payload.get("limit", 100)
         offset = payload.get("offset", 0)
 
-        store = _get_passthrough_store()
+        store = await _get_passthrough_store()
         result = await store.alist_namespaces(
             prefix=prefix,
             suffix=suffix,

langgraph_api/metadata.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+from collections import defaultdict
 from datetime import UTC, datetime
 
 import langgraph.version
@@ -36,8 +37,8 @@ PLAN = "enterprise" if plus_features_enabled() else "developer"
 USER_API_URL = os.getenv("LANGGRAPH_API_URL", None)
 
 LOGS: list[dict] = []
-RUN_COUNTER = 0
-NODE_COUNTER = 0
+RUN_COUNTER = defaultdict(int)
+NODE_COUNTER = defaultdict(int)
 FROM_TIMESTAMP = datetime.now(UTC).isoformat()
 
 if (
@@ -49,14 +50,12 @@ else:
     METADATA_ENDPOINT = "https://api.smith.langchain.com/v1/metadata/submit"
 
 
-def incr_runs(*, incr: int = 1) -> None:
-    global RUN_COUNTER
-    RUN_COUNTER += incr
+def incr_runs(*, graph_id: str | None = None, incr: int = 1) -> None:
+    RUN_COUNTER[graph_id] += incr
 
 
-def incr_nodes(_, *, incr: int = 1) -> None:
-    global NODE_COUNTER
-    NODE_COUNTER += incr
+def incr_nodes(*_, graph_id: str | None = None, incr: int = 1) -> None:
+    NODE_COUNTER[graph_id] += incr
 
 
 def append_log(log: dict) -> None:
@@ -89,13 +88,23 @@ async def metadata_loop() -> None:
         # we don't need a lock as long as there's no awaits in this block
         from_timestamp = FROM_TIMESTAMP
         to_timestamp = datetime.now(UTC).isoformat()
-        nodes = NODE_COUNTER
-        runs = RUN_COUNTER
+        nodes = NODE_COUNTER.copy()
+        runs = RUN_COUNTER.copy()
         logs = LOGS.copy()
         LOGS.clear()
-        RUN_COUNTER = 0
-        NODE_COUNTER = 0
+        RUN_COUNTER.clear()
+        NODE_COUNTER.clear()
         FROM_TIMESTAMP = to_timestamp
+        graph_measures = {
+            f"langgraph.platform.graph_runs.{graph_id}": runs.get(graph_id, 0)
+            for graph_id in runs
+        }
+        graph_measures.update(
+            {
+                f"langgraph.platform.graph_nodes.{graph_id}": nodes.get(graph_id, 0)
+                for graph_id in nodes
+            }
+        )
 
         payload = {
             "license_key": LANGGRAPH_CLOUD_LICENSE_KEY,
@@ -120,8 +129,9 @@ async def metadata_loop() -> None:
                 "user_app.uses_store_ttl": str(USES_STORE_TTL),
             },
             "measures": {
-                "langgraph.platform.runs": runs,
-                "langgraph.platform.nodes": nodes,
+                "langgraph.platform.runs": sum(runs.values()),
+                "langgraph.platform.nodes": sum(nodes.values()),
+                **graph_measures,
             },
             "logs": logs,
         }
@@ -134,8 +144,10 @@ async def metadata_loop() -> None:
             )
         except Exception as e:
             # retry on next iteration
-            incr_runs(incr=runs)
-            incr_nodes("", incr=nodes)
+            for graph_id, incr in runs.items():
+                incr_runs(graph_id=graph_id, incr=incr)
+            for graph_id, incr in nodes.items():
+                incr_nodes(graph_id=graph_id, incr=incr)
             FROM_TIMESTAMP = from_timestamp
             await logger.ainfo("Metadata submission skipped.", error=str(e))
         await asyncio.sleep(INTERVAL)