llama-stack 0.0.42__py3-none-any.whl → 0.3.4__py3-none-any.whl

llama_stack/providers/utils/kvstore/kvstore.py

@@ -4,12 +4,28 @@
 # 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
+# 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 __future__ import annotations
+
+from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType
+
+from .api import KVStore
+from .config import KVStoreConfig
 
 
 def kvstore_dependencies():
-    return ["aiosqlite", "psycopg2-binary", "redis"]
+    """
+    Returns all possible kvstore dependencies for registry/provider specifications.
+
+    NOTE: For specific kvstore implementations, use config.pip_packages instead.
+    This function returns the union of all dependencies for cases where the specific
+    kvstore type is not known at declaration time (e.g., provider registries).
+    """
+    return ["aiosqlite", "psycopg2-binary", "redis", "pymongo"]
 
 
 class InmemoryKVStoreImpl(KVStore):
@@ -19,31 +35,61 @@ class InmemoryKVStoreImpl(KVStore):
     async def initialize(self) -> None:
         pass
 
-    async def get(self, key: str) -> Optional[str]:
+    async def get(self, key: str) -> str | None:
         return self._store.get(key)
 
     async def set(self, key: str, value: str) -> None:
         self._store[key] = value
 
-    async def range(self, start_key: str, end_key: str) -> List[str]:
-        return [
-            self._store[key]
-            for key in self._store.keys()
-            if key >= start_key and key < end_key
-        ]
+    async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
+        return [self._store[key] for key in self._store.keys() if key >= start_key and key < end_key]
+
+    async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
+        """Get all keys in the given range."""
+        return [key for key in self._store.keys() if key >= start_key and key < end_key]
+
+    async def delete(self, key: str) -> None:
+        del self._store[key]
 
 
-async def kvstore_impl(config: KVStoreConfig) -> KVStore:
-    if config.type == KVStoreType.redis.value:
+_KVSTORE_BACKENDS: dict[str, KVStoreConfig] = {}
+
+
+def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None:
+    """Register the set of available KV store backends for reference resolution."""
+    global _KVSTORE_BACKENDS
+
+    _KVSTORE_BACKENDS.clear()
+    for name, cfg in backends.items():
+        _KVSTORE_BACKENDS[name] = cfg
+
+
+async def kvstore_impl(reference: KVStoreReference) -> KVStore:
+    backend_name = reference.backend
+
+    backend_config = _KVSTORE_BACKENDS.get(backend_name)
+    if backend_config is None:
+        raise ValueError(f"Unknown KVStore backend '{backend_name}'. Registered backends: {sorted(_KVSTORE_BACKENDS)}")
+
+    config = backend_config.model_copy()
+    config.namespace = reference.namespace
+
+    if config.type == StorageBackendType.KV_REDIS.value:
         from .redis import RedisKVStoreImpl
 
         impl = RedisKVStoreImpl(config)
-    elif config.type == KVStoreType.sqlite.value:
+    elif config.type == StorageBackendType.KV_SQLITE.value:
         from .sqlite import SqliteKVStoreImpl
 
         impl = SqliteKVStoreImpl(config)
-    elif config.type == KVStoreType.postgres.value:
-        raise NotImplementedError()
+    elif config.type == StorageBackendType.KV_POSTGRES.value:
+        from .postgres import PostgresKVStoreImpl
+
+        impl = PostgresKVStoreImpl(config)
+    elif config.type == StorageBackendType.KV_MONGODB.value:
+        from .mongodb import MongoDBKVStoreImpl
+
+        impl = MongoDBKVStoreImpl(config)
     else:
         raise ValueError(f"Unknown kvstore type {config.type}")
 

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

@@ -0,0 +1,9 @@
+# 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 .mongodb import MongoDBKVStoreImpl
+
+__all__ = ["MongoDBKVStoreImpl"]

llama_stack/providers/utils/kvstore/mongodb/mongodb.py

