Merge branch 'main' into enable-vector-stores-files-api-tests

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

@@ -12,6 +12,7 @@ from typing import Any
 from numpy.typing import NDArray
 from pymilvus import DataType, Function, FunctionType, MilvusClient
 
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
 from llama_stack.apis.files.files import Files
 from llama_stack.apis.inference import Inference, InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
@@ -321,11 +322,11 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
             return self.cache[vector_db_id]
 
         if self.vector_db_store is None:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         vector_db = await self.vector_db_store.get_vector_db(vector_db_id)
         if not vector_db:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         index = VectorDBWithIndex(
             vector_db=vector_db,
@@ -348,7 +349,7 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
     ) -> None:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         await index.insert_chunks(chunks)
 
@@ -360,7 +361,7 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
     ) -> QueryChunksResponse:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         if params and params.get("mode") == "keyword":
             # Check if this is inline Milvus (Milvus-Lite)
@@ -376,7 +377,7 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP
         """Delete a chunk from a milvus vector store."""
         index = await self._get_and_cache_vector_db_index(store_id)
         if not index:
-            raise ValueError(f"Vector DB {store_id} not found")
+            raise VectorStoreNotFoundError(store_id)
 
         for chunk_id in chunk_ids:
             # Use the index's delete_chunk method

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

@@ -13,6 +13,7 @@ from psycopg2 import sql
 from psycopg2.extras import Json, execute_values
 from pydantic import BaseModel, TypeAdapter
 
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
 from llama_stack.apis.files.files import Files
 from llama_stack.apis.inference import InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
@@ -275,7 +276,7 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco
         """Delete a chunk from a PostgreSQL vector store."""
         index = await self._get_and_cache_vector_db_index(store_id)
         if not index:
-            raise ValueError(f"Vector DB {store_id} not found")
+            raise VectorStoreNotFoundError(store_id)
 
         for chunk_id in chunk_ids:
             # Use the index's delete_chunk method

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

@@ -12,6 +12,7 @@ from numpy.typing import NDArray
 from qdrant_client import AsyncQdrantClient, models
 from qdrant_client.models import PointStruct
 
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
 from llama_stack.apis.inference import InterleavedContent
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import (
@@ -173,7 +174,7 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
 
         vector_db = await self.vector_db_store.get_vector_db(vector_db_id)
         if not vector_db:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         index = VectorDBWithIndex(
             vector_db=vector_db,
@@ -191,7 +192,7 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
     ) -> None:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         await index.insert_chunks(chunks)
 
@@ -203,7 +204,7 @@ class QdrantVectorIOAdapter(VectorIO, VectorDBsProtocolPrivate):
     ) -> QueryChunksResponse:
         index = await self._get_and_cache_vector_db_index(vector_db_id)
         if not index:
-            raise ValueError(f"Vector DB {vector_db_id} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         return await index.query_chunks(query, params)
 

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

@@ -14,6 +14,7 @@ from weaviate.classes.init import Auth
 from weaviate.classes.query import Filter
 
 from llama_stack.apis.common.content_types import InterleavedContent
+from llama_stack.apis.common.errors import VectorStoreNotFoundError
 from llama_stack.apis.files.files import Files
 from llama_stack.apis.vector_dbs import VectorDB
 from llama_stack.apis.vector_io import Chunk, QueryChunksResponse, VectorIO
@@ -259,7 +260,7 @@ class WeaviateVectorIOAdapter(
 
         vector_db = await self.vector_db_store.get_vector_db(sanitized_collection_name)
         if not vector_db:
-            raise ValueError(f"Vector DB {sanitized_collection_name} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         client = self._get_client()
         if not client.collections.exists(vector_db.identifier):
@@ -282,7 +283,7 @@ class WeaviateVectorIOAdapter(
         sanitized_collection_name = sanitize_collection_name(vector_db_id, weaviate_format=True)
         index = await self._get_and_cache_vector_db_index(sanitized_collection_name)
         if not index:
-            raise ValueError(f"Vector DB {sanitized_collection_name} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         await index.insert_chunks(chunks)
 
@@ -295,7 +296,7 @@ class WeaviateVectorIOAdapter(
         sanitized_collection_name = sanitize_collection_name(vector_db_id, weaviate_format=True)
         index = await self._get_and_cache_vector_db_index(sanitized_collection_name)
         if not index:
-            raise ValueError(f"Vector DB {sanitized_collection_name} not found")
+            raise VectorStoreNotFoundError(vector_db_id)
 
         return await index.query_chunks(query, params)