langgraph-runtime-inmem 0.6.7__py3-none-any.whl → 0.6.10__py3-none-any.whl

langgraph_runtime_inmem/__init__.py

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

langgraph_runtime_inmem/inmem_stream.py

@@ -42,6 +42,7 @@ class ContextQueue(asyncio.Queue):
 class StreamManager:
     def __init__(self):
         self.queues = defaultdict(list)  # Dict[UUID, List[asyncio.Queue]]
+        self.control_keys = defaultdict()
         self.control_queues = defaultdict(list)
 
         self.message_stores = defaultdict(list)  # Dict[UUID, List[Message]]
@@ -51,6 +52,14 @@ class StreamManager:
         run_id = _ensure_uuid(run_id)
         return self.queues[run_id]
 
+    def get_control_queues(self, run_id: UUID | str) -> list[asyncio.Queue]:
+        run_id = _ensure_uuid(run_id)
+        return self.control_queues[run_id]
+
+    def get_control_key(self, run_id: UUID | str) -> Message | None:
+        run_id = _ensure_uuid(run_id)
+        return self.control_keys.get(run_id)
+
     async def put(
         self, run_id: UUID | str, message: Message, resumable: bool = False
     ) -> None:
@@ -61,8 +70,10 @@ class StreamManager:
             self.message_stores[run_id].append(message)
         topic = message.topic.decode()
         if "control" in topic:
-            self.control_queues[run_id].append(message)
-        queues = self.queues.get(run_id, [])
+            self.control_keys[run_id] = message
+            queues = self.control_queues[run_id]
+        else:
+            queues = self.queues[run_id]
         coros = [queue.put(message) for queue in queues]
         results = await asyncio.gather(*coros, return_exceptions=True)
         for result in results:
@@ -73,14 +84,12 @@ class StreamManager:
         run_id = _ensure_uuid(run_id)
         queue = ContextQueue()
         self.queues[run_id].append(queue)
-        for control_msg in self.control_queues[run_id]:
-            try:
-                await queue.put(control_msg)
-            except Exception:
-                logger.exception(
-                    f"Failed to put control message in queue: {control_msg}"
-                )
+        return queue
 
+    async def add_control_queue(self, run_id: UUID | str) -> asyncio.Queue:
+        run_id = _ensure_uuid(run_id)
+        queue = ContextQueue()
+        self.control_queues[run_id].append(queue)
         return queue
 
     async def remove_queue(self, run_id: UUID | str, queue: asyncio.Queue):
@@ -89,8 +98,13 @@ class StreamManager:
             self.queues[run_id].remove(queue)
             if not self.queues[run_id]:
                 del self.queues[run_id]
-                if run_id in self.message_stores:
-                    del self.message_stores[run_id]
+
+    async def remove_control_queue(self, run_id: UUID | str, queue: asyncio.Queue):
+        run_id = _ensure_uuid(run_id)
+        if run_id in self.control_queues:
+            self.control_queues[run_id].remove(queue)
+            if not self.control_queues[run_id]:
+                del self.control_queues[run_id]
 
     def restore_messages(
         self, run_id: UUID | str, message_id: str | None

langgraph_runtime_inmem/ops.py

@@ -157,9 +157,6 @@ class Assistants(Authenticated):
             and (not filters or _check_filter_match(assistant["metadata"], filters))
         ]
 
