hermes_core/cluster.py

"""动态语义聚类 — scope 分配、合并、拆分."""

import uuid
import sqlite3
import numpy as np
from datetime import datetime, timezone
from hermes_core.types import Scope, Dimension
from hermes_core.db import (
    get_active_scopes, get_active_records, upsert_scope, get_scope, get_records_by_scope
)
from hermes_core.embedder import Embedder


DEFAULT_MATCH_THRESHOLD = 0.70
DEFAULT_MERGE_THRESHOLD = 0.90
DEFAULT_SPLIT_THRESHOLD = 0.60


def _new_scope_id() -> str:
    return f"scope_{uuid.uuid4().hex[:8]}"


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def assign_scope(conn: sqlite3.Connection, embedder: Embedder,
                 scope_desc: str) -> tuple[str, str]:
    """为一条 scope_desc 分配 scope。匹配现有 scope 或创建新 scope。

    Returns:
        (scope_id, scope_label)
    """
    vec = embedder.encode(scope_desc)
    scopes = get_active_scopes(conn)

    best_score = 0.0
    best_scope = None

    for scope in scopes:
        if scope.centroid is None:
            continue
        sim = embedder.cosine_similarity(vec, scope.centroid)
        if sim > best_score:
            best_score = sim
            best_scope = scope

    if best_score >= DEFAULT_MATCH_THRESHOLD and best_scope is not None:
        return best_scope.id, best_scope.label

    # 不匹配任何现有 scope → 创建新 scope
    sid = _new_scope_id()
    new_scope = Scope(
        id=sid,
        label=scope_desc,
        centroid=vec,
        record_count=0,
        created_at=_now(),
        last_activity=_now(),
    )
    upsert_scope(conn, new_scope)
    return sid, scope_desc


def recluster_scope(conn: sqlite3.Connection, embedder: Embedder,
                    scope_id: str) -> None:
    """重新计算 scope 的聚类中心和内聚度。"""
    records = get_records_by_scope(conn, scope_id)
    if not records:
        return

    scope = get_scope(conn, scope_id)
    if scope is None:
        return

    labels = [r.scope_label for r in records]
    vecs = embedder.encode_batch(labels)
    centroid = np.mean(vecs, axis=0).tolist()

    # 计算内聚度：所有 label 与 centroid 的平均相似度
    similarities = [embedder.cosine_similarity(v, centroid) for v in vecs]
    coherence = float(np.mean(similarities)) if similarities else 1.0

    scope.centroid = centroid
    scope.record_count = len(records)
    scope.coherence = coherence
    scope.last_activity = _now()
    upsert_scope(conn, scope)


def check_merge(conn: sqlite3.Connection, embedder: Embedder,
                threshold: float = DEFAULT_MERGE_THRESHOLD) -> list[tuple[str, str]]:
    """检查所有 active scope 两两之间是否应该合并。

    Returns:
        [(scope_id_a, scope_id_b), ...] 需要合并的 scope 对
    """
    scopes = get_active_scopes(conn)
    pairs = []
    for i in range(len(scopes)):
        for j in range(i + 1, len(scopes)):
            a, b = scopes[i], scopes[j]
            if a.centroid is None or b.centroid is None:
                continue
            sim = embedder.cosine_similarity(a.centroid, b.centroid)
            if sim >= threshold:
                pairs.append((a.id, b.id))
    return pairs


def check_split(conn: sqlite3.Connection, embedder: Embedder,
                scope_id: str, threshold: float = DEFAULT_SPLIT_THRESHOLD) -> bool:
    """检查一个 scope 内聚度是否过低，需要拆分。

    Returns:
        True if split is needed
    """
    recluster_scope(conn, embedder, scope_id)
    scope = get_scope(conn, scope_id)
    if scope is None:
        return False
    return scope.coherence < threshold