langgraph-runtime-inmem 0.6.12__py3-none-any.whl → 0.8.0__py3-none-any.whl

langgraph_runtime_inmem/__init__.py

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

langgraph_runtime_inmem/ops.py

@@ -34,10 +34,12 @@ if typing.TYPE_CHECKING:
     from langgraph_api.config import ThreadTTLConfig
     from langgraph_api.schema import (
         Assistant,
+        AssistantSelectField,
         Checkpoint,
         Config,
         Context,
         Cron,
+        CronSelectField,
         DeprecatedInterrupt,
         IfNotExists,
         MetadataInput,
@@ -46,14 +48,17 @@ if typing.TYPE_CHECKING:
         OnConflictBehavior,
         QueueStats,
         Run,
+        RunSelectField,
         RunStatus,
         StreamMode,
         Thread,
+        ThreadSelectField,
         ThreadStatus,
         ThreadUpdateResponse,
     )
     from langgraph_api.schema import Interrupt as InterruptSchema
     from langgraph_api.serde import Fragment
+    from langgraph_api.utils import AsyncConnectionProto
 
 
 logger = structlog.stdlib.get_logger(__name__)
@@ -136,6 +141,7 @@ class Assistants(Authenticated):
         offset: int,
         sort_by: str | None = None,
         sort_order: str | None = None,
