langgraph-api 0.0.1__py3-none-any.whl

langgraph_api/api/threads.py

@@ -0,0 +1,263 @@
+from uuid import uuid4
+
+from starlette.responses import Response
+from starlette.routing import BaseRoute
+
+from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
+from langgraph_api.state import state_snapshot_to_thread_state
+from langgraph_api.utils import fetchone, validate_uuid
+from langgraph_api.validation import (
+    ThreadCreate,
+    ThreadPatch,
+    ThreadSearchRequest,
+    ThreadStateCheckpointRequest,
+    ThreadStateSearch,
+    ThreadStateUpdate,
+)
+from langgraph_storage.database import connect
+from langgraph_storage.ops import Threads
+from langgraph_storage.retry import retry_db
+
+
+@retry_db
+async def create_thread(
+    request: ApiRequest,
+):
+    """Create a thread."""
+    payload = await request.json(ThreadCreate)
+    if thread_id := payload.get("thread_id"):
+        validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    async with connect() as conn:
+        iter = await Threads.put(
+            conn,
+            thread_id or uuid4(),
+            metadata=payload.get("metadata"),
+            if_exists=payload.get("if_exists") or "raise",
+        )
+    return ApiResponse(await fetchone(iter, not_found_code=409))
+
+
+@retry_db
+async def search_threads(
+    request: ApiRequest,
+):
+    """List threads."""
+    payload = await request.json(ThreadSearchRequest)
+    async with connect() as conn:
+        iter = await Threads.search(
+            conn,
+            status=payload.get("status"),
+            values=payload.get("values"),
+            metadata=payload.get("metadata"),
+            limit=payload.get("limit") or 10,
+            offset=payload.get("offset") or 0,
+        )
+    return ApiResponse([thread async for thread in iter])
+
+
+@retry_db
+async def get_thread_state(
+    request: ApiRequest,
+):
+    """Get state for a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    subgraphs = request.query_params.get("subgraphs") in ("true", "True")
+    async with connect() as conn:
+        state = state_snapshot_to_thread_state(
+            await Threads.State.get(
+                conn, {"configurable": {"thread_id": thread_id}}, subgraphs=subgraphs
+            )
+        )
+    return ApiResponse(state)
+
+
+@retry_db
+async def get_thread_state_at_checkpoint(
+    request: ApiRequest,
+):
+    """Get state for a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    checkpoint_id = request.path_params["checkpoint_id"]
+    async with connect() as conn:
+        state = state_snapshot_to_thread_state(
+            await Threads.State.get(
+                conn,
+                {
+                    "configurable": {
+                        "thread_id": thread_id,
+                        "checkpoint_id": checkpoint_id,
+                    }
+                },
+                subgraphs=request.query_params.get("subgraphs") in ("true", "True"),
+            )
+        )
+    return ApiResponse(state)
+
+
+@retry_db
+async def get_thread_state_at_checkpoint_post(
+    request: ApiRequest,
+):
+    """Get state for a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    payload = await request.json(ThreadStateCheckpointRequest)
+    async with connect() as conn:
+        state = state_snapshot_to_thread_state(
+            await Threads.State.get(
+                conn,
+                {"configurable": {"thread_id": thread_id, **payload["checkpoint"]}},
+                subgraphs=payload.get("subgraphs", False),
+            )
+        )
+    return ApiResponse(state)
+
+
+@retry_db
+async def update_thread_state(
+    request: ApiRequest,
+):
+    """Add state to a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    payload = await request.json(ThreadStateUpdate)
+    config = {"configurable": {"thread_id": thread_id}}
+    if payload.get("checkpoint_id"):
+        config["configurable"]["checkpoint_id"] = payload["checkpoint_id"]
+    if payload.get("checkpoint"):
+        config["configurable"].update(payload["checkpoint"])
+    try:
+        if user_id := request.user.display_name:
+            config["configurable"]["user_id"] = user_id
+    except AssertionError:
+        pass
+    async with connect() as conn:
+        inserted = await Threads.State.post(
+            conn,
+            config,
+            payload.get("values"),
+            payload.get("as_node"),
+        )
+    return ApiResponse(inserted)
+
+
+@retry_db
+async def get_thread_history(
+    request: ApiRequest,
+):
+    """Get all past states for a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    limit = request.query_params.get("limit", 10)
+    before = request.query_params.get("before")
+    config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
+    async with connect() as conn:
+        states = [
+            state_snapshot_to_thread_state(c)
+            for c in await Threads.State.list(
+                conn, config=config, limit=limit, before=before
+            )
+        ]
+    return ApiResponse(states)
+
+
+@retry_db
+async def get_thread_history_post(
+    request: ApiRequest,
+):
+    """Get all past states for a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    payload = await request.json(ThreadStateSearch)
+    config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
+    config["configurable"].update(payload.get("checkpoint", {}))
+    async with connect() as conn:
+        states = [
+            state_snapshot_to_thread_state(c)
+            for c in await Threads.State.list(
+                conn,
+                config=config,
+                limit=int(payload.get("limit") or 10),
+                before=payload.get("before"),
+                metadata=payload.get("metadata"),
+            )
+        ]
+    return ApiResponse(states)
+
+
+@retry_db
+async def get_thread(
+    request: ApiRequest,
+):
+    """Get a thread by ID."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    async with connect() as conn:
+        thread = await Threads.get(conn, thread_id)
+    return ApiResponse(await fetchone(thread))
+
+
+@retry_db
+async def patch_thread(
+    request: ApiRequest,
+):
+    """Update a thread."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    payload = await request.json(ThreadPatch)
+    async with connect() as conn:
+        thread = await Threads.patch(conn, thread_id, metadata=payload["metadata"])
+    return ApiResponse(await fetchone(thread))
+
+
+@retry_db
+async def delete_thread(request: ApiRequest):
+    """Delete a thread by ID."""
+    thread_id = request.path_params["thread_id"]
+    validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+    async with connect() as conn:
+        tid = await Threads.delete(conn, thread_id)
+    await fetchone(tid)
+    return Response(status_code=204)
+
+
+@retry_db
+async def copy_thread(request: ApiRequest):
+    thread_id = request.path_params["thread_id"]
+    async with connect() as conn:
+        iter = await Threads.copy(conn, thread_id)
+    return ApiResponse(await fetchone(iter, not_found_code=409))
+
+
+threads_routes: list[BaseRoute] = [
+    ApiRoute("/threads", endpoint=create_thread, methods=["POST"]),
+    ApiRoute("/threads/search", endpoint=search_threads, methods=["POST"]),
+    ApiRoute("/threads/{thread_id}", endpoint=get_thread, methods=["GET"]),
+    ApiRoute("/threads/{thread_id}", endpoint=patch_thread, methods=["PATCH"]),
+    ApiRoute("/threads/{thread_id}", endpoint=delete_thread, methods=["DELETE"]),
+    ApiRoute("/threads/{thread_id}/state", endpoint=get_thread_state, methods=["GET"]),
+    ApiRoute(
+        "/threads/{thread_id}/state", endpoint=update_thread_state, methods=["POST"]
+    ),
+    ApiRoute(
+        "/threads/{thread_id}/history", endpoint=get_thread_history, methods=["GET"]
+    ),
+    ApiRoute("/threads/{thread_id}/copy", endpoint=copy_thread, methods=["POST"]),
+    ApiRoute(
+        "/threads/{thread_id}/history",
+        endpoint=get_thread_history_post,
+        methods=["POST"],
+    ),
+    ApiRoute(
+        "/threads/{thread_id}/state/{checkpoint_id}",
+        endpoint=get_thread_state_at_checkpoint,
+        methods=["GET"],
+    ),
+    ApiRoute(
+        "/threads/{thread_id}/state/checkpoint",
+        endpoint=get_thread_state_at_checkpoint_post,
+        methods=["POST"],
+    ),
+]

