From 2aba2c1236f7a060b94ae19c74e621dd9e2d523a Mon Sep 17 00:00:00 2001 From: Francisco Arceo Date: Wed, 23 Jul 2025 13:35:48 -0400 Subject: [PATCH 1/5] chore: Moving vector store and vector store files helper methods to openai_vector_store_mixin (#2863) # What does this PR do? Moving vector store and vector store files helper methods to `openai_vector_store_mixin.py` ## Test Plan The tests are already supported in the CI and tests the inline providers and current integration tests. Note that the `vector_index` fixture will be test `milvus_vec_adapter`, `faiss_vec_adapter`, and `sqlite_vec_adapter` in `tests/unit/providers/vector_io/test_vector_io_openai_vector_stores.py`. Additionally, the integration tests in `integration-vector-io-tests.yml` runs `tests/integration/vector_io` tests for the following providers: ```python vector-io-provider: ["inline::faiss", "inline::sqlite-vec", "inline::milvus", "remote::chromadb", "remote::pgvector"] ``` Signed-off-by: Francisco Javier Arceo --- .../providers/inline/vector_io/faiss/faiss.py | 45 ----- .../inline/vector_io/sqlite_vec/sqlite_vec.py | 135 ------------- .../remote/vector_io/milvus/milvus.py | 184 ------------------ .../remote/vector_io/pgvector/pgvector.py | 122 ------------ .../utils/memory/openai_vector_store_mixin.py | 56 ++++-- 5 files changed, 38 insertions(+), 504 deletions(-) diff --git a/llama_stack/providers/inline/vector_io/faiss/faiss.py b/llama_stack/providers/inline/vector_io/faiss/faiss.py index 2a1370c56..a57b4a4ee 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/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 diff --git a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py b/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py index 060b5b15c..f2598cc7c 100644 --- a/llama_stack/providers/inline/vector_io/sqlite_vec/sqlite_vec.py +++ b/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: diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index f301942cb..dc4852821 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/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 diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index 7fdd8af9b..3aeb3f30d 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/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 diff --git a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py index c11de396b..f178e9299 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/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, From 8353ad498117a1bcc2a40406faa7734bb18f47ea Mon Sep 17 00:00:00 2001 From: Mark Campbell Date: Wed, 23 Jul 2025 19:25:12 +0100 Subject: [PATCH 2/5] fix: search mode validation for rag query (#2857) # What does this PR do? I noticed a few issues with my implementation of the search mode validation for RagQuery. This PR replaces the check for search mode in RagQuery with a Literal. There were issues before with ``` TypeError: Object of type RAGSearchMode is not JSON serializable ``` When using ``` query_config = RAGQueryConfig(max_chunks=6, mode="vector").model_dump() ``` It also fixes the fact that despite user input "vector" was always the used search mode. ## Test Plan Verify that a chosen search mode works when using Rag Query or use below agent config: ``` agent = Agent( client, model=model_id, instructions="You are a helpful assistant", tools=[ { "name": "builtin::rag/knowledge_search", "args": { "vector_db_ids": [vector_db_id], "query_config": { "mode": "keyword", "max_chunks": 6 } }, } ], ) ``` Running Unit Tests: ``` uv sync --extra dev uv run pytest tests/unit/rag/test_rag_query.py -v ``` --- llama_stack/apis/tools/rag_tool.py | 4 ++-- tests/unit/rag/test_rag_query.py | 16 ++++++++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/llama_stack/apis/tools/rag_tool.py b/llama_stack/apis/tools/rag_tool.py index cfaa49488..1d5e7b6cb 100644 --- a/llama_stack/apis/tools/rag_tool.py +++ b/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 diff --git a/tests/unit/rag/test_rag_query.py b/tests/unit/rag/test_rag_query.py index a9149541a..05ccecb99 100644 --- a/tests/unit/rag/test_rag_query.py +++ b/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") From 2fcfb0f0b51fef70df1e3b358aae0762774952c8 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Wed, 23 Jul 2025 11:40:59 -0700 Subject: [PATCH 3/5] fix: bring back dell template (#2880) This template is definitely needed since it (and related docker, which will push soon) is used by folks at Dell. --- llama_stack/templates/dell/__init__.py | 7 + llama_stack/templates/dell/build.yaml | 35 ++++ llama_stack/templates/dell/dell.py | 142 ++++++++++++++ llama_stack/templates/dell/doc_template.md | 178 ++++++++++++++++++ .../templates/dell/run-with-safety.yaml | 131 +++++++++++++ llama_stack/templates/dell/run.yaml | 122 ++++++++++++ 6 files changed, 615 insertions(+) create mode 100644 llama_stack/templates/dell/__init__.py create mode 100644 llama_stack/templates/dell/build.yaml create mode 100644 llama_stack/templates/dell/dell.py create mode 100644 llama_stack/templates/dell/doc_template.md create mode 100644 llama_stack/templates/dell/run-with-safety.yaml create mode 100644 llama_stack/templates/dell/run.yaml diff --git a/llama_stack/templates/dell/__init__.py b/llama_stack/templates/dell/__init__.py new file mode 100644 index 000000000..143add56e --- /dev/null +++ b/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 diff --git a/llama_stack/templates/dell/build.yaml b/llama_stack/templates/dell/build.yaml new file mode 100644 index 000000000..ff8d58a08 --- /dev/null +++ b/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] diff --git a/llama_stack/templates/dell/dell.py b/llama_stack/templates/dell/dell.py new file mode 100644 index 000000000..5a6f52a89 --- /dev/null +++ b/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", + ), + }, + ) diff --git a/llama_stack/templates/dell/doc_template.md b/llama_stack/templates/dell/doc_template.md new file mode 100644 index 000000000..6bdd7f81c --- /dev/null +++ b/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 +``` diff --git a/llama_stack/templates/dell/run-with-safety.yaml b/llama_stack/templates/dell/run-with-safety.yaml new file mode 100644 index 000000000..768fad4fa --- /dev/null +++ b/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 diff --git a/llama_stack/templates/dell/run.yaml b/llama_stack/templates/dell/run.yaml new file mode 100644 index 000000000..de2ada009 --- /dev/null +++ b/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 From 51affe5783f9d55a20db854cca499e897aa3a1e4 Mon Sep 17 00:00:00 2001 From: IAN MILLER <75687988+r3v5@users.noreply.github.com> Date: Wed, 23 Jul 2025 19:50:20 +0100 Subject: [PATCH 4/5] fix: fixed test_access_control.py unit test (#2876) # What does this PR do? I fixed test_access_policy() function providing provider_model_id in each register model endpoint to pass assertions. Initially I faced this issue: ``` tests/unit/server/test_quota.py::test_authenticated_quota_allows_up_to_limit tests/unit/server/test_quota.py::test_authenticated_quota_blocks_after_limit tests/unit/server/test_quota.py::test_anonymous_quota_allows_up_to_limit tests/unit/server/test_quota.py::test_anonymous_quota_blocks_after_limit /Users/iamiller/GitHub/llama-stack/.venv/lib/python3.12/site-packages/aiosqlite/core.py:105: DeprecationWarning: The default datetime adapter is deprecated as of Python 3.12; see the sqlite3 documentation for suggested replacement recipes result = function() -- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html ============================================================================== short test summary info =============================================================================== FAILED tests/unit/server/test_access_control.py::test_access_policy - AssertionError: assert 'test_provider/model-1' == 'model-1' ==================================================================== 1 failed, 436 passed, 194 warnings in 20.09s ==================================================================== ``` After resolved, all works: ``` -- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html ========================================================================= 437 passed, 194 warnings in 19.41s ========================================================================= ``` ## Test Plan Run ` ./scripts/unit-tests.sh` --- tests/unit/server/test_access_control.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/unit/server/test_access_control.py b/tests/unit/server/test_access_control.py index fb9c6f95e..d6a420c13 100644 --- a/tests/unit/server/test_access_control.py +++ b/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") From 427136bb63a5bc066bc58a6cce95ab01a546e05e Mon Sep 17 00:00:00 2001 From: Derek Higgins Date: Wed, 23 Jul 2025 19:54:54 +0100 Subject: [PATCH 5/5] fix: cleanup after build_container.sh (#2869) - rm TEMP_DIR when build_container.sh succeeds - prevents multiple temp directories with Containerfile being left in /tmp Signed-off-by: Derek Higgins --- llama_stack/distribution/build_container.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/llama_stack/distribution/build_container.sh b/llama_stack/distribution/build_container.sh index 74776dd7d..1b1fcc33c 100755 --- a/llama_stack/distribution/build_container.sh +++ b/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!"