+        select: list[AssistantSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> tuple[AsyncIterator[Assistant], int]:
         metadata = metadata if metadata is not None else {}
@@ -178,6 +184,7 @@ class Assistants(Authenticated):
             else:
                 filtered_assistants.sort(key=lambda x: x.get(sort_by), reverse=reverse)
         else:
+            sort_by = "created_at"
             # Default sorting by created_at in descending order
             filtered_assistants.sort(key=lambda x: x["created_at"], reverse=True)
 
@@ -187,14 +194,21 @@ class Assistants(Authenticated):
 
         async def assistant_iterator() -> AsyncIterator[Assistant]:
             for assistant in paginated_assistants:
-                yield assistant
+                if select:
+                    # Filter to only selected fields
+                    filtered_assistant = {
+                        k: v for k, v in assistant.items() if k in select
+                    }
+                    yield filtered_assistant
+                else:
+                    yield assistant
 
         return assistant_iterator(), cur
 
     @staticmethod
     async def get(
         conn: InMemConnectionProto,
-        assistant_id: UUID,
+        assistant_id: UUID | str,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Assistant]:
         """Get an assistant by ID."""
@@ -217,7 +231,7 @@ class Assistants(Authenticated):
     @staticmethod
     async def put(
         conn: InMemConnectionProto,
-        assistant_id: UUID,
+        assistant_id: UUID | str,
         *,
         graph_id: str,
         config: Config,
@@ -597,6 +611,37 @@ class Assistants(Authenticated):
 
         return _yield_versions()
 
+    @staticmethod
+    async def count(
+        conn: InMemConnectionProto,
+        *,
+        graph_id: str | None = None,
+        metadata: MetadataInput = None,
+        ctx: Auth.types.BaseAuthContext | None = None,
+    ) -> int:
+        """Get count of assistants."""
+        metadata = metadata if metadata is not None else {}
+        filters = await Assistants.handle_event(
+            ctx,
+            "search",
+            Auth.types.AssistantsSearch(
+                graph_id=graph_id, metadata=metadata, limit=0, offset=0
+            ),
+        )
+
+        count = 0
+        for assistant in conn.store["assistants"]:
+            if (
+                (not graph_id or assistant["graph_id"] == graph_id)
+                and (
+                    not metadata or is_jsonb_contained(assistant["metadata"], metadata)
+                )
+                and (not filters or _check_filter_match(assistant["metadata"], filters))
+            ):
+                count += 1
+
+        return count
+
 
 def is_jsonb_contained(superset: dict[str, Any], subset: dict[str, Any]) -> bool:
     """
@@ -680,7 +725,7 @@ def _patch_interrupt(
             "value": interrupt.value,
             "resumable": interrupt.resumable,
             "ns": interrupt.ns,
-            "when": interrupt.when,
+            "when": interrupt.when,  # type: ignore[unresolved-attribute]
         }
 
 
@@ -698,6 +743,7 @@ class Threads(Authenticated):
         offset: int,
         sort_by: str | None = None,
         sort_order: str | None = None,
+        select: list[ThreadSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> tuple[AsyncIterator[Thread], int]:
         threads = conn.store["threads"]
@@ -747,6 +793,7 @@ class Threads(Authenticated):
                 filtered_threads, key=lambda x: x.get(sort_by), reverse=reverse
             )
         else:
+            sort_by = "created_at"
             # Default sorting by created_at in descending order
             sorted_threads = sorted(
                 filtered_threads, key=lambda x: x["updated_at"], reverse=True
@@ -758,7 +805,12 @@ class Threads(Authenticated):
 
         async def thread_iterator() -> AsyncIterator[Thread]:
             for thread in paginated_threads:
-                yield thread
+                if select:
+                    # Filter to only selected fields
+                    filtered_thread = {k: v for k, v in thread.items() if k in select}
+                    yield filtered_thread
+                else:
+                    yield thread
 
         return thread_iterator(), cursor
 
@@ -822,7 +874,7 @@ class Threads(Authenticated):
     @staticmethod
     async def put(
         conn: InMemConnectionProto,
-        thread_id: UUID,
+        thread_id: UUID | str,
         *,
         metadata: MetadataInput,
         if_exists: OnConflictBehavior,
@@ -1001,6 +1053,7 @@ class Threads(Authenticated):
         thread_id: UUID,
         run_id: UUID,
         run_status: RunStatus | Literal["rollback"],
+        graph_id: str,
         checkpoint: CheckpointPayload | None = None,
         exception: BaseException | None = None,
     ) -> None:
@@ -1080,6 +1133,7 @@ class Threads(Authenticated):
                 final_thread_status = "busy"
             else:
                 final_thread_status = base_thread_status
+        thread["metadata"]["graph_id"] = graph_id
         thread.update(
             {
                 "updated_at": now,
@@ -1273,7 +1327,15 @@ class Threads(Authenticated):
             metadata = thread.get("metadata", {})
             thread_config = thread.get("config", {})
 
-            if graph_id := metadata.get("graph_id"):
+            # Fallback to graph_id from run if not in thread metadata
+            graph_id = metadata.get("graph_id")
+            if not graph_id:
+                for run in conn.store["runs"]:
+                    if run["thread_id"] == thread_id:
+                        graph_id = run["kwargs"]["config"]["configurable"]["graph_id"]
+                        break
+
+            if graph_id:
                 # format latest checkpoint for response
                 checkpointer.latest_iter = checkpoint
                 async with get_graph(
@@ -1351,7 +1413,15 @@ class Threads(Authenticated):
                     detail=f"Thread {thread_id} has in-flight runs: {pending_runs}",
                 )
 
-            if graph_id := metadata.get("graph_id"):
+            # Fallback to graph_id from run if not in thread metadata
+            graph_id = metadata.get("graph_id")
+            if not graph_id:
+                for run in conn.store["runs"]:
+                    if run["thread_id"] == thread_id:
+                        graph_id = run["kwargs"]["config"]["configurable"]["graph_id"]
+                        break
+
+            if graph_id:
                 config["configurable"].setdefault("graph_id", graph_id)
 
                 checkpointer.latest_iter = checkpoint
@@ -1389,7 +1459,11 @@ class Threads(Authenticated):
                         checkpoint_id=next_config["configurable"]["checkpoint_id"],
                     )
             else:
-                raise HTTPException(status_code=400, detail="Thread has no graph ID.")
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Thread '{thread['thread_id']}' has no assigned graph ID. This usually occurs when no runs have been made on this particular thread."
+                    " This operation requires a graph ID. Please ensure a run has been made for the thread or manually update the thread metadata (by setting the 'graph_id' field) before running this operation.",
+                )
 
         @staticmethod
         async def bulk(
@@ -1471,7 +1545,11 @@ class Threads(Authenticated):
                         checkpoint=next_config["configurable"],
                     )
             else:
-                raise HTTPException(status_code=400, detail="Thread has no graph ID")
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Thread '{thread['thread_id']}' has no assigned graph ID. This usually occurs when no runs have been made on this particular thread."
+                    " This operation requires a graph ID. Please ensure a run has been made for the thread or manually update the thread metadata (by setting the 'graph_id' field) before running this operation.",
+                )
 
         @staticmethod
         async def list(
@@ -1533,6 +1611,53 @@ class Threads(Authenticated):
 
             return []
 
+    @staticmethod
+    async def count(
+        conn: InMemConnectionProto,
+        *,
+        metadata: MetadataInput = None,
+        values: MetadataInput = None,
+        status: ThreadStatus | None = None,
+        ctx: Auth.types.BaseAuthContext | None = None,
+    ) -> int:
+        """Get count of threads."""
+        threads = conn.store["threads"]
+        metadata = metadata if metadata is not None else {}
+        values = values if values is not None else {}
+        filters = await Threads.handle_event(
+            ctx,
+            "search",
+            Auth.types.ThreadsSearch(
+                metadata=metadata,
+                values=values,
+                status=status,
+                limit=0,
+                offset=0,
+            ),
+        )
+
+        count = 0
+        for thread in threads:
+            if filters and not _check_filter_match(thread["metadata"], filters):
+                continue
+
+            if metadata and not is_jsonb_contained(thread["metadata"], metadata):
+                continue
+
+            if (
+                values
+                and "values" in thread
+                and not is_jsonb_contained(thread["values"], values)
+            ):
+                continue
+
+            if status and thread.get("status") != status:
+                continue
+
+            count += 1
+
+        return count
+
 
 RUN_LOCK = asyncio.Lock()
 
@@ -1678,9 +1803,9 @@ class Runs(Authenticated):
             await stream_manager.remove_control_queue(run_id, queue)
 
     @staticmethod
-    async def sweep(conn: InMemConnectionProto) -> list[UUID]:
+    async def sweep() -> None:
         """Sweep runs that are no longer running"""
-        return []
+        pass
 
     @staticmethod
     def _merge_jsonb(*objects: dict) -> dict:
@@ -1698,7 +1823,7 @@ class Runs(Authenticated):
 
     @staticmethod
     async def put(
-        conn: InMemConnectionProto,
+        conn: InMemConnectionProto | AsyncConnectionProto,
         assistant_id: UUID,
         kwargs: dict,
         *,
@@ -1714,6 +1839,7 @@ class Runs(Authenticated):
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Run]:
         """Create a run."""
+        from langgraph_api.config import FF_RICH_THREADS
         from langgraph_api.schema import Run, Thread
 
         assistant_id = _ensure_uuid(assistant_id)
@@ -1759,34 +1885,50 @@ class Runs(Authenticated):
             # Create new thread
             if thread_id is None:
                 thread_id = uuid4()
-            thread = Thread(
-                thread_id=thread_id,
-                status="busy",
-                metadata={
-                    "graph_id": assistant["graph_id"],
-                    "assistant_id": str(assistant_id),
-                    **(config.get("metadata") or {}),
-                    **metadata,
-                },
-                config=Runs._merge_jsonb(
-                    assistant["config"],
-                    config,
-                    {
-                        "configurable": Runs._merge_jsonb(
-                            Runs._get_configurable(assistant["config"]),
-                            Runs._get_configurable(config),
-                        )
+            if FF_RICH_THREADS:
+                thread = Thread(
+                    thread_id=thread_id,
+                    status="busy",
+                    metadata={
+                        "graph_id": assistant["graph_id"],
+                        "assistant_id": str(assistant_id),
+                        **(config.get("metadata") or {}),
+                        **metadata,
                     },
-                ),
-                created_at=datetime.now(UTC),
-                updated_at=datetime.now(UTC),
-                values=b"",
-            )
+                    config=Runs._merge_jsonb(
+                        assistant["config"],
+                        config,
+                        {
+                            "configurable": Runs._merge_jsonb(
+                                Runs._get_configurable(assistant["config"]),
+                                Runs._get_configurable(config),
+                            )
+                        },
+                    ),
+                    created_at=datetime.now(UTC),
+                    updated_at=datetime.now(UTC),
+                    values=b"",
+                )
+            else:
+                thread = Thread(
+                    thread_id=thread_id,
+                    status="idle",
+                    metadata={
+                        "graph_id": assistant["graph_id"],
+                        "assistant_id": str(assistant_id),
+                        **(config.get("metadata") or {}),
+                        **metadata,
+                    },
+                    config={},
+                    created_at=datetime.now(UTC),
+                    updated_at=datetime.now(UTC),
+                    values=b"",
+                )
             await logger.ainfo("Creating thread", thread_id=thread_id)
             conn.store["threads"].append(thread)
         elif existing_thread:
             # Update existing thread
-            if existing_thread["status"] != "busy":
+            if FF_RICH_THREADS and existing_thread["status"] != "busy":
                 existing_thread["status"] = "busy"
                 existing_thread["metadata"] = Runs._merge_jsonb(
                     existing_thread["metadata"],
@@ -2025,8 +2167,8 @@ class Runs(Authenticated):
 
     @staticmethod
     async def cancel(
-        conn: InMemConnectionProto,
-        run_ids: Sequence[UUID] | None = None,
+        conn: InMemConnectionProto | AsyncConnectionProto,
+        run_ids: Sequence[UUID | str] | None = None,
         *,
         action: Literal["interrupt", "rollback"] = "interrupt",
         thread_id: UUID | None = None,
@@ -2190,6 +2332,7 @@ class Runs(Authenticated):
         limit: int = 10,
         offset: int = 0,
         status: RunStatus | None = None,
+        select: list[RunSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Run]:
         """List all runs by thread."""
@@ -2217,7 +2360,12 @@ class Runs(Authenticated):
 
         async def _return():
             for run in sliced_runs:
-                yield run
+                if select:
+                    # Filter to only selected fields
+                    filtered_run = {k: v for k, v in run.items() if k in select}
+                    yield filtered_run
+                else:
+                    yield run
 
         return _return()
 
@@ -2389,7 +2537,7 @@ class Runs(Authenticated):
 
         @staticmethod
         async def publish(
-            run_id: UUID,
+            run_id: UUID | str,
             event: str,
             message: bytes,
             *,
@@ -2490,10 +2638,24 @@ class Crons:
         thread_id: UUID | None,
         limit: int,
         offset: int,
+        select: list[CronSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
-    ) -> AsyncIterator[Cron]:
+        sort_by: str | None = None,
+        sort_order: Literal["asc", "desc"] | None = None,
+    ) -> tuple[AsyncIterator[Cron], int]:
         raise NotImplementedError
 
+    @staticmethod
+    async def count(
+        conn: InMemConnectionProto,
+        *,
+        assistant_id: UUID | None = None,
+        thread_id: UUID | None = None,
+        ctx: Auth.types.BaseAuthContext | None = None,
+    ) -> int:
+        """Get count of crons."""
+        raise NotImplementedError("The in-mem server does not implement Crons.")
+
 
 async def cancel_run(
     thread_id: UUID, run_id: UUID, ctx: Auth.types.BaseAuthContext | None = None

langgraph_runtime_inmem/queue.py

@@ -128,8 +128,7 @@ async def queue():
                         # sweep runs if needed
                         if do_sweep:
                             last_sweep_secs = loop.time()
-                            run_ids = await ops.Runs.sweep(conn)
-                            logger.info("Swept runs", run_ids=run_ids)
+                            await ops.Runs.sweep()
             except Exception as exc:
                 # keep trying to run the scheduler indefinitely
                 logger.exception("Background worker scheduler failed", exc_info=exc)

{langgraph_runtime_inmem-0.6.12.dist-info → langgraph_runtime_inmem-0.8.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langgraph-runtime-inmem
-Version: 0.6.12
+Version: 0.8.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.6.12.dist-info → langgraph_runtime_inmem-0.8.0.dist-info}/RECORD

@@ -1,13 +1,13 @@
-langgraph_runtime_inmem/__init__.py,sha256=Y4pMrNRANpl0WNkm4rJgbGxM0sdgeJNJ17xTUA5Do0w,311
+langgraph_runtime_inmem/__init__.py,sha256=ZCa8w1LAQPb9BIGUaaEwW1hak0vidcMAUArcD_C1lyU,310
 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=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=lhvPaQRUmEKumeWJUMOcSmYpZkLbOWvAJT-wYzu5NrM,91701
-langgraph_runtime_inmem/queue.py,sha256=nqfgz7j_Jkh5Ek5-RsHB2Uvwbxguu9IUPkGXIxvFPns,10037
+langgraph_runtime_inmem/ops.py,sha256=pV2IHiis4ydYFtroOfX-FcQjhXvk5O1ESuy2kk7mbyI,98064
+langgraph_runtime_inmem/queue.py,sha256=33qfFKPhQicZ1qiibllYb-bTFzUNSN2c4bffPACP5es,9952
 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.12.dist-info/METADATA,sha256=kVDdIT1O23mKzQMYtfFwpy9F4IYZzLYt60T12AczxeQ,566
-langgraph_runtime_inmem-0.6.12.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-langgraph_runtime_inmem-0.6.12.dist-info/RECORD,,
+langgraph_runtime_inmem-0.8.0.dist-info/METADATA,sha256=hg_MVwCcdeCjZrYul77k4J5y5ByIUDOtTqqUOQbyqDs,565
+langgraph_runtime_inmem-0.8.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+langgraph_runtime_inmem-0.8.0.dist-info/RECORD,,