@@ -0,0 +1,82 @@
+# 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 pymongo import AsyncMongoClient
+from pymongo.asynchronous.collection import AsyncCollection
+
+from llama_stack.log import get_logger
+from llama_stack.providers.utils.kvstore import KVStore
+
+from ..config import MongoDBKVStoreConfig
+
+log = get_logger(name=__name__, category="providers::utils")
+
+
+class MongoDBKVStoreImpl(KVStore):
+    def __init__(self, config: MongoDBKVStoreConfig):
+        self.config = config
+        self.conn: AsyncMongoClient | None = None
+
+    @property
+    def collection(self) -> AsyncCollection:
+        if self.conn is None:
+            raise RuntimeError("MongoDB connection is not initialized")
+        return self.conn[self.config.db][self.config.collection_name]
+
+    async def initialize(self) -> None:
+        try:
+            conn_creds = {
+                "host": self.config.host,
+                "port": self.config.port,
+                "username": self.config.user,
+                "password": self.config.password,
+            }
+            conn_creds = {k: v for k, v in conn_creds.items() if v is not None}
+            self.conn = AsyncMongoClient(**conn_creds)
+        except Exception as e:
+            log.exception("Could not connect to MongoDB database server")
+            raise RuntimeError("Could not connect to MongoDB database server") from e
+
+    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: str, expiration: datetime | None = None) -> None:
+        key = self._namespaced_key(key)
+        update_query = {"$set": {"value": value, "expiration": expiration}}
+        await self.collection.update_one({"key": key}, update_query, upsert=True)
+
+    async def get(self, key: str) -> str | None:
+        key = self._namespaced_key(key)
+        query = {"key": key}
+        result = await self.collection.find_one(query, {"value": 1, "_id": 0})
+        return result["value"] if result else None
+
+    async def delete(self, key: str) -> None:
+        key = self._namespaced_key(key)
+        await self.collection.delete_one({"key": key})
+
+    async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        query = {
+            "key": {"$gte": start_key, "$lt": end_key},
+        }
+        cursor = self.collection.find(query, {"value": 1, "_id": 0}).sort("key", 1)
+        result = []
+        async for doc in cursor:
+            result.append(doc["value"])
+        return result
+
+    async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+        query = {"key": {"$gte": start_key, "$lt": end_key}}
+        cursor = self.collection.find(query, {"key": 1, "_id": 0}).sort("key", 1)
+        return [doc["key"] for doc in cursor]

llama_stack/providers/utils/kvstore/postgres/__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 .postgres import PostgresKVStoreImpl  # noqa: F401 F403

llama_stack/providers/utils/kvstore/postgres/postgres.py

