Merge branch 'main' into auto_instrument_1

Author: iamemilio

pyproject.toml

@@ -356,6 +356,10 @@ exclude = [
 module = [
     "yaml",
     "fire",
+    "redis.asyncio",
+    "psycopg2",
+    "psycopg2.extras",
+    "psycopg2.extensions",
     "torchtune.*",
     "fairscale.*",
     "torchvision.*",

src/llama_stack/core/conversations/conversations.py

@@ -11,10 +11,9 @@
 from pydantic import BaseModel, TypeAdapter
 
 from llama_stack.core.datatypes import AccessRule, StackRunConfig
+from llama_stack.core.storage.sqlstore.authorized_sqlstore import AuthorizedSqlStore
+from llama_stack.core.storage.sqlstore.sqlstore import sqlstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.sqlstore.api import ColumnDefinition, ColumnType
-from llama_stack.providers.utils.sqlstore.authorized_sqlstore import AuthorizedSqlStore
-from llama_stack.providers.utils.sqlstore.sqlstore import sqlstore_impl
 from llama_stack_api import (
     Conversation,
     ConversationDeletedResource,
@@ -25,6 +24,7 @@
     Conversations,
     Metadata,
 )
+from llama_stack_api.internal.sqlstore import ColumnDefinition, ColumnType
 
 logger = get_logger(name=__name__, category="openai_conversations")
 

src/llama_stack/core/prompts/prompts.py

@@ -10,7 +10,7 @@
 from pydantic import BaseModel
 
 from llama_stack.core.datatypes import StackRunConfig
-from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl
+from llama_stack.core.storage.kvstore import KVStore, kvstore_impl
 from llama_stack_api import ListPromptsResponse, Prompt, Prompts
 
 

src/llama_stack/core/server/quota.py

@@ -11,9 +11,9 @@
 from starlette.types import ASGIApp, Receive, Scope, Send
 
 from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendType
+from llama_stack.core.storage.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore.api import KVStore
-from llama_stack.providers.utils.kvstore.kvstore import _KVSTORE_BACKENDS, kvstore_impl
+from llama_stack_api.internal.kvstore import KVStore
 
 logger = get_logger(name=__name__, category="core::server")
 

src/llama_stack/core/stack.py

@@ -385,8 +385,8 @@ def _initialize_storage(run_config: StackRunConfig):
         else:
             raise ValueError(f"Unknown storage backend type: {type}")
 
-    from llama_stack.providers.utils.kvstore.kvstore import register_kvstore_backends
-    from llama_stack.providers.utils.sqlstore.sqlstore import register_sqlstore_backends
+    from llama_stack.core.storage.kvstore.kvstore import register_kvstore_backends
+    from llama_stack.core.storage.sqlstore.sqlstore import register_sqlstore_backends
 
     register_kvstore_backends(kv_backends)
     register_sqlstore_backends(sql_backends)

src/llama_stack/providers/utils/kvstore/__init__.py renamed to src/llama_stack/core/storage/kvstore/__init__.py

@@ -4,4 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+from llama_stack_api.internal.kvstore import KVStore as KVStore
+
 from .kvstore import *  # noqa: F401, F403

src/llama_stack/providers/utils/kvstore/config.py renamed to src/llama_stack/core/storage/kvstore/config.py

@@ -13,11 +13,19 @@
 
 import asyncio
 from collections import defaultdict
+from datetime import datetime
+from typing import cast
 
-from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig, StorageBackendType
+from llama_stack.core.storage.datatypes import KVStoreReference, StorageBackendConfig
+from llama_stack_api.internal.kvstore import KVStore
 
-from .api import KVStore
-from .config import KVStoreConfig
+from .config import (
+    KVStoreConfig,
+    MongoDBKVStoreConfig,
+    PostgresKVStoreConfig,
+    RedisKVStoreConfig,
+    SqliteKVStoreConfig,
+)
 
 
 def kvstore_dependencies():
@@ -33,15 +41,15 @@ def kvstore_dependencies():
 
 class InmemoryKVStoreImpl(KVStore):
     def __init__(self):
-        self._store = {}
+        self._store: dict[str, str] = {}
 
     async def initialize(self) -> None:
         pass
 
     async def get(self, key: str) -> str | None:
         return self._store.get(key)
 
-    async def set(self, key: str, value: str) -> None:
+    async def set(self, key: str, value: str, expiration: datetime | None = None) -> None:
         self._store[key] = value
 
     async def values_in_range(self, start_key: str, end_key: str) -> list[str]:
@@ -70,7 +78,8 @@ def register_kvstore_backends(backends: dict[str, StorageBackendConfig]) -> None
     _KVSTORE_INSTANCES.clear()
     _KVSTORE_LOCKS.clear()
     for name, cfg in backends.items():
-        _KVSTORE_BACKENDS[name] = cfg
+        typed_cfg = cast(KVStoreConfig, cfg)
+        _KVSTORE_BACKENDS[name] = typed_cfg
 
 
 async def kvstore_impl(reference: KVStoreReference) -> KVStore:
@@ -94,19 +103,20 @@ async def kvstore_impl(reference: KVStoreReference) -> KVStore:
         config = backend_config.model_copy()
         config.namespace = reference.namespace
 
-        if config.type == StorageBackendType.KV_REDIS.value:
+        impl: KVStore
+        if isinstance(config, RedisKVStoreConfig):
             from .redis import RedisKVStoreImpl
 
             impl = RedisKVStoreImpl(config)
-        elif config.type == StorageBackendType.KV_SQLITE.value:
+        elif isinstance(config, SqliteKVStoreConfig):
             from .sqlite import SqliteKVStoreImpl
 
             impl = SqliteKVStoreImpl(config)
-        elif config.type == StorageBackendType.KV_POSTGRES.value:
+        elif isinstance(config, PostgresKVStoreConfig):
             from .postgres import PostgresKVStoreImpl
 
             impl = PostgresKVStoreImpl(config)
-        elif config.type == StorageBackendType.KV_MONGODB.value:
+        elif isinstance(config, MongoDBKVStoreConfig):
             from .mongodb import MongoDBKVStoreImpl
 
             impl = MongoDBKVStoreImpl(config)

src/llama_stack/providers/utils/kvstore/kvstore.py renamed to src/llama_stack/core/storage/kvstore/kvstore.py

@@ -9,8 +9,8 @@
 from pymongo import AsyncMongoClient
 from pymongo.asynchronous.collection import AsyncCollection
 
+from llama_stack.core.storage.kvstore import KVStore
 from llama_stack.log import get_logger
-from llama_stack.providers.utils.kvstore import KVStore
 
 from ..config import MongoDBKVStoreConfig