langgraph-api 0.4.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

langgraph_api/api/store.py

@@ -5,7 +5,13 @@ from starlette.responses import Response
 from starlette.routing import BaseRoute
 
 from langgraph_api.auth.custom import handle_event as _handle_event
+from langgraph_api.encryption.middleware import (
+    decrypt_response,
+    decrypt_responses,
+    encrypt_request,
+)
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
+from langgraph_api.schema import STORE_ENCRYPTION_FIELDS
 from langgraph_api.store import get_store
 from langgraph_api.utils import get_auth_ctx
 from langgraph_api.validation import (
@@ -48,6 +54,7 @@ async def handle_event(
 async def put_item(request: ApiRequest):
     """Store or update an item."""
     payload = await request.json(StorePutRequest)
+    payload = await encrypt_request(payload, "store", STORE_ENCRYPTION_FIELDS)
     namespace = tuple(payload["namespace"]) if payload.get("namespace") else ()
     if err := _validate_namespace(namespace):
         return err
@@ -78,7 +85,11 @@ async def get_item(request: ApiRequest):
     }
     await handle_event("get", handler_payload)
     result = await (await get_store()).aget(namespace, key)
-    return ApiResponse(result.dict() if result is not None else None)
+    if result is None:
+        return ApiResponse(None)
+    return ApiResponse(
+        await decrypt_response(result.dict(), "store", STORE_ENCRYPTION_FIELDS)
+    )
 
 
 @retry_db
@@ -125,7 +136,13 @@ async def search_items(request: ApiRequest):
         offset=handler_payload["offset"],
         query=handler_payload["query"],
     )
-    return ApiResponse({"items": [item.dict() for item in items]})
+    return ApiResponse(
+        {
+            "items": await decrypt_responses(
+                [item.dict() for item in items], "store", STORE_ENCRYPTION_FIELDS
+            )
+        }
+    )
 
 
 @retry_db

langgraph_api/api/threads.py

@@ -1,11 +1,23 @@
+from typing import get_args
 from uuid import uuid4
 
 from starlette.exceptions import HTTPException
 from starlette.responses import Response
 from starlette.routing import BaseRoute
 
