diff --git a/llama_stack/providers/inline/vector_io/faiss/faiss.py b/llama_stack/providers/inline/vector_io/faiss/faiss.py index e4233e158..405c134e5 100644 --- a/llama_stack/providers/inline/vector_io/faiss/faiss.py +++ b/llama_stack/providers/inline/vector_io/faiss/faiss.py @@ -200,15 +200,10 @@ class FaissIndex(EmbeddingIndex): class FaissVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolPrivate): def __init__(self, config: FaissVectorIOConfig, inference_api: Inference, files_api: Files | None) -> None: + super().__init__(files_api=files_api, kvstore=None) self.config = config self.inference_api = inference_api - self.files_api = files_api self.cache: dict[str, VectorDBWithIndex] = {} - self.kvstore: KVStore | None = None - self.openai_vector_stores: dict[str, dict[str, Any]] = {} - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 async def initialize(self) -> None: self.kvstore = await kvstore_impl(self.config.kvstore) 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 dd6fa4e8d..26231a9b7 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 @@ -410,15 +410,10 @@ class SQLiteVecVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoc """ def __init__(self, config, inference_api: Inference, files_api: Files | None) -> None: + super().__init__(files_api=files_api, kvstore=None) self.config = config self.inference_api = inference_api - self.files_api = files_api self.cache: dict[str, VectorDBWithIndex] = {} - self.openai_vector_stores: dict[str, dict[str, Any]] = {} - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 - self.kvstore: KVStore | None = None async def initialize(self) -> None: self.kvstore = await kvstore_impl(self.config.kvstore) diff --git a/llama_stack/providers/remote/vector_io/chroma/chroma.py b/llama_stack/providers/remote/vector_io/chroma/chroma.py index ba7512a6b..511123d6e 100644 --- a/llama_stack/providers/remote/vector_io/chroma/chroma.py +++ b/llama_stack/providers/remote/vector_io/chroma/chroma.py @@ -140,14 +140,13 @@ class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP inference_api: Api.inference, files_api: Files | None, ) -> None: + super().__init__(files_api=files_api, kvstore=None) log.info(f"Initializing ChromaVectorIOAdapter with url: {config}") self.config = config self.inference_api = inference_api self.client = None self.cache = {} - self.kvstore: KVStore | None = None self.vector_db_store = None - self.files_api = files_api async def initialize(self) -> None: self.kvstore = await kvstore_impl(self.config.kvstore) @@ -166,9 +165,6 @@ class ChromaVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP log.info(f"Connecting to Chroma local db at: {self.config.db_path}") self.client = chromadb.PersistentClient(path=self.config.db_path) self.openai_vector_stores = await self._load_openai_vector_stores() - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 async def shutdown(self) -> None: pass diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index c31a61264..0acc90595 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -309,17 +309,12 @@ class MilvusVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP inference_api: Inference, files_api: Files | None, ) -> None: + super().__init__(files_api=files_api, kvstore=None) self.config = config self.cache = {} self.client = None self.inference_api = inference_api - self.files_api = files_api - self.kvstore: KVStore | None = None self.vector_db_store = None - self.openai_vector_stores: dict[str, dict[str, Any]] = {} - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 self.metadata_collection_name = "openai_vector_stores_metadata" async def initialize(self) -> None: diff --git a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py index 63f9c1c3d..dfdfef6eb 100644 --- a/llama_stack/providers/remote/vector_io/pgvector/pgvector.py +++ b/llama_stack/providers/remote/vector_io/pgvector/pgvector.py @@ -4,7 +4,6 @@ # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import asyncio import heapq from typing import Any @@ -346,17 +345,12 @@ class PGVectorVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtoco inference_api: Api.inference, files_api: Files | None = None, ) -> None: + super().__init__(files_api=files_api, kvstore=None) self.config = config self.inference_api = inference_api self.conn = None self.cache = {} - self.files_api = files_api - self.kvstore: KVStore | None = None self.vector_db_store = None - self.openai_vector_stores: dict[str, dict[str, Any]] = {} - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 self.metadata_collection_name = "openai_vector_stores_metadata" async def initialize(self) -> None: diff --git a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py index b59b1fbae..6b386840c 100644 --- a/llama_stack/providers/remote/vector_io/qdrant/qdrant.py +++ b/llama_stack/providers/remote/vector_io/qdrant/qdrant.py @@ -27,7 +27,7 @@ from llama_stack.apis.vector_io import ( from llama_stack.log import get_logger from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate from llama_stack.providers.inline.vector_io.qdrant import QdrantVectorIOConfig as InlineQdrantVectorIOConfig -from llama_stack.providers.utils.kvstore import KVStore, kvstore_impl +from llama_stack.providers.utils.kvstore import kvstore_impl from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin from llama_stack.providers.utils.memory.vector_store import ( ChunkForDeletion, @@ -162,17 +162,12 @@ class QdrantVectorIOAdapter(OpenAIVectorStoreMixin, VectorIO, VectorDBsProtocolP inference_api: Api.inference, files_api: Files | None = None, ) -> None: + super().__init__(files_api=files_api, kvstore=None) self.config = config self.client: AsyncQdrantClient = None self.cache = {} self.inference_api = inference_api - self.files_api = files_api self.vector_db_store = None - self.kvstore: KVStore | None = None - self.openai_vector_stores: dict[str, dict[str, Any]] = {} - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 self._qdrant_lock = asyncio.Lock() async def initialize(self) -> None: diff --git a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py index 6b1730150..54ac6f8d3 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py +++ b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py @@ -3,7 +3,6 @@ # # This source code is licensed under the terms described in the LICENSE file in # the root directory of this source tree. -import asyncio import json from typing import Any @@ -285,17 +284,12 @@ class WeaviateVectorIOAdapter( inference_api: Api.inference, files_api: Files | None, ) -> None: + super().__init__(files_api=files_api, kvstore=None) self.config = config self.inference_api = inference_api self.client_cache = {} self.cache = {} - self.files_api = files_api - self.kvstore: KVStore | None = None self.vector_db_store = None - self.openai_vector_stores: dict[str, dict[str, Any]] = {} - self.openai_file_batches: dict[str, dict[str, Any]] = {} - self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} - self._last_file_batch_cleanup_time = 0 self.metadata_collection_name = "openai_vector_stores_metadata" def _get_client(self) -> weaviate.WeaviateClient: 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 95d790690..44d0e21ca 100644 --- a/llama_stack/providers/utils/memory/openai_vector_store_mixin.py +++ b/llama_stack/providers/utils/memory/openai_vector_store_mixin.py @@ -71,16 +71,15 @@ class OpenAIVectorStoreMixin(ABC): an openai_vector_stores in-memory cache. """ - # These should be provided by the implementing class - openai_vector_stores: dict[str, dict[str, Any]] - openai_file_batches: dict[str, dict[str, Any]] - files_api: Files | None - # KV store for persisting OpenAI vector store metadata - kvstore: KVStore | None - # Track last cleanup time to throttle cleanup operations - _last_file_batch_cleanup_time: int - # Track running file batch processing tasks - _file_batch_tasks: dict[str, asyncio.Task[None]] + # Implementing classes should call super().__init__() in their __init__ method + # to properly initialize the mixin attributes. + def __init__(self, files_api: Files | None = None, kvstore: KVStore | None = None): + self.openai_vector_stores: dict[str, dict[str, Any]] = {} + self.openai_file_batches: dict[str, dict[str, Any]] = {} + self.files_api = files_api + self.kvstore = kvstore + self._last_file_batch_cleanup_time = 0 + self._file_batch_tasks: dict[str, asyncio.Task[None]] = {} async def _save_openai_vector_store(self, store_id: str, store_info: dict[str, Any]) -> None: """Save vector store metadata to persistent storage."""