feat: Add support for memory-only kvstore

llama_stack/providers/utils/kvstore/config.py

@@ -53,7 +53,7 @@ class SqliteKVStoreConfig(CommonConfig):
     type: Literal["sqlite"] = KVStoreType.sqlite.value
     db_path: str = Field(
         default=(RUNTIME_BASE_DIR / "kvstore.db").as_posix(),
-        description="File path for the sqlite database",
+        description="File path for the sqlite database. Use ':memory:' for an in-memory database",
     )
 
     @classmethod

llama_stack/providers/utils/kvstore/sqlite/sqlite.py

@@ -21,67 +21,79 @@ class SqliteKVStoreImpl(KVStore):
     def __init__(self, config: SqliteKVStoreConfig):
         self.db_path = config.db_path
         self.table_name = "kvstore"
+        self._conn = None
 
     def __str__(self):
         return f"SqliteKVStoreImpl(db_path={self.db_path}, table_name={self.table_name})"
 
+    def _is_memory_db(self) -> bool:
+        """Check if this is an in-memory database."""
+        return self.db_path == ":memory:" or "mode=memory" in self.db_path
+
     async def initialize(self):
-        os.makedirs(os.path.dirname(self.db_path), exist_ok=True)
-        async with aiosqlite.connect(self.db_path) as db:
-            await db.execute(
-                f"""
-                CREATE TABLE IF NOT EXISTS {self.table_name} (
-                    key TEXT PRIMARY KEY,
-                    value TEXT,
-                    expiration TIMESTAMP
-                )
-            """
+        # Skip directory creation for in-memory databases and file: URIs
+        if not self._is_memory_db() and not self.db_path.startswith("file:"):
+            db_dir = os.path.dirname(self.db_path)
+            if db_dir:  # Only create if there's a directory component
+                os.makedirs(db_dir, exist_ok=True)
+
+        # Create persistent connection for all databases
+        self._conn = await aiosqlite.connect(self.db_path)
+        await self._conn.execute(
+            f"""
+            CREATE TABLE IF NOT EXISTS {self.table_name} (
+                key TEXT PRIMARY KEY,
+                value TEXT,
+                expiration TIMESTAMP
             )
-            await db.commit()
+        """
+        )
+        await self._conn.commit()
+
+    async def close(self):
+        """Close the persistent connection."""
+        if self._conn:
+            await self._conn.close()
+            self._conn = None
 
     async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
-        async with aiosqlite.connect(self.db_path) as db:
-            await db.execute(
-                f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
-                (key, value, expiration),
-            )
-            await db.commit()
+        await self._conn.execute(
+            f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
+            (key, value, expiration),
+        )
+        await self._conn.commit()
 
     async def get(self, key: str) -> str | None:
-        async with aiosqlite.connect(self.db_path) as db:
-            async with db.execute(f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)) as cursor:
-                row = await cursor.fetchone()
-                if row is None:
-                    return None
-                value, expiration = row
-                if not isinstance(value, str):
-                    logger.warning(f"Expected string value for key {key}, got {type(value)}, returning None")
-                    return None
-                return value
+        async with self._conn.execute(f"SELECT value, expiration FROM {self.table_name} WHERE key = ?", (key,)) as cursor:
+            row = await cursor.fetchone()
+            if row is None:
+                return None
+            value, expiration = row
+            if not isinstance(value, str):
+                logger.warning(f"Expected string value for key {key}, got {type(value)}, returning None")
+                return None
+            return value
 
     async def delete(self, key: str) -> None:
-        async with aiosqlite.connect(self.db_path) as db:
-            await db.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
-            await db.commit()
+        await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
+        await self._conn.commit()
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
-        async with aiosqlite.connect(self.db_path) as db:
-            async with db.execute(
-                f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
-                (start_key, end_key),
-            ) as cursor:
-                result = []
-                async for row in cursor:
-                    _, value, _ = row
-                    result.append(value)
-                return result
+        async with self._conn.execute(
+            f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
+            (start_key, end_key),
+        ) as cursor:
+            result = []
+            async for row in cursor:
+                _, value, _ = row
+                result.append(value)
+            return result
 
     async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
         """Get all keys in the given range."""
