diff --git a/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/cache/sqlite/__init__.py new file mode 100644 index 000000000..780e88282 --- /dev/null +++ b/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 diff --git a/libs/checkpoint/langgraph/cache/memory/__init__.py b/libs/checkpoint/langgraph/cache/memory/__init__.py new file mode 100644 index 000000000..d93744e80 --- /dev/null +++ b/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) diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 6de17548a..7a4829e03 100644 --- a/libs/langgraph/tests/conftest.py +++ b/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") diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index bca22ceb2..d26a6e408 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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())}} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 63c0a4b3e..0c815f49e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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())}}