langgraph-api 0.4.0__py3-none-any.whl → 0.4.7__py3-none-any.whl

langgraph_api/__init__.py

@@ -1 +1 @@
-__version__ = "0.4.0"
+__version__ = "0.4.7"

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 (
@@ -293,10 +294,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)