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()