kb-memory 2.1.0__py3-none-any.whl

kb/embedder.py

@@ -0,0 +1,134 @@
+"""嵌入器：嵌入模型延迟加载 + 设备支持 + 归一化向量输出（实例化不加载模型）。
+
+English: Embedder: lazy-loaded embedding model, device support, normalized vectors
+(construction never loads the model).
+"""
+import gc
+import threading
+import time
+
+
+class EmbeddingUnavailableError(RuntimeError):
+    """嵌入模型不可用（下载/加载/推理失败）；API 层据此返回 503 结构化错误。
+    English: Embedding model unavailable (download/load/encode failure); the API layer
+    maps this to a structured 503."""
+
+    def __init__(self, model_name: str, cause: str):
+        self.model_name = model_name
+        self.cause = cause
+        super().__init__(f"嵌入模型不可用：{model_name}（{cause}）")
+
+
+class Embedder:
+    """SentenceTransformer 封装；构造只存参数，首次 embed 才加载模型。
+
+    English: SentenceTransformer wrapper; the constructor only stores parameters,
+    and the model is loaded on the first embed call.
+    """
+
+    def __init__(self, model_name: str, device: str = "cpu",
+                 idle_unload_seconds: int = 0):
+        """只存参数不加载模型；_model 延迟到首次使用。
+        English: Store parameters only; defer loading until first use."""
+        self.model_name = model_name
+        self.device = device
+        self.idle_unload_seconds = max(0, int(idle_unload_seconds or 0))
+        self._model = None
+        self._state = "unloaded"   # unloaded | loading | ok | error
+        self._error = ""
+        self._last_used = time.monotonic()
+        self._lock = threading.RLock()
+
+    def status(self) -> dict:
+        """返回可放进 healthz 的轻量状态（不触发加载）。
+        English: Return a lightweight status suitable for healthz (never triggers loading)."""
+        return {"state": self._state, "model": self.model_name,
+                "error": self._error,
+                "idle_unload_seconds": self.idle_unload_seconds}
+
+    def _mark_error(self, exc: Exception) -> EmbeddingUnavailableError:
+        """记录失败状态并包装成统一异常。
+        English: Record the failure state and wrap it in the unified exception."""
+        self._state = "error"
+        self._error = str(exc)[:300]
+        return EmbeddingUnavailableError(self.model_name, self._error)
+
+    def _ensure_loaded(self):
+        """首次使用时加载模型；优先离线（缓存命中，不联网），失败再在线下载；cuda 下 fp16。
+
+        加载失败时状态置 error，并抛 EmbeddingUnavailableError（不把底层
+        OSError/HTTPError 裸抛给调用方）。
+        English: Load the model on first use — offline first (cache hit, no network),
+        then online download; fp16 under cuda. On failure the state becomes error and
+        EmbeddingUnavailableError is raised instead of leaking OSError/HTTPError.
+        """
+        with self._lock:
+            if self._model is not None:
+                self._last_used = time.monotonic()
+                return
+            self._state = "loading"
+            try:
+                from sentence_transformers import SentenceTransformer
+                try:
+                    self._model = SentenceTransformer(
+                        self.model_name, device=self.device, local_files_only=True)
+                except Exception:
+                    # 缓存未命中（首次部署新模型），走在线下载；HF_ENDPOINT 镜像由环境变量提供
+                    self._model = SentenceTransformer(self.model_name, device=self.device)
+                if self.device == "cuda":
+                    self._model.half()
+                self._state = "ok"
+                self._error = ""
+                self._last_used = time.monotonic()
+            except Exception as exc:  # noqa: BLE001 - 统一转换为可读错误
+                raise self._mark_error(exc) from exc
+
+    def unload(self) -> bool:
+        """卸载模型释放内存；未加载时返回 False。
+        English: Unload the model to free memory; returns False when not loaded."""
+        with self._lock:
+            if self._model is None:
+                return False
+            self._model = None
+            self._state = "unloaded"
+            self._error = ""
+            gc.collect()
+            if self.device == "cuda":
+                try:
+                    import torch
+                    if torch.cuda.is_available():
+                        torch.cuda.empty_cache()
+                except Exception:  # noqa: BLE001 - 释放显存失败不影响卸载
+                    pass
+            return True
+
+    def unload_if_idle(self, now: float | None = None) -> bool:
+        """空闲超过阈值时卸载模型；返回是否真正卸载。
+        English: Unload when idle beyond the threshold; returns whether it actually unloaded."""
+        if self.idle_unload_seconds <= 0 or self._model is None:
+            return False
+        now = time.monotonic() if now is None else now
+        if now - self._last_used < self.idle_unload_seconds:
+            return False
+        return self.unload()
+
+    def embed_texts(self, texts: list[str]) -> list[list[float]]:
+        """批量编码，normalize_embeddings=True（单位向量）。
+        English: Batch encode with normalize_embeddings=True (unit vectors)."""
+        try:
+            with self._lock:
+                self._ensure_loaded()
+                vecs = self._model.encode(
+                    texts, normalize_embeddings=True, convert_to_numpy=True
+                )
+                self._last_used = time.monotonic()
+        except EmbeddingUnavailableError:
+            raise
+        except Exception as exc:  # noqa: BLE001 - 推理失败同样视为模型不可用
+            raise self._mark_error(exc) from exc
+        return [v.tolist() for v in vecs]
+
+    def embed_query(self, query: str) -> list[float]:
+        """单条查询编码，同上归一化。
+        English: Encode a single query with the same normalization."""
+        return self.embed_texts([query])[0]

