Merge branch 'main' into add-to-provider-gen

llama_stack/apis/tools/rag_tool.py

@@ -4,7 +4,7 @@
 # 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 enum import Enum, StrEnum
 from typing import Annotated, Any, Literal, Protocol
 
 from pydantic import BaseModel, Field, field_validator
@@ -88,7 +88,7 @@ class RAGQueryGenerator(Enum):
 
 
 @json_schema_type
-class RAGSearchMode(Enum):
+class RAGSearchMode(StrEnum):
     """
     Search modes for RAG query retrieval:
     - VECTOR: Uses vector similarity search for semantic matching

llama_stack/distribution/build_container.sh

@@ -336,7 +336,7 @@ $CONTAINER_BINARY build \
   "$BUILD_CONTEXT_DIR"
 
 # clean up tmp/configs
-rm -f "$BUILD_CONTEXT_DIR/run.yaml"
+rm -rf "$BUILD_CONTEXT_DIR/run.yaml" "$TEMP_DIR"
 set +x
 
 echo "Success!"

llama_stack/providers/inline/vector_io/faiss/faiss.py

@@ -260,48 +260,3 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPr
             raise ValueError(f"Vector DB {vector_db_id} not found")
 
         return await index.query_chunks(query, params)
-
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file data to kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=key, value=json.dumps(file_info))
-        content_key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=content_key, value=json.dumps(file_contents))
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        stored_data = await self.kvstore.get(key)
-        return json.loads(stored_data) if stored_data else {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}"
-        stored_data = await self.kvstore.get(key)
-        return json.loads(stored_data) if stored_data else []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in kvstore."""
-        assert self.kvstore is not None
-        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
-        await self.kvstore.set(key=key, value=json.dumps(file_info))
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store data from kvstore."""
-        assert self.kvstore is not None
-
-        keys_to_delete = [
-            f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}",
-            f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}",
-        ]
-        for key in keys_to_delete:
-            try:
-                await self.kvstore.delete(key)
-            except Exception as e:
-                logger.warning(f"Failed to delete key {key}: {e}")
-                continue

llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 import logging
 import re
 import sqlite3
@@ -506,140 +505,6 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc
         await self.cache[vector_db_id].index.delete()
         del self.cache[vector_db_id]
 
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to SQLite database."""
-
-        def _create_or_store():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                # Create a table to persist OpenAI vector store files.
-                cur.execute("""
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files (
-                        store_id TEXT,
-                        file_id TEXT,
-                        metadata TEXT,
-                        PRIMARY KEY (store_id, file_id)
-                    );
-                """)
-                cur.execute("""
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files_contents (
-                        store_id TEXT,
-                        file_id TEXT,
-                        contents TEXT,
-                        PRIMARY KEY (store_id, file_id)
-                    );
-                """)
-                connection.commit()
-                cur.execute(
-                    "INSERT OR REPLACE INTO openai_vector_store_files (store_id, file_id, metadata) VALUES (?, ?, ?)",
-                    (store_id, file_id, json.dumps(file_info)),
-                )
-                cur.execute(
-                    "INSERT OR REPLACE INTO openai_vector_store_files_contents (store_id, file_id, contents) VALUES (?, ?, ?)",
-                    (store_id, file_id, json.dumps(file_contents)),
-                )
-                connection.commit()
-            except Exception as e:
-                logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}")
-                raise
-            finally:
-                cur.close()
-                connection.close()
-
-        try:
-            await asyncio.to_thread(_create_or_store)
-        except Exception as e:
-            logger.error(f"Error saving openai vector store file {store_id} {file_id}: {e}")
-            raise
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from SQLite database."""
-
-        def _load():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "SELECT metadata FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                if row is None:
-                    return None
-                (metadata,) = row
-                return metadata
-            finally:
-                cur.close()
-                connection.close()
-
-        stored_data = await asyncio.to_thread(_load)
-        return json.loads(stored_data) if stored_data else {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from SQLite database."""
-
-        def _load():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                if row is None:
-                    return None
-                (contents,) = row
-                return contents
-            finally:
-                cur.close()
-                connection.close()
-
-        stored_contents = await asyncio.to_thread(_load)
-        return json.loads(stored_contents) if stored_contents else []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in SQLite database."""
-
-        def _update():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "UPDATE openai_vector_store_files SET metadata = ? WHERE store_id = ? AND file_id = ?",
-                    (json.dumps(file_info), store_id, file_id),
-                )
-                connection.commit()
-            finally:
-                cur.close()
-                connection.close()
-
-        await asyncio.to_thread(_update)
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from SQLite database."""
-
-        def _delete():
-            connection = _create_sqlite_connection(self.config.db_path)
-            cur = connection.cursor()
-            try:
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files WHERE store_id = ? AND file_id = ?", (store_id, file_id)
-                )
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files_contents WHERE store_id = ? AND file_id = ?",
-                    (store_id, file_id),
-                )
-                connection.commit()
-            finally:
-                cur.close()
-                connection.close()
-
-        await asyncio.to_thread(_delete)
-
     async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:

llama_stack/providers/remote/vector_io/milvus/milvus.py

@@ -5,7 +5,6 @@
 # the root directory of this source tree.
 
 import asyncio
-import json
 import logging
 import os
 import re
@@ -370,186 +369,3 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
                 )
 
         return await index.query_chunks(query, params)
-
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to Milvus database."""
-        if store_id not in self.openai_vector_stores:
-            store_info = await self._load_openai_vector_stores(store_id)
-            if not store_info:
-                logger.error(f"OpenAI vector store {store_id} not found")
-                raise ValueError(f"No vector store found with id {store_id}")
-
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                file_schema = MilvusClient.create_schema(
-                    auto_id=False,
-                    enable_dynamic_field=True,
-                    description="Metadata for OpenAI vector store files",
-                )
-                file_schema.add_field(
-                    field_name="store_file_id", datatype=DataType.VARCHAR, is_primary=True, max_length=512
-                )
-                file_schema.add_field(field_name="store_id", datatype=DataType.VARCHAR, max_length=512)
-                file_schema.add_field(field_name="file_id", datatype=DataType.VARCHAR, max_length=512)
-                file_schema.add_field(field_name="file_info", datatype=DataType.VARCHAR, max_length=65535)
-
-                await asyncio.to_thread(
-                    self.client.create_collection,
-                    collection_name="openai_vector_store_files",
-                    schema=file_schema,
-                )
-
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files_contents"):
-                content_schema = MilvusClient.create_schema(
-                    auto_id=False,
-                    enable_dynamic_field=True,
-                    description="Contents for OpenAI vector store files",
-                )
-                content_schema.add_field(
-                    field_name="chunk_id", datatype=DataType.VARCHAR, is_primary=True, max_length=1024
-                )
-                content_schema.add_field(field_name="store_file_id", datatype=DataType.VARCHAR, max_length=1024)
-                content_schema.add_field(field_name="store_id", datatype=DataType.VARCHAR, max_length=512)
-                content_schema.add_field(field_name="file_id", datatype=DataType.VARCHAR, max_length=512)
-                content_schema.add_field(field_name="content", datatype=DataType.VARCHAR, max_length=65535)
-
-                await asyncio.to_thread(
-                    self.client.create_collection,
-                    collection_name="openai_vector_store_files_contents",
-                    schema=content_schema,
-                )
-
-            file_data = [
-                {
-                    "store_file_id": f"{store_id}_{file_id}",
-                    "store_id": store_id,
-                    "file_id": file_id,
-                    "file_info": json.dumps(file_info),
-                }
-            ]
-            await asyncio.to_thread(
-                self.client.upsert,
-                collection_name="openai_vector_store_files",
-                data=file_data,
-            )
-
-            # Save file contents
-            contents_data = [
-                {
-                    "chunk_id": content.get("chunk_metadata").get("chunk_id"),
-                    "store_file_id": f"{store_id}_{file_id}",
-                    "store_id": store_id,
-                    "file_id": file_id,
-                    "content": json.dumps(content),
-                }
-                for content in file_contents
-            ]
-            await asyncio.to_thread(
-                self.client.upsert,
-                collection_name="openai_vector_store_files_contents",
-                data=contents_data,
-            )
-
-        except Exception as e:
-            logger.error(f"Error saving openai vector store file {file_id} for store {store_id}: {e}")
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                return {}
-
-            query_filter = f"store_file_id == '{store_id}_{file_id}'"
-            results = await asyncio.to_thread(
-                self.client.query,
-                collection_name="openai_vector_store_files",
-                filter=query_filter,
-                output_fields=["file_info"],
-            )
-
-            if results:
-                try:
-                    return json.loads(results[0]["file_info"])
-                except json.JSONDecodeError as e:
-                    logger.error(f"Failed to decode file_info for store {store_id}, file {file_id}: {e}")
-                    return {}
-            return {}
-        except Exception as e:
-            logger.error(f"Error loading openai vector store file {file_id} for store {store_id}: {e}")
-            return {}
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                return
-
-            file_data = [
-                {
-                    "store_file_id": f"{store_id}_{file_id}",
-                    "store_id": store_id,
-                    "file_id": file_id,
-                    "file_info": json.dumps(file_info),
-                }
-            ]
-            await asyncio.to_thread(
-                self.client.upsert,
-                collection_name="openai_vector_store_files",
-                data=file_data,
-            )
-        except Exception as e:
-            logger.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}")
-            raise
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files_contents"):
-                return []
-
-            query_filter = (
-                f"store_id == '{store_id}' AND file_id == '{file_id}' AND store_file_id == '{store_id}_{file_id}'"
-            )
-            results = await asyncio.to_thread(
-                self.client.query,
-                collection_name="openai_vector_store_files_contents",
-                filter=query_filter,
-                output_fields=["chunk_id", "store_id", "file_id", "content"],
-            )
-
-            contents = []
-            for result in results:
-                try:
-                    content = json.loads(result["content"])
-                    contents.append(content)
-                except json.JSONDecodeError as e:
-                    logger.error(f"Failed to decode content for store {store_id}, file {file_id}: {e}")
-            return contents
-        except Exception as e:
-            logger.error(f"Error loading openai vector store file contents for {file_id} in store {store_id}: {e}")
-            return []
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from Milvus database."""
-        try:
-            if not await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files"):
-                return
-
-            query_filter = f"store_file_id in ['{store_id}_{file_id}']"
-            await asyncio.to_thread(
-                self.client.delete,
-                collection_name="openai_vector_store_files",
-                filter=query_filter,
-            )
-            if await asyncio.to_thread(self.client.has_collection, "openai_vector_store_files_contents"):
-                await asyncio.to_thread(
-                    self.client.delete,
-                    collection_name="openai_vector_store_files_contents",
-                    filter=query_filter,
-                )
-
-        except Exception as e:
-            logger.error(f"Error deleting openai vector store file {file_id} for store {store_id}: {e}")
-            raise

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