langgraph_api/asyncio.py

@@ -0,0 +1,201 @@
+import asyncio
+from collections.abc import AsyncIterator, Coroutine
+from contextlib import AbstractAsyncContextManager
+from functools import partial
+from typing import Any, Generic, TypeVar
+
+import structlog
+
+T = TypeVar("T")
+
+logger = structlog.stdlib.get_logger(__name__)
+
+
+async def sleep_if_not_done(delay: float, done: asyncio.Event) -> None:
+    try:
+        await asyncio.wait_for(done.wait(), delay)
+    except TimeoutError:
+        pass
+
+
+class ValueEvent(asyncio.Event):
+    def set(self, value: Any = True) -> None:
+        """Set the internal flag to true. All coroutines waiting for it to
+        become set are awakened. Coroutine that call wait() once the flag is
+        true will not block at all.
+        """
+        if not self._value:
+            self._value = value
+
+            for fut in self._waiters:
+                if not fut.done():
+                    fut.set_result(value)
+
+    async def wait(self):
+        """Block until the internal flag is set.
+
+        If the internal flag is set on entry, return value
+        immediately.  Otherwise, block until another coroutine calls
+        set() to set the flag, then return the value.
+        """
+        if self._value:
+            return self._value
+
+        fut = self._get_loop().create_future()
+        self._waiters.append(fut)
+        try:
+            return await fut
+        finally:
+            self._waiters.remove(fut)
+
+
+async def wait_if_not_done(coro: Coroutine[Any, Any, T], done: ValueEvent) -> T:
+    """Wait for the coroutine to finish or the event to be set."""
+    try:
+        async with asyncio.TaskGroup() as tg:
+            coro_task = tg.create_task(coro)
+            done_task = tg.create_task(done.wait())
+            coro_task.add_done_callback(lambda _: done_task.cancel())
+            done_task.add_done_callback(lambda _: coro_task.cancel(done._value))
+            try:
+                return await coro_task
+            except asyncio.CancelledError as e:
+                if e.args and isinstance(e.args[0], Exception):
+                    raise e.args[0] from None
+                raise
+    except ExceptionGroup as e:
+        raise e.exceptions[0] from None
+
+
+PENDING_TASKS = set()
+
+
+def _create_task_done_callback(
+    ignore_exceptions: tuple[Exception, ...], task: asyncio.Task
+) -> None:
+    PENDING_TASKS.remove(task)
+    try:
+        if exc := task.exception():
+            if not isinstance(exc, ignore_exceptions):
+                logger.exception("Background task failed", exc_info=exc)
+    except asyncio.CancelledError:
+        pass
+
+
+def create_task(
+    coro: Coroutine[Any, Any, T], ignore_exceptions: tuple[Exception, ...] = ()
+) -> asyncio.Task[T]:
+    """Create a new task in the current task group and return it."""
+    task = asyncio.create_task(coro)
+    PENDING_TASKS.add(task)
+    task.add_done_callback(partial(_create_task_done_callback, ignore_exceptions))
+    return task
+
+
+class SimpleTaskGroup(AbstractAsyncContextManager["SimpleTaskGroup"]):
+    """An async task group that can be configured to wait and/or cancel tasks on exit.
+
+    asyncio.TaskGroup and anyio.TaskGroup both expect enter and exit to be called
+    in the same asyncio task, which is not true for our use case, where exit is
+    shielded from cancellation."""
+
+    tasks: set[asyncio.Task]
+
+    def __init__(
+        self, *coros: Coroutine[Any, Any, T], cancel: bool = False, wait: bool = True
+    ) -> None:
+        self.tasks = set()
+        self.cancel = cancel
+        self.wait = wait
+        for coro in coros:
+            self.create_task(coro)
+
+    def _create_task_done_callback(
+        self, ignore_exceptions: tuple[Exception, ...], task: asyncio.Task
+    ) -> None:
+        try:
+            self.tasks.remove(task)
+        except AttributeError:
+            pass
+        try:
+            if exc := task.exception():
+                if not isinstance(exc, ignore_exceptions):
+                    logger.exception("Background task failed", exc_info=exc)
+        except asyncio.CancelledError:
+            pass
+
+    def create_task(
+        self,
+        coro: Coroutine[Any, Any, T],
+        ignore_exceptions: tuple[Exception, ...] = (),
+    ) -> asyncio.Task[T]:
+        """Create a new task in the current task group and return it."""
+        task = asyncio.create_task(coro)
+        self.tasks.add(task)
+        task.add_done_callback(
+            partial(self._create_task_done_callback, ignore_exceptions)
+        )
+        return task
+
+    async def __aexit__(self, exc_type, exc_value, traceback) -> None:
+        tasks = self.tasks
+        # break reference cycles between tasks and task group
+        del self.tasks
+        # cancel all tasks
+        if self.cancel:
+            for task in tasks:
+                task.cancel()
+        # wait for all tasks
+        if self.wait:
+            await asyncio.gather(*tasks, return_exceptions=True)
+
+
+def to_aiter(*args: T) -> AsyncIterator[T]:
+    async def agen():
+        for arg in args:
+            yield arg
+
+    return agen()
+
+
+V = TypeVar("V")
+
+
+class aclosing(Generic[V], AbstractAsyncContextManager):
+    """Async context manager for safely finalizing an asynchronously cleaned-up
+    resource such as an async generator, calling its ``aclose()`` method.
+
+    Code like this:
+
+        async with aclosing(<module>.fetch(<arguments>)) as agen:
+            <block>
+
+    is equivalent to this:
+
+        agen = <module>.fetch(<arguments>)
+        try:
+            <block>
+        finally:
+            await agen.aclose()
+
+    """
+
+    def __init__(self, thing: V):
+        self.thing = thing
+
+    async def __aenter__(self) -> V:
+        return self.thing
+
+    async def __aexit__(self, *exc_info):
+        await self.thing.aclose()
+
+
+async def aclosing_aiter(aiter: AsyncIterator[T]) -> AsyncIterator[T]:
+    if hasattr(aiter, "__aenter__"):
+        async with aiter:
+            async for item in aiter:
+                yield item
+    else:
+        async with aclosing(aiter):
+            async for item in aiter:
+                yield item

