131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
"""世谱匹配 / 预算 / 注入分组单元测试。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from app.modules.lorebook import (
|
||
LorebookBook,
|
||
LorebookEntry,
|
||
MatchedEntry,
|
||
build_injections,
|
||
match_entries,
|
||
select_within_budget,
|
||
)
|
||
|
||
|
||
def _entry(eid: str, **kw) -> LorebookEntry:
|
||
return LorebookEntry(id=eid, **kw)
|
||
|
||
|
||
def _book(entries: list[LorebookEntry], **kw) -> LorebookBook:
|
||
return LorebookBook(id="lorebook_x", entries=entries, **kw)
|
||
|
||
|
||
def test_keyword_match() -> None:
|
||
book = _book([_entry("1", keys=["猫"], content="她喜欢猫。")])
|
||
assert [m.entry.id for m in match_entries(book, "今天看到一只猫")] == ["1"]
|
||
|
||
|
||
def test_no_match_when_keyword_absent() -> None:
|
||
book = _book([_entry("1", keys=["猫"], content="她喜欢猫。")])
|
||
assert match_entries(book, "今天天气不错") == []
|
||
|
||
|
||
def test_case_insensitive_by_default() -> None:
|
||
book = _book([_entry("1", keys=["cat"], content="cat info")])
|
||
assert match_entries(book, "I have a CAT") != []
|
||
|
||
|
||
def test_case_sensitive() -> None:
|
||
book = _book([_entry("1", keys=["Cat"], content="x", case_sensitive=True)])
|
||
assert match_entries(book, "a cat here") == []
|
||
assert [m.entry.id for m in match_entries(book, "a Cat here")] == ["1"]
|
||
|
||
|
||
def test_constant_always_matched() -> None:
|
||
book = _book([_entry("1", keys=["不会命中"], content="恒定设定", constant=True)])
|
||
matched = match_entries(book, "无关文本")
|
||
assert [m.entry.id for m in matched] == ["1"]
|
||
assert matched[0].reason == "constant"
|
||
|
||
|
||
def test_disabled_entry_skipped() -> None:
|
||
book = _book([_entry("1", keys=["猫"], content="x", enabled=False)])
|
||
assert match_entries(book, "一只猫") == []
|
||
|
||
|
||
def test_selective_requires_both_keys() -> None:
|
||
book = _book(
|
||
[
|
||
_entry(
|
||
"1",
|
||
keys=["猫"],
|
||
secondary_keys=["黑"],
|
||
content="黑猫设定",
|
||
selective=True,
|
||
)
|
||
]
|
||
)
|
||
assert match_entries(book, "看到一只猫") == []
|
||
assert [m.entry.id for m in match_entries(book, "看到一只黑猫")] == ["1"]
|
||
|
||
|
||
def test_regex_match() -> None:
|
||
book = _book([_entry("1", keys=[r"\b\d{4}\b"], content="年份", use_regex=True)])
|
||
assert [m.entry.id for m in match_entries(book, "生于 1999 年")] == ["1"]
|
||
assert match_entries(book, "生于 99 年") == []
|
||
|
||
|
||
def test_recursive_scanning() -> None:
|
||
# 条目 1 命中关键词，其 content 含触发条目 2 的关键词
|
||
book = _book(
|
||
[
|
||
_entry("1", keys=["学校"], content="她在学校认识小明。"),
|
||
_entry("2", keys=["小明"], content="小明是同班同学。"),
|
||
],
|
||
recursive_scanning=True,
|
||
)
|
||
matched = match_entries(book, "今天去了学校")
|
||
ids = {m.entry.id for m in matched}
|
||
assert ids == {"1", "2"}
|
||
assert any(m.reason == "recursive" for m in matched if m.entry.id == "2")
|
||
|
||
|
||
def test_budget_keeps_lower_priority() -> None:
|
||
book = _book(
|
||
[
|
||
_entry("a", keys=["k"], content="aaaa", priority=10),
|
||
_entry("b", keys=["k"], content="bbbb", priority=1),
|
||
],
|
||
token_budget=5,
|
||
)
|
||
matched = match_entries(book, "k")
|
||
kept = select_within_budget(matched, token_budget=5)
|
||
# 仅 b（cost=4）能进预算；a（cost=4，但 priority 低）在 b 之后超预算
|
||
assert [m.entry.id for m in kept] == ["b"]
|
||
|
||
|
||
def test_budget_none_keeps_all_sorted_by_order() -> None:
|
||
book = _book(
|
||
[
|
||
_entry("a", keys=["k"], content="x", priority=5, insertion_order=20),
|
||
_entry("b", keys=["k"], content="y", priority=1, insertion_order=10),
|
||
]
|
||
)
|
||
kept = select_within_budget(match_entries(book, "k"))
|
||
assert [m.entry.id for m in kept] == ["b", "a"]
|
||
|
||
|
||
def test_build_injections_groups_by_position() -> None:
|
||
matched = [
|
||
MatchedEntry(entry=_entry("a", keys=["k"], content="前文", position="before_char")),
|
||
MatchedEntry(entry=_entry("b", keys=["k"], content="后文1", position="after_char")),
|
||
MatchedEntry(entry=_entry("c", keys=["k"], content="后文2", position="after_char")),
|
||
]
|
||
inj = build_injections(matched)
|
||
assert inj == {"before_char": "前文", "after_char": "后文1\n后文2"}
|
||
|
||
|
||
def test_unknown_position_falls_back_to_after_char() -> None:
|
||
matched = [MatchedEntry(entry=_entry("a", keys=["k"], content="x", position="weird"))]
|
||
assert build_injections(matched) == {"after_char": "x"}
|