@@ -265,125 +265,3 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
         index = PGVectorIndex(vector_db, vector_db.embedding_dimension, self.conn)
         self.cache[vector_db_id] = VectorDBWithIndex(vector_db, index, self.inference_api)
         return self.cache[vector_db_id]
-
-    # OpenAI Vector Stores File operations are not supported in PGVector
-    async def _save_openai_vector_store_file(
-        self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
-    ) -> None:
-        """Save vector store file metadata to Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files (
-                        store_id TEXT,
-                        file_id TEXT,
-                        metadata JSONB,
-                        PRIMARY KEY (store_id, file_id)
-                    )
-                    """
-                )
-                cur.execute(
-                    """
-                    CREATE TABLE IF NOT EXISTS openai_vector_store_files_contents (
-                        store_id TEXT,
-                        file_id TEXT,
-                        contents JSONB,
-                        PRIMARY KEY (store_id, file_id)
-                    )
-                    """
-                )
-                # Insert file metadata
-                files_query = sql.SQL(
-                    """
-                    INSERT INTO openai_vector_store_files (store_id, file_id, metadata)
-                    VALUES %s
-                    ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata
-                    """
-                )
-                files_values = [(store_id, file_id, Json(file_info))]
-                execute_values(cur, files_query, files_values, template="(%s, %s, %s)")
-                # Insert file contents
-                contents_query = sql.SQL(
-                    """
-                    INSERT INTO openai_vector_store_files_contents (store_id, file_id, contents)
-                    VALUES %s
-                    ON CONFLICT (store_id, file_id) DO UPDATE SET contents = EXCLUDED.contents
-                    """
-                )
-                contents_values = [(store_id, file_id, Json(file_contents))]
-                execute_values(cur, contents_query, contents_values, template="(%s, %s, %s)")
-        except Exception as e:
-            log.error(f"Error saving openai vector store file {file_id} for store {store_id}: {e}")
-            raise
-
-    async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
-        """Load vector store file metadata from Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    "SELECT metadata FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                return row[0] if row and row[0] is not None else {}
-        except Exception as e:
-            log.error(f"Error loading openai vector store file {file_id} for store {store_id}: {e}")
-            return {}
-
-    async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
-        """Load vector store file contents from Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    "SELECT contents FROM openai_vector_store_files_contents WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-                row = cur.fetchone()
-                return row[0] if row and row[0] is not None else []
-        except Exception as e:
-            log.error(f"Error loading openai vector store file contents for {file_id} in store {store_id}: {e}")
-            return []
-
-    async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
-        """Update vector store file metadata in Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                query = sql.SQL(
-                    """
-                    INSERT INTO openai_vector_store_files (store_id, file_id, metadata)
-                    VALUES %s
-                    ON CONFLICT (store_id, file_id) DO UPDATE SET metadata = EXCLUDED.metadata
-                    """
-                )
-                values = [(store_id, file_id, Json(file_info))]
-                execute_values(cur, query, values, template="(%s, %s, %s)")
-        except Exception as e:
-            log.error(f"Error updating openai vector store file {file_id} for store {store_id}: {e}")
-            raise
-
-    async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
-        """Delete vector store file metadata from Postgres database."""
-        if self.conn is None:
-            raise RuntimeError("PostgreSQL connection is not initialized")
-        try:
-            with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur:
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-                cur.execute(
-                    "DELETE FROM openai_vector_store_files_contents WHERE store_id = %s AND file_id = %s",
-                    (store_id, file_id),
-                )
-        except Exception as e:
-            log.error(f"Error deleting openai vector store file {file_id} for store {store_id}: {e}")
-            raise

llama_stack/providers/utils/memory/openai_vector_store_mixin.py

@@ -66,7 +66,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _save_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
         """Save vector store metadata to persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
         await self.kvstore.set(key=key, value=json.dumps(store_info))
         # update in-memory cache
@@ -74,7 +74,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _load_openai_vector_stores(self) -> dict[str, dict[str, Any]]:
         """Load all vector store metadata from persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         start_key = OPENAI_VECTOR_STORES_PREFIX
         end_key = f"{OPENAI_VECTOR_STORES_PREFIX}\xff"
         stored_data = await self.kvstore.values_in_range(start_key, end_key)
@@ -87,7 +87,7 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _update_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None:
         """Update vector store metadata in persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
         await self.kvstore.set(key=key, value=json.dumps(store_info))
         # update in-memory cache