@@ -0,0 +1,114 @@
+# 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
+
+import psycopg2
+from psycopg2.extras import DictCursor
+
+from llama_stack.log import get_logger
+
+from ..api import KVStore
+from ..config import PostgresKVStoreConfig
+
+log = get_logger(name=__name__, category="providers::utils")
+
+
+class PostgresKVStoreImpl(KVStore):
+    def __init__(self, config: PostgresKVStoreConfig):
+        self.config = config
+        self.conn = None
+        self.cursor = None
+
+    async def initialize(self) -> None:
+        try:
+            self.conn = psycopg2.connect(
+                host=self.config.host,
+                port=self.config.port,
+                database=self.config.db,
+                user=self.config.user,
+                password=self.config.password,
+                sslmode=self.config.ssl_mode,
+                sslrootcert=self.config.ca_cert_path,
+            )
+            self.conn.autocommit = True
+            self.cursor = self.conn.cursor(cursor_factory=DictCursor)
+
+            # Create table if it doesn't exist
+            self.cursor.execute(
+                f"""
+                CREATE TABLE IF NOT EXISTS {self.config.table_name} (
+                    key TEXT PRIMARY KEY,
+                    value TEXT,
+                    expiration TIMESTAMP
+                )
+                """
+            )
+        except Exception as e:
+            log.exception("Could not connect to PostgreSQL database server")
+            raise RuntimeError("Could not connect to PostgreSQL database server") from e
+
+    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: str, expiration: datetime | None = None) -> None:
+        key = self._namespaced_key(key)
+        self.cursor.execute(
+            f"""
+            INSERT INTO {self.config.table_name} (key, value, expiration)
+            VALUES (%s, %s, %s)
+            ON CONFLICT (key) DO UPDATE
+            SET value = EXCLUDED.value, expiration = EXCLUDED.expiration
+            """,
+            (key, value, expiration),
+        )
+
+    async def get(self, key: str) -> str | None:
+        key = self._namespaced_key(key)
+        self.cursor.execute(
+            f"""
+            SELECT value FROM {self.config.table_name}
+            WHERE key = %s
+            AND (expiration IS NULL OR expiration > NOW())
+            """,
+            (key,),
+        )
+        result = self.cursor.fetchone()
+        return result[0] if result else None
+
+    async def delete(self, key: str) -> None:
+        key = self._namespaced_key(key)
+        self.cursor.execute(
+            f"DELETE FROM {self.config.table_name} WHERE key = %s",
+            (key,),
+        )
+
+    async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+
+        self.cursor.execute(
+            f"""
+            SELECT value FROM {self.config.table_name}
+            WHERE key >= %s AND key < %s
+            AND (expiration IS NULL OR expiration > NOW())
+            ORDER BY key
+            """,
+            (start_key, end_key),
+        )
+        return [row[0] for row in self.cursor.fetchall()]
+
+    async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
+        start_key = self._namespaced_key(start_key)
+        end_key = self._namespaced_key(end_key)
+
+        self.cursor.execute(
+            f"SELECT key FROM {self.config.table_name} WHERE key >= %s AND key < %s",
+            (start_key, end_key),
+        )
+        return [row[0] for row in self.cursor.fetchall()]

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

@@ -5,11 +5,10 @@
 # the root directory of this source tree.
 
 from datetime import datetime
-from typing import List, Optional
 
 from redis.asyncio import Redis
 
-from ..api import *  # noqa: F403
+from ..api import KVStore
 from ..config import RedisKVStoreConfig
 
 
@@ -25,28 +24,53 @@ class RedisKVStoreImpl(KVStore):
             return key
         return f"{self.config.namespace}:{key}"
 
