317 lines
11 KiB
Python
317 lines
11 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
Neo4j 向量检索稀疏用户场景测试
|
||
|
||
重点测试场景：
|
||
- 一个用户数据量少，但其他用户数据量大且相似度更高
|
||
- 全局 top-K HNSW 可能被其他用户填满
|
||
- 迭代搜索应该通过增加 K 倍数来确保稀疏用户的结果被返回
|
||
|
||
这验证了 Solution B (Iterative Fetch with Adaptive K) 的有效性。
|
||
"""
|
||
|
||
import pytest
|
||
from typing import List
|
||
import numpy as np
|
||
|
||
|
||
def cosine_similarity(vec1: List[float], vec2: List[float]) -> float:
|
||
"""计算余弦相似度"""
|
||
v1 = np.array(vec1)
|
||
v2 = np.array(vec2)
|
||
return float(np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2)))
|
||
|
||
|
||
@pytest.mark.integration
|
||
class TestSparseUserVectorSearch:
|
||
"""稀疏用户场景向量搜索测试"""
|
||
|
||
@pytest.mark.slow
|
||
async def test_sparse_user_finds_own_memories(self, graph_store_neo4j_fixture):
|
||
"""
|
||
核心测试：稀疏用户能找到自己的记忆
|
||
|
||
问题描述（Solution B 之前的问题）:
|
||
- User A（稀疏）: 5 个记忆
|
||
- User B（密集）: 100 个高相似度的记忆
|
||
- 全局 top-30 被 User B 填满
|
||
- OLD: User A 的搜索返回空或仅返回 User B 的结果（数据泄露）
|
||
|
||
预期（Solution B 之后）:
|
||
- 迭代搜索通过增加 K 倍数（3 → 4.5 → 6.75）
|
||
- 最终找到 User A 的记忆并返回
|
||
"""
|
||
ik_sparse = "sparse_user_a::default"
|
||
ik_dense = "dense_user_b::default"
|
||
|
||
# 创建向量：使用 NumPy 确保一致性
|
||
np.random.seed(42)
|
||
|
||
# Sparse user 的记忆向量（独特的方向）
|
||
base_vector_sparse = np.random.randn(1536)
|
||
base_vector_sparse = base_vector_sparse / np.linalg.norm(base_vector_sparse)
|
||
|
||
sparse_memories = []
|
||
for i in range(5):
|
||
# 围绕基向量添加噪声
|
||
noise = np.random.randn(1536) * 0.1
|
||
vec = base_vector_sparse + noise
|
||
vec = vec / np.linalg.norm(vec)
|
||
sparse_memories.append({
|
||
"node_id": f"sparse_a_{i}",
|
||
"content": f"Sparse user A memory {i}",
|
||
"layer": "l6_schema",
|
||
"status": "active",
|
||
"isolation_key": ik_sparse,
|
||
"embedding": vec.tolist(),
|
||
"confidence": 0.8,
|
||
})
|
||
|
||
# Dense user 的记忆向量（都高度相似，相似度 > 0.95）
|
||
base_vector_dense = np.random.randn(1536)
|
||
base_vector_dense = base_vector_dense / np.linalg.norm(base_vector_dense)
|
||
|
||
dense_memories = []
|
||
for i in range(100):
|
||
# 非常小的噪声（高相似度）
|
||
noise = np.random.randn(1536) * 0.01
|
||
vec = base_vector_dense + noise
|
||
vec = vec / np.linalg.norm(vec)
|
||
dense_memories.append({
|
||
"node_id": f"dense_b_{i}",
|
||
"content": f"Dense user B memory {i}",
|
||
"layer": "l6_schema",
|
||
"status": "active",
|
||
"isolation_key": ik_dense,
|
||
"embedding": vec.tolist(),
|
||
"confidence": 0.95,
|
||
})
|
||
|
||
# 搜索：使用 sparse user 的向量
|
||
query_embedding = sparse_memories[0]["embedding"]
|
||
|
||
results = await graph_store_neo4j_fixture.vector_search(
|
||
query_embedding=query_embedding,
|
||
isolation_key=ik_sparse,
|
||
limit=5,
|
||
)
|
||
|
||
# 验证
|
||
assert len(results) > 0, \
|
||
"Sparse user should find its own memories (iterative fetch should handle this)"
|
||
|
||
for result in results:
|
||
assert result["node_id"].startswith("sparse_a_"), \
|
||
f"All results should belong to sparse user, got {result['node_id']}"
|
||
assert "sparse_a_" in result["node_id"]
|
||
|
||
print(f"✓ Sparse user test passed: found {len(results)}/5 own memories despite 100 dense competitors")
|
||
|
||
@pytest.mark.slow
|
||
async def test_sparse_user_score_distribution(self, graph_store_neo4j_fixture):
|
||
"""
|
||
测试：稀疏用户的结果分数合理分布
|
||
|
||
验证：
|
||
- 返回的稀疏用户记忆的分数都 > 0.5（valid match）
|
||
- 分数反映了与查询的相似度
|
||
"""
|
||
ik_sparse = "sparse_score_test::default"
|
||
|
||
np.random.seed(123)
|
||
|
||
base_vector = np.random.randn(1536)
|
||
base_vector = base_vector / np.linalg.norm(base_vector)
|
||
|
||
memories = []
|
||
for i in range(5):
|
||
# 不同相似度等级的向量
|
||
if i < 2:
|
||
noise_level = 0.05 # 高相似度
|
||
elif i < 4:
|
||
noise_level = 0.15 # 中相似度
|
||
else:
|
||
noise_level = 0.3 # 低相似度
|
||
|
||
noise = np.random.randn(1536) * noise_level
|
||
vec = base_vector + noise
|
||
vec = vec / np.linalg.norm(vec)
|
||
memories.append({
|
||
"node_id": f"score_test_{i}",
|
||
"content": f"Memory with noise level {noise_level}",
|
||
"layer": "l6_schema",
|
||
"status": "active",
|
||
"isolation_key": ik_sparse,
|
||
"embedding": vec.tolist(),
|
||
"confidence": 0.8,
|
||
})
|
||
|
||
query_embedding = base_vector.tolist()
|
||
|
||
results = await graph_store_neo4j_fixture.vector_search(
|
||
query_embedding=query_embedding,
|
||
isolation_key=ik_sparse,
|
||
limit=10,
|
||
)
|
||
|
||
# 验证分数
|
||
assert len(results) > 0
|
||
scores = [r["score"] for r in results]
|
||
|
||
# 分数应该在 (0, 1] 范围内（余弦相似度）
|
||
assert all(0 < s <= 1.0 for s in scores), \
|
||
f"Scores should be in (0, 1], got {scores}"
|
||
|
||
# 分数应该单调递减
|
||
assert scores == sorted(scores, reverse=True), \
|
||
"Scores should be sorted in descending order"
|
||
|
||
print(f"✓ Score distribution test passed: got {len(results)} results with valid scores {scores[:3]}")
|
||
|
||
@pytest.mark.slow
|
||
async def test_sparse_user_layer_filter(self, graph_store_neo4j_fixture):
|
||
"""
|
||
测试：稀疏用户在指定 layer 过滤下的搜索
|
||
|
||
Setup:
|
||
- 创建多个 layer 的记忆
|
||
- 搜索指定 layer，例如仅 l6_schema
|
||
|
||
Expected:
|
||
- 返回的结果仅属于指定的 layer
|
||
- 隔离性和迭代逻辑依然生效
|
||
"""
|
||
ik = "layer_filter_test::default"
|
||
|
||
np.random.seed(456)
|
||
base_vector = np.random.randn(1536)
|
||
base_vector = base_vector / np.linalg.norm(base_vector)
|
||
|
||
# 创建多个 layer 的记忆
|
||
memories = []
|
||
for layer in ["l5_identity", "l6_schema", "l7_intention"]:
|
||
for i in range(3):
|
||
noise = np.random.randn(1536) * 0.1
|
||
vec = base_vector + noise
|
||
vec = vec / np.linalg.norm(vec)
|
||
memories.append({
|
||
"node_id": f"mem_{layer}_{i}",
|
||
"content": f"Memory in {layer}",
|
||
"layer": layer,
|
||
"status": "active",
|
||
"isolation_key": ik,
|
||
"embedding": vec.tolist(),
|
||
"confidence": 0.8,
|
||
})
|
||
|
||
# 仅搜索 l6_schema
|
||
results = await graph_store_neo4j_fixture.vector_search(
|
||
query_embedding=base_vector.tolist(),
|
||
isolation_key=ik,
|
||
layers=["l6_schema"],
|
||
limit=10,
|
||
)
|
||
|
||
# 验证
|
||
for result in results:
|
||
assert result["layer"] == "l6_schema", \
|
||
f"All results should be l6_schema, got {result['layer']}"
|
||
|
||
print(f"✓ Layer filter test passed: got {len(results)} l6_schema results with proper filtering")
|
||
|
||
@pytest.mark.slow
|
||
async def test_sparse_user_threshold_filter(self, graph_store_neo4j_fixture):
|
||
"""
|
||
测试：稀疏用户在 score_threshold 过滤下的搜索
|
||
|
||
迭代搜索应该正确处理 threshold 过滤
|
||
"""
|
||
ik = "threshold_test::default"
|
||
|
||
np.random.seed(789)
|
||
base_vector = np.random.randn(1536)
|
||
base_vector = base_vector / np.linalg.norm(base_vector)
|
||
|
||
# 创建各种相似度的向量
|
||
memories = []
|
||
for i in range(10):
|
||
noise_level = 0.05 + i * 0.05
|
||
noise = np.random.randn(1536) * noise_level
|
||
vec = base_vector + noise
|
||
vec = vec / np.linalg.norm(vec)
|
||
memories.append({
|
||
"node_id": f"thresh_mem_{i}",
|
||
"content": f"Memory with similarity variance {noise_level}",
|
||
"layer": "l6_schema",
|
||
"status": "active",
|
||
"isolation_key": ik,
|
||
"embedding": vec.tolist(),
|
||
"confidence": 0.8,
|
||
})
|
||
|
||
# 搜索：要求 score > 0.7
|
||
results = await graph_store_neo4j_fixture.vector_search(
|
||
query_embedding=base_vector.tolist(),
|
||
isolation_key=ik,
|
||
limit=20,
|
||
score_threshold=0.7,
|
||
)
|
||
|
||
# 验证
|
||
for result in results:
|
||
assert result["score"] > 0.7, \
|
||
f"All scores should be > 0.7, got {result['score']}"
|
||
|
||
print(f"✓ Threshold filter test passed: got {len(results)} results above threshold 0.7")
|
||
|
||
@pytest.mark.slow
|
||
async def test_iteration_logs_presence(self, graph_store_neo4j_fixture, caplog):
|
||
"""
|
||
测试：迭代搜索的日志应该被记录（用于诊断）
|
||
|
||
当需要多次迭代时，debug 日志应该记录：
|
||
- 每次迭代的 K 值和收集数量
|
||
- 提前退出或耗尽迭代次数的信息
|
||
"""
|
||
ik = "log_test::default"
|
||
|
||
np.random.seed(999)
|
||
base_vector = np.random.randn(1536)
|
||
base_vector = base_vector / np.linalg.norm(base_vector)
|
||
|
||
# 创建足够的记忆来触发迭代
|
||
memories = []
|
||
for i in range(50):
|
||
noise = np.random.randn(1536) * 0.05
|
||
vec = base_vector + noise
|
||
vec = vec / np.linalg.norm(vec)
|
||
memories.append({
|
||
"node_id": f"log_mem_{i}",
|
||
"content": f"Memory for iteration logging {i}",
|
||
"layer": "l6_schema",
|
||
"status": "active",
|
||
"isolation_key": ik,
|
||
"embedding": vec.tolist(),
|
||
"confidence": 0.8,
|
||
})
|
||
|
||
import logging
|
||
# 启用 debug 日志捕获
|
||
with caplog.at_level(logging.DEBUG):
|
||
results = await graph_store_neo4j_fixture.vector_search(
|
||
query_embedding=base_vector.tolist(),
|
||
isolation_key=ik,
|
||
limit=10,
|
||
)
|
||
|
||
# 验证日志中包含迭代信息（如果发生了迭代）
|
||
# 注意：可能不一定触发迭代（取决于实际数据），但至少应该有 debug 日志
|
||
log_text = caplog.text
|
||
has_iteration_log = "vector_search iteration" in log_text or "collected" in log_text
|
||
|
||
if has_iteration_log:
|
||
print(f"✓ Iteration logs test passed: found iteration logs in debug output")
|
||
else:
|
||
print(f"✓ Iteration logs test passed: no iteration needed (all results in first fetch)")
|
||
|