@@ -95,38 +95,62 @@ class OpenAIVectorStoreMixin(ABC):
 
     async def _delete_openai_vector_store_from_storage(self, store_id: str) -> None:
         """Delete vector store metadata from persistent storage."""
-        assert self.kvstore is not None
+        assert self.kvstore
         key = f"{OPENAI_VECTOR_STORES_PREFIX}{store_id}"
         await self.kvstore.delete(key)
         # remove from in-memory cache
         self.openai_vector_stores.pop(store_id, None)
 
-    @abstractmethod
     async def _save_openai_vector_store_file(
         self, store_id: str, file_id: str, file_info: dict[str, Any], file_contents: list[dict[str, Any]]
     ) -> None:
         """Save vector store file metadata to persistent storage."""
-        pass
+        assert self.kvstore
+        meta_key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        await self.kvstore.set(key=meta_key, value=json.dumps(file_info))
+        contents_prefix = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}:"
+        for idx, chunk in enumerate(file_contents):
+            await self.kvstore.set(key=f"{contents_prefix}{idx}", value=json.dumps(chunk))
 
-    @abstractmethod
     async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]:
         """Load vector store file metadata from persistent storage."""
-        pass
+        assert self.kvstore
+        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        stored_data = await self.kvstore.get(key)
+        return json.loads(stored_data) if stored_data else {}
 
-    @abstractmethod
     async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]:
         """Load vector store file contents from persistent storage."""
-        pass
+        assert self.kvstore
+        prefix = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}:"
+        end_key = f"{prefix}\xff"
+        raw_items = await self.kvstore.values_in_range(prefix, end_key)
+        return [json.loads(item) for item in raw_items]
 
-    @abstractmethod
     async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None:
         """Update vector store file metadata in persistent storage."""
-        pass
+        assert self.kvstore
+        key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        await self.kvstore.set(key=key, value=json.dumps(file_info))
 
-    @abstractmethod
     async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None:
         """Delete vector store file metadata from persistent storage."""
-        pass
+        assert self.kvstore
+
+        meta_key = f"{OPENAI_VECTOR_STORES_FILES_PREFIX}{store_id}:{file_id}"
+        await self.kvstore.delete(meta_key)
+
+        contents_prefix = f"{OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX}{store_id}:{file_id}:"
+        end_key = f"{contents_prefix}\xff"
+        # load all stored chunk values (values_in_range is implemented by all backends)
+        raw_items = await self.kvstore.values_in_range(contents_prefix, end_key)
+        # delete each chunk by its index suffix
+        for idx in range(len(raw_items)):
+            await self.kvstore.delete(f"{contents_prefix}{idx}")
+
+    async def initialize_openai_vector_stores(self) -> None:
+        """Load existing OpenAI vector stores into the in-memory cache."""
+        self.openai_vector_stores = await self._load_openai_vector_stores()
 
     @abstractmethod
     async def register_vector_db(self, vector_db: VectorDB) -> None:
@@ -138,10 +162,6 @@ class OpenAIVectorStoreMixin(ABC):
         """Unregister a vector database (provider-specific implementation)."""
         pass
 
