kvstore impls for redis / sqlite moved

llama_stack/providers/utils/kvstore/__init__.py

@@ -0,0 +1,5 @@
+# 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.

llama_stack/providers/utils/kvstore/api.py

@@ -0,0 +1,28 @@
+# 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.
+
+from datetime import datetime
+from typing import Any, List, Optional, Protocol
+
+from pydantic import BaseModel
+
+
+class KVStoreValue(BaseModel):
+    key: str
+    value: Any
+    expiration: Optional[datetime] = None
+
+
+class KVStore(Protocol):
+    async def set(
+        self, key: str, value: Any, expiration: Optional[datetime] = None
+    ) -> None: ...
+
+    async def get(self, key: str) -> Optional[KVStoreValue]: ...
+
+    async def delete(self, key: str) -> None: ...
+
+    async def range(self, start_key: str, end_key: str) -> List[KVStoreValue]: ...

llama_stack/providers/utils/kvstore/config.py

@@ -0,0 +1,52 @@
+# 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.
+
+from enum import Enum
+from typing import Literal, Optional, Union
+
+from pydantic import BaseModel
+from typing_extensions import Annotated
+
+
+class KVStoreType(Enum):
+    redis = "redis"
+    sqlite = "sqlite"
+    pgvector = "pgvector"
+
+
+class CommonConfig(BaseModel):
+    namespace: Optional[str] = Field(
+        default=None,
+        description="All keys will be prefixed with this namespace",
+    )
+
+
+class RedisKVStoreImplConfig(CommonConfig):
+    type: Literal[KVStoreType.redis.value] = KVStoreType.redis.value
+    host: str = "localhost"
+    port: int = 6379
+
+
+class SqliteKVStoreImplConfig(CommonConfig):
+    type: Literal[KVStoreType.sqlite.value] = KVStoreType.sqlite.value
+    db_path: str = Field(
+        description="File path for the sqlite database",
+    )
+
+
+class PGVectorKVStoreImplConfig(CommonConfig):
+    type: Literal[KVStoreType.pgvector.value] = KVStoreType.pgvector.value
+    host: str = "localhost"
+    port: int = 5432
+    db: str = "llamastack"
+    user: str
+    password: Optional[str] = None
+
+
+KVStoreConfig = Annotated[
+    Union[RedisKVStoreImplConfig, SqliteKVStoreImplConfig, PGVectorKVStoreImplConfig],
+    Field(discriminator="type"),
+]

llama_stack/providers/utils/kvstore/kvstore.py

@@ -0,0 +1,30 @@
+# 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.
+
+from .api import *  # noqa: F403
+from .config import *  # noqa: F403
+
+
+def kvstore_dependencies():
+    return ["aiosqlite", "psycopg2-binary", "redis"]
+
+
+async def kvstore_impl(config: KVStoreConfig) -> KVStore:
+    if config.type == KVStoreType.redis:
+        from .redis import RedisKVStoreImpl
+
+        impl = RedisKVStoreImpl(config)
+    elif config.type == KVStoreType.sqlite:
+        from .sqlite import SqliteKVStoreImpl
+
+        impl = SqliteKVStoreImpl(config)
+    elif config.type == KVStoreType.pgvector:
+        raise NotImplementedError()
+    else:
+        raise ValueError(f"Unknown kvstore type {config.type}")
+
+    await impl.initialize()
+    return impl

llama_stack/providers/utils/kvstore/redis/__init__.py

@@ -0,0 +1,7 @@
+# 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.
+
+from .redis import RedisKVStoreImpl  # noqa: F401

llama_stack/providers/utils/kvstore/redis/redis.py

@@ -0,0 +1,60 @@
+# 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.
+
+from datetime import datetime, timedelta
+from typing import Any, List, Optional
+
+from redis.asyncio import Redis
+
+from ..api import *  # noqa: F403
+from ..config import RedisKVStoreImplConfig
+
+
+class RedisKVStoreImpl(KVStore):
+    def __init__(self, config: RedisKVStoreImplConfig):
+        self.config = config
+
+    async def initialize(self) -> None:
+        self.redis = Redis.from_url(self.config.url)
+
+    def _namespaced_key(self, key: str) -> str:
+        if not self.config.namespace:
+            return key
+        return f"{self.config.namespace}:{key}"
+
+    async def set(
+        self, key: str, value: Any, expiration: Optional[datetime] = None
+    ) -> None:
+        key = self._namespaced_key(key)
+        await self.redis.set(key, value)
+        if expiration:
+            await self.redis.expireat(key, expiration)
+
+    async def get(self, key: str) -> Optional[KVStoreValue]:
+        key = self._namespaced_key(key)
+        value = await self.redis.get(key)
+        if value is None:
+            return None
+        ttl = await self.redis.ttl(key)
+        expiration = datetime.now() + timedelta(seconds=ttl) if ttl > 0 else None
+        return KVStoreValue(key=key, value=value, expiration=expiration)
+
+    async def delete(self, key: str) -> None:
+        key = self._namespaced_key(key)
+        await self.redis.delete(key)
+
+    async def range(self, start_key: str, end_key: str) -> List[KVStoreValue]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+
+        keys = await self.redis.keys(f"{start_key}*")
+        result = []
+        for key in keys:
+            if key <= end_key:
+                value = await self.get(key)
+                if value:
+                    result.append(value)
+        return result

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

@@ -0,0 +1,7 @@
+# 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.
+
+from .sqlite import SqliteKVStoreImpl  # noqa: F401

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

@@ -0,0 +1,19 @@
+# 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.
+
+from llama_models.schema_utils import json_schema_type
+from pydantic import BaseModel, Field
+
+
+@json_schema_type
+class SqliteControlPlaneConfig(BaseModel):
+    db_path: str = Field(
+        description="File path for the sqlite database",
+    )
+    table_name: str = Field(
+        default="llamastack_control_plane",
+        description="Table into which all the keys will be placed",
+    )

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

@@ -0,0 +1,77 @@
+# 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 json
+from datetime import datetime
+from typing import Any, List, Optional
+
+import aiosqlite
+
+from ..api import *  # noqa: F403
+from ..config import SqliteKVStoreConfig
+
+
+class SqliteKVStoreImpl(KVStore):
+    def __init__(self, config: SqliteKVStoreConfig):
+        self.db_path = config.db_path
+        self.table_name = config.table_name
+
+    async def initialize(self):
+        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
+                )
+            """
+            )
+            await db.commit()
+
+    async def set(
+        self, key: str, value: Any, expiration: Optional[datetime] = 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, json.dumps(value), expiration),
+            )
+            await db.commit()
+
+    async def get(self, key: str) -> Optional[KVStoreValue]:
+        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
+                return KVStoreValue(
+                    key=key, value=json.loads(value), expiration=expiration
+                )
+
+    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()
+
+    async def range(self, start_key: str, end_key: str) -> List[KVStoreValue]:
+        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:
+                    key, value, expiration = row
+                    result.append(
+                        KVStoreValue(
+                            key=key, value=json.loads(value), expiration=expiration
+                        )
+                    )
+                return result