Merge branch 'meta-llama:main' into main

Author: maximgroshev

docs/source/providers/vector_io/inline_chromadb.md

@@ -42,11 +42,15 @@ See [Chroma's documentation](https://docs.trychroma.com/docs/overview/introducti
 | Field | Type | Required | Default | Description |
 |-------|------|----------|---------|-------------|
 | `db_path` | `<class 'str'>` | No | PydanticUndefined |  |
+| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend |
 
 ## Sample Configuration
 
 ```yaml
 db_path: ${env.CHROMADB_PATH}
+kvstore:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/chroma_inline_registry.db
 
 ```
 

docs/source/providers/vector_io/remote_chromadb.md

@@ -41,11 +41,15 @@ See [Chroma's documentation](https://docs.trychroma.com/docs/overview/introducti
 | Field | Type | Required | Default | Description |
 |-------|------|----------|---------|-------------|
 | `url` | `str \| None` | No | PydanticUndefined |  |
+| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig` | No | sqlite | Config for KV store backend |
 
 ## Sample Configuration
 
 ```yaml
 url: ${env.CHROMADB_URL}
+kvstore:
+  type: sqlite
+  db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/chroma_remote_registry.db
 
 ```
 

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

@@ -136,16 +136,20 @@ EOF
 # Add pip dependencies first since llama-stack is what will change most often
 # so we can reuse layers.
 if [ -n "$pip_dependencies" ]; then
+  read -ra pip_args <<< "$pip_dependencies"
+  quoted_deps=$(printf " %q" "${pip_args[@]}")
   add_to_container << EOF
-RUN $MOUNT_CACHE uv pip install $pip_dependencies
+RUN $MOUNT_CACHE uv pip install $quoted_deps
 EOF
 fi
 
 if [ -n "$special_pip_deps" ]; then
   IFS='#' read -ra parts <<<"$special_pip_deps"
   for part in "${parts[@]}"; do
+    read -ra pip_args <<< "$part"
+    quoted_deps=$(printf " %q" "${pip_args[@]}")
     add_to_container <<EOF
-RUN $MOUNT_CACHE uv pip install $part
+RUN $MOUNT_CACHE uv pip install $quoted_deps
 EOF
   done
 fi
@@ -336,7 +340,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/chroma/__init__.py

@@ -16,6 +16,6 @@ async def get_provider_impl(config: ChromaVectorIOConfig, deps: dict[Api, Any]):
         ChromaVectorIOAdapter,
     )
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference])
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl

llama_stack/providers/inline/vector_io/chroma/config.py

@@ -6,12 +6,25 @@
 
 from typing import Any
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
+from llama_stack.providers.utils.kvstore.config import KVStoreConfig, SqliteKVStoreConfig
+from llama_stack.schema_utils import json_schema_type
 
+
+@json_schema_type
 class ChromaVectorIOConfig(BaseModel):
     db_path: str
+    kvstore: KVStoreConfig = Field(description="Config for KV store backend")
 
     @classmethod
-    def sample_run_config(cls, db_path: str = "${env.CHROMADB_PATH}", **kwargs: Any) -> dict[str, Any]:
-        return {"db_path": db_path}
+    def sample_run_config(
+        cls, __distro_dir__: str, db_path: str = "${env.CHROMADB_PATH}", **kwargs: Any
+    ) -> dict[str, Any]:
+        return {
+            "db_path": db_path,
+            "kvstore": SqliteKVStoreConfig.sample_run_config(
+                __distro_dir__=__distro_dir__,
+                db_name="chroma_inline_registry.db",
+            ),
+        }

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

@@ -260,48 +260,3 @@ async def query_chunks(
             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 @@ async def unregister_vector_db(self, vector_db_id: str) -> None:
         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/chroma/__init__.py

@@ -12,6 +12,6 @@
 async def get_adapter_impl(config: ChromaVectorIOConfig, deps: dict[Api, ProviderSpec]):
     from .chroma import ChromaVectorIOAdapter
 
-    impl = ChromaVectorIOAdapter(config, deps[Api.inference])
+    impl = ChromaVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files))
     await impl.initialize()
     return impl