From 6238833185f5e79c716feab8c65d195f230cd832 Mon Sep 17 00:00:00 2001 From: Derek Higgins Date: Mon, 21 Jul 2025 11:13:21 +0100 Subject: [PATCH] feat(vector-io): implement chunk deletion for PGVector provider Implement delete_chunk() method with SQL DELETE operation Add _delete_openai_chunk_from_vector_store() for OpenAI compatibility Fix chunk_id generation to use actual chunk.chunk_id instead of generated IDs (so they can be removed by id) Set pgvector provider to requiring Api.files as optional Update remote_provider_spec() function to support optional_api_dependencies parameter This allows pgvector to work in configurations without a files provider while still supporting file operations when files API is available. Signed-off-by: Derek Higgins --- llama_stack/providers/datatypes.py | 6 +++++- llama_stack/providers/registry/vector_io.py | 1 + .../remote/vector_io/pgvector/__init__.py | 2 +- .../remote/vector_io/pgvector/pgvector.py | 16 +++++++++++++++- 4 files changed, 22 insertions(+), 3 deletions(-) diff --git a/llama_stack/providers/datatypes.py b/llama_stack/providers/datatypes.py index 424380324..24b3b6547 100644 --- a/llama_stack/providers/datatypes.py +++ b/llama_stack/providers/datatypes.py @@ -226,7 +226,10 @@ API responses, specify the adapter here. def remote_provider_spec( - api: Api, adapter: AdapterSpec, api_dependencies: list[Api] | None = None + api: Api, + adapter: AdapterSpec, + api_dependencies: list[Api] | None = None, + optional_api_dependencies: list[Api] | None = None, ) -> RemoteProviderSpec: return RemoteProviderSpec( api=api, @@ -234,6 +237,7 @@ def remote_provider_spec( config_class=adapter.config_class, adapter=adapter, api_dependencies=api_dependencies or [], + optional_api_dependencies=optional_api_dependencies or [], ) diff --git a/llama_stack/providers/registry/vector_io.py b/llama_stack/providers/registry/vector_io.py index e391341b4..063b382df 100644 --- a/llama_stack/providers/registry/vector_io.py +++ b/llama_stack/providers/registry/vector_io.py @@ -410,6 +410,7 @@ See [PGVector's documentation](https://github.com/pgvector/pgvector) for more de """, ), api_dependencies=[Api.inference], + optional_api_dependencies=[Api.files], ), remote_provider_spec( Api.vector_io, diff --git a/llama_stack/providers/remote/vector_io/pgvector/__init__.py b/llama_stack/providers/remote/vector_io/pgvector/__init__.py index 9f528db74..59eef4c81 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/__init__.py +++ b/llama_stack/providers/remote/vector_io/pgvector/__init__.py @@ -12,6 +12,6 @@ from .config import PGVectorVectorIOConfig async def get_adapter_impl(config: PGVectorVectorIOConfig, deps: dict[Api, ProviderSpec]): from .pgvector import PGVectorVectorIOAdapter - impl = PGVectorVectorIOAdapter(config, deps[Api.inference]) + impl = PGVectorVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files, None)) await impl.initialize() return impl diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index 3aeb3f30d..7f01ffce2 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -99,7 +99,7 @@ class PGVectorIndex(EmbeddingIndex): for i, chunk in enumerate(chunks): values.append( ( - f"{chunk.metadata['document_id']}:chunk-{i}", + f"{chunk.chunk_id}", Json(chunk.model_dump()), embeddings[i].tolist(), ) @@ -159,6 +159,11 @@ class PGVectorIndex(EmbeddingIndex): with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur: cur.execute(f"DROP TABLE IF EXISTS {self.table_name}") + async def delete_chunk(self, chunk_id: str) -> None: + """Remove a chunk from the PostgreSQL table.""" + with self.conn.cursor(cursor_factory=psycopg2.extras.DictCursor) as cur: + cur.execute(f"DELETE FROM {self.table_name} WHERE id = %s", (chunk_id,)) + class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate): def __init__( @@ -265,3 +270,12 @@ 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] + + async def _delete_openai_chunk_from_vector_store(self, store_id: str, chunk_id: str) -> None: + """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") + + # Use the index's delete_chunk method + await index.index.delete_chunk(chunk_id)