kuhyx
71e4a6115e
Merges ~/ck3-knowledge and ~/kcd-knowledge, which were 90% duplicate code and
neither of which had ever been committed anywhere.
Works on any wiki, not just MediaWiki:
- sources.py splits the engine-specific parts (enumerate, fetch) behind a Source
protocol. MediaWikiSource is preferred whenever an api.php exists;
SitemapHtmlSource handles custom engines with no API (verified on Fextralife).
Everything downstream was already source-agnostic.
- scripts/detect_wiki.py probes /api.php, /w/api.php then /sitemap.xml and prints
the config to use. The API path and article URL vary independently — Fandom
serves /api.php but /wiki/{Title}.
- Profiles are plain env files, since config.py was already fully env-driven.
Per-wiki behaviour that must be measured rather than assumed:
- Subpage policy. On the CK3 wiki "A/B" is a mod; on the KCD wiki it is the
sequel, and rejecting subpages deletes 441 pages of wanted content.
- Noise sections. "Other languages" tables were 7% of KCD chunks and matched
English queries by repeating every title in eight languages.
- Variant tagging for wikis covering two editions, with scopes that exclude the
other edition but keep both/unknown pages.
Build performance: length bucketing in embed_efficiently is 2.7x (fastembed pads
each batch to its longest member; 184-char median vs 2000-char max). With noise
sections that is 61 min -> 20 min on a 6,140-page wiki at half the workers.
Fandom portable-infoboxes are transformed to tables before markdownify —
otherwise stat values survive but detach from their labels.
100% branch coverage, ruff clean.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
111 lines
3.4 KiB
Python
111 lines
3.4 KiB
Python
"""Tests for index building."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import numpy as np
|
|
|
|
from wikikb import config, embed, indexer
|
|
|
|
_DOC_A = """---
|
|
title: "A"
|
|
url: https://x/A
|
|
revid: 1
|
|
categories: []
|
|
---
|
|
|
|
# A
|
|
|
|
Genius trait gives plus five diplomacy.
|
|
"""
|
|
|
|
_DOC_B = """---
|
|
title: "B"
|
|
url: https://x/B
|
|
revid: 2
|
|
categories: []
|
|
---
|
|
|
|
# B
|
|
|
|
Herculean physique gives health.
|
|
"""
|
|
|
|
|
|
def _embed_fn(fake_model):
|
|
return lambda texts: embed.embed_passages(texts, model=fake_model)
|
|
|
|
|
|
def test_collect_chunks_empty(tmp_kb) -> None:
|
|
assert indexer.collect_chunks(config.CORPUS_DIR) == []
|
|
|
|
|
|
def test_build_index_creates_db_and_embeddings(tmp_kb, fake_model) -> None:
|
|
(config.CORPUS_DIR / "a.md").write_text(_DOC_A, encoding="utf-8")
|
|
(config.CORPUS_DIR / "b.md").write_text(_DOC_B, encoding="utf-8")
|
|
count = indexer.build_index(embed_fn=_embed_fn(fake_model), log=lambda _: None)
|
|
assert count >= 2
|
|
assert config.DB_PATH.exists()
|
|
arr = np.load(config.EMBEDDINGS_PATH)
|
|
assert arr.shape == (count, config.EMBED_DIM)
|
|
|
|
|
|
def test_build_index_empty_corpus(tmp_kb, fake_model) -> None:
|
|
count = indexer.build_index(embed_fn=_embed_fn(fake_model), log=lambda _: None)
|
|
assert count == 0
|
|
arr = np.load(config.EMBEDDINGS_PATH)
|
|
assert arr.shape == (0, config.EMBED_DIM)
|
|
|
|
|
|
def test_build_index_embeds_context_augmented_text(tmp_kb, fake_model) -> None:
|
|
(config.CORPUS_DIR / "a.md").write_text(_DOC_A, encoding="utf-8")
|
|
captured: list[str] = []
|
|
|
|
def capture_embed(texts):
|
|
captured.extend(texts)
|
|
return embed.embed_passages(texts, model=fake_model)
|
|
|
|
indexer.build_index(embed_fn=capture_embed, log=lambda _: None)
|
|
assert captured
|
|
# embedded text is prefixed with the page title (context), body stays clean in DB
|
|
assert all(t.startswith("A. ") for t in captured)
|
|
|
|
|
|
def test_embed_efficiently_returns_vectors_in_original_order(fake_model) -> None:
|
|
# Worthless if it corrupts alignment: embedding row i must still be rowid i+1
|
|
# after the internal length sort and dedupe.
|
|
texts = ["a" * 500, "bb", "a" * 500, "ccc", "d" * 900]
|
|
got = indexer.embed_efficiently(texts, _embed_fn(fake_model))
|
|
want = embed.embed_passages(texts, model=fake_model)
|
|
np.testing.assert_allclose(got, want)
|
|
|
|
|
|
def test_embed_efficiently_embeds_each_distinct_text_once(fake_model) -> None:
|
|
seen: list[list[str]] = []
|
|
|
|
def counting_embed(texts):
|
|
seen.append(list(texts))
|
|
return embed.embed_passages(texts, model=fake_model)
|
|
|
|
got = indexer.embed_efficiently(["dup", "dup", "dup", "other"], counting_embed)
|
|
assert seen == [["dup", "other"]] # deduped, then sorted by length
|
|
assert got.shape[0] == 4
|
|
np.testing.assert_allclose(got[0], got[2])
|
|
|
|
|
|
def test_embed_efficiently_empty(fake_model) -> None:
|
|
assert indexer.embed_efficiently([], _embed_fn(fake_model)).shape == (0, config.EMBED_DIM)
|
|
|
|
|
|
def test_noise_sections_are_not_indexed(tmp_kb, fake_model, monkeypatch) -> None:
|
|
monkeypatch.setattr(config, "NOISE_SECTIONS", frozenset({"other languages"}))
|
|
doc = _DOC_A + "\n## Other languages\n\nlocalized UI names table\n"
|
|
(config.CORPUS_DIR / "a.md").write_text(doc, encoding="utf-8")
|
|
chunks = indexer.collect_chunks(config.CORPUS_DIR)
|
|
assert chunks
|
|
assert not any("localized UI names" in c.text for c in chunks)
|
|
|
|
|
|
def test_main_delegates(monkeypatch) -> None:
|
|
monkeypatch.setattr(indexer, "build_index", lambda: 3)
|
|
assert indexer.main() == 0
|