268 lines
12 KiB
Python
268 lines
12 KiB
Python
import json
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
from sqlmodel import SQLModel, Session, create_engine
|
|
|
|
from app import training_corpus as corpus
|
|
from app.models import KnowledgeBase, KnowledgeFile
|
|
from app.rag import store
|
|
|
|
|
|
@pytest.fixture
|
|
def scope(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(corpus, "CORPUS_DIR", tmp_path / "corpora")
|
|
monkeypatch.setattr(store, "FALLBACK_DIR", tmp_path / "fallback")
|
|
store.FALLBACK_DIR.mkdir()
|
|
engine = create_engine("sqlite://")
|
|
SQLModel.metadata.create_all(engine)
|
|
with Session(engine) as session:
|
|
yield session, tmp_path
|
|
engine.dispose()
|
|
|
|
|
|
def _kb(session, name="资料库", **extra):
|
|
row = KnowledgeBase(name=name, index_status="ready", **extra)
|
|
session.add(row)
|
|
session.commit()
|
|
return row
|
|
|
|
|
|
def _draft(kbs, identifier=1):
|
|
return SimpleNamespace(id=identifier, version=1, kb_ids=[kb.id for kb in kbs], topic="分数教学")
|
|
|
|
|
|
def _chunks(kb, count, text_size=30):
|
|
return [{"id": f"chunk-{index}", "text": f"第{index}条定义和步骤" + "材料" * text_size,
|
|
"meta": {"kb_id": kb.id, "file": "source.txt"}} for index in range(count)]
|
|
|
|
|
|
def test_small_corpus_includes_sources_after_top_40_without_extra_model_calls(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
chunks = _chunks(kb, 55, 0)
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(chunks))
|
|
def no_model(*args):
|
|
raise AssertionError("小库无需额外摘要调用")
|
|
evidence, coverage = corpus.build_corpus(session, _draft([kb]), no_model, None)
|
|
assert len(evidence) == 55 and evidence[-1]["chunk_id"] == "chunk-54"
|
|
assert coverage["scanned_chunks"] == 55
|
|
assert coverage["read_chunks"] == 0 and coverage["queued_raw_chunks"] == 55
|
|
assert coverage["phase"] == coverage["stage"] == "planning"
|
|
assert coverage["complete"] is False
|
|
page = corpus.source_page(_draft([kb]), evidence[-1]["id"])
|
|
assert page["total"] == 1 and page["items"][0]["text"] == chunks[-1]["text"]
|
|
|
|
|
|
def test_full_batches_read_all_libraries_and_preserve_provenance(scope, monkeypatch):
|
|
session, _ = scope
|
|
first, second = _kb(session, "大资料库"), _kb(session, "小资料库")
|
|
data = {first.id: _chunks(first, 70, 300), second.id: _chunks(second, 2, 10)}
|
|
visited, fed, snapshots = [], [], []
|
|
def iterate(kid):
|
|
visited.append(kid)
|
|
return iter(data[kid])
|
|
monkeypatch.setattr(corpus, "iter_chunks", iterate)
|
|
def summarize(evidence, max_chars):
|
|
assert len(json.dumps(evidence, ensure_ascii=False)) <= 16000
|
|
fed.extend(item["chunk_id"] for item in evidence if item["kind"] == "raw")
|
|
return "资料覆盖：" + "、".join(item["chunk_id"] for item in evidence)[:max_chars - 6]
|
|
row = _draft([first, second])
|
|
evidence, coverage = corpus.build_corpus(session, row, summarize, snapshots.append)
|
|
assert visited == [first.id, second.id]
|
|
assert set(item["kb_id"] for item in evidence) == set(row.kb_ids)
|
|
assert coverage["read_chunks"] == 72 and coverage["scanned_chunks"] == 72
|
|
assert len(fed) == 72 and "chunk-69" in fed
|
|
assert len(json.dumps(evidence, ensure_ascii=False)) <= 40000
|
|
seen = set()
|
|
for item in evidence:
|
|
page = corpus.source_page(row, item["id"], 0, 100)
|
|
assert page["total"] == item["source_count"]
|
|
seen.update((source["kb_id"], source["chunk_id"]) for source in page["items"])
|
|
assert len(seen) == 72
|
|
assert any(snapshot["batches_done"] for snapshot in snapshots)
|
|
assert coverage["kbs"] == coverage["per_kb"]
|
|
|
|
|
|
def test_recursive_reduction_is_bounded_and_long_chunk_tail_is_read(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
monkeypatch.setattr(corpus, "BATCH_CHARS", 900)
|
|
monkeypatch.setattr(corpus, "FINAL_CHARS", 1800)
|
|
original = ("引号\"和换行\n" * 350) + "不可丢失的最后一句：禁止省略前提。"
|
|
data = _chunks(kb, 16, 80) + [{"id": "long", "text": original, "meta": {"file": "long.txt"}}]
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(data))
|
|
long_parts, levels = [], []
|
|
def summarize(evidence, max_chars):
|
|
assert len(json.dumps(evidence, ensure_ascii=False)) <= 900
|
|
for item in evidence:
|
|
if item["chunk_id"] == "long":
|
|
long_parts.append(item["text"])
|
|
levels.append(any(item["kind"] == "summary" for item in evidence))
|
|
return "保留资料中概念定义、练习步骤、适用前提与内容限制。" * 5
|
|
row = _draft([kb])
|
|
evidence, coverage = corpus.build_corpus(session, row, summarize, None)
|
|
assert "".join(long_parts) == original
|
|
assert any(levels)
|
|
assert coverage["read_chunks"] == 17
|
|
assert len(json.dumps(evidence, ensure_ascii=False)) <= 1800
|
|
sources = [source for item in evidence for source in corpus.source_page(row, item["id"], 0, 100)["items"]]
|
|
assert next(source for source in sources if source["chunk_id"] == "long")["text"] == original
|
|
|
|
|
|
def test_unindexed_disk_and_registered_files_are_reported(scope, monkeypatch):
|
|
session, directory = scope
|
|
source = directory / "source"
|
|
source.mkdir()
|
|
(source / "source.txt").write_text("已索引", encoding="utf-8")
|
|
(source / "new.txt").write_text("未索引新资料", encoding="utf-8")
|
|
kb = _kb(session, source_dir=str(source))
|
|
session.add(KnowledgeFile(kb_id=kb.id, file_name="registered-only.txt"))
|
|
session.commit()
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(_chunks(kb, 1)))
|
|
_, coverage = corpus.build_corpus(session, _draft([kb]), None, None)
|
|
assert coverage["total_files"] == 3 and coverage["indexed_files"] == 1
|
|
assert coverage["unindexed_files"] == 2 and not coverage["source_complete"]
|
|
assert coverage["per_kb"][0]["unindexed_files"] == ["new.txt", "registered-only.txt"]
|
|
assert coverage["warnings"]
|
|
|
|
|
|
def test_indexing_and_read_failures_never_claim_completion(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
kb.index_status = "indexing"
|
|
session.add(kb)
|
|
session.commit()
|
|
snapshots = []
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: pytest.fail("不能读取正在重建的索引"))
|
|
with pytest.raises(ValueError, match="正在建立索引"):
|
|
corpus.build_corpus(session, _draft([kb]), None, snapshots.append)
|
|
assert snapshots[-1]["phase"] == "failed" and not snapshots[-1]["complete"]
|
|
kb.index_status = "ready"
|
|
session.add(kb)
|
|
session.commit()
|
|
def broken(kid):
|
|
yield from _chunks(kb, 1)
|
|
raise RuntimeError("索引后半部损坏")
|
|
monkeypatch.setattr(corpus, "iter_chunks", broken)
|
|
with pytest.raises(RuntimeError, match="后半部"):
|
|
corpus.build_corpus(session, _draft([kb]), None, snapshots.append)
|
|
assert snapshots[-1]["scanned_chunks"] == 1
|
|
assert snapshots[-1]["read_chunks"] == 0 and snapshots[-1]["phase"] == "failed"
|
|
|
|
|
|
def test_summary_failure_and_overbudget_are_not_silently_truncated(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(_chunks(kb, 70, 150)))
|
|
snapshots = []
|
|
with pytest.raises(ValueError, match="摘要超出"):
|
|
corpus.build_corpus(session, _draft([kb]), lambda evidence, max_chars: "字" * (max_chars + 1), snapshots.append)
|
|
assert snapshots[-1]["read_chunks"] == 0
|
|
assert snapshots[-1]["phase"] == "failed"
|
|
assert not corpus._path(1, 1).exists()
|
|
|
|
|
|
def test_duplicate_chunk_identity_counted_but_conflicting_duplicate_fails(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
data = _chunks(kb, 1)
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(data + data))
|
|
evidence, coverage = corpus.build_corpus(session, _draft([kb]), None, None)
|
|
assert len(evidence) == 1 and coverage["duplicate_chunks"] == 1
|
|
assert coverage["scanned_chunks"] == 2 and coverage["queued_raw_chunks"] == 1
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(data + [{**data[0], "text": "冲突正文"}]))
|
|
with pytest.raises(ValueError, match="正文冲突"):
|
|
corpus.build_corpus(session, _draft([kb]), None, None)
|
|
|
|
|
|
def test_fallback_stream_does_not_materialize_embeddings_and_rejects_broken_tail(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
data = _chunks(kb, 60)
|
|
for item in data:
|
|
item["embedding"] = [0.314159] * 300
|
|
store._write_fallback(kb.id, data)
|
|
monkeypatch.setattr(store, "_read_fallback", lambda kid: pytest.fail("不得整库加载 fallback"))
|
|
original_loads = store.json.loads
|
|
def guarded_loads(value, *args, **kwargs):
|
|
assert "0.314159" not in value
|
|
return original_loads(value, *args, **kwargs)
|
|
monkeypatch.setattr(store.json, "loads", guarded_loads)
|
|
streamed = list(store.iter_chunks(kb.id))
|
|
assert len(streamed) == 60 and streamed[-1]["id"] == "chunk-59"
|
|
assert all("embedding" not in item for item in streamed)
|
|
with store._fallback_path(kb.id).open("a") as handle:
|
|
handle.write("malformed-tail")
|
|
with pytest.raises(RuntimeError, match="多余数据"):
|
|
list(store.iter_chunks(kb.id))
|
|
|
|
|
|
def test_chroma_full_pagination_requests_only_documents_and_metadata(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
data, offsets = _chunks(kb, 61), []
|
|
class Collection:
|
|
def count(self):
|
|
return len(data)
|
|
def get(self, *, include, limit, offset):
|
|
assert include == ["documents", "metadatas"]
|
|
offsets.append(offset)
|
|
page = data[offset:offset + limit]
|
|
return {"ids": [item["id"] for item in page], "documents": [item["text"] for item in page],
|
|
"metadatas": [item["meta"] for item in page]}
|
|
monkeypatch.setattr(store, "get_client", lambda: SimpleNamespace(get_collection=lambda name: Collection()))
|
|
assert len(list(store.iter_chunks(kb.id, page_size=17))) == 61
|
|
assert offsets == [0, 17, 34, 51]
|
|
|
|
|
|
def test_sources_pagination_path_checks_day_expansion_and_delete(scope, monkeypatch):
|
|
session, _ = scope
|
|
kb = _kb(session)
|
|
data = _chunks(kb, 80, 70)
|
|
monkeypatch.setattr(corpus, "iter_chunks", lambda kid: iter(data))
|
|
row = _draft([kb])
|
|
evidence, _ = corpus.build_corpus(session, row, lambda evidence, max_chars: "资料定义和必要前提", None)
|
|
first = evidence[0]
|
|
page = corpus.source_page(row, first["id"], 1, 2)
|
|
assert len(page["items"]) == 2 and page["offset"] == 1 and page["limit"] == 2
|
|
assert page["items"][0]["chunk_id"] == "chunk-1"
|
|
expanded = corpus.day_evidence(row, evidence)
|
|
assert expanded[:len(evidence)] == evidence
|
|
assert any(item["kind"] == "raw" for item in expanded)
|
|
assert len(json.dumps(expanded, ensure_ascii=False)) <= 40000
|
|
for bad in ["../v1.sqlite", "E0", "other"]:
|
|
with pytest.raises(ValueError):
|
|
corpus.source_page(row, bad)
|
|
with pytest.raises(ValueError):
|
|
corpus.source_page(row, first["id"], -1)
|
|
with pytest.raises(ValueError):
|
|
corpus.remove_corpus("../1")
|
|
corpus.remove_corpus(row.id)
|
|
assert not corpus._path(row.id, row.version).exists()
|
|
|
|
|
|
def test_summary_cache_reuses_exact_batches_across_versions_but_not_topic_or_model(scope):
|
|
_, _ = scope
|
|
row = SimpleNamespace(id=1, version=1, owner_user_id=7, organization_id=1,
|
|
topic='分数教学', requested_days=None)
|
|
entries = [{'id': 'E1', 'text': '第一批原文'}]
|
|
calls = []
|
|
def generate():
|
|
calls.append(True)
|
|
return '原文依据与教学步骤'
|
|
assert corpus.cached_summary(row, entries, 100, 'model-a', generate)[1] is False
|
|
row.version = 2
|
|
assert corpus.cached_summary(row, entries, 100, 'model-a', generate)[1] is True
|
|
assert len(calls) == 1
|
|
assert corpus.cached_summary(row, entries, 100, 'model-b', generate)[1] is False
|
|
row.topic = '小数教学'
|
|
assert corpus.cached_summary(row, entries, 100, 'model-a', generate)[1] is False
|
|
entries[0]['text'] = '原文已更新'
|
|
assert corpus.cached_summary(row, entries, 100, 'model-a', generate)[1] is False
|
|
assert len(calls) == 4
|
|
with pytest.raises(ValueError):
|
|
corpus.cached_summary(row, entries, 1, 'model-a', generate)
|
|
corpus.remove_corpus(1)
|
|
assert not (corpus.CORPUS_DIR / '1').exists()
|