-    async def initialize_openai_vector_stores(self) -> None:
-        """Load existing OpenAI vector stores into the in-memory cache."""
-        self.openai_vector_stores = await self._load_openai_vector_stores()
-
     @abstractmethod
     async def insert_chunks(
         self,

llama_stack/templates/dell/__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 .dell import get_distribution_template  # noqa: F401

llama_stack/templates/dell/build.yaml

@@ -0,0 +1,35 @@
+version: 2
+distribution_spec:
+  description: Dell's distribution of Llama Stack. TGI inference via Dell's custom
+    container
+  providers:
+    inference:
+    - remote::tgi
+    - inline::sentence-transformers
+    vector_io:
+    - inline::faiss
+    - remote::chromadb
+    - remote::pgvector
+    safety:
+    - inline::llama-guard
+    agents:
+    - inline::meta-reference
+    telemetry:
+    - inline::meta-reference
+    eval:
+    - inline::meta-reference
+    datasetio:
+    - remote::huggingface
+    - inline::localfs
+    scoring:
+    - inline::basic
+    - inline::llm-as-judge
+    - inline::braintrust
+    tool_runtime:
+    - remote::brave-search
+    - remote::tavily-search
+    - inline::rag-runtime
+image_type: conda
+additional_pip_packages:
+- aiosqlite
+- sqlalchemy[asyncio]

llama_stack/templates/dell/dell.py

@@ -0,0 +1,142 @@
+# 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_stack.apis.models import ModelType
+from llama_stack.distribution.datatypes import (
+    ModelInput,
+    Provider,
+    ShieldInput,
+    ToolGroupInput,
+)
+from llama_stack.providers.inline.inference.sentence_transformers import (
+    SentenceTransformersInferenceConfig,
+)
+from llama_stack.templates.template import DistributionTemplate, RunConfigSettings
+
+
+def get_distribution_template() -> DistributionTemplate:
+    providers = {
+        "inference": ["remote::tgi", "inline::sentence-transformers"],
+        "vector_io": ["inline::faiss", "remote::chromadb", "remote::pgvector"],
+        "safety": ["inline::llama-guard"],
+        "agents": ["inline::meta-reference"],
+        "telemetry": ["inline::meta-reference"],
+        "eval": ["inline::meta-reference"],
+        "datasetio": ["remote::huggingface", "inline::localfs"],
+        "scoring": ["inline::basic", "inline::llm-as-judge", "inline::braintrust"],
+        "tool_runtime": [
+            "remote::brave-search",
+            "remote::tavily-search",
+            "inline::rag-runtime",
+        ],
+    }
+    name = "dell"
+    inference_provider = Provider(
+        provider_id="tgi0",
+        provider_type="remote::tgi",
+        config={
+            "url": "${env.DEH_URL}",
+        },
+    )
+    safety_inference_provider = Provider(
+        provider_id="tgi1",
+        provider_type="remote::tgi",
+        config={
+            "url": "${env.DEH_SAFETY_URL}",
+        },
+    )
+    embedding_provider = Provider(
+        provider_id="sentence-transformers",
+        provider_type="inline::sentence-transformers",
+        config=SentenceTransformersInferenceConfig.sample_run_config(),
+    )
+    chromadb_provider = Provider(
+        provider_id="chromadb",
+        provider_type="remote::chromadb",
+        config={
+            "url": "${env.CHROMA_URL}",
+        },
+    )
+
+    inference_model = ModelInput(
+        model_id="${env.INFERENCE_MODEL}",
+        provider_id="tgi0",
+    )
+    safety_model = ModelInput(
+        model_id="${env.SAFETY_MODEL}",
+        provider_id="tgi1",
+    )
+    embedding_model = ModelInput(
+        model_id="all-MiniLM-L6-v2",
+        provider_id="sentence-transformers",
+        model_type=ModelType.embedding,
+        metadata={
+            "embedding_dimension": 384,
+        },
+    )
+    default_tool_groups = [
+        ToolGroupInput(
+            toolgroup_id="builtin::websearch",
+            provider_id="brave-search",
+        ),
+        ToolGroupInput(
+            toolgroup_id="builtin::rag",
+            provider_id="rag-runtime",
+        ),
+    ]
+
+    return DistributionTemplate(
+        name=name,
+        distro_type="self_hosted",
+        description="Dell's distribution of Llama Stack. TGI inference via Dell's custom container",
+        container_image=None,
+        providers=providers,
+        run_configs={
+            "run.yaml": RunConfigSettings(
+                provider_overrides={
+                    "inference": [inference_provider, embedding_provider],
+                    "vector_io": [chromadb_provider],
+                },
+                default_models=[inference_model, embedding_model],
+                default_tool_groups=default_tool_groups,
+            ),
+            "run-with-safety.yaml": RunConfigSettings(
+                provider_overrides={
+                    "inference": [
+                        inference_provider,
+                        safety_inference_provider,
+                        embedding_provider,
+                    ],
+                    "vector_io": [chromadb_provider],
+                },
+                default_models=[inference_model, safety_model, embedding_model],
+                default_shields=[ShieldInput(shield_id="${env.SAFETY_MODEL}")],
+                default_tool_groups=default_tool_groups,
+            ),
+        },
+        run_config_env_vars={
+            "DEH_URL": (
+                "http://0.0.0.0:8181",
+                "URL for the Dell inference server",
+            ),
+            "DEH_SAFETY_URL": (
+                "http://0.0.0.0:8282",
+                "URL for the Dell safety inference server",
+            ),
+            "CHROMA_URL": (
+                "http://localhost:6601",
+                "URL for the Chroma server",
+            ),
+            "INFERENCE_MODEL": (
+                "meta-llama/Llama-3.2-3B-Instruct",
+                "Inference model loaded into the TGI server",
+            ),
+            "SAFETY_MODEL": (
+                "meta-llama/Llama-Guard-3-1B",
+                "Name of the safety (Llama-Guard) model to use",
+            ),
+        },
+    )

llama_stack/templates/dell/doc_template.md

@@ -0,0 +1,178 @@
+---
+orphan: true
+---
+
+# Dell Distribution of Llama Stack
+
+```{toctree}
+:maxdepth: 2
+:hidden:
+
+self
+```
+
+The `llamastack/distribution-{{ name }}` distribution consists of the following provider configurations.
+
+{{ providers_table }}
+
+You can use this distribution if you have GPUs and want to run an independent TGI or Dell Enterprise Hub container for running inference.
+
+{% if run_config_env_vars %}
+### Environment Variables
+
+The following environment variables can be configured:
+
+{% for var, (default_value, description) in run_config_env_vars.items() %}
+- `{{ var }}`: {{ description }} (default: `{{ default_value }}`)
+{% endfor %}
+{% endif %}
+
+
+## Setting up Inference server using Dell Enterprise Hub's custom TGI container.
+
+NOTE: This is a placeholder to run inference with TGI. This will be updated to use [Dell Enterprise Hub's containers](https://dell.huggingface.co/authenticated/models) once verified.
+
+```bash
+export INFERENCE_PORT=8181
+export DEH_URL=http://0.0.0.0:$INFERENCE_PORT
+export INFERENCE_MODEL=meta-llama/Llama-3.1-8B-Instruct
+export CHROMADB_HOST=localhost
+export CHROMADB_PORT=6601
+export CHROMA_URL=http://$CHROMADB_HOST:$CHROMADB_PORT
+export CUDA_VISIBLE_DEVICES=0
+export LLAMA_STACK_PORT=8321
+
+docker run --rm -it \
+  --pull always \
+  --network host \
+  -v $HOME/.cache/huggingface:/data \
+  -e HF_TOKEN=$HF_TOKEN \
+  -p $INFERENCE_PORT:$INFERENCE_PORT \
+  --gpus $CUDA_VISIBLE_DEVICES \
+  ghcr.io/huggingface/text-generation-inference \
+  --dtype bfloat16 \
+  --usage-stats off \
+  --sharded false \
+  --cuda-memory-fraction 0.7 \
+  --model-id $INFERENCE_MODEL \
+  --port $INFERENCE_PORT --hostname 0.0.0.0
+```
+
+If you are using Llama Stack Safety / Shield APIs, then you will need to also run another instance of a TGI with a corresponding safety model like `meta-llama/Llama-Guard-3-1B` using a script like:
+
+```bash
+export SAFETY_INFERENCE_PORT=8282
+export DEH_SAFETY_URL=http://0.0.0.0:$SAFETY_INFERENCE_PORT
+export SAFETY_MODEL=meta-llama/Llama-Guard-3-1B
+export CUDA_VISIBLE_DEVICES=1
+
+docker run --rm -it \
+  --pull always \
+  --network host \
+  -v $HOME/.cache/huggingface:/data \
+  -e HF_TOKEN=$HF_TOKEN \
+  -p $SAFETY_INFERENCE_PORT:$SAFETY_INFERENCE_PORT \
+  --gpus $CUDA_VISIBLE_DEVICES \
+  ghcr.io/huggingface/text-generation-inference \
+  --dtype bfloat16 \
+  --usage-stats off \
+  --sharded false \
+  --cuda-memory-fraction 0.7 \
+  --model-id $SAFETY_MODEL \
+  --hostname 0.0.0.0 \
+  --port $SAFETY_INFERENCE_PORT
+```
+
+## Dell distribution relies on ChromaDB for vector database usage
+
+You can start a chroma-db easily using docker.
+```bash
+# This is where the indices are persisted
+mkdir -p $HOME/chromadb
+
+podman run --rm -it \
+  --network host \
+  --name chromadb \
+  -v $HOME/chromadb:/chroma/chroma \
+  -e IS_PERSISTENT=TRUE \
+  chromadb/chroma:latest \
+  --port $CHROMADB_PORT \
+  --host $CHROMADB_HOST
+```
+
+## Running Llama Stack
+
+Now you are ready to run Llama Stack with TGI as the inference provider. You can do this via Conda (build code) or Docker which has a pre-built image.
+
+### Via Docker
+
+This method allows you to get started quickly without having to build the distribution code.
+
+```bash
+docker run -it \
+  --pull always \
+  --network host \
+  -p $LLAMA_STACK_PORT:$LLAMA_STACK_PORT \
+  -v $HOME/.llama:/root/.llama \
+  # NOTE: mount the llama-stack directory if testing local changes else not needed
+  -v /home/hjshah/git/llama-stack:/app/llama-stack-source \
+  # localhost/distribution-dell:dev if building / testing locally
+  llamastack/distribution-{{ name }}\
+  --port $LLAMA_STACK_PORT  \
+  --env INFERENCE_MODEL=$INFERENCE_MODEL \
+  --env DEH_URL=$DEH_URL \
+  --env CHROMA_URL=$CHROMA_URL
+
+```
+
+If you are using Llama Stack Safety / Shield APIs, use:
+
+```bash
+# You need a local checkout of llama-stack to run this, get it using
+# git clone https://github.com/meta-llama/llama-stack.git
+cd /path/to/llama-stack
+
+export SAFETY_INFERENCE_PORT=8282
+export DEH_SAFETY_URL=http://0.0.0.0:$SAFETY_INFERENCE_PORT
+export SAFETY_MODEL=meta-llama/Llama-Guard-3-1B
+
+docker run \
+  -it \
+  --pull always \
+  -p $LLAMA_STACK_PORT:$LLAMA_STACK_PORT \
+  -v $HOME/.llama:/root/.llama \
+  -v ./llama_stack/templates/tgi/run-with-safety.yaml:/root/my-run.yaml \
+  llamastack/distribution-{{ name }} \
+  --config /root/my-run.yaml \
+  --port $LLAMA_STACK_PORT \
+  --env INFERENCE_MODEL=$INFERENCE_MODEL \
+  --env DEH_URL=$DEH_URL \
+  --env SAFETY_MODEL=$SAFETY_MODEL \
+  --env DEH_SAFETY_URL=$DEH_SAFETY_URL \
+  --env CHROMA_URL=$CHROMA_URL
+```
+
+### Via Conda
+
+Make sure you have done `pip install llama-stack` and have the Llama Stack CLI available.
+
+```bash
+llama stack build --template {{ name }} --image-type conda
+llama stack run {{ name }}
+  --port $LLAMA_STACK_PORT \
+  --env INFERENCE_MODEL=$INFERENCE_MODEL \
+  --env DEH_URL=$DEH_URL \
+  --env CHROMA_URL=$CHROMA_URL
+```
+
+If you are using Llama Stack Safety / Shield APIs, use:
+
+```bash
+llama stack run ./run-with-safety.yaml \
+  --port $LLAMA_STACK_PORT \
+  --env INFERENCE_MODEL=$INFERENCE_MODEL \
+  --env DEH_URL=$DEH_URL \
+  --env SAFETY_MODEL=$SAFETY_MODEL \
+  --env DEH_SAFETY_URL=$DEH_SAFETY_URL \
+  --env CHROMA_URL=$CHROMA_URL
+```

llama_stack/templates/dell/run-with-safety.yaml

@@ -0,0 +1,131 @@
+version: 2
+image_name: dell
+apis:
+- agents
+- datasetio
+- eval
+- inference
+- safety
+- scoring
+- telemetry
+- tool_runtime
+- vector_io
+providers:
+  inference:
+  - provider_id: tgi0
+    provider_type: remote::tgi
+    config:
+      url: ${env.DEH_URL}
+  - provider_id: tgi1
+    provider_type: remote::tgi
+    config:
+      url: ${env.DEH_SAFETY_URL}
+  - provider_id: sentence-transformers
+    provider_type: inline::sentence-transformers
+    config: {}
+  vector_io:
+  - provider_id: chromadb
+    provider_type: remote::chromadb
+    config:
+      url: ${env.CHROMA_URL}
+  safety:
+  - provider_id: llama-guard
+    provider_type: inline::llama-guard
+    config:
+      excluded_categories: []
+  agents:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      persistence_store:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/agents_store.db
+      responses_store:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/responses_store.db
+  telemetry:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      service_name: "${env.OTEL_SERVICE_NAME:=\u200B}"
+      sinks: ${env.TELEMETRY_SINKS:=console,sqlite}
+      sqlite_db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/trace_store.db
+      otel_exporter_otlp_endpoint: ${env.OTEL_EXPORTER_OTLP_ENDPOINT:=}
+  eval:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      kvstore:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/meta_reference_eval.db
+  datasetio:
+  - provider_id: huggingface
+    provider_type: remote::huggingface
+    config:
+      kvstore:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/huggingface_datasetio.db
+  - provider_id: localfs
+    provider_type: inline::localfs
+    config:
+      kvstore:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/localfs_datasetio.db
+  scoring:
+  - provider_id: basic
+    provider_type: inline::basic
+    config: {}
+  - provider_id: llm-as-judge
+    provider_type: inline::llm-as-judge
+    config: {}
+  - provider_id: braintrust
+    provider_type: inline::braintrust
+    config:
+      openai_api_key: ${env.OPENAI_API_KEY:=}
+  tool_runtime:
+  - provider_id: brave-search
+    provider_type: remote::brave-search
+    config:
+      api_key: ${env.BRAVE_SEARCH_API_KEY:=}
+      max_results: 3
+  - provider_id: tavily-search
+    provider_type: remote::tavily-search
+    config:
+      api_key: ${env.TAVILY_SEARCH_API_KEY:=}
+      max_results: 3
+  - provider_id: rag-runtime
+    provider_type: inline::rag-runtime
+    config: {}
+metadata_store:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db
+inference_store:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db
+models:
+- metadata: {}
+  model_id: ${env.INFERENCE_MODEL}
+  provider_id: tgi0
+  model_type: llm
+- metadata: {}
+  model_id: ${env.SAFETY_MODEL}
+  provider_id: tgi1
+  model_type: llm
+- metadata:
+    embedding_dimension: 384
+  model_id: all-MiniLM-L6-v2
+  provider_id: sentence-transformers
+  model_type: embedding
+shields:
+- shield_id: ${env.SAFETY_MODEL}
+vector_dbs: []
+datasets: []
+scoring_fns: []
+benchmarks: []
+tool_groups:
+- toolgroup_id: builtin::websearch
+  provider_id: brave-search
+- toolgroup_id: builtin::rag
+  provider_id: rag-runtime
+server:
+  port: 8321

llama_stack/templates/dell/run.yaml

@@ -0,0 +1,122 @@
+version: 2
+image_name: dell
+apis:
+- agents
+- datasetio
+- eval
+- inference
+- safety
+- scoring
+- telemetry
+- tool_runtime
+- vector_io
+providers:
+  inference:
+  - provider_id: tgi0
+    provider_type: remote::tgi
+    config:
+      url: ${env.DEH_URL}
+  - provider_id: sentence-transformers
+    provider_type: inline::sentence-transformers
+    config: {}
+  vector_io:
+  - provider_id: chromadb
+    provider_type: remote::chromadb
+    config:
+      url: ${env.CHROMA_URL}
+  safety:
+  - provider_id: llama-guard
+    provider_type: inline::llama-guard
+    config:
+      excluded_categories: []
+  agents:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      persistence_store:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/agents_store.db
+      responses_store:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/responses_store.db
+  telemetry:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      service_name: "${env.OTEL_SERVICE_NAME:=\u200B}"
+      sinks: ${env.TELEMETRY_SINKS:=console,sqlite}
+      sqlite_db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/trace_store.db
+      otel_exporter_otlp_endpoint: ${env.OTEL_EXPORTER_OTLP_ENDPOINT:=}
+  eval:
+  - provider_id: meta-reference
+    provider_type: inline::meta-reference
+    config:
+      kvstore:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/meta_reference_eval.db
+  datasetio:
+  - provider_id: huggingface
+    provider_type: remote::huggingface
+    config:
+      kvstore:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/huggingface_datasetio.db
+  - provider_id: localfs
+    provider_type: inline::localfs
+    config:
+      kvstore:
+        type: sqlite
+        db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/localfs_datasetio.db
+  scoring:
+  - provider_id: basic
+    provider_type: inline::basic
+    config: {}
+  - provider_id: llm-as-judge
+    provider_type: inline::llm-as-judge
+    config: {}
+  - provider_id: braintrust
+    provider_type: inline::braintrust
+    config:
+      openai_api_key: ${env.OPENAI_API_KEY:=}
+  tool_runtime:
+  - provider_id: brave-search
+    provider_type: remote::brave-search
+    config:
+      api_key: ${env.BRAVE_SEARCH_API_KEY:=}
+      max_results: 3
+  - provider_id: tavily-search
+    provider_type: remote::tavily-search
+    config:
+      api_key: ${env.TAVILY_SEARCH_API_KEY:=}
+      max_results: 3
+  - provider_id: rag-runtime
+    provider_type: inline::rag-runtime
+    config: {}
+metadata_store:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/registry.db
+inference_store:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:=~/.llama/distributions/dell}/inference_store.db
+models:
+- metadata: {}
+  model_id: ${env.INFERENCE_MODEL}
+  provider_id: tgi0
+  model_type: llm
+- metadata:
+    embedding_dimension: 384
+  model_id: all-MiniLM-L6-v2
+  provider_id: sentence-transformers
+  model_type: embedding
+shields: []
+vector_dbs: []
+datasets: []
+scoring_fns: []
+benchmarks: []
+tool_groups:
+- toolgroup_id: builtin::websearch
+  provider_id: brave-search
+- toolgroup_id: builtin::rag
+  provider_id: rag-runtime
+server:
+  port: 8321

tests/unit/rag/test_rag_query.py

@@ -65,7 +65,15 @@ class TestRagQuery:
             RAGQueryConfig(mode="invalid_mode")
 
     async def test_query_accepts_valid_modes(self):
-        RAGQueryConfig()  # Test default (vector)
-        RAGQueryConfig(mode="vector")  # Test vector
-        RAGQueryConfig(mode="keyword")  # Test keyword
-        RAGQueryConfig(mode="hybrid")  # Test hybrid
+        default_config = RAGQueryConfig()  # Test default (vector)
+        assert default_config.mode == "vector"
+        vector_config = RAGQueryConfig(mode="vector")  # Test vector
+        assert vector_config.mode == "vector"
+        keyword_config = RAGQueryConfig(mode="keyword")  # Test keyword
+        assert keyword_config.mode == "keyword"
+        hybrid_config = RAGQueryConfig(mode="hybrid")  # Test hybrid
+        assert hybrid_config.mode == "hybrid"
+
+        # Test that invalid mode raises an error
+        with pytest.raises(ValueError):
+            RAGQueryConfig(mode="wrong_mode")

tests/unit/server/test_access_control.py

@@ -285,9 +285,15 @@ async def test_access_policy(mock_get_authenticated_user, test_setup_with_access
             "projects": ["foo", "bar"],
         },
     )
-    await routing_table.register_model("model-1", provider_id="test_provider")
-    await routing_table.register_model("model-2", provider_id="test_provider")
-    await routing_table.register_model("model-3", provider_id="test_provider")
+    await routing_table.register_model(
+        "model-1", provider_model_id="test_provider/model-1", provider_id="test_provider"
+    )
+    await routing_table.register_model(
+        "model-2", provider_model_id="test_provider/model-2", provider_id="test_provider"
+    )
+    await routing_table.register_model(
+        "model-3", provider_model_id="test_provider/model-3", provider_id="test_provider"
+    )
     model = await routing_table.get_model("model-1")
     assert model.identifier == "model-1"
     model = await routing_table.get_model("model-2")