task_start.py

#!/usr/bin/env python3
"""
Memory Engineering V1 — Task Start Hook
用户提出具体任务时调用，召回相关 Rule + Lesson。
用法: python3 task_start.py "<关键词1> <关键词2> ..."
输出: JSON — 匹配到的 rules + lessons（BM25 分数降序）
"""
import json, os, glob, sys, re, math
from datetime import datetime

BASE = os.path.expanduser("~/Desktop/Memory_System/memory")
AGENT = "hermes"

# ─── BM25 Module ──────────────────────────────────────────────────
BM25_K1 = 1.5
BM25_B = 0.75


def tokenize(text: str) -> list[str]:
    """中文按字+词切分，英文按空格切分。返回小写 token 列表。"""
    tokens = []
    for match in re.finditer(r'[一-鿿]+|[a-zA-Z0-9]+', text.lower()):
        token = match.group()
        if re.match(r'[一-鿿]', token):
            # 中文：单字 + 双字组合（bigram）
            tokens.extend(token)  # 单字
            tokens.extend(token[i:i+2] for i in range(len(token)-1))  # 双字
        else:
            tokens.append(token)  # 英文单词保持原样
    return tokens


def compute_idf(documents: list[list[str]]) -> dict[str, float]:
    """计算每个 token 的 IDF 值。N=文档总数，df=包含该 token 的文档数。"""
    N = len(documents)
    df = {}
    for doc in documents:
        for token in set(doc):  # 每个 token 在每个文档中只计 1 次
            df[token] = df.get(token, 0) + 1
    return {token: math.log((N - df_i + 0.5) / (df_i + 0.5) + 1.0)
            for token, df_i in df.items()}


def compute_avgdl(documents: list[list[str]]) -> float:
    """计算平均文档长度。"""
    if not documents:
        return 1.0
    return sum(len(doc) for doc in documents) / len(documents)


def bm25_score(query_tokens: list[str], doc_tokens: list[str],
               idf: dict[str, float], avgdl: float) -> float:
    """计算单个文档的 BM25 分数。"""
    score = 0.0
    doc_len = len(doc_tokens)
    # 统计文档中各 token 的频率
    tf = {}
    for t in doc_tokens:
        tf[t] = tf.get(t, 0) + 1

    for qt in query_tokens:
        if qt not in idf:
            continue
        f = tf.get(qt, 0)
        if f == 0:
            continue
        # BM25 TF component
        numerator = f * (BM25_K1 + 1)
        denominator = f + BM25_K1 * (1 - BM25_B + BM25_B * doc_len / avgdl)
        score += idf[qt] * numerator / denominator
    return score


# ─── Search ───────────────────────────────────────────────────────
def search_memories(keywords: str):
    """BM25 检索 memory 目录，召回匹配关键词的 Rule + Lesson。按分数降序返回。"""
    results = {"rules": [], "lessons": [], "keywords": keywords.split()}
    query_tokens = tokenize(keywords)

    # 搜索路径: hermes/rule, shared/rule, hermes/lesson, shared/lesson
    search_dirs = [
        (os.path.join(BASE, AGENT, "rule"), "rule", "private"),
        (os.path.join(BASE, "shared", "rule"), "rule", "shared"),
        (os.path.join(BASE, AGENT, "lesson"), "lesson", "private"),
        (os.path.join(BASE, "shared", "lesson"), "lesson", "shared"),
    ]

    # Step 1: 收集所有 active memory + 计算文档 tokens
    all_mems = []
    all_docs = []

    for search_dir, mem_type, scope in search_dirs:
        if not os.path.isdir(search_dir):
            continue
        for f in sorted(glob.glob(os.path.join(search_dir, "*.json"))):
            with open(f) as fh:
                mem = json.load(fh)
            if mem.get("status") != "active":
                continue
            text_to_index = json.dumps(mem, ensure_ascii=False).lower()
            doc_tokens = tokenize(text_to_index)
            if not doc_tokens:
                continue
            all_mems.append((mem, mem_type, scope, f))
            all_docs.append(doc_tokens)

    # Step 2: 计算 IDF + avgdl
    if not all_docs:
        return results

    idf = compute_idf(all_docs)
    avgdl = compute_avgdl(all_docs)

    # Step 3: 对每个文档计算 BM25 分数
    scored = []
    for i, (mem, mem_type, scope, filepath) in enumerate(all_mems):
        score = bm25_score(query_tokens, all_docs[i], idf, avgdl)
        if score > 0:
            scored.append((score, mem, mem_type, scope, filepath))

    # Step 4: 按分数降序排列
    scored.sort(key=lambda x: x[0], reverse=True)

    # Step 5: 构建输出
    for score, mem, mem_type, scope, filepath in scored:
        entry = {
            "id": mem.get("id"),
            "file": os.path.basename(filepath),
            "bm25_score": round(score, 4),
            "matched_keywords": [],  # 向后兼容：BM25 不再计算逐词命中，留空列表
            "type": mem_type,
            "scope": scope,
        }
        if mem_type == "rule":
            entry["rule"] = mem.get("rule", "")
            entry["trigger"] = mem.get("trigger", "")
            entry["consequence"] = mem.get("consequence", "")
        elif mem_type == "lesson":
            entry["lesson"] = mem.get("lesson", "")
            entry["impact"] = mem.get("impact", "")
            entry["reinforcement_count"] = mem.get("reinforcement_count", 0)

        if mem_type == "rule":
            results["rules"].append(entry)
        else:
            results["lessons"].append(entry)

    return results


if __name__ == "__main__":
    keywords = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else ""

    if not keywords.strip():
        print(json.dumps({
            "hook": "task_start",
            "agent": AGENT,
            "timestamp": datetime.now().isoformat(),
            "status": "skipped",
            "reason": "no keywords provided",
            "results": {"rules": [], "lessons": []}
        }, ensure_ascii=False, indent=2))
        sys.exit(0)

    results = search_memories(keywords)

    total_hits = len(results["rules"]) + len(results["lessons"])
    output = {
        "hook": "task_start",
        "agent": AGENT,
        "timestamp": datetime.now().isoformat(),
        "status": "ok" if total_hits > 0 else "no_match",
        "stats": {
            "rules_matched": len(results["rules"]),
            "lessons_matched": len(results["lessons"])
        },
        "results": results
    }

    print(json.dumps(output, ensure_ascii=False, indent=2))