初始化:记忆系统源代码上传(已脱敏)
- 排除 .env / *.bak / 内部运维文档(README_INTERNAL.html) - config.py 默认密码已替换为占位符 CHANGE_ME_* - init_db.sql 移除生产数据库用户 GRANT 段 - README.html 数据库用户名已脱敏 - 保留:源码 + 公网 API 文档 + 建表 SQL(无授权语句)
This commit is contained in:
@@ -0,0 +1,113 @@
|
||||
"""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
|
||||
Reference in New Issue
Block a user