langgraph-runtime-inmem 0.12.1__tar.gz → 0.14.0__tar.gz

{langgraph_runtime_inmem-0.12.1 → langgraph_runtime_inmem-0.14.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-runtime-inmem
-Version: 0.12.1
+Version: 0.14.0
 Summary: Inmem implementation for the LangGraph API server.
 Author-email: Will Fu-Hinthorn <will@langchain.dev>
 License: Elastic-2.0

{langgraph_runtime_inmem-0.12.1 → langgraph_runtime_inmem-0.14.0}/langgraph_runtime_inmem/__init__.py

@@ -9,7 +9,7 @@ from langgraph_runtime_inmem import (
     store,
 )
 
-__version__ = "0.12.1"
+__version__ = "0.14.0"
 __all__ = [
     "ops",
     "database",

{langgraph_runtime_inmem-0.12.1 → langgraph_runtime_inmem-0.14.0}/langgraph_runtime_inmem/inmem_stream.py

@@ -108,9 +108,10 @@ class StreamManager:
             thread_id = _ensure_uuid(thread_id)
 
         message.id = _generate_ms_seq_id().encode()
+        # For resumable run streams, embed the generated message ID into the frame
+        topic = message.topic.decode()
         if resumable:
             self.message_stores[thread_id][run_id].append(message)
-        topic = message.topic.decode()
         if "control" in topic:
             self.control_keys[thread_id][run_id] = message
             queues = self.control_queues[thread_id][run_id]

{langgraph_runtime_inmem-0.12.1 → langgraph_runtime_inmem-0.14.0}/langgraph_runtime_inmem/ops.py

@@ -62,9 +62,9 @@ if typing.TYPE_CHECKING:
         ThreadUpdateResponse,
     )
     from langgraph_api.schema import Interrupt as InterruptSchema
-    from langgraph_api.serde import Fragment
     from langgraph_api.utils import AsyncConnectionProto
 
+StreamHandler = ContextQueue
 
 logger = structlog.stdlib.get_logger(__name__)
 
@@ -1720,6 +1720,9 @@ class Threads(Authenticated):
             stream_modes: list[ThreadStreamMode],
         ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
             """Stream the thread output."""
+            from langgraph_api.utils.stream_codec import (
+                decode_stream_message,
+            )
 
             def should_filter_event(event_name: str, message_bytes: bytes) -> bool:
                 """Check if an event should be filtered out based on stream_modes."""
@@ -1738,8 +1741,6 @@ class Threads(Authenticated):
                         pass
                 return True
 
-            from langgraph_api.serde import json_loads
-
             stream_manager = get_stream_manager()
             seen_runs: set[UUID] = set()
             created_queues: list[tuple[UUID, asyncio.Queue]] = []
@@ -1768,35 +1769,26 @@ class Threads(Authenticated):
 
                         # Yield sorted events
                         for message, run_id in all_events:
-                            data = json_loads(message.data)
-                            event_name = data["event"]
-                            message_content = data["message"]
+                            decoded = decode_stream_message(
+                                message.data, channel=message.topic
+                            )
+                            event_bytes = decoded.event_bytes
+                            message_bytes = decoded.message_bytes
 
-                            if event_name == "control":
-                                if message_content == b"done":
+                            if event_bytes == b"control":
+                                if message_bytes == b"done":
                                     event_bytes = b"metadata"
                                     message_bytes = orjson.dumps(
                                         {"status": "run_done", "run_id": run_id}
                                     )
-                                    # Filter events based on stream_modes
-                                    if not should_filter_event(
-                                        "metadata", message_bytes
-                                    ):
-                                        yield (
-                                            event_bytes,
-                                            message_bytes,
-                                            message.id,
-                                        )
-                            else:
-                                event_bytes = event_name.encode()
-                                message_bytes = base64.b64decode(message_content)
-                                # Filter events based on stream_modes
-                                if not should_filter_event(event_name, message_bytes):
-                                    yield (
-                                        event_bytes,
-                                        message_bytes,
-                                        message.id,
-                                    )
+                            if not should_filter_event(
+                                event_bytes.decode("utf-8"), message_bytes
+                            ):
+                                yield (
+                                    event_bytes,
+                                    message_bytes,
+                                    message.id,
+                                )
 
                     # Listen for live messages from all queues
                     while True:
@@ -1813,40 +1805,27 @@ class Threads(Authenticated):
                                 message = await asyncio.wait_for(
                                     queue.get(), timeout=0.2
                                 )
-                                data = json_loads(message.data)
-                                event_name = data["event"]
-                                message_content = data["message"]
-
-                                if event_name == "control":
-                                    if message_content == b"done":
-                                        # Extract run_id from topic
-                                        topic = message.topic.decode()
-                                        run_id = topic.split("run:")[1].split(":")[0]
-                                        event_bytes = b"metadata"
-                                        message_bytes = orjson.dumps(
-                                            {"status": "run_done", "run_id": run_id}
-                                        )
-                                        # Filter events based on stream_modes
-                                        if not should_filter_event(
-                                            "metadata", message_bytes
-                                        ):
-                                            yield (
-                                                event_bytes,
-                                                message_bytes,
-                                                message.id,
-                                            )
-                                else:
-                                    event_bytes = event_name.encode()
-                                    message_bytes = base64.b64decode(message_content)
-                                    # Filter events based on stream_modes
+                                decoded = decode_stream_message(
+                                    message.data, channel=message.topic
+                                )
+                                event = decoded.event_bytes
+                                event_name = event.decode("utf-8")
+                                payload = decoded.message_bytes
+
+                                if event == b"control" and payload == b"done":
+                                    topic = message.topic.decode()
+                                    run_id = topic.split("run:")[1].split(":")[0]
+                                    meta_event = b"metadata"
+                                    meta_payload = orjson.dumps(
+                                        {"status": "run_done", "run_id": run_id}
+                                    )
                                     if not should_filter_event(
-                                        event_name, message_bytes
+                                        "metadata", meta_payload
                                     ):
-                                        yield (
-                                            event_bytes,
-                                            message_bytes,
-                                            message.id,
-                                        )
+                                        yield (meta_event, meta_payload, message.id)
+                                else:
+                                    if not should_filter_event(event_name, payload):
+                                        yield (event, payload, message.id)
 
                             except TimeoutError:
                                 continue
@@ -1882,18 +1861,12 @@ class Threads(Authenticated):
             message: bytes,
         ) -> None:
             """Publish a thread-level event to the thread stream."""
