# -*- 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)")