diff --git a/docs/modules/storages.md b/docs/modules/storages.md index 2416de8c..839e62b7 100644 --- a/docs/modules/storages.md +++ b/docs/modules/storages.md @@ -122,7 +122,7 @@ The `StorageHandler` is tightly integrated with the `RAGEngine` class to support ## Configuration The `StorageHandler` relies on the `StoreConfig` class (defined in `storages_config.py`) to configure its backends: - **`DBConfig`**: Configures relational databases (e.g., SQLite) with settings like `db_name`, `path`, `ip`, and `port`. -- **`VectorStoreConfig`**: Configures vector databases (e.g., FAISS, Qdrant) with settings like `vector_name`, `dimensions`, `index_type`, `qdrant_url`, and `qdrant_collection_name`. +- **`VectorStoreConfig`**: Configures vector databases (e.g., FAISS, Qdrant, Milvus) with settings like `vector_name`, `dimensions`, `index_type`, `qdrant_url`, `qdrant_collection_name`, `milvus_uri`, and `milvus_collection_name`. - **`GraphStoreConfig`**: Configures graph databases (e.g., Neo4j) with settings上午 like `graph_name`, `uri`, `username`, `password`, and `database`. The configuration is validated using Pydantic, ensuring robust type checking and default values. @@ -166,6 +166,28 @@ index = storage_handler.load_index("corpus1") print(index) # {'corpus_id': 'corpus1', 'content': {...}, 'metadata': {...}} ``` +To use Milvus instead of FAISS, install the optional dependency and set `vector_name="milvus"`: + +```bash +pip install "evoagentx[milvus]" +``` + +```python +config = StoreConfig( + dbConfig=DBConfig(db_name="sqlite", path="data/storage.db"), + vectorConfig=VectorStoreConfig( + vector_name="milvus", + dimensions=1536, + milvus_uri="./data/evoagentx_milvus.db", + milvus_collection_name="evoagentx_vectors", + milvus_metric_type="IP", + ), + path="data/index_cache" +) +``` + +`milvus_uri` can point to a Milvus Lite `.db` file, a self-hosted Milvus endpoint such as `http://localhost:19530`, or a Zilliz Cloud endpoint when paired with `milvus_token`. + ## Notes - The `load_memory` and `save_memory` methods are not yet fully implemented and will be developed alongside `LongTermMemory`. - The `StorageHandler` assumes the database schema is managed by `DBStoreBase` and its factory, ensuring compatibility with `TableType` enums. @@ -173,4 +195,4 @@ print(index) # {'corpus_id': 'corpus1', 'content': {...}, 'metadata': {...}} - Error handling is implemented throughout, with logs generated via the `evoagentx.core.logging.logger` module. ## Conclusion -The `StorageHandler` class provides a flexible and extensible interface for managing multiple storage backends in a unified manner. Its integration with `RAGEngine` makes it a key component for RAG pipelines, enabling efficient storage and retrieval of indexed data. By leveraging factory patterns and Pydantic validation, it ensures robustness and scalability for applications requiring complex data management. \ No newline at end of file +The `StorageHandler` class provides a flexible and extensible interface for managing multiple storage backends in a unified manner. Its integration with `RAGEngine` makes it a key component for RAG pipelines, enabling efficient storage and retrieval of indexed data. By leveraging factory patterns and Pydantic validation, it ensures robustness and scalability for applications requiring complex data management. diff --git a/docs/tutorial/rag.md b/docs/tutorial/rag.md index 67d41c00..35e1e46b 100644 --- a/docs/tutorial/rag.md +++ b/docs/tutorial/rag.md @@ -75,6 +75,24 @@ print("RAGEngine is ready to go!") - `RetrievalConfig`: Retrieves the top 3 most similar chunks with a similarity score above 0.3. - **Initialization**: We create the `RAGEngine` instance, ready to process documents. +To use Milvus instead of FAISS, install the optional dependency and replace the vector store configuration: + +```bash +pip install "evoagentx[milvus]" +``` + +```python +vectorConfig=VectorStoreConfig( + vector_name="milvus", + dimensions=1536, + milvus_uri="./data/evoagentx_milvus.db", + milvus_collection_name="rag_tutorial_vectors", + milvus_metric_type="IP", +) +``` + +Use a Milvus Lite `.db` path for local development, or set `milvus_uri` and `milvus_token` for Milvus server or Zilliz Cloud. + For more details on configuration, see the [RAGEngine documentation](../modules/rag.md). ## 2. Indexing and Querying Documents @@ -458,4 +476,4 @@ Congratulations! You’ve built your first RAG system with `RAGEngine`. Here are For a complete example, refer to the [RAGEngine example](../../examples/rag.py). -Happy building with EvoAgentX! \ No newline at end of file +Happy building with EvoAgentX! diff --git a/evoagentx/storages/storages_config.py b/evoagentx/storages/storages_config.py index ef6e4633..bee738cb 100644 --- a/evoagentx/storages/storages_config.py +++ b/evoagentx/storages/storages_config.py @@ -17,14 +17,21 @@ class DBConfig(BaseConfig): class VectorStoreConfig(BaseConfig): """ - Configuration for vector databases, supporting FAISS and Qdrant. + Configuration for vector databases, supporting FAISS, Qdrant, and Milvus. """ - vector_name: str = Field(default="faiss", description="Name of the vector database provider (e.g., 'faiss', 'qdrant')") + vector_name: str = Field(default="faiss", description="Name of the vector database provider (e.g., 'faiss', 'qdrant', 'milvus')") dimensions: Optional[int] = Field(default=1536, description="Dimension of the embedding vectors") index_type: Optional[str] = Field(default="flat_l2", description="Index type for FAISS (e.g., 'flat_l2', 'ivf_flat')") qdrant_url: Optional[str] = Field(default=None, description="URL for Qdrant server (e.g., 'http://localhost:6333')") qdrant_api_key: Optional[str] = Field(default=None, description="API key for Qdrant authentication") qdrant_collection_name: Optional[str] = Field(default="default_collection", description="Name of the Qdrant collection") + milvus_uri: Optional[str] = Field(default="./milvus.db", description="Milvus URI. Use a local .db path for Milvus Lite or an HTTP endpoint for Milvus server/Zilliz Cloud.") + milvus_token: Optional[str] = Field(default=None, description="Token for Milvus server or Zilliz Cloud authentication") + milvus_db_name: Optional[str] = Field(default=None, description="Milvus database name") + milvus_collection_name: Optional[str] = Field(default="evoagentx_vectors", description="Name of the Milvus collection") + milvus_metric_type: Optional[str] = Field(default="IP", description="Milvus vector metric type, such as IP, COSINE, or L2") + milvus_consistency_level: Optional[str] = Field(default="Session", description="Milvus consistency level") + milvus_overwrite: Optional[bool] = Field(default=False, description="Drop and recreate the Milvus collection during initialization") class GraphStoreConfig(BaseConfig): @@ -48,4 +55,4 @@ class StoreConfig(BaseConfig): vectorConfig: Optional[VectorStoreConfig] = Field(None, description="Configuration for the vector store") graphConfig: Optional[GraphStoreConfig] = Field(None, description="Optional configuration for the graph store") # For file storage - path: Optional[str] = Field(default="/index_cache", description="directory path for storing the index") \ No newline at end of file + path: Optional[str] = Field(default="/index_cache", description="directory path for storing the index") diff --git a/evoagentx/storages/vectore_stores/__init__.py b/evoagentx/storages/vectore_stores/__init__.py index f4e2052c..bc920b35 100644 --- a/evoagentx/storages/vectore_stores/__init__.py +++ b/evoagentx/storages/vectore_stores/__init__.py @@ -26,7 +26,14 @@ def create(self, store_type: str, store_config: Dict[str, Any] = None) -> Vector # raise ValueError("Qdrant requires a valid URL") # client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key) # vector_store = QdrantVectorStore(client=client, collection_name=collection_name) + elif store_type == VectorStoreType.MILVUS: + dimensions = store_config.get("dimensions") + if not dimensions or not isinstance(dimensions, int): + raise ValueError("Milvus requires a valid dimension") + from .milvus import MilvusVectorStoreWrapper + + vector_store = MilvusVectorStoreWrapper(**store_config) else: raise ValueError(f"Unsupported vector store type: {store_type}") logger.info(f"Created vector store: {store_type}") - return vector_store \ No newline at end of file + return vector_store diff --git a/evoagentx/storages/vectore_stores/base.py b/evoagentx/storages/vectore_stores/base.py index baeb8afc..44e3c2f3 100644 --- a/evoagentx/storages/vectore_stores/base.py +++ b/evoagentx/storages/vectore_stores/base.py @@ -6,6 +6,7 @@ class VectorStoreType(str, Enum): FAISS = "faiss" + MILVUS = "milvus" class VectorStoreBase(ABC): @@ -19,4 +20,4 @@ def get_vector_store(self) -> BasePydanticVectorStore: @abstractmethod async def aload(self): """Asynchronously load a single node into the vector store.""" - pass \ No newline at end of file + pass diff --git a/evoagentx/storages/vectore_stores/milvus.py b/evoagentx/storages/vectore_stores/milvus.py new file mode 100644 index 00000000..6e7162ab --- /dev/null +++ b/evoagentx/storages/vectore_stores/milvus.py @@ -0,0 +1,340 @@ +import json +import re +from typing import Any, Dict, List, Optional, Sequence + +from llama_index.core.bridge.pydantic import Field, PrivateAttr +from llama_index.core.schema import BaseNode +from llama_index.core.vector_stores.types import ( + BasePydanticVectorStore, + FilterCondition, + FilterOperator, + MetadataFilters, + VectorStoreQuery, + VectorStoreQueryMode, + VectorStoreQueryResult, +) +from llama_index.core.vector_stores.utils import node_to_metadata_dict + +from .base import VectorStoreBase +from evoagentx.core.logging import logger + + +DEFAULT_MILVUS_URI = "./milvus.db" +DEFAULT_COLLECTION_NAME = "evoagentx_vectors" +DEFAULT_METRIC_TYPE = "IP" +DEFAULT_CONSISTENCY_LEVEL = "Session" +DEFAULT_TEXT_ID_FIELD = "id" +DEFAULT_DOC_ID_FIELD = "doc_id" +DEFAULT_EMBEDDING_FIELD = "embedding" +MAX_VARCHAR_LENGTH = 65535 + +_VALID_FIELD_NAME = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") +_RESERVED_FIELDS = { + DEFAULT_TEXT_ID_FIELD, + DEFAULT_DOC_ID_FIELD, + DEFAULT_EMBEDDING_FIELD, +} + + +class MilvusVectorStore(BasePydanticVectorStore): + """LlamaIndex-compatible Milvus vector store backed by MilvusClient.""" + + stores_text: bool = False + uri: str = Field(default=DEFAULT_MILVUS_URI) + token: Optional[str] = Field(default=None) + db_name: Optional[str] = Field(default=None) + collection_name: str = Field(default=DEFAULT_COLLECTION_NAME) + dimensions: int = Field(default=1536) + metric_type: str = Field(default=DEFAULT_METRIC_TYPE) + consistency_level: str = Field(default=DEFAULT_CONSISTENCY_LEVEL) + overwrite: bool = Field(default=False) + text_id_field: str = Field(default=DEFAULT_TEXT_ID_FIELD) + doc_id_field: str = Field(default=DEFAULT_DOC_ID_FIELD) + embedding_field: str = Field(default=DEFAULT_EMBEDDING_FIELD) + + _client: Any = PrivateAttr() + + def __init__(self, **kwargs: Any) -> None: + super().__init__(**kwargs) + self.metric_type = self.metric_type.upper() + self._client = self._create_client() + self._ensure_collection() + + @classmethod + def class_name(cls) -> str: + return "MilvusVectorStore" + + @property + def client(self) -> Any: + return self._client + + def _create_client(self) -> Any: + try: + from pymilvus import MilvusClient + except ImportError as exc: + raise ImportError( + "Milvus vector store requires the optional dependency " + "`pymilvus[milvus-lite]`. Install it with `pip install evoagentx[milvus]`." + ) from exc + + return MilvusClient( + uri=self.uri, + token=self.token or "", + db_name=self.db_name or "", + ) + + def _ensure_collection(self) -> None: + if self.overwrite and self._client.has_collection(self.collection_name): + self._client.drop_collection(self.collection_name) + + if self._client.has_collection(self.collection_name): + self._validate_existing_collection() + return + + from pymilvus import DataType + + schema = self._client.create_schema( + auto_id=False, + enable_dynamic_field=True, + ) + schema.add_field( + field_name=self.text_id_field, + datatype=DataType.VARCHAR, + is_primary=True, + max_length=MAX_VARCHAR_LENGTH, + ) + schema.add_field( + field_name=self.embedding_field, + datatype=DataType.FLOAT_VECTOR, + dim=self.dimensions, + ) + schema.add_field( + field_name=self.doc_id_field, + datatype=DataType.VARCHAR, + max_length=MAX_VARCHAR_LENGTH, + ) + + index_params = self._client.prepare_index_params() + index_params.add_index( + field_name=self.embedding_field, + index_type="AUTOINDEX", + metric_type=self.metric_type, + ) + + self._client.create_collection( + collection_name=self.collection_name, + schema=schema, + index_params=index_params, + consistency_level=self.consistency_level, + ) + + def _validate_existing_collection(self) -> None: + description = self._client.describe_collection(self.collection_name) + fields = {field["name"]: field for field in description["fields"]} + + if self.text_id_field not in fields: + raise ValueError( + f"Milvus collection '{self.collection_name}' is missing primary field " + f"'{self.text_id_field}'." + ) + if self.embedding_field not in fields: + raise ValueError( + f"Milvus collection '{self.collection_name}' is missing vector field " + f"'{self.embedding_field}'." + ) + + vector_dim = int(fields[self.embedding_field]["params"].get("dim", 0)) + if vector_dim != self.dimensions: + raise ValueError( + f"Milvus collection '{self.collection_name}' has dimension {vector_dim}, " + f"but the configured dimension is {self.dimensions}." + ) + + def add( + self, + nodes: Sequence[BaseNode], + **add_kwargs: Any, + ) -> List[str]: + rows = [self._node_to_row(node) for node in nodes] + if rows: + self._client.upsert( + collection_name=self.collection_name, + data=rows, + ) + return [node.node_id for node in nodes] + + def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None: + self._client.delete( + collection_name=self.collection_name, + filter=f"{self.doc_id_field} == {self._format_value(ref_doc_id)}", + ) + + def delete_nodes( + self, + node_ids: Optional[List[str]] = None, + filters: Optional[MetadataFilters] = None, + **delete_kwargs: Any, + ) -> None: + expressions = [] + if node_ids: + ids = ", ".join(self._format_value(node_id) for node_id in node_ids) + expressions.append(f"{self.text_id_field} in [{ids}]") + if filters is not None: + expressions.append(self._to_filter_expression(filters)) + + if not expressions: + return + + self._client.delete( + collection_name=self.collection_name, + filter=" and ".join(f"({expr})" for expr in expressions if expr), + ) + + def clear(self) -> None: + if self._client.has_collection(self.collection_name): + self._client.drop_collection(self.collection_name) + self._ensure_collection() + + def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult: + if query.mode != VectorStoreQueryMode.DEFAULT: + raise ValueError(f"Milvus vector store does not support query mode: {query.mode}") + if query.query_embedding is None: + raise ValueError("Milvus vector store requires a query embedding.") + + filters = [] + if query.node_ids: + ids = ", ".join(self._format_value(node_id) for node_id in query.node_ids) + filters.append(f"{self.text_id_field} in [{ids}]") + if query.filters is not None: + filters.append(self._to_filter_expression(query.filters)) + + result = self._client.search( + collection_name=self.collection_name, + data=[query.query_embedding], + anns_field=self.embedding_field, + limit=query.similarity_top_k, + filter=" and ".join(f"({expr})" for expr in filters if expr), + output_fields=[self.text_id_field], + ) + + hits = result[0] if result else [] + ids = [str(hit.get("id", hit.get(self.text_id_field))) for hit in hits] + similarities = [float(hit["distance"]) for hit in hits] + return VectorStoreQueryResult(ids=ids, similarities=similarities) + + def _node_to_row(self, node: BaseNode) -> Dict[str, Any]: + metadata = node_to_metadata_dict(node, remove_text=True, flat_metadata=False) + metadata.pop("_node_content", None) + + row = { + self.text_id_field: node.node_id, + self.embedding_field: node.get_embedding(), + self.doc_id_field: node.ref_doc_id or "None", + } + row.update(self._metadata_to_dynamic_fields(metadata)) + return row + + def _metadata_to_dynamic_fields(self, metadata: Dict[str, Any]) -> Dict[str, Any]: + fields = {} + for key, value in metadata.items(): + if key in _RESERVED_FIELDS or not _VALID_FIELD_NAME.match(key): + continue + if isinstance(value, (str, int, float, bool)) or value is None: + fields[key] = value + return fields + + def _to_filter_expression(self, filters: MetadataFilters) -> str: + condition = filters.condition or FilterCondition.AND + joiner = " and " if condition == FilterCondition.AND else " or " + expressions = [] + + for filter_item in filters.filters: + if isinstance(filter_item, MetadataFilters): + expressions.append(f"({self._to_filter_expression(filter_item)})") + continue + + field = filter_item.key + if not _VALID_FIELD_NAME.match(field) or field in _RESERVED_FIELDS: + raise ValueError(f"Unsupported metadata filter field: {field}") + expressions.append( + self._metadata_filter_to_expression( + field=field, + operator=filter_item.operator, + value=filter_item.value, + ) + ) + + return joiner.join(expressions) + + def _metadata_filter_to_expression( + self, + field: str, + operator: FilterOperator, + value: Any, + ) -> str: + if operator == FilterOperator.EQ: + return f"{field} == {self._format_value(value)}" + if operator == FilterOperator.NE: + return f"{field} != {self._format_value(value)}" + if operator == FilterOperator.GT: + return f"{field} > {self._format_value(value)}" + if operator == FilterOperator.GTE: + return f"{field} >= {self._format_value(value)}" + if operator == FilterOperator.LT: + return f"{field} < {self._format_value(value)}" + if operator == FilterOperator.LTE: + return f"{field} <= {self._format_value(value)}" + if operator == FilterOperator.IN: + return f"{field} in {self._format_value_list(value)}" + if operator == FilterOperator.NIN: + return f"{field} not in {self._format_value_list(value)}" + + raise ValueError(f"Unsupported metadata filter operator for Milvus: {operator}") + + def _format_value_list(self, value: Any) -> str: + if not isinstance(value, list): + raise ValueError("Milvus 'in' filters require a list value.") + return "[" + ", ".join(self._format_value(item) for item in value) + "]" + + def _format_value(self, value: Any) -> str: + if isinstance(value, bool): + return "true" if value else "false" + if isinstance(value, (int, float)): + return str(value) + if value is None: + return "null" + return json.dumps(str(value)) + + +class MilvusVectorStoreWrapper(VectorStoreBase): + """Wrapper for Milvus vector store.""" + + def __init__( + self, + dimensions: int = 1536, + milvus_uri: str = DEFAULT_MILVUS_URI, + milvus_token: Optional[str] = None, + milvus_db_name: Optional[str] = None, + milvus_collection_name: str = DEFAULT_COLLECTION_NAME, + milvus_metric_type: str = DEFAULT_METRIC_TYPE, + milvus_consistency_level: str = DEFAULT_CONSISTENCY_LEVEL, + milvus_overwrite: bool = False, + **kwargs: Any, + ) -> None: + self.vector_store = MilvusVectorStore( + uri=milvus_uri or DEFAULT_MILVUS_URI, + token=milvus_token, + db_name=milvus_db_name, + collection_name=milvus_collection_name or DEFAULT_COLLECTION_NAME, + dimensions=dimensions, + metric_type=milvus_metric_type or DEFAULT_METRIC_TYPE, + consistency_level=milvus_consistency_level or DEFAULT_CONSISTENCY_LEVEL, + overwrite=bool(milvus_overwrite), + ) + + def get_vector_store(self) -> MilvusVectorStore: + return self.vector_store + + async def aload(self, node: BaseNode) -> None: + self.vector_store.add([node]) + logger.info(f"Inserted node with ID {node.node_id} into Milvus vector store.") diff --git a/pyproject.toml b/pyproject.toml index 07889f67..291f370b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,9 @@ rag = [ "docx2txt", "python-pptx" ] +milvus = [ + "pymilvus[milvus-lite]>=3.0.0" +] tools = [ "docker>=6.1.2", "googlesearch-python", @@ -159,7 +162,8 @@ all = [ "ujson>=5.9.0", "sympy", "antlr4-python3-runtime==4.11", - "matplotlib>=3.10.0" + "matplotlib>=3.10.0", + "pymilvus[milvus-lite]>=3.0.0" ] [project.urls] diff --git a/requirements.txt b/requirements.txt index 35a05b96..1d17db23 100644 --- a/requirements.txt +++ b/requirements.txt @@ -40,6 +40,7 @@ neo4j ollama docx2txt python-pptx +pymilvus[milvus-lite]>=3.0.0 # tools docker>=6.1.2 diff --git a/tests/src/storages/test_milvus_vector_store.py b/tests/src/storages/test_milvus_vector_store.py new file mode 100644 index 00000000..556f54df --- /dev/null +++ b/tests/src/storages/test_milvus_vector_store.py @@ -0,0 +1,216 @@ +"""Milvus vector store integration tests.""" + +# ruff: noqa: E402 + +from pathlib import Path +import importlib.util +import sys +import types +from unittest.mock import patch + +import pytest +from llama_index.core.embeddings import BaseEmbedding +from llama_index.core.schema import TextNode +from llama_index.core.vector_stores.types import MetadataFilters, VectorStoreQuery + +pytest.importorskip("pymilvus") + + +def install_optional_embedding_stubs() -> None: + """Stub unused optional embedding packages when running focused tests.""" + if importlib.util.find_spec("sentence_transformers") is None: + sentence_transformers = types.ModuleType("sentence_transformers") + + class SentenceTransformer: + pass + + sentence_transformers.SentenceTransformer = SentenceTransformer + sys.modules["sentence_transformers"] = sentence_transformers + + if importlib.util.find_spec("voyageai") is None: + voyageai = types.ModuleType("voyageai") + + class AsyncClient: + pass + + voyageai.AsyncClient = AsyncClient + sys.modules["voyageai"] = voyageai + + +install_optional_embedding_stubs() + +from evoagentx.rag.indexings.base import IndexType +from evoagentx.rag.rag import RAGEngine +from evoagentx.rag.rag_config import ( + ChunkerConfig, + EmbeddingConfig, + IndexConfig, + RAGConfig, + ReaderConfig, + RetrievalConfig, +) +from evoagentx.rag.schema import Corpus, Query, TextChunk +from evoagentx.storages.base import StorageHandler +from evoagentx.storages.storages_config import DBConfig, StoreConfig, VectorStoreConfig +from evoagentx.storages.vectore_stores import VectorStoreFactory + + +class DeterministicEmbedding(BaseEmbedding): + """Small deterministic embedding model for vector-store integration tests.""" + + def __init__(self) -> None: + super().__init__(model_name="deterministic-test", embed_batch_size=10) + + def _embed(self, text: str) -> list[float]: + text = text.lower() + if "milvus" in text: + return [1.0, 0.0] + if "faiss" in text: + return [0.0, 1.0] + return [0.5, 0.5] + + def _get_query_embedding(self, query: str) -> list[float]: + return self._embed(query) + + def _get_text_embedding(self, text: str) -> list[float]: + return self._embed(text) + + async def _aget_query_embedding(self, query: str) -> list[float]: + return self._embed(query) + + async def _aget_text_embedding(self, text: str) -> list[float]: + return self._embed(text) + + def _get_text_embeddings(self, texts: list[str]) -> list[list[float]]: + return [self._embed(text) for text in texts] + + async def _aget_text_embeddings(self, texts: list[str]) -> list[list[float]]: + return [self._embed(text) for text in texts] + + @property + def dimensions(self) -> int: + return 2 + + +class DeterministicEmbeddingWrapper: + """Embedding wrapper that avoids external API calls in tests.""" + + def __init__(self) -> None: + self.model_name = "deterministic-test" + self._embedding_model = DeterministicEmbedding() + + def get_embedding_model(self) -> BaseEmbedding: + return self._embedding_model + + def validate_model(self, provider: str, model_name: str) -> bool: + return model_name == self.model_name + + @property + def dimensions(self) -> int: + return 2 + + +def milvus_config(tmp_path: Path, collection_name: str = "evoagentx_test") -> VectorStoreConfig: + return VectorStoreConfig( + vector_name="milvus", + dimensions=2, + milvus_uri=str(tmp_path / "milvus.db"), + milvus_collection_name=collection_name, + milvus_metric_type="IP", + milvus_consistency_level="Strong", + milvus_overwrite=True, + ) + + +def test_milvus_vector_store_add_query_filter_and_delete(tmp_path: Path) -> None: + wrapper = VectorStoreFactory().create( + store_type="milvus", + store_config=milvus_config(tmp_path).model_dump(), + ) + vector_store = wrapper.get_vector_store() + + vector_store.add( + [ + TextNode( + id_="milvus-node", + text="Milvus stores vectors.", + embedding=[1.0, 0.0], + metadata={"topic": "milvus"}, + ), + TextNode( + id_="faiss-node", + text="FAISS stores vectors.", + embedding=[0.0, 1.0], + metadata={"topic": "faiss"}, + ), + ] + ) + + result = vector_store.query( + VectorStoreQuery(query_embedding=[1.0, 0.0], similarity_top_k=2) + ) + assert result.ids == ["milvus-node", "faiss-node"] + assert result.similarities[0] > result.similarities[1] + + filtered = vector_store.query( + VectorStoreQuery( + query_embedding=[1.0, 0.0], + similarity_top_k=2, + filters=MetadataFilters.from_dicts([{"key": "topic", "value": "milvus"}]), + ) + ) + assert filtered.ids == ["milvus-node"] + + vector_store.delete_nodes(["milvus-node"]) + after_delete = vector_store.query( + VectorStoreQuery(query_embedding=[1.0, 0.0], similarity_top_k=2) + ) + assert after_delete.ids == ["faiss-node"] + + +def test_rag_engine_uses_milvus_vector_store(tmp_path: Path) -> None: + store_config = StoreConfig( + dbConfig=DBConfig(db_name="sqlite", path=str(tmp_path / "storage.db")), + vectorConfig=milvus_config(tmp_path, collection_name="evoagentx_rag_test"), + graphConfig=None, + path=str(tmp_path / "index_cache"), + ) + storage_handler = StorageHandler(storageConfig=store_config) + rag_config = RAGConfig( + reader=ReaderConfig(), + chunker=ChunkerConfig(strategy="simple", chunk_size=512, chunk_overlap=0), + embedding=EmbeddingConfig( + provider="openai", + model_name="deterministic-test", + api_key="dummy-key", + ), + index=IndexConfig(index_type="vector"), + retrieval=RetrievalConfig( + retrivel_type="vector", + postprocessor_type="simple", + top_k=1, + similarity_cutoff=None, + ), + ) + + with patch( + "evoagentx.rag.rag.EmbeddingFactory.create", + return_value=DeterministicEmbeddingWrapper(), + ): + rag_engine = RAGEngine(config=rag_config, storage_handler=storage_handler) + corpus = Corpus( + corpus_id="test_corpus", + chunks=[ + TextChunk(text="Milvus provides scalable vector search.", chunk_id="milvus"), + TextChunk(text="FAISS provides local vector search.", chunk_id="faiss"), + ], + ) + rag_engine.add(index_type=IndexType.VECTOR, nodes=corpus, corpus_id="test_corpus") + + result = rag_engine.query( + Query(query_str="Milvus vector database", top_k=1, similarity_cutoff=None), + corpus_id="test_corpus", + ) + + assert len(result.corpus.chunks) == 1 + assert result.corpus.chunks[0].chunk_id == "milvus"