diff --git a/docs/source/providers/vector_io/remote_weaviate.md b/docs/source/providers/vector_io/remote_weaviate.md index 1b76d8375..4873758ca 100644 --- a/docs/source/providers/vector_io/remote_weaviate.md +++ b/docs/source/providers/vector_io/remote_weaviate.md @@ -37,15 +37,15 @@ See [Weaviate's documentation](https://weaviate.io/developers/weaviate) for more | 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 | +| `weaviate_api_key` | `str \| None` | No | | The API key for the Weaviate instance | +| `weaviate_cluster_url` | `str \| None` | No | | 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 +weaviate_api_key: null +weaviate_cluster_url: null kvstore: type: sqlite db_path: ${env.SQLITE_STORE_DIR:=~/.llama/dummy}/weaviate_registry.db diff --git a/llama_stack/apis/vector_io/vector_io.py b/llama_stack/apis/vector_io/vector_io.py index 853c4656c..d09a9f2f3 100644 --- a/llama_stack/apis/vector_io/vector_io.py +++ b/llama_stack/apis/vector_io/vector_io.py @@ -16,7 +16,7 @@ from pydantic import BaseModel, Field from llama_stack.apis.inference import InterleavedContent from llama_stack.apis.vector_dbs import VectorDB from llama_stack.providers.utils.telemetry.trace_protocol import trace_protocol -from llama_stack.providers.utils.vector_io.chunk_utils import generate_chunk_id +from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id from llama_stack.schema_utils import json_schema_type, webmethod from llama_stack.strong_typing.schema import register_schema diff --git a/llama_stack/providers/remote/vector_io/milvus/milvus.py b/llama_stack/providers/remote/vector_io/milvus/milvus.py index dc4852821..281d96792 100644 --- a/llama_stack/providers/remote/vector_io/milvus/milvus.py +++ b/llama_stack/providers/remote/vector_io/milvus/milvus.py @@ -7,7 +7,6 @@ import asyncio import logging import os -import re from typing import Any from numpy.typing import NDArray @@ -30,6 +29,7 @@ from llama_stack.providers.utils.memory.vector_store import ( EmbeddingIndex, VectorDBWithIndex, ) +from llama_stack.providers.utils.vector_io.vector_utils import sanitize_collection_name from .config import MilvusVectorIOConfig as RemoteMilvusVectorIOConfig @@ -43,14 +43,6 @@ OPENAI_VECTOR_STORES_FILES_PREFIX = f"openai_vector_stores_files:milvus:{VERSION OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX = f"openai_vector_stores_files_contents:milvus:{VERSION}::" -def sanitize_collection_name(name: str) -> str: - """ - Sanitize collection name to ensure it only contains numbers, letters, and underscores. - Any other characters are replaced with underscores. - """ - return re.sub(r"[^a-zA-Z0-9_]", "_", name) - - class MilvusIndex(EmbeddingIndex): def __init__( self, client: MilvusClient, collection_name: str, consistency_level="Strong", kvstore: KVStore | None = None diff --git a/llama_stack/providers/remote/vector_io/weaviate/config.py b/llama_stack/providers/remote/vector_io/weaviate/config.py index de615060e..cbbbca758 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/config.py +++ b/llama_stack/providers/remote/vector_io/weaviate/config.py @@ -17,15 +17,15 @@ from llama_stack.schema_utils import json_schema_type @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") + weaviate_api_key: str | None = Field(description="The API key for the Weaviate instance", default=None) + weaviate_cluster_url: str | None = Field(description="The URL of the Weaviate cluster", default=None) kvstore: KVStoreConfig | None = Field(description="Config for KV store backend (SQLite only for now)", default=None) @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", + "weaviate_api_key": None, + "weaviate_cluster_url": None, "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 d7b920229..eaad5a45e 100644 --- a/llama_stack/providers/remote/vector_io/weaviate/weaviate.py +++ b/llama_stack/providers/remote/vector_io/weaviate/weaviate.py @@ -26,6 +26,7 @@ from llama_stack.providers.utils.memory.vector_store import ( EmbeddingIndex, VectorDBWithIndex, ) +from llama_stack.providers.utils.vector_io.vector_utils import sanitize_collection_name from .config import WeaviateVectorIOConfig @@ -42,9 +43,12 @@ OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX = f"openai_vector_stores_files_conten class WeaviateIndex(EmbeddingIndex): def __init__(self, client: weaviate.Client, collection_name: str, kvstore: KVStore | None = None): self.client = client - self.collection_name = collection_name + self.collection_name = sanitize_collection_name(collection_name, weaviate_format=True) self.kvstore = kvstore + async def initialize(self): + pass + async def add_chunks(self, chunks: list[Chunk], embeddings: NDArray): assert len(chunks) == len(embeddings), ( f"Chunk length {len(chunks)} does not match embedding length {len(embeddings)}" @@ -68,7 +72,8 @@ class WeaviateIndex(EmbeddingIndex): collection.data.insert_many(data_objects) async def query_vector(self, embedding: NDArray, k: int, score_threshold: float) -> QueryChunksResponse: - collection = self.client.collections.get(self.collection_name) + sanitized_collection_name = sanitize_collection_name(self.collection_name) + collection = self.client.collections.get(sanitized_collection_name) results = collection.query.near_vector( near_vector=embedding.tolist(), @@ -93,7 +98,8 @@ class WeaviateIndex(EmbeddingIndex): return QueryChunksResponse(chunks=chunks, scores=scores) async def delete(self, chunk_ids: list[str]) -> None: - collection = self.client.collections.get(self.collection_name) + sanitized_collection_name = sanitize_collection_name(self.collection_name) + collection = self.client.collections.get(sanitized_collection_name) collection.data.delete_many(where=Filter.by_property("id").contains_any(chunk_ids)) async def query_keyword( @@ -139,8 +145,7 @@ class WeaviateVectorIOAdapter( self.metadata_collection_name = "openai_vector_stores_metadata" def _get_client(self) -> weaviate.Client: - # if self.config.test_environment: - if True: + if self.config.weaviate_cluster_url is None: key = "local::test" client = weaviate.connect_to_local( host="localhost", @@ -160,8 +165,8 @@ 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 - if self.kvstore is not None: + # Initialize KV store for metadata if configured + if self.config.kvstore is not None: self.kvstore = await kvstore_impl(self.config.kvstore) else: self.kvstore = None @@ -194,11 +199,12 @@ class WeaviateVectorIOAdapter( vector_db: VectorDB, ) -> None: client = self._get_client() - + sanitized_collection_name = sanitize_collection_name(vector_db.identifier) # Create collection if it doesn't exist - if not client.collections.exists(vector_db.identifier): + if not client.collections.exists(sanitized_collection_name): + print(f"creating collection {vector_db}") client.collections.create( - name=vector_db.identifier, + name=sanitized_collection_name, vectorizer_config=wvc.config.Configure.Vectorizer.none(), properties=[ wvc.config.Property( @@ -208,37 +214,40 @@ class WeaviateVectorIOAdapter( ], ) - self.cache[vector_db.identifier] = VectorDBWithIndex( + self.cache[sanitized_collection_name] = VectorDBWithIndex( vector_db, - WeaviateIndex(client=client, collection_name=vector_db.identifier), + WeaviateIndex(client=client, collection_name=sanitized_collection_name), 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") + sanitized_collection_name = sanitize_collection_name(vector_db_id) + + if sanitized_collection_name not in self.cache: + log.warning(f"Vector DB {sanitized_collection_name} not found") return - await self.cache[vector_db_id].index.delete() - del self.cache[vector_db_id] + await self.cache[sanitized_collection_name].index.delete() + del self.cache[sanitized_collection_name] 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] + sanitized_collection_name = sanitize_collection_name(vector_db_id) + if sanitized_collection_name in self.cache: + return self.cache[sanitized_collection_name] - vector_db = await self.vector_db_store.get_vector_db(vector_db_id) + vector_db = await self.vector_db_store.get_vector_db(sanitized_collection_name) if not vector_db: - raise ValueError(f"Vector DB {vector_db_id} not found") + raise ValueError(f"Vector DB {sanitized_collection_name} not found") client = self._get_client() if not client.collections.exists(vector_db.identifier): - raise ValueError(f"Collection with name `{vector_db.identifier}` not found") + raise ValueError(f"Collection with name `{sanitized_collection_name}` not found") index = VectorDBWithIndex( vector_db=vector_db, - index=WeaviateIndex(client=client, collection_name=vector_db.identifier), + index=WeaviateIndex(client=client, collection_name=sanitized_collection_name), inference_api=self.inference_api, ) - self.cache[vector_db_id] = index + self.cache[sanitized_collection_name] = index return index async def insert_chunks( @@ -247,9 +256,10 @@ class WeaviateVectorIOAdapter( chunks: list[Chunk], ttl_seconds: int | None = None, ) -> None: - index = await self._get_and_cache_vector_db_index(vector_db_id) + sanitized_collection_name = sanitize_collection_name(vector_db_id) + index = await self._get_and_cache_vector_db_index(sanitized_collection_name) if not index: - raise ValueError(f"Vector DB {vector_db_id} not found") + raise ValueError(f"Vector DB {sanitized_collection_name} not found") await index.insert_chunks(chunks) @@ -259,8 +269,9 @@ class WeaviateVectorIOAdapter( query: InterleavedContent, params: dict[str, Any] | None = None, ) -> QueryChunksResponse: - index = await self._get_and_cache_vector_db_index(vector_db_id) + sanitized_collection_name = sanitize_collection_name(vector_db_id) + index = await self._get_and_cache_vector_db_index(sanitized_collection_name) if not index: - raise ValueError(f"Vector DB {vector_db_id} not found") + raise ValueError(f"Vector DB {sanitized_collection_name} not found") return await index.query_chunks(query, params) diff --git a/llama_stack/providers/utils/memory/vector_store.py b/llama_stack/providers/utils/memory/vector_store.py index f892d33c6..8c7fdff35 100644 --- a/llama_stack/providers/utils/memory/vector_store.py +++ b/llama_stack/providers/utils/memory/vector_store.py @@ -30,7 +30,7 @@ from llama_stack.providers.datatypes import Api from llama_stack.providers.utils.inference.prompt_adapter import ( interleaved_content_as_str, ) -from llama_stack.providers.utils.vector_io.chunk_utils import generate_chunk_id +from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id log = logging.getLogger(__name__) diff --git a/llama_stack/providers/utils/vector_io/chunk_utils.py b/llama_stack/providers/utils/vector_io/vector_utils.py similarity index 52% rename from llama_stack/providers/utils/vector_io/chunk_utils.py rename to llama_stack/providers/utils/vector_io/vector_utils.py index 01afa6ec8..8123caf98 100644 --- a/llama_stack/providers/utils/vector_io/chunk_utils.py +++ b/llama_stack/providers/utils/vector_io/vector_utils.py @@ -5,6 +5,7 @@ # the root directory of this source tree. import hashlib +import re import uuid @@ -19,3 +20,22 @@ def generate_chunk_id(document_id: str, chunk_text: str, chunk_window: str | Non if chunk_window: hash_input += f":{chunk_window}".encode() return str(uuid.UUID(hashlib.md5(hash_input, usedforsecurity=False).hexdigest())) + + +def proper_case(s: str) -> str: + """Convert a string to proper case (first letter uppercase, rest lowercase).""" + return s[0].upper() + s[1:].lower() if s else s + + +def sanitize_collection_name(name: str, weaviate_format=False) -> str: + """ + Sanitize collection name to ensure it only contains numbers, letters, and underscores. + Any other characters are replaced with underscores. + """ + print(f"Sanitizing collection name: {name} (Weaviate format: {weaviate_format})") + if not weaviate_format: + s = re.sub(r"[^a-zA-Z0-9_]", "_", name) + else: + s = proper_case(re.sub(r"[^a-zA-Z0-9]", "", name)) + print(f"Sanitized collection name from: {name} to: {s}") + return s diff --git a/tests/integration/vector_io/test_openai_vector_stores.py b/tests/integration/vector_io/test_openai_vector_stores.py index 1b784c108..b80d3d77a 100644 --- a/tests/integration/vector_io/test_openai_vector_stores.py +++ b/tests/integration/vector_io/test_openai_vector_stores.py @@ -110,11 +110,11 @@ def test_openai_create_vector_store(compat_client_with_empty_stores, client_with # Create a vector store vector_store = client.vector_stores.create( - name="test_vector_store", metadata={"purpose": "testing", "environment": "integration"} + name="Vs_test_vector_store", metadata={"purpose": "testing", "environment": "integration"} ) assert vector_store is not None - assert vector_store.name == "test_vector_store" + assert vector_store.name == "Vs_test_vector_store" assert vector_store.object == "vector_store" assert vector_store.status in ["completed", "in_progress"] assert vector_store.metadata["purpose"] == "testing" diff --git a/tests/unit/providers/vector_io/test_chunk_utils.py b/tests/unit/providers/vector_io/test_vector_utils.py similarity index 97% rename from tests/unit/providers/vector_io/test_chunk_utils.py rename to tests/unit/providers/vector_io/test_vector_utils.py index 535b76d73..a5d803a82 100644 --- a/tests/unit/providers/vector_io/test_chunk_utils.py +++ b/tests/unit/providers/vector_io/test_vector_utils.py @@ -5,7 +5,7 @@ # the root directory of this source tree. from llama_stack.apis.vector_io import Chunk, ChunkMetadata -from llama_stack.providers.utils.vector_io.chunk_utils import generate_chunk_id +from llama_stack.providers.utils.vector_io.vector_utils import generate_chunk_id # This test is a unit test for the chunk_utils.py helpers. This should only contain # tests which are specific to this file. More general (API-level) tests should be placed in