diff --git a/docs/deprecated.txt b/docs/deprecated.txt index 18618372..06f9648f 100644 --- a/docs/deprecated.txt +++ b/docs/deprecated.txt @@ -1,6 +1,6 @@ from loguru import logger -用英文注释,完善module/class/function docstring,要一句话简洁,不要变更代码 +用英文注释,完善module/class/function docstring,要一句话简洁,不要变更代码逻辑,符合pep和pylint规范,使用list而不是typing.List/Dict,不使用typing.Union 看看代码有什么问题 用英文注释,完善module/class/function docstring,要一句话简洁,代码要简洁,符合pep和pylint规范,使用list而不是typing.List,不使用typing.Union diff --git a/reme_ai/core/vector_store/__init__.py b/reme_ai/core/vector_store/__init__.py new file mode 100644 index 00000000..79500294 --- /dev/null +++ b/reme_ai/core/vector_store/__init__.py @@ -0,0 +1,17 @@ +"""vector store""" + +from .base_vector_store import BaseVectorStore +from .chroma_vector_store import ChromaVectorStore +from .es_vector_store import ESVectorStore +from .local_vector_store import LocalVectorStore +from .pgvector_store import PGVectorStore +from .qdrant_vector_store import QdrantVectorStore + +__all__ = [ + "BaseVectorStore", + "ChromaVectorStore", + "ESVectorStore", + "LocalVectorStore", + "PGVectorStore", + "QdrantVectorStore", +] diff --git a/reme_ai/core/vector_store/base_vector_store.py b/reme_ai/core/vector_store/base_vector_store.py new file mode 100644 index 00000000..e64b6dbf --- /dev/null +++ b/reme_ai/core/vector_store/base_vector_store.py @@ -0,0 +1,92 @@ +"""Base vector store interface for managing vector embeddings and similarity search.""" + +import asyncio +from abc import ABC, abstractmethod +from collections.abc import Callable +from functools import partial + +from reme_ai.core.context import C +from reme_ai.core.embedding import BaseEmbeddingModel +from reme_ai.core.schema import VectorNode + + +class BaseVectorStore(ABC): + """Abstract base class defining the interface for vector storage and retrieval.""" + + def __init__( + self, + collection_name: str, + embedding_model: BaseEmbeddingModel, + **kwargs, + ): + """Initialize the vector store with a collection name and an embedding model.""" + if embedding_model is None: + raise ValueError("embedding_model is required") + self.collection_name: str = collection_name + self.embedding_model: BaseEmbeddingModel = embedding_model + self.kwargs: dict = kwargs + + @staticmethod + async def _run_sync_in_executor(sync_func: Callable, *args, **kwargs): + """Run a synchronous function in the context-defined thread pool executor.""" + loop = asyncio.get_running_loop() + return await loop.run_in_executor(C.thread_pool, partial(sync_func, *args, **kwargs)) + + async def get_node_embedding(self, node: VectorNode) -> VectorNode: + """Generate and assign embedding for a single vector node.""" + return await self.embedding_model.get_node_embedding(node) + + async def get_node_embeddings(self, nodes: list[VectorNode]) -> list[VectorNode]: + """Generate and assign embeddings for multiple vector nodes.""" + return await self.embedding_model.get_node_embeddings(nodes) + + async def get_embedding(self, query: str) -> list[float]: + """Convert a single text query into vector embedding using the configured model.""" + return await self.embedding_model.get_embedding(query) + + async def get_embeddings(self, queries: list[str]) -> list[list[float]]: + """Convert multiple text queries into vector embeddings using the configured model.""" + return await self.embedding_model.get_embeddings(queries) + + @abstractmethod + async def list_collections(self) -> list[str]: + """Retrieve a list of all existing collection names in the store.""" + + @abstractmethod + async def create_collection(self, collection_name: str, **kwargs) -> None: + """Create a new vector collection with the specified name and configuration.""" + + @abstractmethod + async def delete_collection(self, collection_name: str, **kwargs) -> None: + """Permanently remove a collection from the vector store.""" + + @abstractmethod + async def copy_collection(self, collection_name: str, **kwargs) -> None: + """Duplicate the current collection to a new one with the given name.""" + + @abstractmethod + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None: + """Add one or more vector nodes into the current collection.""" + + @abstractmethod + async def search(self, query: str, limit: int = 5, filters: dict | None = None, **kwargs) -> list[VectorNode]: + """Find the most similar vector nodes based on a text query.""" + + @abstractmethod + async def delete(self, vector_ids: str | list[str], **kwargs) -> None: + """Remove specific vectors from the collection using their identifiers.""" + + @abstractmethod + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs) -> None: + """Update the data or metadata of existing vectors in the collection.""" + + @abstractmethod + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]: + """Fetch specific vector nodes from the collection by their IDs.""" + + @abstractmethod + async def list(self, filters: dict | None = None, limit: int | None = None) -> list[VectorNode]: + """Retrieve vectors from the collection that match the given filters.""" + + async def close(self) -> None: + """Release resources and close active connections to the vector store.""" diff --git a/reme_ai/core/vector_store/chroma_vector_store.py b/reme_ai/core/vector_store/chroma_vector_store.py new file mode 100644 index 00000000..24b88d36 --- /dev/null +++ b/reme_ai/core/vector_store/chroma_vector_store.py @@ -0,0 +1,397 @@ +"""ChromaDB vector store implementation for the ReMe framework.""" + +from typing import Any + +from loguru import logger + +from .base_vector_store import BaseVectorStore +from ..context import C +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + +_CHROMADB_IMPORT_ERROR = None + +try: + import chromadb + from chromadb.config import Settings +except ImportError as e: + _CHROMADB_IMPORT_ERROR = e + chromadb = None + Settings = None + + +@C.register_vector_store("chroma") +class ChromaVectorStore(BaseVectorStore): + """ChromaDB-based vector store implementation for local or remote storage.""" + + def __init__( + self, + collection_name: str, + embedding_model: BaseEmbeddingModel, + client: chromadb.ClientAPI | None = None, + host: str | None = None, + port: int | None = None, + path: str | None = None, + api_key: str | None = None, + tenant: str | None = None, + database: str | None = None, + **kwargs, + ): + """Initialize the ChromaDB vector store with the provided configuration.""" + if _CHROMADB_IMPORT_ERROR is not None: + raise ImportError( + "ChromaDB requires extra dependencies. Install with `pip install chromadb`", + ) from _CHROMADB_IMPORT_ERROR + + super().__init__( + collection_name=collection_name, + embedding_model=embedding_model, + **kwargs, + ) + + self.client: chromadb.ClientAPI + self.collection: chromadb.Collection + + if client: + self.client = client + elif api_key and tenant: + logger.info("Initializing ChromaDB Cloud client") + self.client = chromadb.CloudClient( + api_key=api_key, + tenant=tenant, + database=database or "default", + ) + elif host and port: + logger.info(f"Initializing ChromaDB HTTP client at {host}:{port}") + self.client = chromadb.HttpClient(host=host, port=port) + else: + if path is None: + path = "./chroma_db" + logger.info(f"Initializing local ChromaDB at {path}") + self.client = chromadb.PersistentClient( + path=path, + settings=Settings(anonymized_telemetry=False), + ) + + self.collection = self.client.get_or_create_collection( + name=collection_name, + metadata={"hnsw:space": "cosine"}, + ) + + @staticmethod + def _parse_results( + results: dict, + include_score: bool = False, + ) -> list[VectorNode]: + """Convert ChromaDB query results into a list of VectorNode objects.""" + nodes = [] + + ids = results.get("ids", []) + documents = results.get("documents", []) + metadatas = results.get("metadatas", []) + embeddings = results.get("embeddings") if results.get("embeddings") is not None else [] + distances = results.get("distances") if results.get("distances") is not None else [] + + if ids and isinstance(ids[0], list): + ids = ids[0] if ids else [] + documents = documents[0] if documents else [] + metadatas = metadatas[0] if metadatas else [] + embeddings = embeddings[0] if embeddings and len(embeddings) > 0 else [] + distances = distances[0] if distances and len(distances) > 0 else [] + + for i, vector_id in enumerate(ids): + metadata = metadatas[i] if i < len(metadatas) and metadatas[i] else {} + + if include_score and distances and i < len(distances): + metadata["_score"] = 1.0 - distances[i] + + node = VectorNode( + vector_id=vector_id, + content=documents[i] if i < len(documents) and documents[i] else "", + vector=embeddings[i] if len(embeddings) > i else None, + metadata=metadata, + ) + nodes.append(node) + + return nodes + + @staticmethod + def _generate_where_clause(filters: dict | None) -> dict | None: + """Convert the universal filter format to a ChromaDB-compatible where clause.""" + if not filters: + return None + + def convert_condition(k: str, v: Any) -> dict | None: + """Convert a single filter condition to ChromaDB operator format.""" + if v == "*": + return None + if isinstance(v, dict): + chroma_condition = {} + for op, val in v.items(): + mapping = { + "eq": "$eq", + "ne": "$ne", + "gt": "$gt", + "gte": "$gte", + "lt": "$lt", + "lte": "$lte", + "in": "$in", + "nin": "$nin", + } + chroma_op = mapping.get(op, "$eq") + chroma_condition[k] = {chroma_op: val} + return chroma_condition + if isinstance(v, list): + return {k: {"$in": v}} + return {k: {"$eq": v}} + + processed_filters = [] + + for key, value in filters.items(): + if key == "$or": + or_conditions = [] + for condition in value: + or_condition = {} + for sub_key, sub_value in condition.items(): + converted = convert_condition(sub_key, sub_value) + if converted: + or_condition.update(converted) + if or_condition: + or_conditions.append(or_condition) + if len(or_conditions) > 1: + processed_filters.append({"$or": or_conditions}) + elif len(or_conditions) == 1: + processed_filters.append(or_conditions[0]) + + elif key == "$and": + for condition in value: + for sub_key, sub_value in condition.items(): + converted = convert_condition(sub_key, sub_value) + if converted: + processed_filters.append(converted) + elif key == "$not": + continue + else: + converted = convert_condition(key, value) + if converted: + processed_filters.append(converted) + + if not processed_filters: + return None + return processed_filters[0] if len(processed_filters) == 1 else {"$and": processed_filters} + + async def list_collections(self) -> list[str]: + """Retrieve a list of all existing collection names.""" + + def _list(): + return [col.name for col in self.client.list_collections()] + + return await self._run_sync_in_executor(_list) + + async def create_collection(self, collection_name: str, **kwargs): + """Create a new collection with specified distance metrics and metadata.""" + + def _create(): + distance_metric = kwargs.get("distance_metric", "cosine") + metadata = kwargs.get("metadata", {}) + metadata["hnsw:space"] = distance_metric + return self.client.get_or_create_collection(name=collection_name, metadata=metadata) + + new_collection = await self._run_sync_in_executor(_create) + if collection_name == self.collection_name: + self.collection = new_collection + logger.info(f"Created collection {collection_name}") + + async def delete_collection(self, collection_name: str, **kwargs): + """Delete a specified collection from the database.""" + + def _delete(): + try: + self.client.delete_collection(name=collection_name) + return True + except Exception as e: + logger.warning(f"Failed to delete collection {collection_name}: {e}") + return False + + deleted = await self._run_sync_in_executor(_delete) + if deleted and collection_name == self.collection_name: + self.collection = None + logger.info(f"Deleted collection {collection_name}") + + async def copy_collection(self, collection_name: str, **kwargs): + """Copy all data from the current collection to a new collection.""" + + def _copy(): + source_data = self.collection.get(include=["documents", "metadatas", "embeddings"]) + if not source_data["ids"]: + logger.warning(f"Source collection {self.collection_name} is empty") + return + + target_collection = self.client.get_or_create_collection( + name=collection_name, + metadata={"hnsw:space": "cosine"}, + ) + target_collection.add( + ids=source_data["ids"], + documents=source_data["documents"], + metadatas=source_data["metadatas"], + embeddings=source_data["embeddings"], + ) + + await self._run_sync_in_executor(_copy) + logger.info(f"Copied collection {self.collection_name} to {collection_name}") + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Insert vector nodes into the current collection in batches.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + if not nodes: + return + + # Batch generate embeddings for nodes that need them + nodes_without_vectors = [node for node in nodes if node.vector is None] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + # Create a mapping for quick lookup + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes] + else: + nodes_to_insert = nodes + + batch_size = kwargs.get("batch_size", 100) + + def _insert_batch(batch_nodes: list[VectorNode]): + self.collection.add( + ids=[n.vector_id for n in batch_nodes], + documents=[n.content for n in batch_nodes], + embeddings=[n.vector for n in batch_nodes], + metadatas=[n.metadata for n in batch_nodes], + ) + + for i in range(0, len(nodes_to_insert), batch_size): + await self._run_sync_in_executor(_insert_batch, nodes_to_insert[i : i + batch_size]) + logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}") + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[VectorNode]: + """Search for the most similar vector nodes based on a text query.""" + query_vector = await self.get_embedding(query) + where_clause = self._generate_where_clause(filters) + include_embeddings = kwargs.get("include_embeddings", False) + + def _search(): + include: list = ["documents", "metadatas", "distances"] + if include_embeddings: + include.append("embeddings") + return self.collection.query( + query_embeddings=[query_vector], + n_results=limit, + where=where_clause, + include=include, + ) + + results = await self._run_sync_in_executor(_search) + nodes = self._parse_results(results, include_score=True) + + score_threshold = kwargs.get("score_threshold") + if score_threshold is not None: + nodes = [n for n in nodes if n.metadata.get("_score", 0) >= score_threshold] + return nodes + + async def delete(self, vector_ids: str | list[str], **kwargs): + """Delete specific vector nodes by their IDs.""" + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + if not vector_ids: + return + + def _delete(): + self.collection.delete(ids=vector_ids) + + await self._run_sync_in_executor(_delete) + logger.info(f"Deleted {len(vector_ids)} nodes from {self.collection_name}") + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Update existing vector nodes with new content or metadata.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + if not nodes: + return + + # Batch generate embeddings for nodes that need them + nodes_without_vectors = [node for node in nodes if node.vector is None and node.content] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + # Create a mapping for quick lookup + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes] + else: + nodes_to_update = nodes + + def _update(): + self.collection.upsert( + ids=[n.vector_id for n in nodes_to_update], + documents=[n.content for n in nodes_to_update], + embeddings=[n.vector for n in nodes_to_update if n.vector] or None, + metadatas=[n.metadata for n in nodes_to_update], + ) + + await self._run_sync_in_executor(_update) + logger.info(f"Updated {len(nodes_to_update)} nodes in {self.collection_name}") + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None: + """Fetch vector nodes by their IDs from the collection.""" + is_single = isinstance(vector_ids, str) + ids = [vector_ids] if is_single else vector_ids + + def _get(): + return self.collection.get(ids=ids, include=["documents", "metadatas", "embeddings"]) + + results = await self._run_sync_in_executor(_get) + nodes = self._parse_results(results) + return nodes[0] if is_single and nodes else (nodes if not is_single else None) + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + ) -> list[VectorNode]: + """List vector nodes matching optional metadata filters.""" + where_clause = self._generate_where_clause(filters) + + def _list(): + return self.collection.get( + where=where_clause, + limit=limit, + include=["documents", "metadatas", "embeddings"], + ) + + results = await self._run_sync_in_executor(_list) + return self._parse_results(results) + + async def count(self) -> int: + """Return the total number of vectors in the current collection.""" + return await self._run_sync_in_executor(self.collection.count) + + async def reset(self): + """Reset the current collection by clearing all its data.""" + logger.warning(f"Resetting collection {self.collection_name}...") + await self.delete_collection(self.collection_name) + + def _recreate(): + self.collection = self.client.get_or_create_collection( + name=self.collection_name, + metadata={"hnsw:space": "cosine"}, + ) + + await self._run_sync_in_executor(_recreate) + logger.info(f"Collection {self.collection_name} has been reset") + + async def close(self): + """Close the vector store and log the shutdown process.""" + logger.info(f"ChromaDB vector store for collection {self.collection_name} closed") diff --git a/reme_ai/core/vector_store/es_vector_store.py b/reme_ai/core/vector_store/es_vector_store.py new file mode 100644 index 00000000..06dcafa5 --- /dev/null +++ b/reme_ai/core/vector_store/es_vector_store.py @@ -0,0 +1,458 @@ +"""Elasticsearch vector store implementation for ReMe. + +This module provides an Elasticsearch-based vector store that implements the BaseVectorStore +interface for high-performance dense vector storage and retrieval. +""" + +from typing import Any + +from loguru import logger + +from .base_vector_store import BaseVectorStore +from ..context import C +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + +_ELASTICSEARCH_IMPORT_ERROR = None + +try: + from elasticsearch import AsyncElasticsearch + from elasticsearch.helpers import async_bulk +except ImportError as e: + _ELASTICSEARCH_IMPORT_ERROR = e + AsyncElasticsearch = None + async_bulk = None + + +@C.register_vector_store("es") +class ESVectorStore(BaseVectorStore): + """Elasticsearch-based vector store for dense vector storage and kNN search.""" + + def __init__( + self, + collection_name: str, + embedding_model: BaseEmbeddingModel, + hosts: str | list[str] | None = None, + basic_auth: tuple[str, str] | None = None, + cloud_id: str | None = None, + api_key: str | None = None, + verify_certs: bool = True, + headers: dict[str, str] | None = None, + **kwargs, + ): + """Initialize the Elasticsearch client and vector store configuration. + + Args: + collection_name: Name of the Elasticsearch index (converted to lowercase). + embedding_model: Model instance used to generate vector embeddings. + hosts: Connection host(s) for the Elasticsearch cluster. + basic_auth: Credentials for basic authentication. + cloud_id: Deployment ID for Elastic Cloud. + api_key: API key for authentication. + verify_certs: Enable or disable SSL certificate verification. + headers: Custom HTTP headers for requests. + **kwargs: Additional configuration passed to the base class. + """ + if _ELASTICSEARCH_IMPORT_ERROR is not None: + raise ImportError( + "Elasticsearch requires extra dependencies. Install with `pip install elasticsearch`", + ) from _ELASTICSEARCH_IMPORT_ERROR + + # Elasticsearch requires lowercase index names + collection_name = collection_name.lower() + + super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs) + + # Initialize AsyncElasticsearch client + self.client = AsyncElasticsearch( + hosts=hosts, + cloud_id=cloud_id, + api_key=api_key, + basic_auth=basic_auth, + verify_certs=verify_certs, + headers=headers or {}, + ) + + async def list_collections(self) -> list[str]: + """List all available index names in the Elasticsearch cluster.""" + aliases = await self.client.indices.get_alias() + return list(aliases.keys()) + + async def create_collection(self, collection_name: str, **kwargs): + """Create a new index with dense vector mappings for kNN search. + + Args: + collection_name: Name of the index to create. + **kwargs: Settings like dimensions, similarity, shards, and replicas. + """ + collection_name = collection_name.lower() + + if await self.client.indices.exists(index=collection_name): + return + + dimensions = kwargs.get("dimensions", self.embedding_model.dimensions) + similarity = kwargs.get("similarity", "cosine") + number_of_shards = kwargs.get("number_of_shards", 5) + number_of_replicas = kwargs.get("number_of_replicas", 1) + refresh_interval = kwargs.get("refresh_interval", "1s") + + index_settings = { + "settings": { + "index": { + "number_of_replicas": number_of_replicas, + "number_of_shards": number_of_shards, + "refresh_interval": refresh_interval, + }, + }, + "mappings": { + "properties": { + "vector_id": {"type": "keyword"}, + "content": {"type": "text"}, + "vector": { + "type": "dense_vector", + "dims": dimensions, + "index": True, + "similarity": similarity, + }, + "metadata": {"type": "object", "enabled": True}, + }, + }, + } + + if not await self.client.indices.exists(index=collection_name): + await self.client.indices.create(index=collection_name, body=index_settings) + logger.info(f"Created index {collection_name} with dimensions={dimensions}") + else: + logger.info(f"Index {collection_name} already exists") + + async def delete_collection(self, collection_name: str, **kwargs): + """Permanently delete an Elasticsearch index. + + Args: + collection_name: Name of the index to delete. + **kwargs: Additional parameters for the deletion request. + """ + collection_name = collection_name.lower() + + if await self.client.indices.exists(index=collection_name): + await self.client.indices.delete(index=collection_name) + logger.info(f"Deleted index {collection_name}") + else: + logger.warning(f"Index {collection_name} does not exist") + + async def copy_collection(self, collection_name: str, **kwargs): + """Reindex the current collection into a new index with identical mappings. + + Args: + collection_name: Name of the destination index. + **kwargs: Additional parameters for the reindexing process. + """ + collection_name = collection_name.lower() + + current_index = await self.client.indices.get(index=self.collection_name) + current_settings = current_index[self.collection_name] + + settings_to_copy = current_settings.get("settings", {}).copy() + if "index" in settings_to_copy: + index_settings = settings_to_copy["index"].copy() + internal_keys = [ + "uuid", + "creation_date", + "provided_name", + "version", + "store", + "routing", + "replication", + ] + for key in internal_keys: + index_settings.pop(key, None) + settings_to_copy["index"] = index_settings + + await self.client.indices.create( + index=collection_name, + body={ + "settings": settings_to_copy, + "mappings": current_settings.get("mappings", {}), + }, + ) + + await self.client.reindex( + body={ + "source": {"index": self.collection_name}, + "dest": {"index": collection_name}, + }, + ) + + logger.info(f"Copied collection {self.collection_name} to {collection_name}") + + async def insert(self, nodes: VectorNode | list[VectorNode], refresh: bool = True, **kwargs): + """Insert nodes into the index, generating embeddings if missing. + + Args: + nodes: Single or multiple VectorNode objects to index. + refresh: If True, makes the operation visible to search immediately. + **kwargs: Additional insertion options. + """ + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes_without_vectors = [node for node in nodes if node.vector is None] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes] + else: + nodes_to_insert = nodes + + actions = [] + for node in nodes_to_insert: + action = { + "_index": self.collection_name, + "_id": node.vector_id, + "_source": { + "vector_id": node.vector_id, + "content": node.content, + "vector": node.vector, + "metadata": node.metadata, + }, + } + actions.append(action) + + success, failed = await async_bulk(self.client, actions, raise_on_error=False) + + if failed: + logger.warning(f"Failed to insert {len(failed)} documents") + + logger.info(f"Inserted {success} documents into {self.collection_name}") + + if refresh: + await self.client.indices.refresh(index=self.collection_name) + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[VectorNode]: + """Perform a kNN similarity search based on a text query. + + Args: + query: The text to search for. + limit: Maximum number of nearest neighbors to return. + filters: Metadata filters for exact match or 'IN' operations. + **kwargs: Search parameters like num_candidates or score_threshold. + + Returns: + List of VectorNode objects ordered by similarity. + """ + query_vector = await self.get_embedding(query) + num_candidates = kwargs.get("num_candidates", limit * 2) + + search_query: dict = { + "knn": { + "field": "vector", + "query_vector": query_vector, + "k": limit, + "num_candidates": num_candidates, + }, + "size": limit, + } + + if filters: + filter_conditions = [] + for key, value in filters.items(): + if isinstance(value, list): + filter_conditions.append({"terms": {f"metadata.{key}": value}}) + else: + filter_conditions.append({"term": {f"metadata.{key}": value}}) + search_query["knn"]["filter"] = {"bool": {"must": filter_conditions}} + + response = await self.client.search(index=self.collection_name, body=search_query) + + results = [] + for hit in response["hits"]["hits"]: + source = hit["_source"] + node = VectorNode( + vector_id=source.get("vector_id", hit["_id"]), + content=source.get("content", ""), + vector=source.get("vector"), + metadata=source.get("metadata", {}), + ) + node.metadata["_score"] = hit["_score"] + results.append(node) + + return results + + async def delete(self, vector_ids: str | list[str], refresh: bool = True, **kwargs): + """Delete specific vectors from the index by their IDs. + + Args: + vector_ids: Single ID or list of IDs to remove. + refresh: If True, refreshes the index after deletion. + **kwargs: Additional deletion parameters. + """ + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + + actions = [] + for vector_id in vector_ids: + actions.append( + { + "_op_type": "delete", + "_index": self.collection_name, + "_id": vector_id, + }, + ) + + success, failed = await async_bulk( + self.client, + actions, + raise_on_error=False, + raise_on_exception=False, + ) + + if failed: + logger.warning(f"Failed to delete {len(failed)} documents") + + logger.info(f"Deleted {success} documents from {self.collection_name}") + + if refresh: + await self.client.indices.refresh(index=self.collection_name) + + async def update(self, nodes: VectorNode | list[VectorNode], refresh: bool = True, **kwargs): + """Update existing documents with new content or metadata. + + Args: + nodes: Single or multiple VectorNode objects with updated data. + refresh: If True, refreshes the index after update. + **kwargs: Additional update parameters. + """ + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes_without_vectors = [node for node in nodes if node.vector is None and node.content] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes] + else: + nodes_to_update = nodes + + actions = [] + for node in nodes_to_update: + doc = { + "vector_id": node.vector_id, + "content": node.content, + "metadata": node.metadata, + } + if node.vector is not None: + doc["vector"] = node.vector + + actions.append( + { + "_op_type": "update", + "_index": self.collection_name, + "_id": node.vector_id, + "doc": doc, + }, + ) + + success, failed = await async_bulk( + self.client, + actions, + raise_on_error=False, + raise_on_exception=False, + ) + + if failed: + logger.warning(f"Failed to update {len(failed)} documents") + + logger.info(f"Updated {success} documents in {self.collection_name}") + + if refresh: + await self.client.indices.refresh(index=self.collection_name) + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]: + """Fetch documents by their IDs from the current index. + + Args: + vector_ids: Single ID or list of IDs to retrieve. + + Returns: + A single VectorNode or a list of VectorNodes. + """ + single_result = isinstance(vector_ids, str) + if single_result: + vector_ids = [vector_ids] + + response = await self.client.mget( + index=self.collection_name, + body={"ids": vector_ids}, + ) + + results = [] + for doc in response["docs"]: + if doc.get("found"): + source = doc["_source"] + node = VectorNode( + vector_id=source.get("vector_id", doc["_id"]), + content=source.get("content", ""), + vector=source.get("vector"), + metadata=source.get("metadata", {}), + ) + results.append(node) + else: + logger.warning(f"Document with ID {doc['_id']} not found") + + return results[0] if single_result and results else results + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + ) -> list[VectorNode]: + """Retrieve a list of nodes filtered by metadata or limit. + + Args: + filters: Optional metadata filtering criteria. + limit: Maximum number of nodes to return. + + Returns: + A list of matching VectorNode objects. + """ + query: dict[str, Any] = {"query": {"match_all": {}}} + + if filters: + filter_conditions = [] + for key, value in filters.items(): + if isinstance(value, list): + filter_conditions.append({"terms": {f"metadata.{key}": value}}) + else: + filter_conditions.append({"term": {f"metadata.{key}": value}}) + query["query"] = {"bool": {"must": filter_conditions}} + + if limit: + query["size"] = limit + else: + query["size"] = 10000 + + response = await self.client.search(index=self.collection_name, body=query) + + results = [] + for hit in response["hits"]["hits"]: + source = hit["_source"] + node = VectorNode( + vector_id=source.get("vector_id", hit["_id"]), + content=source.get("content", ""), + vector=source.get("vector"), + metadata=source.get("metadata", {}), + ) + results.append(node) + + return results + + async def close(self): + """Terminate the Elasticsearch client session and release resources.""" + await self.client.close() + logger.info("Elasticsearch client connection closed") diff --git a/reme_ai/core/vector_store/local_vector_store.py b/reme_ai/core/vector_store/local_vector_store.py new file mode 100644 index 00000000..25beef8d --- /dev/null +++ b/reme_ai/core/vector_store/local_vector_store.py @@ -0,0 +1,281 @@ +"""Local file system vector store implementation for ReMe.""" + +import json +from pathlib import Path + +from loguru import logger + +from .base_vector_store import BaseVectorStore +from ..context import C +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + + +@C.register_vector_store("local") +class LocalVectorStore(BaseVectorStore): + """Local file system-based vector store using JSON files and manual cosine similarity.""" + + def __init__( + self, + collection_name: str, + embedding_model: BaseEmbeddingModel, + root_path: str = "./local_vector_store", + **kwargs, + ): + """Initialize the local vector store with a root path and collection name.""" + super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs) + self.root_path = Path(root_path) + self.collection_path = self.root_path / collection_name + self.root_path.mkdir(parents=True, exist_ok=True) + + def _get_collection_path(self, collection_name: str) -> Path: + """Get the file system path for a specific collection.""" + return self.root_path / collection_name + + def _get_node_file_path(self, vector_id: str, collection_name: str | None = None) -> Path: + """Get the JSON file path for a specific vector node.""" + col_path = self._get_collection_path(collection_name or self.collection_name) + return col_path / f"{vector_id}.json" + + def _save_node(self, node: VectorNode, collection_name: str | None = None): + """Save a vector node to a JSON file on disk.""" + file_path = self._get_node_file_path(node.vector_id, collection_name) + file_path.parent.mkdir(parents=True, exist_ok=True) + + with open(file_path, "w", encoding="utf-8") as f: + json.dump(node.model_dump(), f, ensure_ascii=False, indent=2) + + def _load_node(self, vector_id: str, collection_name: str | None = None) -> VectorNode | None: + """Load a vector node from a JSON file.""" + file_path = self._get_node_file_path(vector_id, collection_name) + + if not file_path.exists(): + return None + + with open(file_path, "r", encoding="utf-8") as f: + data = json.load(f) + return VectorNode(**data) + + def _load_all_nodes(self, collection_name: str | None = None) -> list[VectorNode]: + """Load all vector nodes existing in a collection.""" + col_path = self._get_collection_path(collection_name or self.collection_name) + + if not col_path.exists(): + return [] + + nodes = [] + for file_path in col_path.glob("*.json"): + try: + with open(file_path, "r", encoding="utf-8") as f: + data = json.load(f) + nodes.append(VectorNode(**data)) + except Exception as e: + logger.warning(f"Failed to load node from {file_path}: {e}") + + return nodes + + @staticmethod + def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float: + """Calculate the cosine similarity between two numeric vectors.""" + if len(vec1) != len(vec2): + raise ValueError(f"Vectors must have same length: {len(vec1)} != {len(vec2)}") + + dot_product = sum(a * b for a, b in zip(vec1, vec2)) + magnitude1 = sum(a * a for a in vec1) ** 0.5 + magnitude2 = sum(b * b for b in vec2) ** 0.5 + + if magnitude1 == 0 or magnitude2 == 0: + return 0.0 + + return dot_product / (magnitude1 * magnitude2) + + @staticmethod + def _match_filters(node: VectorNode, filters: dict | None) -> bool: + """Check if a vector node matches the provided metadata filters.""" + if not filters: + return True + + for key, value in filters.items(): + node_value = node.metadata.get(key) + + if isinstance(value, list): + if node_value not in value: + return False + else: + if node_value != value: + return False + + return True + + async def list_collections(self) -> list[str]: + """List all collection directories in the root path.""" + if not self.root_path.exists(): + return [] + + return [d.name for d in self.root_path.iterdir() if d.is_dir() and not d.name.startswith(".")] + + async def create_collection(self, collection_name: str, **kwargs): + """Create a new collection directory.""" + col_path = self._get_collection_path(collection_name) + col_path.mkdir(parents=True, exist_ok=True) + logger.info(f"Created collection {collection_name} at {col_path}") + + async def delete_collection(self, collection_name: str, **kwargs): + """Delete a collection directory and all its JSON files.""" + col_path = self._get_collection_path(collection_name) + + if not col_path.exists(): + logger.warning(f"Collection {collection_name} does not exist") + return + + for file_path in col_path.glob("*.json"): + file_path.unlink() + + col_path.rmdir() + logger.info(f"Deleted collection {collection_name}") + + async def copy_collection(self, collection_name: str, **kwargs): + """Copy all nodes from the current collection to a new one.""" + source_path = self._get_collection_path(self.collection_name) + target_path = self._get_collection_path(collection_name) + + if not source_path.exists(): + logger.warning(f"Source collection {self.collection_name} does not exist") + return + + target_path.mkdir(parents=True, exist_ok=True) + + for file_path in source_path.glob("*.json"): + target_file = target_path / file_path.name + target_file.write_text(file_path.read_text(encoding="utf-8"), encoding="utf-8") + + logger.info(f"Copied collection {self.collection_name} to {collection_name}") + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Insert vector nodes into the local store, generating embeddings if necessary.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes_without_vectors = [node for node in nodes if node.vector is None] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes] + else: + nodes_to_insert = nodes + + for node in nodes_to_insert: + self._save_node(node) + + logger.info(f"Inserted {len(nodes_to_insert)} nodes into {self.collection_name}") + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[VectorNode]: + """Search for nodes similar to the query using brute-force cosine similarity.""" + query_vector = await self.get_embedding(query) + all_nodes = self._load_all_nodes() + filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)] + + scored_nodes = [] + for node in filtered_nodes: + if node.vector is None: + logger.warning(f"Node {node.vector_id} has no vector, skipping") + continue + + try: + score = self._cosine_similarity(query_vector, node.vector) + scored_nodes.append((node, score)) + except ValueError as e: + logger.warning(f"Failed to calculate similarity for node {node.vector_id}: {e}") + + scored_nodes.sort(key=lambda x: x[1], reverse=True) + + score_threshold = kwargs.get("score_threshold") + if score_threshold is not None: + scored_nodes = [(node, score) for node, score in scored_nodes if score >= score_threshold] + + scored_nodes = scored_nodes[:limit] + results = [] + for node, score in scored_nodes: + node.metadata["_score"] = score + results.append(node) + + return results + + async def delete(self, vector_ids: str | list[str], **kwargs): + """Delete specific vector nodes by their IDs.""" + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + + deleted_count = 0 + for vector_id in vector_ids: + file_path = self._get_node_file_path(vector_id) + if file_path.exists(): + file_path.unlink() + deleted_count += 1 + else: + logger.warning(f"Node {vector_id} does not exist") + + logger.info(f"Deleted {deleted_count} nodes from {self.collection_name}") + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Update existing vector nodes with new data or embeddings.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes_without_vectors = [node for node in nodes if node.vector is None and node.content] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes] + else: + nodes_to_update = nodes + + updated_count = 0 + for node in nodes_to_update: + file_path = self._get_node_file_path(node.vector_id) + if file_path.exists(): + self._save_node(node) + updated_count += 1 + else: + logger.warning(f"Node {node.vector_id} does not exist, skipping update") + + logger.info(f"Updated {updated_count} nodes in {self.collection_name}") + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]: + """Retrieve one or more vector nodes by their unique IDs.""" + is_single = isinstance(vector_ids, str) + ids = [vector_ids] if is_single else vector_ids + + results = [] + for vector_id in ids: + node = self._load_node(vector_id) + if node: + results.append(node) + else: + logger.warning(f"Node {vector_id} not found") + + return results[0] if is_single and results else results + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + ) -> list[VectorNode]: + """List vector nodes in the collection with optional filtering and limits.""" + all_nodes = self._load_all_nodes() + filtered_nodes = [node for node in all_nodes if self._match_filters(node, filters)] + + if limit is not None: + filtered_nodes = filtered_nodes[:limit] + + return filtered_nodes + + async def close(self): + """Close the vector store (no-op for local file system).""" + logger.info("Local vector store closed") diff --git a/reme_ai/core/vector_store/pgvector_store.py b/reme_ai/core/vector_store/pgvector_store.py new file mode 100644 index 00000000..38c5667a --- /dev/null +++ b/reme_ai/core/vector_store/pgvector_store.py @@ -0,0 +1,533 @@ +"""PostgreSQL pgvector implementation for vector storage and retrieval.""" + +import json +from typing import Any + +from loguru import logger + +from .base_vector_store import BaseVectorStore +from ..context import C +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + +_ASYNCPG_IMPORT_ERROR = None + +try: + import asyncpg + from asyncpg import Pool +except ImportError as e: + _ASYNCPG_IMPORT_ERROR = e + asyncpg = None + Pool = None + + +@C.register_vector_store("pgvector") +class PGVectorStore(BaseVectorStore): + """Vector store implementation using PostgreSQL and pgvector for efficient similarity search.""" + + def __init__( + self, + collection_name: str, + embedding_model: BaseEmbeddingModel, + host: str = "localhost", + port: int = 5432, + database: str = "postgres", + user: str = "postgres", + password: str = "", + min_size: int = 1, + max_size: int = 10, + dsn: str | None = None, + use_hnsw: bool = True, + use_diskann: bool = False, + **kwargs, + ): + """Initialize the PGVector store with connection parameters and index settings.""" + if _ASYNCPG_IMPORT_ERROR is not None: + raise ImportError( + "PGVector requires extra dependencies. Install with `pip install asyncpg pgvector`", + ) from _ASYNCPG_IMPORT_ERROR + + super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs) + + self.dsn = dsn + self.host = host + self.port = port + self.database = database + self.user = user + self.password = password + self.min_size = min_size + self.max_size = max_size + self.use_hnsw = use_hnsw + self.use_diskann = use_diskann + self._pool: Pool | None = None + self.embedding_model_dims = embedding_model.dimensions + + async def _get_pool(self) -> Pool: + """Create or return the existing asyncpg connection pool.""" + if self._pool is None: + if self.dsn: + self._pool = await asyncpg.create_pool( + dsn=self.dsn, + min_size=self.min_size, + max_size=self.max_size, + ) + else: + self._pool = await asyncpg.create_pool( + host=self.host, + port=self.port, + database=self.database, + user=self.user, + password=self.password, + min_size=self.min_size, + max_size=self.max_size, + ) + + async with self._pool.acquire() as conn: + await conn.execute("CREATE EXTENSION IF NOT EXISTS vector") + + logger.info(f"PGVector connection pool created for database {self.database}") + + return self._pool + + async def _ensure_collection_exists(self): + """Check if the collection table exists and create it if missing.""" + collections = await self.list_collections() + if self.collection_name not in collections: + await self.create_collection(self.collection_name) + + async def list_collections(self) -> list[str]: + """List all available table names in the current database.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + rows = await conn.fetch( + "SELECT table_name FROM information_schema.tables WHERE table_schema = 'public'", + ) + return [row["table_name"] for row in rows] + + async def create_collection(self, collection_name: str, **kwargs): + """Create a new PostgreSQL table with vector support and appropriate indexing.""" + pool = await self._get_pool() + dimensions = kwargs.get("dimensions", self.embedding_model_dims) + + async with pool.acquire() as conn: + await conn.execute( + f""" + CREATE TABLE IF NOT EXISTS {collection_name} ( + id TEXT PRIMARY KEY, + content TEXT, + vector vector({dimensions}), + metadata JSONB + ) + """, + ) + + if self.use_diskann and dimensions < 2000: + result = await conn.fetchval( + "SELECT 1 FROM pg_extension WHERE extname = 'vectorscale'", + ) + if result: + await conn.execute( + f""" + CREATE INDEX IF NOT EXISTS {collection_name}_diskann_idx + ON {collection_name} + USING diskann (vector) + """, + ) + logger.info(f"Created DiskANN index for collection {collection_name}") + else: + logger.warning("vectorscale extension not available, skipping DiskANN index") + elif self.use_hnsw: + await conn.execute( + f""" + CREATE INDEX IF NOT EXISTS {collection_name}_hnsw_idx + ON {collection_name} + USING hnsw (vector vector_cosine_ops) + """, + ) + logger.info(f"Created HNSW index for collection {collection_name}") + + logger.info(f"Created collection {collection_name} with dimensions={dimensions}") + + async def delete_collection(self, collection_name: str, **kwargs): + """Remove the specified collection table from the database.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {collection_name}") + logger.info(f"Deleted collection {collection_name}") + + async def copy_collection(self, collection_name: str, **kwargs): + """Duplicate the structure and content of the current collection to a new table.""" + pool = await self._get_pool() + + async with pool.acquire() as conn: + columns = await conn.fetch( + """ + SELECT column_name, data_type, udt_name + FROM information_schema.columns + WHERE table_name = $1 AND table_schema = 'public' + """, + self.collection_name, + ) + + if not columns: + raise ValueError(f"Source collection {self.collection_name} does not exist") + + await conn.execute(f"CREATE TABLE {collection_name} AS TABLE {self.collection_name}") + await conn.execute(f"ALTER TABLE {collection_name} ADD PRIMARY KEY (id)") + + if self.use_hnsw: + await conn.execute( + f""" + CREATE INDEX IF NOT EXISTS {collection_name}_hnsw_idx + ON {collection_name} + USING hnsw (vector vector_cosine_ops) + """, + ) + + logger.info(f"Copied collection {self.collection_name} to {collection_name}") + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Insert or upsert vector nodes into the PostgreSQL collection.""" + await self._ensure_collection_exists() + + if isinstance(nodes, VectorNode): + nodes = [nodes] + + if not nodes: + return + + nodes_without_vectors = [node for node in nodes if node.vector is None] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes] + else: + nodes_to_insert = nodes + + pool = await self._get_pool() + data = [ + ( + node.vector_id, + node.content, + f"[{','.join(map(str, node.vector))}]", + json.dumps(node.metadata), + ) + for node in nodes_to_insert + ] + + async with pool.acquire() as conn: + on_conflict = kwargs.get("on_conflict", "update") + + if on_conflict == "update": + await conn.executemany( + f""" + INSERT INTO {self.collection_name} (id, content, vector, metadata) + VALUES ($1, $2, $3::vector, $4::jsonb) + ON CONFLICT (id) DO UPDATE SET + content = EXCLUDED.content, + vector = EXCLUDED.vector, + metadata = EXCLUDED.metadata + """, + data, + ) + elif on_conflict == "ignore": + await conn.executemany( + f""" + INSERT INTO {self.collection_name} (id, content, vector, metadata) + VALUES ($1, $2, $3::vector, $4::jsonb) + ON CONFLICT (id) DO NOTHING + """, + data, + ) + else: + await conn.executemany( + f""" + INSERT INTO {self.collection_name} (id, content, vector, metadata) + VALUES ($1, $2, $3::vector, $4::jsonb) + """, + data, + ) + + logger.info(f"Inserted {len(nodes_to_insert)} documents into {self.collection_name}") + + @staticmethod + def _build_filter_clause(filters: dict | None) -> tuple[str, list]: + """Generate an SQL WHERE clause and parameter list from a filter dictionary.""" + if not filters: + return "", [] + + conditions = [] + params = [] + param_idx = 1 + + for key, value in filters.items(): + if isinstance(value, list): + placeholders = ", ".join([f"${param_idx + i}" for i in range(len(value))]) + conditions.append(f"metadata->>'{key}' IN ({placeholders})") + params.extend([str(v) for v in value]) + param_idx += len(value) + else: + conditions.append(f"metadata->>'{key}' = ${param_idx}") + params.append(str(value)) + param_idx += 1 + + filter_clause = "WHERE " + " AND ".join(conditions) if conditions else "" + return filter_clause, params + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs, + ) -> list[VectorNode]: + """Perform vector similarity search with optional metadata filtering.""" + await self._ensure_collection_exists() + + query_vector = await self.get_embedding(query) + vector_str = f"[{','.join(map(str, query_vector))}]" + pool = await self._get_pool() + + filter_clause, filter_params = self._build_filter_clause(filters) + + if filter_clause: + for i in range(len(filter_params)): + old_idx = i + 1 + new_idx = i + 2 + filter_clause = filter_clause.replace(f"${old_idx}", f"${new_idx}", 1) + + async with pool.acquire() as conn: + sql = f""" + SELECT id, content, vector, metadata, vector <=> $1::vector AS distance + FROM {self.collection_name} + {filter_clause} + ORDER BY distance + LIMIT ${len(filter_params) + 2} + """ + rows = await conn.fetch(sql, vector_str, *filter_params, limit) + + results = [] + score_threshold = kwargs.get("score_threshold") + + for row in rows: + distance = row["distance"] + if score_threshold is not None and distance > score_threshold: + continue + + vector_data = None + if row["vector"]: + vector_str_raw = str(row["vector"]) + if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"): + vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")] + + metadata = row["metadata"] if row["metadata"] else {} + if isinstance(metadata, str): + metadata = json.loads(metadata) + + metadata["_score"] = 1 - distance + metadata["_distance"] = distance + + node = VectorNode( + vector_id=row["id"], + content=row["content"] or "", + vector=vector_data, + metadata=metadata, + ) + results.append(node) + + return results + + async def delete(self, vector_ids: str | list[str], **kwargs): + """Remove specific vector records from the collection by their IDs.""" + await self._ensure_collection_exists() + + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + + if not vector_ids: + return + + pool = await self._get_pool() + async with pool.acquire() as conn: + placeholders = ", ".join([f"${i + 1}" for i in range(len(vector_ids))]) + await conn.execute( + f"DELETE FROM {self.collection_name} WHERE id IN ({placeholders})", + *vector_ids, + ) + + logger.info(f"Deleted {len(vector_ids)} documents from {self.collection_name}") + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Update existing vector nodes with new content, embeddings, or metadata.""" + await self._ensure_collection_exists() + + if isinstance(nodes, VectorNode): + nodes = [nodes] + + if not nodes: + return + + nodes_without_vectors = [node for node in nodes if node.vector is None and node.content] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes] + else: + nodes_to_update = nodes + + pool = await self._get_pool() + async with pool.acquire() as conn: + for node in nodes_to_update: + update_fields = [] + params = [] + idx = 1 + + if node.content: + update_fields.append(f"content = ${idx}") + params.append(node.content) + idx += 1 + + if node.vector: + vector_str = f"[{','.join(map(str, node.vector))}]" + update_fields.append(f"vector = ${idx}::vector") + params.append(vector_str) + idx += 1 + + if node.metadata: + update_fields.append(f"metadata = ${idx}::jsonb") + params.append(json.dumps(node.metadata)) + idx += 1 + + if update_fields: + params.append(node.vector_id) + await conn.execute( + f"UPDATE {self.collection_name} SET {', '.join(update_fields)} WHERE id = ${idx}", + *params, + ) + + logger.info(f"Updated {len(nodes_to_update)} documents in {self.collection_name}") + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None: + """Retrieve vector nodes by their unique identifiers.""" + await self._ensure_collection_exists() + + single_result = isinstance(vector_ids, str) + if single_result: + vector_ids = [vector_ids] + + if not vector_ids: + return [] if not single_result else None + + pool = await self._get_pool() + async with pool.acquire() as conn: + placeholders = ", ".join([f"${i + 1}" for i in range(len(vector_ids))]) + rows = await conn.fetch( + f"SELECT id, content, vector, metadata FROM {self.collection_name} WHERE id IN ({placeholders})", + *vector_ids, + ) + + results = [] + for row in rows: + vector_data = None + if row["vector"]: + vector_str_raw = str(row["vector"]) + if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"): + vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")] + + metadata = row["metadata"] if row["metadata"] else {} + if isinstance(metadata, str): + metadata = json.loads(metadata) + + results.append( + VectorNode( + vector_id=row["id"], + content=row["content"] or "", + vector=vector_data, + metadata=metadata, + ), + ) + + if single_result: + return results[0] if results else None + return results + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + ) -> list[VectorNode]: + """Return a list of vector nodes matching the provided filters and limit.""" + await self._ensure_collection_exists() + + pool = await self._get_pool() + filter_clause, filter_params = self._build_filter_clause(filters) + + limit_clause = "" + if limit: + limit_clause = f"LIMIT ${len(filter_params) + 1}" + filter_params.append(limit) + + async with pool.acquire() as conn: + sql = f""" + SELECT id, content, vector, metadata + FROM {self.collection_name} + {filter_clause} + {limit_clause} + """ + rows = await conn.fetch(sql, *filter_params) + + results = [] + for row in rows: + vector_data = None + if row["vector"]: + vector_str_raw = str(row["vector"]) + if vector_str_raw.startswith("[") and vector_str_raw.endswith("]"): + vector_data = [float(x) for x in vector_str_raw[1:-1].split(",")] + + metadata = row["metadata"] if row["metadata"] else {} + if isinstance(metadata, str): + metadata = json.loads(metadata) + + results.append( + VectorNode( + vector_id=row["id"], + content=row["content"] or "", + vector=vector_data, + metadata=metadata, + ), + ) + + return results + + async def collection_info(self) -> dict[str, Any]: + """Fetch metadata including record count and disk usage for the collection.""" + pool = await self._get_pool() + + async with pool.acquire() as conn: + row = await conn.fetchrow( + f""" + SELECT + '{self.collection_name}' as name, + (SELECT COUNT(*) FROM {self.collection_name}) as row_count, + pg_size_pretty(pg_total_relation_size('{self.collection_name}')) as total_size + """, + ) + + return { + "name": row["name"], + "count": row["row_count"], + "size": row["total_size"], + } + + async def reset(self): + """Purge all data by dropping and recreating the collection table.""" + logger.warning(f"Resetting collection {self.collection_name}...") + await self.delete_collection(self.collection_name) + await self.create_collection(self.collection_name) + + async def close(self): + """Terminate the database connection pool and release associated resources.""" + if self._pool is not None: + await self._pool.close() + self._pool = None + logger.info("PGVector connection pool closed") diff --git a/reme_ai/core/vector_store/qdrant_vector_store.py b/reme_ai/core/vector_store/qdrant_vector_store.py new file mode 100644 index 00000000..97eb5b61 --- /dev/null +++ b/reme_ai/core/vector_store/qdrant_vector_store.py @@ -0,0 +1,444 @@ +"""Qdrant vector store implementation for the ReMe project.""" + +from typing import Any + +from loguru import logger + +from .base_vector_store import BaseVectorStore +from ..context import C +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + +_QDRANT_IMPORT_ERROR = None + +try: + from qdrant_client import AsyncQdrantClient + from qdrant_client.models import ( + Distance, + FieldCondition, + Filter, + MatchValue, + PointIdsList, + PointStruct, + Range, + VectorParams, + ) +except ImportError as e: + _QDRANT_IMPORT_ERROR = e + AsyncQdrantClient = None + Distance = None + FieldCondition = None + Filter = None + MatchValue = None + PointIdsList = None + PointStruct = None + Range = None + VectorParams = None + + +@C.register_vector_store("qdrant") +class QdrantVectorStore(BaseVectorStore): + """Vector store implementation using Qdrant for dense vector search.""" + + def __init__( + self, + collection_name: str, + embedding_model: BaseEmbeddingModel, + host: str | None = None, + port: int = 6333, + path: str | None = None, + url: str | None = None, + api_key: str | None = None, + https: bool | None = None, + grpc_port: int = 6334, + prefer_grpc: bool = False, + distance: str = "cosine", + on_disk: bool = False, + **kwargs: Any, + ): + """Initialize the Qdrant client and collection configuration. + + Args: + collection_name: Name of the collection. + embedding_model: Model used for generating vector embeddings. + host: Server host address. + port: HTTP port for the server. + path: Local storage path for on-disk/in-memory mode. + url: Full connection URL. + api_key: Authentication key for Qdrant Cloud. + https: Use secure connection if True. + grpc_port: gRPC interface port. + prefer_grpc: Use gRPC instead of HTTP if True. + distance: Metric for similarity (cosine, euclid, dot). + on_disk: Enable persistent storage for vectors. + **kwargs: Additional client configuration. + """ + if _QDRANT_IMPORT_ERROR is not None: + raise ImportError( + "Qdrant requires extra dependencies. Install with `pip install qdrant-client`", + ) from _QDRANT_IMPORT_ERROR + + super().__init__(collection_name=collection_name, embedding_model=embedding_model, **kwargs) + + self.client = AsyncQdrantClient( + host=host, + port=port, + path=path, + url=url, + api_key=api_key, + https=https, + grpc_port=grpc_port, + prefer_grpc=prefer_grpc, + **kwargs, + ) + + self.is_local = path is not None + distance_map = { + "cosine": Distance.COSINE, + "euclid": Distance.EUCLID, + "dot": Distance.DOT, + } + self.distance = distance_map.get(distance.lower(), Distance.COSINE) + self.on_disk = on_disk + + async def list_collections(self) -> list[str]: + """Retrieve names of all existing collections in the Qdrant instance.""" + collections = await self.client.get_collections() + return [collection.name for collection in collections.collections] + + async def create_collection(self, collection_name: str, **kwargs: Any): + """Create a new collection with the specified vector configuration. + + Args: + collection_name: Name of the collection to create. + **kwargs: Overrides for dimensions, distance, or on_disk settings. + """ + collections = await self.list_collections() + if collection_name in collections: + logger.info(f"Collection {collection_name} already exists") + return + + dimensions = kwargs.get("dimensions", self.embedding_model.dimensions) + distance = kwargs.get("distance", self.distance) + on_disk = kwargs.get("on_disk", self.on_disk) + + await self.client.create_collection( + collection_name=collection_name, + vectors_config=VectorParams( + size=dimensions, + distance=distance, + on_disk=on_disk, + ), + ) + + logger.info(f"Created collection {collection_name} with dimensions={dimensions}") + + if not self.is_local: + await self._create_payload_indexes(collection_name) + + async def _create_payload_indexes(self, collection_name: str): + """Create keyword indexes for common metadata fields to optimize filtering.""" + common_fields = ["user_id", "agent_id", "run_id", "actor_id", "source"] + + for field in common_fields: + try: + await self.client.create_payload_index( + collection_name=collection_name, + field_name=field, + field_schema="keyword", + ) + logger.debug(f"Created index for {field} in collection {collection_name}") + except Exception as e: + logger.debug(f"Index for {field} might already exist: {e}") + + async def delete_collection(self, collection_name: str, **kwargs: Any): + """Permanently remove a collection from the Qdrant instance.""" + collections = await self.list_collections() + if collection_name in collections: + await self.client.delete_collection(collection_name=collection_name) + logger.info(f"Deleted collection {collection_name}") + else: + logger.warning(f"Collection {collection_name} does not exist") + + async def copy_collection(self, collection_name: str, **kwargs: Any): + """Duplicate an existing collection to a new one including all data.""" + collection_info = await self.client.get_collection(collection_name=self.collection_name) + + await self.client.create_collection( + collection_name=collection_name, + vectors_config=collection_info.config.params.vectors, + ) + + offset = None + batch_size = 100 + + while True: + records, next_offset = await self.client.scroll( + collection_name=self.collection_name, + limit=batch_size, + offset=offset, + with_payload=True, + with_vectors=True, + ) + + if not records: + break + + points = [ + PointStruct( + id=record.id, + vector=record.vector, + payload=record.payload, + ) + for record in records + ] + + await self.client.upsert( + collection_name=collection_name, + points=points, + ) + + offset = next_offset + if offset is None: + break + + logger.info(f"Copied collection {self.collection_name} to {collection_name}") + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs: Any): + """Insert vector nodes into the collection, generating embeddings as needed.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes_without_vectors = [node for node in nodes if node.vector is None] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_insert = [vector_map.get(n.vector_id, n) if n.vector is None else n for n in nodes] + else: + nodes_to_insert = nodes + + points = [] + for node in nodes_to_insert: + try: + point_id = int(node.vector_id) + except ValueError: + point_id = abs(hash(node.vector_id)) % (10**18) + + point = PointStruct( + id=point_id, + vector=node.vector, + payload={ + "vector_id": node.vector_id, + "content": node.content, + "metadata": node.metadata, + }, + ) + points.append(point) + + wait = kwargs.get("wait", True) + await self.client.upsert( + collection_name=self.collection_name, + points=points, + wait=wait, + ) + + logger.info(f"Inserted {len(points)} documents into {self.collection_name}") + + @staticmethod + def _create_filter(filters: dict) -> Filter | None: + """Convert a dictionary of filter conditions into a Qdrant Filter object.""" + if not filters: + return None + + conditions = [] + for key, value in filters.items(): + if isinstance(value, dict) and ("gte" in value or "lte" in value): + range_params = {} + if "gte" in value: + range_params["gte"] = value["gte"] + if "lte" in value: + range_params["lte"] = value["lte"] + conditions.append( + FieldCondition( + key=f"metadata.{key}", + range=Range(**range_params), + ), + ) + elif isinstance(value, list): + conditions.append( + FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value[0])), + ) + else: + conditions.append( + FieldCondition(key=f"metadata.{key}", match=MatchValue(value=value)), + ) + + return Filter(must=conditions) if conditions else None + + async def search( + self, + query: str, + limit: int = 5, + filters: dict | None = None, + **kwargs: Any, + ) -> list[VectorNode]: + """Search for the most similar vectors based on a text query.""" + query_vector = await self.get_embedding(query) + query_filter = self._create_filter(filters) if filters else None + score_threshold = kwargs.get("score_threshold", None) + + results = await self.client.query_points( + collection_name=self.collection_name, + query=query_vector, + query_filter=query_filter, + limit=limit, + score_threshold=score_threshold, + ) + + nodes = [] + for point in results.points: + payload = point.payload or {} + node = VectorNode( + vector_id=payload.get("vector_id", str(point.id)), + content=payload.get("content", ""), + vector=point.vector if hasattr(point, "vector") else None, + metadata=payload.get("metadata", {}), + ) + node.metadata["_score"] = point.score + nodes.append(node) + + return nodes + + async def delete(self, vector_ids: str | list[str], **kwargs: Any): + """Delete specific vectors from the collection using their IDs.""" + if isinstance(vector_ids, str): + vector_ids = [vector_ids] + + point_ids = [] + for vector_id in vector_ids: + try: + point_id = int(vector_id) + except ValueError: + point_id = abs(hash(vector_id)) % (10**18) + point_ids.append(point_id) + + wait = kwargs.get("wait", True) + await self.client.delete( + collection_name=self.collection_name, + points_selector=PointIdsList(points=point_ids), + wait=wait, + ) + + logger.info(f"Deleted {len(point_ids)} documents from {self.collection_name}") + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs: Any): + """Update existing vector nodes with new content or metadata.""" + if isinstance(nodes, VectorNode): + nodes = [nodes] + + nodes_without_vectors = [node for node in nodes if node.vector is None and node.content] + if nodes_without_vectors: + nodes_with_vectors = await self.get_node_embeddings(nodes_without_vectors) + vector_map = {n.vector_id: n for n in nodes_with_vectors} + nodes_to_update = [vector_map.get(n.vector_id, n) if n.vector is None and n.content else n for n in nodes] + else: + nodes_to_update = nodes + + points = [] + for node in nodes_to_update: + try: + point_id = int(node.vector_id) + except ValueError: + point_id = abs(hash(node.vector_id)) % (10**18) + + point = PointStruct( + id=point_id, + vector=node.vector, + payload={ + "vector_id": node.vector_id, + "content": node.content, + "metadata": node.metadata, + }, + ) + points.append(point) + + wait = kwargs.get("wait", True) + await self.client.upsert( + collection_name=self.collection_name, + points=points, + wait=wait, + ) + + logger.info(f"Updated {len(points)} documents in {self.collection_name}") + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode]: + """Retrieve vector nodes by their IDs from the collection.""" + single_result = isinstance(vector_ids, str) + if single_result: + vector_ids = [vector_ids] + + point_ids = [] + for vector_id in vector_ids: + try: + point_id = int(vector_id) + except ValueError: + point_id = abs(hash(vector_id)) % (10**18) + point_ids.append(point_id) + + points = await self.client.retrieve( + collection_name=self.collection_name, + ids=point_ids, + with_payload=True, + with_vectors=True, + ) + + results = [] + for point in points: + if point: + payload = point.payload or {} + node = VectorNode( + vector_id=payload.get("vector_id", str(point.id)), + content=payload.get("content", ""), + vector=point.vector, + metadata=payload.get("metadata", {}), + ) + results.append(node) + else: + logger.warning("Point not found") + + return results[0] if single_result and results else results + + async def list( + self, + filters: dict | None = None, + limit: int | None = None, + ) -> list[VectorNode]: + """List all vector nodes in the collection matching the filter criteria.""" + scroll_filter = self._create_filter(filters) if filters else None + + limit = limit or 10000 + records, _ = await self.client.scroll( + collection_name=self.collection_name, + scroll_filter=scroll_filter, + limit=limit, + with_payload=True, + with_vectors=True, + ) + + results = [] + for record in records: + payload = record.payload or {} + node = VectorNode( + vector_id=payload.get("vector_id", str(record.id)), + content=payload.get("content", ""), + vector=record.vector, + metadata=payload.get("metadata", {}), + ) + results.append(node) + + return results + + async def close(self): + """Close the AsyncQdrantClient connection and release resources.""" + await self.client.close() + logger.info("Qdrant client connection closed") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py new file mode 100644 index 00000000..76de21e1 --- /dev/null +++ b/tests/test_vector_store.py @@ -0,0 +1,1417 @@ +# pylint: disable=too-many-lines +"""Unified test suite for vector store implementations. + +This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore, +PGVectorStore, QdrantVectorStore, and ChromaVectorStore implementations. Tests can be +run for specific vector stores or all implementations. + +Usage: + python test_vector_store.py --local # Test LocalVectorStore only + python test_vector_store.py --es # Test ESVectorStore only + python test_vector_store.py --pgvector # Test PGVectorStore only + python test_vector_store.py --qdrant # Test QdrantVectorStore only + python test_vector_store.py --chroma # Test ChromaVectorStore only + python test_vector_store.py --all # Test all vector stores + +""" + +import argparse +import asyncio +import shutil +from pathlib import Path +from typing import List + +from loguru import logger + +from reme_ai.core.embedding import OpenAIEmbeddingModel +from reme_ai.core.schema import VectorNode +from reme_ai.core.vector_store import ( + BaseVectorStore, + ChromaVectorStore, + LocalVectorStore, + ESVectorStore, + PGVectorStore, + QdrantVectorStore, +) + + +# ==================== Configuration ==================== + + +class TestConfig: + """Configuration for test execution.""" + + # LocalVectorStore settings + LOCAL_ROOT_PATH = "./test_vector_store_local" + + # ESVectorStore settings + ES_HOSTS = "http://11.160.132.46:8200" + ES_BASIC_AUTH = None # Set to ("username", "password") if authentication is required + + # QdrantVectorStore settings + QDRANT_PATH = None # "./test_vector_store_qdrant" # For local mode + QDRANT_HOST = None # Set to host address for remote mode (e.g., "localhost") + QDRANT_PORT = None # Set to port for remote mode (e.g., 6333) + QDRANT_URL = "http://11.160.132.46:6333" # Alternative to host/port (e.g., http://localhost:6333) + QDRANT_API_KEY = None # Set for Qdrant Cloud authentication + + # PGVectorStore settings + PG_DSN = "postgresql://localhost/postgres" # PostgreSQL connection string + PG_MIN_SIZE = 1 # Minimum connections in pool + PG_MAX_SIZE = 5 # Maximum connections in pool + PG_USE_HNSW = True # Use HNSW index for faster search + PG_USE_DISKANN = False # Use DiskANN index (requires vectorscale extension) + + # ChromaVectorStore settings + CHROMA_PATH = "./test_vector_store_chroma" # For local persistent mode + CHROMA_HOST = None # Set to host address for remote mode (e.g., "localhost") + CHROMA_PORT = None # Set to port for remote mode (e.g., 8000) + CHROMA_API_KEY = None # Set for ChromaDB Cloud authentication + CHROMA_TENANT = None # Set for ChromaDB Cloud tenant + CHROMA_DATABASE = None # Set for ChromaDB Cloud database + + # Embedding model settings + EMBEDDING_MODEL_NAME = "text-embedding-v4" + EMBEDDING_DIMENSIONS = 64 + + # Test collection naming + TEST_COLLECTION_PREFIX = "test_vector_store" + + +# ==================== Sample Data Generator ==================== + + +class SampleDataGenerator: + """Generator for sample test data.""" + + @staticmethod + def create_sample_nodes(prefix: str = "") -> List[VectorNode]: + """Create sample VectorNode instances for testing. + + Args: + prefix: Optional prefix for vector_id to avoid conflicts + + Returns: + List[VectorNode]: List of sample nodes with diverse metadata + """ + id_prefix = f"{prefix}_" if prefix else "" + return [ + VectorNode( + vector_id=f"{id_prefix}node1", + content="Artificial intelligence is a technology that simulates human intelligence.", + metadata={ + "node_type": "tech", + "category": "AI", + "source": "research", + "priority": "high", + "year": "2023", + "department": "engineering", + "language": "english", + "status": "published", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node2", + content="Machine learning is a subset of artificial intelligence.", + metadata={ + "node_type": "tech", + "category": "ML", + "source": "research", + "priority": "high", + "year": "2022", + "department": "engineering", + "language": "english", + "status": "published", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node3", + content="Deep learning uses neural networks with multiple layers.", + metadata={ + "node_type": "tech_new", + "category": "DL", + "source": "blog", + "priority": "medium", + "year": "2024", + "department": "marketing", + "language": "chinese", + "status": "draft", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node4", + content="I love eating delicious seafood, especially fresh fish.", + metadata={ + "node_type": "food", + "category": "preference", + "source": "personal", + "priority": "low", + "year": "2023", + "department": "lifestyle", + "language": "english", + "status": "published", + }, + ), + VectorNode( + vector_id=f"{id_prefix}node5", + content="Natural language processing enables computers to understand human language.", + metadata={ + "node_type": "tech", + "category": "NLP", + "source": "research", + "priority": "high", + "year": "2024", + "department": "engineering", + "language": "english", + "status": "review", + }, + ), + ] + + +# ==================== Vector Store Factory ==================== + + +def get_store_type(store: BaseVectorStore) -> str: + """Get the type identifier of a vector store instance. + + Args: + store: Vector store instance + + Returns: + str: Type identifier ("local", "es", "pgvector", "qdrant", or "chroma") + """ + if isinstance(store, LocalVectorStore): + return "local" + elif isinstance(store, QdrantVectorStore): + return "qdrant" + elif isinstance(store, ESVectorStore): + return "es" + elif isinstance(store, PGVectorStore): + return "pgvector" + elif isinstance(store, ChromaVectorStore): + return "chroma" + else: + raise ValueError(f"Unknown vector store type: {type(store)}") + + +def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStore: + """Create a vector store instance based on type. + + Args: + store_type: Type of vector store ("local", "es", or "qdrant") + collection_name: Name of the collection + + Returns: + BaseVectorStore: Initialized vector store instance + """ + config = TestConfig() + + # Initialize embedding model + embedding_model = OpenAIEmbeddingModel( + model_name=config.EMBEDDING_MODEL_NAME, + dimensions=config.EMBEDDING_DIMENSIONS, + ) + + if store_type == "local": + return LocalVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + root_path=config.LOCAL_ROOT_PATH, + ) + elif store_type == "es": + return ESVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + hosts=config.ES_HOSTS, + basic_auth=config.ES_BASIC_AUTH, + ) + elif store_type == "qdrant": + return QdrantVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + path=config.QDRANT_PATH, + host=config.QDRANT_HOST, + port=config.QDRANT_PORT, + url=config.QDRANT_URL, + api_key=config.QDRANT_API_KEY, + distance="cosine", + on_disk=False, + ) + elif store_type == "pgvector": + return PGVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + dsn=config.PG_DSN, + min_size=config.PG_MIN_SIZE, + max_size=config.PG_MAX_SIZE, + use_hnsw=config.PG_USE_HNSW, + use_diskann=config.PG_USE_DISKANN, + ) + elif store_type == "chroma": + return ChromaVectorStore( + collection_name=collection_name, + embedding_model=embedding_model, + path=config.CHROMA_PATH, + host=config.CHROMA_HOST, + port=config.CHROMA_PORT, + api_key=config.CHROMA_API_KEY, + tenant=config.CHROMA_TENANT, + database=config.CHROMA_DATABASE, + ) + else: + raise ValueError(f"Unknown store type: {store_type}") + + +# ==================== Test Functions ==================== + + +async def test_create_collection(store: BaseVectorStore, _store_name: str): + """Test collection creation.""" + logger.info("=" * 20 + " CREATE COLLECTION TEST " + "=" * 20) + + # Clean up if exists + collections = await store.list_collections() + if store.collection_name in collections: + await store.delete_collection(store.collection_name) + logger.info(f"Cleaned up existing collection: {store.collection_name}") + + # Create collection + await store.create_collection(store.collection_name) + + # Verify creation + collections = await store.list_collections() + assert store.collection_name in collections, "Collection should exist after creation" + logger.info(f"✓ Created collection: {store.collection_name}") + + +async def test_insert(store: BaseVectorStore, _store_name: str) -> List[VectorNode]: + """Test node insertion (single and batch).""" + logger.info("=" * 20 + " INSERT TEST " + "=" * 20) + + # Test single node insertion + single_node = VectorNode( + vector_id="test_single_insert", + content="This is a single node insertion test", + metadata={"test_type": "single_insert"}, + ) + await store.insert(single_node) + logger.info("✓ Inserted single node") + + # Test batch insertion + sample_nodes = SampleDataGenerator.create_sample_nodes("test") + await store.insert(sample_nodes) + logger.info(f"✓ Batch inserted {len(sample_nodes)} nodes") + + # Verify total insertions + all_nodes = await store.list(limit=20) + assert len(all_nodes) >= len(sample_nodes) + 1, "Should have at least sample nodes + single node" + logger.info(f"✓ Total nodes in collection: {len(all_nodes)}") + + return sample_nodes + + +async def test_search(store: BaseVectorStore, _store_name: str): + """Test basic vector search.""" + logger.info("=" * 20 + " SEARCH TEST " + "=" * 20) + + results = await store.search( + query="What is artificial intelligence?", + limit=3, + ) + + logger.info(f"Search returned {len(results)} results") + for i, r in enumerate(results, 1): + score = r.metadata.get("_score", "N/A") + logger.info(f" Result {i}: {r.content[:60]}... (score: {score})") + + assert len(results) > 0, "Search should return results" + logger.info("✓ Basic search test passed") + + +async def test_search_with_single_filter(store: BaseVectorStore, _store_name: str): + """Test vector search with single metadata filter.""" + logger.info("=" * 20 + " SINGLE FILTER SEARCH TEST " + "=" * 20) + + # Test single value filter + filters = {"node_type": "tech"} + results = await store.search( + query="What is artificial intelligence?", + limit=5, + filters=filters, + ) + + logger.info(f"Filtered search (node_type=tech) returned {len(results)} results") + for i, r in enumerate(results, 1): + node_type = r.metadata.get("node_type") + logger.info(f" Result {i}: type={node_type}, content={r.content[:50]}...") + assert node_type == "tech", "Result should have node_type='tech'" + + logger.info("✓ Single filter search test passed") + + +async def test_search_with_list_filter(store: BaseVectorStore, _store_name: str): + """Test vector search with list filter (IN operation).""" + logger.info("=" * 20 + " LIST FILTER SEARCH TEST " + "=" * 20) + + # Test list filter (IN operation) + filters = {"node_type": ["tech", "tech_new"]} + results = await store.search( + query="What is artificial intelligence?", + limit=5, + filters=filters, + ) + + logger.info(f"Filtered search (node_type IN [tech, tech_new]) returned {len(results)} results") + for i, r in enumerate(results, 1): + node_type = r.metadata.get("node_type") + logger.info(f" Result {i}: type={node_type}, content={r.content[:50]}...") + assert node_type in ["tech", "tech_new"], "Result should have node_type in [tech, tech_new]" + + logger.info("✓ List filter search test passed") + + +async def test_search_with_multiple_filters(store: BaseVectorStore, _store_name: str): + """Test vector search with multiple metadata filters (AND operation).""" + logger.info("=" * 20 + " MULTIPLE FILTERS SEARCH TEST " + "=" * 20) + + # Test multiple filters (AND operation) + filters = { + "node_type": ["tech", "tech_new"], + "source": "research", + } + results = await store.search( + query="What is artificial intelligence?", + limit=5, + filters=filters, + ) + + logger.info( + f"Multi-filter search (node_type IN [tech, tech_new] AND source=research) " f"returned {len(results)} results", + ) + for i, r in enumerate(results, 1): + node_type = r.metadata.get("node_type") + source = r.metadata.get("source") + logger.info(f" Result {i}: type={node_type}, source={source}, content={r.content[:40]}...") + assert node_type in ["tech", "tech_new"], "Result should have node_type in [tech, tech_new]" + assert source == "research", "Result should have source='research'" + + logger.info("✓ Multiple filters search test passed") + + +async def test_get_by_id(store: BaseVectorStore, _store_name: str): + """Test retrieving nodes by vector_id (single and batch).""" + logger.info("=" * 20 + " GET BY ID TEST " + "=" * 20) + + # Test single ID retrieval + target_id = "test_node1" + result = await store.get(target_id) + + assert isinstance(result, VectorNode), "Should return a VectorNode for single ID" + assert result.vector_id == target_id, f"Result should have vector_id={target_id}" + logger.info(f"✓ Retrieved single node: {result.vector_id}") + + # Test batch retrieval (small batch) + target_ids = ["test_node1", "test_node2"] + results = await store.get(target_ids) + + assert isinstance(results, list), "Should return a list for multiple IDs" + assert len(results) == 2, f"Should return 2 results, got {len(results)}" + result_ids = {r.vector_id for r in results} + assert result_ids == set(target_ids), f"Result IDs should match {target_ids}" + logger.info(f"✓ Batch retrieved {len(results)} nodes") + + # Test larger batch retrieval + large_batch_ids = ["test_node1", "test_node2", "test_node3", "test_node5"] + large_results = await store.get(large_batch_ids) + assert isinstance(large_results, list), "Should return a list for batch IDs" + assert len(large_results) >= 3, "Should return at least 3 results" + logger.info(f"✓ Large batch retrieved {len(large_results)} nodes") + + +async def test_list_all(store: BaseVectorStore, _store_name: str): + """Test listing all nodes in collection.""" + logger.info("=" * 20 + " LIST ALL TEST " + "=" * 20) + + results = await store.list(limit=10) + + logger.info(f"Collection contains {len(results)} nodes") + for i, node in enumerate(results, 1): + logger.info(f" Node {i}: id={node.vector_id}, content={node.content[:50]}...") + + assert len(results) > 0, "Collection should contain nodes" + logger.info("✓ List all nodes test passed") + + +async def test_list_with_filters(store: BaseVectorStore, _store_name: str): + """Test listing nodes with metadata filters.""" + logger.info("=" * 20 + " LIST WITH FILTERS TEST " + "=" * 20) + + filters = {"category": "AI"} + results = await store.list(filters=filters, limit=10) + + logger.info(f"Filtered list (category=AI) returned {len(results)} nodes") + for i, node in enumerate(results, 1): + category = node.metadata.get("category") + logger.info(f" Node {i}: category={category}, id={node.vector_id}") + assert category == "AI", "All nodes should have category=AI" + + logger.info("✓ List with filters test passed") + + +async def test_update(store: BaseVectorStore, _store_name: str): + """Test updating existing nodes (single and batch).""" + logger.info("=" * 20 + " UPDATE TEST " + "=" * 20) + + # Test single node update + updated_node = VectorNode( + vector_id="test_node2", + content="Machine learning is a powerful subset of AI that learns from data.", + metadata={ + "node_type": "tech", + "category": "ML", + "updated": "true", + "update_timestamp": "2024-12-26", + }, + ) + + await store.update(updated_node) + + # Verify single update + result = await store.get("test_node2") + assert "updated" in result.metadata, "Updated metadata should be present" + logger.info(f"✓ Updated single node: {result.vector_id}") + logger.info(f" New content: {result.content[:60]}...") + + # Test batch update (update multiple nodes at once) + batch_update_nodes = [ + VectorNode( + vector_id="test_node1", + content="Artificial intelligence is evolving rapidly with new breakthroughs.", + metadata={ + "node_type": "tech", + "category": "AI", + "batch_updated": "true", + "update_timestamp": "2024-12-31", + }, + ), + VectorNode( + vector_id="test_node3", + content="Deep learning revolutionizes neural network architectures.", + metadata={ + "node_type": "tech_new", + "category": "DL", + "batch_updated": "true", + "update_timestamp": "2024-12-31", + }, + ), + ] + + await store.update(batch_update_nodes) + logger.info(f"✓ Batch updated {len(batch_update_nodes)} nodes") + + # Verify batch updates + results = await store.get(["test_node1", "test_node3"]) + for r in results: + assert r.metadata.get("batch_updated") == "true", f"Node {r.vector_id} should have batch_updated metadata" + logger.info(f"✓ Verified batch update for {len(results)} nodes") + + +async def test_delete(store: BaseVectorStore, _store_name: str): + """Test deleting nodes (single and batch).""" + logger.info("=" * 20 + " DELETE TEST " + "=" * 20) + + # Test single node deletion + node_to_delete = "test_node4" + await store.delete(node_to_delete) + + # Verify single deletion - try to get the deleted node + try: + result = await store.get(node_to_delete) + # If result is empty list or None, deletion was successful + if isinstance(result, list): + assert len(result) == 0, "Deleted node should not be retrievable" + else: + assert result is None, "Deleted node should not be retrievable" + except Exception: + pass # Expected if node doesn't exist + + logger.info(f"✓ Deleted single node: {node_to_delete}") + + # Test batch deletion - first insert some nodes to delete + batch_delete_nodes = [ + VectorNode( + vector_id=f"delete_test_{i}", + content=f"Node to be deleted {i}", + metadata={"test_type": "delete_batch"}, + ) + for i in range(5) + ] + await store.insert(batch_delete_nodes) + logger.info(f"✓ Inserted {len(batch_delete_nodes)} nodes for batch delete test") + + # Batch delete + delete_ids = [f"delete_test_{i}" for i in range(5)] + await store.delete(delete_ids) + logger.info(f"✓ Batch deleted {len(delete_ids)} nodes") + + # Verify batch deletion + try: + results = await store.get(delete_ids) + if isinstance(results, list): + assert len(results) == 0, "All deleted nodes should not be retrievable" + except Exception: + pass # Expected if nodes don't exist + logger.info("✓ Verified batch deletion") + + +async def test_copy_collection(store: BaseVectorStore, store_name: str): + """Test copying a collection.""" + logger.info("=" * 20 + " COPY COLLECTION TEST " + "=" * 20) + + config = TestConfig() + copy_collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_name}_copy" + + # Elasticsearch and PostgreSQL require lowercase table/index names + store_type = get_store_type(store) + if store_type in ("es", "pgvector"): + copy_collection_name = copy_collection_name.lower() + + # Clean up if exists + collections = await store.list_collections() + if copy_collection_name in collections: + await store.delete_collection(copy_collection_name) + + # Copy collection + await store.copy_collection(copy_collection_name) + + # Verify copy + collections = await store.list_collections() + assert copy_collection_name in collections, "Copied collection should exist" + logger.info(f"✓ Copied collection to: {copy_collection_name}") + + # Verify content in copied collection + copied_store = create_vector_store(store_type, copy_collection_name) + copied_nodes = await copied_store.list() + logger.info(f"✓ Copied collection has {len(copied_nodes)} nodes") + await copied_store.close() + + # Clean up copied collection + await store.delete_collection(copy_collection_name) + logger.info("✓ Cleaned up copied collection") + + +async def test_list_collections(store: BaseVectorStore, _store_name: str): + """Test listing all collections.""" + logger.info("=" * 20 + " LIST COLLECTIONS TEST " + "=" * 20) + + collections = await store.list_collections() + + logger.info(f"Found {len(collections)} collections") + config = TestConfig() + test_collections = [c for c in collections if c.startswith(config.TEST_COLLECTION_PREFIX)] + logger.info(f" Test collections: {test_collections}") + + assert store.collection_name in collections, "Main test collection should be listed" + logger.info("✓ List collections test passed") + + +async def test_delete_collection(store: BaseVectorStore, _store_name: str): + """Test deleting a collection.""" + logger.info("=" * 20 + " DELETE COLLECTION TEST " + "=" * 20) + + await store.delete_collection(store.collection_name) + + # Verify deletion + collections = await store.list_collections() + assert store.collection_name not in collections, "Collection should not exist after deletion" + logger.info(f"✓ Deleted collection: {store.collection_name}") + + +async def test_cosine_similarity(store_name: str): + """Test manual cosine similarity calculation (LocalVectorStore only).""" + if store_name != "LocalVectorStore": + logger.info("=" * 20 + " COSINE SIMILARITY TEST (SKIPPED) " + "=" * 20) + logger.info("⊘ Skipped: Only applicable to LocalVectorStore") + return + + logger.info("=" * 20 + " COSINE SIMILARITY TEST " + "=" * 20) + + vec1 = [1.0, 0.0, 0.0] + vec2 = [0.0, 1.0, 0.0] + vec3 = [1.0, 0.0, 0.0] + + # Test perpendicular vectors (similarity = 0) + sim1 = LocalVectorStore._cosine_similarity(vec1, vec2) # pylint: disable=protected-access + logger.info(f"Similarity between perpendicular vectors: {sim1:.4f}") + assert abs(sim1) < 0.0001, "Perpendicular vectors should have similarity close to 0" + + # Test identical vectors (similarity = 1) + sim2 = LocalVectorStore._cosine_similarity(vec1, vec3) # pylint: disable=protected-access + logger.info(f"Similarity between identical vectors: {sim2:.4f}") + assert abs(sim2 - 1.0) < 0.0001, "Identical vectors should have similarity close to 1" + + # Test with real-world like vectors + vec4 = [0.5, 0.5, 0.5] + vec5 = [0.6, 0.4, 0.5] + sim3 = LocalVectorStore._cosine_similarity(vec4, vec5) # pylint: disable=protected-access + logger.info(f"Similarity between similar vectors: {sim3:.4f}") + assert sim3 > 0.9, "Similar vectors should have high similarity" + + logger.info("✓ Cosine similarity tests passed") + + +async def test_batch_operations(store: BaseVectorStore, _store_name: str): + """Test large-scale batch insert, update, and delete operations. + + This test validates the efficiency and correctness of batch operations + by processing 100 nodes at once, which is more realistic for production use. + """ + logger.info("=" * 20 + " BATCH OPERATIONS TEST " + "=" * 20) + + # Create a large batch of nodes (100 nodes) + batch_nodes = [] + for i in range(100): + batch_nodes.append( + VectorNode( + vector_id=f"batch_node_{i}", + content=f"This is batch test content number {i} about various topics in technology and science.", + metadata={ + "batch_id": str(i // 10), # Group into batches of 10 + "index": str(i), + "category": ["tech", "science", "business"][i % 3], + "priority": ["high", "medium", "low"][i % 3], + }, + ), + ) + + # Batch insert + await store.insert(batch_nodes) + logger.info(f"✓ Inserted {len(batch_nodes)} nodes in batch") + + # Verify batch insert + results = await store.list(limit=150) + assert len(results) >= 100, f"Should have at least 100 nodes, got {len(results)}" + logger.info(f"✓ Verified batch insert: {len(results)} total nodes") + + # Batch update (update first 20 nodes) + update_nodes = [] + for i in range(20): + update_nodes.append( + VectorNode( + vector_id=f"batch_node_{i}", + content=f"UPDATED: This is updated batch content {i}", + metadata={ + "batch_id": str(i // 10), + "index": str(i), + "updated": "true", + "update_timestamp": "2024-12-31", + }, + ), + ) + + await store.update(update_nodes) + logger.info(f"✓ Updated {len(update_nodes)} nodes in batch") + + # Verify updates + updated_results = await store.list(filters={"updated": "true"}, limit=50) + assert len(updated_results) >= 20, "Should have at least 20 updated nodes" + logger.info(f"✓ Verified batch update: {len(updated_results)} updated nodes") + + # Batch delete (delete nodes with batch_id >= 5) + delete_ids = [f"batch_node_{i}" for i in range(50, 100)] + await store.delete(delete_ids) + logger.info(f"✓ Deleted {len(delete_ids)} nodes in batch") + + # Verify deletions + remaining = await store.list(limit=150) + batch_nodes_remaining = [n for n in remaining if n.vector_id.startswith("batch_node_")] + assert len(batch_nodes_remaining) <= 50, "Should have at most 50 batch nodes remaining" + logger.info(f"✓ Verified batch delete: {len(batch_nodes_remaining)} nodes remaining") + + +async def test_complex_metadata_queries(store: BaseVectorStore, _store_name: str): + """Test complex metadata filtering with nested conditions.""" + logger.info("=" * 20 + " COMPLEX METADATA QUERIES TEST " + "=" * 20) + + # Insert nodes with rich metadata + complex_nodes = [ + VectorNode( + vector_id="complex_1", + content="Advanced neural networks for computer vision applications", + metadata={ + "domain": "AI", + "subdomain": "computer_vision", + "year": "2024", + "citations": "150", + "impact_factor": "high", + "tags": "neural_networks,vision,deep_learning", + }, + ), + VectorNode( + vector_id="complex_2", + content="Natural language processing with transformer models", + metadata={ + "domain": "AI", + "subdomain": "nlp", + "year": "2023", + "citations": "200", + "impact_factor": "high", + "tags": "transformers,nlp,language_models", + }, + ), + VectorNode( + vector_id="complex_3", + content="Reinforcement learning for robotics control", + metadata={ + "domain": "AI", + "subdomain": "robotics", + "year": "2024", + "citations": "80", + "impact_factor": "medium", + "tags": "reinforcement_learning,robotics,control", + }, + ), + VectorNode( + vector_id="complex_4", + content="Quantum computing algorithms and applications", + metadata={ + "domain": "quantum", + "subdomain": "algorithms", + "year": "2024", + "citations": "50", + "impact_factor": "medium", + "tags": "quantum,algorithms,computing", + }, + ), + ] + + await store.insert(complex_nodes) + logger.info(f"✓ Inserted {len(complex_nodes)} nodes with complex metadata") + + # Test 1: Multiple field filters with list values + filters_1 = { + "domain": "AI", + "year": ["2023", "2024"], + "impact_factor": "high", + } + results_1 = await store.search( + query="artificial intelligence research", + limit=10, + filters=filters_1, + ) + logger.info(f"Test 1 - AI + high impact + recent years: {len(results_1)} results") + for r in results_1: + assert r.metadata.get("domain") == "AI" + assert r.metadata.get("impact_factor") == "high" + assert r.metadata.get("year") in ["2023", "2024"] + + # Test 2: List filter with multiple subdomains + filters_2 = { + "subdomain": ["nlp", "computer_vision"], + } + results_2 = await store.search( + query="deep learning applications", + limit=10, + filters=filters_2, + ) + logger.info(f"Test 2 - NLP or Computer Vision: {len(results_2)} results") + for r in results_2: + assert r.metadata.get("subdomain") in ["nlp", "computer_vision"] + + # Test 3: Year-based filtering + filters_3 = { + "year": "2024", + } + results_3 = await store.list(filters=filters_3, limit=10) + logger.info(f"Test 3 - Year 2024 only: {len(results_3)} results") + for r in results_3: + assert r.metadata.get("year") == "2024" + + logger.info("✓ Complex metadata queries test passed") + + +async def test_edge_cases(store: BaseVectorStore, _store_name: str): + """Test edge cases and boundary conditions.""" + logger.info("=" * 20 + " EDGE CASES TEST " + "=" * 20) + + # Test 1: Empty content + edge_node_1 = VectorNode( + vector_id="edge_empty_content", + content="", + metadata={"type": "empty"}, + ) + try: + await store.insert([edge_node_1]) + logger.info("✓ Handled empty content node") + except Exception as e: + logger.info(f"⊘ Empty content not supported: {e}") + + # Test 2: Very long content + edge_node_2 = VectorNode( + vector_id="edge_long_content", + content="A" * 10000, # 10k characters + metadata={"type": "long_content"}, + ) + await store.insert([edge_node_2]) + result = await store.get("edge_long_content") + assert len(result.content) == 10000 + logger.info("✓ Handled very long content (10k chars)") + + # Test 3: Special characters in content + edge_node_3 = VectorNode( + vector_id="edge_special_chars", + content="Special chars: @#$%^&*()[]{}|\\;:'\",.<>?/~`+=−×÷", + metadata={"type": "special_chars"}, + ) + await store.insert([edge_node_3]) + result = await store.get("edge_special_chars") + assert "@#$%^&*()" in result.content + logger.info("✓ Handled special characters in content") + + # Test 4: Unicode and emoji content + edge_node_4 = VectorNode( + vector_id="edge_unicode", + content="Unicode test: 你好世界 🌍 مرحبا العالم Привет мир", + metadata={"type": "unicode", "language": "multi"}, + ) + await store.insert([edge_node_4]) + result = await store.get("edge_unicode") + assert "你好世界" in result.content + assert "🌍" in result.content + logger.info("✓ Handled unicode and emoji content") + + # Test 5: Search with empty query + try: + results = await store.search(query="", limit=5) + logger.info(f"✓ Empty query returned {len(results)} results") + except Exception as e: + logger.info(f"⊘ Empty query not supported: {e}") + + # Test 6: Search with very high limit + results = await store.search(query="test", limit=1000) + logger.info(f"✓ High limit search returned {len(results)} results") + + # Test 7: Get non-existent ID + result = await store.get("non_existent_id_12345") + if isinstance(result, list): + assert len(result) == 0, "Non-existent ID should return empty list" + else: + assert result is None, "Non-existent ID should return None" + logger.info("✓ Handled non-existent ID gracefully") + + # Test 8: Metadata with empty string values + edge_node_5 = VectorNode( + vector_id="edge_empty_metadata", + content="Testing empty string values in metadata", + metadata={"field1": "value1", "field2": "", "field3": "value3"}, + ) + await store.insert([edge_node_5]) + logger.info("✓ Handled empty string values in metadata") + + logger.info("✓ Edge cases test passed") + + +async def test_concurrent_operations(store: BaseVectorStore, _store_name: str): + """Test concurrent read/write operations.""" + logger.info("=" * 20 + " CONCURRENT OPERATIONS TEST " + "=" * 20) + + # Prepare concurrent insert nodes + concurrent_nodes = [ + VectorNode( + vector_id=f"concurrent_{i}", + content=f"Concurrent test content {i}", + metadata={"thread_id": str(i % 5), "index": str(i)}, + ) + for i in range(50) + ] + + # Test concurrent inserts + insert_tasks = [] + for i in range(0, 50, 10): + batch = concurrent_nodes[i : i + 10] + insert_tasks.append(store.insert(batch)) + + await asyncio.gather(*insert_tasks) + logger.info("✓ Completed concurrent inserts") + + # Test concurrent searches + search_tasks = [store.search(query=f"concurrent test {i}", limit=5) for i in range(10)] + search_results = await asyncio.gather(*search_tasks) + logger.info(f"✓ Completed {len(search_results)} concurrent searches") + + # Test concurrent reads + get_tasks = [store.get(f"concurrent_{i}") for i in range(0, 50, 5)] + get_results = await asyncio.gather(*get_tasks) + logger.info(f"✓ Completed {len(get_results)} concurrent reads") + + # Test batch updates (using batch update instead of concurrent individual updates) + update_nodes = [ + VectorNode( + vector_id=f"concurrent_{i}", + content=f"UPDATED concurrent content {i}", + metadata={"thread_id": str(i % 5), "updated": "true"}, + ) + for i in range(0, 20, 2) + ] + await store.update(update_nodes) + logger.info(f"✓ Completed batch update of {len(update_nodes)} nodes") + + logger.info("✓ Concurrent operations test passed") + + +async def test_search_relevance_ranking(store: BaseVectorStore, _store_name: str): + """Test search result relevance and ranking.""" + logger.info("=" * 20 + " SEARCH RELEVANCE RANKING TEST " + "=" * 20) + + # Insert nodes with varying relevance + relevance_nodes = [ + VectorNode( + vector_id="relevance_exact", + content="Machine learning is a subset of artificial intelligence focused on learning from data.", + metadata={"relevance": "exact"}, + ), + VectorNode( + vector_id="relevance_high", + content="Artificial intelligence and machine learning are transforming technology.", + metadata={"relevance": "high"}, + ), + VectorNode( + vector_id="relevance_medium", + content="Deep learning uses neural networks for pattern recognition.", + metadata={"relevance": "medium"}, + ), + VectorNode( + vector_id="relevance_low", + content="Software engineering best practices for code quality.", + metadata={"relevance": "low"}, + ), + VectorNode( + vector_id="relevance_none", + content="Cooking recipes for delicious Italian pasta dishes.", + metadata={"relevance": "none"}, + ), + ] + + await store.insert(relevance_nodes) + logger.info(f"✓ Inserted {len(relevance_nodes)} nodes with varying relevance") + + # Search with specific query + query = "What is machine learning and artificial intelligence?" + results = await store.search(query=query, limit=5) + + logger.info(f"Search results for: '{query}'") + for i, result in enumerate(results, 1): + score = result.metadata.get("_score", "N/A") + relevance = result.metadata.get("relevance", "unknown") + logger.info(f" {i}. [{relevance}] score={score}: {result.content[:60]}...") + + # Verify that more relevant results appear first + if len(results) >= 2: + # The exact match should be in top results + top_ids = [r.vector_id for r in results[:3]] + assert ( + "relevance_exact" in top_ids or "relevance_high" in top_ids + ), "Most relevant results should appear in top 3" + logger.info("✓ Relevance ranking verified") + + # Test with different query + query2 = "neural networks deep learning" + results2 = await store.search(query=query2, limit=5) + logger.info(f"\nSearch results for: '{query2}'") + for i, result in enumerate(results2, 1): + score = result.metadata.get("_score", "N/A") + logger.info(f" {i}. score={score}: {result.content[:60]}...") + + logger.info("✓ Search relevance ranking test passed") + + +async def test_metadata_statistics(store: BaseVectorStore, _store_name: str): + """Test metadata aggregation and statistics.""" + logger.info("=" * 20 + " METADATA STATISTICS TEST " + "=" * 20) + + # Get all nodes and analyze metadata + all_nodes = await store.list(limit=500) + logger.info(f"Total nodes in collection: {len(all_nodes)}") + + # Count by category + category_counts = {} + for node in all_nodes: + category = node.metadata.get("category", "unknown") + category_counts[category] = category_counts.get(category, 0) + 1 + + logger.info("Category distribution:") + for category, count in sorted(category_counts.items()): + logger.info(f" {category}: {count}") + + # Count by node_type + type_counts = {} + for node in all_nodes: + node_type = node.metadata.get("node_type", "unknown") + type_counts[node_type] = type_counts.get(node_type, 0) + 1 + + logger.info("Node type distribution:") + for node_type, count in sorted(type_counts.items()): + logger.info(f" {node_type}: {count}") + + # Verify we can filter by each category + for category in category_counts: + if category != "unknown": + filtered = await store.list(filters={"category": category}, limit=100) + logger.info(f"✓ Filter by category '{category}': {len(filtered)} results") + + logger.info("✓ Metadata statistics test passed") + + +async def test_update_metadata_only(store: BaseVectorStore, _store_name: str): + """Test updating only metadata without changing content.""" + logger.info("=" * 20 + " UPDATE METADATA ONLY TEST " + "=" * 20) + + # Get an existing node + original = await store.get("test_node1") + original_content = original.content + + # Update with same content but different metadata + updated_node = VectorNode( + vector_id="test_node1", + content=original_content, # Keep same content + metadata={ + **original.metadata, + "metadata_updated": "true", + "update_count": "1", + "last_modified": "2024-12-31", + }, + ) + + await store.update(updated_node) + logger.info("✓ Updated metadata without changing content") + + # Verify update + result = await store.get("test_node1") + assert result.content == original_content, "Content should remain unchanged" + assert result.metadata.get("metadata_updated") == "true", "Metadata should be updated" + logger.info("✓ Verified metadata-only update") + + # Update metadata again + updated_node_2 = VectorNode( + vector_id="test_node1", + content=original_content, + metadata={ + **result.metadata, + "update_count": "2", + "last_modified": "2024-12-31T12:00:00", + }, + ) + await store.update(updated_node_2) + + result_2 = await store.get("test_node1") + assert result_2.metadata.get("update_count") == "2", "Metadata should be updated again" + logger.info("✓ Multiple metadata updates successful") + + logger.info("✓ Update metadata only test passed") + + +async def test_filter_combinations(store: BaseVectorStore, _store_name: str): + """Test various filter combinations and edge cases.""" + logger.info("=" * 20 + " FILTER COMBINATIONS TEST " + "=" * 20) + + # Test 1: Empty filter (should return all results) + results_1 = await store.search(query="technology", filters={}, limit=10) + logger.info(f"Test 1 - Empty filter: {len(results_1)} results") + + # Test 2: Single value filter + results_2 = await store.search( + query="technology", + filters={"node_type": "tech"}, + limit=10, + ) + logger.info(f"Test 2 - Single value filter: {len(results_2)} results") + for r in results_2: + assert r.metadata.get("node_type") == "tech" + + # Test 3: List filter with single item + results_3 = await store.search( + query="technology", + filters={"node_type": ["tech"]}, + limit=10, + ) + logger.info(f"Test 3 - List filter (single item): {len(results_3)} results") + + # Test 4: List filter with multiple items + results_4 = await store.search( + query="technology", + filters={"category": ["AI", "ML", "DL"]}, + limit=10, + ) + logger.info(f"Test 4 - List filter (multiple items): {len(results_4)} results") + for r in results_4: + assert r.metadata.get("category") in ["AI", "ML", "DL"] + + # Test 5: Multiple filters (AND operation) + results_5 = await store.search( + query="technology", + filters={ + "node_type": ["tech", "tech_new"], + "source": "research", + "priority": "high", + }, + limit=10, + ) + logger.info(f"Test 5 - Multiple filters (AND): {len(results_5)} results") + for r in results_5: + assert r.metadata.get("node_type") in ["tech", "tech_new"] + assert r.metadata.get("source") == "research" + assert r.metadata.get("priority") == "high" + + # Test 6: Filter with non-existent value + results_6 = await store.search( + query="technology", + filters={"category": "NON_EXISTENT_CATEGORY"}, + limit=10, + ) + logger.info(f"Test 6 - Non-existent filter value: {len(results_6)} results") + assert len(results_6) == 0, "Should return no results for non-existent filter value" + + # Test 7: List operation with filters + list_results = await store.list( + filters={"node_type": "tech", "priority": "high"}, + limit=20, + ) + logger.info(f"Test 7 - List with filters: {len(list_results)} results") + for r in list_results: + assert r.metadata.get("node_type") == "tech" + assert r.metadata.get("priority") == "high" + + logger.info("✓ Filter combinations test passed") + + +# ==================== Test Runner ==================== + + +async def run_all_tests_for_store(store_type: str, store_name: str): + """Run all tests for a specific vector store type. + + Args: + store_type: Type of vector store ("local" or "es") + store_name: Display name for the vector store + """ + logger.info(f"\n\n{'#' * 60}") + logger.info(f"# Running all tests for: {store_name}") + logger.info(f"{'#' * 60}") + + config = TestConfig() + collection_name = f"{config.TEST_COLLECTION_PREFIX}_{store_type}_main" + + # Create vector store instance + store = create_vector_store(store_type, collection_name) + + try: + # Run cosine similarity test first (only for LocalVectorStore) + await test_cosine_similarity(store_name) + + # ========== Basic Tests ========== + logger.info(f"\n{'#' * 60}") + logger.info("# BASIC FUNCTIONALITY TESTS") + logger.info(f"{'#' * 60}") + + await test_create_collection(store, store_name) + await test_insert(store, store_name) + await test_search(store, store_name) + await test_search_with_single_filter(store, store_name) + await test_search_with_list_filter(store, store_name) + await test_search_with_multiple_filters(store, store_name) + await test_get_by_id(store, store_name) + await test_list_all(store, store_name) + await test_list_with_filters(store, store_name) + await test_update(store, store_name) + await test_delete(store, store_name) + + # ========== Advanced Tests ========== + logger.info(f"\n{'#' * 60}") + logger.info("# ADVANCED FUNCTIONALITY TESTS") + logger.info(f"{'#' * 60}") + + await test_batch_operations(store, store_name) + await test_complex_metadata_queries(store, store_name) + await test_edge_cases(store, store_name) + await test_concurrent_operations(store, store_name) + await test_search_relevance_ranking(store, store_name) + await test_metadata_statistics(store, store_name) + await test_update_metadata_only(store, store_name) + await test_filter_combinations(store, store_name) + + # ========== Collection Management Tests ========== + logger.info(f"\n{'#' * 60}") + logger.info("# COLLECTION MANAGEMENT TESTS") + logger.info(f"{'#' * 60}") + + await test_list_collections(store, store_name) + await test_copy_collection(store, store_name) + await test_delete_collection(store, store_name) + + logger.info(f"\n{'=' * 60}") + logger.info(f"✓ All tests passed for {store_name}!") + logger.info(f"{'=' * 60}") + + except Exception as e: + logger.error(f"Test failed: {e}") + raise + finally: + # Cleanup + await cleanup_store(store, store_type) + + +async def cleanup_store(store: BaseVectorStore, store_type: str): + """Clean up test resources for a vector store. + + Args: + store: Vector store instance + store_type: Type of vector store ("local" or "es") + """ + logger.info("=" * 20 + " CLEANUP " + "=" * 20) + + try: + # Clean up test collections + config = TestConfig() + collections = await store.list_collections() + test_collections = [c for c in collections if c.startswith(config.TEST_COLLECTION_PREFIX)] + + for collection in test_collections: + try: + await store.delete_collection(collection) + logger.info(f"Deleted test collection: {collection}") + except Exception as e: + logger.warning(f"Failed to delete collection {collection}: {e}") + + # Close connections + await store.close() + + # Clean up local directory if LocalVectorStore + if store_type == "local": + test_dir = Path(config.LOCAL_ROOT_PATH) + if test_dir.exists(): + shutil.rmtree(test_dir) + logger.info(f"Cleaned up local directory: {config.LOCAL_ROOT_PATH}") + + # Clean up local directory if ChromaVectorStore + if store_type == "chroma" and config.CHROMA_PATH: + test_dir = Path(config.CHROMA_PATH) + if test_dir.exists(): + shutil.rmtree(test_dir) + logger.info(f"Cleaned up chroma directory: {config.CHROMA_PATH}") + + logger.info("✓ Cleanup completed") + except Exception as e: + logger.error(f"Cleanup error: {e}") + + +# ==================== Main Entry Point ==================== + + +async def main(): + """Main entry point for running tests.""" + parser = argparse.ArgumentParser( + description="Run vector store tests", + formatter_class=argparse.RawDescriptionHelpFormatter, + epilog=""" +Examples: + python test_vector_store.py --local # Test LocalVectorStore only + python test_vector_store.py --es # Test ESVectorStore only + python test_vector_store.py --pgvector # Test PGVectorStore only + python test_vector_store.py --qdrant # Test QdrantVectorStore only + python test_vector_store.py --chroma # Test ChromaVectorStore only + python test_vector_store.py --all # Test all vector stores + """, + ) + parser.add_argument( + "--local", + action="store_true", + help="Test LocalVectorStore", + ) + parser.add_argument( + "--es", + action="store_true", + help="Test ESVectorStore", + ) + parser.add_argument( + "--qdrant", + action="store_true", + help="Test QdrantVectorStore", + ) + parser.add_argument( + "--pgvector", + action="store_true", + help="Test PGVectorStore", + ) + parser.add_argument( + "--chroma", + action="store_true", + help="Test ChromaVectorStore", + ) + parser.add_argument( + "--all", + action="store_true", + help="Run tests for all available vector stores", + ) + + args = parser.parse_args() + + # Determine which vector stores to test + stores_to_test = [] + + if args.all: + stores_to_test = [ + ("local", "LocalVectorStore"), + ("es", "ESVectorStore"), + ("pgvector", "PGVectorStore"), + ("qdrant", "QdrantVectorStore"), + ("chroma", "ChromaVectorStore"), + ] + else: + # Build list based on individual flags + if args.local: + stores_to_test.append(("local", "LocalVectorStore")) + if args.es: + stores_to_test.append(("es", "ESVectorStore")) + if args.pgvector: + stores_to_test.append(("pgvector", "PGVectorStore")) + if args.qdrant: + stores_to_test.append(("qdrant", "QdrantVectorStore")) + if args.chroma: + stores_to_test.append(("chroma", "ChromaVectorStore")) + + if not stores_to_test: + # Default to all vector stores if no argument provided + stores_to_test = [ + ("local", "LocalVectorStore"), + ("es", "ESVectorStore"), + ("pgvector", "PGVectorStore"), + ("qdrant", "QdrantVectorStore"), + ("chroma", "ChromaVectorStore"), + ] + print("No vector store specified, defaulting to test all vector stores") + print( + "Use --local/--es/--pgvector/--qdrant/--chroma to test specific ones\n", + ) + + # Run tests for each vector store + for store_type, store_name in stores_to_test: + try: + await run_all_tests_for_store(store_type, store_name) + except Exception as e: + logger.error(f"\n✗ FAILED: {store_name} tests failed with error:") + logger.error(f" {type(e).__name__}: {e}") + raise + + # Final summary + print(f"\n\n{'#' * 60}") + print("# TEST SUMMARY") + print(f"{'#' * 60}") + print(f"✓ All tests passed for {len(stores_to_test)} vector store(s):") + for _, store_name in stores_to_test: + print(f" - {store_name}") + print(f"{'#' * 60}\n") + + +if __name__ == "__main__": + asyncio.run(main())