SQL Alchemy based persistent cache implementation. Using it with sqlite backend is the easiest option for local workflows.

Why SQLite: performance wise it is a terrible choice and any sane deployment should use something like memcached or redis. And SQL (e.g. postgres) is just an overkill. However sqlite is a zero-conf option that does not require any daemon running. So it is very easy to use for local workflows when we don't care that much about performance.

PiperOrigin-RevId: 829383146

Author: kibergus

genai_processors/sql_cache.py

@@ -0,0 +1,243 @@
+# Copyright 2025 DeepMind Technologies Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""SQLite-based cache for processors using SQLAlchemy.
+
+This is a persistent cache suitable for small Parts, for example metadata
+extracted using constrained decoding. It might not scale for Parts containing
+large amounts of data e.g. video frames.
+
+Cache is very handy for writing long-running agents or during development as it
+provides a lightweight way to resume execution from the point of failure. By
+wrapping all LLM calls and other heavy logic in a cache, while keeping the rest
+of the code idempotent, restarting the agent from the beginning will promptly
+catch up to the place where it has previously failed. During development one can
+force the changed code to be rerun by altering key_prefix e.g. by appending code
+version to it.
+"""
+
+import asyncio
+from collections.abc import Callable
+import contextlib
+import datetime
+import json
+from typing import AsyncIterator
+
+from genai_processors import cache
+from genai_processors import cache_base
+from genai_processors import content_api
+import sqlalchemy
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.ext.asyncio import create_async_engine
+import sqlalchemy.orm
+from typing_extensions import override
+
+
+ProcessorContent = content_api.ProcessorContent
+
+_Base = sqlalchemy.orm.declarative_base()
+
+
+class _ContentCacheEntry(_Base):
+  """SQLAlchemy model for the cache table."""
+
+  __tablename__ = 'content_cache'
+
+  key = sqlalchemy.Column(sqlalchemy.String, primary_key=True)
+  value = sqlalchemy.Column(sqlalchemy.LargeBinary)
+  expires_at = sqlalchemy.Column(sqlalchemy.DateTime(timezone=True), index=True)
+
+
+@contextlib.asynccontextmanager
+async def sql_cache(
+    db_url: str,
+    ttl_hours: float | None = 12,
+    hash_fn: (
+        Callable[[content_api.ProcessorContentTypes], str | None] | None
+    ) = None,
+) -> AsyncIterator['SqlCache']:
+  """Context manager that creates an SqlCache instance.
+
+  Args:
+    db_url: SQLAlchemy database URL.
+    ttl_hours: Time-to-live for cache items in hours. If None, the cache items
+      never expire.
+    hash_fn: Function to convert a content_api.ProcessorContentTypes query into
+      a string key. If None, `cache.default_processor_content_hash` is used.
+
+  Yields:
+    A SqlCache instance.
+  """
+  engine = create_async_engine(db_url)
+
+  async with engine.begin() as conn:
+    await conn.run_sync(_Base.metadata.create_all)
+
+  async with AsyncSession(engine) as session:
+    yield SqlCache(
+        session=session,
+        ttl_hours=ttl_hours,
+        hash_fn=hash_fn or cache.default_processor_content_hash,
+        lock=asyncio.Lock(),
+    )
+
+
+class SqlCache(cache_base.CacheBase):
+  """An SQLAlchemy based persistent content cache."""
+
+  def __init__(
+      self,
+      *,
+      session: AsyncSession,
+      ttl_hours: float | None,
+      hash_fn: Callable[[content_api.ProcessorContentTypes], str | None],
+      lock: asyncio.Lock,
+  ):
+    """Prefer using sql_cache() factory to construct the cache."""
+    self._hash_fn = hash_fn
+    self._ttl = (
+        datetime.timedelta(hours=ttl_hours) if ttl_hours is not None else None
+    )
+    self._session = session
+    self._lock = lock
+
+  @property
+  @override
+  def hash_fn(
+      self,
+  ) -> Callable[[content_api.ProcessorContentTypes], str | None]:
+    return self._hash_fn
+
+  @override
+  def with_key_prefix(self, prefix: str) -> 'SqlCache':
+    """Creates a new SqlCache instance with a key prefix.
+
+    Args:
+      prefix: String to prepend to generated string keys.
+
+    Returns:
+      A new SqlCache instance with the given prefix.
+    """
+    # This creates a new instance but shares the same session.
+    return SqlCache(
+        session=self._session,
+        ttl_hours=self._ttl.total_seconds() / 3600
+        if self._ttl is not None
+        else None,
+        hash_fn=cache.prefixed_hash_fn(self._hash_fn, prefix),
+        lock=self._lock,
+    )
+
+  @override
+  async def lookup(
+      self,
+      query: content_api.ProcessorContentTypes | None = None,
+      *,
+      key: str | None = None,
+  ) -> content_api.ProcessorContent | cache_base.CacheMissT:
+    query_key = key if key is not None else self._hash_fn(query)
+
+    async with self._lock:
+      item = await self._session.get(_ContentCacheEntry, query_key)
+      if item is None:
+        return cache_base.CacheMiss
+
+      if self._ttl is not None:
+        # Ensure item.expires_at is offset-aware for comparison
+        # Assuming stored times are UTC, add UTC timezone info if none.
+        expires_at = item.expires_at
+        if expires_at.tzinfo is None:
+          expires_at = expires_at.replace(tzinfo=datetime.timezone.utc)
+
+        if expires_at < datetime.datetime.now(datetime.timezone.utc):
+          await self._remove_by_string_key(query_key)
+          return cache_base.CacheMiss
+
+      try:
+        return _deserialize_content(item.value)
+      except Exception:  # pylint: disable=broad-exception-caught
+        await self._remove_by_string_key(query_key)
+        return cache_base.CacheMiss
+
+  @override
+  async def put(
+      self,
+      *,
+      query: content_api.ProcessorContentTypes | None = None,
+      key: str | None = None,
+      value: content_api.ProcessorContentTypes,
+  ) -> None:
+    query_key = key if key is not None else self._hash_fn(query)
+
+    data_to_cache_bytes = _serialize_content(
+        content_api.ProcessorContent(value)
+    )
+
+    expires_at = None
+    if self._ttl is not None:
+      expires_at = datetime.datetime.now(datetime.timezone.utc) + self._ttl
+
+    item = _ContentCacheEntry()
+    item.key = query_key
+    item.value = data_to_cache_bytes
+    item.expires_at = expires_at
+    async with self._lock:
+      self._session.add(item)
+      await self._cleanup_expired()
+      await self._session.commit()
+
+  async def _remove_by_string_key(self, string_key: str) -> None:
+    """Internal helper to remove by the actual string key."""
+    item = await self._session.get(_ContentCacheEntry, string_key)
+    if item:
+      await self._session.delete(item)
+      await self._session.commit()
+
+  @override
+  async def remove(self, query: content_api.ProcessorContentTypes) -> None:
+    query_key = self._hash_fn(query)
+    if query_key is None:
+      return
+    async with self._lock:
+      await self._remove_by_string_key(query_key)
+
+  async def _cleanup_expired(self) -> None:
+    """Removes expired items from the cache."""
+    if self._ttl is None:
+      return
+    now = datetime.datetime.now(datetime.timezone.utc)
+    expired_items = await self._session.execute(
+        sqlalchemy.select(_ContentCacheEntry).where(
+            _ContentCacheEntry.expires_at < now
+        )
+    )
+    for item in expired_items.scalars():
+      await self._session.delete(item)
+
+
+def _serialize_content(value: ProcessorContent) -> bytes:
+  """Serializes ProcessorContent to bytes (via JSON)."""
+  list_of_part_dicts_val = [part.to_dict() for part in value.all_parts]
+  json_string_val = json.dumps(list_of_part_dicts_val)
+  return json_string_val.encode('utf-8')
+
+
+def _deserialize_content(data_bytes: bytes) -> ProcessorContent:
+  """Deserializer for ProcessorContent from bytes (via JSON)."""
+  json_string_val = data_bytes.decode('utf-8')
+  list_of_part_dicts_val = json.loads(json_string_val)
+  return ProcessorContent([
+      content_api.ProcessorPart.from_dict(data=pd)
+      for pd in list_of_part_dicts_val
+  ])

genai_processors/tests/sql_cache_test.py

@@ -0,0 +1,118 @@
+import asyncio
+import datetime
+import os
+import tempfile
+import unittest
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from genai_processors import cache_base
+from genai_processors import content_api
+from genai_processors import sql_cache
+import sqlalchemy
+
+# Test Content
+TEST_QUERY = content_api.ProcessorContent('test query')
+TEST_VALUE = content_api.ProcessorContent('test value')
+TEST_QUERY_2 = content_api.ProcessorContent('test query 2')
+TEST_VALUE_2 = content_api.ProcessorContent('test value 2')
+
+
+class SqlCacheTest(parameterized.TestCase, unittest.IsolatedAsyncioTestCase):
+
+  def setUp(self):
+    super().setUp()
+    with tempfile.NamedTemporaryFile(suffix='.sqlite', delete=False) as tmp:
+      self.db_url = f'sqlite+aiosqlite:///{tmp.name}'
+    self.addCleanup(os.remove, tmp.name)
+
+  async def test_cache_put_and_lookup(self):
+    """Tests basic put and lookup functionality."""
+    async with sql_cache.sql_cache(self.db_url) as cache:
+      await cache.put(TEST_QUERY, TEST_VALUE)
+      result = await cache.lookup(TEST_QUERY)
+      self.assertEqual(result, TEST_VALUE)
+
+  async def test_cache_miss(self):
+    """Tests cache miss for a non-existent key."""
+    async with sql_cache.sql_cache(self.db_url) as cache:
+      result = await cache.lookup(TEST_QUERY)
+      self.assertIs(result, cache_base.CacheMiss)
+
+  async def test_cache_ttl(self):
+    """Tests that cache entries expire after TTL."""
+    async with sql_cache.sql_cache(
+        self.db_url, ttl_hours=0.0001 / 60 / 60
+    ) as cache:  # Very short TTL
+      await cache.put(TEST_QUERY, TEST_VALUE)
+      await asyncio.sleep(0.0002)  # Wait for TTL to expire
+      result = await cache.lookup(TEST_QUERY)
+      self.assertIs(result, cache_base.CacheMiss)
+
+  async def test_cache_remove(self):
+    """Tests removing an item from the cache."""
+    async with sql_cache.sql_cache(self.db_url) as cache:
+      await cache.put(TEST_QUERY, TEST_VALUE)
+      await cache.remove(TEST_QUERY)
+      result = await cache.lookup(TEST_QUERY)
+      self.assertIs(result, cache_base.CacheMiss)
+
+  async def test_with_key_prefix(self):
+    """Tests that with_key_prefix creates a namespace."""
+    async with sql_cache.sql_cache(self.db_url) as cache1:
+      cache2 = cache1.with_key_prefix('prefix:')
+
+      await cache1.put(TEST_QUERY, TEST_VALUE)
+      await cache2.put(TEST_QUERY, TEST_VALUE_2)
+
+      result1 = await cache1.lookup(TEST_QUERY)
+      result2 = await cache2.lookup(TEST_QUERY)
+
+      self.assertEqual(result1, TEST_VALUE)
+      self.assertEqual(result2, TEST_VALUE_2)
+      self.assertNotEqual(result1, result2)
+
+  async def test_different_content_types(self):
+    """Tests caching with different content types."""
+    async with sql_cache.sql_cache(self.db_url) as cache:
+      query1 = content_api.ProcessorContent('text query')
+      value1 = content_api.ProcessorContent(['list ', 'of ', 'strings'])
+      query2 = content_api.ProcessorContent(
+          [content_api.ProcessorPart(b'imagedata', mimetype='image/png')]
+      )
+
+      value2 = content_api.ProcessorContent({'a': 1, 'b': True})
+
+      await cache.put(query1, value1)
+      await cache.put(query2, value2)
+
+      self.assertEqual(await cache.lookup(query1), value1)
+      self.assertEqual(await cache.lookup(query2), value2)
+
+  async def test_cleanup_expired(self):
+    """Tests that the _cleanup_expired method removes old entries."""
+    async with sql_cache.sql_cache(self.db_url, ttl_hours=0.0001) as cache:
+      await cache.put(TEST_QUERY, TEST_VALUE)
+      # Mock datetime to control time
+      with mock.patch('genai_processors.sql_cache.datetime') as mock_datetime:
+        mock_datetime.datetime.now.return_value = datetime.datetime.now(
+            datetime.timezone.utc
+        ) + datetime.timedelta(hours=1)
+        mock_datetime.timedelta.side_effect = datetime.timedelta
+        mock_datetime.timezone.utc = datetime.timezone.utc
+        await cache._cleanup_expired()
+
+      # Check that the entry is gone
+      async def _lookup():
+        stmt = sqlalchemy.select(sql_cache._ContentCacheEntry).where(
+            sql_cache._ContentCacheEntry.key == cache._hash_fn(TEST_QUERY)
+        )
+        result = await cache._session.execute(stmt)
+        return result.fetchone()
+
+      result = await _lookup()
+      self.assertIsNone(result)
+
+if __name__ == '__main__':
+  absltest.main()