初始化提交

backend/app/modules/lorebook/budget.py

@@ -0,0 +1,58 @@
+"""世谱预算裁剪与注入分组。
+
+按 token 预算裁剪（priority 值小优先保留），再按 insertion_order 排定顺序，
+最后按 position 分组为可注入的文本块。纯逻辑，无 IO。
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from app.modules.lorebook.types import MatchedEntry
+
+TokenCounter = Callable[[str], int]
+
+_POSITIONS = ("before_char", "after_char")
+
+
+def _default_token_counter(text: str) -> int:
+    """P0 以字符数近似 token；后续接神识 tokenizer。"""
+    return len(text)
+
+
+def select_within_budget(
+    matched: list[MatchedEntry],
+    token_budget: int | None = None,
+    token_counter: TokenCounter | None = None,
+) -> list[MatchedEntry]:
+    """按 priority 选入预算内条目，再按 insertion_order 排定注入顺序。"""
+    counter = token_counter or _default_token_counter
+    if token_budget is None:
+        return sorted(matched, key=lambda m: m.entry.insertion_order)
+
+    by_priority = sorted(matched, key=lambda m: m.entry.priority)
+    kept: list[MatchedEntry] = []
+    total = 0
+    for m in by_priority:
+        cost = counter(m.entry.content)
+        if cost == 0:
+            continue
+        if total + cost > token_budget:
+            continue
+        kept.append(m)
+        total += cost
+    kept.sort(key=lambda m: m.entry.insertion_order)
+    return kept
+
+
+def build_injections(matched: list[MatchedEntry]) -> dict[str, str]:
+    """按 position 分组，合并同组条目 content 为单段文本。
+
+    返回 ``{"before_char": "...", "after_char": "..."}``，空组不出现在结果中。
+    集成层据此把 before_char 插在道体 personality_summary 之前、after_char 之后。
+    """
+    groups: dict[str, list[str]] = {pos: [] for pos in _POSITIONS}
+    for m in matched:
+        pos = m.entry.position if m.entry.position in groups else "after_char"
+        groups[pos].append(m.entry.content)
+    return {pos: "\n".join(parts) for pos, parts in groups.items() if parts}