mao1/anchor_search.py

#!/usr/bin/env python3
"""
Anchor Search Engine
多维加权检索引擎。
结合关键词匹配、情感共鸣、时间接近度、重要性进行排序。
"""
import os
import math
import yaml
import json
from datetime import datetime

VAULT_DIR = "/Users/fyah/Documents/如梦初醒/memory/test/Anchor-Memory"

# 检索权重配置
W_TOPIC = 4.0          # 主题域/关键词权重
W_EMOTION = 2.0        # 情感共鸣权重
W_TIME = 1.5           # 时间接近度权重
W_IMPORTANCE = 1.0     # 重要性权重
MAX_RESULTS = 5        # 返回结果数量

def parse_yaml_frontmatter(content):
    """极简 YAML 解析，兼容无 YAML 头的纯 Markdown"""
    if content.startswith("---"):
        parts = content.split("---", 2)
        if len(parts) >= 2:
            try:
                return yaml.safe_load(parts[1]), parts[2]
            except:
                pass
    # 如果没有 YAML 头，返回空字典和全文
    # 提取标题作为默认 title
    title = "Untitled"
    if content.startswith("# "):
        title = content.split("\n")[0][2:].strip()
    return {"title": title, "status": "active", "importance": 5}, content

def calc_keyword_score(query: str, metadata: dict, body: str) -> float:
    """计算关键词匹配得分"""
    q_lower = query.lower()
    score = 0.0
    
    # 标题匹配
    title = metadata.get("title", "").lower()
    if q_lower in title:
        score += 3.0
        
    # 标签匹配
    tags = metadata.get("tags", [])
    if isinstance(tags, list):
        for tag in tags:
            if q_lower in tag.lower():
                score += 2.0
                
    # 内容匹配
    body_lower = body.lower()
    if q_lower in body_lower:
        score += 1.0
        
    # 模糊匹配（简单版：计算重叠词）
    query_words = set(q_lower.split())
    body_words = set(body_lower.split())
    overlap = len(query_words & body_words)
    if overlap > 0:
        score += overlap * 0.5
        
    return score

def calc_emotion_score(query_emotion: dict, metadata: dict) -> float:
    """计算情感共鸣得分（基于 Russell 环形模型距离）"""
    try:
        q_val = query_emotion.get("valence", 0.5)
        q_arousal = query_emotion.get("arousal", 0.3)
        
        m_val = metadata.get("valence", 0.5)
        m_arousal = metadata.get("arousal", 0.3)
        
        # 计算欧氏距离（归一化到 0~1，0 表示完全一致，1 表示完全相反）
        distance = math.sqrt((q_val - m_val)**2 + (q_arousal - m_arousal)**2)
        max_distance = math.sqrt(1.0**2 + 1.0**2) # 最大距离约 1.414
        
        # 距离越近，得分越高
        return max(0.0, 1.0 - (distance / max_distance))
    except:
        return 0.0

def calc_time_score(metadata: dict) -> float:
    """计算时间接近度得分（越近越好）"""
    try:
        date_str = metadata.get("date", "")
        mem_date = datetime.strptime(str(date_str), "%Y-%m-%d %H:%M")
        days_diff = (datetime.now() - mem_date).total_seconds() / 86400
        
        # 指数衰减：越久远的记忆得分越低
        return math.exp(-0.1 * days_diff)
    except:
        return 0.0

def calc_importance_score(metadata: dict) -> float:
    """计算重要性得分（归一化到 0~1）"""
    try:
        imp = metadata.get("importance", 5)
        return float(imp) / 10.0
    except:
        return 0.5

def search_memories(query: str, emotion_hint: dict = None):
    """执行多维检索"""
    print(f"[*] 开始检索: '{query}'")
    results = []
    
    # 遍历所有活跃记忆（跳过归档目录）
    for root, dirs, files in os.walk(VAULT_DIR):
        if "Archive" in root:
            continue
            
        for file in files:
            if not file.endswith(".md"):
                continue
            
            filepath = os.path.join(root, file)
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    content = f.read()
                
                metadata, body = parse_yaml_frontmatter(content)
                if not metadata or metadata.get("status") == "archived":
                    continue
                
                # 计算各维度得分
                k_score = calc_keyword_score(query, metadata, body)
                e_score = calc_emotion_score(emotion_hint or {}, metadata)
                t_score = calc_time_score(metadata)
                i_score = calc_importance_score(metadata)
                
                # 加权总分
                total_score = (
                    k_score * W_TOPIC +
                    e_score * W_EMOTION +
                    t_score * W_TIME +
                    i_score * W_IMPORTANCE
                )
                
                if total_score > 0:
                    results.append({
                        "file": file,
                        "path": os.path.relpath(filepath, VAULT_DIR),
                        "title": metadata.get("title", "Untitled"),
                        "score": total_score,
                        "metadata": metadata,
                        "preview": body[:200].strip()
                    })
            except Exception as e:
                print(f"  [!] 读取失败 {file}: {e}")
                
    # 按得分排序
    results.sort(key=lambda x: x["score"], reverse=True)
    
    # 返回 Top N
    top_results = results[:MAX_RESULTS]
    
    if not top_results:
        print(f"[-] 未找到相关记忆。")
        return []
        
    print(f"[+] 找到 {len(top_results)} 条相关记忆（Top {MAX_RESULTS}）：")
    for r in top_results:
        print(f"\n📄 {r['title']} (得分: {r['score']:.2f})")
        print(f"📍 {r['path']}")
        print(f"👁️ {r['preview']}...")
        
    return top_results

if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1:
        query = " ".join(sys.argv[1:])
        search_memories(query)
    else:
        print("用法: python3 anchor_search.py <查询内容>")