kb/eval.py

@@ -0,0 +1,136 @@
+"""最小评测基准：中文 QA 数据集 → Recall@1/@5 + MRR（N26，A3.5 spec §3.5）。
+
+评测流程：
+  独立 KB_DATA_DIR（CLI 层创建临时目录，不碰生产库）
+    → 导入全部 corpus（每条独立成块，互为干扰项）
+    → 逐条 question 检索（指定 mode/top_k，治理默认关）
+    → 命中判定：corpus 记录 id ∈ top-k 结果
+    → 指标：Recall@1 / Recall@5 / MRR = mean(1/rank)
+    → 报告：总体 + 分难度 + 未命中清单 + 平均延迟
+
+run_eval 与 service 解耦（鸭子类型：add_memory / search），可注入替身单测。
+
+English: Minimal evaluation benchmark: Chinese QA dataset → Recall@1/@5 + MRR (N26, A3.5 spec §3.5).
+Evaluation flow: isolated KB_DATA_DIR → import all corpus → retrieve per question → hit judgment
+→ metrics (Recall@1 / Recall@5 / MRR) → report. run_eval is decoupled from service
+(duck-typed add_memory / search), so a test stub can be injected for unit tests.
+"""
+import json
+import time
+from pathlib import Path
+
+DIFFICULTY_TAGS = ("keyword", "semantic", "distractor")
+
+
+class EvalDatasetError(Exception):
+    """数据集加载失败（文件缺失 / JSONL 损坏 / 必填字段缺失）。
+    English: Dataset load failure (missing file / corrupt JSONL / missing required fields)."""
+
+
+def load_dataset(path) -> list[dict]:
+    """加载 JSONL 数据集；每行 {"qid", "question", "corpus", "tags"}。
+
+    字段缺失 / JSON 损坏 / 文件不存在抛 EvalDatasetError。
+    English: Load a JSONL dataset; each line is {"qid", "question", "corpus", "tags"}.
+    Missing fields / corrupt JSON / a nonexistent file raise EvalDatasetError."""
+    path = Path(path)
+    if not path.exists():
+        raise EvalDatasetError(f"数据集不存在: {path}")
+    items = []
+    try:
+        for line_no, line in enumerate(
+                path.read_text(encoding="utf-8").splitlines(), start=1):
+            line = line.strip()
+            if not line:
+                continue
+            item = json.loads(line)
+            for field in ("qid", "question", "corpus"):
+                if field not in item:
+                    raise EvalDatasetError(
+                        f"第 {line_no} 行缺字段 {field}: {item}")
+            items.append(item)
+    except json.JSONDecodeError as e:
+        raise EvalDatasetError(f"JSONL 解析失败（{path}）: {e}")
+    return items
+
+
+def compute_metrics(ranks: list[int | None]) -> dict:
+    """由 1-based 排名列表计算 Recall@1 / Recall@5 / MRR。
+
+    ranks: 每条问题的目标记录排名（1 起）；None=未进 top-k。
+    空列表安全返回全 0。
+    English: Compute Recall@1 / Recall@5 / MRR from a 1-based rank list.
+    ranks: the target-record rank per question (starting at 1); None = not in top-k.
+    An empty list safely returns all zeros."""
+    n = len(ranks)
+    if n == 0:
+        return {"recall_at_1": 0.0, "recall_at_5": 0.0, "mrr": 0.0}
+    hit1 = sum(1 for r in ranks if r == 1)
+    hit5 = sum(1 for r in ranks if r is not None and r <= 5)
+    mrr = sum(1 / r for r in ranks if r is not None) / n
+    return {"recall_at_1": hit1 / n, "recall_at_5": hit5 / n, "mrr": mrr}
+
+
+def _difficulty_of(item: dict) -> str | None:
+    """提取难度标签（keyword/semantic/distractor）；无则 None。
+    English: Extract the difficulty tag (keyword/semantic/distractor); None if absent."""
+    for t in item.get("tags") or []:
+        if t in DIFFICULTY_TAGS:
+            return t
+    return None
+
+
+def run_eval(service, dataset: list[dict], top_k: int = 5,
+             mode: str = "hybrid") -> dict:
+    """端到端评测：写入 corpus → 逐条检索 → 指标汇总。
+
+    service: KBService 或替身（需 add_memory(content, tags=...) / search(...)）。
+    返回报告 dict：{count, top_k, mode, recall_at_1/5, mrr, latency_ms_avg,
+                    by_difficulty, misses}。
+    English: End-to-end evaluation: write corpus → retrieve per question → aggregate metrics.
+    service: a KBService or stub (needs add_memory(content, tags=...) / search(...)).
+    Returns a report dict: {count, top_k, mode, recall_at_1/5, mrr, latency_ms_avg,
+    by_difficulty, misses}."""
+    corpus_ids: dict[int, str] = {}  # qid → record_id
+    for item in dataset:
+        rec = service.add_memory(item["corpus"],
+                                 tags=list(item.get("tags") or ["eval"]))
+        corpus_ids[item["qid"]] = rec.id
+
+    ranks: list[int | None] = []
+    latency_total = 0.0
+    misses = []
+    for item in dataset:
+        t0 = time.perf_counter()
+        results = service.search(item["question"], top_k=top_k, mode=mode)
+        latency_total += (time.perf_counter() - t0) * 1000.0
+        target_id = corpus_ids[item["qid"]]
+        rank = None
+        for i, r in enumerate(results, start=1):
+            if r["id"] == target_id:
+                rank = i
+                break
+        ranks.append(rank)
+        if rank is None:
+            misses.append({"qid": item["qid"], "question": item["question"]})
+
+    overall = compute_metrics(ranks)
+    by_difficulty: dict[str, dict] = {}
+    for tag in DIFFICULTY_TAGS:
+        sub_ranks = [r for item, r in zip(dataset, ranks)
+                     if _difficulty_of(item) == tag]
+        if sub_ranks:
+            by_difficulty[tag] = compute_metrics(sub_ranks)
+            by_difficulty[tag]["count"] = len(sub_ranks)
+
+    return {
+        "count": len(dataset),
+        "top_k": top_k,
+        "mode": mode,
+        "recall_at_1": overall["recall_at_1"],
+        "recall_at_5": overall["recall_at_5"],
+        "mrr": overall["mrr"],
+        "latency_ms_avg": latency_total / max(1, len(dataset)),
+        "by_difficulty": by_difficulty,
+        "misses": misses,
+    }

kb/governance.py

@@ -0,0 +1,188 @@
+"""记忆治理 A3：访问频率衰减纯函数模块（N21b，TASK-0068）+ 语义去重（N22a，TASK-0069）。
+
+纯函数优先：输入标量输出分数，不依赖 Record/Store/配置，可直接单测。
+衰减公式（A3 spec §3.1）：
+  decay_factor = exp(-λ * days_since_last_accessed) * (1 + γ * log₂(1 + access_count))
+  final_score = rrf_score * decay_factor
+
+语义去重（A3 spec §3.2，卡片简化为 409 拦截）：
+  check_duplicate(content, store, embedder, threshold) → (existing_id|None, similarity)
+  写入前向量检索 top1，相似度 ≥ 阈值视为重复，service 层抛 DuplicateError，api 层返回 409。
+
+默认 KB_DECAY_ENABLED=false / KB_DEDUP_ENABLED=false，零行为变化。
+BM25 路径不受衰减影响（仅向量/RRF 排序受影响）。
+"""
+import math
+from datetime import datetime
+
+
+class DuplicateError(Exception):
+    """语义去重命中异常（N22a/TASK-0069）。
+
+    service.add_memory 在 dedup_enabled 且 check_duplicate 命中时抛出；
+    api.create_memory 捕获后返回 409 + duplicate_of 字段。
+    """
+
+    def __init__(self, existing_id: str, similarity: float):
+        self.existing_id = existing_id
+        self.similarity = similarity
+        super().__init__(f"语义重复（相似度 {similarity:.3f}），已有记录 {existing_id}")
+
+
+def compute_decay_factor(days_since_last_accessed: float,
+                         access_count: int,
+                         lambda_: float = 0.02,
+                         gamma: float = 0.3) -> float:
+    """计算衰减因子（A3 spec §3.1）。
+
+    decay_factor = exp(-λ * days) * (1 + γ * log₂(1 + access_count))
+
+    参数：
+        days_since_last_accessed: 距上次访问的天数（>=0，last_accessed 为空时用 created_at）
+        access_count: 累计访问次数（>=0，新记录=0）
+        lambda_: 衰减速率 λ（/天），默认 0.02（半衰期≈35天）
+        gamma: 高频访问加权系数 γ，默认 0.3（access_count=10→约2.0倍）
+
+    返回：
+        衰减因子（>0）；高频访问记忆获得加权，长期未访问记忆降权。
+    """
+    if days_since_last_accessed < 0:
+        days_since_last_accessed = 0.0
+    if access_count < 0:
+        access_count = 0
+    recency = math.exp(-lambda_ * days_since_last_accessed)
+    frequency = 1.0 + gamma * math.log2(1.0 + access_count)
+    return recency * frequency
+
+
+def apply_decay(rrf_score: float, decay_factor: float) -> float:
+    """对 RRF 融合分应用衰减因子（A3 spec §3.1）。
+
+    final_score = rrf_score * decay_factor
+
+    参数：
+        rrf_score: RRF 融合分（>=0）
+        decay_factor: 衰减因子（compute_decay_factor 的返回值）
+
+    返回：
+        衰减后的最终排序分。
+    """
+    return rrf_score * decay_factor
+
+
+def days_since(last_accessed: str, created_at: str,
+               now: datetime | None = None) -> float:
+    """计算距上次访问的天数；last_accessed 为空时用 created_at（A3 spec §3.1）。
+
+    参数：
+        last_accessed: 上次访问时间（ISO 格式字符串，空串=""表示从未访问）
+        created_at: 记录创建时间（ISO 格式字符串，从未访问时用此替代）
+        now: 当前时间（None=datetime.now()）
+
+    返回：
+        天数差（>=0，未来时间视为0）。解析失败时返回0.0（安全降级）。
+    """
+    if now is None:
+        now = datetime.now()
+    ref_str = last_accessed if last_accessed else created_at
+    if not ref_str:
+        return 0.0
+    try:
+        ref = datetime.fromisoformat(ref_str)
+    except (ValueError, TypeError):
+        return 0.0
+    delta = (now - ref).total_seconds() / 86400.0
+    return max(0.0, delta)
+
+
+def freshness_boost(days_since_updated: float,
+                    beta: float = 0.05,
+                    alpha: float = 0.3) -> float:
+    """计算新鲜度加权因子（A3 spec §3.3）。
+
+    recency = exp(-β * days_since_updated)
+    boost = 1 + α * recency
+
+    与衰减（§3.1）正交：衰减看 last_accessed（访问冷热），新鲜度看 updated_at
+    （内容新旧），两者独立可叠加（正交相乘）。
+
+    参数：
+        days_since_updated: 距上次更新的天数（>=0，用 updated_at）
+        beta: 新鲜度衰减速率 β（/天），默认 0.05（半衰期≈14天）
+        alpha: 新鲜度加权上限系数 α，默认 0.3（boost 范围 [1, 1.3]）
+
+    返回：
+        新鲜度加权因子（范围 [1, 1+α]，新记忆接近 1+α，旧记忆接近 1）。
+    """
+    if days_since_updated < 0:
+        days_since_updated = 0.0
+    recency = math.exp(-beta * days_since_updated)
+    return 1.0 + alpha * recency
+
+
+def compute_stats(records, now: datetime | None = None) -> dict:
+    """计算治理统计（A3 spec §4.2，TASK-0070）：纯函数，输入记录迭代器输出统计 dict。
+
+    返回：
+        total_count: 总记录数
+        avg_access_count: 平均 access_count（保留2位小数）
+        stale_90d_count: 超 90 天未命中数（last_accessed 为空时用 created_at）
+
+    access_count/last_accessed 用 getattr 兼容 TASK-0067 未合入（未合入时均为 0/""）。
+    """
+    if now is None:
+        now = datetime.now()
+    total = 0
+    access_sum = 0
+    stale_90d = 0
+    for rec in records:
+        total += 1
+        ac = getattr(rec, "access_count", 0) or 0
+        access_sum += ac
+        last_accessed = getattr(rec, "last_accessed", "") or ""
+        created_at = getattr(rec, "created_at", "") or ""
+        days = days_since(last_accessed, str(created_at), now)
+        if days > 90:
+            stale_90d += 1
+    avg_access = (access_sum / total) if total > 0 else 0.0
+    return {
+        "total_count": total,
+        "avg_access_count": round(avg_access, 2),
+        "stale_90d_count": stale_90d,
+    }
+
+
+# ---- 语义去重（N22a/TASK-0069）----
+
+def check_duplicate(content: str, store, embedder,
+                    threshold: float = 0.92) -> tuple[str | None, float]:
+    """语义去重检查（A3 spec §3.2，top1 拦截）。
+
+    对新内容做向量检索 top1，若最高余弦相似度 ≥ threshold 则视为重复，
+    返回 (existing_id, similarity)；否则返回 (None, 0.0)。
+
+    参数：
+        content: 待写入的新内容文本
+        store: ChromaStore 实例（有 query(embedding, top_k) 方法）
+        embedder: Embedder 实例（有 embed_query(text) 方法）
+        threshold: 余弦相似度阈值，默认 0.92（≥此值视为重复）
+
+    返回：
+        (existing_id, similarity) 命中重复；(None, 0.0) 无重复。
+
+    异常降级：嵌入或检索失败时返回 (None, 0.0)（不阻塞写入，记 WARNING 由调用方处理）。
+    """
+    import logging
+    logger = logging.getLogger("kb.governance")
+    try:
+        vec = embedder.embed_query(content)
+        hits = store.query(vec, top_k=1)
+        if not hits:
+            return (None, 0.0)
+        _rec, score = hits[0]
+        if score >= threshold:
+            return (_rec.id, float(score))
+        return (None, 0.0)
+    except Exception as e:
+        logger.warning("语义去重检查失败，降级为不拦截: %s", e)
+        return (None, 0.0)