From e13a2a64fe3006ee72d2de1f82918c18596b15fb Mon Sep 17 00:00:00 2001 From: Francisco Javier Arceo Date: Sat, 26 Jul 2025 00:33:37 -0400 Subject: [PATCH] WIP Signed-off-by: Francisco Javier Arceo --- .../providers/vector_io/remote_weaviate.md | 10 +++ .../remote/vector_io/weaviate/__init__.py | 2 +- .../remote/vector_io/weaviate/config.py | 12 +-- .../remote/vector_io/weaviate/weaviate.py | 74 +++++++++---------- pyproject.toml | 1 + .../vector_io/test_openai_vector_stores.py | 1 + uv.lock | 66 +++++++++++++++++ 7 files changed, 119 insertions(+), 47 deletions(-) diff --git a/docs/source/providers/vector_io/remote_weaviate.md b/docs/source/providers/vector_io/remote_weaviate.md index d930515d5..1b76d8375 100644 --- a/docs/source/providers/vector_io/remote_weaviate.md +++ b/docs/source/providers/vector_io/remote_weaviate.md @@ -33,9 +33,19 @@ To install Weaviate see the [Weaviate quickstart documentation](https://weaviate See [Weaviate's documentation](https://weaviate.io/developers/weaviate) for more details about Weaviate in general. +## Configuration + +| Field | Type | Required | Default | Description | +|-------|------|----------|---------|-------------| +| `weaviate_api_key` | `` | No | PydanticUndefined | The API key for the Weaviate instance | +| `weaviate_cluster_url` | `` | No | PydanticUndefined | The URL of the Weaviate cluster | +| `kvstore` | `utils.kvstore.config.RedisKVStoreConfig \| utils.kvstore.config.SqliteKVStoreConfig \| utils.kvstore.config.PostgresKVStoreConfig \| utils.kvstore.config.MongoDBKVStoreConfig, annotation=NoneType, required=False, default='sqlite', discriminator='type'` | No | | Config for KV store backend (SQLite only for now) | + ## Sample Configuration ```yaml +weaviate_api_key: dummy-api-key-for-testing +weaviate_cluster_url: http://localhost:8080 kvstore: type: sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/weaviate_registry.db diff --git a/llama_stack/providers/remote/vector_io/weaviate/__init__.py b/llama_stack/providers/remote/vector_io/weaviate/__init__.py index 22e116c22..9272b21e2 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/__init__.py +++ b/llama_stack/providers/remote/vector_io/weaviate/__init__.py @@ -12,6 +12,6 @@ from .config import WeaviateVectorIOConfig async def get_adapter_impl(config: WeaviateVectorIOConfig, deps: dict[Api, ProviderSpec]): from .weaviate import WeaviateVectorIOAdapter - impl = WeaviateVectorIOAdapter(config, deps[Api.inference]) + impl = WeaviateVectorIOAdapter(config, deps[Api.inference], deps.get(Api.files, None)) await impl.initialize() return impl diff --git a/llama_stack/providers/remote/vector_io/weaviate/config.py b/llama_stack/providers/remote/vector_io/weaviate/config.py index 4283b8d3b..de615060e 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/config.py +++ b/llama_stack/providers/remote/vector_io/weaviate/config.py @@ -12,18 +12,20 @@ from llama_stack.providers.utils.kvstore.config import ( KVStoreConfig, SqliteKVStoreConfig, ) +from llama_stack.schema_utils import json_schema_type -class WeaviateRequestProviderData(BaseModel): - weaviate_api_key: str - weaviate_cluster_url: str +@json_schema_type +class WeaviateVectorIOConfig(BaseModel): + weaviate_api_key: str = Field(description="The API key for the Weaviate instance") + weaviate_cluster_url: str = Field(description="The URL of the Weaviate cluster") kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) - -class WeaviateVectorIOConfig(BaseModel): @classmethod def sample_run_config(cls, __distro_dir__: str, **kwargs: Any) -> dict[str, Any]: return { + "weaviate_api_key": "dummy-api-key-for-testing", + "weaviate_cluster_url": "http://localhost:8080", "kvstore": SqliteKVStoreConfig.sample_run_config( __distro_dir__=__distro_dir__, db_name="weaviate_registry.db", diff --git a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py index 35bb40454..adcadbc29 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py +++ b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py @@ -21,12 +21,13 @@ from llama_stack.distribution.request_headers import NeedsRequestProviderData from llama_stack.providers.datatypes import Api, VectorDBsProtocolPrivate from llama_stack.providers.utils.kvstore import kvstore_impl from llama_stack.providers.utils.kvstore.api import KVStore +from llama_stack.providers.utils.memory.openai_vector_store_mixin import OpenAIVectorStoreMixin from llama_stack.providers.utils.memory.vector_store import ( EmbeddingIndex, VectorDBWithIndex, ) -from .config import WeaviateRequestProviderData, WeaviateVectorIOConfig +from .config import WeaviateVectorIOConfig log = logging.getLogger(__name__) @@ -116,6 +117,7 @@ class WeaviateIndex(EmbeddingIndex): class WeaviateVectorIOAdapter( + OpenAIVectorStoreMixin, VectorIO, NeedsRequestProviderData, VectorDBsProtocolPrivate, @@ -137,17 +139,13 @@ class WeaviateVectorIOAdapter( self.metadata_collection_name = "openai_vector_stores_metadata" def _get_client(self) -> weaviate.Client: - provider_data = self.get_request_provider_data() - assert provider_data is not None, "Request provider data must be set" - assert isinstance(provider_data, WeaviateRequestProviderData) - - key = f"{provider_data.weaviate_cluster_url}::{provider_data.weaviate_api_key}" + key = f"{self.config.weaviate_cluster_url}::{self.config.weaviate_api_key}" if key in self.client_cache: return self.client_cache[key] client = weaviate.connect_to_weaviate_cloud( - cluster_url=provider_data.weaviate_cluster_url, - auth_credentials=Auth.api_key(provider_data.weaviate_api_key), + cluster_url=self.config.weaviate_cluster_url, + auth_credentials=Auth.api_key(self.config.weaviate_api_key), ) self.client_cache[key] = client return client @@ -155,24 +153,29 @@ class WeaviateVectorIOAdapter( async def initialize(self) -> None: """Set up KV store and load existing vector DBs and OpenAI vector stores.""" # Initialize KV store for metadata - self.kvstore = await kvstore_impl(self.config.kvstore) + if self.kvstore is not None: + self.kvstore = await kvstore_impl(self.config.kvstore) + else: + self.kvstore = None + log.info("No kvstore configured, registry will not persist across restarts") # Load existing vector DB definitions - start_key = VECTOR_DBS_PREFIX - end_key = f"{VECTOR_DBS_PREFIX}\xff" - stored = await self.kvstore.values_in_range(start_key, end_key) - for raw in stored: - vector_db = VectorDB.model_validate_json(raw) - client = self._get_client() - idx = WeaviateIndex(client=client, collection_name=vector_db.identifier, kvstore=self.kvstore) - self.cache[vector_db.identifier] = VectorDBWithIndex( - vector_db=vector_db, - index=idx, - inference_api=self.inference_api, - ) + if self.kvstore is not None: + start_key = VECTOR_DBS_PREFIX + end_key = f"{VECTOR_DBS_PREFIX}\xff" + stored = await self.kvstore.values_in_range(start_key, end_key) + for raw in stored: + vector_db = VectorDB.model_validate_json(raw) + client = self._get_client() + idx = WeaviateIndex(client=client, collection_name=vector_db.identifier, kvstore=self.kvstore) + self.cache[vector_db.identifier] = VectorDBWithIndex( + vector_db=vector_db, + index=idx, + inference_api=self.inference_api, + ) - # Load OpenAI vector stores metadata into cache - await self.initialize_openai_vector_stores() + # Load OpenAI vector stores metadata into cache + await self.initialize_openai_vector_stores() async def shutdown(self) -> None: for client in self.client_cache.values(): @@ -203,6 +206,13 @@ class WeaviateVectorIOAdapter( self.inference_api, ) + async def unregister_vector_db(self, vector_db_id: str) -> None: + if vector_db_id not in self.cache: + log.warning(f"Vector DB {vector_db_id} not found") + return + await self.cache[vector_db_id].index.delete() + del self.cache[vector_db_id] + async def _get_and_cache_vector_db_index(self, vector_db_id: str) -> VectorDBWithIndex | None: if vector_db_id in self.cache: return self.cache[vector_db_id] @@ -246,21 +256,3 @@ class WeaviateVectorIOAdapter( raise ValueError(f"Vector DB {vector_db_id} not found") return await index.query_chunks(query, params) - - # OpenAI Vector Stores File operations are not supported in Weaviate - 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: - raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate") - - async def _load_openai_vector_store_file(self, store_id: str, file_id: str) -> dict[str, Any]: - raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate") - - async def _load_openai_vector_store_file_contents(self, store_id: str, file_id: str) -> list[dict[str, Any]]: - raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate") - - async def _update_openai_vector_store_file(self, store_id: str, file_id: str, file_info: dict[str, Any]) -> None: - raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate") - - async def _delete_openai_vector_store_file_from_storage(self, store_id: str, file_id: str) -> None: - raise NotImplementedError("OpenAI Vector Stores API is not supported in Weaviate") diff --git a/pyproject.toml b/pyproject.toml index ceb9d0b44..293cdcf5e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,6 +114,7 @@ test = [ "sqlalchemy[asyncio]>=2.0.41", "requests", "pymilvus>=2.5.12", + "weaviate-client>=4.16.4", ] docs = [ "setuptools", diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index 9771ab290..1b784c108 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -29,6 +29,7 @@ def skip_if_provider_doesnt_support_openai_vector_stores(client_with_models): "inline::chromadb", "remote::pgvector", "remote::chromadb", + "remote::weaviate", ]: return diff --git a/uv.lock b/uv.lock index 091873dac..79c05712d 100644 --- a/uv.lock +++ b/uv.lock @@ -192,6 +192,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fc/30/d4986a882011f9df997a55e6becd864812ccfcd821d64aac8570ee39f719/attrs-25.1.0-py3-none-any.whl", hash = "sha256:c75a69e28a550a7e93789579c22aa26b0f5b83b75dc4e08fe092980051e1090a", size = 63152 }, ] +[[package]] +name = "authlib" +version = "1.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/8e/a1/d8d1c6f8bc922c0b87ae0d933a8ed57be1bef6970894ed79c2852a153cd3/authlib-1.6.1.tar.gz", hash = "sha256:4dffdbb1460ba6ec8c17981a4c67af7d8af131231b5a36a88a1e8c80c111cdfd", size = 159988 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/58/cc6a08053f822f98f334d38a27687b69c6655fb05cd74a7a5e70a2aeed95/authlib-1.6.1-py2.py3-none-any.whl", hash = "sha256:e9d2031c34c6309373ab845afc24168fe9e93dc52d252631f52642f21f5ed06e", size = 239299 }, +] + [[package]] name = "autoevals" version = "0.0.122" @@ -693,6 +705,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6e/c6/ac0b6c1e2d138f1002bcf799d330bd6d85084fece321e662a14223794041/Deprecated-1.2.18-py2.py3-none-any.whl", hash = "sha256:bd5011788200372a32418f888e326a09ff80d0214bd961147cfed01b5c018eec", size = 9998 }, ] +[[package]] +name = "deprecation" +version = "2.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "packaging" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/d3/8ae2869247df154b64c1884d7346d412fed0c49df84db635aab2d1c40e62/deprecation-2.1.0.tar.gz", hash = "sha256:72b3bde64e5d778694b0cf68178aed03d15e15477116add3fb773e581f9518ff", size = 173788 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/c3/253a89ee03fc9b9682f1541728eb66db7db22148cd94f89ab22528cd1e1b/deprecation-2.1.0-py2.py3-none-any.whl", hash = "sha256:a10811591210e1fb0e768a8c25517cabeabcba6f0bf96564f8ff45189f90b14a", size = 11178 }, +] + [[package]] name = "dill" version = "0.3.8" @@ -1003,6 +1027,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/96/44759eca966720d0f3e1b105c43f8ad4590c97bf8eb3cd489656e9590baa/grpcio-1.67.1-cp313-cp313-win_amd64.whl", hash = "sha256:fa0c739ad8b1996bd24823950e3cb5152ae91fca1c09cc791190bf1627ffefba", size = 4346042 }, ] +[[package]] +name = "grpcio-health-checking" +version = "1.67.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "grpcio" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/64/dd/e3b339fa44dc75b501a1a22cb88f1af5b1f8c964488f19c4de4cfbbf05ba/grpcio_health_checking-1.67.1.tar.gz", hash = "sha256:ca90fa76a6afbb4fda71d734cb9767819bba14928b91e308cffbb0c311eb941e", size = 16775 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/8d/7a9878dca6616b48093d71c52d0bc79cb2dd1a2698ff6f5ce7406306de12/grpcio_health_checking-1.67.1-py3-none-any.whl", hash = "sha256:93753da5062152660aef2286c9b261e07dd87124a65e4dc9fbd47d1ce966b39d", size = 18924 }, +] + [[package]] name = "grpcio-tools" version = "1.67.1" @@ -1608,6 +1645,7 @@ test = [ { name = "torchvision", version = "0.21.0", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or sys_platform == 'darwin'" }, { name = "torchvision", version = "0.21.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "(platform_machine != 'aarch64' and sys_platform == 'linux') or (sys_platform != 'darwin' and sys_platform != 'linux')" }, { name = "transformers" }, + { name = "weaviate-client" }, ] unit = [ { name = "aiohttp" }, @@ -1719,6 +1757,7 @@ test = [ { name = "torch", specifier = ">=2.6.0", index = "https://download.pytorch.org/whl/cpu" }, { name = "torchvision", specifier = ">=0.21.0", index = "https://download.pytorch.org/whl/cpu" }, { name = "transformers" }, + { name = "weaviate-client", specifier = ">=4.16.4" }, ] unit = [ { name = "aiohttp" }, @@ -4393,6 +4432,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/63/9a/0962b05b308494e3202d3f794a6e85abe471fe3cafdbcf95c2e8c713aabd/uvloop-0.21.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a5c39f217ab3c663dc699c04cbd50c13813e31d917642d459fdcec07555cc553", size = 4660018 }, ] +[[package]] +name = "validators" +version = "0.35.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/66/a435d9ae49850b2f071f7ebd8119dd4e84872b01630d6736761e6e7fd847/validators-0.35.0.tar.gz", hash = "sha256:992d6c48a4e77c81f1b4daba10d16c3a9bb0dbb79b3a19ea847ff0928e70497a", size = 73399 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fa/6e/3e955517e22cbdd565f2f8b2e73d52528b14b8bcfdb04f62466b071de847/validators-0.35.0-py3-none-any.whl", hash = "sha256:e8c947097eae7892cb3d26868d637f79f47b4a0554bc6b80065dfe5aac3705dd", size = 44712 }, +] + [[package]] name = "virtualenv" version = "20.29.2" @@ -4470,6 +4518,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/fd/84/fd2ba7aafacbad3c4201d395674fc6348826569da3c0937e75505ead3528/wcwidth-0.2.13-py2.py3-none-any.whl", hash = "sha256:3da69048e4540d84af32131829ff948f1e022c1c6bdb8d6102117aac784f6859", size = 34166 }, ] +[[package]] +name = "weaviate-client" +version = "4.16.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "authlib" }, + { name = "deprecation" }, + { name = "grpcio" }, + { name = "grpcio-health-checking" }, + { name = "httpx" }, + { name = "pydantic" }, + { name = "validators" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/6a/be2257fbd890674204c393782e75b794cf0939a8dab8a4f1c0ab38d1663a/weaviate_client-4.16.4.tar.gz", hash = "sha256:50efc819274640564bbeca91405aef1674170e3d93f0bd7a9496b5b37d0fe746", size = 779617 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/a5/414c1452b786a27be1d1314c3051abf82af921ca20ada16df55c0d19997d/weaviate_client-4.16.4-py3-none-any.whl", hash = "sha256:299bab2ae822271ddc40005a7cd95f332e8243c7cff38df3bb9f4aca23221dcc", size = 597160 }, +] + [[package]] name = "websocket-client" version = "1.8.0"