+ 
+
+ English · + 简体中文 · + 繁體中文 · + 日本語 · + 한국어 · + Español · + Português · + Français · + Deutsch +
+ +
+ 
Source: Bahdanau and Vaswani.
+The model focuses on relevant input tokens.
+| Type | Description |
|---|---|
| Self-Attention | Tokens attend to tokens. |
The Transformer uses multi-head attention.
+
+
+
+""",
+ encoding="utf-8",
+ )
+ archived.write_text(source.read_text(encoding="utf-8"), encoding="utf-8")
+
+ markdown = parse_html_markdown(source, archived)
+
+ assert extract_title(markdown, "fallback") == "Web Page - Attention Mechanism"
+ assert "### Metadata" not in markdown
+ assert "Raw HTML" not in markdown
+ assert "Navigation noise" not in markdown
+ assert "### Attention Mechanism" in markdown
+ assert "#### Overview" in markdown
+ assert "| Type | Description |" in markdown
+ assert "" in markdown
+ assert "Image note: Transformer architecture." in markdown
+ assert (tmp_path / "raw" / "assets" / "2026-05-15_attention" / "img-001.png").exists()
+ manifest = (tmp_path / "raw" / "assets" / "2026-05-15_attention" / "manifest.json").read_text(encoding="utf-8")
+ assert '"original_src": "arch.png"' in manifest
+ assert '"section": "Architecture"' in manifest
+
+
+def test_parse_html_markdown_keeps_remote_images_as_urls(tmp_path: Path) -> None:
+ source = tmp_path / "remote.htm"
+ archived = tmp_path / "raw" / "2026-05-15_remote.htm"
+ archived.parent.mkdir()
+ source.write_text(
+ 'Intro.
',
+ encoding="utf-8",
+ )
+ archived.write_text(source.read_text(encoding="utf-8"), encoding="utf-8")
+
+ markdown = parse_html_markdown(source, archived)
+
+ assert "
| This article has multiple issues. |
For other uses, see Knowledge graph (disambiguation).
+Knowledge graph is a graph-structured knowledge base used to represent entities, facts, and relationships for retrieval and reasoning systems.
+The term has been used in several database and semantic web contexts.
+
Small. Fast. Reliable.SQL Reference explains statements, expressions, and data manipulation concepts for database users.
+The reference lists supported statement syntax.
+Hello.
", encoding="utf-8") + archived.write_text(source.read_text(encoding="utf-8"), encoding="utf-8") + + document = parse_document(source, archived, _config()) + + assert document.title == "Web Page - HTML Page" + assert document.metadata["extension"] == ".html" + assert "### HTML Page" in document.markdown_content diff --git a/tests/test_image_parser.py b/tests/test_image_parser.py new file mode 100644 index 0000000..c013499 --- /dev/null +++ b/tests/test_image_parser.py @@ -0,0 +1,95 @@ +from pathlib import Path + +from heta.config.schema import HetaConfig, InsertPlanningConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.kb.image_parser import build_image_markdown +from heta.kb.parser import parse_document +from heta.kb.text import extract_title + + +def _config() -> HetaConfig: + return HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + ) + + +def _custom_without_multimodal_config() -> HetaConfig: + return HetaConfig( + version=1, + llm=LLMConfig( + provider="custom", + api_key="sk-test", + chat_api_key="sk-chat", + chat_model="chat-model", + chat_base_url="http://chat.local/v1", + embedding_api_key="sk-embedding", + embedding_model="embedding-model", + embedding_base_url="http://embedding.local/v1", + ), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + ) + + +def test_build_image_markdown_uses_compact_retrieval_sections() -> None: + markdown = build_image_markdown( + title="Image - Architecture Diagram", + source_name="diagram.png", + image_path="../../raw/diagram.png", + summary="A system architecture diagram.", + visual_facts="Scene/type: diagram. Main subject: service pipeline.", + visible_text="API Gateway", + interpretation_keywords="Represents a backend data flow. keywords: API, pipeline.", + ) + + assert extract_title(markdown, "fallback") == "Image - Architecture Diagram" + assert "" in markdown + assert "### Visual Facts" in markdown + assert "### Visible Text" in markdown + assert "### Interpretation and Keywords" in markdown + assert "## Related Pages" in markdown + assert "## Source" in markdown + + +def test_parse_document_accepts_image_branch(monkeypatch, tmp_path: Path) -> None: + source = tmp_path / "diagram.png" + archived = tmp_path / "raw_diagram.png" + source.write_bytes(b"png") + archived.write_bytes(b"png") + + monkeypatch.setattr( + "heta.kb.parser.parse_image_markdown", + lambda source_path, archived_path, config: build_image_markdown( + title="Image - Diagram", + source_name=archived_path.name, + image_path="../../raw/raw_diagram.png", + summary="A diagram.", + visual_facts="A simple diagram.", + visible_text="None detected.", + interpretation_keywords="diagram, test", + ), + ) + + document = parse_document(source, archived, _config()) + + assert document.title == "Image - Diagram" + assert document.source_name == "raw_diagram.png" + assert document.metadata["extension"] == ".png" + assert "### Visual Facts" in document.markdown_content + + +def test_image_requires_multimodal_when_custom_skips_it(tmp_path: Path) -> None: + source = tmp_path / "diagram.png" + source.write_bytes(b"png") + + try: + parse_document(source, source, _custom_without_multimodal_config()) + except ValueError as exc: + assert "requires a multimodal model" in str(exc) + assert "heta init" in str(exc) + else: + raise AssertionError("image parsing should require multimodal config") diff --git a/tests/test_kb_insert.py b/tests/test_kb_insert.py index bd9a8a0..7231bc6 100644 --- a/tests/test_kb_insert.py +++ b/tests/test_kb_insert.py @@ -2,25 +2,37 @@ import pytest -from heta.config.schema import HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.config.schema import ( + DynamicInsertConfig, + InsertPlanningConfig, + HetaConfig, + LLMConfig, + MinerUConfig, + VectorIndexConfig, +) from heta.kb.discovery import collect_insert_files -from heta.kb.models import FileChange -from heta.kb.insert import insert_paths +from heta.kb.models import FileChange, ParsedDocument +from heta.kb.insert import _ensure_code_raw_links, insert_paths from heta.kb.text import frontmatter_page, slugify, summarize -from heta.kb.wiki import normalize_wiki_pages +from heta.kb.wiki import normalize_wiki_pages, repair_broken_wiki_links -def _config(mineru: MinerUConfig | None = None) -> HetaConfig: +def _config(mineru: MinerUConfig | None = None, *, dynamic_insert: bool = True) -> HetaConfig: return HetaConfig( version=1, llm=LLMConfig(provider="qwen", api_key="sk-test"), mineru=mineru or MinerUConfig.disabled(), vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + dynamic_insert=DynamicInsertConfig(enable=dynamic_insert), ) -def _fake_agent(monkeypatch) -> None: +def _fake_agent(monkeypatch, calls: list[list[str]] | None = None) -> None: def run_merge_agent(*, task_id, documents, root_dir, config): + assert len(documents) == 1 + if calls is not None: + calls.append([document.source_name for document in documents]) pages = root_dir / "pages" pages.mkdir(parents=True, exist_ok=True) added = [] @@ -96,6 +108,156 @@ def test_insert_same_title_updates_existing_page(monkeypatch, tmp_path: Path) -> assert "## Imported Update" in page.read_text(encoding="utf-8") +def test_insert_multiple_files_runs_agent_sequentially(monkeypatch, tmp_path: Path) -> None: + calls: list[list[str]] = [] + progress = [] + _fake_agent(monkeypatch, calls) + first = tmp_path / "alpha.md" + second = tmp_path / "beta.md" + first.write_text("# Alpha\n\nFirst details.", encoding="utf-8") + second.write_text("# Beta\n\nSecond details.", encoding="utf-8") + + result = insert_paths( + [first, second], + _config(), + base_dir=tmp_path, + on_progress=progress.append, + ) + + wiki = tmp_path / "workspace" / "kb" / "wiki" + assert calls[0][0].endswith("_alpha.md") + assert calls[1][0].endswith("_beta.md") + assert (wiki / "pages" / "1-alpha.md").exists() + assert (wiki / "pages" / "2-beta.md").exists() + assert [change.path for change in result.added] == ["pages/1-alpha.md", "pages/2-beta.md"] + assert progress[0].percent == 1 + merge_percents = [event.percent for event in progress if event.phase == "merge"] + assert 50 in merge_percents + assert 99 in merge_percents + assert progress[-1].percent == 100 + assert progress[-1].phase == "done" + + +def test_insert_defaults_to_static_pages(monkeypatch, tmp_path: Path) -> None: + def fail_agent(**kwargs): + raise AssertionError("dynamic agent should not run in static insert mode") + + monkeypatch.setattr("heta.kb.insert.run_merge_agent", fail_agent) + monkeypatch.setattr( + "heta.kb.static_insert.generate_summary", + lambda *, document, config: f"Summary for {document.title}.", + ) + source = tmp_path / "manual.md" + source.write_text("# Main Heading\n\n## Sub Heading\n\nBody text.", encoding="utf-8") + + result = insert_paths([source], _config(dynamic_insert=False), base_dir=tmp_path) + + wiki = tmp_path / "workspace" / "kb" / "wiki" + page = wiki / "pages" / "1-main-heading.md" + text = page.read_text(encoding="utf-8") + assert result.added[0].path == "pages/1-main-heading.md" + assert "Summary for Main Heading." in text + assert "### Main Heading" in text + assert "#### Sub Heading" in text + assert "## Related Pages\n\n- None yet" in text + assert "- " + result.raw_files[0].name in text + assert "[[Main Heading]]" in (wiki / "index.md").read_text(encoding="utf-8") + assert "Created static page: Main Heading" in (wiki / "log.md").read_text(encoding="utf-8") + + +def test_insert_reports_vector_sync_error_without_rolling_back(monkeypatch, tmp_path: Path) -> None: + monkeypatch.setattr( + "heta.kb.static_insert.generate_summary", + lambda *, document, config: f"Summary for {document.title}.", + ) + monkeypatch.setattr( + "heta.kb.insert.sync_wiki_vector_index", + lambda **kwargs: (_ for _ in ()).throw(RuntimeError("embedding unavailable")), + ) + source = tmp_path / "manual.md" + source.write_text("# Main Heading\n\nBody text.", encoding="utf-8") + config = HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig(enable=True), + insert_planning=InsertPlanningConfig.enabled(), + dynamic_insert=DynamicInsertConfig.disabled(), + ) + + result = insert_paths([source], config, base_dir=tmp_path) + + page = tmp_path / "workspace" / "kb" / "wiki" / "pages" / "1-main-heading.md" + assert result.commit_id + assert page.exists() + assert result.vector_index_error == "embedding unavailable" + + +def test_insert_continues_when_agent_makes_no_wiki_changes(monkeypatch, tmp_path: Path) -> None: + calls: list[str] = [] + + def run_merge_agent(*, task_id, documents, root_dir, config): + document = documents[0] + calls.append(document.source_name) + if "beta" in document.source_name: + return {"added": [], "updated": [], "deleted": []} + + pages = root_dir / "pages" + pages.mkdir(parents=True, exist_ok=True) + page = pages / f"{slugify(document.title)}.md" + page.write_text( + frontmatter_page( + document.title, + document.source_name, + summarize(document.markdown_content), + document.markdown_content, + ), + encoding="utf-8", + ) + return {"added": [FileChange("added", document.title, f"pages/{page.name}")], "updated": [], "deleted": []} + + monkeypatch.setattr("heta.kb.insert.run_merge_agent", run_merge_agent) + first = tmp_path / "alpha.md" + second = tmp_path / "beta.md" + third = tmp_path / "gamma.md" + first.write_text("# Alpha\n\nFirst details.", encoding="utf-8") + second.write_text("# Beta\n\nSecond details.", encoding="utf-8") + third.write_text("# Gamma\n\nThird details.", encoding="utf-8") + + result = insert_paths([first, second, third], _config(), base_dir=tmp_path) + + wiki = tmp_path / "workspace" / "kb" / "wiki" + assert len(calls) == 3 + assert [change.path for change in result.added] == ["pages/1-alpha.md", "pages/2-gamma.md"] + assert result.skipped_documents == [calls[1]] + assert (wiki / "pages" / "1-alpha.md").exists() + assert not (wiki / "pages" / "2-beta.md").exists() + assert (wiki / "pages" / "2-gamma.md").exists() + assert "Skipped no-op merge" not in (wiki / "log.md").read_text(encoding="utf-8") + + +def test_ensure_code_raw_links_restores_agent_dropped_raw_link(tmp_path: Path) -> None: + wiki = tmp_path / "wiki" + page = wiki / "pages" / "1-code-demo.md" + page.parent.mkdir(parents=True) + page.write_text( + frontmatter_page("Code - demo.py", "2026-05-15_demo.py", "Summary.", "### File Overview\n- language: python"), + encoding="utf-8", + ) + document = ParsedDocument( + source_path=tmp_path / "demo.py", + archived_path=tmp_path / "raw" / "2026-05-15_demo.py", + title="Code - demo.py", + markdown_content="", + source_name="2026-05-15_demo.py", + metadata={"extension": ".py"}, + ) + + _ensure_code_raw_links(wiki, document, [FileChange("added", "Code - demo.py", "pages/1-code-demo.md")]) + + assert "[Raw source](<../../raw/2026-05-15_demo.py>)" in page.read_text(encoding="utf-8") + + def test_pdf_requires_mineru_when_disabled(tmp_path: Path) -> None: source = tmp_path / "paper.pdf" source.write_bytes(b"%PDF") @@ -104,6 +266,60 @@ def test_pdf_requires_mineru_when_disabled(tmp_path: Path) -> None: collect_insert_files([source], _config()) +def test_office_requires_mineru_when_disabled(tmp_path: Path) -> None: + source = tmp_path / "deck.pptx" + source.write_bytes(b"pptx") + + with pytest.raises(ValueError, match="requires MinerU"): + collect_insert_files([source], _config()) + + +def test_collect_insert_files_accepts_office_when_mineru_enabled(tmp_path: Path) -> None: + files = [] + for name in ["notes.doc", "notes.docx", "deck.ppt", "deck.pptx", "sheet.xls", "sheet.xlsx"]: + file = tmp_path / name + file.write_bytes(b"office") + files.append(file) + + collected = collect_insert_files( + [tmp_path], + _config(MinerUConfig(enable=True, provider="cloud", api_key="mineru-token", endpoint=None)), + ) + + assert collected == sorted(files) + + +def test_collect_insert_files_accepts_common_images(tmp_path: Path) -> None: + image = tmp_path / "diagram.png" + image.write_bytes(b"png") + + files = collect_insert_files([image], _config()) + + assert files == [image] + + +def test_collect_insert_files_accepts_audio_and_video(tmp_path: Path) -> None: + audio = tmp_path / "meeting.mp3" + video = tmp_path / "demo.mp4" + audio.write_bytes(b"mp3") + video.write_bytes(b"mp4") + + files = collect_insert_files([audio, video], _config()) + + assert files == [audio, video] + + +def test_collect_insert_files_accepts_code_and_html(tmp_path: Path) -> None: + code = tmp_path / "module.py" + html = tmp_path / "index.html" + code.write_text("def run():\n pass\n", encoding="utf-8") + html.write_text("", encoding="utf-8") + + files = collect_insert_files([tmp_path], _config()) + + assert files == [html, code] + + def test_collect_directory_skips_workspace(tmp_path: Path) -> None: source = tmp_path / "a.md" workspace_file = tmp_path / "workspace" / "kb" / "wiki" / "pages" / "old.md" @@ -142,3 +358,30 @@ def test_normalize_wiki_pages_assigns_max_plus_one_without_reusing_deleted_ids(t assert "- [1] [[Old]] (pages/1-old.md) — Old summary." in index assert "- [3] [[Existing]] (pages/3-existing.md) — Existing summary." in index assert "- [4] [[New Topic]] (pages/4-new-topic.md) — New summary." in index + + +def test_repair_broken_wiki_links_downgrades_missing_targets(tmp_path: Path) -> None: + wiki = tmp_path / "wiki" + pages = wiki / "pages" + pages.mkdir(parents=True) + page = pages / "1-topic.md" + page.write_text( + frontmatter_page( + "Topic", + "source.md", + "Summary.", + "See [[Existing Topic]] and [[Missing Topic]].", + ), + encoding="utf-8", + ) + (pages / "2-existing-topic.md").write_text( + frontmatter_page("Existing Topic", "source.md", "Summary.", "Body."), + encoding="utf-8", + ) + + repair_broken_wiki_links(wiki) + + text = page.read_text(encoding="utf-8") + assert "[[Existing Topic]]" in text + assert "[[Missing Topic]]" not in text + assert "Missing Topic" in text diff --git a/tests/test_kb_invalidate.py b/tests/test_kb_invalidate.py new file mode 100644 index 0000000..cc1e4d1 --- /dev/null +++ b/tests/test_kb_invalidate.py @@ -0,0 +1,140 @@ +"""Tests for heta.mem.kb_invalidate.""" + +from __future__ import annotations + +import time +import uuid +from pathlib import Path + +import pytest + +from heta.mem.db import get_connection, init_db +from heta.mem.kb_invalidate import delete_all_insights, delete_insights_by_paths +from heta.mem.kb_store import insert_insight_embedding, insert_kb_insight +from heta.mem.meta_store import insert_meta +from heta.mem.models import KBInsight, MemoryMeta + + +@pytest.fixture() +def conn(tmp_path: Path): + db = tmp_path / "test_mem.sqlite3" + c = get_connection(db, with_vec=True) + init_db(c) + yield c + c.close() + + +def _now() -> int: + return int(time.time()) + + +def _insert_insight(conn, source_paths, insight_text: str = "fact") -> str: + """source_paths can be a single str (for backward-compat tests) or a list.""" + if isinstance(source_paths, str): + source_paths = [source_paths] + mid = str(uuid.uuid4()) + insert_meta(conn, MemoryMeta( + memory_id=mid, memory_type="kb_insight", session_id=None, + origin="kb_insight", created_at=_now(), last_access_at=_now(), + )) + insert_kb_insight(conn, KBInsight( + memory_id=mid, insight=insight_text, question="q", + source_paths=source_paths, wiki_id=None, heading_path=None, + created_at=_now(), + )) + # 1024-dim float embedding (matches EMBEDDING_DIM) + from heta.mem.client import EMBEDDING_DIM + insert_insight_embedding(conn, mid, [0.0] * EMBEDDING_DIM) + conn.commit() + return mid + + +def _count(conn, table: str) -> int: + return conn.execute(f"SELECT COUNT(*) FROM {table}").fetchone()[0] + + +# ── tests ───────────────────────────────────────────────────────────────────── + + +def test_delete_by_paths_removes_matching_rows(conn): + _insert_insight(conn, "pages/1-foo.md") + _insert_insight(conn, "pages/1-foo.md") # 2 insights on same page + _insert_insight(conn, "pages/2-bar.md") + + deleted = delete_insights_by_paths(conn, ["pages/1-foo.md"]) + + assert deleted == 2 + assert _count(conn, "kb_insight") == 1 + assert _count(conn, "kb_insight_vec") == 1 + assert _count(conn, "memory_meta") == 1 + + +def test_delete_by_paths_leaves_other_pages_untouched(conn): + _insert_insight(conn, "pages/1-foo.md") + bar_id = _insert_insight(conn, "pages/2-bar.md") + + delete_insights_by_paths(conn, ["pages/1-foo.md"]) + + rows = conn.execute("SELECT memory_id FROM kb_insight").fetchall() + assert [r[0] for r in rows] == [bar_id] + + +def test_delete_by_paths_empty_input(conn): + _insert_insight(conn, "pages/1-foo.md") + assert delete_insights_by_paths(conn, []) == 0 + assert _count(conn, "kb_insight") == 1 + + +def test_delete_by_paths_no_match(conn): + _insert_insight(conn, "pages/1-foo.md") + assert delete_insights_by_paths(conn, ["pages/does-not-exist.md"]) == 0 + assert _count(conn, "kb_insight") == 1 + + +def test_delete_all_clears_everything(conn): + _insert_insight(conn, "pages/1-foo.md") + _insert_insight(conn, "pages/2-bar.md") + _insert_insight(conn, "pages/3-baz.md") + + deleted = delete_all_insights(conn) + + assert deleted == 3 + assert _count(conn, "kb_insight") == 0 + assert _count(conn, "kb_insight_vec") == 0 + assert _count(conn, "memory_meta") == 0 + + +def test_delete_all_on_empty_db_returns_zero(conn): + assert delete_all_insights(conn) == 0 + + +def test_delete_by_paths_invalidates_multi_source_insight(conn): + """An insight derived from multiple pages dies when ANY of its sources changes.""" + multi = _insert_insight(conn, ["pages/1-foo.md", "pages/2-bar.md"]) + solo = _insert_insight(conn, ["pages/3-baz.md"]) + + deleted = delete_insights_by_paths(conn, ["pages/2-bar.md"]) + + assert deleted == 1 + remaining = [r[0] for r in conn.execute("SELECT memory_id FROM kb_insight").fetchall()] + assert multi not in remaining + assert solo in remaining + # both rows in kb_insight_source for the multi insight should be gone + assert _count(conn, "kb_insight_source") == 1 + + +def test_delete_by_paths_preserves_other_memory_types(conn): + """Deleting kb_insight by path must not touch L1/L2/etc.""" + _insert_insight(conn, "pages/1-foo.md") + # an unrelated memory_meta row (e.g. L2) + other = str(uuid.uuid4()) + insert_meta(conn, MemoryMeta( + memory_id=other, memory_type="L2", session_id=None, + origin="extracted", created_at=_now(), last_access_at=_now(), + )) + + delete_insights_by_paths(conn, ["pages/1-foo.md"]) + + assert _count(conn, "memory_meta") == 1 + remaining = conn.execute("SELECT memory_id FROM memory_meta").fetchone() + assert remaining[0] == other diff --git a/tests/test_memory_ingest.py b/tests/test_memory_ingest.py new file mode 100644 index 0000000..f59d066 --- /dev/null +++ b/tests/test_memory_ingest.py @@ -0,0 +1,384 @@ +"""Tests for the heta remember ingestion pipeline (pipeline.py). + +All LLM and embedding calls are mocked so tests run offline and deterministically. +""" + +from __future__ import annotations + +import json +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +from heta.config.schema import InsertPlanningConfig, HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.mem.client import EMBEDDING_DIM +from heta.mem.db import get_connection, init_db +from heta.mem.pipeline import remember + + +# ── constants ───────────────────────────────────────────────────────────────── + +FAKE_EMB = [0.01] * EMBEDDING_DIM # deterministic dummy embedding + +EPISODE_DICT = { + "who": ["user"], + "what": "参加了技术分享会", + "where_loc": "公司会议室", + "when_text": "上周", + "when_resolved": "2026-W19", + "when_precision": "week", + "why": None, + "summary": "用户上周在公司会议室参加了技术分享会", +} + +FACT_DICT = { + "subject": "用户", + "predicate": "居住在", + "object": "北京朝阳区", + "object_type": "literal", + "when_text": None, + "when_resolved": None, + "when_precision": None, +} + + +# ── fixtures ────────────────────────────────────────────────────────────────── + +@pytest.fixture() +def config() -> HetaConfig: + return HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig.enabled(), + insert_planning=InsertPlanningConfig.enabled(), + ) + + +@pytest.fixture() +def tmp_db(tmp_path: Path) -> Path: + return tmp_path / "mem.sqlite3" + + +# ── patch helper ────────────────────────────────────────────────────────────── + +@contextmanager +def _patch_pipeline( + tmp_db: Path, + episodes: list[dict] | None = None, + facts: list[dict] | None = None, + conflicts: list[str] | None = None, # memory_ids to deprecate +): + """Patch all external I/O in pipeline.remember so tests run offline.""" + if episodes is None: + episodes = [] + if facts is None: + facts = [] + if conflicts is None: + conflicts = [] + + mock_client = MagicMock() + + def _open_conn(path, *, with_vec=False): + c = get_connection(tmp_db, with_vec=True) + init_db(c) + return c + + with ( + patch("heta.mem.pipeline.ensure_mem_dir"), + patch("heta.mem.pipeline.db_path", return_value=tmp_db), + patch("heta.mem.pipeline.get_connection", side_effect=_open_conn), + patch("heta.mem.pipeline.init_db"), + patch("heta.mem.pipeline.build_client", return_value=(mock_client, "mock-llm")), + patch("heta.mem.pipeline.build_embedding_client", return_value=(mock_client, "mock-emb")), + patch("heta.mem.pipeline.extract_episodes", return_value=episodes), + patch("heta.mem.pipeline.extract_facts", return_value=facts), + patch("heta.mem.pipeline.embed_text", return_value=FAKE_EMB), + patch("heta.mem.pipeline.detect_conflicts", return_value=(conflicts, FAKE_EMB)), + ): + yield + + +def _open(tmp_db: Path): + """Open a fresh read connection to the tmp DB after pipeline closes it.""" + return get_connection(tmp_db, with_vec=True) + + +# ── basic ingestion ─────────────────────────────────────────────────────────── + +def test_remember_creates_session_and_l0_turn(config, tmp_db) -> None: + with _patch_pipeline(tmp_db): + result = remember("hello world", config) + + conn = _open(tmp_db) + sessions = conn.execute("SELECT * FROM session WHERE session_id = ?", (result.session_id,)).fetchall() + turns = conn.execute("SELECT * FROM l0_turn WHERE session_id = ?", (result.session_id,)).fetchall() + conn.close() + + assert len(sessions) == 1 + assert sessions[0]["session_id"] == result.session_id + assert sessions[0]["ended_at"] is not None # session was closed + assert len(turns) == 1 + assert turns[0]["text_content"] == "hello world" + assert turns[0]["role"] == "user" + + +def test_remember_l0_turn_indexed_in_fts(config, tmp_db) -> None: + with _patch_pipeline(tmp_db): + remember("爬山打羽毛球", config) + + conn = _open(tmp_db) + rows = conn.execute( + "SELECT text_content FROM l0_turn_fts WHERE text_content MATCH ?", + ('"爬山打羽毛球"',), + ).fetchall() + conn.close() + + assert len(rows) == 1 + + +def test_remember_empty_extraction_succeeds(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, episodes=[], facts=[]): + result = remember("no events here", config) + + assert result.l1_count == 0 + assert result.l2_count == 0 + + +def test_remember_returns_correct_counts(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, episodes=[EPISODE_DICT], facts=[FACT_DICT, FACT_DICT]): + result = remember("some text", config) + + assert result.l1_count == 1 + assert result.l2_count == 2 + assert result.session_id != "" + assert result.elapsed_s >= 0 + + +# ── L1 episode storage ──────────────────────────────────────────────────────── + +def test_remember_l1_episode_stored_correctly(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, episodes=[EPISODE_DICT]): + result = remember("some text", config) + + conn = _open(tmp_db) + row = conn.execute( + "SELECT * FROM l1_episodic e JOIN memory_meta m ON e.memory_id = m.memory_id " + "WHERE m.session_id = ?", + (result.session_id,), + ).fetchone() + conn.close() + + assert row is not None + assert row["what"] == EPISODE_DICT["what"] + assert row["where_loc"] == EPISODE_DICT["where_loc"] + assert row["when_text"] == EPISODE_DICT["when_text"] + assert row["when_resolved"] == EPISODE_DICT["when_resolved"] + assert row["when_precision"] == EPISODE_DICT["when_precision"] + assert row["summary"] == EPISODE_DICT["summary"] + assert json.loads(row["who"]) == EPISODE_DICT["who"] + assert row["memory_type"] == "L1" + assert row["status"] == "active" + + +def test_remember_l1_episode_embedding_inserted(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, episodes=[EPISODE_DICT]): + result = remember("some text", config) + + conn = _open(tmp_db) + meta = conn.execute( + "SELECT memory_id FROM memory_meta WHERE session_id = ? AND memory_type = 'L1'", + (result.session_id,), + ).fetchone() + vec_row = conn.execute( + "SELECT memory_id FROM l1_episode_vec WHERE memory_id = ?", + (meta["memory_id"],), + ).fetchone() + conn.close() + + assert vec_row is not None + + +def test_remember_l1_who_defaults_to_user_when_missing(config, tmp_db) -> None: + ep = {**EPISODE_DICT} + del ep["who"] + with _patch_pipeline(tmp_db, episodes=[ep]): + remember("some text", config) + + conn = _open(tmp_db) + row = conn.execute("SELECT who FROM l1_episodic").fetchone() + conn.close() + + assert json.loads(row["who"]) == ["user"] + + +# ── L2 fact storage ─────────────────────────────────────────────────────────── + +def test_remember_l2_fact_stored_correctly(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, facts=[FACT_DICT]): + result = remember("some text", config) + + conn = _open(tmp_db) + row = conn.execute( + "SELECT * FROM l2_semantic s JOIN memory_meta m ON s.memory_id = m.memory_id " + "WHERE m.session_id = ?", + (result.session_id,), + ).fetchone() + conn.close() + + assert row is not None + assert row["subject"] == FACT_DICT["subject"] + assert row["predicate"] == FACT_DICT["predicate"] + assert row["object"] == FACT_DICT["object"] + assert row["object_type"] == "literal" + assert row["t_valid_end"] is None # active, not yet expired + assert row["status"] == "active" + assert row["memory_type"] == "L2" + + +def test_remember_l2_fact_embedding_inserted(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, facts=[FACT_DICT]): + result = remember("some text", config) + + conn = _open(tmp_db) + meta = conn.execute( + "SELECT memory_id FROM memory_meta WHERE session_id = ? AND memory_type = 'L2'", + (result.session_id,), + ).fetchone() + vec_row = conn.execute( + "SELECT memory_id FROM l2_fact_vec WHERE memory_id = ?", + (meta["memory_id"],), + ).fetchone() + conn.close() + + assert vec_row is not None + + +def test_remember_l2_object_type_list_coerced_to_string(config, tmp_db) -> None: + """LLM occasionally returns object_type as a list; pipeline should normalise it.""" + fact = {**FACT_DICT, "object_type": ["literal", "extra"]} + with _patch_pipeline(tmp_db, facts=[fact]): + remember("some text", config) + + conn = _open(tmp_db) + row = conn.execute("SELECT object_type FROM l2_semantic").fetchone() + conn.close() + + assert isinstance(row["object_type"], str) + assert row["object_type"] == "literal" + + +# ── conflict resolution ─────────────────────────────────────────────────────── + +def test_remember_conflict_deprecates_old_fact(config, tmp_db) -> None: + """When detect_conflicts returns an old memory_id, that fact is expired + deprecated.""" + # session 1: insert a fact directly so we have an old_id to conflict + from heta.mem.l2_store import insert_fact, insert_fact_embedding + from heta.mem.meta_store import insert_meta + from heta.mem.models import L2Semantic, MemoryMeta + from heta.mem.session_store import create_session, close_session + from heta.mem.models import Session + import time, uuid + + old_id = str(uuid.uuid4()) + now = int(time.time()) + + setup_conn = get_connection(tmp_db, with_vec=True) + init_db(setup_conn) + sid0 = str(uuid.uuid4()) + create_session(setup_conn, Session(session_id=sid0, started_at=now)) + close_session(setup_conn, sid0, now) + insert_meta(setup_conn, MemoryMeta( + memory_id=old_id, memory_type="L2", session_id=sid0, + origin="extracted", created_at=now, last_access_at=now, + )) + insert_fact(setup_conn, L2Semantic( + memory_id=old_id, subject="用户", predicate="居住在", object="北京朝阳区", + object_type="literal", fact_text="用户 居住在 北京朝阳区", + t_valid_start=now, + )) + insert_fact_embedding(setup_conn, old_id, FAKE_EMB) + setup_conn.close() + + # session 2: new fact conflicts with old_id + new_fact = {**FACT_DICT, "object": "北京海淀区"} + with _patch_pipeline(tmp_db, facts=[new_fact], conflicts=[old_id]): + remember("搬家了", config) + + conn = _open(tmp_db) + old_row = conn.execute( + "SELECT s.t_valid_end, m.status, m.deprecated_by " + "FROM l2_semantic s JOIN memory_meta m ON s.memory_id = m.memory_id " + "WHERE s.memory_id = ?", + (old_id,), + ).fetchone() + active_rows = conn.execute( + "SELECT object FROM l2_semantic WHERE t_valid_end IS NULL" + ).fetchall() + conn.close() + + assert old_row["t_valid_end"] is not None # expired + assert old_row["status"] == "deprecated" + assert old_row["deprecated_by"] is not None # FK to new fact + assert len(active_rows) == 1 + assert active_rows[0]["object"] == "北京海淀区" + + +def test_remember_no_conflict_keeps_both_facts(config, tmp_db) -> None: + """When detect_conflicts returns [], both old and new facts remain active.""" + fact_a = {**FACT_DICT, "predicate": "喜欢", "object": "爬山"} + fact_b = {**FACT_DICT, "predicate": "喜欢", "object": "羽毛球"} + + with _patch_pipeline(tmp_db, facts=[fact_a], conflicts=[]): + remember("喜欢爬山", config) + with _patch_pipeline(tmp_db, facts=[fact_b], conflicts=[]): + remember("喜欢羽毛球", config) + + conn = _open(tmp_db) + active = conn.execute( + "SELECT object FROM l2_semantic WHERE t_valid_end IS NULL ORDER BY object" + ).fetchall() + conn.close() + + objects = {r["object"] for r in active} + assert objects == {"爬山", "羽毛球"} + + +def test_remember_detect_conflicts_receives_session_id(config, tmp_db) -> None: + """detect_conflicts must be called with the current session_id so same-session + facts are excluded from conflict candidates.""" + captured_kwargs: dict = {} + + def _fake_detect_conflicts(**kwargs): + captured_kwargs.update(kwargs) + return ([], FAKE_EMB) + + with ( + _patch_pipeline(tmp_db, facts=[FACT_DICT]), + patch("heta.mem.pipeline.detect_conflicts", side_effect=_fake_detect_conflicts), + ): + result = remember("some text", config) + + assert "session_id" in captured_kwargs + assert captured_kwargs["session_id"] == result.session_id + + +# ── multiple sessions ───────────────────────────────────────────────────────── + +def test_remember_multiple_sessions_accumulate(config, tmp_db) -> None: + with _patch_pipeline(tmp_db, facts=[FACT_DICT]): + r1 = remember("first", config) + with _patch_pipeline(tmp_db, facts=[FACT_DICT]): + r2 = remember("second", config) + + assert r1.session_id != r2.session_id + + conn = _open(tmp_db) + n_sessions = conn.execute("SELECT COUNT(*) FROM session").fetchone()[0] + n_turns = conn.execute("SELECT COUNT(*) FROM l0_turn").fetchone()[0] + conn.close() + + assert n_sessions == 2 + assert n_turns == 2 diff --git a/tests/test_mineru_parser.py b/tests/test_mineru_parser.py new file mode 100644 index 0000000..697e762 --- /dev/null +++ b/tests/test_mineru_parser.py @@ -0,0 +1,160 @@ +from __future__ import annotations + +import zipfile +from io import BytesIO +from pathlib import Path + +from heta.config.schema import HetaConfig, InsertPlanningConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.kb.parser import parse_document + + +class _Response: + def __init__( + self, + *, + status_code: int = 200, + payload: dict | None = None, + content: bytes = b"", + text: str = "", + headers: dict[str, str] | None = None, + ) -> None: + self.status_code = status_code + self._payload = payload or {} + self.content = content + self.text = text + self.headers = headers or {} + + def json(self) -> dict: + return self._payload + + +def _config() -> HetaConfig: + return HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig(enable=True, provider="cloud", api_key="mineru-token", endpoint=None), + vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + ) + + +def test_parse_document_accepts_office_via_mineru_cloud(monkeypatch, tmp_path: Path) -> None: + source = tmp_path / "slides.pptx" + archived = tmp_path / "2026-05-15_slides.pptx" + source.write_bytes(b"pptx") + archived.write_bytes(b"pptx") + zip_bytes = _mineru_zip("# Slides\n\nParsed by MinerU.") + requests_seen: list[tuple[str, str]] = [] + + def post(url, **kwargs): + requests_seen.append(("POST", url)) + assert kwargs["headers"]["Authorization"] == "Bearer mineru-token" + assert kwargs["json"]["files"] == [{"name": "2026-05-15_slides.pptx", "data_id": "2026-05-15_slides"}] + assert kwargs["json"]["model_version"] == "vlm" + return _Response(payload={"code": 0, "data": {"batch_id": "batch-1", "file_urls": ["https://upload"]}}) + + def put(url, **kwargs): + requests_seen.append(("PUT", url)) + assert url == "https://upload" + return _Response(status_code=200) + + def get(url, **kwargs): + requests_seen.append(("GET", url)) + if url.endswith("/batch-1"): + assert kwargs["headers"]["Authorization"] == "Bearer mineru-token" + return _Response( + payload={ + "code": 0, + "data": { + "extract_result": [ + { + "file_name": "2026-05-15_slides.pptx", + "state": "done", + "full_zip_url": "https://result.zip", + } + ] + }, + } + ) + assert url == "https://result.zip" + return _Response(content=zip_bytes) + + monkeypatch.setattr("heta.kb.parser.requests.post", post) + monkeypatch.setattr("heta.kb.parser.requests.put", put) + monkeypatch.setattr("heta.kb.parser.requests.get", get) + + document = parse_document(source, archived, _config()) + + assert document.title == "Slides" + assert document.metadata["extension"] == ".pptx" + assert document.markdown_content == "# Slides\n\nParsed by MinerU." + assert [method for method, _ in requests_seen] == ["POST", "PUT", "GET", "GET"] + + +def test_parse_document_uses_local_mineru_zip_artifacts(monkeypatch, tmp_path: Path) -> None: + source = tmp_path / "paper.pdf" + archived = tmp_path / "2026-05-15_paper.pdf" + source.write_bytes(b"%PDF") + archived.write_bytes(b"%PDF") + zip_bytes = _mineru_zip( + "# Paper\n\n\n", + content_list=[ + { + "type": "image", + "img_path": "images/figure.jpg", + "bbox": [1, 2, 3, 4], + "page_idx": 2, + } + ], + images={"images/figure.jpg": b"jpg"}, + ) + config = HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig(enable=True, provider="local", api_key=None, endpoint="http://127.0.0.1:8000"), + vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + ) + + def post(url, **kwargs): + assert url == "http://127.0.0.1:8000/file_parse" + assert "files" in kwargs["files"] + assert kwargs["data"]["response_format_zip"] == "true" + assert kwargs["data"]["return_content_list"] == "true" + return _Response(content=zip_bytes, headers={"content-type": "application/zip"}) + + monkeypatch.setattr("heta.kb.parser.requests.post", post) + + document = parse_document( + source, + archived, + config, + original_name="paper.pdf", + page_offset=20, + base_dir=tmp_path, + ) + + assert "../../raw/parsed/2026-05-15_paper/images/figure.jpg" in document.markdown_content + assert "Source: paper.pdf, page 23, bbox [1, 2, 3, 4]" in document.markdown_content + parsed = tmp_path / "workspace" / "kb" / "raw" / "parsed" / "2026-05-15_paper" + assert (parsed / "full.md").exists() + assert (parsed / "content_list.json").exists() + assert (parsed / "images" / "figure.jpg").read_bytes() == b"jpg" + + +def _mineru_zip( + markdown: str, + *, + content_list: list[dict] | None = None, + images: dict[str, bytes] | None = None, +) -> bytes: + buffer = BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("full.md", markdown) + if content_list is not None: + import json + + archive.writestr("demo_content_list.json", json.dumps(content_list)) + for path, data in (images or {}).items(): + archive.writestr(path, data) + return buffer.getvalue() diff --git a/tests/test_pdf_plan.py b/tests/test_pdf_plan.py index 1dee6dd..0a76351 100644 --- a/tests/test_pdf_plan.py +++ b/tests/test_pdf_plan.py @@ -3,7 +3,7 @@ from pypdf import PdfWriter -from heta.config.schema import HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.config.schema import InsertPlanningConfig, HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig from heta.kb import paths from heta.kb.pdf_plan import PDF_PAGE_THRESHOLD, SplitUnit, estimate_pdf_pages, plan_insert_files @@ -17,11 +17,11 @@ def test_plan_insert_files_splits_large_pdf_and_keeps_original(tmp_path: Path) - assert len(plans) == 1 assert plans[0].enabled is True assert plans[0].page_count == PDF_PAGE_THRESHOLD + 1 - assert plans[0].parts == 3 - assert len(prepared) == 3 + assert plans[0].parts == 5 + assert len(prepared) == 5 assert (paths.raw_dir(tmp_path) / "originals").exists() assert all(item.archived_path.exists() for item in prepared) - assert [estimate_pdf_pages(item.archived_path) for item in prepared] == [40, 40, 1] + assert [estimate_pdf_pages(item.archived_path) for item in prepared] == [20, 20, 20, 20, 1] assert all(item.original_path is not None for item in prepared) assert all(item.metadata_path is not None and item.metadata_path.exists() for item in prepared) @@ -60,12 +60,12 @@ def test_plan_insert_files_uses_agent_split_plan(monkeypatch, tmp_path: Path) -> assert plans[0].document_type == "textbook" assert plans[0].split_strategy == "chapter" - assert [item.page_start for item in prepared] == [1, 31, 71] - assert [item.page_end for item in prepared] == [30, 70, 81] + assert [item.page_start for item in prepared] == [1, 21, 31, 51, 71] + assert [item.page_end for item in prepared] == [20, 30, 50, 70, 81] metadata = json.loads(prepared[0].metadata_path.read_text(encoding="utf-8")) assert metadata["original"].endswith("large.pdf") assert metadata["start_page"] == 1 - assert metadata["end_page"] == 30 + assert metadata["end_page"] == 20 assert metadata["split_strategy"] == "chapter" @@ -84,7 +84,7 @@ def test_plan_insert_files_falls_back_when_agent_plan_is_invalid(monkeypatch, tm prepared, plans = plan_insert_files([source], config=_config(), base_dir=tmp_path) assert plans[0].split_strategy == "fixed_page_window" - assert [estimate_pdf_pages(item.archived_path) for item in prepared] == [40, 40, 1] + assert [estimate_pdf_pages(item.archived_path) for item in prepared] == [20, 20, 20, 20, 1] def test_plan_insert_files_fills_pages_missing_from_agent_plan(monkeypatch, tmp_path: Path) -> None: @@ -102,7 +102,7 @@ def test_plan_insert_files_fills_pages_missing_from_agent_plan(monkeypatch, tmp_ prepared, plans = plan_insert_files([source], config=_config(), base_dir=tmp_path) assert plans[0].split_strategy == "section" - assert [(item.page_start, item.page_end) for item in prepared] == [(1, 40), (41, 60), (61, 81)] + assert [(item.page_start, item.page_end) for item in prepared] == [(1, 20), (21, 40), (41, 60), (61, 80), (81, 81)] def _write_pdf(path: Path, pages: int) -> None: @@ -119,4 +119,5 @@ def _config() -> HetaConfig: llm=LLMConfig(provider="qwen", api_key="sk-test"), mineru=MinerUConfig.disabled(), vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), ) diff --git a/tests/test_provider_clients.py b/tests/test_provider_clients.py new file mode 100644 index 0000000..2ebe73b --- /dev/null +++ b/tests/test_provider_clients.py @@ -0,0 +1,83 @@ +from heta.config.schema import HetaConfig, InsertPlanningConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.providers import clients + + +def _config() -> HetaConfig: + return HetaConfig( + version=1, + llm=LLMConfig( + provider="custom", + api_key="legacy-key", + chat_api_key="chat-key", + chat_model="chat-model", + chat_base_url="http://chat.local/v1", + multimodal_api_key="mm-key", + multimodal_model="mm-model", + multimodal_base_url="http://mm.local/v1", + embedding_api_key="embedding-key", + embedding_model="embedding-model", + embedding_base_url="http://embedding.local/v1", + ), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + ) + + +def test_provider_clients_use_capability_specific_api_keys(monkeypatch) -> None: + seen: list[dict] = [] + + class FakeOpenAI: + def __init__(self, **kwargs): + seen.append(kwargs) + + monkeypatch.setattr(clients, "OpenAI", FakeOpenAI) + config = _config() + + chat = clients.build_chat_client(config) + multimodal = clients.build_multimodal_client(config) + embedding = clients.build_embedding_client(config) + + assert chat.model == "chat-model" + assert multimodal.model == "mm-model" + assert embedding.model == "embedding-model" + assert seen == [ + {"api_key": "chat-key", "timeout": 60, "base_url": "http://chat.local/v1"}, + {"api_key": "mm-key", "timeout": 300, "base_url": "http://mm.local/v1"}, + {"api_key": "embedding-key", "timeout": 120, "base_url": "http://embedding.local/v1"}, + ] + + +def test_extra_body_prefers_explicit_config() -> None: + config = _config() + explicit = HetaConfig( + version=config.version, + llm=LLMConfig( + provider="custom", + api_key="legacy-key", + chat_api_key="chat-key", + chat_model="chat-model", + chat_base_url="http://chat.local/v1", + chat_extra_body={"enable_thinking": False}, + embedding_api_key="embedding-key", + embedding_model="embedding-model", + embedding_base_url="http://embedding.local/v1", + ), + mineru=config.mineru, + vector_index=config.vector_index, + insert_planning=config.insert_planning, + ) + + assert clients.extra_body(explicit) == {"enable_thinking": False} + + +def test_extra_body_keeps_qwen_default() -> None: + config = HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig(enable=False), + insert_planning=InsertPlanningConfig.enabled(), + ) + + assert clients.extra_body(config) == {"enable_thinking": False} diff --git a/tests/test_query.py b/tests/test_query.py index a76eb16..51396d9 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -2,12 +2,13 @@ import pytest -from heta.config.schema import HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.config.schema import InsertPlanningConfig, HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig from heta.kb import paths from heta.kb.text import frontmatter_page +from heta.query.agent import _parse_final_answer, _vector_match_map from heta.query.models import QueryResult, QuerySource, VectorMatch from heta.query.pipeline import run_wiki_query -from heta.query.tools import format_vector_matches, read_page, source_from_page_path +from heta.query.tools import format_vector_matches, read_page, read_raw, source_from_page_path def _config(vector_enabled: bool = False) -> HetaConfig: @@ -16,6 +17,7 @@ def _config(vector_enabled: bool = False) -> HetaConfig: llm=LLMConfig(provider="qwen", api_key="sk-test"), mineru=MinerUConfig.disabled(), vector_index=VectorIndexConfig(enable=vector_enabled), + insert_planning=InsertPlanningConfig.enabled(), ) @@ -51,6 +53,16 @@ def test_read_page_is_limited_to_pages(tmp_path: Path) -> None: assert read_page("index.md", tmp_path).startswith("error:") +def test_read_raw_is_limited_to_raw_directory(tmp_path: Path) -> None: + raw = paths.raw_dir(tmp_path) + raw.mkdir(parents=True) + (raw / "module.py").write_text("def run():\n return True\n", encoding="utf-8") + + assert "def run" in read_raw("raw/module.py", tmp_path) + assert "def run" in read_raw("../../raw/module.py", tmp_path) + assert read_raw("../heta.yaml", tmp_path).startswith("error:") + + def test_source_from_page_path_reads_frontmatter_and_wiki_id(tmp_path: Path) -> None: page = paths.pages_dir(tmp_path) / "12-hetagen.md" page.parent.mkdir(parents=True) @@ -61,6 +73,79 @@ def test_source_from_page_path_reads_frontmatter_and_wiki_id(tmp_path: Path) -> assert source == QuerySource(12, "HetaGen", "pages/12-hetagen.md", "Content") +def test_query_sources_include_validated_vector_chunks_only(tmp_path: Path) -> None: + pages = paths.pages_dir(tmp_path) + pages.mkdir(parents=True) + (pages / "8-image.md").write_text(frontmatter_page("Image", "image.png", "Image summary.", "Body."), encoding="utf-8") + (pages / "10-audio.md").write_text( + frontmatter_page("Audio", "audio.mp3", "Audio summary.", "Transcript."), + encoding="utf-8", + ) + vector_matches = _vector_match_map( + [ + VectorMatch(8, "8-image.md", "pages/8-image.md", "8:abc", "Content > Visible Text", "image text", 0.8), + VectorMatch(10, "10-audio.md", "pages/10-audio.md", "10:def", "Content > Transcript", "hello", 0.9), + ] + ) + + final = _parse_final_answer( + text=( + '{"answer": "The audio says hello.", "used_sources": [' + '{"path": "pages/10-audio.md", "heading_path": "Content > Transcript"},' + '{"path": "pages/8-image.md", "heading_path": "Content > Missing"}' + "]}" + ), + read_paths=set(), + vector_matches=vector_matches, + base_dir=tmp_path, + ) + + assert final.answer == "The audio says hello." + assert final.sources == [QuerySource(10, "Audio", "pages/10-audio.md", "Content > Transcript")] + + +def test_query_sources_accept_read_pages_without_vector_heading(tmp_path: Path) -> None: + pages = paths.pages_dir(tmp_path) + pages.mkdir(parents=True) + (pages / "10-audio.md").write_text( + frontmatter_page("Audio", "audio.mp3", "Audio summary.", "Transcript."), + encoding="utf-8", + ) + + final = _parse_final_answer( + text='{"answer": "From the full page.", "used_sources": [{"path": "pages/10-audio.md"}]}', + read_paths={"pages/10-audio.md"}, + vector_matches={}, + base_dir=tmp_path, + ) + + assert final.sources == [QuerySource(10, "Audio", "pages/10-audio.md")] + + +def test_query_sources_reject_raw_used_sources(tmp_path: Path) -> None: + final = _parse_final_answer( + text='{"answer": "Raw helped.", "used_sources": [{"path": "raw/module.py"}]}', + read_paths=set(), + vector_matches={}, + base_dir=tmp_path, + ) + + assert final.sources == [] + + +def test_parse_final_answer_marks_non_json_response_invalid(tmp_path: Path) -> None: + final = _parse_final_answer( + text="This is not JSON.", + read_paths=set(), + vector_matches={}, + base_dir=tmp_path, + ) + + assert final.answer == "This is not JSON." + assert final.sources == [] + assert not final.valid_json + + def test_format_vector_matches_includes_chunk_identity() -> None: text = format_vector_matches( [ diff --git a/tests/test_status.py b/tests/test_status.py index 3d231dd..e342824 100644 --- a/tests/test_status.py +++ b/tests/test_status.py @@ -1,7 +1,7 @@ from pathlib import Path from heta.cli.status import build_status_summary -from heta.config.schema import HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.config.schema import InsertPlanningConfig, HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig def test_status_summary_counts_kb_and_wiki_pages(tmp_path: Path) -> None: @@ -19,12 +19,15 @@ def test_status_summary_counts_kb_and_wiki_pages(tmp_path: Path) -> None: llm=LLMConfig(provider="qwen", api_key="sk-test"), mineru=MinerUConfig(enable=True, provider="local", api_key=None, endpoint="http://127.0.0.1:8000"), vector_index=VectorIndexConfig.enabled(), + insert_planning=InsertPlanningConfig.enabled(), ) summary = build_status_summary(config, tmp_path) assert summary.llm_provider == "qwen" assert summary.mineru == "local (http://127.0.0.1:8000)" + assert summary.insert_planning == "enabled" + assert summary.dynamic_insert == "disabled" assert summary.kb_files == 2 assert summary.wiki_pages == 1 assert summary.heta_space == tmp_path @@ -37,6 +40,8 @@ def test_status_summary_handles_missing_config_and_workspace(tmp_path: Path) -> assert summary.llm_provider == "not configured" assert summary.mineru == "not configured" + assert summary.insert_planning == "not configured" + assert summary.dynamic_insert == "not configured" assert summary.kb_files == 0 assert summary.wiki_pages == 0 assert summary.heta_used_bytes == 0 diff --git a/tests/test_validators.py b/tests/test_validators.py index a62d58c..507bfae 100644 --- a/tests/test_validators.py +++ b/tests/test_validators.py @@ -18,6 +18,14 @@ def test_validate_llm_non_200_fails(get: Mock) -> None: assert validate_llm("qwen", "bad-key") is False +@patch("heta.providers.llm.requests.get") +def test_validate_custom_llm_uses_base_url_models(get: Mock) -> None: + get.return_value.status_code = 200 + + assert validate_llm("custom", "sk-test", "http://llm.local/v1") is True + assert get.call_args.args[0] == "http://llm.local/v1/models" + + @patch("heta.providers.mineru.requests.post") def test_validate_mineru_cloud_success(post: Mock) -> None: post.return_value.status_code = 200 diff --git a/tests/test_vector_index.py b/tests/test_vector_index.py index e945d1e..975d2aa 100644 --- a/tests/test_vector_index.py +++ b/tests/test_vector_index.py @@ -5,9 +5,19 @@ import sqlite_vec -from heta.config.schema import HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig +from heta.config.schema import InsertPlanningConfig, HetaConfig, LLMConfig, MinerUConfig, VectorIndexConfig from heta.kb import paths -from heta.kb.vector_index import _ensure_schema, _insert_chunk, chunk_wiki_page, search_wiki_vector_index +from heta.kb.models import FileChange +from heta.kb.vector_index import ( + _ensure_schema, + _fts_terms, + _insert_chunk, + chunk_wiki_page, + search_wiki_fts_index, + search_wiki_hybrid_index, + search_wiki_vector_index, + sync_wiki_vector_index, +) def test_chunk_wiki_page_uses_heading_path_and_page_context(tmp_path: Path) -> None: @@ -77,6 +87,7 @@ def test_search_wiki_vector_index_returns_ranked_chunks(monkeypatch, tmp_path: P llm=LLMConfig(provider="qwen", api_key="sk-test"), mineru=MinerUConfig.disabled(), vector_index=VectorIndexConfig.enabled(), + insert_planning=InsertPlanningConfig.enabled(), ) results = search_wiki_vector_index(query="table synthesis", config=config, top_k=3, base_dir=tmp_path) @@ -88,8 +99,183 @@ def test_search_wiki_vector_index_returns_ranked_chunks(monkeypatch, tmp_path: P assert results[0].score == 1.0 +def test_search_wiki_fts_index_matches_mixed_chinese_and_codes(tmp_path: Path) -> None: + db_path = paths.vector_db_path(tmp_path) + db_path.parent.mkdir(parents=True) + conn = sqlite3.connect(db_path) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + _ensure_schema(conn) + target = chunk_wiki_page( + _write_page( + tmp_path, + "1-flow-control.md", + "ZXDOC-A-10-20-30 温度控制策略", + "结构图和部件说明。", + "温度控制策略结构图包含文档代码 ZXDOC-A-10-20-30-ALPHA。", + ) + )[0] + other = chunk_wiki_page( + _write_page( + tmp_path, + "2-power.md", + "ZXDOC-B-40-50-60 存储模块", + "存储说明。", + "主存储模块说明。", + ) + )[0] + _insert_chunk(conn, target, [1.0] + [0.0] * 1023) + _insert_chunk(conn, other, [0.0, 1.0] + [0.0] * 1022) + conn.commit() + conn.close() + + results = search_wiki_fts_index(query="温度控制策略 10-20-30", top_k=3, base_dir=tmp_path) + + assert results + assert results[0].wiki_id == 1 + assert results[0].retrieval == "fts" + + +def test_search_wiki_hybrid_index_fuses_vector_and_fts(monkeypatch, tmp_path: Path) -> None: + db_path = paths.vector_db_path(tmp_path) + db_path.parent.mkdir(parents=True) + conn = sqlite3.connect(db_path) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + _ensure_schema(conn) + vector_chunk = chunk_wiki_page( + _write_page( + tmp_path, + "1-semantic.md", + "HetaGen", + "Structured content generation.", + "HetaGen supports table synthesis.", + ) + )[0] + keyword_chunk = chunk_wiki_page( + _write_page( + tmp_path, + "2-code.md", + "ZXDOC-A-10-20-30", + "结构图。", + "文档代码 ZXDOC-A-10-20-30-ALPHA。", + ) + )[0] + _insert_chunk(conn, vector_chunk, [1.0] + [0.0] * 1023) + _insert_chunk(conn, keyword_chunk, [0.0, 1.0] + [0.0] * 1022) + conn.commit() + conn.close() + + monkeypatch.setattr("heta.kb.vector_index._embed_texts", lambda texts, config: [[1.0] + [0.0] * 1023]) + config = HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig.enabled(), + insert_planning=InsertPlanningConfig.enabled(), + ) + + results = search_wiki_hybrid_index( + query="table synthesis ZXDOC-A-10-20-30", + config=config, + top_k=2, + base_dir=tmp_path, + ) + + assert {result.wiki_id for result in results} == {1, 2} + + +def test_fts_terms_normalizes_mixed_width_punctuation() -> None: + assert _fts_terms("zxdoc－a－10－20 温度控制，结构图") == [ + "ZXDOC-A-10-20", + "温度控制", + "结构图", + ] + + +def test_search_wiki_fts_index_backfills_existing_chunks(tmp_path: Path) -> None: + db_path = paths.vector_db_path(tmp_path) + db_path.parent.mkdir(parents=True) + conn = sqlite3.connect(db_path) + conn.enable_load_extension(True) + sqlite_vec.load(conn) + conn.enable_load_extension(False) + _ensure_schema(conn) + conn.execute( + """ + INSERT INTO wiki_chunks ( + wiki_id, page_name, chunk_id, heading_path, content, content_hash + ) + VALUES (?, ?, ?, ?, ?, ?) + """, + ( + 7, + "7-existing.md", + "7:existing", + "Content", + "Page: ZXDOC-A-10-20 结构图\n\n已有 chunk 应该自动回填到 FTS。", + "existing", + ), + ) + conn.commit() + conn.close() + + results = search_wiki_fts_index(query="10-20 结构图", top_k=1, base_dir=tmp_path) + + assert len(results) == 1 + assert results[0].wiki_id == 7 + conn = sqlite3.connect(db_path) + try: + assert conn.execute("SELECT count(*) FROM wiki_chunk_fts").fetchone()[0] == 1 + finally: + conn.close() + + +def test_sync_wiki_vector_index_deduplicates_repeated_page_changes(monkeypatch, tmp_path: Path) -> None: + _write_page( + tmp_path, + "1-hetagen.md", + "HetaGen", + "Structured content generation.", + """ +### Capabilities + +HetaGen supports table synthesis. + +### Query + +HetaGen can answer structured questions. +""", + ) + monkeypatch.setattr("heta.kb.vector_index._embed_texts", lambda texts, config: [[1.0] + [0.0] * 1023 for _ in texts]) + config = HetaConfig( + version=1, + llm=LLMConfig(provider="qwen", api_key="sk-test"), + mineru=MinerUConfig.disabled(), + vector_index=VectorIndexConfig.enabled(), + insert_planning=InsertPlanningConfig.enabled(), + ) + + sync_wiki_vector_index( + changes=[ + FileChange("added", "HetaGen", "pages/1-hetagen.md"), + FileChange("updated", "HetaGen", "pages/1-hetagen.md"), + ], + config=config, + base_dir=tmp_path, + ) + + conn = sqlite3.connect(paths.vector_db_path(tmp_path)) + try: + assert conn.execute("SELECT count(*) FROM wiki_chunks").fetchone()[0] == 2 + finally: + conn.close() + + def _write_page(tmp_path: Path, name: str, title: str, summary: str, content: str) -> Path: page = paths.pages_dir(tmp_path) / name - page.parent.mkdir(parents=True) + page.parent.mkdir(parents=True, exist_ok=True) page.write_text(frontmatter_page(title, "source.md", summary, content), encoding="utf-8") return page