-        async with aiosqlite.connect(self.db_path) as db:
-            cursor = await db.execute(
-                f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
-                (start_key, end_key),
-            )
-            rows = await cursor.fetchall()
-            return [row[0] for row in rows]
+        cursor = await self._conn.execute(
+            f"SELECT key FROM {self.table_name} WHERE key >= ? AND key <= ?",
+            (start_key, end_key),
+        )
+        rows = await cursor.fetchall()
+        return [row[0] for row in rows]

tests/unit/utils/kvstore/test_sqlite_memory.py

@@ -0,0 +1,121 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import pytest
+
+from llama_stack.providers.utils.kvstore.config import SqliteKVStoreConfig
+from llama_stack.providers.utils.kvstore.sqlite import SqliteKVStoreImpl
+
+
+@pytest.mark.asyncio
+async def test_memory_kvstore_basic_operations():
+    """Test basic CRUD operations with :memory: database."""
+    config = SqliteKVStoreConfig(db_path=":memory:")
+    kvstore = SqliteKVStoreImpl(config)
+    await kvstore.initialize()
+
+    try:
+        # Test set and get
+        await kvstore.set("test_key", "test_value")
+        value = await kvstore.get("test_key")
+        assert value == "test_value"
+
+        # Test get non-existent key
+        value = await kvstore.get("non_existent")
+        assert value is None
+
+        # Test delete
+        await kvstore.delete("test_key")
+        value = await kvstore.get("test_key")
+        assert value is None
+    finally:
+        await kvstore.close()
+
+
+@pytest.mark.asyncio
+async def test_memory_kvstore_range_operations():
+    """Test range operations with :memory: database."""
+    config = SqliteKVStoreConfig(db_path=":memory:")
+    kvstore = SqliteKVStoreImpl(config)
+    await kvstore.initialize()
+
+    try:
+        # Set multiple keys
+        await kvstore.set("key_1", "value_1")
+        await kvstore.set("key_2", "value_2")
+        await kvstore.set("key_3", "value_3")
+        await kvstore.set("key_4", "value_4")
+
+        # Test values_in_range
+        values = await kvstore.values_in_range("key_1", "key_3")
+        assert len(values) == 3
+        assert "value_1" in values
+        assert "value_2" in values
+        assert "value_3" in values
+
+        # Test keys_in_range
+        keys = await kvstore.keys_in_range("key_2", "key_4")
+        assert len(keys) == 3
+        assert "key_2" in keys
+        assert "key_3" in keys
+        assert "key_4" in keys
+    finally:
+        await kvstore.close()
+
+
+@pytest.mark.asyncio
+async def test_memory_kvstore_multiple_instances():
+    """Test that multiple :memory: instances are independent."""
+    config1 = SqliteKVStoreConfig(db_path=":memory:")
+    kvstore1 = SqliteKVStoreImpl(config1)
+    await kvstore1.initialize()
+
+    config2 = SqliteKVStoreConfig(db_path=":memory:")
+    kvstore2 = SqliteKVStoreImpl(config2)
+    await kvstore2.initialize()
+
+    try:
+        # Set value in first instance
+        await kvstore1.set("shared_key", "value_1")
+
+        # Verify second instance doesn't have the value
+        value = await kvstore2.get("shared_key")
+        assert value is None
+
+        # Set different value in second instance
+        await kvstore2.set("shared_key", "value_2")
+
+        # Verify instances remain independent
+        value1 = await kvstore1.get("shared_key")
+        value2 = await kvstore2.get("shared_key")
+        assert value1 == "value_1"
+        assert value2 == "value_2"
+    finally:
+        await kvstore1.close()
+        await kvstore2.close()
+
+
+@pytest.mark.asyncio
+async def test_memory_kvstore_persistence_behavior():
+    """Test that :memory: database doesn't persist across instances."""
+    config = SqliteKVStoreConfig(db_path=":memory:")
+
+    # First instance
+    kvstore1 = SqliteKVStoreImpl(config)
+    await kvstore1.initialize()
+    await kvstore1.set("test_key", "test_value")
+    await kvstore1.close()
+
+    # Create new instance with same config
+    kvstore2 = SqliteKVStoreImpl(config)
+    await kvstore2.initialize()
+
+    try:
+        # Data should not persist
+        value = await kvstore2.get("test_key")
+        assert value is None
+    finally:
+        await kvstore2.close()