-    async def set(
-        self, key: str, value: str, expiration: Optional[datetime] = None
-    ) -> None:
+    async def set(self, key: str, value: str, expiration: datetime | None = 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[str]:
+    async def get(self, key: str) -> str | None:
         key = self._namespaced_key(key)
         value = await self.redis.get(key)
         if value is None:
             return None
-        ttl = await self.redis.ttl(key)
+        await self.redis.ttl(key)
         return value
 
     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[str]:
+    async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
         start_key = self._namespaced_key(start_key)
         end_key = self._namespaced_key(end_key)
+        cursor = 0
+        pattern = start_key + "*"  # Match all keys starting with start_key prefix
+        matching_keys = []
+        while True:
+            cursor, keys = await self.redis.scan(cursor, match=pattern, count=1000)
 
-        return await self.redis.zrangebylex(start_key, end_key)
+            for key in keys:
+                key_str = key.decode("utf-8") if isinstance(key, bytes) else key
+                if start_key <= key_str <= end_key:
+                    matching_keys.append(key)
+
+            if cursor == 0:
+                break
+
+        # Then fetch all values in a single MGET call
+        if matching_keys:
+            values = await self.redis.mget(matching_keys)
+            return [
+                value.decode("utf-8") if isinstance(value, bytes) else value for value in values if value is not None
+            ]
+
+        return []
+
+    async def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
+        """Get all keys in the given range."""
+        matching_keys = await self.redis.zrangebylex(self.namespace, f"[{start_key}", f"[{end_key}")
+        if not matching_keys:
+            return []
+        return [k.decode("utf-8") for k in matching_keys]

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

@@ -4,9 +4,10 @@
 # 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
 
+from llama_stack.schema_utils import json_schema_type
+
 
 @json_schema_type
 class SqliteControlPlaneConfig(BaseModel):

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

@@ -5,25 +5,43 @@
 # the root directory of this source tree.
 
 import os
-
 from datetime import datetime
-from typing import List, Optional
 
 import aiosqlite
 
-from ..api import *  # noqa: F403
+from llama_stack.log import get_logger
+
+from ..api import KVStore
 from ..config import SqliteKVStoreConfig
 
+logger = get_logger(name=__name__, category="providers::utils")
+
 
 class SqliteKVStoreImpl(KVStore):
     def __init__(self, config: SqliteKVStoreConfig):
         self.db_path = config.db_path
         self.table_name = "kvstore"
+        self._conn: aiosqlite.Connection | None = 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(
+        # 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)
+
+        # Only use persistent connection for in-memory databases
+        # File-based databases use connection-per-operation to avoid hangs
+        if self._is_memory_db():
+            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,
@@ -32,37 +50,88 @@ class SqliteKVStoreImpl(KVStore):
                 )
             """
             )
-            await db.commit()
+            await self._conn.commit()
+        else:
+            # For file-based databases, just create the table
+            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: str, expiration: Optional[datetime] = None
-    ) -> None:
-        async with aiosqlite.connect(self.db_path) as db:
-            await db.execute(
+    async def shutdown(self):
+        """Close the persistent connection (only for in-memory databases)."""
+        if self._conn:
+            await self._conn.close()
+            self._conn = None
+
+    async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
+        if self._conn:
+            # In-memory database with persistent connection
+            await self._conn.execute(
                 f"INSERT OR REPLACE INTO {self.table_name} (key, value, expiration) VALUES (?, ?, ?)",
                 (key, value, expiration),
             )
-            await db.commit()
+            await self._conn.commit()
+        else:
+            # File-based database with connection per operation
+            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()
 
-    async def get(self, key: str) -> Optional[str]:
-        async with aiosqlite.connect(self.db_path) as db:
-            async with db.execute(
+    async def get(self, key: str) -> str | None:
+        if self._conn:
+            # In-memory database with persistent connection
+            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
+        else:
+            # File-based database with connection per operation
+            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 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()
+        if self._conn:
+            # In-memory database with persistent connection
+            await self._conn.execute(f"DELETE FROM {self.table_name} WHERE key = ?", (key,))
+            await self._conn.commit()
+        else:
+            # File-based database with connection per operation
+            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[str]:
-        async with aiosqlite.connect(self.db_path) as db:
-            async with db.execute(
+    async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
+        if self._conn:
+            # In-memory database with persistent connection
+            async with self._conn.execute(
                 f"SELECT key, value, expiration FROM {self.table_name} WHERE key >= ? AND key <= ?",
                 (start_key, end_key),
             ) as cursor:
@@ -71,3 +140,35 @@ class SqliteKVStoreImpl(KVStore):
                     _, value, _ = row
                     result.append(value)
                 return result
+        else:
+            # File-based database with connection per operation
+            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 def keys_in_range(self, start_key: str, end_key: str) -> list[str]:
+        """Get all keys in the given range."""
+        if self._conn:
+            # In-memory database with persistent connection
+            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]
+        else:
+            # File-based database with connection per operation
+            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]

llama_stack/providers/utils/memory/file_utils.py

@@ -8,7 +8,7 @@ import base64
 import mimetypes
 import os
 
-from llama_models.llama3.api.datatypes import URL
+from llama_stack.apis.common.content_types import URL
 
 
 def data_url_from_file(file_path: str) -> URL: