llama-stack 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

llama_stack/providers/remote/vector_io/pgvector/pgvector.py

@@ -368,6 +368,22 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
             log.exception("Could not connect to PGVector database server")
             raise RuntimeError("Could not connect to PGVector database server") from e
 
+        # Load existing vector stores from KV store into cache
+        start_key = VECTOR_DBS_PREFIX
+        end_key = f"{VECTOR_DBS_PREFIX}\xff"
+        stored_vector_stores = await self.kvstore.values_in_range(start_key, end_key)
+        for vector_store_data in stored_vector_stores:
+            vector_store = VectorStore.model_validate_json(vector_store_data)
+            pgvector_index = PGVectorIndex(
+                vector_store=vector_store,
+                dimension=vector_store.embedding_dimension,
+                conn=self.conn,
+                kvstore=self.kvstore,
+            )
+            await pgvector_index.initialize()
+            index = VectorStoreWithIndex(vector_store, index=pgvector_index, inference_api=self.inference_api)
+            self.cache[vector_store.identifier] = index
+
     async def shutdown(self) -> None:
         if self.conn is not None:
             self.conn.close()
@@ -377,7 +393,13 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
 
     async def register_vector_store(self, vector_store: VectorStore) -> None:
         # Persist vector DB metadata in the KV store
-        assert self.kvstore is not None
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before registering vector stores.")
+
+        # Save to kvstore for persistence
+        key = f"{VECTOR_DBS_PREFIX}{vector_store.identifier}"
+        await self.kvstore.set(key=key, value=vector_store.model_dump_json())
+
         # Upsert model metadata in Postgres
         upsert_models(self.conn, [(vector_store.identifier, vector_store)])
 
@@ -396,7 +418,8 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
             del self.cache[vector_store_id]
 
         # Delete vector DB metadata from KV store
-        assert self.kvstore is not None
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before unregistering vector stores.")
         await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}")
 
     async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
@@ -413,13 +436,16 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt
         if vector_store_id in self.cache:
             return self.cache[vector_store_id]
 
-        if self.vector_store_table is None:
-            raise VectorStoreNotFoundError(vector_store_id)
+        # Try to load from kvstore
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before using vector stores.")
 
-        vector_store = await self.vector_store_table.get_vector_store(vector_store_id)
-        if not vector_store:
+        key = f"{VECTOR_DBS_PREFIX}{vector_store_id}"
+        vector_store_data = await self.kvstore.get(key)
+        if not vector_store_data:
             raise VectorStoreNotFoundError(vector_store_id)
 
+        vector_store = VectorStore.model_validate_json(vector_store_data)
         index = PGVectorIndex(vector_store, vector_store.embedding_dimension, self.conn)
         await index.initialize()
         self.cache[vector_store_id] = VectorStoreWithIndex(vector_store, index, self.inference_api)

llama_stack/providers/remote/vector_io/qdrant/qdrant.py

@@ -183,7 +183,8 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc
         await super().shutdown()
 
     async def register_vector_store(self, vector_store: VectorStore) -> None:
-        assert self.kvstore is not None
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before registering vector stores.")
         key = f"{VECTOR_DBS_PREFIX}{vector_store.identifier}"
         await self.kvstore.set(key=key, value=vector_store.model_dump_json())
 
@@ -200,20 +201,24 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc
             await self.cache[vector_store_id].index.delete()
             del self.cache[vector_store_id]
 
-        assert self.kvstore is not None
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before using vector stores.")
         await self.kvstore.delete(f"{VECTOR_DBS_PREFIX}{vector_store_id}")
 
     async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex | None:
         if vector_store_id in self.cache:
             return self.cache[vector_store_id]
 
-        if self.vector_store_table is None:
-            raise ValueError(f"Vector DB not found {vector_store_id}")
+        # Try to load from kvstore
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before using vector stores.")
 
-        vector_store = await self.vector_store_table.get_vector_store(vector_store_id)
-        if not vector_store:
+        key = f"{VECTOR_DBS_PREFIX}{vector_store_id}"
+        vector_store_data = await self.kvstore.get(key)
+        if not vector_store_data:
             raise VectorStoreNotFoundError(vector_store_id)
 
