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

langgraph_api/api/assistants.py

@@ -24,6 +24,7 @@ from langgraph_api.utils import (
     validate_select_columns,
     validate_uuid,
 )
+from langgraph_api.utils.headers import get_configurable_headers
 from langgraph_api.validation import (
     AssistantCountRequest,
     AssistantCreate,
@@ -240,6 +241,9 @@ async def get_assistant_graph(
         assistant_ = await Assistants.get(conn, assistant_id)
         assistant = await fetchone(assistant_)
     config = await ajson_loads(assistant["config"])
+    configurable = config.setdefault("configurable", {})
+    configurable.update(get_configurable_headers(request.headers))
+
     async with get_graph(
         assistant["graph_id"],
         config,
@@ -294,6 +298,8 @@ async def get_assistant_subgraphs(
         assistant_ = await Assistants.get(conn, assistant_id)
         assistant = await fetchone(assistant_)
         config = await ajson_loads(assistant["config"])
+        configurable = config.setdefault("configurable", {})
+        configurable.update(get_configurable_headers(request.headers))
         async with get_graph(
             assistant["graph_id"],
             config,
@@ -340,6 +346,8 @@ async def get_assistant_schemas(
         # TODO Implementa  cache so we can de-dent and release this connection.
         assistant = await fetchone(assistant_)
         config = await ajson_loads(assistant["config"])
+        configurable = config.setdefault("configurable", {})
+        configurable.update(get_configurable_headers(request.headers))
         async with get_graph(
             assistant["graph_id"],
             config,

langgraph_api/api/runs.py

@@ -4,11 +4,12 @@ from typing import Literal, cast
 from uuid import uuid4
 
 import orjson
+import structlog
 from starlette.exceptions import HTTPException
 from starlette.responses import Response, StreamingResponse
 
 from langgraph_api import config
-from langgraph_api.asyncio import ValueEvent, aclosing
+from langgraph_api.asyncio import ValueEvent
 from langgraph_api.models.run import create_valid_run
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
 from langgraph_api.schema import CRON_FIELDS, RUN_FIELDS
@@ -34,6 +35,8 @@ from langgraph_runtime.database import connect
 from langgraph_runtime.ops import Crons, Runs, Threads
 from langgraph_runtime.retry import retry_db
 
+logger = structlog.stdlib.get_logger(__name__)
+
 
 @retry_db
 async def create_run(request: ApiRequest):
@@ -101,9 +104,7 @@ async def stream_run(
     payload = await request.json(RunCreateStateful)
     on_disconnect = payload.get("on_disconnect", "continue")
     run_id = uuid7()
-    sub = asyncio.create_task(Runs.Stream.subscribe(run_id, thread_id))
-
-    try:
+    async with await Runs.Stream.subscribe(run_id, thread_id) as sub:
         async with connect() as conn:
             run = await create_valid_run(
                 conn,
@@ -113,25 +114,20 @@ async def stream_run(
                 run_id=run_id,
                 request_start_time=request.scope.get("request_start_time_ms"),
             )
-    except Exception:
-        if not sub.cancelled():
-            handle = await sub
-            await handle.__aexit__(None, None, None)
-        raise
 
-    return EventSourceResponse(
-        Runs.Stream.join(
-            run["run_id"],
-            thread_id=thread_id,
-            cancel_on_disconnect=on_disconnect == "cancel",
-            stream_channel=await sub,
-            last_event_id=None,
-        ),
-        headers={
-            "Location": f"/threads/{thread_id}/runs/{run['run_id']}/stream",
-            "Content-Location": f"/threads/{thread_id}/runs/{run['run_id']}",
-        },
-    )
+        return EventSourceResponse(
+            Runs.Stream.join(
+                run["run_id"],
+                thread_id=thread_id,
+                cancel_on_disconnect=on_disconnect == "cancel",
+                stream_channel=sub,
+                last_event_id=None,
+            ),
+            headers={
+                "Location": f"/threads/{thread_id}/runs/{run['run_id']}/stream",
+                "Content-Location": f"/threads/{thread_id}/runs/{run['run_id']}",
+            },
+        )
 
 
 async def stream_run_stateless(
@@ -143,8 +139,7 @@ async def stream_run_stateless(
     on_disconnect = payload.get("on_disconnect", "continue")
     run_id = uuid7()
     thread_id = uuid4()
-    sub = asyncio.create_task(Runs.Stream.subscribe(run_id, thread_id))
-    try:
+    async with await Runs.Stream.subscribe(run_id, thread_id) as sub:
         async with connect() as conn:
             run = await create_valid_run(
                 conn,
@@ -155,26 +150,21 @@ async def stream_run_stateless(
                 request_start_time=request.scope.get("request_start_time_ms"),
                 temporary=True,
             )
-    except Exception:
-        if not sub.cancelled():
-            handle = await sub
-            await handle.__aexit__(None, None, None)
-        raise
 
-    return EventSourceResponse(
-        Runs.Stream.join(
-            run["run_id"],
-            thread_id=run["thread_id"],
-            ignore_404=True,
-            cancel_on_disconnect=on_disconnect == "cancel",
-            stream_channel=await sub,
-            last_event_id=None,
-        ),
-        headers={
-            "Location": f"/runs/{run['run_id']}/stream",
-            "Content-Location": f"/runs/{run['run_id']}",
-        },
-    )
+        return EventSourceResponse(
+            Runs.Stream.join(
+                run["run_id"],
+                thread_id=run["thread_id"],
+                ignore_404=True,
+                cancel_on_disconnect=on_disconnect == "cancel",
+                stream_channel=sub,
+                last_event_id=None,
+            ),
+            headers={
+                "Location": f"/runs/{run['run_id']}/stream",
+                "Content-Location": f"/runs/{run['run_id']}",
+            },
+        )
 
 
 @retry_db
@@ -184,9 +174,7 @@ async def wait_run(request: ApiRequest):
     payload = await request.json(RunCreateStateful)
     on_disconnect = payload.get("on_disconnect", "continue")
     run_id = uuid7()
-    sub = asyncio.create_task(Runs.Stream.subscribe(run_id, thread_id))
-
-    try:
+    async with await Runs.Stream.subscribe(run_id, thread_id) as sub:
         async with connect() as conn:
             run = await create_valid_run(
                 conn,
@@ -196,25 +184,17 @@ async def wait_run(request: ApiRequest):
                 run_id=run_id,
                 request_start_time=request.scope.get("request_start_time_ms"),
             )
-    except Exception:
-        if not sub.cancelled():
-            handle = await sub
-            await handle.__aexit__(None, None, None)
-        raise
 
-    last_chunk = ValueEvent()
+        last_chunk = ValueEvent()
 
-    async def consume():
-        vchunk: bytes | None = None
-        async with aclosing(
-            Runs.Stream.join(
+        async def consume():
+            vchunk: bytes | None = None
+            async for mode, chunk, _ in Runs.Stream.join(
                 run["run_id"],
                 thread_id=run["thread_id"],
-                stream_channel=await sub,
+                stream_channel=sub,
                 cancel_on_disconnect=on_disconnect == "cancel",
-            )
-        ) as stream:
-            async for mode, chunk, _ in stream:
+            ):
                 if (
                     mode == b"values"
                     or mode == b"updates"
@@ -223,43 +203,47 @@ async def wait_run(request: ApiRequest):
                     vchunk = chunk
                 elif mode == b"error":
                     vchunk = orjson.dumps({"__error__": orjson.Fragment(chunk)})
-        if vchunk is not None:
-            last_chunk.set(vchunk)
-        else:
-            async with connect() as conn:
-                thread_iter = await Threads.get(conn, thread_id)
+            if vchunk is not None:
+                last_chunk.set(vchunk)
+            else:
+                async with connect() as conn:
+                    thread_iter = await Threads.get(conn, thread_id)
+                    try:
+                        thread = await anext(thread_iter)
+                        last_chunk.set(thread["values"])
+                    except StopAsyncIteration:
+                        await logger.awarning(
+                            f"No checkpoint found for thread {thread_id}",
+                            thread_id=thread_id,
+                        )
+                        last_chunk.set(b"{}")
+
+        # keep the connection open by sending whitespace every 5 seconds
+        # leading whitespace will be ignored by json parsers
+        async def body() -> AsyncIterator[bytes]:
+            stream = asyncio.create_task(consume())
+            while True:
                 try:
-                    thread = await anext(thread_iter)
-                    last_chunk.set(thread["values"])
-                except StopAsyncIteration:
-                    last_chunk.set(b"{}")
-
-    # keep the connection open by sending whitespace every 5 seconds
-    # leading whitespace will be ignored by json parsers
-    async def body() -> AsyncIterator[bytes]:
-        stream = asyncio.create_task(consume())
-        while True:
-            try:
-                if stream.done():
-                    # raise stream exception if any
-                    stream.result()
-                yield await asyncio.wait_for(last_chunk.wait(), timeout=5)
-                break
-            except TimeoutError:
-                yield b"\n"
-            except asyncio.CancelledError:
-                stream.cancel()
-                await stream
-                raise
-
-    return StreamingResponse(
-        body(),
-        media_type="application/json",
-        headers={
-            "Location": f"/threads/{thread_id}/runs/{run['run_id']}/join",
-            "Content-Location": f"/threads/{thread_id}/runs/{run['run_id']}",
-        },
-    )
+                    if stream.done():
+                        # raise stream exception if any
+                        stream.result()
+                    yield await asyncio.wait_for(last_chunk.wait(), timeout=5)
+                    break
+                except TimeoutError:
+                    yield b"\n"
+                except asyncio.CancelledError:
+                    stream.cancel()
+                    await stream
+                    raise
+
+        return StreamingResponse(
+            body(),
+            media_type="application/json",
+            headers={
+                "Location": f"/threads/{thread_id}/runs/{run['run_id']}/join",
+                "Content-Location": f"/threads/{thread_id}/runs/{run['run_id']}",
+            },
+        )
 
 
 @retry_db
@@ -270,9 +254,7 @@ async def wait_run_stateless(request: ApiRequest):
     on_disconnect = payload.get("on_disconnect", "continue")
     run_id = uuid7()
     thread_id = uuid4()
-    sub = asyncio.create_task(Runs.Stream.subscribe(run_id, thread_id))
-
-    try:
+    async with await Runs.Stream.subscribe(run_id, thread_id) as sub:
         async with connect() as conn:
             run = await create_valid_run(
                 conn,
@@ -283,25 +265,18 @@ async def wait_run_stateless(request: ApiRequest):
                 request_start_time=request.scope.get("request_start_time_ms"),
                 temporary=True,
             )
-    except Exception:
-        if not sub.cancelled():
-            handle = await sub
-            await handle.__aexit__(None, None, None)
-        raise
-    last_chunk = ValueEvent()
-
-    async def consume():
-        vchunk: bytes | None = None
-        async with aclosing(
-            Runs.Stream.join(
+
+        last_chunk = ValueEvent()
+
+        async def consume():
+            vchunk: bytes | None = None
+            async for mode, chunk, _ in Runs.Stream.join(
                 run["run_id"],
                 thread_id=run["thread_id"],
-                stream_channel=await sub,
+                stream_channel=sub,
                 ignore_404=True,
                 cancel_on_disconnect=on_disconnect == "cancel",
-            )
-        ) as stream:
-            async for mode, chunk, _ in stream:
+            ):
                 if (
                     mode == b"values"
                     or mode == b"updates"
@@ -310,38 +285,43 @@ async def wait_run_stateless(request: ApiRequest):
                     vchunk = chunk
                 elif mode == b"error":
                     vchunk = orjson.dumps({"__error__": orjson.Fragment(chunk)})
-        if vchunk is not None:
-            last_chunk.set(vchunk)
-        else:
-            # we can't fetch the thread (it was deleted), so just return empty values
-            last_chunk.set(b"{}")
-
-    # keep the connection open by sending whitespace every 5 seconds
-    # leading whitespace will be ignored by json parsers
-    async def body() -> AsyncIterator[bytes]:
-        stream = asyncio.create_task(consume())
-        while True:
-            try:
-                if stream.done():
-                    # raise stream exception if any
-                    stream.result()
-                yield await asyncio.wait_for(last_chunk.wait(), timeout=5)
-                break
-            except TimeoutError:
-                yield b"\n"
-            except asyncio.CancelledError:
-                stream.cancel("Run stream cancelled")
-                await stream
-                raise
-
-    return StreamingResponse(
-        body(),
-        media_type="application/json",
-        headers={
-            "Location": f"/threads/{run['thread_id']}/runs/{run['run_id']}/join",
-            "Content-Location": f"/threads/{run['thread_id']}/runs/{run['run_id']}",
-        },
-    )
+            if vchunk is not None:
+                last_chunk.set(vchunk)
+            else:
+                # we can't fetch the thread (it was deleted), so just return empty values
+                await logger.awarning(
+                    "No checkpoint emitted for stateless run",
+                    run_id=run["run_id"],
+                    thread_id=run["thread_id"],
+                )
+                last_chunk.set(b"{}")
+
+        # keep the connection open by sending whitespace every 5 seconds
+        # leading whitespace will be ignored by json parsers
+        async def body() -> AsyncIterator[bytes]:
+            stream = asyncio.create_task(consume())
+            while True:
+                try:
+                    if stream.done():
+                        # raise stream exception if any
+                        stream.result()
+                    yield await asyncio.wait_for(last_chunk.wait(), timeout=5)
+                    break
+                except TimeoutError:
+                    yield b"\n"
+                except asyncio.CancelledError:
+                    stream.cancel("Run stream cancelled")
+                    await stream
+                    raise
+
+        return StreamingResponse(
+            body(),
+            media_type="application/json",
+            headers={
+                "Location": f"/threads/{run['thread_id']}/runs/{run['run_id']}/join",
+                "Content-Location": f"/threads/{run['thread_id']}/runs/{run['run_id']}",
+            },
+        )
 
 
 @retry_db

langgraph_api/api/threads.py

@@ -1,3 +1,4 @@
+from typing import get_args
 from uuid import uuid4
 
 from starlette.exceptions import HTTPException
@@ -5,7 +6,7 @@ from starlette.responses import Response
 from starlette.routing import BaseRoute
 
 from langgraph_api.route import ApiRequest, ApiResponse, ApiRoute
-from langgraph_api.schema import THREAD_FIELDS
+from langgraph_api.schema import 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,6 +16,7 @@ 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,
@@ -46,12 +48,17 @@ async def create_thread(
             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:
@@ -76,6 +83,7 @@ async def search_threads(
             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"),
@@ -113,10 +121,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 +142,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 +168,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 +203,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 +227,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 +254,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)
@@ -293,10 +322,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,
         ),
     )
 

langgraph_api/asyncio.py

@@ -162,7 +162,8 @@ class SimpleTaskGroup(AbstractAsyncContextManager["SimpleTaskGroup"]):
         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.wait = wait
         if taskset:

langgraph_api/feature_flags.py

@@ -6,3 +6,4 @@ LANGGRAPH_PY_MINOR = tuple(map(int, __version__.split(".")[:2]))
 OMIT_PENDING_SENDS = LANGGRAPH_PY_MINOR >= (0, 5)
 USE_RUNTIME_CONTEXT_API = LANGGRAPH_PY_MINOR >= (0, 6)
 USE_NEW_INTERRUPTS = LANGGRAPH_PY_MINOR >= (0, 6)
+USE_DURABILITY = LANGGRAPH_PY_MINOR >= (0, 6)

langgraph_api/logging.py

@@ -69,9 +69,12 @@ class AddApiVersion:
     def __call__(
         self, logger: logging.Logger, method_name: str, event_dict: EventDict
     ) -> EventDict:
-        from langgraph_api import __version__
+        try:
+            from langgraph_api import __version__
 
-        event_dict["langgraph_api_version"] = __version__
+            event_dict["langgraph_api_version"] = __version__
+        except ImportError:
+            pass
         return event_dict