File size: 12,267 Bytes
979853c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 | import pytest
import numpy as np
from unittest.mock import patch, AsyncMock
from lightrag.utils import EmbeddingFunc
from lightrag.kg.postgres_impl import (
PGVectorStorage,
PostgreSQLDB,
_safe_index_name,
)
from lightrag.exceptions import DataMigrationError
from lightrag.namespace import NameSpace
# Mock PostgreSQLDB
@pytest.fixture
def mock_pg_db():
"""Mock PostgreSQL database connection"""
db = AsyncMock()
db.workspace = "test_workspace"
db.vector_index_type = None
# Mock query responses: list for search queries (multirows=True), dict for DDL checks
async def mock_query(sql, params=None, multirows=False, **kwargs):
if multirows:
return []
return {"exists": False, "count": 0}
# Mock for execute that mimics PostgreSQLDB.execute() behavior
async def mock_execute(sql, data=None, **kwargs):
return None
db.query = AsyncMock(side_effect=mock_query)
db.execute = AsyncMock(side_effect=mock_execute)
return db
# Mock get_data_init_lock to avoid async lock issues in tests
@pytest.fixture(autouse=True)
def mock_data_init_lock():
with patch("lightrag.kg.postgres_impl.get_data_init_lock") as mock_lock:
mock_lock_ctx = AsyncMock()
mock_lock.return_value = mock_lock_ctx
yield mock_lock
# Mock ClientManager
@pytest.fixture
def mock_client_manager(mock_pg_db):
with patch("lightrag.kg.postgres_impl.ClientManager") as mock_manager:
mock_manager.get_client = AsyncMock(return_value=mock_pg_db)
mock_manager.release_client = AsyncMock()
yield mock_manager
# Mock Embedding function
@pytest.fixture
def mock_embedding_func():
async def embed_func(texts, **kwargs):
return np.array([[0.1] * 768 for _ in texts])
# Note: EmbeddingFunc in this version of lightrag supports model_name
func = EmbeddingFunc(embedding_dim=768, func=embed_func, model_name="test_model")
return func
@pytest.mark.asyncio
async def test_postgres_halfvec_table_creation(
mock_client_manager, mock_pg_db, mock_embedding_func
):
"""Test if table is created with HALFVEC type when HNSW_HALFVEC is selected"""
# Set index type to HNSW_HALFVEC
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
config = {
"embedding_batch_num": 10,
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
}
storage = PGVectorStorage(
namespace=NameSpace.VECTOR_STORE_CHUNKS,
global_config=config,
embedding_func=mock_embedding_func,
workspace="test_ws",
)
# Mock table doesn't exist
mock_pg_db.check_table_exists = AsyncMock(return_value=False)
# Initialize storage (should trigger table creation)
await storage.initialize()
# Verify table creation SQL contains HALFVEC(768)
create_table_calls = [
call
for call in mock_pg_db.execute.call_args_list
if "CREATE TABLE" in call[0][0]
]
assert len(create_table_calls) > 0
create_sql = create_table_calls[0][0][0]
assert "HALFVEC(768)" in create_sql
assert "VECTOR(768)" not in create_sql
@pytest.mark.asyncio
async def test_postgres_vector_table_creation_default(
mock_client_manager, mock_pg_db, mock_embedding_func
):
"""Test if table is created with default VECTOR type when other index type is selected"""
# Set index type to HNSW (default)
mock_pg_db.vector_index_type = "HNSW"
config = {
"embedding_batch_num": 10,
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
}
storage = PGVectorStorage(
namespace=NameSpace.VECTOR_STORE_CHUNKS,
global_config=config,
embedding_func=mock_embedding_func,
workspace="test_ws",
)
# Mock table doesn't exist
mock_pg_db.check_table_exists = AsyncMock(return_value=False)
# Initialize storage (should trigger table creation)
await storage.initialize()
# Verify table creation SQL contains VECTOR(768)
create_table_calls = [
call
for call in mock_pg_db.execute.call_args_list
if "CREATE TABLE" in call[0][0]
]
assert len(create_table_calls) > 0
create_sql = create_table_calls[0][0][0]
assert "VECTOR(768)" in create_sql
assert "HALFVEC(768)" not in create_sql
# Namespaces that use vector search SQL templates (query path)
QUERY_NAMESPACES = [
NameSpace.VECTOR_STORE_CHUNKS,
NameSpace.VECTOR_STORE_ENTITIES,
NameSpace.VECTOR_STORE_RELATIONSHIPS,
]
@pytest.mark.asyncio
@pytest.mark.parametrize("namespace", QUERY_NAMESPACES)
async def test_query_uses_halfvec_cast_when_hnsw_halfvec(
mock_client_manager, mock_pg_db, mock_embedding_func, namespace
):
"""When HNSW_HALFVEC is set, generated search SQL uses ::halfvec (not ::vector)."""
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
mock_pg_db.check_table_exists = AsyncMock(return_value=True)
config = {
"embedding_batch_num": 10,
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
}
storage = PGVectorStorage(
namespace=namespace,
global_config=config,
embedding_func=mock_embedding_func,
workspace="test_ws",
)
await storage.initialize()
query_embedding = [0.1] * 768
await storage.query("test query", top_k=5, query_embedding=query_embedding)
assert mock_pg_db.query.called
call_args = mock_pg_db.query.call_args
sql = call_args[0][0]
assert "::halfvec" in sql
assert "::vector" not in sql
@pytest.mark.asyncio
@pytest.mark.parametrize("namespace", QUERY_NAMESPACES)
async def test_query_uses_vector_cast_when_hnsw_default(
mock_client_manager, mock_pg_db, mock_embedding_func, namespace
):
"""When HNSW (default) is set, generated search SQL uses ::vector (not ::halfvec)."""
mock_pg_db.vector_index_type = "HNSW"
mock_pg_db.check_table_exists = AsyncMock(return_value=True)
config = {
"embedding_batch_num": 10,
"vector_db_storage_cls_kwargs": {"cosine_better_than_threshold": 0.8},
}
storage = PGVectorStorage(
namespace=namespace,
global_config=config,
embedding_func=mock_embedding_func,
workspace="test_ws",
)
await storage.initialize()
query_embedding = [0.1] * 768
await storage.query("test query", top_k=5, query_embedding=query_embedding)
assert mock_pg_db.query.called
call_args = mock_pg_db.query.call_args
sql = call_args[0][0]
assert "::vector" in sql
assert "::halfvec" not in sql
# ---------------------------------------------------------------------------
# Index switching: old conflicting indexes are dropped
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_create_vector_index_drops_old_indexes_when_switching(mock_pg_db):
"""Switching from HNSW to HNSW_HALFVEC drops the old hnsw_cosine index."""
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
mock_pg_db.hnsw_m = 16
mock_pg_db.hnsw_ef = 64
mock_pg_db.ivfflat_lists = 100
mock_pg_db.vchordrq_build_options = ""
table_name = "lightrag_vdb_chunks_test"
async def mock_query(sql, params=None, multirows=False, **kwargs):
if "pg_indexes" in sql:
return None
return None
mock_pg_db.query = AsyncMock(side_effect=mock_query)
mock_pg_db.execute = AsyncMock()
# Call the real method with mock_pg_db as self
await PostgreSQLDB._create_vector_index(mock_pg_db, table_name, 3072)
execute_calls = [call[0][0] for call in mock_pg_db.execute.call_args_list]
old_hnsw_name = _safe_index_name(table_name, "hnsw_cosine")
old_ivfflat_name = _safe_index_name(table_name, "ivfflat_cosine")
old_vchordrq_name = _safe_index_name(table_name, "vchordrq_cosine")
drop_calls = [c for c in execute_calls if "DROP INDEX IF EXISTS" in c]
dropped_names = {c.split("DROP INDEX IF EXISTS ")[1].strip() for c in drop_calls}
assert old_hnsw_name in dropped_names
assert old_ivfflat_name in dropped_names
assert old_vchordrq_name in dropped_names
new_index_name = _safe_index_name(table_name, "hnsw_halfvec_cosine")
assert new_index_name not in dropped_names
alter_calls = [c for c in execute_calls if "ALTER TABLE" in c]
assert any("HALFVEC(3072)" in c for c in alter_calls)
create_calls = [c for c in execute_calls if "CREATE INDEX" in c]
assert any("halfvec_cosine_ops" in c for c in create_calls)
@pytest.mark.asyncio
async def test_create_vector_index_no_drop_when_index_exists(mock_pg_db):
"""If the target index already exists, no DROP or CREATE is issued."""
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
mock_pg_db.hnsw_m = 16
mock_pg_db.hnsw_ef = 64
mock_pg_db.ivfflat_lists = 100
mock_pg_db.vchordrq_build_options = ""
table_name = "lightrag_vdb_chunks_test"
async def mock_query(sql, params=None, multirows=False, **kwargs):
if "pg_indexes" in sql:
return {"?column?": 1}
return None
mock_pg_db.query = AsyncMock(side_effect=mock_query)
mock_pg_db.execute = AsyncMock()
await PostgreSQLDB._create_vector_index(mock_pg_db, table_name, 3072)
execute_calls = [call[0][0] for call in mock_pg_db.execute.call_args_list]
assert not any("DROP INDEX" in c for c in execute_calls)
assert not any("CREATE INDEX" in c for c in execute_calls)
# ---------------------------------------------------------------------------
# HalfVector dimension detection in setup_table
# ---------------------------------------------------------------------------
class _MockHalfVector:
"""Mimics pgvector.halfvec.HalfVector for testing dimension detection."""
def __init__(self, dim: int):
self._dim = dim
def dimensions(self) -> int:
return self._dim
def to_list(self):
return [0.0] * self._dim
@pytest.mark.asyncio
async def test_setup_table_detects_halfvector_dimension_mismatch(mock_pg_db):
"""DataMigrationError is raised when a HalfVector column has a different dimension."""
table_name = "lightrag_vdb_chunks_new"
legacy_table = "lightrag_vdb_chunks"
mock_pg_db.check_table_exists = AsyncMock(
side_effect=lambda t: t.lower() == legacy_table.lower()
)
call_count = 0
async def mock_query(sql, params=None, multirows=False, **kwargs):
nonlocal call_count
call_count += 1
if "COUNT(*)" in sql:
return {"count": 5}
if "content_vector" in sql:
return {"content_vector": _MockHalfVector(1024)}
return None
mock_pg_db.query = AsyncMock(side_effect=mock_query)
mock_pg_db.execute = AsyncMock()
with pytest.raises(DataMigrationError, match="Dimension mismatch"):
await PGVectorStorage.setup_table(
db=mock_pg_db,
table_name=table_name,
workspace="test_ws",
embedding_dim=768,
legacy_table_name=legacy_table,
base_table=legacy_table,
)
@pytest.mark.asyncio
async def test_setup_table_accepts_matching_halfvector_dimension(mock_pg_db):
"""No error when HalfVector dimension matches the expected embedding_dim."""
table_name = "lightrag_vdb_chunks_new"
legacy_table = "lightrag_vdb_chunks"
mock_pg_db.check_table_exists = AsyncMock(
side_effect=lambda t: t.lower() == legacy_table.lower()
)
mock_pg_db.vector_index_type = "HNSW_HALFVEC"
async def mock_query(sql, params=None, multirows=False, **kwargs):
if "COUNT(*)" in sql:
return {"count": 5}
if "content_vector" in sql:
return {"content_vector": _MockHalfVector(768)}
if multirows:
return []
return None
mock_pg_db.query = AsyncMock(side_effect=mock_query)
mock_pg_db.execute = AsyncMock()
with patch.object(PGVectorStorage, "_pg_create_table", new_callable=AsyncMock):
await PGVectorStorage.setup_table(
db=mock_pg_db,
table_name=table_name,
workspace="test_ws",
embedding_dim=768,
legacy_table_name=legacy_table,
base_table=legacy_table,
)
|