-            from langgraph_api.serde import json_dumpb
+            from langgraph_api.utils.stream_codec import STREAM_CODEC
 
             topic = f"thread:{thread_id}:stream".encode()
 
             stream_manager = get_stream_manager()
-            # Send to thread stream topic
-            payload = json_dumpb(
-                {
-                    "event": event,
-                    "message": message,
-                }
-            )
+            payload = STREAM_CODEC.encode(event, message)
             await stream_manager.put_thread(
                 str(thread_id), Message(topic=topic, data=payload)
             )
@@ -2065,6 +2038,7 @@ class Runs(Authenticated):
         This method should be called as a context manager by a worker executing a run.
         """
         from langgraph_api.asyncio import SimpleTaskGroup, ValueEvent
+        from langgraph_api.utils.stream_codec import STREAM_CODEC
 
         stream_manager = get_stream_manager()
         # Get control queue for this run (normal queue is created during run creation)
@@ -2084,10 +2058,10 @@ class Runs(Authenticated):
             )
             await stream_manager.put(run_id, thread_id, control_message)
 
-            # Signal done to all subscribers
+            # Signal done to all subscribers using stream codec
             stream_message = Message(
                 topic=f"run:{run_id}:stream".encode(),
-                data={"event": "control", "message": b"done"},
+                data=STREAM_CODEC.encode("control", b"done"),
             )
             await stream_manager.put(
                 run_id, thread_id, stream_message, resumable=resumable
@@ -2399,68 +2373,6 @@ class Runs(Authenticated):
 
         return _yield_deleted()
 
-    @staticmethod
-    async def join(
-        run_id: UUID,
-        *,
-        thread_id: UUID,
-        ctx: Auth.types.BaseAuthContext | None = None,
-    ) -> Fragment:
-        """Wait for a run to complete. If already done, return immediately.
-
-        Returns:
-            the final state of the run.
-        """
-        from langgraph_api.serde import Fragment
-        from langgraph_api.utils import fetchone
-
-        async with connect() as conn:
-            # Validate ownership
-            thread_iter = await Threads.get(conn, thread_id, ctx=ctx)
-            await fetchone(thread_iter)
-        last_chunk: bytes | None = None
-        # wait for the run to complete
-        # Rely on this join's auth
-        async with await Runs.Stream.subscribe(run_id, thread_id) as sub:
-            async for mode, chunk, _ in Runs.Stream.join(
-                run_id,
-                thread_id=thread_id,
-                ctx=ctx,
-                ignore_404=True,
-                stream_channel=sub,
-                stream_mode=["values", "updates", "error"],
-            ):
-                if mode == b"values":
-                    last_chunk = chunk
-                elif mode == b"updates" and b"__interrupt__" in chunk:
-                    last_chunk = chunk
-                elif mode == b"error":
-                    last_chunk = orjson.dumps({"__error__": orjson.Fragment(chunk)})
-        # if we received a final chunk, return it
-        if last_chunk is not None:
-            # ie. if the run completed while we were waiting for it
-            return Fragment(last_chunk)
-        else:
-            # otherwise, the run had already finished, so fetch the state from thread
-            async with connect() as conn:
-                thread_iter = await Threads.get(conn, thread_id, ctx=ctx)
-                thread = await fetchone(thread_iter)
-                if thread["status"] == "error":
-                    return Fragment(
-                        orjson.dumps({"__error__": orjson.Fragment(thread["error"])})
-                    )
-                if thread["status"] == "interrupted":
-                    # Get an interrupt for the thread. There is the case where there are multiple interrupts for the same run and we may not show the same
-                    # interrupt, but we'll always show one. Long term we should show all of them.
-                    try:
-                        interrupt_map = thread["interrupts"]
-                        interrupt = [next(iter(interrupt_map.values()))[0]]
-                        return Fragment(orjson.dumps({"__interrupt__": interrupt}))
-                    except Exception:
-                        # No interrupt, but status is interrupted from a before/after block. Default back to values.
-                        pass
-                return thread["values"]
-
     @staticmethod
     async def cancel(
         conn: InMemConnectionProto | AsyncConnectionProto,
@@ -2720,39 +2632,28 @@ class Runs(Authenticated):
         ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
             """Stream the run output."""
             from langgraph_api.asyncio import create_task
-            from langgraph_api.serde import json_loads
+            from langgraph_api.serde import json_dumpb
+            from langgraph_api.utils.stream_codec import decode_stream_message
 
             queue = stream_channel
             try:
                 async with connect() as conn:
-                    filters = await Runs.handle_event(
-                        ctx,
-                        "read",
-                        Auth.types.ThreadsRead(thread_id=thread_id),
-                    )
-                    if filters:
-                        thread = await Threads._get_with_filters(
-                            cast(InMemConnectionProto, conn), thread_id, filters
-                        )
-                        if not thread:
-                            raise WrappedHTTPException(
-                                HTTPException(
-                                    status_code=404, detail="Thread not found"
-                                )
-                            )
+                    try:
+                        await Runs.Stream.check_run_stream_auth(run_id, thread_id, ctx)
+                    except HTTPException as e:
+                        raise WrappedHTTPException(e) from None
                     run = await Runs.get(conn, run_id, thread_id=thread_id, ctx=ctx)
 
                     for message in get_stream_manager().restore_messages(
                         run_id, thread_id, last_event_id
                     ):
                         data, id = message.data, message.id
-
-                        data = json_loads(data)
-                        mode = data["event"]
-                        message = data["message"]
+                        decoded = decode_stream_message(data, channel=message.topic)
+                        mode = decoded.event_bytes.decode("utf-8")
+                        payload = decoded.message_bytes
 
                         if mode == "control":
-                            if message == b"done":
+                            if payload == b"done":
                                 return
                         elif (
                             not stream_mode
@@ -2765,7 +2666,7 @@ class Runs(Authenticated):
                                 and mode.startswith("messages")
                             )
                         ):
-                            yield mode.encode(), base64.b64decode(message), id
+                            yield mode.encode(), payload, id
                             logger.debug(
                                 "Replayed run event",
                                 run_id=str(run_id),
@@ -2779,13 +2680,12 @@ class Runs(Authenticated):
                             # Wait for messages with a timeout
                             message = await asyncio.wait_for(queue.get(), timeout=0.5)
                             data, id = message.data, message.id
-
-                            data = json_loads(data)
-                            mode = data["event"]
-                            message = data["message"]
+                            decoded = decode_stream_message(data, channel=message.topic)
+                            mode = decoded.event_bytes.decode("utf-8")
+                            payload = decoded.message_bytes
 
                             if mode == "control":
-                                if message == b"done":
+                                if payload == b"done":
                                     break
                             elif (
                                 not stream_mode
@@ -2798,13 +2698,13 @@ class Runs(Authenticated):
                                     and mode.startswith("messages")
                                 )
                             ):
-                                yield mode.encode(), base64.b64decode(message), id
+                                yield mode.encode(), payload, id
                                 logger.debug(
                                     "Streamed run event",
                                     run_id=str(run_id),
                                     stream_mode=mode,
                                     message_id=id,
-                                    data=message,
+                                    data=payload,
                                 )
                         except TimeoutError:
                             # Check if the run is still pending
@@ -2818,8 +2718,10 @@ class Runs(Authenticated):
                             elif run is None:
                                 yield (
                                     b"error",
-                                    HTTPException(
-                                        status_code=404, detail="Run not found"
+                                    json_dumpb(
+                                        HTTPException(
+                                            status_code=404, detail="Run not found"
+                                        )
                                     ),
                                     None,
                                 )
@@ -2836,6 +2738,25 @@ class Runs(Authenticated):
                 stream_manager = get_stream_manager()
                 await stream_manager.remove_queue(run_id, thread_id, queue)
 
+        @staticmethod
+        async def check_run_stream_auth(
+            run_id: UUID,
+            thread_id: UUID,
+            ctx: Auth.types.BaseAuthContext | None = None,
+        ) -> None:
+            async with connect() as conn:
+                filters = await Runs.handle_event(
+                    ctx,
+                    "read",
+                    Auth.types.ThreadsRead(thread_id=thread_id),
+                )
+                if filters:
+                    thread = await Threads._get_with_filters(
+                        cast(InMemConnectionProto, conn), thread_id, filters
+                    )
+                    if not thread:
+                        raise HTTPException(status_code=404, detail="Thread not found")
+
         @staticmethod
         async def publish(
             run_id: UUID | str,
@@ -2846,18 +2767,13 @@ class Runs(Authenticated):
             resumable: bool = False,
         ) -> None:
             """Publish a message to all subscribers of the run stream."""
-            from langgraph_api.serde import json_dumpb
+            from langgraph_api.utils.stream_codec import STREAM_CODEC
 
             topic = f"run:{run_id}:stream".encode()
 
             stream_manager = get_stream_manager()
-            # Send to all queues subscribed to this run_id
-            payload = json_dumpb(
-                {
-                    "event": event,
-                    "message": message,
-                }
-            )
+            # Send to all queues subscribed to this run_id using protocol frame
+            payload = STREAM_CODEC.encode(event, message)
             await stream_manager.put(
                 run_id, thread_id, Message(topic=topic, data=payload), resumable
             )
@@ -3037,6 +2953,7 @@ async def _empty_generator():
 
 
 __all__ = [
+    "StreamHandler",
     "Assistants",
     "Crons",
     "Runs",