-        # Get total count before sorting and pagination
-        total_count = len(filtered_assistants)
-
         # Sort based on sort_by and sort_order
         sort_by = sort_by.lower() if sort_by else None
         if sort_by and sort_by in (
@@ -186,12 +183,13 @@ class Assistants(Authenticated):
 
         # Apply pagination
         paginated_assistants = filtered_assistants[offset : offset + limit]
+        cur = offset + limit if len(filtered_assistants) > offset + limit else None
 
         async def assistant_iterator() -> AsyncIterator[Assistant]:
             for assistant in paginated_assistants:
                 yield assistant
 
-        return assistant_iterator(), total_count
+        return assistant_iterator(), cur
 
     @staticmethod
     async def get(
@@ -724,7 +722,7 @@ class Threads(Authenticated):
         else:
             # Default sorting by created_at in descending order
             sorted_threads = sorted(
-                filtered_threads, key=lambda x: x["created_at"], reverse=True
+                filtered_threads, key=lambda x: x["updated_at"], reverse=True
             )
 
         # Apply limit and offset
@@ -1453,7 +1451,7 @@ class Threads(Authenticated):
             conn: InMemConnectionProto,
             *,
             config: Config,
-            limit: int = 10,
+            limit: int = 1,
             before: str | Checkpoint | None = None,
             metadata: MetadataInput = None,
             ctx: Auth.types.BaseAuthContext | None = None,
@@ -1628,7 +1626,7 @@ class Runs(Authenticated):
 
         stream_manager = get_stream_manager()
         # Get queue for this run
-        queue = await Runs.Stream.subscribe(run_id)
+        queue = await stream_manager.add_control_queue(run_id)
 
         async with SimpleTaskGroup(cancel=True, taskgroup_name="Runs.enter") as tg:
             done = ValueEvent()
@@ -1636,16 +1634,21 @@ class Runs(Authenticated):
 
             # Give done event to caller
             yield done
-            # Signal done to all subscribers
+            # Store the control message for late subscribers
             control_message = Message(
                 topic=f"run:{run_id}:control".encode(), data=b"done"
             )
-
-            # Store the control message for late subscribers
             await stream_manager.put(run_id, control_message)
-            stream_manager.control_queues[run_id].append(control_message)
-            # Clean up this queue
-            await stream_manager.remove_queue(run_id, queue)
+
+            # Signal done to all subscribers
+            stream_message = Message(
+                topic=f"run:{run_id}:stream".encode(),
+                data={"event": "control", "message": b"done"},
+            )
+            await stream_manager.put(run_id, stream_message)
+
+            # Remove the queue
+            await stream_manager.remove_control_queue(run_id, queue)
 
     @staticmethod
     async def sweep(conn: InMemConnectionProto) -> list[UUID]:
@@ -1981,6 +1984,16 @@ class Runs(Authenticated):
                     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
@@ -2201,8 +2214,6 @@ class Runs(Authenticated):
         @staticmethod
         async def subscribe(
             run_id: UUID,
-            *,
-            stream_mode: StreamMode | None = None,
         ) -> asyncio.Queue:
             """Subscribe to the run stream, returning a queue."""
             stream_manager = get_stream_manager()
@@ -2222,20 +2233,18 @@ class Runs(Authenticated):
             ignore_404: bool = False,
             cancel_on_disconnect: bool = False,
             stream_channel: asyncio.Queue | None = None,
-            stream_mode: list[StreamMode] | StreamMode,
+            stream_mode: list[StreamMode] | StreamMode | None = None,
             last_event_id: str | None = None,
             ctx: Auth.types.BaseAuthContext | None = None,
         ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
             """Stream the run output."""
             from langgraph_api.asyncio import create_task
-
-            if stream_mode and not isinstance(stream_mode, list):
-                stream_mode = [stream_mode]
+            from langgraph_api.serde import json_loads
 
             queue = (
                 stream_channel
                 if stream_channel
-                else await Runs.Stream.subscribe(run_id, stream_mode=stream_mode)
+                else await Runs.Stream.subscribe(run_id)
             )
 
             try:
@@ -2256,53 +2265,71 @@ class Runs(Authenticated):
                                 )
                             )
                     run = await Runs.get(conn, run_id, thread_id=thread_id, ctx=ctx)
-                    channel_prefix = f"run:{run_id}:stream:"
-                    len_prefix = len(channel_prefix.encode())
 
                     for message in get_stream_manager().restore_messages(
                         run_id, last_event_id
                     ):
