115 lines
3.7 KiB
Python
115 lines
3.7 KiB
Python
"""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"]
|