This commit is contained in:
1 parent
a7a57f9a03
commit
491f31cead
7 files changed
+4813
-4601
No files matched your search
+1385
-1385
File diff suppressed because it is too large.
Load diff
@@ -20,7 +20,7 @@ from ..models import (
|
||||
|
||||
|
||||
|
||||
User, Conversation, Message, AIModel, ModelProvider, TokenUsageLog,
|
||||
User, Conversation, Message, AIModel, ModelProvider, TokenUsageLog, ConversationSummary,
|
||||
|
||||
|
||||
|
||||
@@ -1230,3 +1230,163 @@ def _get_conversation_context(db: Session, conv_id: str) -> dict:
|
||||
|
||||
|
||||
return context
|
||||
|
||||
|
||||
|
||||
# ===== Phase 2 上下文引擎升级:分层记忆 + 压缩 + 隔离 =====
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _estimate_tokens(text):
|
||||
"""粗略 token 估算:中英文混排约 2 字符/token。"""
|
||||
if not text:
|
||||
return 0
|
||||
return max(1, len(text) // 2)
|
||||
|
||||
|
||||
def _build_chat_history(db, conv_id, max_tokens=6000, max_messages=40,
|
||||
exclude_latest_user=True, include_summary=True):
|
||||
"""会话层上下文引擎:预算感知的历史组装 + 可选压缩摘要前置。
|
||||
|
||||
- 若 conversation_summaries 存在摘要,以 system 角色「对话摘要」前置(压缩层)。
|
||||
- 按 token 估算从新到旧累加,超预算即截断(隔离:不溢出上下文窗口)。
|
||||
- exclude_latest_user=True 时丢弃最新一条 user 消息(当前轮次另行发送)。
|
||||
"""
|
||||
summary_block = None
|
||||
if include_summary:
|
||||
try:
|
||||
summ = db.query(ConversationSummary).filter(
|
||||
ConversationSummary.conversation_id == conv_id
|
||||
).first()
|
||||
if summ and summ.summary_text:
|
||||
summary_block = {
|
||||
"role": "system",
|
||||
"content": "【对话历史摘要,请据此理解用户此前的意图与已确认结论】\n" + summ.summary_text,
|
||||
}
|
||||
except Exception:
|
||||
pass
|
||||
rows = db.query(Message).filter(
|
||||
Message.conversation_id == conv_id,
|
||||
Message.role.in_(["user", "assistant"]),
|
||||
).order_by(desc(Message.created_at)).limit(max_messages).all()
|
||||
if exclude_latest_user:
|
||||
for i, m in enumerate(rows):
|
||||
if m.role == "user":
|
||||
rows = rows[:i] + rows[i + 1:]
|
||||
break
|
||||
kept = []
|
||||
used = 0
|
||||
for m in rows: # 新->旧
|
||||
t = _estimate_tokens(m.content or "")
|
||||
if kept and used + t > max_tokens:
|
||||
break
|
||||
kept.append({"role": m.role, "content": m.content})
|
||||
used += t
|
||||
kept.reverse() # 时间正序
|
||||
out = []
|
||||
if summary_block:
|
||||
out.append(summary_block)
|
||||
out.extend(kept)
|
||||
return out
|
||||
|
||||
|
||||
async def _summarize_via_llm(prompt):
|
||||
"""调用默认文本模型生成对话摘要;失败返回 None(压缩为可选项,不阻塞)。"""
|
||||
try:
|
||||
from ..capability_registry import resolve_auto_model
|
||||
from app.api.chat_utils import _proxy_for_url
|
||||
sdb = SessionLocal()
|
||||
try:
|
||||
ai_model = resolve_auto_model("text", sdb)
|
||||
if not ai_model:
|
||||
return None
|
||||
provider = sdb.query(ModelProvider).filter(
|
||||
ModelProvider.id == ai_model.provider_id
|
||||
).first()
|
||||
api_key = ai_model.api_key or (provider.api_key if provider else None)
|
||||
if not provider or not provider.api_base_url or not api_key:
|
||||
return None
|
||||
api_url = f"{provider.api_base_url.rstrip('/')}/v1/chat/completions"
|
||||
body = {
|
||||
"model": ai_model.name,
|
||||
"messages": [
|
||||
{"role": "system", "content": "你是专业的对话摘要助手,输出简洁中文摘要。"},
|
||||
{"role": "user", "content": prompt},
|
||||
],
|
||||
"max_tokens": 512,
|
||||
"temperature": 0.2,
|
||||
}
|
||||
async with httpx.AsyncClient(proxy=_proxy_for_url(api_url), trust_env=False, timeout=60) as client:
|
||||
resp = await client.post(
|
||||
api_url,
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json=body,
|
||||
)
|
||||
if resp.status_code != 200:
|
||||
return None
|
||||
data = resp.json()
|
||||
return data.get("choices", [{}])[0].get("message", {}).get("content", "").strip()
|
||||
finally:
|
||||
sdb.close()
|
||||
except Exception as e:
|
||||
logger.warning(f"[context] summarize llm failed: {e}")
|
||||
return None
|
||||
|
||||
|
||||
async def _maybe_compress_history(conv_id):
|
||||
"""后台异步压缩:将受保护窗口外的旧消息摘要落库 conversation_summaries。
|
||||
|
||||
失败静默回滚,绝不阻塞用户本次响应(fire-and-forget)。
|
||||
"""
|
||||
sdb = None
|
||||
try:
|
||||
sdb = SessionLocal()
|
||||
all_msgs = sdb.query(Message).filter(
|
||||
Message.conversation_id == conv_id,
|
||||
Message.role.in_(["user", "assistant"]),
|
||||
).order_by(asc(Message.created_at)).all()
|
||||
PROTECTED = 20
|
||||
if len(all_msgs) <= PROTECTED:
|
||||
return
|
||||
compressible = all_msgs[: len(all_msgs) - PROTECTED]
|
||||
if len(compressible) < 8:
|
||||
return
|
||||
chunk = compressible[:30]
|
||||
text = "\n".join(f"{m.role}: {m.content}" for m in chunk)
|
||||
existing = sdb.query(ConversationSummary).filter(
|
||||
ConversationSummary.conversation_id == conv_id
|
||||
).first()
|
||||
prior = existing.summary_text if existing else ""
|
||||
prompt = (
|
||||
"请把以下对话片段浓缩为简明摘要,保留用户核心诉求、已确认结论、关键事实与待办。"
|
||||
"若提供『已有摘要』,请合并而非重复。\n\n"
|
||||
f"已有摘要:\n{prior}\n\n待摘要片段:\n{text}"
|
||||
)
|
||||
summary = await _summarize_via_llm(prompt)
|
||||
if not summary:
|
||||
return
|
||||
if existing:
|
||||
existing.summary_text = summary
|
||||
existing.token_count = _estimate_tokens(summary)
|
||||
existing.updated_at = datetime.now()
|
||||
else:
|
||||
sdb.add(ConversationSummary(
|
||||
conversation_id=conv_id,
|
||||
summary_text=summary,
|
||||
token_count=_estimate_tokens(summary),
|
||||
updated_at=datetime.now(),
|
||||
))
|
||||
sdb.commit()
|
||||
except Exception as e:
|
||||
logger.warning(f"[context] compress history failed conv={conv_id}: {e}")
|
||||
if sdb is not None:
|
||||
try:
|
||||
sdb.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
if sdb is not None:
|
||||
try:
|
||||
sdb.close()
|
||||
except Exception:
|
||||
pass
|
||||
+2009
-2045
File diff suppressed because it is too large.
Load diff
+1136
-1127
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,114 @@
|
||||
"""Phase 2 context-engine regression: _build_chat_history budget/layering/summary.
|
||||
|
||||
Pure unit test with a fake DB session — no network, no real model. Guards the
|
||||
layered-memory + compression + isolation behaviour against regressions.
|
||||
"""
|
||||
from app.api.chat_context import _estimate_tokens, _build_chat_history
|
||||
|
||||
|
||||
class _Msg:
|
||||
def __init__(self, role, content, cid="c1"):
|
||||
self.role = role
|
||||
self.content = content
|
||||
self.conversation_id = cid
|
||||
|
||||
|
||||
class _Summary:
|
||||
def __init__(self, text):
|
||||
self.summary_text = text
|
||||
|
||||
|
||||
class _Query:
|
||||
def __init__(self, rows):
|
||||
self._rows = list(rows)
|
||||
|
||||
def filter(self, *a, **k):
|
||||
return self
|
||||
|
||||
def order_by(self, *a, **k):
|
||||
# test pre-sorts newest-first to match desc(created_at)
|
||||
return self
|
||||
|
||||
def limit(self, n=None):
|
||||
return self
|
||||
|
||||
def all(self):
|
||||
return self._rows
|
||||
|
||||
def first(self):
|
||||
return self._rows[0] if self._rows else None
|
||||
|
||||
|
||||
class _DB:
|
||||
def __init__(self, messages, summary=None):
|
||||
self._messages = messages
|
||||
self._summary = summary
|
||||
|
||||
def query(self, model):
|
||||
if getattr(model, "__name__", "") == "ConversationSummary":
|
||||
return _Query([self._summary] if self._summary else [])
|
||||
return _Query(self._messages)
|
||||
|
||||
|
||||
def _newest_first(msgs):
|
||||
# _build_chat_history expects desc(created_at) -> newest first.
|
||||
# Test authors write the list already newest-first, so return as-is.
|
||||
return list(msgs)
|
||||
|
||||
|
||||
def test_estimate_tokens():
|
||||
assert _estimate_tokens("") == 0
|
||||
assert _estimate_tokens(None) == 0
|
||||
assert _estimate_tokens("abcd") == 2 # 4//2
|
||||
assert _estimate_tokens("a" * 5) == 2 # max(1, 5//2)=2
|
||||
# just ensure positive monotonic
|
||||
assert _estimate_tokens("x" * 100) == 50
|
||||
|
||||
|
||||
def test_history_excludes_latest_user():
|
||||
# newest-first: [current user, prev assistant, older user]
|
||||
msgs = _newest_first([
|
||||
_Msg("user", "current question"),
|
||||
_Msg("assistant", "previous answer"),
|
||||
_Msg("user", "older question"),
|
||||
])
|
||||
db = _DB(msgs)
|
||||
out = _build_chat_history(db, "c1", max_tokens=10000, exclude_latest_user=True)
|
||||
roles = [m["role"] for m in out]
|
||||
# current user dropped; chronological order = [older user, previous assistant]
|
||||
assert roles == ["user", "assistant"]
|
||||
assert out[0]["content"] == "older question"
|
||||
|
||||
|
||||
def test_summary_injected_as_system_first():
|
||||
msgs = _newest_first([
|
||||
_Msg("user", "hi"),
|
||||
_Msg("assistant", "hello"),
|
||||
])
|
||||
db = _DB(msgs, summary=_Summary("用户想做图片生成"))
|
||||
out = _build_chat_history(db, "c1", max_tokens=10000, exclude_latest_user=True)
|
||||
assert out[0]["role"] == "system"
|
||||
assert out[0]["content"].startswith("【对话历史摘要")
|
||||
assert "用户想做图片生成" in out[0]["content"]
|
||||
# after summary: chronological history (assistant then user 'hi' excluded? exclude_latest_user drops 'hi')
|
||||
assert out[1]["role"] == "assistant"
|
||||
|
||||
|
||||
def test_budget_truncation():
|
||||
# 5 assistant messages each ~200 tokens; budget small -> truncated
|
||||
msgs = _newest_first([_Msg("assistant", "x" * 400) for _ in range(5)])
|
||||
db = _DB(msgs)
|
||||
out = _build_chat_history(db, "c1", max_tokens=500, exclude_latest_user=False)
|
||||
# each msg ~200 tokens; budget 500 -> keep newest ~2
|
||||
assert len(out) <= 3
|
||||
total = sum(_estimate_tokens(m["content"]) for m in out)
|
||||
assert total <= 500 + 200 # within budget (+1 msg slack)
|
||||
|
||||
|
||||
def test_no_summary_when_absent():
|
||||
msgs = _newest_first([_Msg("user", "hi"), _Msg("assistant", "hello")])
|
||||
db = _DB(msgs)
|
||||
out = _build_chat_history(db, "c1", max_tokens=10000, exclude_latest_user=True)
|
||||
# 'hi' excluded; only assistant remains, no system summary
|
||||
assert all(m["role"] != "system" for m in out)
|
||||
assert [m["role"] for m in out] == ["assistant"]
|
||||
Reference in new issue
Block a user