-                        topic, data, id = message.topic, message.data, message.id
-                        if topic.decode() == f"run:{run_id}:control":
-                            if data == b"done":
+                        data, id = message.data, message.id
+
+                        data = json_loads(data)
+                        mode = data["event"]
+                        message = data["message"]
+
+                        if mode == "control":
+                            if message == b"done":
                                 return
-                        else:
-                            mode = topic[len_prefix:]
-                            if mode == b"updates" and "updates" not in stream_mode:
-                                continue
-                            else:
-                                yield mode, data, id
-                                logger.debug(
-                                    "Replayed run event",
-                                    run_id=str(run_id),
-                                    message_id=id,
-                                    stream_mode=mode,
-                                    data=data,
+                        elif (
+                            not stream_mode
+                            or mode in stream_mode
+                            or (
+                                (
+                                    "messages" in stream_mode
+                                    or "messages-tuple" in stream_mode
                                 )
+                                and mode.startswith("messages")
+                            )
+                        ):
+                            yield mode.encode(), base64.b64decode(message), id
+                            logger.debug(
+                                "Replayed run event",
+                                run_id=str(run_id),
+                                message_id=id,
+                                stream_mode=mode,
+                                data=data,
+                            )
 
                     while True:
                         try:
                             # Wait for messages with a timeout
                             message = await asyncio.wait_for(queue.get(), timeout=0.5)
-                            topic, data, id = message.topic, message.data, message.id
+                            data, id = message.data, message.id
 
-                            if topic.decode() == f"run:{run_id}:control":
-                                if data == b"done":
+                            data = json_loads(data)
+                            mode = data["event"]
+                            message = data["message"]
+
+                            if mode == "control":
+                                if message == b"done":
                                     break
-                            else:
-                                # Extract mode from topic
-                                mode = topic[len_prefix:]
-                                if mode == b"updates" and "updates" not in stream_mode:
-                                    continue
-                                else:
-                                    yield mode, data, id
-                                    logger.debug(
-                                        "Streamed run event",
-                                        run_id=str(run_id),
-                                        stream_mode=mode,
-                                        message_id=id,
-                                        data=data,
+                            elif (
+                                not stream_mode
+                                or mode in stream_mode
+                                or (
+                                    (
+                                        "messages" in stream_mode
+                                        or "messages-tuple" in stream_mode
                                     )
+                                    and mode.startswith("messages")
+                                )
+                            ):
+                                yield mode.encode(), base64.b64decode(message), id
+                                logger.debug(
+                                    "Streamed run event",
+                                    run_id=str(run_id),
+                                    stream_mode=mode,
+                                    message_id=id,
+                                    data=message,
+                                )
                         except TimeoutError:
                             # Check if the run is still pending
                             run_iter = await Runs.get(
@@ -2342,12 +2369,20 @@ class Runs(Authenticated):
             resumable: bool = False,
         ) -> None:
             """Publish a message to all subscribers of the run stream."""
-            topic = f"run:{run_id}:stream:{event}".encode()
+            from langgraph_api.serde import json_dumpb
+
+            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,
+                }
+            )
             await stream_manager.put(
-                run_id, Message(topic=topic, data=message), resumable
+                run_id, Message(topic=topic, data=payload), resumable
             )
 
 
