From 86926ee339216075b9926396d3c2cb51028da2b2 Mon Sep 17 00:00:00 2001 From: Ashwin Bharambe Date: Mon, 20 Oct 2025 19:40:59 -0700 Subject: [PATCH] revert insert_chunks and query_chunks --- .../providers/inline/vector_io/faiss/faiss.py | 12 ++++++------ .../inline/vector_io/sqlite_vec/sqlite_vec.py | 12 ++++++------ .../providers/remote/vector_io/chroma/chroma.py | 12 ++++++------ .../providers/remote/vector_io/milvus/milvus.py | 12 ++++++------ .../remote/vector_io/pgvector/pgvector.py | 8 ++++---- .../providers/remote/vector_io/qdrant/qdrant.py | 12 ++++++------ .../remote/vector_io/weaviate/weaviate.py | 12 ++++++------ .../vector_io/test_openai_vector_stores.py | 12 ++++++------ tests/integration/vector_io/test_vector_io.py | 14 +++++++------- 9 files changed, 53 insertions(+), 53 deletions(-) diff --git a/llama_stack/providers/inline/vector_io/faiss/faiss.py b/llama_stack/providers/inline/vector_io/faiss/faiss.py index cc7614ce8..7b4c146a8 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/llama_stack/providers/inline/vector_io/faiss/faiss.py @@ -249,19 +249,19 @@ class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoco del self.cache[vector_store_id] await self.kvstore.delete(f"{VECTOR_DBS_PREFIX}{vector_store_id}") - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = self.cache.get(vector_store_id) + async def insert_chunks(self, vector_db_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: + index = self.cache.get(vector_db_id) if index is None: - raise ValueError(f"Vector DB {vector_store_id} not found. found: {self.cache.keys()}") + raise ValueError(f"Vector DB {vector_db_id} not found. found: {self.cache.keys()}") await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None + self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = self.cache.get(vector_store_id) + index = self.cache.get(vector_db_id) if index is None: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) return await index.query_chunks(query, params) 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 393bbf9c5..1dcb9647f 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 @@ -448,20 +448,20 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresPro await self.cache[vector_store_id].index.delete() del self.cache[vector_store_id] - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = await self._get_and_cache_vector_store_index(vector_store_id) + 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_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) # The VectorStoreWithIndex helper is expected to compute embeddings via the inference_api # and then call our index's add_chunks. await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: Any, params: dict[str, Any] | None = None + self, vector_db_id: str, query: Any, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_store_index(vector_store_id) + index = await self._get_and_cache_vector_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) return await index.query_chunks(query, params) async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None: diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index a4fd15f77..2663ad43e 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -169,20 +169,20 @@ class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc await self.cache[vector_store_id].index.delete() del self.cache[vector_store_id] - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = await self._get_and_cache_vector_store_index(vector_store_id) + 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_store_index(vector_db_id) if index is None: - raise ValueError(f"Vector DB {vector_store_id} not found in Chroma") + raise ValueError(f"Vector DB {vector_db_id} not found in Chroma") await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None + self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_store_index(vector_store_id) + index = await self._get_and_cache_vector_store_index(vector_db_id) if index is None: - raise ValueError(f"Vector DB {vector_store_id} not found in Chroma") + raise ValueError(f"Vector DB {vector_db_id} not found in Chroma") return await index.query_chunks(query, params) diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index ace9ab1c4..cccf13816 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -348,19 +348,19 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc await self.cache[vector_store_id].index.delete() del self.cache[vector_store_id] - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = await self._get_and_cache_vector_store_index(vector_store_id) + 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_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None + self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_store_index(vector_store_id) + index = await self._get_and_cache_vector_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) return await index.query_chunks(query, params) async def delete_chunks(self, store_id: str, chunks_for_deletion: list[ChunkForDeletion]) -> None: diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index 29cfd673f..f28bd3cd9 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -399,14 +399,14 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProt assert self.kvstore is not None await self.kvstore.delete(key=f"{VECTOR_DBS_PREFIX}{vector_store_id}") - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = await self._get_and_cache_vector_store_index(vector_store_id) + 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_store_index(vector_db_id) await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None + self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_store_index(vector_store_id) + index = await self._get_and_cache_vector_store_index(vector_db_id) return await index.query_chunks(query, params) async def _get_and_cache_vector_store_index(self, vector_store_id: str) -> VectorStoreWithIndex: diff --git a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py index 266e9bf58..93d0894a6 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py @@ -222,19 +222,19 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorStoresProtoc self.cache[vector_store_id] = index return index - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = await self._get_and_cache_vector_store_index(vector_store_id) + 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_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None + self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_store_index(vector_store_id) + index = await self._get_and_cache_vector_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) return await index.query_chunks(query, params) diff --git a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py index 7813f6e5c..66922aa3f 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py +++ b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py @@ -366,19 +366,19 @@ class WeaviateVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, NeedsRequestProv self.cache[vector_store_id] = index return index - async def insert_chunks(self, vector_store_id: str, chunks: list[Chunk], ttl_seconds: int | None = None) -> None: - index = await self._get_and_cache_vector_store_index(vector_store_id) + 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_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) await index.insert_chunks(chunks) async def query_chunks( - self, vector_store_id: str, query: InterleavedContent, params: dict[str, Any] | None = None + self, vector_db_id: str, query: InterleavedContent, params: dict[str, Any] | None = None ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_store_index(vector_store_id) + index = await self._get_and_cache_vector_store_index(vector_db_id) if not index: - raise VectorStoreNotFoundError(vector_store_id) + raise VectorStoreNotFoundError(vector_db_id) return await index.query_chunks(query, params) diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index f2131c614..626faf42d 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -367,7 +367,7 @@ def test_openai_vector_store_with_chunks( # Insert chunks using the native LlamaStack API (since OpenAI API doesn't have direct chunk insertion) llama_client.vector_io.insert( - vector_store_id=vector_store.id, + vector_db_id=vector_store.id, chunks=sample_chunks, ) @@ -434,7 +434,7 @@ def test_openai_vector_store_search_relevance( # Insert chunks using native API llama_client.vector_io.insert( - vector_store_id=vector_store.id, + vector_db_id=vector_store.id, chunks=sample_chunks, ) @@ -484,7 +484,7 @@ def test_openai_vector_store_search_with_ranking_options( # Insert chunks llama_client.vector_io.insert( - vector_store_id=vector_store.id, + vector_db_id=vector_store.id, chunks=sample_chunks, ) @@ -544,7 +544,7 @@ def test_openai_vector_store_search_with_high_score_filter( # Insert chunks llama_client.vector_io.insert( - vector_store_id=vector_store.id, + vector_db_id=vector_store.id, chunks=sample_chunks, ) @@ -610,7 +610,7 @@ def test_openai_vector_store_search_with_max_num_results( # Insert chunks llama_client.vector_io.insert( - vector_store_id=vector_store.id, + vector_db_id=vector_store.id, chunks=sample_chunks, ) @@ -1175,7 +1175,7 @@ def test_openai_vector_store_search_modes( ) client_with_models.vector_io.insert( - vector_store_id=vector_store.id, + vector_db_id=vector_store.id, chunks=sample_chunks, ) query = "Python programming language" diff --git a/tests/integration/vector_io/test_vector_io.py b/tests/integration/vector_io/test_vector_io.py index a312456b9..1f67ddb24 100644 --- a/tests/integration/vector_io/test_vector_io.py +++ b/tests/integration/vector_io/test_vector_io.py @@ -123,12 +123,12 @@ def test_insert_chunks( actual_vector_store_id = create_response.id client_with_empty_registry.vector_io.insert( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, chunks=sample_chunks, ) response = client_with_empty_registry.vector_io.query( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, query="What is the capital of France?", ) assert response is not None @@ -137,7 +137,7 @@ def test_insert_chunks( query, expected_doc_id = test_case response = client_with_empty_registry.vector_io.query( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, query=query, ) assert response is not None @@ -174,13 +174,13 @@ def test_insert_chunks_with_precomputed_embeddings( ] client_with_empty_registry.vector_io.insert( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, chunks=chunks_with_embeddings, ) provider = [p.provider_id for p in client_with_empty_registry.providers.list() if p.api == "vector_io"][0] response = client_with_empty_registry.vector_io.query( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, query="precomputed embedding test", params=vector_io_provider_params_dict.get(provider, None), ) @@ -224,13 +224,13 @@ def test_query_returns_valid_object_when_identical_to_embedding_in_vdb( ] client_with_empty_registry.vector_io.insert( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, chunks=chunks_with_embeddings, ) provider = [p.provider_id for p in client_with_empty_registry.providers.list() if p.api == "vector_io"][0] response = client_with_empty_registry.vector_io.query( - vector_store_id=actual_vector_store_id, + vector_db_id=actual_vector_store_id, query="duplicate", params=vector_io_provider_params_dict.get(provider, None), )