langgraph_api/auth/langsmith/backend.py

@@ -0,0 +1,67 @@
+from typing import NotRequired, TypedDict
+
+from starlette.authentication import (
+    AuthCredentials,
+    AuthenticationBackend,
+    AuthenticationError,
+    BaseUser,
+    SimpleUser,
+)
+from starlette.requests import HTTPConnection
+
+from langgraph_api.auth.langsmith.client import auth_client
+from langgraph_api.config import (
+    LANGSMITH_AUTH_VERIFY_TENANT_ID,
+    LANGSMITH_TENANT_ID,
+)
+
+
+class AuthDict(TypedDict):
+    organization_id: str
+    tenant_id: str
+    user_id: NotRequired[str]
+    user_email: NotRequired[str]
+
+
+class LangsmithAuthBackend(AuthenticationBackend):
+    async def authenticate(
+        self, conn: HTTPConnection
+    ) -> tuple[AuthCredentials, BaseUser] | None:
+        headers = [
+            ("Authorization", conn.headers.get("Authorization")),
+            ("X-Tenant-Id", conn.headers.get("x-tenant-id")),
+            ("X-Api-Key", conn.headers.get("x-api-key")),
+            ("X-Service-Key", conn.headers.get("x-service-key")),
+        ]
+        if not any(h[1] for h in headers):
+            raise AuthenticationError("Missing authentication headers")
+        async with auth_client() as auth:
+            if not LANGSMITH_AUTH_VERIFY_TENANT_ID and not conn.headers.get(
+                "x-api-key"
+            ):
+                # when LANGSMITH_AUTH_VERIFY_TENANT_ID is false, we allow
+                # any valid bearer token to pass through
+                # api key auth is always required to match the tenant id
+                res = await auth.get(
+                    "/auth/verify", headers=[h for h in headers if h[1] is not None]
+                )
+            else:
+                res = await auth.get(
+                    "/auth/public", headers=[h for h in headers if h[1] is not None]
+                )
+            if res.status_code == 401:
+                raise AuthenticationError("Invalid token")
+            elif res.status_code == 403:
+                raise AuthenticationError("Forbidden")
+            else:
+                res.raise_for_status()
+                auth_dict: AuthDict = res.json()
+
+            # If tenant id verification is disabled, the bearer token requests
+            # are not required to match the tenant id. Api key requests are
+            # always required to match the tenant id.
+            if LANGSMITH_AUTH_VERIFY_TENANT_ID or conn.headers.get("x-api-key"):
+                if auth_dict["tenant_id"] != LANGSMITH_TENANT_ID:
+                    raise AuthenticationError("Invalid tenant ID")
+
+        return AuthCredentials(["authenticated"]), SimpleUser(auth_dict.get("user_id"))

