164 lines
6.6 KiB
Python
164 lines
6.6 KiB
Python
"""会话历史：数据库保存完整对话，每次请求按 token 预算带最近的轮次，最近一轮总是带上
|
|
|
|
裁剪在服务层组装请求时做一次，意图判断、问答和 screen agent 看到的是同一段历史
|
|
每轮是 {user, assistant, outcome}，有轮次记录时另带 changes、queries、clarifications（字段同 Turn）：
|
|
用户原话和回复转成对话消息，记录由 history_records 放进本次请求的 JSON，两者共用一个预算
|
|
"""
|
|
|
|
import json
|
|
from functools import cache
|
|
|
|
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
|
|
from langchain_core.messages.utils import count_tokens_approximately
|
|
from loguru import logger
|
|
|
|
from app.schemas.base import JsonObject
|
|
from app.schemas.screen import Screen
|
|
|
|
RECORD_KEYS = ("changes", "queries", "clarifications")
|
|
# 放进请求的单条 SQL 上限（字符）；更长的不带，原文仍在会话记录里
|
|
RECORD_SQL_MAX_CHARS = 2000
|
|
# 每条消息在请求里的固定开销（角色、分隔），计入历史预算
|
|
MESSAGE_OVERHEAD = 4
|
|
# 最近一轮超过预算被截短时，回复和用户原话各自至少保留的 token
|
|
MIN_KEPT_TOKENS = 200
|
|
|
|
|
|
@cache
|
|
def _encoding():
|
|
"""o200k 近似各家模型的 token；编码表第一次使用时由 tiktoken 下载并缓存，拿不到时退回按字符估算"""
|
|
try:
|
|
import tiktoken
|
|
|
|
return tiktoken.get_encoding("o200k_base")
|
|
except Exception: # noqa: BLE001 计数只用来裁剪历史，不能让请求失败
|
|
logger.exception("tiktoken 编码表加载失败，历史预算改按字符估算，中文会低估")
|
|
return None
|
|
|
|
|
|
def count_tokens(text: str) -> int:
|
|
encoding = _encoding()
|
|
if encoding is None:
|
|
return count_tokens_approximately([HumanMessage(content=text)])
|
|
return len(encoding.encode(text, disallowed_special=()))
|
|
|
|
|
|
def _record(turn: JsonObject, ago: int, present: set[str] | None) -> JsonObject | None:
|
|
"""一轮放进请求的记录，没有的字段不出现；present 是当前大屏的组件 id，None 时按都不在算（只用于计数）
|
|
|
|
过长的 SQL 截断了就不能执行，这条查询整条不带，原文仍在会话记录里
|
|
"""
|
|
record: JsonObject = {key: turn[key] for key in RECORD_KEYS if turn.get(key)}
|
|
if "queries" in record:
|
|
record["queries"] = [
|
|
query
|
|
for query in record["queries"]
|
|
if len(query["sql"]) <= RECORD_SQL_MAX_CHARS
|
|
] or None
|
|
if "changes" in record:
|
|
record["changes"] = [
|
|
{**change, "on_screen": _on_screen(change, present or set())}
|
|
for change in record["changes"]
|
|
]
|
|
record = {key: value for key, value in record.items() if value}
|
|
return {"ago": ago, **record} if record else None
|
|
|
|
|
|
def _on_screen(change: JsonObject, present: set[str]) -> bool:
|
|
if change["type"] == "canvas":
|
|
return True
|
|
return change["node_id"] in present
|
|
|
|
|
|
def turn_tokens(turn: JsonObject) -> int:
|
|
"""一轮在请求里的大小：两条消息（含每条的固定开销）和放进请求的记录"""
|
|
tokens = (
|
|
count_tokens(turn["user"])
|
|
+ count_tokens(turn["assistant"])
|
|
+ MESSAGE_OVERHEAD * 2
|
|
)
|
|
if record := _record(turn, ago=1, present=None):
|
|
tokens += count_tokens(json.dumps(record, ensure_ascii=False))
|
|
return tokens
|
|
|
|
|
|
def history_messages(history: list[JsonObject]) -> list[BaseMessage]:
|
|
"""每轮转成一条用户消息和一条助手消息，按时间从早到晚排列；记录不进消息"""
|
|
messages: list[BaseMessage] = []
|
|
for turn in history:
|
|
messages.append(HumanMessage(content=turn["user"]))
|
|
messages.append(AIMessage(content=turn["assistant"]))
|
|
return messages
|
|
|
|
|
|
def trim_history(history: list[JsonObject], max_tokens: int) -> list[JsonObject]:
|
|
"""从最近一轮往前保留，消息和记录合计不超过 max_tokens；更早的轮次放不下就整轮不带
|
|
|
|
最近一轮总是保留：它单独就超过预算时截短再带上（见 _fit），不能因为一条长回答把整段历史丢掉
|
|
"""
|
|
if not history:
|
|
return []
|
|
latest = _fit(history[-1], max_tokens)
|
|
used = turn_tokens(latest)
|
|
kept = [latest]
|
|
for turn in reversed(history[:-1]):
|
|
used += turn_tokens(turn)
|
|
if used > max_tokens:
|
|
break
|
|
kept.append(turn)
|
|
return kept[::-1]
|
|
|
|
|
|
def _fit(turn: JsonObject, max_tokens: int) -> JsonObject:
|
|
"""一轮超过预算时先截回复、再截用户原话，都是去掉中间、保留开头和结尾（提议和结论通常在回复末尾）
|
|
|
|
记录不截：查询最多 10 条、每条 SQL 不超过 RECORD_SQL_MAX_CHARS，默认预算下放得下
|
|
"""
|
|
for key in ("assistant", "user"):
|
|
over = turn_tokens(turn) - max_tokens
|
|
if over <= 0:
|
|
break
|
|
keep = max(count_tokens(turn[key]) - over, MIN_KEPT_TOKENS)
|
|
turn = {**turn, key: truncate_middle(turn[key], keep)}
|
|
return turn
|
|
|
|
|
|
def truncate_middle(text: str, max_tokens: int) -> str:
|
|
"""去掉中间、保留开头和结尾各一半，中间写明省略了多少字；结果不超过 max_tokens"""
|
|
total = count_tokens(text)
|
|
if total <= max_tokens:
|
|
return text
|
|
keep = len(text) * max_tokens // total
|
|
while keep > 0:
|
|
head, tail = text[: keep // 2], text[len(text) - (keep - keep // 2) :]
|
|
shortened = f"{head}\n…（中间省略 {len(text) - keep} 字）…\n{tail}"
|
|
if count_tokens(shortened) <= max_tokens:
|
|
return shortened
|
|
keep = keep * 9 // 10
|
|
return f"…（省略 {len(text)} 字）…"
|
|
|
|
|
|
def intent_history(history: list[JsonObject]) -> list[JsonObject]:
|
|
"""意图判断看的历史：原话、回复、结果和澄清问答；SQL 和修改记录不给"""
|
|
projected = []
|
|
for turn in history:
|
|
item = {key: turn[key] for key in ("user", "assistant", "outcome")}
|
|
if turn.get("clarifications"):
|
|
item["clarifications"] = turn["clarifications"]
|
|
projected.append(item)
|
|
return projected
|
|
|
|
|
|
def history_records(history: list[JsonObject], screen: Screen) -> list[JsonObject]:
|
|
"""保留下来的各轮的记录，最近的在前；ago 是往前数第几轮，1 是上一轮
|
|
|
|
on_screen 按当前大屏现算：只表示这个 id 现在还在大屏上，不代表那次修改已经生效
|
|
澄清恢复时可能换了大屏，所以在节点里调用，不在组装请求时算好
|
|
"""
|
|
present = {node.id for node in screen.nodes}
|
|
records = []
|
|
for ago, turn in enumerate(reversed(history), start=1):
|
|
if record := _record(turn, ago, present):
|
|
records.append(record)
|
|
return records
|