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,
|
|||
|
|
)
|