langgraph_api/auth/langsmith/client.py

@@ -0,0 +1,145 @@
+from collections.abc import AsyncGenerator
+from contextlib import asynccontextmanager
+from typing import Any
+
+import httpx
+from httpx._types import HeaderTypes, QueryParamTypes, RequestData
+from tenacity import retry
+from tenacity.retry import retry_if_exception
+from tenacity.stop import stop_after_attempt
+from tenacity.wait import wait_exponential_jitter
+
+from langgraph_api.config import LANGSMITH_AUTH_ENDPOINT
+
+_client: "JsonHttpClient"
+
+
+def is_retriable_error(exception: Exception) -> bool:
+    if isinstance(exception, httpx.TransportError):
+        return True
+    if isinstance(exception, httpx.HTTPStatusError):
+        if exception.response.status_code > 499:
+            return True
+
+    return False
+
+
+retry_httpx = retry(
+    reraise=True,
+    retry=retry_if_exception(is_retriable_error),
+    wait=wait_exponential_jitter(),
+    stop=stop_after_attempt(3),
+)
+
+
+class JsonHttpClient:
+    """HTTPX client for JSON requests, with retries."""
+
+    def __init__(self, client: httpx.AsyncClient) -> None:
+        """Initialize the auth client."""
+        self.client = client
+
+    async def _get(
+        self,
+        path: str,
+        *,
+        params: QueryParamTypes | None = None,
+        headers: HeaderTypes | None = None,
+    ) -> httpx.Response:
+        return await self.client.get(path, params=params, headers=headers)
+
+    @retry_httpx
+    async def get(
+        self,
+        path: str,
+        *,
+        params: QueryParamTypes | None = None,
+        headers: HeaderTypes | None = None,
+    ) -> httpx.Response:
+        return await self.client.get(path, params=params, headers=headers)
+
+    async def _post(
+        self,
+        path: str,
+        *,
+        data: RequestData | None = None,
+        json: Any | None = None,
+        params: QueryParamTypes | None = None,
+        headers: HeaderTypes | None = None,
+    ) -> httpx.Response:
+        return await self.client.post(
+            path, data=data, json=json, params=params, headers=headers
+        )
+
+    @retry_httpx
+    async def post(
+        self,
+        path: str,
+        *,
+        data: RequestData | None = None,
+        json: Any | None = None,
+        params: QueryParamTypes | None = None,
+        headers: HeaderTypes | None = None,
+    ) -> httpx.Response:
+        return await self.client.post(
+            path, data=data, json=json, params=params, headers=headers
+        )
+
+
+def create_client() -> JsonHttpClient:
+    """Create the auth http client."""
+    return JsonHttpClient(
+        httpx.AsyncClient(
+            transport=httpx.AsyncHTTPTransport(
+                retries=5,  # this applies only to ConnectError, ConnectTimeout
+                limits=httpx.Limits(
+                    max_keepalive_connections=40,
+                    keepalive_expiry=240.0,
+                ),
+            ),
+            timeout=httpx.Timeout(2.0),
+            base_url=LANGSMITH_AUTH_ENDPOINT,
+        )
+    )
+
+
+async def close_auth_client() -> None:
+    """Close the auth http client."""
+    global _client
+    try:
+        await _client.client.aclose()
+    except NameError:
+        pass
+
+
+async def initialize_auth_client() -> None:
+    """Initialize the auth http client."""
+    await close_auth_client()
+    global _client
+    _client = create_client()
+
+
+@asynccontextmanager
+async def auth_client() -> AsyncGenerator[JsonHttpClient, None]:
+    """Get the auth http client."""
+    # pytest does something funny with event loops,
+    # so we can't use a global pool for tests
+    if LANGSMITH_AUTH_ENDPOINT.startswith("http://localhost"):
+        client = create_client()
+        try:
+            yield client
+        finally:
+            await client.client.aclose()
+    else:
+        try:
+            if not _client.client.is_closed:
+                found = True
+            else:
+                found = False
+        except NameError:
+            found = False
+        if found:
+            yield _client
+        else:
+            await initialize_auth_client()
+            yield _client