+        vector_store = VectorStore.model_validate_json(vector_store_data)
         index = VectorStoreWithIndex(
             vector_store=vector_store,
             index=QdrantIndex(client=self.client, collection_name=vector_store.identifier),

llama_stack/providers/remote/vector_io/weaviate/weaviate.py

@@ -346,13 +346,16 @@ class WeaviateVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, NeedsRequestProv
         if vector_store_id in self.cache:
             return self.cache[vector_store_id]
 
-        if self.vector_store_table is None:
-            raise VectorStoreNotFoundError(vector_store_id)
+        # Try to load from kvstore
+        if self.kvstore is None:
+            raise RuntimeError("KVStore not initialized. Call initialize() before using vector stores.")
 
-        vector_store = await self.vector_store_table.get_vector_store(vector_store_id)
-        if not vector_store:
+        key = f"{VECTOR_DBS_PREFIX}{vector_store_id}"
+        vector_store_data = await self.kvstore.get(key)
+        if not vector_store_data:
             raise VectorStoreNotFoundError(vector_store_id)
 
+        vector_store = VectorStore.model_validate_json(vector_store_data)
         client = self._get_client()
         sanitized_collection_name = sanitize_collection_name(vector_store.identifier, weaviate_format=True)
         if not client.collections.exists(sanitized_collection_name):

llama_stack/providers/utils/inference/embedding_mixin.py

@@ -46,8 +46,7 @@ class SentenceTransformerEmbeddingMixin:
             raise ValueError("Empty list not supported")
 
         # Get the model and generate embeddings
-        model_obj = await self.model_store.get_model(params.model)
-        embedding_model = await self._load_sentence_transformer_model(model_obj.provider_resource_id)
+        embedding_model = await self._load_sentence_transformer_model(params.model)
         embeddings = await asyncio.to_thread(embedding_model.encode, input_list, show_progress_bar=False)
 
         # Convert embeddings to the requested format

llama_stack/providers/utils/inference/inference_store.py

@@ -35,6 +35,7 @@ class InferenceStore:
         self.reference = reference
         self.sql_store = None
         self.policy = policy
+        self.enable_write_queue = True
 
         # Async write queue and worker control
         self._queue: asyncio.Queue[tuple[OpenAIChatCompletion, list[OpenAIMessageParam]]] | None = None
@@ -47,14 +48,13 @@ class InferenceStore:
         base_store = sqlstore_impl(self.reference)
         self.sql_store = AuthorizedSqlStore(base_store, self.policy)
 
-        # Disable write queue for SQLite to avoid concurrency issues
-        backend_name = self.reference.backend
-        backend_config = _SQLSTORE_BACKENDS.get(backend_name)
-        if backend_config is None:
-            raise ValueError(
-                f"Unregistered SQL backend '{backend_name}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
-            )
-        self.enable_write_queue = backend_config.type != StorageBackendType.SQL_SQLITE
+        # Disable write queue for SQLite since WAL mode handles concurrency
+        # Keep it enabled for other backends (like Postgres) for performance
+        backend_config = _SQLSTORE_BACKENDS.get(self.reference.backend)
+        if backend_config and backend_config.type == StorageBackendType.SQL_SQLITE:
+            self.enable_write_queue = False
+            logger.debug("Write queue disabled for SQLite (WAL mode handles concurrency)")
+
         await self.sql_store.create_table(
             "chat_completions",
             {
@@ -70,8 +70,9 @@ class InferenceStore:
             self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
             for _ in range(self._num_writers):
                 self._worker_tasks.append(asyncio.create_task(self._worker_loop()))
-        else:
-            logger.info("Write queue disabled for SQLite to avoid concurrency issues")
+            logger.debug(
+                f"Inference store write queue enabled with {self._num_writers} writers, max queue size {self._max_write_queue_size}"
+            )
 
     async def shutdown(self) -> None:
         if not self._worker_tasks:
@@ -93,10 +94,29 @@ class InferenceStore:
         if self.enable_write_queue and self._queue is not None:
             await self._queue.join()
 
+    async def _ensure_workers_started(self) -> None:
+        """Ensure the async write queue workers run on the current loop."""
+        if not self.enable_write_queue:
+            return
+
+        if self._queue is None:
+            self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
+            logger.debug(
+                f"Inference store write queue created with max size {self._max_write_queue_size} "
+                f"and {self._num_writers} writers"
+            )
+
+        if not self._worker_tasks:
+            loop = asyncio.get_running_loop()
+            for _ in range(self._num_writers):
+                task = loop.create_task(self._worker_loop())
+                self._worker_tasks.append(task)
+
     async def store_chat_completion(
         self, chat_completion: OpenAIChatCompletion, input_messages: list[OpenAIMessageParam]
     ) -> None:
         if self.enable_write_queue:
+            await self._ensure_workers_started()
             if self._queue is None:
                 raise ValueError("Inference store is not initialized")
             try:

llama_stack/providers/utils/inference/model_registry.py

@@ -20,7 +20,7 @@ logger = get_logger(name=__name__, category="providers::utils")
 
 
 class RemoteInferenceProviderConfig(BaseModel):
-    allowed_models: list[str] | None = Field(  # TODO: make this non-optional and give a list() default
+    allowed_models: list[str] | None = Field(
         default=None,
         description="List of models that should be registered with the model registry. If None, all models are allowed.",
     )

llama_stack/providers/utils/inference/openai_mixin.py

@@ -82,9 +82,6 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
     # This is set in list_models() and used in check_model_availability()
     _model_cache: dict[str, Model] = {}
 
-    # List of allowed models for this provider, if empty all models allowed
-    allowed_models: list[str] = []
-
     # Optional field name in provider data to look for API key, which takes precedence
     provider_data_api_key_field: str | None = None
 
@@ -191,6 +188,19 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
 
         return api_key
 
+    def _validate_model_allowed(self, provider_model_id: str) -> None:
+        """
+        Validate that the model is in the allowed_models list if configured.
+
+        :param provider_model_id: The provider-specific model ID to validate
+        :raises ValueError: If the model is not in the allowed_models list
+        """
+        if self.config.allowed_models is not None and provider_model_id not in self.config.allowed_models:
+            raise ValueError(
+                f"Model '{provider_model_id}' is not in the allowed models list. "
+                f"Allowed models: {self.config.allowed_models}"
+            )
+
     async def _get_provider_model_id(self, model: str) -> str:
         """
         Get the provider-specific model ID from the model store.
@@ -201,8 +211,11 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         :param model: The registered model name/identifier
         :return: The provider-specific model ID (e.g., "gpt-4")
         """
-        # Look up the registered model to get the provider-specific model ID
         # self.model_store is injected by the distribution system at runtime
+        if not await self.model_store.has_model(model):  # type: ignore[attr-defined]
+            return model
+
+        # Look up the registered model to get the provider-specific model ID
         model_obj: Model = await self.model_store.get_model(model)  # type: ignore[attr-defined]
         # provider_resource_id is str | None, but we expect it to be str for OpenAI calls
         if model_obj.provider_resource_id is None:
@@ -234,8 +247,11 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         Direct OpenAI completion API call.
         """
         # TODO: fix openai_completion to return type compatible with OpenAI's API response
+        provider_model_id = await self._get_provider_model_id(params.model)
+        self._validate_model_allowed(provider_model_id)
+
         completion_kwargs = await prepare_openai_completion_params(
-            model=await self._get_provider_model_id(params.model),
+            model=provider_model_id,
             prompt=params.prompt,
             best_of=params.best_of,
             echo=params.echo,
@@ -267,6 +283,9 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         """
         Direct OpenAI chat completion API call.
         """
+        provider_model_id = await self._get_provider_model_id(params.model)
+        self._validate_model_allowed(provider_model_id)
+
         messages = params.messages
 
         if self.download_images:
@@ -288,7 +307,7 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
             messages = [await _localize_image_url(m) for m in messages]
 
         request_params = await prepare_openai_completion_params(
-            model=await self._get_provider_model_id(params.model),
+            model=provider_model_id,
             messages=messages,
             frequency_penalty=params.frequency_penalty,
             function_call=params.function_call,
@@ -326,9 +345,13 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         """
         Direct OpenAI embeddings API call.
         """
-        # Prepare request parameters
-        request_params = {
-            "model": await self._get_provider_model_id(params.model),
+        provider_model_id = await self._get_provider_model_id(params.model)
+        self._validate_model_allowed(provider_model_id)
+
+        # Build request params conditionally to avoid NotGiven/Omit type mismatch
+        # The OpenAI SDK uses Omit in signatures but NOT_GIVEN has type NotGiven
+        request_params: dict[str, Any] = {
+            "model": provider_model_id,
             "input": params.input,
             "encoding_format": params.encoding_format if params.encoding_format is not None else NOT_GIVEN,
             "dimensions": params.dimensions if params.dimensions is not None else NOT_GIVEN,
@@ -413,7 +436,7 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
         for provider_model_id in provider_models_ids:
             if not isinstance(provider_model_id, str):
                 raise ValueError(f"Model ID {provider_model_id} from list_provider_model_ids() is not a string")
-            if self.allowed_models and provider_model_id not in self.allowed_models:
+            if self.config.allowed_models is not None and provider_model_id not in self.config.allowed_models:
                 logger.info(f"Skipping model {provider_model_id} as it is not in the allowed models list")
                 continue
             if metadata := self.embedding_model_metadata.get(provider_model_id):

llama_stack/providers/utils/responses/responses_store.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 import asyncio
-from typing import Any
 
 from llama_stack.apis.agents import (
     Order,
@@ -55,28 +54,19 @@ class ResponsesStore:
 
         self.policy = policy
         self.sql_store = None
-        self.enable_write_queue = True
-
-        # Async write queue and worker control
-        self._queue: (
-            asyncio.Queue[tuple[OpenAIResponseObject, list[OpenAIResponseInput], list[OpenAIMessageParam]]] | None
-        ) = None
-        self._worker_tasks: list[asyncio.Task[Any]] = []
-        self._max_write_queue_size: int = self.reference.max_write_queue_size
-        self._num_writers: int = max(1, self.reference.num_writers)
 
     async def initialize(self):
         """Create the necessary tables if they don't exist."""
         base_store = sqlstore_impl(self.reference)
         self.sql_store = AuthorizedSqlStore(base_store, self.policy)
 
+        # Disable write queue for SQLite since WAL mode handles concurrency
+        # Keep it enabled for other backends (like Postgres) for performance
         backend_config = _SQLSTORE_BACKENDS.get(self.reference.backend)
-        if backend_config is None:
-            raise ValueError(
-                f"Unregistered SQL backend '{self.reference.backend}'. Registered backends: {sorted(_SQLSTORE_BACKENDS)}"
-            )
-        if backend_config.type == StorageBackendType.SQL_SQLITE:
+        if backend_config and backend_config.type == StorageBackendType.SQL_SQLITE:
             self.enable_write_queue = False
+            logger.debug("Write queue disabled for SQLite (WAL mode handles concurrency)")
+
         await self.sql_store.create_table(
             "openai_responses",
             {
@@ -99,28 +89,16 @@ class ResponsesStore:
             self._queue = asyncio.Queue(maxsize=self._max_write_queue_size)
             for _ in range(self._num_writers):
                 self._worker_tasks.append(asyncio.create_task(self._worker_loop()))
-        else:
-            logger.debug("Write queue disabled for SQLite to avoid concurrency issues")
+            logger.debug(
+                f"Responses store write queue enabled with {self._num_writers} writers, max queue size {self._max_write_queue_size}"
+            )
 
     async def shutdown(self) -> None:
-        if not self._worker_tasks:
-            return
-        if self._queue is not None:
-            await self._queue.join()
-        for t in self._worker_tasks:
-            if not t.done():
-                t.cancel()
-        for t in self._worker_tasks:
-            try:
-                await t
-            except asyncio.CancelledError:
-                pass
-        self._worker_tasks.clear()
+        return
 
     async def flush(self) -> None:
-        """Wait for all queued writes to complete. Useful for testing."""
-        if self.enable_write_queue and self._queue is not None:
-            await self._queue.join()
+        """Maintained for compatibility; no-op now that writes are synchronous."""
+        return
 
     async def store_response_object(
         self,
@@ -128,31 +106,7 @@ class ResponsesStore:
         input: list[OpenAIResponseInput],
         messages: list[OpenAIMessageParam],
     ) -> None:
-        if self.enable_write_queue:
-            if self._queue is None:
-                raise ValueError("Responses store is not initialized")
-            try:
-                self._queue.put_nowait((response_object, input, messages))
-            except asyncio.QueueFull:
-                logger.warning(f"Write queue full; adding response id={getattr(response_object, 'id', '<unknown>')}")
-                await self._queue.put((response_object, input, messages))
-        else:
-            await self._write_response_object(response_object, input, messages)
-
-    async def _worker_loop(self) -> None:
-        assert self._queue is not None
-        while True:
-            try:
-                item = await self._queue.get()
-            except asyncio.CancelledError:
-                break
-            response_object, input, messages = item
-            try:
-                await self._write_response_object(response_object, input, messages)
-            except Exception as e:  # noqa: BLE001
-                logger.error(f"Error writing response object: {e}")
-            finally:
-                self._queue.task_done()
+        await self._write_response_object(response_object, input, messages)
 
     async def _write_response_object(
         self,

llama_stack/providers/utils/sqlstore/authorized_sqlstore.py

@@ -45,8 +45,13 @@ def _enhance_item_with_access_control(item: Mapping[str, Any], current_user: Use
         enhanced["owner_principal"] = current_user.principal
         enhanced["access_attributes"] = current_user.attributes
     else:
-        enhanced["owner_principal"] = None
-        enhanced["access_attributes"] = None
+        # IMPORTANT: Use empty string and null value (not None) to match public access filter
+        # The public access filter in _get_public_access_conditions() expects:
+        # - owner_principal = '' (empty string)
+        # - access_attributes = null (JSON null, which serializes to the string 'null')
+        # Setting them to None (SQL NULL) will cause rows to be filtered out on read.
+        enhanced["owner_principal"] = ""
+        enhanced["access_attributes"] = None  # Pydantic/JSON will serialize this as JSON null
     return enhanced
 
 
@@ -188,8 +193,9 @@ class AuthorizedSqlStore:
             enhanced_data["owner_principal"] = current_user.principal
             enhanced_data["access_attributes"] = current_user.attributes
         else:
-            enhanced_data["owner_principal"] = None
-            enhanced_data["access_attributes"] = None
+            # IMPORTANT: Use empty string for owner_principal to match public access filter
+            enhanced_data["owner_principal"] = ""
+            enhanced_data["access_attributes"] = None  # Will serialize as JSON null
 
         await self.sql_store.update(table, enhanced_data, where)
 
@@ -245,14 +251,24 @@ class AuthorizedSqlStore:
             raise ValueError(f"Unsupported database type: {self.database_type}")
 
     def _get_public_access_conditions(self) -> list[str]:
-        """Get the SQL conditions for public access."""
-        # Public records are records that have no owner_principal or access_attributes
+        """Get the SQL conditions for public access.
+
+        Public records are those with:
+        - owner_principal = '' (empty string)
+        - access_attributes is either SQL NULL or JSON null
+
+        Note: Different databases serialize None differently:
+        - SQLite: None → JSON null (text = 'null')
+        - Postgres: None → SQL NULL (IS NULL)
+        """
         conditions = ["owner_principal = ''"]
         if self.database_type == StorageBackendType.SQL_POSTGRES.value:
-            # Postgres stores JSON null as 'null'
-            conditions.append("access_attributes::text = 'null'")
+            # Accept both SQL NULL and JSON null for Postgres compatibility
+            # This handles both old rows (SQL NULL) and new rows (JSON null)
+            conditions.append("(access_attributes IS NULL OR access_attributes::text = 'null')")
         elif self.database_type == StorageBackendType.SQL_SQLITE.value:
-            conditions.append("access_attributes = 'null'")
+            # SQLite serializes None as JSON null
+            conditions.append("(access_attributes IS NULL OR access_attributes = 'null')")
         else:
             raise ValueError(f"Unsupported database type: {self.database_type}")
         return conditions

llama_stack/providers/utils/sqlstore/sqlalchemy_sqlstore.py

@@ -17,6 +17,7 @@ from sqlalchemy import (
     String,
     Table,
     Text,
+    event,
     inspect,
     select,
     text,
@@ -75,7 +76,36 @@ class SqlAlchemySqlStoreImpl(SqlStore):
         self.metadata = MetaData()
 
     def create_engine(self) -> AsyncEngine:
-        return create_async_engine(self.config.engine_str, pool_pre_ping=True)
+        # Configure connection args for better concurrency support
+        connect_args = {}
+        if "sqlite" in self.config.engine_str:
+            # SQLite-specific optimizations for concurrent access
+            # With WAL mode, most locks resolve in milliseconds, but allow up to 5s for edge cases
+            connect_args["timeout"] = 5.0
+            connect_args["check_same_thread"] = False  # Allow usage across asyncio tasks
+
+        engine = create_async_engine(
+            self.config.engine_str,
+            pool_pre_ping=True,
+            connect_args=connect_args,
+        )
+
+        # Enable WAL mode for SQLite to support concurrent readers and writers
+        if "sqlite" in self.config.engine_str:
+
+            @event.listens_for(engine.sync_engine, "connect")
+            def set_sqlite_pragma(dbapi_conn, connection_record):
+                cursor = dbapi_conn.cursor()
+                # Enable Write-Ahead Logging for better concurrency
+                cursor.execute("PRAGMA journal_mode=WAL")
+                # Set busy timeout to 5 seconds (retry instead of immediate failure)
+                # With WAL mode, locks should be brief; if we hit 5s there's a bigger issue
+                cursor.execute("PRAGMA busy_timeout=5000")
+                # Use NORMAL synchronous mode for better performance (still safe with WAL)
+                cursor.execute("PRAGMA synchronous=NORMAL")
+                cursor.close()
+
+        return engine
 
     async def create_table(
         self,