111 lines
3.7 KiB
Python
111 lines
3.7 KiB
Python
"""大周天记忆/世界抽取纯函数测试。"""
|
|
|
|
from app.modules.extraction.extract import (
|
|
MemoryCandidateDraft,
|
|
WorldEntityDraft,
|
|
parse_extraction,
|
|
)
|
|
|
|
|
|
def test_parse_standard_object() -> None:
|
|
raw = (
|
|
'{"memories":[{"type":"preference","title":"喜欢简洁","content":"用户偏好简洁",'
|
|
'"importance":0.8,"sensitivity":"low","confidence":0.7,"reason":"多次提及"}],'
|
|
'"entities":[{"type":"project","name":"项目X","attributes":{"status":"active"},'
|
|
'"state_type":"inference","confidence":0.6}]}'
|
|
)
|
|
memories, entities = parse_extraction(raw)
|
|
assert len(memories) == 1
|
|
assert memories[0].title == "喜欢简洁"
|
|
assert memories[0].type == "preference"
|
|
assert memories[0].importance == 0.8
|
|
assert len(entities) == 1
|
|
assert entities[0].name == "项目X"
|
|
assert entities[0].type == "project"
|
|
assert entities[0].attributes == {"status": "active"}
|
|
assert entities[0].state_type == "inference"
|
|
|
|
|
|
def test_parse_markdown_block() -> None:
|
|
raw = '```json\n{"memories":[],"entities":[]}\n```'
|
|
memories, entities = parse_extraction(raw)
|
|
assert memories == []
|
|
assert entities == []
|
|
|
|
|
|
def test_parse_text_mixed() -> None:
|
|
raw = '抽取结果：\n{"memories":[{"title":"t","content":"c"}],"entities":[]}\n以上'
|
|
memories, entities = parse_extraction(raw)
|
|
assert len(memories) == 1
|
|
assert entities == []
|
|
|
|
|
|
def test_parse_invalid_returns_empty() -> None:
|
|
assert parse_extraction("不是JSON") == ([], [])
|
|
assert parse_extraction("[]") == ([], []) # 数组不是对象
|
|
|
|
|
|
def test_missing_keys_returns_empty_lists() -> None:
|
|
memories, entities = parse_extraction('{"memories":[]}')
|
|
assert memories == []
|
|
assert entities == []
|
|
|
|
|
|
def test_skip_memory_missing_title_or_content() -> None:
|
|
raw = '{"memories":[{"title":"t"},{"content":"c"},{"title":"t","content":"c"}]}'
|
|
memories, _ = parse_extraction(raw)
|
|
assert len(memories) == 1
|
|
|
|
|
|
def test_skip_entity_missing_name() -> None:
|
|
raw = '{"entities":[{"type":"project"},{"name":"有名字"}]}'
|
|
_, entities = parse_extraction(raw)
|
|
assert len(entities) == 1
|
|
assert entities[0].name == "有名字"
|
|
|
|
|
|
def test_memory_safe_defaults() -> None:
|
|
raw = '{"memories":[{"title":"t","content":"c"}]}'
|
|
m = parse_extraction(raw)[0][0]
|
|
assert m.sensitivity == "low"
|
|
assert m.type == "other"
|
|
assert m.confidence == 0.5
|
|
assert m.importance == 0.5
|
|
|
|
|
|
def test_entity_safe_defaults() -> None:
|
|
raw = '{"entities":[{"name":"n"}]}'
|
|
e = parse_extraction(raw)[1][0]
|
|
assert e.state_type == "inference" # 非 fact，降权
|
|
assert e.type == "concept"
|
|
assert e.confidence == 0.5
|
|
assert e.attributes == {}
|
|
|
|
|
|
def test_clamp_values() -> None:
|
|
raw = (
|
|
'{"memories":[{"title":"t","content":"c","importance":1.5,"confidence":-0.2}],'
|
|
'"entities":[{"name":"n","confidence":2.0}]}'
|
|
)
|
|
memories, entities = parse_extraction(raw)
|
|
assert memories[0].importance == 1.0
|
|
assert memories[0].confidence == 0.0
|
|
assert entities[0].confidence == 1.0
|
|
|
|
|
|
def test_invalid_sensitivity_defaults() -> None:
|
|
raw = '{"memories":[{"title":"t","content":"c","sensitivity":"unknown"}]}'
|
|
assert parse_extraction(raw)[0][0].sensitivity == "low"
|
|
|
|
|
|
def test_invalid_state_type_defaults() -> None:
|
|
raw = '{"entities":[{"name":"n","state_type":"certain"}]}'
|
|
assert parse_extraction(raw)[1][0].state_type == "inference"
|
|
|
|
|
|
def test_draft_types() -> None:
|
|
raw = '{"memories":[{"title":"t","content":"c"}],"entities":[{"name":"n"}]}'
|
|
memories, entities = parse_extraction(raw)
|
|
assert isinstance(memories[0], MemoryCandidateDraft)
|
|
assert isinstance(entities[0], WorldEntityDraft)
|