+from langgraph_api.encryption.middleware import (
+    decrypt_response,
+    decrypt_responses,
+    encrypt_request,
+)
+from langgraph_api.feature_flags import FF_USE_CORE_API
+from langgraph_api.grpc.ops import Threads as GrpcThreads
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
-from langgraph_api.schema import THREAD_FIELDS
+from langgraph_api.schema import (
+    THREAD_ENCRYPTION_FIELDS,
+    THREAD_FIELDS,
+    ThreadStreamMode,
+)
 from langgraph_api.sse import EventSourceResponse
 from langgraph_api.state import state_snapshot_to_thread_state
 from langgraph_api.utils import (
@@ -15,10 +27,12 @@ from langgraph_api.utils import (
     validate_stream_id,
     validate_uuid,
 )
+from langgraph_api.utils.headers import get_configurable_headers
 from langgraph_api.validation import (
     ThreadCountRequest,
     ThreadCreate,
     ThreadPatch,
+    ThreadPruneRequest,
     ThreadSearchRequest,
     ThreadStateCheckpointRequest,
     ThreadStateSearch,
@@ -28,6 +42,8 @@ from langgraph_runtime.database import connect
 from langgraph_runtime.ops import Threads
 from langgraph_runtime.retry import retry_db
 
+CrudThreads = GrpcThreads if FF_USE_CORE_API else Threads
+
 
 @retry_db
 async def create_thread(
@@ -37,28 +53,56 @@ async def create_thread(
     payload = await request.json(ThreadCreate)
     if thread_id := payload.get("thread_id"):
         validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+
+    # Validate keep_latest TTL requires core API
+    ttl = payload.get("ttl")
+    if ttl and ttl.get("strategy") == "keep_latest" and not FF_USE_CORE_API:
+        raise HTTPException(
+            status_code=422,
+            detail="keep_latest TTL strategy requires FF_USE_CORE_API=true",
+        )
+
+    # Encrypt metadata before storing
+    encrypted_payload = await encrypt_request(
+        payload,
+        "thread",
+        ["metadata"],
+    )
+
     async with connect() as conn:
         thread_id = thread_id or str(uuid4())
-        iter = await Threads.put(
+        iter = await CrudThreads.put(
             conn,
             thread_id,
-            metadata=payload.get("metadata"),
+            metadata=encrypted_payload.get("metadata") or {},
             if_exists=payload.get("if_exists") or "raise",
             ttl=payload.get("ttl"),
         )
-
+        config = {
+            "configurable": {
+                **get_configurable_headers(request.headers),
+                "thread_id": thread_id,
+            }
+        }
         if supersteps := payload.get("supersteps"):
             try:
                 await Threads.State.bulk(
                     conn,
-                    config={"configurable": {"thread_id": thread_id}},
+                    config=config,
                     supersteps=supersteps,
                 )
             except HTTPException as e:
                 detail = f"Thread {thread_id} was created, but there were problems updating the state: {e.detail}"
                 raise HTTPException(status_code=201, detail=detail) from e
 
-    return ApiResponse(await fetchone(iter, not_found_code=409))
+    # Decrypt thread fields in response
+    thread = await fetchone(iter, not_found_code=409)
+    thread = await decrypt_response(
+        thread,
+        "thread",
+        THREAD_ENCRYPTION_FIELDS,
+    )
+    return ApiResponse(thread)
 
 
 @retry_db
@@ -70,12 +114,14 @@ async def search_threads(
     select = validate_select_columns(payload.get("select") or None, THREAD_FIELDS)
     limit = int(payload.get("limit") or 10)
     offset = int(payload.get("offset") or 0)
+
     async with connect() as conn:
-        threads_iter, next_offset = await Threads.search(
+        threads_iter, next_offset = await CrudThreads.search(
             conn,
             status=payload.get("status"),
             values=payload.get("values"),
             metadata=payload.get("metadata"),
+            ids=payload.get("ids"),
             limit=limit,
             offset=offset,
             sort_by=payload.get("sort_by"),
@@ -85,7 +131,15 @@ async def search_threads(
     threads, response_headers = await get_pagination_headers(
         threads_iter, next_offset, offset
     )
-    return ApiResponse(threads, headers=response_headers)
+
+    # Decrypt metadata, values, interrupts, and error in all returned threads
+    decrypted_threads = await decrypt_responses(
+        threads,
+        "thread",
+        THREAD_ENCRYPTION_FIELDS,
+    )
+
+    return ApiResponse(decrypted_threads, headers=response_headers)
 
 
 @retry_db
@@ -95,7 +149,7 @@ async def count_threads(
     """Count threads."""
     payload = await request.json(ThreadCountRequest)
     async with connect() as conn:
-        count = await Threads.count(
+        count = await CrudThreads.count(
             conn,
             status=payload.get("status"),
             values=payload.get("values"),
@@ -113,10 +167,14 @@ async def get_thread_state(
     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:
+        config = {
+            "configurable": {
+                **get_configurable_headers(request.headers),
+                "thread_id": thread_id,
+            }
+        }
         state = state_snapshot_to_thread_state(
-            await Threads.State.get(
-                conn, {"configurable": {"thread_id": thread_id}}, subgraphs=subgraphs
-            )
+            await Threads.State.get(conn, config=config, subgraphs=subgraphs)
         )
     return ApiResponse(state)
 
@@ -130,15 +188,17 @@ async def get_thread_state_at_checkpoint(
     validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
     checkpoint_id = request.path_params["checkpoint_id"]
     async with connect() as conn:
+        config = {
+            "configurable": {
+                **get_configurable_headers(request.headers),
+                "thread_id": thread_id,
+                "checkpoint_id": checkpoint_id,
+            }
+        }
         state = state_snapshot_to_thread_state(
             await Threads.State.get(
                 conn,
-                {
-                    "configurable": {
-                        "thread_id": thread_id,
-                        "checkpoint_id": checkpoint_id,
-                    }
-                },
+                config=config,
                 subgraphs=request.query_params.get("subgraphs") in ("true", "True"),
             )
         )
@@ -154,10 +214,17 @@ async def get_thread_state_at_checkpoint_post(
     validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
     payload = await request.json(ThreadStateCheckpointRequest)
     async with connect() as conn:
+        config = {
+            "configurable": {
+                **payload["checkpoint"],
+                **get_configurable_headers(request.headers),
+                "thread_id": thread_id,
+            }
+        }
         state = state_snapshot_to_thread_state(
             await Threads.State.get(
                 conn,
-                {"configurable": {"thread_id": thread_id, **payload["checkpoint"]}},
+                config=config,
                 subgraphs=payload.get("subgraphs", False),
             )
         )
@@ -182,6 +249,7 @@ async def update_thread_state(
             config["configurable"]["user_id"] = user_id
     except AssertionError:
         pass
+    config["configurable"].update(get_configurable_headers(request.headers))
     async with connect() as conn:
         inserted = await Threads.State.post(
             conn,
@@ -205,7 +273,13 @@ async def get_thread_history(
     except ValueError:
         raise HTTPException(status_code=422, detail=f"Invalid limit {limit_}") from None
     before = request.query_params.get("before")
-    config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
+    config = {
+        "configurable": {
+            "thread_id": thread_id,
+            "checkpoint_ns": "",
+            **get_configurable_headers(request.headers),
+        }
+    }
     async with connect() as conn:
         states = [
             state_snapshot_to_thread_state(c)
@@ -226,6 +300,7 @@ async def get_thread_history_post(
     payload = await request.json(ThreadStateSearch)
     config = {"configurable": {"thread_id": thread_id, "checkpoint_ns": ""}}
     config["configurable"].update(payload.get("checkpoint", {}))
+    config["configurable"].update(get_configurable_headers(request.headers))
     async with connect() as conn:
         states = [
             state_snapshot_to_thread_state(c)
@@ -247,9 +322,23 @@ async def get_thread(
     """Get a thread by ID."""
     thread_id = request.path_params["thread_id"]
     validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
+
+    # Parse include parameter for optional fields (e.g., ttl)
+    include_param = request.query_params.get("include", "")
+    include_fields = [f.strip() for f in include_param.split(",") if f.strip()]
+    include_ttl = "ttl" in include_fields
+
     async with connect() as conn:
-        thread = await Threads.get(conn, thread_id)
-    return ApiResponse(await fetchone(thread))
+        thread = await CrudThreads.get(conn, thread_id, include_ttl=include_ttl)
+
+    # Decrypt metadata, values, interrupts, and error in response
+    thread_data = await fetchone(thread)
+    thread_data = await decrypt_response(
+        thread_data,
+        "thread",
+        THREAD_ENCRYPTION_FIELDS,
+    )
+    return ApiResponse(thread_data)
 
 
 @retry_db
@@ -260,9 +349,37 @@ async def patch_thread(
     thread_id = request.path_params["thread_id"]
     validate_uuid(thread_id, "Invalid thread ID: must be a UUID")
     payload = await request.json(ThreadPatch)
+
+    # Validate keep_latest TTL requires core API
+    ttl = payload.get("ttl")
+    if ttl and ttl.get("strategy") == "keep_latest" and not FF_USE_CORE_API:
+        raise HTTPException(
+            status_code=422,
+            detail="keep_latest TTL strategy requires FF_USE_CORE_API=true",
+        )
+
+    # Encrypt metadata before storing
+    encrypted_payload = await encrypt_request(
+        payload,
+        "thread",
+        ["metadata"],
+    )
+
     async with connect() as conn:
-        thread = await Threads.patch(conn, thread_id, metadata=payload["metadata"])
-    return ApiResponse(await fetchone(thread))
+        thread = await CrudThreads.patch(
+            conn,
+            thread_id,
+            metadata=encrypted_payload.get("metadata") or {},
+            ttl=payload.get("ttl"),
+        )
+    thread_data = await fetchone(thread)
+    # Decrypt metadata, values, interrupts, and error in response
+    thread_data = await decrypt_response(
+        thread_data,
+        "thread",
+        THREAD_ENCRYPTION_FIELDS,
+    )
+    return ApiResponse(thread_data)
 
 
 @retry_db
@@ -271,17 +388,60 @@ async def delete_thread(request: ApiRequest):
     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)
+        tid = await CrudThreads.delete(conn, thread_id)
     await fetchone(tid)
     return Response(status_code=204)
 
 
+@retry_db
+async def prune_threads(request: ApiRequest):
+    """Prune threads by ID."""
+    payload = await request.json(ThreadPruneRequest)
+    thread_ids = payload.get("thread_ids", [])
+    strategy = payload.get("strategy", "delete")
+
+    # Validate each thread_id is a valid UUID
+    for tid in thread_ids:
+        validate_uuid(tid, "Invalid thread ID: must be a UUID")
+
+    # Validate strategy
+    if strategy not in ("delete", "keep_latest"):
+        raise HTTPException(
+            status_code=422,
+            detail=f"Invalid strategy: {strategy}. Expected 'delete' or 'keep_latest'.",
+        )
+
+    # Empty list is a no-op, return early
+    if not thread_ids:
+        return ApiResponse({"pruned_count": 0})
+
+    if not FF_USE_CORE_API:
+        raise HTTPException(
+            status_code=422,
+            detail="Thread prune requires FF_USE_CORE_API=true",
+        )
+
+    pruned_count = await CrudThreads.prune(
+        thread_ids=thread_ids,
+        strategy=strategy,
+    )
+
+    return ApiResponse({"pruned_count": pruned_count})
+
+
 @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))
+        iter = await CrudThreads.copy(conn, thread_id)
+    thread_data = await fetchone(iter, not_found_code=409)
+    # Decrypt metadata, values, interrupts, and error in response
+    thread_data = await decrypt_response(
+        thread_data,
+        "thread",
+        THREAD_ENCRYPTION_FIELDS,
+    )
+    return ApiResponse(thread_data)
 
 
 @retry_db
@@ -293,10 +453,31 @@ async def join_thread_stream(request: ApiRequest):
     validate_stream_id(
         last_event_id, "Invalid last-event-id: must be a valid Redis stream ID"
     )
+
+    # Parse stream_modes parameter - can be single string or comma-separated list
+    stream_modes_param = request.query_params.get("stream_modes")
+    if stream_modes_param:
+        if "," in stream_modes_param:
+            # Handle comma-separated list
+            stream_modes = [mode.strip() for mode in stream_modes_param.split(",")]
+        else:
+            # Handle single value
+            stream_modes = [stream_modes_param]
+        # Validate each mode
+        for mode in stream_modes:
+            if mode not in get_args(ThreadStreamMode):
+                raise HTTPException(
+                    status_code=422, detail=f"Invalid stream mode: {mode}"
+                )
+    else:
+        # Default to run_modes
+        stream_modes = ["run_modes"]
+
     return EventSourceResponse(
         Threads.Stream.join(
             thread_id,
             last_event_id=last_event_id,
+            stream_modes=stream_modes,
         ),
     )
 
@@ -305,6 +486,7 @@ threads_routes: list[BaseRoute] = [
     ApiRoute("/threads", endpoint=create_thread, methods=["POST"]),
     ApiRoute("/threads/search", endpoint=search_threads, methods=["POST"]),
     ApiRoute("/threads/count", endpoint=count_threads, methods=["POST"]),
+    ApiRoute("/threads/prune", endpoint=prune_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"]),

langgraph_api/asgi_transport.py

@@ -25,7 +25,7 @@ def is_running_trio() -> bool:
         # sniffio is a dependency of trio.
 
         # See https://github.com/python-trio/trio/issues/2802
-        import sniffio
+        import sniffio  # type: ignore[unresolved-import]
 
         if sniffio.current_async_library() == "trio":
             return True
@@ -84,7 +84,8 @@ class ASGITransport(ASGITransportBase):
     ) -> Response:
         from langgraph_api.asyncio import call_soon_in_main_loop
 
-        assert isinstance(request.stream, AsyncByteStream)
+        if not isinstance(request.stream, AsyncByteStream):
+            raise ValueError("Request stream must be an AsyncByteStream")
 
         # ASGI scope.
         scope = {
@@ -133,14 +134,15 @@ class ASGITransport(ASGITransportBase):
             nonlocal status_code, response_headers, response_started
 
             if message["type"] == "http.response.start":
-                assert not response_started
-
+                if response_started:
+                    raise RuntimeError("Response already 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()
+                if response_complete.is_set():
+                    raise RuntimeError("Response already complete")
                 body = message.get("body", b"")
                 more_body = message.get("more_body", False)
 
@@ -152,7 +154,7 @@ class ASGITransport(ASGITransportBase):
 
         try:
             await call_soon_in_main_loop(self.app(scope, receive, send))
-        except Exception:  # noqa: PIE-786
+        except Exception:
             if self.raise_app_exceptions:
                 raise
 
@@ -162,9 +164,12 @@ class ASGITransport(ASGITransportBase):
             if response_headers is None:
                 response_headers = {}
 
-        assert response_complete.is_set()
-        assert status_code is not None
-        assert response_headers is not None
+        if not response_complete.is_set():
+            raise RuntimeError("Response not complete")
+        if status_code is None:
+            raise RuntimeError("Status code not set")
+        if response_headers is None:
+            raise RuntimeError("Response headers not set")
 
         stream = ASGIResponseStream(body_parts)
 

langgraph_api/asyncio.py

@@ -115,11 +115,16 @@ def create_task(
 
 
 def run_coroutine_threadsafe(
-    coro: Coroutine[Any, Any, T], ignore_exceptions: tuple[type[Exception], ...] = ()
+    coro: Coroutine[Any, Any, T],
+    ignore_exceptions: tuple[type[Exception], ...] = (),
+    *,
+    loop: asyncio.AbstractEventLoop | None = None,
 ) -> concurrent.futures.Future[T] | concurrent.futures.Future[None]:
-    if _MAIN_LOOP is None:
+    if loop is None:
+        loop = _MAIN_LOOP
+    if loop is None:
         raise RuntimeError("No event loop set")
-    future = asyncio.run_coroutine_threadsafe(coro, _MAIN_LOOP)
+    future = asyncio.run_coroutine_threadsafe(coro, loop)
     future.add_done_callback(partial(_create_task_done_callback, ignore_exceptions))
     return future
 
@@ -158,12 +163,15 @@ class SimpleTaskGroup(AbstractAsyncContextManager["SimpleTaskGroup"]):
         self,
         *coros: Coroutine[Any, Any, T],
         cancel: bool = False,
+        cancel_event: asyncio.Event | None = None,
         wait: bool = True,
         taskset: set[asyncio.Task] | None = None,
         taskgroup_name: str | None = None,
     ) -> None:
-        self.tasks = taskset if taskset is not None else set()
+        # Copy the taskset to avoid modifying the original set unintentionally (like in lifespan)
+        self.tasks = taskset.copy() if taskset is not None else set()
         self.cancel = cancel
+        self.cancel_event = cancel_event
         self.wait = wait
         if taskset:
             for task in tuple(taskset):
@@ -180,6 +188,8 @@ class SimpleTaskGroup(AbstractAsyncContextManager["SimpleTaskGroup"]):
         try:
             if (exc := task.exception()) and not isinstance(exc, ignore_exceptions):
                 logger.exception("asyncio.task failed in task group", exc_info=exc)
+                if self.cancel_event:
+                    self.cancel_event.set()
         except asyncio.CancelledError:
             pass