@@ -2356,20 +2391,18 @@ async def listen_for_cancellation(queue: asyncio.Queue, run_id: UUID, done: Valu
     from langgraph_api.errors import UserInterrupt, UserRollback
 
     stream_manager = get_stream_manager()
-    control_key = f"run:{run_id}:control"
 
-    if existing_queue := stream_manager.control_queues.get(run_id):
-        for message in existing_queue:
-            payload = message.data
-            if payload == b"rollback":
-                done.set(UserRollback())
-            elif payload == b"interrupt":
-                done.set(UserInterrupt())
+    if control_key := stream_manager.get_control_key(run_id):
+        payload = control_key.data
+        if payload == b"rollback":
+            done.set(UserRollback())
+        elif payload == b"interrupt":
+            done.set(UserInterrupt())
 
     while not done.is_set():
         try:
             # This task gets cancelled when Runs.enter exits anyway,
-            # so we can have a pretty length timeout here
+            # so we can have a pretty lengthy timeout here
             message = await asyncio.wait_for(queue.get(), timeout=240)
             payload = message.data
             if payload == b"rollback":
@@ -2379,10 +2412,6 @@ async def listen_for_cancellation(queue: asyncio.Queue, run_id: UUID, done: Valu
             elif payload == b"done":
                 done.set()
                 break
-
-            # Store control messages for late subscribers
-            if message.topic.decode() == control_key:
-                stream_manager.control_queues[run_id].append(message)
         except TimeoutError:
             break
 

{langgraph_runtime_inmem-0.6.7.dist-info → langgraph_runtime_inmem-0.6.10.dist-info}/METADATA

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

{langgraph_runtime_inmem-0.6.7.dist-info → langgraph_runtime_inmem-0.6.10.dist-info}/RECORD

@@ -1,13 +1,13 @@
-langgraph_runtime_inmem/__init__.py,sha256=IO8abGzSnjJCCIfUCYcUfMMQUyqIIQ6oHu2-2yJJCg0,310
+langgraph_runtime_inmem/__init__.py,sha256=3R6DOc_wOQRDAv1aNNMMGjRofpJi6B_qIBkZ373AG44,311
 langgraph_runtime_inmem/checkpoint.py,sha256=nc1G8DqVdIu-ibjKTqXfbPfMbAsKjPObKqegrSzo6Po,4432
 langgraph_runtime_inmem/database.py,sha256=G_6L2khpRDSpS2Vs_SujzHayODcwG5V2IhFP7LLBXgw,6349
-langgraph_runtime_inmem/inmem_stream.py,sha256=65z_2mBNJ0-yJsXWnlYwRc71039_y6Sa0MN8fL_U3Ko,4581
+langgraph_runtime_inmem/inmem_stream.py,sha256=UWk1srLF44HZPPbRdArGGhsy0MY0UOJKSIxBSO7Hosc,5138
 langgraph_runtime_inmem/lifespan.py,sha256=t0w2MX2dGxe8yNtSX97Z-d2pFpllSLS4s1rh2GJDw5M,3557
 langgraph_runtime_inmem/metrics.py,sha256=HhO0RC2bMDTDyGBNvnd2ooLebLA8P1u5oq978Kp_nAA,392
-langgraph_runtime_inmem/ops.py,sha256=kx0f6pKqYTypbcoMByM6L0YBF4O0b1TEYpOqdBnJMao,89354
+langgraph_runtime_inmem/ops.py,sha256=CSH5vi7AsaeaWSngZ_DCtsPY-M7ah3cz8acTVi_UbUw,90559
 langgraph_runtime_inmem/queue.py,sha256=nqfgz7j_Jkh5Ek5-RsHB2Uvwbxguu9IUPkGXIxvFPns,10037
 langgraph_runtime_inmem/retry.py,sha256=XmldOP4e_H5s264CagJRVnQMDFcEJR_dldVR1Hm5XvM,763
 langgraph_runtime_inmem/store.py,sha256=rTfL1JJvd-j4xjTrL8qDcynaWF6gUJ9-GDVwH0NBD_I,3506
-langgraph_runtime_inmem-0.6.7.dist-info/METADATA,sha256=83cPJjL678QXBM5qzhYS5CRQZUObUqyGjQRLMI-1TBA,565
-langgraph_runtime_inmem-0.6.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langgraph_runtime_inmem-0.6.7.dist-info/RECORD,,
+langgraph_runtime_inmem-0.6.10.dist-info/METADATA,sha256=5YlbIvGy2fpU9gydQgbaSfu34loSdodD2PUtPlcn_1c,566
+langgraph_runtime_inmem-0.6.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langgraph_runtime_inmem-0.6.10.dist-info/RECORD,,