zqf
b1f93588f1
- 排除 .env / *.bak / 内部运维文档（README_INTERNAL.html） - config.py 默认密码已替换为占位符 CHANGE_ME_* - init_db.sql 移除生产数据库用户 GRANT 段 - README.html 数据库用户名已脱敏 - 保留：源码 + 公网 API 文档 + 建表 SQL（无授权语句）
114 lines
3.5 KiB
Python
114 lines
3.5 KiB
Python
"""Vector Search Layer - Pure NumPy Implementation"""
|
|
import struct
|
|
import logging
|
|
from typing import List, Tuple, Optional
|
|
|
|
import numpy as np
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def embedding_to_bytes(embedding: List[float]) -> bytes:
|
|
"""Convert float list to compact bytes for MySQL BLOB storage."""
|
|
return struct.pack(f"{len(embedding)}f", *embedding)
|
|
|
|
|
|
def bytes_to_embedding(data: bytes) -> np.ndarray:
|
|
"""Convert bytes back to numpy array."""
|
|
n = len(data) // 4
|
|
return np.array(struct.unpack(f"{n}f", data), dtype=np.float32)
|
|
|
|
|
|
def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
|
|
"""Compute cosine similarity between two vectors."""
|
|
norm_a = np.linalg.norm(a)
|
|
norm_b = np.linalg.norm(b)
|
|
if norm_a == 0 or norm_b == 0:
|
|
return 0.0
|
|
return float(np.dot(a, b) / (norm_a * norm_b))
|
|
|
|
|
|
def search_similar(
|
|
query_embedding: List[float],
|
|
stored_items: List[dict],
|
|
top_k: int = 10,
|
|
min_score: float = 0.0,
|
|
) -> List[Tuple[str, float, dict]]:
|
|
"""
|
|
Search for similar items using cosine similarity.
|
|
|
|
Args:
|
|
query_embedding: Query vector
|
|
stored_items: List of dicts with 'id', 'embedding' (bytes), 'content'
|
|
top_k: Max results
|
|
min_score: Minimum similarity score
|
|
|
|
Returns:
|
|
List of (id, score, item_dict) tuples sorted by score descending
|
|
"""
|
|
if not stored_items:
|
|
return []
|
|
|
|
query_vec = np.array(query_embedding, dtype=np.float32)
|
|
query_norm = np.linalg.norm(query_vec)
|
|
if query_norm == 0:
|
|
return []
|
|
|
|
scores = []
|
|
for item in stored_items:
|
|
emb_bytes = item.get("embedding")
|
|
if emb_bytes is None:
|
|
continue
|
|
try:
|
|
stored_vec = bytes_to_embedding(emb_bytes)
|
|
score = float(np.dot(query_vec, stored_vec) / (query_norm * np.linalg.norm(stored_vec)))
|
|
if score >= min_score:
|
|
scores.append((item["id"], score, item))
|
|
except Exception as e:
|
|
logger.warning(f"Failed to compute similarity for item {item.get('id')}: {e}")
|
|
continue
|
|
|
|
scores.sort(key=lambda x: x[1], reverse=True)
|
|
return scores[:top_k]
|
|
|
|
def find_duplicates(
|
|
new_embedding: List[float],
|
|
stored_items: List[dict],
|
|
threshold: float = 0.85,
|
|
) -> List[Tuple[str, float, dict]]:
|
|
"""
|
|
Find items whose similarity to new_embedding exceeds threshold.
|
|
Used for deduplication before storing new memories.
|
|
|
|
Args:
|
|
new_embedding: The embedding of the new memory to check
|
|
stored_items: Existing memories with 'id', 'embedding' (bytes), 'content'
|
|
threshold: Similarity threshold above which items are considered duplicates
|
|
|
|
Returns:
|
|
List of (id, score, item_dict) for items exceeding threshold, sorted by score desc
|
|
"""
|
|
if not stored_items:
|
|
return []
|
|
|
|
query_vec = np.array(new_embedding, dtype=np.float32)
|
|
query_norm = np.linalg.norm(query_vec)
|
|
if query_norm == 0:
|
|
return []
|
|
|
|
duplicates = []
|
|
for item in stored_items:
|
|
emb_bytes = item.get("embedding")
|
|
if emb_bytes is None:
|
|
continue
|
|
try:
|
|
stored_vec = bytes_to_embedding(emb_bytes)
|
|
score = float(np.dot(query_vec, stored_vec) / (query_norm * np.linalg.norm(stored_vec)))
|
|
if score >= threshold:
|
|
duplicates.append((item["id"], score, item))
|
|
except Exception as e:
|
|
continue
|
|
|
|
duplicates.sort(key=lambda x: x[1], reverse=True)
|
|
return duplicates
|