406 lines
17 KiB
Python
406 lines
17 KiB
Python
from __future__ import annotations
|
||
|
||
import copy
|
||
import hashlib
|
||
import math
|
||
import re
|
||
from collections.abc import Iterable
|
||
from dataclasses import dataclass, field
|
||
from datetime import UTC, datetime
|
||
from typing import Any
|
||
|
||
from server.timeutil import fmt_dt
|
||
|
||
LAYERS = ("fixed", "project", "session", "retrieval")
|
||
|
||
DEFAULT_SYSTEM_PROMPT = (
|
||
"你是 APS 排产智能体。你只拥有提议权：一切写操作必须经门禁确认；"
|
||
"来源与提交语义必须区分，中间状态不得当作已完成事实；"
|
||
"失败必须显式报告，不得用泛化确认掩盖。"
|
||
)
|
||
|
||
DEFAULT_WINDOW = 8
|
||
|
||
_CJK_RE = re.compile(r"[\u4e00-\u9fff]")
|
||
_MESSAGE_ID_FALLBACK = "m{idx}"
|
||
|
||
|
||
def _now_local() -> datetime:
|
||
"""本地无时区时间（与仓库 fmt_dt 语义一致）。"""
|
||
return datetime.now(UTC).astimezone().replace(tzinfo=None)
|
||
|
||
|
||
def estimate_tokens(text: str) -> int:
|
||
"""确定性 token 估算：CJK 字符按 1 token，其余按 4 字符 1 token。
|
||
|
||
仅用于预算控制（可预测、确定性），不等同真实分词器。
|
||
"""
|
||
if not text:
|
||
return 0
|
||
cjk = len(_CJK_RE.findall(text))
|
||
other = len(text) - cjk
|
||
return cjk + math.ceil(other / 4)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class ContextBudget:
|
||
"""四层 token 预算；各层独立限额，总预算 = 各层之和。"""
|
||
|
||
fixed: int = 1500 # L-固定：系统提示
|
||
project: int = 3000 # L-项目：项目资料摘要
|
||
session: int = 6000 # L-会话：滚动摘要 + 钉住 + 最近 N 轮
|
||
retrieval: int = 4000 # L-检索：RAG 命中 + 引用解析结果
|
||
|
||
def layer_limit(self, layer: str) -> int:
|
||
if layer not in LAYERS:
|
||
raise ValueError(f"unknown layer: {layer!r}")
|
||
return getattr(self, layer)
|
||
|
||
def total(self) -> int:
|
||
return self.fixed + self.project + self.session + self.retrieval
|
||
|
||
@classmethod
|
||
def from_dict(cls, values: dict[str, Any]) -> ContextBudget:
|
||
return cls(**{k: int(v) for k, v in (values or {}).items() if k in LAYERS})
|
||
|
||
|
||
class ContextBudgetExceeded(Exception):
|
||
"""某层上下文超出预算：可预测拒绝（fail closed），携带层/用量/限额。"""
|
||
|
||
def __init__(self, layer: str, required: int, budget: int) -> None:
|
||
super().__init__(f"context layer {layer!r} needs {required} tokens, budget {budget}")
|
||
self.layer = layer
|
||
self.required = required
|
||
self.budget = budget
|
||
|
||
|
||
@dataclass
|
||
class ContextAssembly:
|
||
"""四层组装结果：各层文本可检查，整体可注入 prompt。"""
|
||
|
||
layers: dict[str, str] # layer -> 渲染文本
|
||
usage: dict[str, int] # layer -> 估算 token
|
||
budgets: dict[str, int] # layer -> 预算
|
||
pinned: list[dict[str, Any]] = field(default_factory=list) # 已生效钉住项
|
||
summary: dict[str, Any] | None = None # 当前滚动摘要（含溯源）
|
||
|
||
def layer(self, name: str) -> str:
|
||
if name not in self.layers:
|
||
raise ValueError(f"unknown layer: {name!r}")
|
||
return self.layers[name]
|
||
|
||
def text(self) -> str:
|
||
"""按固定/项目/会话/检索顺序拼接为单一 prompt 文本。"""
|
||
return "\n\n".join(f"[{name}] {self.layers[name]}" for name in LAYERS)
|
||
|
||
def tokens(self) -> int:
|
||
return sum(self.usage.values())
|
||
|
||
|
||
# ---------------- 策略状态（钉住 / 滚动摘要） ----------------
|
||
|
||
def _policies(store_data: dict[str, Any]) -> dict[str, dict[str, Any]]:
|
||
return store_data.setdefault("contextPolicies", {})
|
||
|
||
|
||
def get_policy(store_data: dict[str, Any], session_id: str) -> dict[str, Any]:
|
||
"""读会话上下文策略（钉住项/滚动摘要）；不写时返回空策略。"""
|
||
return copy.deepcopy(_policies(store_data).get(session_id) or {})
|
||
|
||
|
||
def _message_id(message: dict[str, Any], idx: int) -> str:
|
||
return str(message.get("id") or _MESSAGE_ID_FALLBACK.format(idx=idx))
|
||
|
||
|
||
def _find_message(store_data: dict[str, Any], session_id: str, message_id: str) -> int:
|
||
messages = (store_data.get("messages") or {}).get(session_id) or []
|
||
for idx, message in enumerate(messages):
|
||
if isinstance(message, dict) and _message_id(message, idx) == message_id:
|
||
return idx
|
||
return -1
|
||
|
||
|
||
def pin_message(store_data: dict[str, Any], session_id: str, message_id: str,
|
||
*, note: str | None = None, by: str | None = None,
|
||
now: datetime | None = None) -> dict[str, Any]:
|
||
"""把关键消息钉进会话上下文：钉住项永不被摘要/窗口淘汰。
|
||
|
||
- 消息必须存在（fail closed：不存在显式 ValueError）
|
||
- 幂等：重复钉同一消息只保留一条
|
||
Returns:
|
||
新钉住项 {messageId, note, pinnedAt, pinnedBy}
|
||
"""
|
||
idx = _find_message(store_data, session_id, message_id)
|
||
if idx < 0:
|
||
raise ValueError(f"message not found: {message_id!r} in session {session_id!r}")
|
||
policies = _policies(store_data)
|
||
policy = policies.setdefault(session_id, {"pinned": [], "summary": None, "summaries": []})
|
||
pinned = policy.setdefault("pinned", [])
|
||
for item in pinned:
|
||
if item.get("messageId") == message_id:
|
||
return copy.deepcopy(item)
|
||
entry = {
|
||
"messageId": message_id,
|
||
"note": note,
|
||
"pinnedAt": fmt_dt(now or _now_local()),
|
||
"pinnedBy": by,
|
||
}
|
||
pinned.append(entry)
|
||
return copy.deepcopy(entry)
|
||
|
||
|
||
def unpin_message(store_data: dict[str, Any], session_id: str, message_id: str) -> bool:
|
||
"""解除钉住；返回是否实际移除。"""
|
||
policies = _policies(store_data)
|
||
policy = policies.get(session_id)
|
||
if not policy:
|
||
return False
|
||
pinned = policy.get("pinned") or []
|
||
kept = [item for item in pinned if item.get("messageId") != message_id]
|
||
if len(kept) == len(pinned):
|
||
return False
|
||
policy["pinned"] = kept
|
||
return True
|
||
|
||
|
||
def scroll_summary(store_data: dict[str, Any], session_id: str, text: str, *,
|
||
source_message_ids: list[str] | tuple[str, ...] | None = None,
|
||
decision_points: Iterable[str] = (),
|
||
open_items: Iterable[str] = (),
|
||
key_numbers: Iterable[str] = (),
|
||
by: str | None = None, now: datetime | None = None,
|
||
summarizer: Any | None = None) -> dict[str, Any]:
|
||
"""滚动压缩：把超出窗口的对话折叠为结构化摘要，可追溯到原消息 id。
|
||
|
||
新摘要继承前一条摘要的溯源链（supersedes），sourceMessageIds 记录
|
||
本次压缩覆盖的原消息 id——「聊了三小时依然记得开头的验收标准」。
|
||
Args:
|
||
summarizer: 压缩生成方标记（如 "llm"/"deterministic"，来自
|
||
summarizer.compress_scroll_summary 的 mode）。默认 None 保持纯确定性
|
||
行为（现有调用不变，不写入新字段）；传入时在摘要记录中落
|
||
compressedBy 溯源，便于审计压缩来自 LLM 还是确定性回退。
|
||
LLM 实际调用是异步的，在 summarizer.compress_scroll_summary 完成，
|
||
本函数负责把压缩结果按既有溯源链入库。
|
||
Returns:
|
||
新摘要 {summaryId, text, sourceMessageIds, decisionPoints,
|
||
openItems, keyNumbers, createdAt, createdBy, supersedes}
|
||
（传入 summarizer 时另含 compressedBy）
|
||
"""
|
||
if not text or not str(text).strip():
|
||
raise ValueError("scroll summary text must not be empty")
|
||
policies = _policies(store_data)
|
||
policy = policies.setdefault(session_id, {"pinned": [], "summary": None, "summaries": []})
|
||
previous = policy.get("summary")
|
||
generation = int((previous or {}).get("generation") or 0) + 1
|
||
digest = hashlib.sha256(
|
||
f"{session_id}|{generation}|{text}".encode()
|
||
).hexdigest()[:12]
|
||
summary = {
|
||
"summaryId": f"sum_{digest}",
|
||
"text": str(text).strip(),
|
||
"sourceMessageIds": list(source_message_ids or []),
|
||
"decisionPoints": list(decision_points or []),
|
||
"openItems": list(open_items or []),
|
||
"keyNumbers": list(key_numbers or []),
|
||
"createdAt": fmt_dt(now or _now_local()),
|
||
"createdBy": by,
|
||
"supersedes": (previous or {}).get("summaryId"),
|
||
"generation": generation,
|
||
}
|
||
if summarizer is not None:
|
||
summary["compressedBy"] = str(summarizer)
|
||
history = policy.setdefault("summaries", [])
|
||
if previous:
|
||
history.append(copy.deepcopy(previous))
|
||
policy["summary"] = summary
|
||
return copy.deepcopy(summary)
|
||
|
||
|
||
def _pinned_entries(store_data: dict[str, Any], session_id: str) -> list[dict[str, Any]]:
|
||
"""把钉住项解析为可注入内容：消息仍存在 -> 全文；被删 -> 显式 missing。"""
|
||
policy = _policies(store_data).get(session_id) or {}
|
||
messages = (store_data.get("messages") or {}).get(session_id) or []
|
||
by_id = {_message_id(m, idx): m for idx, m in enumerate(messages) if isinstance(m, dict)}
|
||
out: list[dict[str, Any]] = []
|
||
for item in policy.get("pinned") or []:
|
||
message = by_id.get(item.get("messageId"))
|
||
if message is None:
|
||
out.append({**copy.deepcopy(item), "missing": True, "text": None})
|
||
else:
|
||
out.append({
|
||
**copy.deepcopy(item),
|
||
"missing": False,
|
||
"role": message.get("role") or "user",
|
||
"text": str(message.get("text") or message.get("content") or ""),
|
||
})
|
||
return out
|
||
|
||
|
||
def _normalize_message(message: dict[str, Any], idx: int) -> dict[str, Any] | None:
|
||
role = str(message.get("role") or "")
|
||
if role in ("assistant", "bot", "ai"):
|
||
role = "agent"
|
||
if role not in ("user", "agent"):
|
||
return None
|
||
text = str(message.get("text") or message.get("content") or "").strip()
|
||
if not text:
|
||
return None
|
||
return {"messageId": _message_id(message, idx), "role": role, "text": text}
|
||
|
||
|
||
# ---------------- 四层组装 ----------------
|
||
|
||
def _fixed_layer(store_data: dict[str, Any]) -> str:
|
||
prompt = str(store_data.get("systemPrompt") or DEFAULT_SYSTEM_PROMPT)
|
||
return prompt.strip()
|
||
|
||
|
||
def _project_layer(store_data: dict[str, Any], session_id: str) -> str:
|
||
"""L-项目：会话归属项目的资料摘要（scope/共享上下文/基准版本）。"""
|
||
session = None
|
||
for row in store_data.get("sessions") or []:
|
||
if row.get("id") == session_id:
|
||
session = row
|
||
break
|
||
if session is None:
|
||
raise ValueError(f"session not found: {session_id!r}")
|
||
project_id = session.get("projectId") or "__personal__"
|
||
lines = [f"项目 ID：{project_id}"]
|
||
if project_id == "__personal__":
|
||
lines.append("作用域：个人空间（未绑定项目）")
|
||
return "\n".join(lines)
|
||
project = None
|
||
for row in store_data.get("projects") or []:
|
||
if row.get("id") == project_id:
|
||
project = row
|
||
break
|
||
if project is None:
|
||
raise ValueError(f"project not accessible: {project_id!r} for session {session_id!r}")
|
||
lines = [f"项目：{project.get('name') or project_id}",
|
||
f"作用域：{project.get('scopeLabel') or '未设定'}"]
|
||
shared = project.get("sharedContext") or []
|
||
if shared:
|
||
lines.append("共享上下文：" + "；".join(str(item) for item in shared))
|
||
for key in ("paramVersion", "baseVersion", "parametersVersion", "baselineVersion"):
|
||
if project.get(key):
|
||
lines.append(f"{key}：{project[key]}")
|
||
if project.get("archived"):
|
||
lines.append("注意：项目已归档，只读。")
|
||
return "\n".join(lines)
|
||
|
||
|
||
def _session_layer(store_data: dict[str, Any], session_id: str,
|
||
*, window: int = DEFAULT_WINDOW) -> tuple[str, list[dict[str, Any]], dict[str, Any] | None]:
|
||
"""L-会话：滚动摘要 + 钉住消息 + 最近 N 轮对话原文（滑动窗口）。"""
|
||
policy = _policies(store_data).get(session_id) or {}
|
||
summary = policy.get("summary")
|
||
pinned = _pinned_entries(store_data, session_id)
|
||
messages = (store_data.get("messages") or {}).get(session_id) or []
|
||
normalized = [item for item in (
|
||
_normalize_message(m, idx) for idx, m in enumerate(messages) if isinstance(m, dict)
|
||
) if item is not None]
|
||
|
||
parts: list[str] = []
|
||
if summary:
|
||
src = "、".join(str(sid) for sid in (summary.get("sourceMessageIds") or []))
|
||
parts.append("-- 滚动摘要 --")
|
||
parts.append(summary.get("text") or "")
|
||
parts.append(f"（摘要 {summary.get('summaryId')}，来源消息：{src}）")
|
||
if pinned:
|
||
parts.append("-- 钉住消息 --")
|
||
for item in pinned:
|
||
if item.get("missing"):
|
||
parts.append(f"[钉住 {item['messageId']} 已删除]"
|
||
+ (f"（备注：{item.get('note')}）" if item.get("note") else ""))
|
||
else:
|
||
who = "用户" if item.get("role") == "user" else "助手"
|
||
note = f"（备注：{item.get('note')}）" if item.get("note") else ""
|
||
parts.append(f"{who}：{item.get('text')}{note}")
|
||
if normalized:
|
||
parts.append("-- 最近对话 --")
|
||
pinned_ids = {item["messageId"] for item in pinned}
|
||
recent = [m for m in normalized[-window:] if m["messageId"] not in pinned_ids]
|
||
for message in recent:
|
||
who = "用户" if message["role"] == "user" else "助手"
|
||
parts.append(f"{who}：{message['text']}")
|
||
return "\n".join(parts), pinned, summary
|
||
|
||
|
||
def _retrieval_layer(store_data: dict[str, Any], query: str | None) -> str:
|
||
"""L-检索：RAG 命中 + 会话引用解析结果（按需注入）。"""
|
||
parts: list[str] = []
|
||
results = store_data.get("ragResults") or store_data.get("retrievalResults")
|
||
if results is None and query and store_data.get("ragAssets"):
|
||
try:
|
||
from server.knowledge.retrieval import hybrid_search
|
||
|
||
units = []
|
||
for asset in store_data.get("ragAssets") or []:
|
||
units.append({
|
||
"assetId": asset.get("assetId"), "title": asset.get("title"),
|
||
"kind": asset.get("kind"), "version": asset.get("version"),
|
||
"tags": asset.get("tags") or [],
|
||
"content": asset.get("content") or "",
|
||
})
|
||
results = hybrid_search(units, query)
|
||
except Exception: # noqa: BLE001 - 检索失败降级为空，绝不阻断上下文组装
|
||
results = []
|
||
if results:
|
||
parts.append("-- RAG 命中 --")
|
||
for hit in results:
|
||
title = hit.get("title") or hit.get("assetId") or "?"
|
||
version = f" v{hit.get('version')}" if hit.get("version") else ""
|
||
snippet = str(hit.get("snippet") or hit.get("content") or "")[:200]
|
||
parts.append(f"[{hit.get('assetId')}] {title}{version}：{snippet}")
|
||
refs = store_data.get("resolvedRefs")
|
||
if refs:
|
||
parts.append("-- 引用解析 --")
|
||
for ref in refs:
|
||
target = ref.get("target_id") or ref.get("targetId") or "?"
|
||
snapshot_at = ref.get("snapshot_at") or ref.get("snapshotAt") or "?"
|
||
parts.append(f"[引用 {ref.get('kind')}:{target} @{snapshot_at}]")
|
||
return "\n".join(parts) if parts else "（无检索命中）"
|
||
|
||
|
||
def assemble_context(store_data: dict[str, Any], session_id: str, query: str | None,
|
||
budget: ContextBudget, *, window: int = DEFAULT_WINDOW) -> ContextAssembly:
|
||
"""四层组装上下文；任何一层超预算 -> ContextBudgetExceeded（可预测拒绝）。
|
||
|
||
Args:
|
||
store_data: 工作区快照（sessions/projects/messages/contextPolicies/ragResults...）
|
||
session_id: 当前会话 ID
|
||
query: 用户问题（用于按需检索；可为 None）
|
||
budget: 四层 token 预算
|
||
window: 会话层最近 N 轮窗口
|
||
Returns:
|
||
ContextAssembly：layers/usage/budgets 可检查，text() 可注入 prompt
|
||
Raises:
|
||
ContextBudgetExceeded: 某层估算 token 超预算（fail closed）
|
||
ValueError: 会话/项目不存在
|
||
"""
|
||
layers: dict[str, str] = {}
|
||
usage: dict[str, int] = {}
|
||
budgets: dict[str, int] = {}
|
||
pinned: list[dict[str, Any]] = []
|
||
summary: dict[str, Any] | None = None
|
||
for layer in LAYERS:
|
||
if layer == "fixed":
|
||
text = _fixed_layer(store_data)
|
||
elif layer == "project":
|
||
text = _project_layer(store_data, session_id)
|
||
elif layer == "session":
|
||
text, pinned, summary = _session_layer(store_data, session_id, window=window)
|
||
else:
|
||
text = _retrieval_layer(store_data, query)
|
||
limit = budget.layer_limit(layer)
|
||
used = estimate_tokens(text)
|
||
if used > limit:
|
||
raise ContextBudgetExceeded(layer=layer, required=used, budget=limit)
|
||
layers[layer] = text
|
||
usage[layer] = used
|
||
budgets[layer] = limit
|
||
return ContextAssembly(
|
||
layers=layers, usage=usage, budgets=budgets,
|
||
pinned=pinned, summary=summary,
|
||
)
|