mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
* feat(seekdb): add Seekdb file and vector stores with pyseekdb>=1.2.0 * refactor(seekdb): add pyseekdb_conn and remote-only host/port config * refactor(embedding): remove env fallbacks from BaseEmbeddingModel; pass credentials in tests * refactor(seekdb): drop tenant from client kwargs; default database test and empty password * fix(deps): gate pyseekdb to Python >=3.11 for CI 3.10 compatibility * fix(seekdb): satisfy pre-commit pylint and formatting for seekdb stores
2070 lines
77 KiB
Python
2070 lines
77 KiB
Python
# pylint: disable=too-many-lines
|
|
"""Unified test suite for vector store implementations.
|
|
|
|
This module provides comprehensive test coverage for LocalVectorStore, ESVectorStore,
|
|
PGVectorStore, QdrantVectorStore, ChromaVectorStore, ObVecVectorStore, HologresVectorStore,
|
|
ZvecVectorStore, and SeekdbVectorStore 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 --obvec # Test ObVecVectorStore only (needs seekdb / OceanBase)
|
|
python test_vector_store.py --hologres # Test HologresVectorStore only
|
|
python test_vector_store.py --zvec # Test ZvecVectorStore only
|
|
python test_vector_store.py --seekdb # Test SeekdbVectorStore (pyseekdb; Python >=3.11)
|
|
python test_vector_store.py --all # Test all vector stores
|
|
"""
|
|
|
|
import argparse
|
|
import asyncio
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
from loguru import logger
|
|
|
|
from reme.core.embedding import OpenAIEmbeddingModel
|
|
from reme.core.schema import VectorNode
|
|
from reme.core.utils import load_env, cosine_similarity
|
|
from reme.core.vector_store import (
|
|
BaseVectorStore,
|
|
ChromaVectorStore,
|
|
HologresVectorStore,
|
|
LocalVectorStore,
|
|
ESVectorStore,
|
|
ObVecVectorStore,
|
|
PGVectorStore,
|
|
QdrantVectorStore,
|
|
ZvecVectorStore,
|
|
)
|
|
|
|
try:
|
|
from reme.core.vector_store import SeekdbVectorStore
|
|
|
|
SEEKDB_AVAILABLE = True
|
|
except ImportError:
|
|
SeekdbVectorStore = None
|
|
SEEKDB_AVAILABLE = False
|
|
|
|
load_env()
|
|
|
|
|
|
def _search_score_for_log(metadata: dict) -> object:
|
|
"""Similarity score for log lines (implementations use ``metadata['score']``)."""
|
|
return metadata.get("score", metadata.get("_score", "N/A"))
|
|
|
|
|
|
# ==================== 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)
|
|
|
|
# SeekdbVectorStore: embedded by default; set SEEKDB_HOST (and optional SEEKDB_PORT) for remote
|
|
SEEKDB_PATH = None # e.g. "./test_vector_store_seekdb"; None => temp dir
|
|
SEEKDB_HOST = os.environ.get("SEEKDB_HOST")
|
|
SEEKDB_PORT = int(os.environ["SEEKDB_PORT"]) if os.environ.get("SEEKDB_PORT") else None
|
|
SEEKDB_USER = os.environ.get("SEEKDB_USER", "root")
|
|
SEEKDB_PASSWORD = os.environ.get("SEEKDB_PASSWORD", "")
|
|
SEEKDB_DATABASE = os.environ.get("SEEKDB_DATABASE", "test")
|
|
|
|
# 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
|
|
|
|
# ObVecVectorStore: seekdb docker often uses user `root` + ROOT_PASSWORD; OceanBase
|
|
# multi-tenant commonly uses `root@<tenant>` (see pyobvector defaults).
|
|
# OBVEC_PASSWORD default `root` matches docker-compose.obvec.yml only—override if your
|
|
# seekdb uses another ROOT_PASSWORD (e.g. another compose stack on the same port).
|
|
OBVEC_URI = os.environ.get("OBVEC_URI", "127.0.0.1:2881")
|
|
OBVEC_USER = os.environ.get("OBVEC_USER", "root")
|
|
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
|
|
# ZvecVectorStore settings
|
|
ZVEC_PATH = "./test_vector_store_zvec" # For local persistent mode
|
|
|
|
# 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 ====================
|
|
|
|
|
|
# pylint: disable=too-many-return-statements
|
|
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", "chroma", "obvec", "zvec", "hologres", or "seekdb"
|
|
)
|
|
"""
|
|
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"
|
|
elif isinstance(store, ZvecVectorStore):
|
|
return "zvec"
|
|
elif isinstance(store, HologresVectorStore):
|
|
return "hologres"
|
|
elif SeekdbVectorStore is not None and isinstance(store, SeekdbVectorStore):
|
|
return "seekdb"
|
|
else:
|
|
raise ValueError(f"Unknown vector store type: {type(store)}")
|
|
|
|
|
|
# pylint: disable=too-many-return-statements
|
|
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", "obvec", "zvec", "hologres", or "seekdb"
|
|
)
|
|
collection_name: Name of the collection
|
|
|
|
Returns:
|
|
BaseVectorStore: Initialized vector store instance
|
|
"""
|
|
config = TestConfig()
|
|
|
|
# api_key/base_url from env (Application layer does the same via embedding_api_key in application.py)
|
|
embedding_api_key = os.environ.get("EMBEDDING_API_KEY") or os.environ.get("OPENAI_API_KEY")
|
|
embedding_base_url = os.environ.get("EMBEDDING_BASE_URL") or os.environ.get("OPENAI_BASE_URL")
|
|
# use_dimensions=True so API output matches HNSW dim / vector column
|
|
embedding_model = OpenAIEmbeddingModel(
|
|
model_name=config.EMBEDDING_MODEL_NAME,
|
|
dimensions=config.EMBEDDING_DIMENSIONS,
|
|
use_dimensions=True,
|
|
api_key=embedding_api_key or None,
|
|
base_url=embedding_base_url or None,
|
|
)
|
|
|
|
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,
|
|
)
|
|
elif store_type == "zvec":
|
|
return ZvecVectorStore(
|
|
collection_name=collection_name,
|
|
embedding_model=embedding_model,
|
|
db_path=config.ZVEC_PATH or tempfile.mkdtemp(prefix="test_zvec_"),
|
|
dimension=config.EMBEDDING_DIMENSIONS,
|
|
distance="cosine",
|
|
)
|
|
elif store_type == "hologres":
|
|
kwargs = {
|
|
"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,
|
|
}
|
|
if config.HOLOGRES_DSN:
|
|
kwargs["dsn"] = config.HOLOGRES_DSN
|
|
return HologresVectorStore(**kwargs)
|
|
elif store_type == "seekdb":
|
|
if SeekdbVectorStore is None:
|
|
raise ImportError("SeekdbVectorStore not available; install pyseekdb (Python >=3.11)")
|
|
remote = bool(config.SEEKDB_HOST and config.SEEKDB_HOST.strip())
|
|
db_path = config.SEEKDB_PATH or tempfile.mkdtemp(prefix="test_seekdb_")
|
|
kw: dict = {
|
|
"collection_name": collection_name,
|
|
"embedding_model": embedding_model,
|
|
"db_path": db_path,
|
|
"database": config.SEEKDB_DATABASE,
|
|
"distance": "cosine",
|
|
}
|
|
if remote:
|
|
kw["host"] = config.SEEKDB_HOST.strip()
|
|
kw["port"] = config.SEEKDB_PORT
|
|
kw["user"] = config.SEEKDB_USER
|
|
kw["password"] = config.SEEKDB_PASSWORD
|
|
else:
|
|
kw["path"] = str(Path(db_path) / "seekdb.db")
|
|
return SeekdbVectorStore(**kw)
|
|
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 = _search_score_for_log(r.metadata)
|
|
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_exact_match_filter(store: BaseVectorStore, _store_name: str):
|
|
"""Test vector search with exact match filter."""
|
|
logger.info("=" * 20 + " EXACT MATCH FILTER SEARCH TEST " + "=" * 20)
|
|
|
|
# Test exact match 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("✓ Exact match 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 exact match filters (AND operation)
|
|
filters = {
|
|
"node_type": "tech",
|
|
"source": "research",
|
|
"priority": "high",
|
|
}
|
|
results = await store.search(
|
|
query="What is artificial intelligence?",
|
|
limit=5,
|
|
filters=filters,
|
|
)
|
|
|
|
logger.info(
|
|
f"Multi-filter search (node_type=tech AND source=research AND priority=high) "
|
|
f"returned {len(results)} results",
|
|
)
|
|
for i, r in enumerate(results, 1):
|
|
node_type = r.metadata.get("node_type")
|
|
source = r.metadata.get("source")
|
|
priority = r.metadata.get("priority")
|
|
logger.info(f" Result {i}: type={node_type}, source={source}, priority={priority}")
|
|
assert node_type == "tech", "Result should have node_type='tech'"
|
|
assert source == "research", "Result should have source='research'"
|
|
assert priority == "high", "Result should have priority='high'"
|
|
|
|
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, PostgreSQL and OceanBase require lowercase table/index names
|
|
store_type = get_store_type(store)
|
|
if store_type in ("es", "pgvector", "obvec", "hologres"):
|
|
copy_collection_name = copy_collection_name.lower()
|
|
# seekdb uses collection names as-is
|
|
|
|
# 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)
|
|
await copied_store.start()
|
|
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 = 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 = 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 = 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 exact match filters
|
|
filters_1 = {
|
|
"domain": "AI",
|
|
"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: {len(results_1)} results")
|
|
for r in results_1:
|
|
assert r.metadata.get("domain") == "AI"
|
|
assert r.metadata.get("impact_factor") == "high"
|
|
|
|
# Test 2: Single exact match filter
|
|
filters_2 = {
|
|
"subdomain": "nlp",
|
|
}
|
|
results_2 = await store.search(
|
|
query="deep learning applications",
|
|
limit=10,
|
|
filters=filters_2,
|
|
)
|
|
logger.info(f"Test 2 - NLP subdomain: {len(results_2)} results")
|
|
for r in results_2:
|
|
assert r.metadata.get("subdomain") == "nlp"
|
|
|
|
# Test 3: Year-based exact match 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 = _search_score_for_log(result.metadata)
|
|
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 = _search_score_for_log(result.metadata)
|
|
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 exact match filter
|
|
results_2 = await store.search(
|
|
query="technology",
|
|
filters={"node_type": "tech"},
|
|
limit=10,
|
|
)
|
|
logger.info(f"Test 2 - Single exact match filter: {len(results_2)} results")
|
|
for r in results_2:
|
|
assert r.metadata.get("node_type") == "tech"
|
|
|
|
# Test 3: Multiple exact match filters (AND operation)
|
|
results_3 = await store.search(
|
|
query="technology",
|
|
filters={
|
|
"node_type": "tech",
|
|
"source": "research",
|
|
"priority": "high",
|
|
},
|
|
limit=10,
|
|
)
|
|
logger.info(f"Test 3 - Multiple exact match filters (AND): {len(results_3)} results")
|
|
for r in results_3:
|
|
assert r.metadata.get("node_type") == "tech"
|
|
assert r.metadata.get("source") == "research"
|
|
assert r.metadata.get("priority") == "high"
|
|
|
|
# Test 4: Filter with non-existent value
|
|
results_4 = await store.search(
|
|
query="technology",
|
|
filters={"category": "NON_EXISTENT_CATEGORY"},
|
|
limit=10,
|
|
)
|
|
logger.info(f"Test 4 - Non-existent filter value: {len(results_4)} results")
|
|
assert len(results_4) == 0, "Should return no results for non-existent filter value"
|
|
|
|
# Test 5: List operation with multiple exact match filters
|
|
list_results = await store.list(
|
|
filters={"node_type": "tech", "priority": "high"},
|
|
limit=20,
|
|
)
|
|
logger.info(f"Test 5 - List with multiple 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")
|
|
|
|
|
|
async def test_range_query_filters(store: BaseVectorStore, _store_name: str):
|
|
"""Test range query filters using the new [start, end] syntax."""
|
|
logger.info("=" * 20 + " RANGE QUERY FILTERS TEST " + "=" * 20)
|
|
|
|
# Clean up any existing test data first
|
|
try:
|
|
existing_nodes = await store.list(filters={"test_type": "range_query_test"})
|
|
if existing_nodes:
|
|
await store.delete([node.vector_id for node in existing_nodes])
|
|
logger.info(f"Cleaned up {len(existing_nodes)} existing test nodes")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to clean up existing nodes: {e}")
|
|
|
|
# Create test nodes with numeric metadata for range queries
|
|
import time
|
|
|
|
base_timestamp = int(time.time())
|
|
test_nodes = []
|
|
|
|
for i in range(20):
|
|
node = VectorNode(
|
|
vector_id=f"range_node_{i}",
|
|
content=f"Test content for range query node {i}",
|
|
metadata={
|
|
"test_type": "range_query_test",
|
|
"timestamp": base_timestamp + i * 1000, # Each node is 1000 seconds apart
|
|
"rating": 50 + i * 2, # Ratings from 50 to 88
|
|
"priority": i % 3, # 0, 1, or 2
|
|
"category": ["tech", "science", "business"][i % 3],
|
|
},
|
|
)
|
|
test_nodes.append(node)
|
|
|
|
# Insert test nodes
|
|
await store.insert(test_nodes)
|
|
logger.info(f"Inserted {len(test_nodes)} test nodes with numeric metadata")
|
|
|
|
# Test 1: Range query on timestamp field
|
|
start_time = base_timestamp + 5000
|
|
end_time = base_timestamp + 15000
|
|
results_1 = await store.search(
|
|
query="test content",
|
|
limit=20,
|
|
filters={
|
|
"timestamp": [start_time, end_time], # Range query: >= start_time AND <= end_time
|
|
},
|
|
)
|
|
logger.info(f"Test 1 - Timestamp range [{start_time}, {end_time}]: {len(results_1)} results")
|
|
|
|
# Verify all results are within range
|
|
for r in results_1:
|
|
ts = r.metadata.get("timestamp")
|
|
assert ts >= start_time, f"Timestamp {ts} should be >= {start_time}"
|
|
assert ts <= end_time, f"Timestamp {ts} should be <= {end_time}"
|
|
logger.debug(f" Node {r.vector_id}: timestamp={ts}")
|
|
|
|
# Expected nodes: range_node_5 to range_node_15 (11 nodes)
|
|
assert len(results_1) >= 10, f"Expected at least 10 results, got {len(results_1)}"
|
|
logger.info("✓ Timestamp range query validated")
|
|
|
|
# Test 2: Range query on rating field
|
|
results_2 = await store.search(
|
|
query="test content",
|
|
limit=20,
|
|
filters={
|
|
"rating": [60, 80], # Range query: rating >= 60 AND rating <= 80
|
|
},
|
|
)
|
|
logger.info(f"Test 2 - Rating range [60, 80]: {len(results_2)} results")
|
|
|
|
# Verify all results are within rating range
|
|
for r in results_2:
|
|
rating = r.metadata.get("rating")
|
|
assert rating >= 60, f"Rating {rating} should be >= 60"
|
|
assert rating <= 80, f"Rating {rating} should be <= 80"
|
|
logger.debug(f" Node {r.vector_id}: rating={rating}")
|
|
|
|
# Expected: ratings from 60 to 80 (nodes 5-15)
|
|
assert len(results_2) >= 10, f"Expected at least 10 results, got {len(results_2)}"
|
|
logger.info("✓ Rating range query validated")
|
|
|
|
# Test 3: Combine range query with exact match filter
|
|
results_3 = await store.search(
|
|
query="test content",
|
|
limit=20,
|
|
filters={
|
|
"timestamp": [start_time, end_time],
|
|
"category": "tech", # Exact match
|
|
},
|
|
)
|
|
logger.info(
|
|
f"Test 3 - Timestamp range + exact match (category=tech): {len(results_3)} results",
|
|
)
|
|
|
|
# Verify filters
|
|
for r in results_3:
|
|
ts = r.metadata.get("timestamp")
|
|
category = r.metadata.get("category")
|
|
assert start_time <= ts <= end_time, "Timestamp should be in range"
|
|
assert category == "tech", f"Category should be 'tech', got '{category}'"
|
|
logger.debug(f" Node {r.vector_id}: timestamp={ts}, category={category}")
|
|
|
|
# Expected: nodes within range AND category=tech
|
|
assert len(results_3) >= 3, f"Expected at least 3 results, got {len(results_3)}"
|
|
logger.info("✓ Combined range + exact match query validated")
|
|
|
|
# Test 4: Multiple range queries
|
|
results_4 = await store.search(
|
|
query="test content",
|
|
limit=20,
|
|
filters={
|
|
"timestamp": [base_timestamp + 8000, base_timestamp + 12000],
|
|
"rating": [65, 75],
|
|
},
|
|
)
|
|
logger.info(f"Test 4 - Multiple range queries: {len(results_4)} results")
|
|
|
|
# Verify both ranges
|
|
for r in results_4:
|
|
ts = r.metadata.get("timestamp")
|
|
rating = r.metadata.get("rating")
|
|
assert base_timestamp + 8000 <= ts <= base_timestamp + 12000, "Timestamp out of range"
|
|
assert 65 <= rating <= 75, f"Rating {rating} out of range [65, 75]"
|
|
logger.debug(f" Node {r.vector_id}: timestamp={ts}, rating={rating}")
|
|
|
|
# Expected: nodes 8-12 (5 nodes) with overlapping ranges
|
|
assert len(results_4) >= 3, f"Expected at least 3 results, got {len(results_4)}"
|
|
logger.info("✓ Multiple range queries validated")
|
|
|
|
# Test 5: Range query with list operation
|
|
results_5 = await store.list(
|
|
filters={
|
|
"rating": [60, 70],
|
|
"test_type": "range_query_test",
|
|
},
|
|
limit=20,
|
|
)
|
|
logger.info(f"Test 5 - Range query in list operation: {len(results_5)} results")
|
|
|
|
# Verify rating range in list results
|
|
for r in results_5:
|
|
rating = r.metadata.get("rating")
|
|
assert 60 <= rating <= 70, f"Rating {rating} should be in range [60, 70]"
|
|
|
|
logger.info("✓ Range query in list operation validated")
|
|
|
|
# Test 6: Edge case - exact boundary values
|
|
results_6 = await store.list(
|
|
filters={
|
|
"rating": [60, 60], # Exact match using range syntax
|
|
"test_type": "range_query_test",
|
|
},
|
|
limit=20,
|
|
)
|
|
logger.info(f"Test 6 - Exact value using range syntax [60, 60]: {len(results_6)} results")
|
|
|
|
# Should return exactly one node (range_node_5 with rating=60)
|
|
for r in results_6:
|
|
rating = r.metadata.get("rating")
|
|
assert rating == 60, f"Rating should be exactly 60, got {rating}"
|
|
|
|
logger.info("✓ Boundary value range query validated")
|
|
|
|
# Test 7: Range query with sorting
|
|
results_7 = await store.list(
|
|
filters={
|
|
"rating": [60, 80],
|
|
"test_type": "range_query_test",
|
|
},
|
|
sort_key="rating",
|
|
reverse=True,
|
|
limit=5,
|
|
)
|
|
logger.info(f"Test 7 - Range query with sorting: {len(results_7)} results")
|
|
|
|
# Verify results are sorted and within range
|
|
for i in range(len(results_7) - 1):
|
|
rating1 = results_7[i].metadata.get("rating")
|
|
rating2 = results_7[i + 1].metadata.get("rating")
|
|
assert rating1 >= rating2, f"Results not sorted: {rating1} < {rating2}"
|
|
assert 60 <= rating1 <= 80, "Rating out of range"
|
|
|
|
logger.info("✓ Range query with sorting validated")
|
|
|
|
# Clean up test data
|
|
await store.delete([node.vector_id for node in test_nodes])
|
|
logger.info("Cleaned up test nodes")
|
|
|
|
logger.info("✓ Range query filters test passed")
|
|
|
|
|
|
async def test_string_range_queries(store: BaseVectorStore, store_name: str):
|
|
"""Test range queries with string values (e.g., date strings, timestamps)."""
|
|
logger.info("=" * 20 + " STRING RANGE QUERIES TEST " + "=" * 20)
|
|
|
|
# Skip this test for stores that don't support string range queries properly
|
|
# Qdrant and ChromaDB only support numeric range queries, not string range queries
|
|
if store_name not in ["PGVectorStore", "LocalVectorStore", "ESVectorStore"]:
|
|
logger.info(f"Skipping string range query test for {store_name}")
|
|
return
|
|
|
|
# Clean up any existing test data first
|
|
try:
|
|
existing_nodes = await store.list(filters={"test_type": "string_range_test"})
|
|
if existing_nodes:
|
|
await store.delete([node.vector_id for node in existing_nodes])
|
|
logger.info(f"Cleaned up {len(existing_nodes)} existing test nodes")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to clean up existing nodes: {e}")
|
|
|
|
# Create test nodes with string date metadata
|
|
test_nodes = []
|
|
dates = [
|
|
"2024-01-01",
|
|
"2024-01-15",
|
|
"2024-02-01",
|
|
"2024-02-15",
|
|
"2024-03-01",
|
|
"2024-03-15",
|
|
"2024-04-01",
|
|
]
|
|
|
|
for i, date in enumerate(dates):
|
|
node = VectorNode(
|
|
vector_id=f"string_range_node_{i}",
|
|
content=f"Test content for date {date}",
|
|
metadata={
|
|
"test_type": "string_range_test",
|
|
"date": date,
|
|
"index": i,
|
|
},
|
|
)
|
|
test_nodes.append(node)
|
|
|
|
# Insert test nodes
|
|
await store.insert(test_nodes)
|
|
logger.info(f"Inserted {len(test_nodes)} test nodes with string dates")
|
|
|
|
# Test 1: String range query on date field
|
|
try:
|
|
results = await store.search(
|
|
query="test content",
|
|
limit=20,
|
|
filters={
|
|
"date": ["2024-02-01", "2024-03-15"], # Range query on string dates
|
|
},
|
|
)
|
|
logger.info(f"Test 1 - String date range ['2024-02-01', '2024-03-15']: {len(results)} results")
|
|
|
|
# Verify all results are within range
|
|
for r in results:
|
|
date = r.metadata.get("date")
|
|
assert date >= "2024-02-01", f"Date {date} should be >= '2024-02-01'"
|
|
assert date <= "2024-03-15", f"Date {date} should be <= '2024-03-15'"
|
|
logger.debug(f" Node {r.vector_id}: date={date}")
|
|
|
|
assert len(results) >= 3, f"Expected at least 3 results, got {len(results)}"
|
|
logger.info("✓ String range query validated")
|
|
except Exception as e:
|
|
# For PGVector, this might fail on older implementations
|
|
if "PGVector" in store_name:
|
|
logger.warning(f"String range query failed for PGVector (expected if not updated): {e}")
|
|
else:
|
|
raise
|
|
|
|
# Clean up test data
|
|
await store.delete([node.vector_id for node in test_nodes])
|
|
logger.info("Cleaned up test nodes")
|
|
|
|
logger.info("✓ String range queries test passed")
|
|
|
|
|
|
async def test_sql_injection_protection(store: BaseVectorStore, store_name: str):
|
|
"""Test SQL injection protection in filter keys and collection names."""
|
|
logger.info("=" * 20 + " SQL INJECTION PROTECTION TEST " + "=" * 20)
|
|
|
|
# This test is only relevant for SQL-based stores
|
|
if store_name not in ["PGVectorStore"]:
|
|
logger.info(f"Skipping SQL injection test for {store_name}")
|
|
return
|
|
|
|
# Test 1: Invalid collection name (SQL injection attempt)
|
|
try:
|
|
embedding_model = OpenAIEmbeddingModel()
|
|
|
|
# This should raise ValueError due to invalid table name
|
|
try:
|
|
_ = PGVectorStore(
|
|
collection_name="test'; DROP TABLE users; --",
|
|
db_path=".",
|
|
embedding_model=embedding_model,
|
|
)
|
|
logger.error("❌ FAILED: Invalid collection name was accepted (SQL injection risk!)")
|
|
assert False, "Should have raised ValueError for invalid collection name"
|
|
except ValueError as e:
|
|
logger.info(f"✓ Invalid collection name rejected: {e}")
|
|
|
|
# Test 2: Invalid metadata key in filters
|
|
try:
|
|
_ = await store.search(
|
|
query="test",
|
|
filters={
|
|
"normal_key": "value",
|
|
"bad'; DROP TABLE users; --": "value",
|
|
},
|
|
)
|
|
logger.error("❌ FAILED: Invalid metadata key was accepted (SQL injection risk!)")
|
|
assert False, "Should have raised ValueError for invalid metadata key"
|
|
except ValueError as e:
|
|
logger.info(f"✓ Invalid metadata key rejected: {e}")
|
|
|
|
logger.info("✓ SQL injection protection validated")
|
|
|
|
except Exception as e:
|
|
logger.error(f"SQL injection protection test failed: {e}")
|
|
raise
|
|
|
|
logger.info("✓ SQL injection protection test passed")
|
|
|
|
|
|
async def test_list_with_sorting(store: BaseVectorStore, _store_name: str):
|
|
"""Test list operation with sorting by timestamp to get most recent top 10 items."""
|
|
logger.info("=" * 20 + " LIST WITH SORTING TEST " + "=" * 20)
|
|
|
|
# Clean up any existing test data first
|
|
try:
|
|
existing_nodes = await store.list(filters={"test_type": "timestamp_test"})
|
|
if existing_nodes:
|
|
await store.delete([node.vector_id for node in existing_nodes])
|
|
logger.info(f"Cleaned up {len(existing_nodes)} existing test nodes")
|
|
except Exception as e:
|
|
logger.warning(f"Failed to clean up existing nodes: {e}")
|
|
|
|
# Create test nodes with timestamps
|
|
import time
|
|
|
|
test_nodes = []
|
|
base_timestamp = int(time.time())
|
|
|
|
for i in range(15):
|
|
node = VectorNode(
|
|
vector_id=f"timestamp_node_{i}",
|
|
content=f"Test content for node {i} with timestamp",
|
|
metadata={
|
|
"test_type": "timestamp_test",
|
|
"timestamp": base_timestamp - (14 - i) * 3600, # Each node is 1 hour newer
|
|
"index": i,
|
|
"created_at": base_timestamp - (14 - i) * 3600,
|
|
},
|
|
)
|
|
test_nodes.append(node)
|
|
|
|
# Insert test nodes
|
|
await store.insert(test_nodes)
|
|
logger.info(f"Inserted {len(test_nodes)} test nodes with timestamps")
|
|
|
|
# Test 1: Get all items sorted by timestamp (ascending)
|
|
results_asc = await store.list(
|
|
filters={"test_type": "timestamp_test"},
|
|
sort_key="timestamp",
|
|
reverse=False,
|
|
)
|
|
logger.info(f"Test 1 - Ascending order: {len(results_asc)} results")
|
|
|
|
# Verify ascending order
|
|
for i in range(len(results_asc) - 1):
|
|
ts1 = results_asc[i].metadata.get("timestamp", 0)
|
|
ts2 = results_asc[i + 1].metadata.get("timestamp", 0)
|
|
assert ts1 <= ts2, f"Results not in ascending order: {ts1} > {ts2}"
|
|
|
|
logger.info(
|
|
f"First item timestamp: {results_asc[0].metadata.get('timestamp')}, "
|
|
f"index: {results_asc[0].metadata.get('index')}",
|
|
)
|
|
logger.info(
|
|
f"Last item timestamp: {results_asc[-1].metadata.get('timestamp')}, "
|
|
f"index: {results_asc[-1].metadata.get('index')}",
|
|
)
|
|
|
|
# Test 2: Get top 10 most recent items (descending order)
|
|
results_desc = await store.list(
|
|
filters={"test_type": "timestamp_test"},
|
|
limit=10,
|
|
sort_key="timestamp",
|
|
reverse=True,
|
|
)
|
|
logger.info(f"Test 2 - Top 10 most recent (descending order): {len(results_desc)} results")
|
|
|
|
# Verify we got exactly 10 results
|
|
assert len(results_desc) == 10, f"Expected 10 results, got {len(results_desc)}"
|
|
|
|
# Log the actual results for debugging
|
|
logger.info("Top 10 results (should be index 14 to 5):")
|
|
for i, node in enumerate(results_desc):
|
|
logger.info(f" {i}: index={node.metadata.get('index')}, timestamp={node.metadata.get('timestamp')}")
|
|
|
|
# Verify descending order
|
|
for i in range(len(results_desc) - 1):
|
|
ts1 = results_desc[i].metadata.get("timestamp", 0)
|
|
ts2 = results_desc[i + 1].metadata.get("timestamp", 0)
|
|
assert ts1 >= ts2, f"Results not in descending order: {ts1} < {ts2}"
|
|
|
|
# Verify we got the most recent items (index 5-14)
|
|
for node in results_desc:
|
|
index = node.metadata.get("index", -1)
|
|
assert index >= 5, f"Top 10 should have index >= 5, got index {index}"
|
|
|
|
logger.info(
|
|
f"Most recent item - timestamp: {results_desc[0].metadata.get('timestamp')}, "
|
|
f"index: {results_desc[0].metadata.get('index')}",
|
|
)
|
|
logger.info(
|
|
f"10th most recent item - timestamp: {results_desc[-1].metadata.get('timestamp')}, "
|
|
f"index: {results_desc[-1].metadata.get('index')}",
|
|
)
|
|
|
|
# Test 3: Get top 5 most recent with additional filter
|
|
results_limited = await store.list(
|
|
filters={"test_type": "timestamp_test"},
|
|
limit=5,
|
|
sort_key="timestamp",
|
|
reverse=True,
|
|
)
|
|
logger.info(f"Test 3 - Top 5 most recent: {len(results_limited)} results")
|
|
assert len(results_limited) == 5, f"Expected 5 results, got {len(results_limited)}"
|
|
|
|
# Verify these are the 5 most recent
|
|
for i, node in enumerate(results_limited):
|
|
expected_index = 14 - i # Should be 14, 13, 12, 11, 10
|
|
actual_index = node.metadata.get("index", -1)
|
|
assert actual_index == expected_index, f"Expected index {expected_index}, got {actual_index}"
|
|
|
|
# Test 4: Sort by different key (created_at)
|
|
results_created = await store.list(
|
|
filters={"test_type": "timestamp_test"},
|
|
limit=3,
|
|
sort_key="created_at",
|
|
reverse=True,
|
|
)
|
|
logger.info(f"Test 4 - Top 3 by created_at: {len(results_created)} results")
|
|
assert len(results_created) == 3, f"Expected 3 results, got {len(results_created)}"
|
|
|
|
# Verify sorting by created_at
|
|
for i in range(len(results_created) - 1):
|
|
ts1 = results_created[i].metadata.get("created_at", 0)
|
|
ts2 = results_created[i + 1].metadata.get("created_at", 0)
|
|
assert ts1 >= ts2, f"Results not sorted by created_at: {ts1} < {ts2}"
|
|
|
|
# Clean up test data
|
|
await store.delete([node.vector_id for node in test_nodes])
|
|
logger.info("Cleaned up test nodes")
|
|
|
|
logger.info("✓ List with sorting 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:
|
|
# Initialize the store (connect to database, create client, etc.)
|
|
await store.start()
|
|
|
|
# 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_exact_match_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)
|
|
await test_range_query_filters(store, store_name)
|
|
await test_string_range_queries(store, store_name)
|
|
await test_sql_injection_protection(store, store_name)
|
|
await test_list_with_sorting(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: Backend key (e.g. ``"local"``, ``"obvec"``)
|
|
"""
|
|
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}")
|
|
|
|
# ObVecVectorStore uses a temp db_path per run (reserved for local sidecar files).
|
|
if store_type == "obvec":
|
|
obvec_dir = getattr(store, "db_path", None)
|
|
if obvec_dir and Path(obvec_dir).exists():
|
|
shutil.rmtree(obvec_dir, ignore_errors=True)
|
|
logger.info(f"Cleaned up obvec temp directory: {obvec_dir}")
|
|
|
|
# Clean up local directory if ZvecVectorStore
|
|
if store_type == "zvec" and config.ZVEC_PATH:
|
|
test_dir = Path(config.ZVEC_PATH)
|
|
if test_dir.exists():
|
|
shutil.rmtree(test_dir)
|
|
logger.info(f"Cleaned up zvec directory: {config.ZVEC_PATH}")
|
|
|
|
# Clean up temp directory if SeekdbVectorStore used temp dir
|
|
if store_type == "seekdb" and config.SEEKDB_PATH is None and getattr(store, "db_path", None):
|
|
test_dir = Path(store.db_path)
|
|
if test_dir.exists() and "test_seekdb_" in str(test_dir):
|
|
shutil.rmtree(test_dir, ignore_errors=True)
|
|
logger.info(f"Cleaned up seekdb temp directory: {test_dir}")
|
|
|
|
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 --obvec # Test ObVecVectorStore (seekdb / OceanBase)
|
|
python test_vector_store.py --hologres # Test HologresVectorStore
|
|
python test_vector_store.py --zvec # Test ZvecVectorStore
|
|
python test_vector_store.py --seekdb # Test SeekdbVectorStore 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(
|
|
"--obvec",
|
|
action="store_true",
|
|
help="Test ObVecVectorStore",
|
|
)
|
|
parser.add_argument(
|
|
"--hologres",
|
|
action="store_true",
|
|
help="Test HologresVectorStore",
|
|
)
|
|
parser.add_argument(
|
|
"--zvec",
|
|
action="store_true",
|
|
help="Test ZvecVectorStore",
|
|
)
|
|
parser.add_argument(
|
|
"--seekdb",
|
|
action="store_true",
|
|
help="Test SeekdbVectorStore (requires pyseekdb on Python >=3.11)",
|
|
)
|
|
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"),
|
|
("obvec", "ObVecVectorStore"),
|
|
("hologres", "HologresVectorStore"),
|
|
("zvec", "ZvecVectorStore"),
|
|
]
|
|
if SEEKDB_AVAILABLE:
|
|
stores_to_test.append(("seekdb", "SeekdbVectorStore"))
|
|
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 args.obvec:
|
|
stores_to_test.append(("obvec", "ObVecVectorStore"))
|
|
if args.hologres:
|
|
stores_to_test.append(("hologres", "HologresVectorStore"))
|
|
if args.zvec:
|
|
stores_to_test.append(("zvec", "ZvecVectorStore"))
|
|
if args.seekdb:
|
|
if not SEEKDB_AVAILABLE:
|
|
raise ImportError(
|
|
"seekdb tests require pyseekdb (Python >=3.11); install: pip install 'pyseekdb>=1.2.0'",
|
|
)
|
|
stores_to_test.append(("seekdb", "SeekdbVectorStore"))
|
|
|
|
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"),
|
|
("obvec", "ObVecVectorStore"),
|
|
("hologres", "HologresVectorStore"),
|
|
("zvec", "ZvecVectorStore"),
|
|
]
|
|
if SEEKDB_AVAILABLE:
|
|
stores_to_test.append(("seekdb", "SeekdbVectorStore"))
|
|
print("No vector store specified, defaulting to test all vector stores")
|
|
print(
|
|
"Use --local/--es/--pgvector/--qdrant/--chroma/--obvec/--zvec/--hologres/--seekdb "
|
|
"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())
|