diff --git a/README.md b/README.md index f2f413c0..89813e4b 100644 --- a/README.md +++ b/README.md @@ -506,7 +506,7 @@ async def main(): "dimensions": 1024, }, default_vector_store_config={ - "backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec + "backend": "local", # Supports local/chroma/qdrant/elasticsearch/obvec/hologres }, ) await reme.start() diff --git a/docs/index.md b/docs/index.md index d7e2e716..d0e9d9d2 100644 --- a/docs/index.md +++ b/docs/index.md @@ -139,7 +139,7 @@ response = requests.post("http://localhost:8002/retrieve_task_memory", json={ ## 📚 Resources - **[Installation Guide](installation.md)**, **[Quick Start](quick_start.md)**: Get started quickly with practical examples -- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local, Elasticsearch, Qdrant, ChromaDB, or ObVec (OceanBase / seekdb via pyobvector) storage and usage +- **[Vector Storage Setup](vector_store_api_guide.md)**: Configure local, Elasticsearch, Qdrant, ChromaDB, ObVec (OceanBase / seekdb via pyobvector) or Hologres storage and usage - **[MCP Guide](mcp_quick_start.md)**: Create MCP services - **[Personal Memory](personal_memory/personal_memory.md)**, **[Task Memory](task_memory/task_memory.md)** & **[Tool Memory](tool_memory/tool_memory.md)**: Operators used in personal memory, task memory and tool memory. You can modify the config to customize the pipelines. - **[Example Collection](./cookbook/appworld/quickstart.md)**: Real use cases and best practices diff --git a/docs/vector_store_api_guide.md b/docs/vector_store_api_guide.md index b0a06eab..ec973c3c 100644 --- a/docs/vector_store_api_guide.md +++ b/docs/vector_store_api_guide.md @@ -34,6 +34,7 @@ FlowLLM provides multiple Vector Store implementations tailored to different use - **ChromaVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/chroma_vector_store.py)): Based on ChromaDB, providing persistent storage and metadata filtering capabilities. - **EsVectorStore** ([source code](https://github.com/flowllm-ai/flowllm/blob/main/flowllm/core/vector_store/es_vector_store.py)): Built on Elasticsearch, enabling powerful combined full-text and vector search functionalities. - **ObVecVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/obvec_vector_store.py)): Uses [pyobvector](https://pypi.org/project/pyobvector/) against **OceanBase** or **seekdb** (MySQL-compatible wire protocol). Suitable when you already run OceanBase/seekdb or need a SQL-native vector table with HNSW-style ANN search and JSON metadata filters. +- **HologresVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/hologres_store.py)): Uses [asyncpg](https://pypi.org/project/asyncpg/) against **Hologres** (PostgreSQL-compatible). Leverages native `float4[]` vector storage with built-in HGraph index for approximate nearest neighbor search. Suitable when you already run Hologres or need high-performance vector search with JSONB metadata filtering in a PostgreSQL-compatible environment. All Vector Store implementations inherit from **BaseVectorStore** ([source code](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/base_vector_store.py)) in ReMe, ensuring a consistent interface specification. @@ -131,6 +132,20 @@ docker run -d --name reme_seekdb -p 2881:2881 -e ROOT_PASSWORD= python tests/test_vector_store.py --obvec ``` +### HologresVectorStore Configuration + +- **host**: Hologres host address (default: `localhost`). +- **port**: Hologres port (default: `80`). +- **database**: Database name (default: `postgres`). +- **user**: Database user (default: `postgres`). +- **password**: Database password. +- **schema**: PostgreSQL schema name (default: `public`). +- **min_size**: Minimum connections in pool (default: `1`). +- **max_size**: Maximum connections in pool (default: `10`). +- **dsn**: Full DSN connection string. When provided, overrides `host`, `port`, `database`, `user`, and `password`. +- **distance_method**: Distance method for the HGraph index: `Cosine`, `InnerProduct`, or `Euclidean` (default: `Cosine`). +- **collection_name**: Table name for the collection (from `VectorStoreConfig`, default `reme`). + ## Configuration File Examples Configure Vector Store in `flowllm/config/default.yaml` under the `vector_store` section. The basic structure is as follows: @@ -151,7 +166,7 @@ vector_store.default.params.= ### Configuration Field Descriptions -- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`. +- **`backend`** (required): Vector store backend type. Options: `local`, `memory`, `chroma`, `qdrant`, `elasticsearch`, `obvec`, `hologres`. - **`embedding_model`** (required): Name of the embedding model configuration, referencing the `embedding_model` section. - **`params`** (optional): Dictionary of backend-specific parameters passed to the vector store constructor. @@ -347,6 +362,36 @@ vector_stores.default.password=your-root-password ReMe service YAML uses the key `vector_stores` (plural); CLI overrides use the same nested paths. +#### 7. HologresVectorStore Configuration + +**Implementation**: [`reme/core/vector_store/hologres_store.py`](https://github.com/agentscope-ai/ReMe/blob/main/reme/core/vector_store/hologres_store.py) + +**Example (Hologres instance)**: + +```yaml +vector_stores: + default: + backend: hologres + embedding_model: default + collection_name: reme + host: "your-hologres-host" + port: 80 + database: "postgres" + user: "postgres" + password: "your-password" + schema: "public" + distance_method: "Cosine" +``` + +```shell +vector_stores.default.backend=hologres +vector_stores.default.host=your-hologres-host +vector_stores.default.port=80 +vector_stores.default.user=postgres +vector_stores.default.password=your-password +vector_stores.default.database=postgres +``` + ### Complete Configuration Example Below is a complete `default.yaml` example including both embedding model and vector store configurations: @@ -404,9 +449,10 @@ Two types of metadata filtering are supported: - **Development & Testing**: Use MemoryVectorStore or LocalVectorStore—no additional services required. - **Small-Scale Applications**: Use LocalVectorStore or ChromaVectorStore for simplicity and ease of use. -- **Production Environments**: Use QdrantVectorStore, EsVectorStore, or ObVecVectorStore (OceanBase/seekdb) for high performance and scalability, depending on your existing infrastructure. +- **Production Environments**: Use QdrantVectorStore, EsVectorStore, ObVecVectorStore (OceanBase/seekdb), or HologresVectorStore for high performance and scalability, depending on your existing infrastructure. - **Hybrid Search**: Use EsVectorStore to combine vector search with full-text search capabilities. - **OceanBase / seekdb**: Use ObVecVectorStore when you standardize on pyobvector and SQL-accessible vector tables. +- **Hologres**: Use HologresVectorStore when you run Hologres and need native HGraph-indexed vector search with PostgreSQL-compatible SQL and JSONB metadata filtering. ## Important Notes diff --git a/reme/core/vector_store/__init__.py b/reme/core/vector_store/__init__.py index 8429b911..c0d1cd55 100644 --- a/reme/core/vector_store/__init__.py +++ b/reme/core/vector_store/__init__.py @@ -3,6 +3,7 @@ from .base_vector_store import BaseVectorStore from .chroma_vector_store import ChromaVectorStore from .es_vector_store import ESVectorStore +from .hologres_store import HologresVectorStore from .local_vector_store import LocalVectorStore from .obvec_vector_store import ObVecVectorStore from .pgvector_store import PGVectorStore @@ -13,6 +14,7 @@ __all__ = [ "BaseVectorStore", "ChromaVectorStore", "ESVectorStore", + "HologresVectorStore", "LocalVectorStore", "ObVecVectorStore", "PGVectorStore", @@ -21,6 +23,7 @@ __all__ = [ R.vector_stores.register("chroma")(ChromaVectorStore) R.vector_stores.register("es")(ESVectorStore) +R.vector_stores.register("hologres")(HologresVectorStore) R.vector_stores.register("local")(LocalVectorStore) R.vector_stores.register("obvec")(ObVecVectorStore) R.vector_stores.register("pgvector")(PGVectorStore) diff --git a/reme/core/vector_store/hologres_store.py b/reme/core/vector_store/hologres_store.py new file mode 100644 index 00000000..270dc1bf --- /dev/null +++ b/reme/core/vector_store/hologres_store.py @@ -0,0 +1,633 @@ +"""Hologres implementation for vector storage and retrieval.""" + +import json +import re +from pathlib import Path +from typing import Any + +from loguru import logger + +from .base_vector_store import BaseVectorStore +from ..embedding import BaseEmbeddingModel +from ..schema import VectorNode + +_ASYNCPG_IMPORT_ERROR: Exception | None = None + +try: + import asyncpg + from asyncpg import Pool +except Exception as e: + _ASYNCPG_IMPORT_ERROR = e + asyncpg = None + Pool = None + + +class HologresVectorStore(BaseVectorStore): + """Vector store implementation using Hologres for efficient similarity search. + + Hologres uses native float4[] arrays for vector storage with built-in + HGraph index for approximate nearest neighbor search, unlike pgvector + which requires an extension. + """ + + @staticmethod + def _validate_table_name(name: str) -> None: + """Validate table name to prevent SQL injection.""" + if not name: + raise ValueError("Table name cannot be empty") + if len(name) > 63: + raise ValueError(f"Table name too long: {len(name)} characters (max 63)") + if not re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", name): + raise ValueError( + f"Invalid table name: {name}. Must start with letter or underscore, " + "and contain only alphanumeric characters and underscores.", + ) + + def __init__( + self, + collection_name: str, + db_path: str | Path, + embedding_model: BaseEmbeddingModel, + host: str = "localhost", + port: int = 80, + database: str = "postgres", + user: str = "postgres", + password: str = "", + schema: str = "public", + min_size: int = 1, + max_size: int = 10, + dsn: str | None = None, + distance_method: str = "Cosine", + **kwargs, + ): + """Initialize the Hologres vector store with connection parameters. + + Args: + collection_name: Name of the collection (table). + db_path: Database path (used by base class). + embedding_model: Embedding model for generating vectors. + host: Hologres host address. + port: Hologres port (default 80 for Hologres). + database: Database name. + user: Database user. + password: Database password. + schema: PostgreSQL schema name (default "public"). + min_size: Minimum connections in pool. + max_size: Maximum connections in pool. + dsn: Full DSN connection string (overrides individual params). + distance_method: Distance method for HGraph index (Cosine, InnerProduct, Euclidean). + """ + if _ASYNCPG_IMPORT_ERROR is not None: + raise ImportError( + "Hologres vector store requires asyncpg. Install with `pip install asyncpg`", + ) from _ASYNCPG_IMPORT_ERROR + + self._validate_table_name(collection_name) + self._validate_table_name(schema) + + super().__init__( + collection_name=collection_name, + db_path=db_path, + embedding_model=embedding_model, + **kwargs, + ) + + self.dsn = dsn + self.host = host + self.port = port + self.database = database + self.user = user + self.password = password + self.schema = schema + self.min_size = min_size + self.max_size = max_size + self.distance_method = distance_method + self._pool: Pool | None = None + self.embedding_model_dims = embedding_model.dimensions + + @property + def _qualified_name(self) -> str: + """Return the schema-qualified table name (e.g. 'my_schema.my_table').""" + return f"{self.schema}.{self.collection_name}" + + def _qualify(self, table_name: str) -> str: + """Return a schema-qualified name for an arbitrary table.""" + return f"{self.schema}.{table_name}" + + @staticmethod + async def _hologres_reset(conn): + """Custom reset for Hologres connections.""" + await conn.execute( + """ + SELECT pg_advisory_unlock_all(); + CLOSE ALL; + RESET ALL; + """, + ) + + 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, + reset=self._hologres_reset, + ) + 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, + reset=self._hologres_reset, + ) + + # Ensure schema exists + async with self._pool.acquire() as conn: + await conn.execute(f"CREATE SCHEMA IF NOT EXISTS {self.schema}") + + logger.info(f"Hologres connection pool created for database {self.database}") + + return self._pool + + @staticmethod + def _vector_to_pg_array(vector: list[float]) -> str: + """Convert a Python list of floats to PostgreSQL array literal format.""" + return "{" + ",".join(map(str, vector)) + "}" + + @staticmethod + def _pg_array_to_vector(pg_array) -> list[float] | None: + """Convert a PostgreSQL array result to a Python list of floats.""" + if pg_array is None: + return None + if isinstance(pg_array, list): + return [float(x) for x in pg_array] + # Handle string format like {1.0,2.0,3.0} + raw = str(pg_array) + if raw.startswith("{") and raw.endswith("}"): + return [float(x) for x in raw[1:-1].split(",")] + return None + + async def list_collections(self) -> list[str]: + """List all available table names in the current schema.""" + 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 = $1", + self.schema, + ) + return [row["table_name"] for row in rows] + + async def create_collection(self, collection_name: str, **kwargs): + """Create a new Hologres table with vector support and HGraph index.""" + self._validate_table_name(collection_name) + pool = await self._get_pool() + dimensions = kwargs.get("dimensions", self.embedding_model_dims) + qualified = self._qualify(collection_name) + + async with pool.acquire() as conn: + create_sql = f""" + CREATE TABLE IF NOT EXISTS {qualified} ( + id TEXT PRIMARY KEY, + content TEXT, + vector float4[] CHECK (array_ndims(vector) = 1 AND array_length(vector, 1) = {dimensions}), + metadata JSONB + ) + WITH ( + vectors = '{{ + "vector": {{ + "algorithm": "HGraph", + "distance_method": "{self.distance_method}", + "builder_params": {{ + "base_quantization_type": "rabitq", + "rabitq_use_fht":true, + "graph_storage_type": "compressed", + "max_total_size_to_merge_mb": 4096, + "max_degree": 64, + "ef_construction": 400, + "precise_quantization_type": "fp32", + "use_reorder": true + }} + }} + }}' + ) + """ + await conn.execute(create_sql) + + logger.info(f"Created Hologres collection {qualified} with dimensions={dimensions}") + + async def delete_collection(self, collection_name: str, **kwargs): + """Remove the specified collection table from the database.""" + self._validate_table_name(collection_name) + pool = await self._get_pool() + qualified = self._qualify(collection_name) + async with pool.acquire() as conn: + await conn.execute(f"DROP TABLE IF EXISTS {qualified}") + logger.info(f"Deleted collection {qualified}") + + async def copy_collection(self, collection_name: str, **kwargs): + """Duplicate the structure and content of the current collection to a new table.""" + self._validate_table_name(collection_name) + pool = await self._get_pool() + qualified_src = self._qualified_name + qualified_dst = self._qualify(collection_name) + + 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 = $2 + """, + self.collection_name, + self.schema, + ) + + if not columns: + raise ValueError(f"Source collection {qualified_src} does not exist") + + # Create new table with primary key, then add data + await conn.execute( + f""" + SET hg_experimental_enable_create_table_like_properties = true; + CALL hg_create_table_like('{qualified_dst}', 'select * from {qualified_src}') + """, + ) + await conn.execute(f"INSERT INTO {qualified_dst} SELECT * FROM {qualified_src} ;") + + logger.info(f"Copied collection {qualified_src} to {qualified_dst}") + + async def insert(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Insert or upsert vector nodes into the Hologres collection.""" + 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, + 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._qualified_name} (id, content, vector, metadata) + VALUES ($1, $2, $3::float4[], $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._qualified_name} (id, content, vector, metadata) + VALUES ($1, $2, $3::float4[], $4::jsonb) + ON CONFLICT (id) DO NOTHING + """, + data, + ) + else: + await conn.executemany( + f""" + INSERT INTO {self._qualified_name} (id, content, vector, metadata) + VALUES ($1, $2, $3::float4[], $4::jsonb) + """, + data, + ) + + logger.info(f"Inserted {len(nodes_to_insert)} documents into {self._qualified_name}") + + @staticmethod + def _build_filter_clause(filters: dict | None) -> tuple[str, list]: + """Generate an SQL WHERE clause and parameter list from a filter dictionary. + + Supports two filter formats: + 1. Range query: {"field": [start_value, end_value]} + 2. Exact match: {"field": value} + """ + if not filters: + return "", [] + + conditions = [] + params = [] + param_idx = 1 + + for key, value in filters.items(): + if not key.replace("_", "").replace(".", "").isalnum(): + raise ValueError( + f"Invalid metadata key: {key}. Only alphanumeric characters, underscore and dot are allowed.", + ) + + if isinstance(value, list) and len(value) == 2: + if isinstance(value[0], (int, float)) and isinstance(value[1], (int, float)): + conditions.append( + f"(metadata->>'{key}')::numeric >= ${param_idx} AND " + f"(metadata->>'{key}')::numeric <= ${param_idx + 1}", + ) + else: + conditions.append(f"metadata->>'{key}' >= ${param_idx} AND metadata->>'{key}' <= ${param_idx + 1}") + params.extend([value[0], value[1]]) + param_idx += 2 + 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 using Hologres approx_cosine_distance.""" + query_vector = await self.get_embedding(query) + vector_str = self._vector_to_pg_array(query_vector) + pool = await self._get_pool() + + filter_clause, filter_params = self._build_filter_clause(filters) + + # filter_params use $1..$N, limit uses $(N+1) + limit_placeholder = f"${len(filter_params) + 1}" + + async with pool.acquire() as conn: + sql = f""" + SELECT id, content, vector, metadata, + approx_cosine_distance(vector, '{vector_str}') AS distance + FROM {self._qualified_name} + {filter_clause} + ORDER BY distance DESC + LIMIT {limit_placeholder} + """ + rows = await conn.fetch(sql, *filter_params, limit) + + results = [] + score_threshold = kwargs.get("score_threshold") + + for row in rows: + distance = float(row["distance"]) + # approx_cosine_distance returns cosine similarity (higher = more similar) + score = distance + if score_threshold is not None and score < score_threshold: + continue + + vector_data = self._pg_array_to_vector(row["vector"]) + + metadata = row["metadata"] if row["metadata"] else {} + if isinstance(metadata, str): + metadata = json.loads(metadata) + + metadata["score"] = score + metadata["_distance"] = 1 - score + + 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.""" + 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._qualified_name} WHERE id IN ({placeholders})", + *vector_ids, + ) + + logger.info(f"Deleted {len(vector_ids)} documents from {self._qualified_name}") + + async def delete_all(self, **kwargs): + """Remove all vectors from the collection.""" + pool = await self._get_pool() + async with pool.acquire() as conn: + result = await conn.execute(f"DELETE FROM {self._qualified_name}") + + logger.info(f"Deleted all documents from {self._qualified_name} result={result}") + + async def update(self, nodes: VectorNode | list[VectorNode], **kwargs): + """Update existing vector nodes with new content, embeddings, or metadata.""" + 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: + update_fields.append(f"vector = ${idx}::float4[]") + params.append(node.vector) + 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._qualified_name} SET {', '.join(update_fields)} WHERE id = ${idx}", + *params, + ) + + logger.info(f"Updated {len(nodes_to_update)} documents in {self._qualified_name}") + + async def get(self, vector_ids: str | list[str]) -> VectorNode | list[VectorNode] | None: + """Retrieve vector nodes by their unique identifiers.""" + 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._qualified_name} WHERE id IN ({placeholders})", + *vector_ids, + ) + + results = [] + for row in rows: + vector_data = self._pg_array_to_vector(row["vector"]) + + 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, + sort_key: str | None = None, + reverse: bool = False, + ) -> list[VectorNode]: + """Return a list of vector nodes matching the provided filters and limit. + + Args: + filters: Dictionary of filter conditions to match vectors + limit: Maximum number of vectors to return + sort_key: Key to sort the results by (e.g., field name in metadata). None for no sorting + reverse: If True, sort in descending order; if False, sort in ascending order + """ + pool = await self._get_pool() + filter_clause, filter_params = self._build_filter_clause(filters) + + order_clause = "" + if sort_key: + order_direction = "DESC" if reverse else "ASC" + order_clause = f"ORDER BY metadata->>'{sort_key}' {order_direction}" + + 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._qualified_name} + {filter_clause} + {order_clause} + {limit_clause} + """ + rows = await conn.fetch(sql, *filter_params) + + results = [] + for row in rows: + vector_data = self._pg_array_to_vector(row["vector"]) + + 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() + qualified = self._qualified_name + + async with pool.acquire() as conn: + count = await conn.fetchval(f"SELECT COUNT(*) FROM {qualified}") + size = await conn.fetchval(f"SELECT pg_size_pretty(pg_total_relation_size('{qualified}'))") + + return { + "name": qualified, + "count": count, + "size": size, + } + + async def reset(self): + """Purge all data by dropping and recreating the collection table.""" + logger.warning(f"Resetting collection {self._qualified_name}...") + await self.delete_collection(self.collection_name) + await self.create_collection(self.collection_name) + + async def reset_collection(self, collection_name: str): + """Reset collection with table name validation.""" + self._validate_table_name(collection_name) + self.collection_name = collection_name + await self.create_collection(collection_name) + logger.info(f"Collection reset to {self._qualified_name}") + + async def start(self) -> None: + """Initialize the PGVector store. + + Creates the connection pool and ensures the collection table exists. + """ + await self._get_pool() + await super().start() + logger.info(f"Hologres collection {self._qualified_name} initialized") + + async def close(self): + """Terminate the database connection pool.""" + if self._pool is not None: + await self._pool.close() + self._pool = None + logger.info("Hologres connection pool closed") diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 7f7264be..9fab656b 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -2,8 +2,8 @@ """Unified test suite for vector store implementations. This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore, -PGVectorStore, QdrantVectorStore, ChromaVectorStore, and ObVecVectorStore implementations. -Tests can be run for specific vector stores or all implementations. +PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, and HologresVectorStore +implementations. Tests can be run for specific vector stores or all implementations. Usage: python test_vector_store.py --local # Test LocalVectorStore only @@ -12,6 +12,7 @@ Usage: python test_vector_store.py --qdrant # Test QdrantVectorStore only python test_vector_store.py --chroma # Test ChromaVectorStore only python test_vector_store.py --obvec # Test ObVecVectorStore only (needs seekdb / OceanBase) + python test_vector_store.py --hologres # Test HologresVectorStore only python test_vector_store.py --all # Test all vector stores """ @@ -31,6 +32,7 @@ from reme.core.utils import load_env, cosine_similarity from reme.core.vector_store import ( BaseVectorStore, ChromaVectorStore, + HologresVectorStore, LocalVectorStore, ESVectorStore, ObVecVectorStore, @@ -90,6 +92,20 @@ class TestConfig: OBVEC_PASSWORD = os.environ.get("OBVEC_PASSWORD", "root") OBVEC_DATABASE = os.environ.get("OBVEC_DATABASE", "test") + # HologresVectorStore settings + HOLOGRES_DSN = os.environ.get( + "HOLOGRES_DSN", + "", + ) # Full DSN connection string (overrides host/port/database/user/password) + HOLOGRES_HOST = os.environ.get("HOLOGRES_HOST", "localhost") + HOLOGRES_PORT = int(os.environ.get("HOLOGRES_PORT", "80")) + HOLOGRES_DATABASE = os.environ.get("HOLOGRES_DATABASE", "postgres") + HOLOGRES_USER = os.environ.get("HOLOGRES_USER", "postgres") + HOLOGRES_PASSWORD = os.environ.get("HOLOGRES_PASSWORD", "") + HOLOGRES_SCHEMA = os.environ.get("HOLOGRES_SCHEMA", "public") + HOLOGRES_MIN_SIZE = 1 + HOLOGRES_MAX_SIZE = 5 + # Embedding model settings EMBEDDING_MODEL_NAME = "text-embedding-v4" EMBEDDING_DIMENSIONS = 64 @@ -192,6 +208,17 @@ class SampleDataGenerator: # ==================== Vector Store Factory ==================== +_STORE_TYPE_MAP = { + LocalVectorStore: "local", + QdrantVectorStore: "qdrant", + ESVectorStore: "es", + PGVectorStore: "pgvector", + ChromaVectorStore: "chroma", + ObVecVectorStore: "obvec", + HologresVectorStore: "hologres", +} + + def get_store_type(store: BaseVectorStore) -> str: """Get the type identifier of a vector store instance. @@ -199,29 +226,107 @@ def get_store_type(store: BaseVectorStore) -> str: store: Vector store instance Returns: - str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", or "obvec") + str: Type identifier ("local", "es", "pgvector", "qdrant", "chroma", "obvec", or "hologres") """ - 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" - elif isinstance(store, ObVecVectorStore): - return "obvec" - else: - raise ValueError(f"Unknown vector store type: {type(store)}") + for cls, name in _STORE_TYPE_MAP.items(): + if isinstance(store, cls): + return name + raise ValueError(f"Unknown vector store type: {type(store)}") + + +def _build_store_kwargs(store_type, config, collection_name, embedding_model): + """Build keyword arguments for creating a vector store instance.""" + store_kwargs_map = { + "local": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": config.LOCAL_ROOT_PATH, + }, + "es": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": tempfile.mkdtemp(prefix="test_es_"), + "hosts": config.ES_HOSTS, + "basic_auth": config.ES_BASIC_AUTH, + }, + "qdrant": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": config.QDRANT_PATH or tempfile.mkdtemp(prefix="test_qdrant_"), + "host": config.QDRANT_HOST, + "port": config.QDRANT_PORT, + "url": config.QDRANT_URL, + "api_key": config.QDRANT_API_KEY, + "distance": "cosine", + "on_disk": False, + }, + "pgvector": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": tempfile.mkdtemp(prefix="test_pgvector_"), + "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, + }, + "chroma": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_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, + }, + "obvec": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": tempfile.mkdtemp(prefix="test_obvec_"), + "uri": config.OBVEC_URI, + "user": config.OBVEC_USER, + "password": config.OBVEC_PASSWORD, + "database": config.OBVEC_DATABASE, + "index_metric": "cosine", + "index_ef_search": 100, + }, + "hologres": lambda: { + "collection_name": collection_name, + "embedding_model": embedding_model, + "db_path": tempfile.mkdtemp(prefix="test_hologres_"), + "host": config.HOLOGRES_HOST, + "port": config.HOLOGRES_PORT, + "database": config.HOLOGRES_DATABASE, + "user": config.HOLOGRES_USER, + "password": config.HOLOGRES_PASSWORD, + "schema": config.HOLOGRES_SCHEMA, + "min_size": config.HOLOGRES_MIN_SIZE, + "max_size": config.HOLOGRES_MAX_SIZE, + **({"dsn": config.HOLOGRES_DSN} if config.HOLOGRES_DSN else {}), + }, + } + if store_type not in store_kwargs_map: + raise ValueError(f"Unknown store type: {store_type}") + return store_kwargs_map[store_type]() + + +_STORE_CLASS_MAP = { + "local": LocalVectorStore, + "es": ESVectorStore, + "qdrant": QdrantVectorStore, + "pgvector": PGVectorStore, + "chroma": ChromaVectorStore, + "obvec": ObVecVectorStore, + "hologres": HologresVectorStore, +} 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", "pgvector", "qdrant", "chroma", or "obvec") + store_type: Type of vector store ("local", "es", "pgvector", "qdrant", "chroma", "obvec", or "hologres") collection_name: Name of the collection Returns: @@ -235,68 +340,9 @@ def create_vector_store(store_type: str, collection_name: str) -> BaseVectorStor dimensions=config.EMBEDDING_DIMENSIONS, ) - if store_type == "local": - return LocalVectorStore( - collection_name=collection_name, - embedding_model=embedding_model, - db_path=config.LOCAL_ROOT_PATH, - ) - elif store_type == "es": - return ESVectorStore( - collection_name=collection_name, - embedding_model=embedding_model, - db_path=tempfile.mkdtemp(prefix="test_es_"), - hosts=config.ES_HOSTS, - basic_auth=config.ES_BASIC_AUTH, - ) - elif store_type == "qdrant": - return QdrantVectorStore( - collection_name=collection_name, - embedding_model=embedding_model, - db_path=config.QDRANT_PATH or tempfile.mkdtemp(prefix="test_qdrant_"), - 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, - db_path=tempfile.mkdtemp(prefix="test_pgvector_"), - 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, - db_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, - ) - elif store_type == "obvec": - return ObVecVectorStore( - collection_name=collection_name, - embedding_model=embedding_model, - db_path=tempfile.mkdtemp(prefix="test_obvec_"), - uri=config.OBVEC_URI, - user=config.OBVEC_USER, - password=config.OBVEC_PASSWORD, - database=config.OBVEC_DATABASE, - index_metric="cosine", - index_ef_search=100, - ) - else: - raise ValueError(f"Unknown store type: {store_type}") + kwargs = _build_store_kwargs(store_type, config, collection_name, embedding_model) + store_cls = _STORE_CLASS_MAP[store_type] + return store_cls(**kwargs) # ==================== Test Functions ==================== @@ -614,7 +660,7 @@ async def test_copy_collection(store: BaseVectorStore, store_name: str): # Elasticsearch, PostgreSQL and OceanBase require lowercase table/index names store_type = get_store_type(store) - if store_type in ("es", "pgvector", "obvec"): + if store_type in ("es", "pgvector", "obvec", "hologres"): copy_collection_name = copy_collection_name.lower() # Clean up if exists @@ -1811,6 +1857,7 @@ Examples: python test_vector_store.py --qdrant # Test QdrantVectorStore only python test_vector_store.py --chroma # Test ChromaVectorStore only python test_vector_store.py --obvec # Test ObVecVectorStore (seekdb / OceanBase) + python test_vector_store.py --hologres # Test HologresVectorStore python test_vector_store.py --all # Test all vector stores """, ) @@ -1844,6 +1891,11 @@ Examples: action="store_true", help="Test ObVecVectorStore", ) + parser.add_argument( + "--hologres", + action="store_true", + help="Test HologresVectorStore", + ) parser.add_argument( "--all", action="store_true", @@ -1863,6 +1915,7 @@ Examples: ("qdrant", "QdrantVectorStore"), ("chroma", "ChromaVectorStore"), ("obvec", "ObVecVectorStore"), + ("hologres", "HologresVectorStore"), ] else: # Build list based on individual flags @@ -1878,6 +1931,8 @@ Examples: stores_to_test.append(("chroma", "ChromaVectorStore")) if args.obvec: stores_to_test.append(("obvec", "ObVecVectorStore")) + if args.hologres: + stores_to_test.append(("hologres", "HologresVectorStore")) if not stores_to_test: # Default to all vector stores if no argument provided @@ -1888,10 +1943,11 @@ Examples: ("qdrant", "QdrantVectorStore"), ("chroma", "ChromaVectorStore"), ("obvec", "ObVecVectorStore"), + ("hologres", "HologresVectorStore"), ] print("No vector store specified, defaulting to test all vector stores") print( - "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec to test specific ones\n", + "Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--hologres to test specific ones\n", ) # Run tests for each vector store