Move FileCache to sqlite package, add InMemoryCache

libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import asyncio
+import datetime
+import sqlite3
+import threading
+from collections.abc import Mapping, Sequence
+from typing import Generic
+
+from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class SqliteCache(BaseCache[ValueT], Generic[ValueT]):
+    """File-based cache using SQLite."""
+
+    def __init__(
+        self,
+        *,
+        path: str,
+        serde: SerializerProtocol | None = None,
+    ) -> None:
+        """Initialize the cache with a file path."""
+        super().__init__(serde=serde)
+        # SQLite backing store
+        self._conn = sqlite3.connect(
+            path,
+            check_same_thread=False,
+        )
+        # Serialize access to the shared connection across threads
+        self._lock = threading.RLock()
+        # Better concurrency & atomicity
+        self._conn.execute("PRAGMA journal_mode=WAL;")
+        # Schema: key -> (expiry, encoding, value)
+        self._conn.execute(
+            """CREATE TABLE IF NOT EXISTS cache (
+                ns TEXT,
+                key TEXT,
+                expiry REAL,
+                encoding TEXT NOT NULL,
+                val BLOB NOT NULL,
+                PRIMARY KEY (ns, key)
+            )"""
+        )
+        self._conn.commit()
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        with self._lock, self._conn:
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            if not keys:
+                return {}
+            placeholders = ",".join("(?, ?)" for _ in keys)
+            params: list[str] = []
+            for ns_tuple, key in keys:
+                params.extend((",".join(ns_tuple), key))
+            cursor = self._conn.execute(
+                f"SELECT ns, key, expiry, encoding, val FROM cache WHERE (ns, key) IN ({placeholders})",
+                tuple(params),
+            )
+            values: dict[FullKey, ValueT] = {}
+            rows = cursor.fetchall()
+            for ns, key, expiry, encoding, raw in rows:
+                if expiry is not None and now > expiry:
+                    # purge expired entry
+                    self._conn.execute(
+                        "DELETE FROM cache WHERE (ns, key) = (?, ?)", (ns, key)
+                    )
+                    continue
+                values[(tuple(ns.split(",")), key)] = self.serde.loads_typed(
+                    (encoding, raw)
+                )
+            return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return await asyncio.to_thread(self.get, keys)
+
+    def set(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys and TTLs."""
+        with self._lock, self._conn:
+            now = datetime.datetime.now(datetime.timezone.utc)
+            for key, (value, ttl) in mapping.items():
+                if ttl is not None:
+                    delta = datetime.timedelta(seconds=ttl)
+                    expiry: float | None = (now + delta).timestamp()
+                else:
+                    expiry = None
+                encoding, raw = self.serde.dumps_typed(value)
+                self._conn.execute(
+                    "INSERT OR REPLACE INTO cache (ns, key, expiry, encoding, val) VALUES (?, ?, ?, ?, ?)",
+                    (",".join(key[0]), key[1], expiry, encoding, raw),
+                )
+
+    async def aset(self, mapping: Mapping[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys and TTLs."""
+        await asyncio.to_thread(self.set, mapping)
+
+    def delete(self, keys: Sequence[Namespace]) -> None:
+        """Delete the cached values for the given namespaces."""
+        if not keys:
+            return
+        with self._lock, self._conn:
+            placeholders = ",".join("?" for _ in keys)
+            self._conn.execute(
+                f"DELETE FROM cache WHERE (ns) IN ({placeholders})",
+                tuple(",".join(key) for key in keys),
+            )
+
+    async def adelete(self, keys: Sequence[Namespace]) -> None:
+        """Asynchronously delete the cached values for the given namespaces."""
+        await asyncio.to_thread(self.delete, keys)
+
+    def __del__(self) -> None:
+        try:
+            self._conn.close()
+        except Exception:
+            pass

libs/checkpoint/langgraph/cache/memory/__init__.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import datetime
+import threading
+from collections.abc import Sequence
+from typing import Generic
+
+from langgraph.cache.base import BaseCache, FullKey, Namespace, ValueT
+from langgraph.checkpoint.serde.base import SerializerProtocol
+
+
+class InMemoryCache(BaseCache[ValueT], Generic[ValueT]):
+    def __init__(self, *, serde: SerializerProtocol | None = None):
+        super().__init__(serde=serde)
+        self._cache: dict[Namespace, dict[str, tuple[str, bytes, int | None]]] = {}
+        self._lock = threading.RLock()
+
+    def get(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Get the cached values for the given keys."""
+        with self._lock:
+            if not keys:
+                return {}
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            values: dict[FullKey, ValueT] = {}
+            for ns_tuple, key in keys:
+                ns = Namespace(ns_tuple)
+                if ns in self._cache and key in self._cache[ns]:
+                    enc, val, expiry = self._cache[ns][key]
+                    if expiry is None or now < expiry:
+                        values[(ns, key)] = self.serde.loads_typed((enc, val))
+                    else:
+                        del self._cache[ns][key]
+            return values
+
+    async def aget(self, keys: Sequence[FullKey]) -> dict[FullKey, ValueT]:
+        """Asynchronously get the cached values for the given keys."""
+        return self.get(keys)
+
+    def set(self, keys: dict[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Set the cached values for the given keys."""
+        with self._lock:
+            now = datetime.datetime.now(datetime.timezone.utc).timestamp()
+            for (ns, key), (value, ttl) in keys.items():
+                if ttl is not None:
+                    delta = datetime.timedelta(seconds=ttl)
+                    expiry: float | None = (now + delta).timestamp()
+                else:
+                    expiry = None
+                if ns not in self._cache:
+                    self._cache[ns] = {}
+                self._cache[ns][key] = (
+                    *self.serde.dumps_typed(value),
+                    expiry,
+                )
+
+    async def aset(self, keys: dict[FullKey, tuple[ValueT, int | None]]) -> None:
+        """Asynchronously set the cached values for the given keys."""
+        self.set(keys)
+
+    def delete(self, keys: Sequence[Namespace]) -> None:
+        """Delete the cached values for the given namespaces."""
+        with self._lock:
+            for ns in keys:
+                if ns in self._cache:
+                    del self._cache[ns]
+
+    async def adelete(self, keys: Sequence[Namespace]) -> None:
+        """Asynchronously delete the cached values for the given namespaces."""
+        self.delete(keys)

libs/langgraph/tests/conftest.py

@@ -1,6 +1,4 @@
-import os
 import sys
-import tempfile
 from collections.abc import AsyncIterator, Iterator
 from contextlib import asynccontextmanager
 from typing import Optional
@@ -14,7 +12,8 @@ from psycopg_pool import AsyncConnectionPool, ConnectionPool
 from pytest_mock import MockerFixture
 
 from langgraph.cache.base import BaseCache
-from langgraph.cache.file import FileCache
+from langgraph.cache.memory import InMemoryCache
+from langgraph.cache.sqlite import SqliteCache
 from langgraph.checkpoint.base import BaseCheckpointSaver
 from langgraph.checkpoint.postgres import PostgresSaver, ShallowPostgresSaver
 from langgraph.checkpoint.postgres.aio import (
@@ -365,18 +364,14 @@ async def _store_postgres_aio_pool():
             await conn.execute(f"DROP DATABASE {database}")
 
 
-@pytest.fixture(scope="function")
-def file_cache() -> Iterator[BaseCache]:
-    _, path = tempfile.mkstemp()
-    os.remove(path)
-    try:
-        yield FileCache(path=path)
-    finally:
-        # Cleanup the file
-        try:
-            os.remove(path)
-        except OSError:
-            pass
+@pytest.fixture(scope="function", params=["sqlite", "memory"])
+def cache(request: pytest.FixtureRequest) -> Iterator[BaseCache]:
+    if request.param == "sqlite":
+        yield SqliteCache(path=":memory:")
+    elif request.param == "memory":
+        yield InMemoryCache()
+    else:
+        raise ValueError(f"Unknown cache type: {request.param}")
 
 
 @pytest.fixture(scope="function")

libs/langgraph/tests/test_pregel.py

@@ -3568,7 +3568,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
 
 @pytest.mark.parametrize("with_cache", [True, False])
 def test_in_one_fan_out_state_graph_waiting_edge_multiple(
-    with_cache: bool, file_cache: BaseCache
+    with_cache: bool, cache: BaseCache
 ) -> None:
     def sorted_add(
         x: list[str], y: Union[list[str], list[tuple[str, str]]]
@@ -3634,7 +3634,7 @@ def test_in_one_fan_out_state_graph_waiting_edge_multiple(
     workflow.add_conditional_edges("decider", decider_cond)
     workflow.set_finish_point("qa")
 
-    app = workflow.compile(cache=file_cache)
+    app = workflow.compile(cache=cache)
 
     assert app.invoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: analyzed: query: what is weather in sf",
@@ -6628,7 +6628,7 @@ def test_multiple_interrupts_functional(
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
 def test_multiple_interrupts_functional_cache(
-    request: pytest.FixtureRequest, checkpointer_name: str, file_cache: BaseCache
+    request: pytest.FixtureRequest, checkpointer_name: str, cache: BaseCache
 ):
     """Test multiple interrupts with functional API."""
     checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
@@ -6642,7 +6642,7 @@ def test_multiple_interrupts_functional_cache(
         counter += 1
         return 2 * x
 
-    @entrypoint(checkpointer=checkpointer, cache=file_cache)
+    @entrypoint(checkpointer=checkpointer, cache=cache)
     def graph(state: dict) -> dict:
         """React tool."""
 
@@ -6683,7 +6683,7 @@ def test_multiple_interrupts_functional_cache(
     assert counter == 3
 
     # clear cache
-    double.clear_cache(file_cache)
+    double.clear_cache(cache)
 
     # should recompute now
     configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}

libs/langgraph/tests/test_pregel_async.py

@@ -5399,7 +5399,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_plus_regular(
 
 @pytest.mark.parametrize("with_cache", [True, False])
 async def test_in_one_fan_out_state_graph_waiting_edge_multiple(
-    with_cache: bool, file_cache: BaseCache
+    with_cache: bool, cache: BaseCache
 ) -> None:
     def sorted_add(
         x: list[str], y: Union[list[str], list[tuple[str, str]]]
@@ -5465,7 +5465,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_multiple(
     workflow.add_conditional_edges("decider", decider_cond)
     workflow.set_finish_point("qa")
 
-    app = workflow.compile(cache=file_cache)
+    app = workflow.compile(cache=cache)
 
     assert await app.ainvoke({"query": "what is weather in sf"}) == {
         "query": "analyzed: query: analyzed: query: what is weather in sf",
@@ -7551,7 +7551,7 @@ async def test_multiple_interrupts_functional(checkpointer_name: str) -> None:
 @NEEDS_CONTEXTVARS
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_multiple_interrupts_functional_cache(
-    checkpointer_name: str, file_cache: BaseCache
+    checkpointer_name: str, cache: BaseCache
 ):
     """Test multiple interrupts with functional API."""
     async with awith_checkpointer(checkpointer_name) as checkpointer:
@@ -7564,7 +7564,7 @@ async def test_multiple_interrupts_functional_cache(
             counter += 1
             return 2 * x
 
-        @entrypoint(checkpointer=checkpointer, cache=file_cache)
+        @entrypoint(checkpointer=checkpointer, cache=cache)
         def graph(state: dict) -> dict:
             """React tool."""
 
@@ -7604,7 +7604,7 @@ async def test_multiple_interrupts_functional_cache(
         assert counter == 3
 
         # clear the cache
-        await double.aclear_cache(file_cache)
+        await double.aclear_cache(cache)
 
         # now should recompute
         configurable = {"configurable": {"thread_id": str(uuid.uuid4())}}