Consolidate CK3/KCD wiki knowledge bases into one generic wiki-kb
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>
This commit is contained in:
commit
71e4a6115e
19
.gitignore
vendored
Normal file
19
.gitignore
vendored
Normal file
@ -0,0 +1,19 @@
|
||||
# Generated / large / regenerable artifacts
|
||||
.venv/
|
||||
corpus/
|
||||
index/
|
||||
manifest.json
|
||||
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
.pytest_cache/
|
||||
.coverage
|
||||
htmlcov/
|
||||
.ruff_cache/
|
||||
|
||||
# fastembed / model cache (if it lands locally)
|
||||
.fastembed_cache/
|
||||
|
||||
# Build artifacts
|
||||
*.egg-info/
|
||||
21
LICENSE
Normal file
21
LICENSE
Normal file
@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2026 Krzysztof Rudnicki
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
18
NOTICE.md
Normal file
18
NOTICE.md
Normal file
@ -0,0 +1,18 @@
|
||||
# NOTICE
|
||||
|
||||
This project creates a **personal, offline** searchable mirror of the
|
||||
[Crusader Kings III Wiki](https://ck3.paradoxwikis.com) for use as a local
|
||||
retrieval/RAG knowledge base.
|
||||
|
||||
- Wiki text content is authored by the CK3 wiki community and is generally
|
||||
available under [CC BY-SA 4.0](https://creativecommons.org/licenses/by-sa/4.0/).
|
||||
Original content and attribution belong to the wiki and its contributors.
|
||||
- Crusader Kings III and related names are trademarks of Paradox Interactive.
|
||||
This project is unofficial and not affiliated with or endorsed by Paradox.
|
||||
- The scraper is intended for **personal, non-redistributive** use. It fetches
|
||||
through the public MediaWiki API with a descriptive User-Agent, `maxlag`,
|
||||
single-flight ~1 req/s rate limiting, and exponential backoff — following
|
||||
[MediaWiki API:Etiquette](https://www.mediawiki.org/wiki/API:Etiquette).
|
||||
- The generated `corpus/` and `index/` are **not** committed to git and are not
|
||||
meant to be republished. Each mirrored page retains a link back to its source
|
||||
URL.
|
||||
137
README.md
Normal file
137
README.md
Normal file
@ -0,0 +1,137 @@
|
||||
# wiki-kb
|
||||
|
||||
Turn a wiki into a local, offline, searchable knowledge base your LLM can query as an
|
||||
MCP tool.
|
||||
|
||||
**scrape** (MediaWiki API, or sitemap + HTML extraction for wikis with no API) →
|
||||
**corpus/** markdown → **index** (SQLite FTS5 + local `bge-small` embeddings) →
|
||||
**hybrid search** (BM25 + cosine, fused with Reciprocal Rank Fusion) → **MCP stdio
|
||||
server**.
|
||||
|
||||
Everything runs locally. No API keys, no cloud, no torch — `fastembed` uses ONNX on
|
||||
CPU.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
./install.sh # venv (Python 3.12) + deps
|
||||
export WIKI_KB_CONTACT="you@example.com" # wiki etiquette requires a real contact
|
||||
|
||||
python scripts/detect_wiki.py https://some.wiki.example.com # what kind of wiki is it?
|
||||
./run.sh --profile ck3 scrape # mirror (resumable, ~1 req/s)
|
||||
./run.sh --profile ck3 index # build the search index
|
||||
./run.sh --profile ck3 search "domain limit by rank"
|
||||
./run.sh check # ruff + pytest (100% branch coverage)
|
||||
```
|
||||
|
||||
Register with Claude Code:
|
||||
|
||||
```bash
|
||||
WIKI_KB_NAMESPACE=ck3 claude mcp add -s user ck3 -- "$PWD/.venv/bin/python" -m wikikb.server
|
||||
```
|
||||
|
||||
## Any wiki, not just MediaWiki
|
||||
|
||||
`scripts/detect_wiki.py <host>` probes `/api.php`, then `/w/api.php`, then
|
||||
`/sitemap.xml`, and prints the source kind, endpoint, article URL pattern, engine,
|
||||
page count, licence, and the env vars to configure with. Verified against three
|
||||
engines:
|
||||
|
||||
| Host | Source | Endpoint | Article URL |
|
||||
|---|---|---|---|
|
||||
| ck3.paradoxwikis.com | mediawiki | `/api.php` | `{host}/{title}` |
|
||||
| kingdom-come-deliverance.fandom.com | mediawiki | `/api.php` | `{host}/wiki/{title}` |
|
||||
| kingdomcomedeliverance2.wiki.fextralife.com | **html** | `/sitemap.xml` | `{host}/{Title+With+Plus}` |
|
||||
|
||||
Note the API path and the article URL vary *independently* — Fandom serves the API at
|
||||
`/api.php` but articles at `/wiki/{Title}`. Hardcoding one layout breaks on the next
|
||||
wiki.
|
||||
|
||||
An API is always preferred: it gives a real page list, server-rendered HTML with
|
||||
templates expanded, categories, revision ids, and `maxlag` so the server can push
|
||||
back. The HTML fallback genuinely gives up all of those — `detect_wiki.py` says so
|
||||
rather than pretending otherwise.
|
||||
|
||||
## Profiles
|
||||
|
||||
`config.py` reads every setting from `WIKI_KB_*` environment variables, so a profile
|
||||
is just an env file in `profiles/` — no plugin system. Each sets its own
|
||||
`WIKI_KB_ROOT`, which is what keeps several wikis' corpora and indexes apart.
|
||||
|
||||
See `profiles/README.md`. Shipped: `ck3`, `kcd`, `kcd-fextralife`.
|
||||
|
||||
## Things that are per-wiki and must be measured, not assumed
|
||||
|
||||
- **Are subpages noise?** On the CK3 wiki `A/B` means a total-conversion mod (reject).
|
||||
On the KCD wiki it means *the sequel* — `Alchemy/KCD2` — and rejecting subpages
|
||||
deletes 441 pages of exactly the content you want (keep). One API pass decides it:
|
||||
`Counter(t.rsplit("/", 1)[-1] for t in titles if "/" in t)`.
|
||||
- **Which sections are boilerplate?** `WIKI_KB_NOISE_SECTIONS` drops them from the
|
||||
index while `*_get_page` still serves the whole page. On the KCD wiki the "Other
|
||||
languages" tables were 7% of chunks and matched English queries because they repeat
|
||||
every page's title in eight languages.
|
||||
- **Does the wiki cover two editions at once?** `WIKI_KB_VARIANTS="kcd2:KCD2,kcd1:KCD"`
|
||||
tags every page and lets searches scope to one. A scope excludes the *other* edition
|
||||
but always keeps pages tagged `both`/`unknown` — filtering *to* one tag is a recall
|
||||
hole, since ~9% of KCD pages carry no edition category yet matter to both.
|
||||
|
||||
Every search hit carries its `source` and `variant`, so a result can never be silently
|
||||
attributed to the wrong wiki or the wrong edition.
|
||||
|
||||
## Build performance
|
||||
|
||||
The index build is CPU-heavy. Two optimisations matter, and the first is large:
|
||||
|
||||
- **Length bucketing** (`indexer.embed_efficiently`) — fastembed pads every batch to
|
||||
its longest member. Wiki chunks are wildly uneven (184-char median, 2000-char max,
|
||||
53% under 200 chars), so most of the compute was padding. Sorting by length measured
|
||||
**2.7x faster**. The sort is undone before the vectors are returned: embedding row
|
||||
`i` must stay aligned with chunk rowid `i+1`.
|
||||
- **Noise sections** cut ~17% of chunks outright.
|
||||
|
||||
Together: 61 min → 20 min on a 6,140-page wiki, at *half* the worker count.
|
||||
|
||||
**`nice` does not protect a running game.** Measured with a GPU-bound game: an
|
||||
8-worker build dropped GPU utilisation from 96.6% to ~50% while the game's *total* CPU
|
||||
barely moved — it was never denied cycles, its render thread was delayed getting them.
|
||||
Core reservation did not measurably help either. The real lever is finishing sooner.
|
||||
|
||||
```bash
|
||||
WIKI_KB_RESERVED_CORES="0-3" # cores to keep free (set "" to use all)
|
||||
WIKI_KB_NICE=10
|
||||
WIKI_KB_EMBED_PARALLEL=4 # 4 workers ~= 20 min on 6k pages; 8 is worse for games
|
||||
```
|
||||
|
||||
`build_index` writes only at the end and is not resumable — make sure a build fits in
|
||||
the time you have.
|
||||
|
||||
## Etiquette and licensing
|
||||
|
||||
- Descriptive User-Agent with a real contact (`WIKI_KB_CONTACT`), `maxlag`, ~1 req/s
|
||||
single-flight, exponential backoff, resumable. Per
|
||||
[MediaWiki API:Etiquette](https://www.mediawiki.org/wiki/API:Etiquette).
|
||||
- Cloudflare bot challenges arrive as **HTTP 200 with an HTML body**, so they pass
|
||||
`raise_for_status()`. They are treated as retryable. The User-Agent is never spoofed
|
||||
to evade one.
|
||||
- `corpus/`, `index/` and `manifest.json` are gitignored. Mirrored content belongs to
|
||||
the wiki and its contributors — most are CC BY-SA. Each mirrored page keeps a link
|
||||
back to its source URL. Check a site's `robots.txt` and terms before mirroring it;
|
||||
`detect_wiki.py` surfaces AI-crawler directives it finds.
|
||||
|
||||
## Layout
|
||||
|
||||
```
|
||||
wikikb/
|
||||
config.py endpoints, paths, noise filters, variants, CPU settings — all env-driven
|
||||
sources.py MediaWikiSource | SitemapHtmlSource + host detection
|
||||
scraper.py politeness, noise filtering, resume, corpus + manifest
|
||||
html2md.py parse-HTML -> markdown (incl. Fandom portable-infobox -> tables)
|
||||
chunker.py heading-scoped, overlapping windows
|
||||
embed.py fastembed bge-small, L2-normalised
|
||||
indexer.py FTS5 table + row-aligned embeddings.npy
|
||||
search.py BM25 (OR terms) + cosine -> RRF fusion
|
||||
variant.py edition tagging for wikis covering two editions
|
||||
server.py FastMCP stdio: <ns>_search / _get_page / _list_pages
|
||||
profiles/ one env file per wiki
|
||||
scripts/ detect_wiki.py + the three acceptance checks
|
||||
```
|
||||
50
install.sh
Executable file
50
install.sh
Executable file
@ -0,0 +1,50 @@
|
||||
#!/bin/bash
|
||||
|
||||
# ============================================================================
|
||||
# install.sh — create an isolated venv (Python 3.12) and install wiki-kb
|
||||
# ============================================================================
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
readonly VENV_DIR="${SCRIPT_DIR}/.venv"
|
||||
|
||||
# Python 3.12 is pinned because fastembed/onnxruntime wheels lag on newer
|
||||
# interpreters. Prefer pyenv's 3.12, fall back to any python3.12 on PATH.
|
||||
find_python() {
|
||||
local candidates=(
|
||||
"${HOME}/.pyenv/versions/3.12.8/bin/python3.12"
|
||||
"$(command -v python3.12 || true)"
|
||||
)
|
||||
local p
|
||||
for p in "${candidates[@]}"; do
|
||||
if [[ -n "$p" && -x "$p" ]]; then
|
||||
printf '%s\n' "$p"
|
||||
return 0
|
||||
fi
|
||||
done
|
||||
echo "Error: python3.12 not found (need it for fastembed/onnxruntime wheels)." >&2
|
||||
echo "Install it, e.g.: pyenv install 3.12.8" >&2
|
||||
return 1
|
||||
}
|
||||
|
||||
main() {
|
||||
local py
|
||||
py="$(find_python)"
|
||||
echo "Using interpreter: $py ($("$py" --version))"
|
||||
|
||||
if [[ ! -d "$VENV_DIR" ]]; then
|
||||
echo "Creating venv at $VENV_DIR"
|
||||
"$py" -m venv "$VENV_DIR"
|
||||
fi
|
||||
|
||||
# shellcheck disable=SC1091
|
||||
source "${VENV_DIR}/bin/activate"
|
||||
python -m pip install --quiet --upgrade pip
|
||||
echo "Installing wiki-kb (+dev extras)..."
|
||||
python -m pip install --quiet -e "${SCRIPT_DIR}[dev]"
|
||||
|
||||
echo "Done. Activate with: source ${VENV_DIR}/bin/activate"
|
||||
}
|
||||
|
||||
main "$@"
|
||||
16
profiles/README.md
Normal file
16
profiles/README.md
Normal file
@ -0,0 +1,16 @@
|
||||
# Wiki profiles
|
||||
|
||||
`config.py` reads every setting from `WIKI_KB_*` environment variables, so a profile
|
||||
is just an env file — no plugin system, no registry, no code.
|
||||
|
||||
```bash
|
||||
./run.sh --profile kcd scrape
|
||||
./run.sh --profile kcd search "best longsword stats"
|
||||
```
|
||||
|
||||
Create a new one by running `python scripts/detect_wiki.py <host>`, which prints the
|
||||
host/API/page-URL lines to paste in, then adding the per-wiki noise filters you find
|
||||
by inspecting the corpus (see the two examples here).
|
||||
|
||||
Each profile MUST set `WIKI_KB_ROOT` to its own directory — that is what keeps several
|
||||
wikis' corpora, indexes and manifests from overwriting each other.
|
||||
10
profiles/ck3.env
Normal file
10
profiles/ck3.env
Normal file
@ -0,0 +1,10 @@
|
||||
# Crusader Kings III wiki — MediaWiki, articles at {host}/{Title}.
|
||||
WIKI_KB_HOST=https://ck3.paradoxwikis.com
|
||||
WIKI_KB_NAMESPACE=ck3
|
||||
WIKI_KB_ROOT=$HOME/.local/share/wiki-kb/ck3
|
||||
|
||||
# Total-conversion mods live in ns0 and swamp retrieval (~33% of the corpus).
|
||||
# On THIS wiki they are subpages, so rejecting "/" titles is correct — it is NOT
|
||||
# correct in general; see profiles/kcd.env for a wiki where "/" means the sequel.
|
||||
WIKI_KB_SUBPAGES=reject
|
||||
WIKI_KB_MOD_PREFIXES=ck3agot,lotr,elder kings,after the end,the fallen eagle,fallen eagle,godherja,princes of darkness,rice,realms in exile,a game of thrones,when the world stopped making sense,more bookmarks
|
||||
10
profiles/kcd-fextralife.env
Normal file
10
profiles/kcd-fextralife.env
Normal file
@ -0,0 +1,10 @@
|
||||
# Fextralife KCD2 wiki — NO API (custom engine). Sitemap + HTML content extraction.
|
||||
# Loses revision ids, maxlag and categories; see references/build-guide.md.
|
||||
# Check the site's robots.txt and terms before enabling this one.
|
||||
WIKI_KB_HOST=https://kingdomcomedeliverance2.wiki.fextralife.com
|
||||
WIKI_KB_SOURCE=html
|
||||
WIKI_KB_SITEMAP=https://kingdomcomedeliverance2.wiki.fextralife.com/sitemap.xml
|
||||
WIKI_KB_NAMESPACE=kcd
|
||||
WIKI_KB_ROOT=$HOME/.local/share/wiki-kb/kcd
|
||||
WIKI_KB_SUBPAGES=keep
|
||||
WIKI_KB_MOD_PREFIXES=
|
||||
22
profiles/kcd.env
Normal file
22
profiles/kcd.env
Normal file
@ -0,0 +1,22 @@
|
||||
# Kingdom Come: Deliverance wiki (Fandom) — MediaWiki, articles at {host}/wiki/{Title}.
|
||||
# Covers BOTH games in one namespace; the sequel's pages are "/KCD2" subpages.
|
||||
WIKI_KB_HOST=https://kingdom-come-deliverance.fandom.com
|
||||
WIKI_KB_API_URL=https://kingdom-come-deliverance.fandom.com/api.php
|
||||
WIKI_KB_PAGE_URL=https://kingdom-come-deliverance.fandom.com/wiki/{title}
|
||||
WIKI_KB_NAMESPACE=kcd
|
||||
WIKI_KB_ROOT=$HOME/.local/share/wiki-kb/kcd
|
||||
|
||||
# Subpages are STRUCTURAL here ("Alchemy/KCD2" is the sequel's Alchemy page), so
|
||||
# rejecting them outright would delete all 441 KCD2 pages. Drop only "/loc", which
|
||||
# holds the same in-game book text repeated in ~8 languages (371 pages, ~32k chars each).
|
||||
WIKI_KB_SUBPAGES=keep
|
||||
WIKI_KB_NOISE_TAILS=loc
|
||||
WIKI_KB_MOD_PREFIXES=
|
||||
|
||||
# Per-page tables of localized UI names: 7% of chunks, and they match English queries
|
||||
# because they repeat every page's title in every language.
|
||||
WIKI_KB_NOISE_SECTIONS=other languages
|
||||
|
||||
# Two-edition wiki: tag each page kcd1/kcd2/both/unknown and default searches to kcd2.
|
||||
WIKI_KB_VARIANTS=kcd2:KCD2,kcd1:KCD
|
||||
WIKI_KB_DEFAULT_VARIANT=kcd2
|
||||
48
pyproject.toml
Normal file
48
pyproject.toml
Normal file
@ -0,0 +1,48 @@
|
||||
[project]
|
||||
name = "wiki-kb"
|
||||
version = "0.1.0"
|
||||
description = "Local searchable mirror of the Crusader Kings III wiki (RAG + MCP)."
|
||||
requires-python = ">=3.12,<3.13"
|
||||
dependencies = [
|
||||
"httpx>=0.27",
|
||||
"beautifulsoup4>=4.12",
|
||||
"markdownify>=0.13",
|
||||
"numpy>=1.26",
|
||||
"fastembed>=0.3",
|
||||
"mcp>=1.2",
|
||||
]
|
||||
|
||||
[project.optional-dependencies]
|
||||
dev = [
|
||||
"pytest>=8.0",
|
||||
"pytest-cov>=5.0",
|
||||
"respx>=0.21",
|
||||
"ruff>=0.6",
|
||||
]
|
||||
|
||||
[build-system]
|
||||
requires = ["setuptools>=68"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools]
|
||||
packages = ["wikikb"]
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
target-version = "py312"
|
||||
|
||||
[tool.ruff.lint]
|
||||
select = ["E", "F", "I", "UP", "B", "SIM", "PTH", "RUF"]
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
addopts = "-q"
|
||||
testpaths = ["tests"]
|
||||
|
||||
[tool.coverage.run]
|
||||
branch = true
|
||||
source = ["wikikb"]
|
||||
|
||||
[tool.coverage.report]
|
||||
show_missing = true
|
||||
fail_under = 100
|
||||
exclude_lines = ["if __name__ == .__main__.:", "pragma: no cover", "\\.\\.\\."]
|
||||
78
run.sh
Executable file
78
run.sh
Executable file
@ -0,0 +1,78 @@
|
||||
#!/bin/bash
|
||||
|
||||
# ============================================================================
|
||||
# run.sh — orchestrate a wiki knowledge base
|
||||
# ./run.sh install create venv + install deps
|
||||
# ./run.sh scrape [--limit N] mirror the wiki into corpus/
|
||||
# ./run.sh index build FTS5 + embeddings index
|
||||
# ./run.sh search "<query>" hybrid search from the CLI
|
||||
# ./run.sh serve run the MCP server (stdio)
|
||||
# ./run.sh check ruff + pytest (with coverage)
|
||||
# ============================================================================
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
readonly SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
readonly VENV_PY="${SCRIPT_DIR}/.venv/bin/python"
|
||||
|
||||
need_venv() {
|
||||
if [[ ! -x "$VENV_PY" ]]; then
|
||||
echo "venv missing — run: ./run.sh install" >&2
|
||||
exit 1
|
||||
fi
|
||||
}
|
||||
|
||||
usage() {
|
||||
grep -E '^# \./run\.sh' "${BASH_SOURCE[0]}" | sed 's/^# //'
|
||||
exit "${1:-0}"
|
||||
}
|
||||
|
||||
# A profile is just an env file (see profiles/README.md) — config.py reads every
|
||||
# setting from WIKI_KB_* variables, so no plugin machinery is needed.
|
||||
# A profile is just an env file (see profiles/README.md) — config.py reads every
|
||||
# setting from WIKI_KB_* variables, so no plugin machinery is needed. Parsed line by
|
||||
# line rather than sourced: values legitimately contain spaces ("other languages"),
|
||||
# which `source` would try to execute.
|
||||
load_profile() {
|
||||
local name="$1" file="${SCRIPT_DIR}/profiles/${1}.env" line key value
|
||||
[[ -f "$file" ]] || {
|
||||
echo "No such profile: $name (available: $(cd "${SCRIPT_DIR}/profiles" && \
|
||||
ls *.env 2>/dev/null | sed 's/\.env$//' | tr '\n' ' '))" >&2
|
||||
exit 1
|
||||
}
|
||||
while IFS= read -r line || [[ -n "$line" ]]; do
|
||||
[[ -z "$line" || "$line" == \#* ]] && continue
|
||||
key="${line%%=*}"
|
||||
value="${line#*=}"
|
||||
value="${value//\$HOME/$HOME}" # only $HOME is expanded, deliberately
|
||||
export "$key=$value"
|
||||
done < "$file"
|
||||
echo "profile: $name -> ${WIKI_KB_HOST:-?} (root ${WIKI_KB_ROOT:-?})" >&2
|
||||
}
|
||||
|
||||
main() {
|
||||
# --profile <name> may precede the command; it only sets env vars.
|
||||
if [[ "${1:-}" == "--profile" ]]; then
|
||||
[[ $# -ge 2 ]] || { echo "--profile needs a name" >&2; exit 1; }
|
||||
load_profile "$2"
|
||||
shift 2
|
||||
fi
|
||||
local cmd="${1:-}"
|
||||
[[ $# -gt 0 ]] && shift || true
|
||||
case "$cmd" in
|
||||
install) exec "${SCRIPT_DIR}/install.sh" ;;
|
||||
scrape) need_venv; exec "$VENV_PY" -u -m wikikb.scraper "$@" ;;
|
||||
index) need_venv; exec "$VENV_PY" -u -m wikikb.indexer "$@" ;;
|
||||
search) need_venv; exec "$VENV_PY" -u -m wikikb.search "$@" ;;
|
||||
serve) need_venv; exec "$VENV_PY" -u -m wikikb.server "$@" ;;
|
||||
check)
|
||||
need_venv
|
||||
"$VENV_PY" -m ruff check "${SCRIPT_DIR}/wikikb" "${SCRIPT_DIR}/tests"
|
||||
"$VENV_PY" -m pytest --cov=wikikb --cov-report=term-missing "${SCRIPT_DIR}/tests"
|
||||
;;
|
||||
""|-h|--help) usage 0 ;;
|
||||
*) echo "Unknown command: $cmd" >&2; usage 1 ;;
|
||||
esac
|
||||
}
|
||||
|
||||
main "$@"
|
||||
76
scripts/detect_wiki.py
Normal file
76
scripts/detect_wiki.py
Normal file
@ -0,0 +1,76 @@
|
||||
#!/usr/bin/env python
|
||||
"""Probe a wiki host and report how (or whether) it can be mirrored.
|
||||
|
||||
Phase 0 of the build, as a command:
|
||||
|
||||
python scripts/detect_wiki.py https://some.wiki.example.com
|
||||
|
||||
Reports the source kind, the endpoint, the article URL pattern, and — for
|
||||
MediaWiki — the engine version and page count. Exits non-zero if the host cannot
|
||||
be mirrored at all, so it can gate a scaffold.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import httpx
|
||||
|
||||
from wikikb import config
|
||||
from wikikb.sources import MEDIAWIKI, SitemapHtmlSource, detect_source_kind
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = sys.argv[1:] if argv is None else argv
|
||||
if not args:
|
||||
print("usage: detect_wiki.py <https://host>", file=sys.stderr)
|
||||
return 2
|
||||
host = args[0].rstrip("/")
|
||||
|
||||
with httpx.Client(
|
||||
headers={"User-Agent": config.USER_AGENT}, timeout=30, follow_redirects=True
|
||||
) as client:
|
||||
try:
|
||||
kind, endpoint = detect_source_kind(client.get, host)
|
||||
except RuntimeError as exc:
|
||||
print(f"[FAIL] {exc}")
|
||||
return 1
|
||||
|
||||
print(f"[ OK ] source kind : {kind}")
|
||||
print(f" endpoint : {endpoint}")
|
||||
|
||||
if kind == MEDIAWIKI:
|
||||
info = client.get(
|
||||
f"{endpoint}?action=query&meta=siteinfo&siprop=general|statistics&format=json"
|
||||
).json()["query"]
|
||||
general, stats = info["general"], info["statistics"]
|
||||
base = general.get("base", "")
|
||||
pattern = "{host}/wiki/{title}" if "/wiki/" in base else "{host}/{title}"
|
||||
print(f" engine : {general.get('generator')}")
|
||||
print(f" content : {stats.get('articles')} articles")
|
||||
print(f" article URL : {pattern}")
|
||||
print(f" licence : {info.get('rightsinfo', {}).get('text', 'unknown')}")
|
||||
print("\nScaffold with:")
|
||||
print(f" WIKI_KB_HOST={host} WIKI_KB_API_URL={endpoint}")
|
||||
if pattern.endswith("/wiki/{title}"):
|
||||
print(f" WIKI_KB_PAGE_URL='{host}/wiki/{{title}}'")
|
||||
else:
|
||||
urls = list(SitemapHtmlSource(client.get, sitemap_url=endpoint).iter_titles())
|
||||
print(f" content : {len(urls)} sitemap URLs")
|
||||
print(f" sample : {urls[:5]}")
|
||||
print("\n[WARN] No API. Revision ids and maxlag are unavailable, so fidelity")
|
||||
print(" checks fall back to content diffing and rate limiting is purely")
|
||||
print(" self-imposed. Check the site's robots.txt and terms before use.")
|
||||
print("\nScaffold with:")
|
||||
print(f" WIKI_KB_HOST={host} WIKI_KB_SOURCE=html WIKI_KB_SITEMAP={endpoint}")
|
||||
|
||||
robots = client.get(f"{host}/robots.txt")
|
||||
if robots.status_code == 200 and "gptbot" in robots.text.lower():
|
||||
print("\n[NOTE] robots.txt names an AI crawler (e.g. GPTBot). It does not name")
|
||||
print(" this tool, but it is the operators' clearest statement about AI")
|
||||
print(" ingestion of their content. Your call — read it before scraping.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
75
scripts/verify_fidelity.py
Normal file
75
scripts/verify_fidelity.py
Normal file
@ -0,0 +1,75 @@
|
||||
#!/usr/bin/env python
|
||||
"""Verify our local mirror matches the LIVE wiki (no corruption / staleness).
|
||||
|
||||
For each key page it re-fetches the live page, runs it through the SAME
|
||||
HTML->markdown pipeline used at scrape time, and checks:
|
||||
1. the freshly-converted live body is byte-identical to our stored copy
|
||||
(proves faithful capture + no drift since the scrape), and
|
||||
2. specific factual tokens (trait names, modifier numbers) appear in BOTH the
|
||||
live page and our local copy — printed side by side so "the same" is visible.
|
||||
|
||||
DONE = every page's key facts match in live & local (byte-identical is reported
|
||||
too, and flagged as WARN if the wiki was edited since the scrape).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import difflib
|
||||
import sys
|
||||
|
||||
import httpx
|
||||
|
||||
from wikikb import config, html2md
|
||||
from wikikb.scraper import Throttle, fetch_page, load_manifest, sanitize_filename
|
||||
|
||||
# page -> factual tokens that must appear identically in live + local
|
||||
FACTS: dict[str, list[str]] = {
|
||||
"Traits": ["Genius", "Intelligent", "Herculean", "Beautiful", "+5"],
|
||||
"Attributes": ["Diplomacy", "Stewardship", "Prowess", "Domain Taxes"],
|
||||
"Succession": ["Primogeniture", "Partition"],
|
||||
}
|
||||
|
||||
|
||||
def _stored_text(title: str) -> str:
|
||||
entry = load_manifest().get(title, {})
|
||||
filename = entry.get("filename") or sanitize_filename(title)
|
||||
return (config.CORPUS_DIR / filename).read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def _line_with(text: str, token: str) -> str:
|
||||
return next((ln.strip() for ln in text.splitlines() if token in ln), "")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ok = True
|
||||
with httpx.Client(headers={"User-Agent": config.USER_AGENT}, timeout=30) as client:
|
||||
throttle = Throttle()
|
||||
for title, facts in FACTS.items():
|
||||
page = fetch_page(client, title, throttle=throttle)
|
||||
if page is None:
|
||||
print(f"[FAIL] {title}: could not fetch live page")
|
||||
ok = False
|
||||
continue
|
||||
stored = _stored_text(page.title)
|
||||
fresh = html2md.html_to_markdown(page.html)
|
||||
|
||||
identical = fresh in stored # stored = frontmatter + "# Title" + this body
|
||||
ratio = difflib.SequenceMatcher(None, fresh, stored).ratio()
|
||||
tag = "PASS" if identical else "WARN"
|
||||
note = "byte-identical to stored copy" if identical else f"differs (ratio {ratio:.3f})"
|
||||
print(f"[{tag}] {page.title} (rev {page.revid}): freshly-converted live body {note}")
|
||||
|
||||
for fact in facts:
|
||||
same = fact in fresh and fact in stored
|
||||
ok &= same
|
||||
print(
|
||||
f" [{'PASS' if same else 'FAIL'}] {fact!r} in live & local -> {same}"
|
||||
f" | live: {_line_with(fresh, fact)[:80]!r}"
|
||||
)
|
||||
|
||||
print(f"\nFIDELITY VERIFICATION: {'PASS' if ok else 'FAIL'}")
|
||||
return 0 if ok else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
68
scripts/verify_mcp.py
Normal file
68
scripts/verify_mcp.py
Normal file
@ -0,0 +1,68 @@
|
||||
#!/usr/bin/env python
|
||||
"""Verify the ck3 MCP server end-to-end over real stdio JSON-RPC.
|
||||
|
||||
DONE = the server starts, completes the MCP handshake, advertises the three
|
||||
tools, and returns sensible results for wiki_search and wiki_get_page. Exits
|
||||
non-zero on any failure. Uses the same interface Claude Code will use, so this
|
||||
confirms the server works without needing a session reload.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import anyio
|
||||
from mcp import ClientSession, StdioServerParameters
|
||||
from mcp.client.stdio import stdio_client
|
||||
|
||||
|
||||
def _text(result: object) -> str:
|
||||
"""Flatten a CallToolResult into a searchable string."""
|
||||
parts: list[str] = []
|
||||
structured = getattr(result, "structuredContent", None)
|
||||
if structured is not None:
|
||||
parts.append(str(structured))
|
||||
for block in getattr(result, "content", []) or []:
|
||||
if hasattr(block, "text"):
|
||||
parts.append(block.text)
|
||||
return " ".join(parts)
|
||||
|
||||
|
||||
async def run() -> int:
|
||||
params = StdioServerParameters(command=sys.executable, args=["-m", "wikikb.server"])
|
||||
ok = True
|
||||
async with stdio_client(params) as (read, write), ClientSession(read, write) as session:
|
||||
await session.initialize()
|
||||
|
||||
tools = {t.name for t in (await session.list_tools()).tools}
|
||||
want = {"wiki_search", "wiki_get_page", "wiki_list_pages"}
|
||||
t_ok = want <= tools
|
||||
ok &= t_ok
|
||||
print(f"[{'PASS' if t_ok else 'FAIL'}] tools/list -> {sorted(tools)}")
|
||||
|
||||
search_text = _text(await session.call_tool("wiki_search", {"query": "domain limit rank"}))
|
||||
s_ok = (
|
||||
any(k in search_text for k in ("Titles", "Attributes"))
|
||||
and "domain" in search_text.lower()
|
||||
)
|
||||
ok &= s_ok
|
||||
print(f"[{'PASS' if s_ok else 'FAIL'}] wiki_search -> {search_text[:80]!r}...")
|
||||
|
||||
page_text = _text(await session.call_tool("wiki_get_page", {"title": "Traits"}))
|
||||
g_ok = "Genius" in page_text and "Herculean" in page_text
|
||||
ok &= g_ok
|
||||
print(
|
||||
f"[{'PASS' if g_ok else 'FAIL'}] wiki_get_page('Traits') has Genius+Herculean -> {g_ok}"
|
||||
)
|
||||
|
||||
list_text = _text(await session.call_tool("wiki_list_pages", {"prefix": "Tr"}))
|
||||
l_ok = "Traits" in list_text
|
||||
ok &= l_ok
|
||||
print(f"[{'PASS' if l_ok else 'FAIL'}] wiki_list_pages('Tr') includes Traits -> {l_ok}")
|
||||
|
||||
print(f"\nMCP VERIFICATION: {'PASS' if ok else 'FAIL'}")
|
||||
return 0 if ok else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(anyio.run(run))
|
||||
58
scripts/verify_search.py
Normal file
58
scripts/verify_search.py
Normal file
@ -0,0 +1,58 @@
|
||||
#!/usr/bin/env python
|
||||
"""Verify search quality against concrete acceptance criteria.
|
||||
|
||||
DONE = every case's expected wiki page appears in the top-3 results, AND no
|
||||
mod/subpage title appears in the top-5 of any probe query. Exits non-zero on
|
||||
any failure so it can gate a build.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
from wikikb.scraper import is_core_title
|
||||
from wikikb.search import search
|
||||
|
||||
# (query, one-of-these titles must be in the top 3)
|
||||
CASES: list[tuple[str, set[str]]] = [
|
||||
("domain limit rank", {"Titles", "Attributes", "Domain"}),
|
||||
("genius intelligent trait skill bonus", {"Traits"}),
|
||||
("primogeniture succession heir inheritance", {"Succession"}),
|
||||
("congenital beauty attraction opinion", {"Traits"}),
|
||||
("health aging death lifespan", {"Attributes", "Health"}),
|
||||
]
|
||||
|
||||
# queries whose top-5 must contain NO mod/subpage results
|
||||
NO_MOD_QUERIES = [
|
||||
"domain limit",
|
||||
"succession law",
|
||||
"development growth",
|
||||
"knight martial army",
|
||||
"culture names",
|
||||
]
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ok = True
|
||||
print("== relevance (expected page in top-3) ==")
|
||||
for query, expected in CASES:
|
||||
top3 = [r.title for r in search(query, k=5)[:3]]
|
||||
hit = any(t in expected for t in top3)
|
||||
ok &= hit
|
||||
print(
|
||||
f" [{'PASS' if hit else 'FAIL'}] {query!r} -> {top3} (want one of {sorted(expected)})"
|
||||
)
|
||||
|
||||
print("== no mod/subpage pollution (top-5) ==")
|
||||
for query in NO_MOD_QUERIES:
|
||||
mods = [r.title for r in search(query, k=5) if not is_core_title(r.title)]
|
||||
clean = not mods
|
||||
ok &= clean
|
||||
print(f" [{'PASS' if clean else 'FAIL'}] {query!r} -> {'clean' if clean else mods}")
|
||||
|
||||
print(f"\nSEARCH VERIFICATION: {'PASS' if ok else 'FAIL'}")
|
||||
return 0 if ok else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
47
tests/conftest.py
Normal file
47
tests/conftest.py
Normal file
@ -0,0 +1,47 @@
|
||||
"""Shared fixtures: temp KB paths and a deterministic fake embedding model."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
from collections.abc import Iterable
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from wikikb import config
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tmp_kb(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
|
||||
"""Point all config paths at an isolated temp tree."""
|
||||
corpus = tmp_path / "corpus"
|
||||
index = tmp_path / "index"
|
||||
corpus.mkdir()
|
||||
index.mkdir()
|
||||
monkeypatch.setattr(config, "CORPUS_DIR", corpus)
|
||||
monkeypatch.setattr(config, "INDEX_DIR", index)
|
||||
monkeypatch.setattr(config, "DB_PATH", index / "wiki.db")
|
||||
monkeypatch.setattr(config, "EMBEDDINGS_PATH", index / "embeddings.npy")
|
||||
monkeypatch.setattr(config, "MANIFEST_PATH", tmp_path / "manifest.json")
|
||||
return tmp_path
|
||||
|
||||
|
||||
class FakeEmbedding:
|
||||
"""Deterministic, offline stand-in for fastembed's TextEmbedding."""
|
||||
|
||||
def _vec(self, text: str) -> np.ndarray:
|
||||
digest = hashlib.sha256(text.encode()).digest()
|
||||
raw = (digest * (config.EMBED_DIM // len(digest) + 1))[: config.EMBED_DIM]
|
||||
return np.frombuffer(raw, dtype=np.uint8).astype(np.float32)
|
||||
|
||||
def embed(self, texts: Iterable[str], **_kwargs: object) -> list[np.ndarray]:
|
||||
return [self._vec(t) for t in texts]
|
||||
|
||||
def query_embed(self, texts: Iterable[str], **_kwargs: object) -> list[np.ndarray]:
|
||||
return [self._vec(t) for t in texts]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_model() -> FakeEmbedding:
|
||||
return FakeEmbedding()
|
||||
83
tests/test_chunker.py
Normal file
83
tests/test_chunker.py
Normal file
@ -0,0 +1,83 @@
|
||||
"""Tests for markdown chunking."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from wikikb import chunker
|
||||
|
||||
_DOC = """---
|
||||
title: "Traits"
|
||||
url: https://ck3.paradoxwikis.com/Traits
|
||||
revid: 1
|
||||
categories: ["Character"]
|
||||
---
|
||||
|
||||
# Traits
|
||||
|
||||
Intro paragraph.
|
||||
|
||||
## Congenital
|
||||
|
||||
Genius gives +5 to all skills.
|
||||
|
||||
## Physique
|
||||
|
||||
Herculean gives +1 health.
|
||||
"""
|
||||
|
||||
|
||||
def test_parse_frontmatter_extracts_meta_and_body() -> None:
|
||||
meta, body = chunker.parse_frontmatter(_DOC)
|
||||
assert meta["title"] == "Traits"
|
||||
assert meta["url"].endswith("/Traits")
|
||||
assert body.lstrip().startswith("# Traits")
|
||||
|
||||
|
||||
def test_parse_frontmatter_absent() -> None:
|
||||
meta, body = chunker.parse_frontmatter("no frontmatter here")
|
||||
assert meta == {}
|
||||
assert body == "no frontmatter here"
|
||||
|
||||
|
||||
def test_chunk_document_sections() -> None:
|
||||
chunks = chunker.chunk_document(_DOC)
|
||||
sections = {c.section for c in chunks}
|
||||
assert {"Traits", "Congenital", "Physique"} <= sections
|
||||
assert all(c.title == "Traits" for c in chunks)
|
||||
assert any("Genius" in c.text for c in chunks)
|
||||
|
||||
|
||||
def test_window_splits_long_sections_with_overlap() -> None:
|
||||
long = "# T\n\n" + ("word " * 2000)
|
||||
chunks = chunker.chunk_document(long, target_chars=200, overlap_chars=50)
|
||||
assert len(chunks) > 1
|
||||
assert all(len(c.text) <= 200 for c in chunks)
|
||||
|
||||
|
||||
def test_window_short_text_single_chunk() -> None:
|
||||
assert chunker._window("short", size=100, overlap=10) == ["short"]
|
||||
|
||||
|
||||
def test_parse_frontmatter_ignores_lines_without_colon() -> None:
|
||||
text = '---\ntitle: "X"\nnocolonhere\n---\n\nbody'
|
||||
meta, _ = chunker.parse_frontmatter(text)
|
||||
assert meta == {"title": "X"}
|
||||
|
||||
|
||||
def test_split_sections_drops_empty_text_sections() -> None:
|
||||
# "H1" has only blank lines under it -> filtered out; "Overview" and "H2" kept.
|
||||
body = "text0\n# H1\n\n\n# H2\nx"
|
||||
sections = chunker._split_sections(body)
|
||||
headings = [h for h, _ in sections]
|
||||
assert headings == ["Overview", "H2"]
|
||||
|
||||
|
||||
def test_split_sections_body_ending_with_heading() -> None:
|
||||
# trailing heading leaves the buffer empty at loop end (no final flush)
|
||||
assert chunker._split_sections("text\n# H") == [("Overview", "text")]
|
||||
|
||||
|
||||
def test_chunk_document_skips_whitespace_only_windows() -> None:
|
||||
doc = "# H\n\nA" + " " * 300 + "B"
|
||||
chunks = chunker.chunk_document(doc, target_chars=100, overlap_chars=10)
|
||||
assert chunks # some windows are all-whitespace and are skipped
|
||||
assert all(c.text.strip() for c in chunks)
|
||||
53
tests/test_cpu.py
Normal file
53
tests/test_cpu.py
Normal file
@ -0,0 +1,53 @@
|
||||
"""Tests for CPU isolation and core-spec parsing."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
from wikikb import config, cpu
|
||||
|
||||
|
||||
def test_parse_core_spec() -> None:
|
||||
assert config._parse_core_spec("0-3", 24) == {0, 1, 2, 3}
|
||||
assert config._parse_core_spec("8,10-11", 24) == {8, 10, 11}
|
||||
assert config._parse_core_spec("", 24) == set()
|
||||
assert config._parse_core_spec("20-100", 24) == {20, 21, 22, 23} # clamped
|
||||
|
||||
|
||||
def test_allowed_cores_excludes_reserved() -> None:
|
||||
assert config._allowed_cores(8, {0, 1}) == [2, 3, 4, 5, 6, 7]
|
||||
|
||||
|
||||
def test_allowed_cores_falls_back_when_all_reserved() -> None:
|
||||
assert config._allowed_cores(4, {0, 1, 2, 3}) == [0, 1, 2, 3]
|
||||
|
||||
|
||||
def test_isolate_sets_affinity_and_nice(monkeypatch) -> None:
|
||||
seen: dict[str, object] = {}
|
||||
monkeypatch.setattr(
|
||||
os, "sched_setaffinity", lambda pid, cores: seen.__setitem__("aff", (pid, set(cores)))
|
||||
)
|
||||
monkeypatch.setattr(os, "nice", lambda inc: seen.__setitem__("nice", inc))
|
||||
assert cpu.isolate(allowed=[4, 5, 6], nice_increment=7) == [4, 5, 6]
|
||||
assert seen["aff"] == (0, {4, 5, 6})
|
||||
assert seen["nice"] == 7
|
||||
|
||||
|
||||
def test_isolate_defaults_use_config(monkeypatch) -> None:
|
||||
seen: dict[str, object] = {}
|
||||
monkeypatch.setattr(
|
||||
os, "sched_setaffinity", lambda _pid, cores: seen.__setitem__("c", set(cores))
|
||||
)
|
||||
monkeypatch.setattr(os, "nice", lambda inc: seen.__setitem__("n", inc))
|
||||
assert cpu.isolate() == list(config.ALLOWED_CORES)
|
||||
assert seen["c"] == set(config.ALLOWED_CORES)
|
||||
assert seen["n"] == config.NICE_INCREMENT
|
||||
|
||||
|
||||
def test_isolate_no_affinity_and_zero_nice(monkeypatch) -> None:
|
||||
monkeypatch.delattr(os, "sched_setaffinity", raising=False)
|
||||
called = {"nice": False}
|
||||
monkeypatch.setattr(os, "nice", lambda _inc: called.__setitem__("nice", True))
|
||||
# nice_increment 0 -> nice() skipped; missing sched_setaffinity -> affinity skipped
|
||||
assert cpu.isolate(allowed=[0], nice_increment=0) == [0]
|
||||
assert called["nice"] is False
|
||||
36
tests/test_embed.py
Normal file
36
tests/test_embed.py
Normal file
@ -0,0 +1,36 @@
|
||||
"""Tests for the embedding wrapper (using the fake model fixture)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
|
||||
from wikikb import config, embed
|
||||
|
||||
|
||||
def test_embed_passages_shape_and_normalization(fake_model) -> None:
|
||||
vectors = embed.embed_passages(["a", "b", "c"], model=fake_model)
|
||||
assert vectors.shape == (3, config.EMBED_DIM)
|
||||
norms = np.linalg.norm(vectors, axis=1)
|
||||
assert np.allclose(norms, 1.0, atol=1e-5)
|
||||
|
||||
|
||||
def test_embed_passages_empty(fake_model) -> None:
|
||||
vectors = embed.embed_passages([], model=fake_model)
|
||||
assert vectors.shape == (0, config.EMBED_DIM)
|
||||
|
||||
|
||||
def test_embed_query_normalized(fake_model) -> None:
|
||||
vec = embed.embed_query("domain limit", model=fake_model)
|
||||
assert vec.shape == (config.EMBED_DIM,)
|
||||
assert np.isclose(np.linalg.norm(vec), 1.0, atol=1e-5)
|
||||
|
||||
|
||||
def test_get_model_returns_cached(monkeypatch, fake_model) -> None:
|
||||
monkeypatch.setattr(embed, "_model", fake_model)
|
||||
assert embed.get_model() is fake_model
|
||||
|
||||
|
||||
def test_normalize_handles_zero_vector() -> None:
|
||||
zero = np.zeros((1, config.EMBED_DIM), dtype=np.float32)
|
||||
out = embed._normalize(zero)
|
||||
assert np.all(out == 0.0) # zero norm replaced by 1, stays zero
|
||||
181
tests/test_html2md.py
Normal file
181
tests/test_html2md.py
Normal file
@ -0,0 +1,181 @@
|
||||
"""Tests for HTML -> markdown conversion."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from wikikb import html2md
|
||||
|
||||
_HTML = """
|
||||
<div class="mw-parser-output">
|
||||
<div class="mw-editsection">[edit]</div>
|
||||
<p>A <b>count</b> can hold up to 4 counties.</p>
|
||||
<table><tr><th>Trait</th><th>Diplomacy</th></tr>
|
||||
<tr><td>Genius</td><td>+5</td></tr></table>
|
||||
<div class="navbox">navigation</div>
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
||||
def test_clean_html_strips_chrome_keeps_content() -> None:
|
||||
cleaned = html2md.clean_html(_HTML)
|
||||
assert "editsection" not in cleaned
|
||||
assert "navbox" not in cleaned
|
||||
assert "count" in cleaned
|
||||
|
||||
|
||||
def test_clean_html_without_parser_output_wrapper() -> None:
|
||||
cleaned = html2md.clean_html("<p>bare</p>")
|
||||
assert "bare" in cleaned
|
||||
|
||||
|
||||
def test_html_to_markdown_preserves_table_numbers() -> None:
|
||||
md = html2md.html_to_markdown(_HTML)
|
||||
assert "Genius" in md
|
||||
assert "+5" in md
|
||||
assert "\n\n\n" not in md # blank lines collapsed
|
||||
|
||||
|
||||
def test_build_document_frontmatter_and_quote_escaping() -> None:
|
||||
doc = html2md.build_document(
|
||||
title='War of the "Roses"',
|
||||
url="https://wiki.example/War",
|
||||
revid=42,
|
||||
categories=["Warfare", "1.19"],
|
||||
html="<p>hello</p>",
|
||||
)
|
||||
assert doc.startswith("---\n")
|
||||
assert 'title: "War of the \\"Roses\\""' in doc
|
||||
assert "revid: 42" in doc
|
||||
assert 'categories: ["Warfare", "1.19"]' in doc
|
||||
assert "# War of the" in doc
|
||||
assert "hello" in doc
|
||||
|
||||
|
||||
# A cut-down Fandom PortableInfobox: labels are <h3> siblings of their values, not
|
||||
# table cells, and the tier variants are inlined behind tabs.
|
||||
_INFOBOX = """
|
||||
<div class="mw-parser-output">
|
||||
<aside class="portable-infobox">
|
||||
<h2 class="pi-item pi-title">Common longsword</h2>
|
||||
<section class="pi-item pi-group">
|
||||
<h2 class="pi-item pi-header">Information</h2>
|
||||
<div class="pi-item pi-data"><h3 class="pi-data-label">Classification</h3>
|
||||
<div class="pi-data-value">Weapons</div></div>
|
||||
</section>
|
||||
<h2 class="pi-item pi-header">Attributes</h2>
|
||||
<section class="pi-item pi-panel wds-tabber">
|
||||
<div class="wds-tabs__wrapper"><ul class="wds-tabs">
|
||||
<li class="wds-tabs__tab"><div class="wds-tabs__tab-label">I</div></li>
|
||||
<li class="wds-tabs__tab"><div class="wds-tabs__tab-label">II</div></li>
|
||||
</ul></div>
|
||||
<div class="wds-tab__content">
|
||||
<section class="pi-item pi-smart-group">
|
||||
<section class="pi-smart-group-head">
|
||||
<h3 class="pi-smart-data-label">Stab dmg.</h3>
|
||||
<h3 class="pi-smart-data-label">Slash dmg.</h3>
|
||||
</section>
|
||||
<section class="pi-smart-group-body">
|
||||
<div class="pi-smart-data-value">74</div>
|
||||
<div class="pi-smart-data-value">71</div>
|
||||
</section>
|
||||
</section>
|
||||
</div>
|
||||
<div class="wds-tab__content">
|
||||
<section class="pi-item pi-smart-group">
|
||||
<section class="pi-smart-group-head">
|
||||
<h3 class="pi-smart-data-label">Stab dmg.</h3>
|
||||
<h3 class="pi-smart-data-label">Slash dmg.</h3>
|
||||
</section>
|
||||
<section class="pi-smart-group-body">
|
||||
<div class="pi-smart-data-value">92</div>
|
||||
<div class="pi-smart-data-value">88</div>
|
||||
</section>
|
||||
</section>
|
||||
</div>
|
||||
</section>
|
||||
</aside>
|
||||
<p>The lead sentence.</p>
|
||||
</div>
|
||||
"""
|
||||
|
||||
|
||||
def test_infobox_becomes_a_table_with_labels_bound_to_values() -> None:
|
||||
md = html2md.html_to_markdown(_INFOBOX)
|
||||
# The point of the transform: not merely that "74" survives, but that it stays
|
||||
# attached to "Stab dmg." Plain markdownify emits headings then orphaned numbers.
|
||||
assert "| Stab dmg. | 74 |" in md
|
||||
assert "| Slash dmg. | 71 |" in md
|
||||
assert "| Classification | Weapons |" in md
|
||||
|
||||
|
||||
def test_infobox_tier_tabs_are_labelled() -> None:
|
||||
md = html2md.html_to_markdown(_INFOBOX)
|
||||
# Without the tab label, both tiers' stat blocks are indistinguishable.
|
||||
assert "## Infobox: Attributes (I)" in md
|
||||
assert "## Infobox: Attributes (II)" in md
|
||||
assert "| Stab dmg. | 92 |" in md
|
||||
|
||||
|
||||
def test_infobox_lead_text_gets_its_own_section() -> None:
|
||||
md = html2md.html_to_markdown(_INFOBOX)
|
||||
# The lead paragraph must not be filed under the last infobox group.
|
||||
assert md.index("## Overview") < md.index("The lead sentence.")
|
||||
|
||||
|
||||
def test_infobox_without_a_group_header_is_named_plainly() -> None:
|
||||
html = (
|
||||
'<aside class="portable-infobox"><div class="pi-item pi-data">'
|
||||
'<h3 class="pi-data-label">ID</h3><div class="pi-data-value">4</div></div></aside>'
|
||||
)
|
||||
md = html2md.html_to_markdown(html)
|
||||
assert "## Infobox\n" in md
|
||||
assert "| ID | 4 |" in md
|
||||
|
||||
|
||||
def test_empty_infobox_is_dropped() -> None:
|
||||
# Theme/dummy infoboxes hold no data and would otherwise leave a stray heading.
|
||||
md = html2md.html_to_markdown(
|
||||
'<div><aside class="portable-infobox pi-theme-dummy"></aside><p>body</p></div>'
|
||||
)
|
||||
assert "Infobox" not in md
|
||||
assert "body" in md
|
||||
|
||||
|
||||
def test_mismatched_smart_group_keeps_the_pairs_that_line_up() -> None:
|
||||
html = (
|
||||
'<aside class="portable-infobox"><section class="pi-smart-group">'
|
||||
'<section class="pi-smart-group-head"><h3 class="pi-smart-data-label">A</h3>'
|
||||
'<h3 class="pi-smart-data-label">B</h3></section>'
|
||||
'<section class="pi-smart-group-body"><div class="pi-smart-data-value">1</div>'
|
||||
"</section></section></aside>"
|
||||
)
|
||||
md = html2md.html_to_markdown(html)
|
||||
assert "| A | 1 |" in md
|
||||
assert "| B |" not in md
|
||||
|
||||
|
||||
def test_incomplete_pi_data_is_skipped_without_losing_the_rest() -> None:
|
||||
# A hand-edited infobox row can carry a label with no value. Skip that row and
|
||||
# keep going rather than emitting a label bound to nothing.
|
||||
html = (
|
||||
'<aside class="portable-infobox">'
|
||||
'<div class="pi-item pi-data"><h3 class="pi-data-label">Orphan</h3></div>'
|
||||
'<div class="pi-item pi-data"><h3 class="pi-data-label">Weight</h3>'
|
||||
'<div class="pi-data-value">3.9</div></div>'
|
||||
"</aside>"
|
||||
)
|
||||
md = html2md.html_to_markdown(html)
|
||||
assert "| Weight | 3.9 |" in md
|
||||
assert "Orphan" not in md
|
||||
|
||||
|
||||
def test_page_tabber_switcher_is_stripped() -> None:
|
||||
# It links to the other game's page; keeping it would put a bare "KCD" link at
|
||||
# the top of every dual-game article.
|
||||
html = (
|
||||
'<div class="mw-parser-output">'
|
||||
'<div class="pagetabber-absolute"><a href="/wiki/Alchemy">KCD</a></div>'
|
||||
"<p>real content</p></div>"
|
||||
)
|
||||
md = html2md.html_to_markdown(html)
|
||||
assert "real content" in md
|
||||
assert "/wiki/Alchemy" not in md
|
||||
110
tests/test_indexer.py
Normal file
110
tests/test_indexer.py
Normal file
@ -0,0 +1,110 @@
|
||||
"""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
|
||||
294
tests/test_scraper.py
Normal file
294
tests/test_scraper.py
Normal file
@ -0,0 +1,294 @@
|
||||
"""Tests for the MediaWiki scraper (HTTP mocked with respx)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from httpx import Response
|
||||
|
||||
from wikikb import config, scraper
|
||||
from wikikb.scraper import Throttle
|
||||
|
||||
|
||||
def _noop(_: float) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _throttle() -> Throttle:
|
||||
return Throttle(interval=0.0, sleep=_noop, clock=lambda: 0.0)
|
||||
|
||||
|
||||
def test_throttle_sleeps_then_skips() -> None:
|
||||
calls: list[float] = []
|
||||
times = iter([0.0, 0.0, 100.0, 100.0])
|
||||
throttle = Throttle(interval=10.0, sleep=calls.append, clock=lambda: next(times))
|
||||
throttle.wait() # now=0, last=0 -> delay 10 -> sleep
|
||||
throttle.wait() # now=100, last=0 -> delay negative -> no sleep
|
||||
assert calls == [10.0]
|
||||
|
||||
|
||||
def test_sanitize_filename() -> None:
|
||||
assert scraper.sanitize_filename("A/B: C?") == "A_B__C.md"
|
||||
assert scraper.sanitize_filename("...") == "untitled.md"
|
||||
|
||||
|
||||
def test_is_core_title() -> None:
|
||||
assert scraper.is_core_title("Traits")
|
||||
assert scraper.is_core_title("Domain limit")
|
||||
assert not scraper.is_core_title("CK3AGOT/Dragon") # subpage
|
||||
assert not scraper.is_core_title("LotR: Realms in Exile") # mod prefix
|
||||
assert not scraper.is_core_title("Elder Kings 2 overview") # mod prefix
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_api_get_success() -> None:
|
||||
respx.get(config.API_URL).mock(return_value=Response(200, json={"ok": 1}))
|
||||
with httpx.Client() as client:
|
||||
data = scraper._api_get(client, {"action": "query"}, throttle=_throttle(), sleep=_noop)
|
||||
assert data == {"ok": 1}
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_api_get_retries_on_503() -> None:
|
||||
route = respx.get(config.API_URL).mock(
|
||||
side_effect=[Response(503), Response(200, json={"ok": 2})]
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
data = scraper._api_get(client, {}, throttle=_throttle(), sleep=_noop)
|
||||
assert data == {"ok": 2}
|
||||
assert route.call_count == 2
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_api_get_retries_on_maxlag() -> None:
|
||||
respx.get(config.API_URL).mock(
|
||||
side_effect=[
|
||||
Response(200, json={"error": {"code": "maxlag"}}),
|
||||
Response(200, json={"ok": 3}),
|
||||
]
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
data = scraper._api_get(client, {}, throttle=_throttle(), sleep=_noop)
|
||||
assert data == {"ok": 3}
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_api_get_retries_on_bot_challenge() -> None:
|
||||
# Cloudflare serves a challenge as 200 + HTML, so status checks pass and only
|
||||
# the JSON decode fails. Without this the challenge page reaches the corpus.
|
||||
challenge = Response(200, html="<html><title>Just a moment...</title></html>")
|
||||
respx.get(config.API_URL).mock(side_effect=[challenge, Response(200, json={"ok": 4})])
|
||||
with httpx.Client() as client:
|
||||
assert scraper._api_get(client, {}, throttle=_throttle(), sleep=_noop) == {"ok": 4}
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_api_get_exhausts_retries_on_persistent_challenge() -> None:
|
||||
respx.get(config.API_URL).mock(
|
||||
return_value=Response(200, html="<html><title>Just a moment...</title></html>")
|
||||
)
|
||||
with httpx.Client() as client, pytest.raises(RuntimeError):
|
||||
scraper._api_get(client, {}, throttle=_throttle(), sleep=_noop)
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_api_get_exhausts_retries() -> None:
|
||||
respx.get(config.API_URL).mock(return_value=Response(503))
|
||||
with httpx.Client() as client, pytest.raises(RuntimeError):
|
||||
scraper._api_get(client, {}, throttle=_throttle(), sleep=_noop)
|
||||
|
||||
|
||||
def test_manifest_roundtrip(tmp_kb) -> None:
|
||||
assert scraper.load_manifest() == {}
|
||||
scraper.save_manifest({"A": {"revid": 1}})
|
||||
assert scraper.load_manifest() == {"A": {"revid": 1}}
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_scrape_writes(tmp_kb) -> None:
|
||||
respx.get(config.API_URL).mock(
|
||||
side_effect=[
|
||||
Response(200, json={"query": {"allpages": [{"title": "Army"}]}}),
|
||||
Response(
|
||||
200,
|
||||
json={
|
||||
"parse": {
|
||||
"title": "Army",
|
||||
"revid": 7,
|
||||
"categories": [],
|
||||
"text": "<p>Levies</p>",
|
||||
}
|
||||
},
|
||||
),
|
||||
]
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
result = scraper.scrape(client=client, throttle=_throttle(), log=lambda _: None)
|
||||
assert result == {"written": 1, "skipped": 0, "missing": 0}
|
||||
assert (config.CORPUS_DIR / "Army.md").exists()
|
||||
assert "Army" in scraper.load_manifest()
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_scrape_skips_existing(tmp_kb) -> None:
|
||||
(config.CORPUS_DIR / "Army.md").write_text("x", encoding="utf-8")
|
||||
scraper.save_manifest({"Army": {"revid": 1, "filename": "Army.md"}})
|
||||
respx.get(config.API_URL).mock(
|
||||
return_value=Response(200, json={"query": {"allpages": [{"title": "Army"}]}})
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
result = scraper.scrape(client=client, throttle=_throttle(), log=lambda _: None)
|
||||
assert result == {"written": 0, "skipped": 1, "missing": 0}
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_scrape_missing(tmp_kb) -> None:
|
||||
respx.get(config.API_URL).mock(
|
||||
side_effect=[
|
||||
Response(200, json={"query": {"allpages": [{"title": "Ghost"}]}}),
|
||||
Response(200, json={"error": {"code": "missingtitle"}}),
|
||||
]
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
result = scraper.scrape(client=client, throttle=_throttle(), log=lambda _: None)
|
||||
assert result == {"written": 0, "skipped": 0, "missing": 1}
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_scrape_periodic_save_and_owns_client(tmp_kb) -> None:
|
||||
titles = [{"title": f"P{i}"} for i in range(25)]
|
||||
responses = [Response(200, json={"query": {"allpages": titles}})]
|
||||
responses += [
|
||||
Response(
|
||||
200,
|
||||
json={
|
||||
"parse": {
|
||||
"title": f"P{i}",
|
||||
"revid": i,
|
||||
"categories": [],
|
||||
"text": f"<p>{i}</p>",
|
||||
}
|
||||
},
|
||||
)
|
||||
for i in range(25)
|
||||
]
|
||||
respx.get(config.API_URL).mock(side_effect=responses)
|
||||
result = scraper.scrape(throttle=_throttle(), log=lambda _: None) # client=None
|
||||
assert result["written"] == 25
|
||||
assert len(scraper.load_manifest()) == 25
|
||||
|
||||
|
||||
def test_main_delegates(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(scraper, "scrape", lambda **_: {"written": 1, "skipped": 0, "missing": 0})
|
||||
assert scraper.main(["--limit", "5"]) == 0
|
||||
|
||||
|
||||
def test_build_source_picks_mediawiki_by_default() -> None:
|
||||
with httpx.Client() as client:
|
||||
source = scraper.build_source(client, throttle=_throttle())
|
||||
assert source.name == "mediawiki"
|
||||
|
||||
|
||||
def test_build_source_picks_html_when_configured(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(config, "SOURCE_KIND", "html")
|
||||
with httpx.Client() as client:
|
||||
source = scraper.build_source(client, throttle=_throttle())
|
||||
assert source.name == "html"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_get_returns_404_without_retrying() -> None:
|
||||
# For the sitemap source a missing page is a normal outcome, not a transient
|
||||
# failure, so 404 must come back rather than burn the retry budget.
|
||||
route = respx.get("https://w.example/Gone").mock(return_value=Response(404))
|
||||
with httpx.Client() as client:
|
||||
assert (
|
||||
scraper._get(client, "https://w.example/Gone", throttle=_throttle()).status_code == 404
|
||||
)
|
||||
assert route.call_count == 1
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_get_retries_then_succeeds() -> None:
|
||||
respx.get("https://w.example/P").mock(side_effect=[Response(503), Response(200, text="ok")])
|
||||
with httpx.Client() as client:
|
||||
got = scraper._get(client, "https://w.example/P", throttle=_throttle(), sleep=_noop)
|
||||
assert got.text == "ok"
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_get_exhausts_retries() -> None:
|
||||
respx.get("https://w.example/P").mock(return_value=Response(503))
|
||||
with httpx.Client() as client, pytest.raises(RuntimeError):
|
||||
scraper._get(client, "https://w.example/P", throttle=_throttle(), sleep=_noop)
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_scrape_skips_noise_titles(tmp_kb) -> None:
|
||||
# Filtering moved out of the source and into scrape(), so it applies identically
|
||||
# whichever source enumerated the titles.
|
||||
respx.get(config.API_URL).mock(
|
||||
side_effect=[
|
||||
Response(
|
||||
200,
|
||||
json={"query": {"allpages": [{"title": "CK3AGOT/Dragon"}, {"title": "Army"}]}},
|
||||
),
|
||||
Response(
|
||||
200,
|
||||
json={"parse": {"title": "Army", "revid": 1, "categories": [], "text": "<p>x</p>"}},
|
||||
),
|
||||
]
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
result = scraper.scrape(client=client, throttle=_throttle(), log=lambda _: None)
|
||||
assert result == {"written": 1, "skipped": 0, "missing": 0}
|
||||
assert sorted(scraper.load_manifest()) == ["Army"]
|
||||
|
||||
|
||||
def test_is_core_title_subpage_policy_reject(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# CK3-style wiki: subpages are total-conversion mods.
|
||||
monkeypatch.setattr(config, "SUBPAGE_POLICY", "reject")
|
||||
monkeypatch.setattr(config, "NOISE_SUBPAGE_TAILS", ())
|
||||
assert scraper.is_core_title("Traits")
|
||||
assert not scraper.is_core_title("CK3AGOT/Dragon")
|
||||
|
||||
|
||||
def test_is_core_title_subpage_policy_keep(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# KCD-style wiki: subpages ARE the sequel's pages. Rejecting them would delete
|
||||
# 441 pages of exactly the content the user wants.
|
||||
monkeypatch.setattr(config, "SUBPAGE_POLICY", "keep")
|
||||
monkeypatch.setattr(config, "NOISE_SUBPAGE_TAILS", ("loc",))
|
||||
monkeypatch.setattr(config, "MOD_TITLE_PREFIXES", ())
|
||||
assert scraper.is_core_title("Alchemy/KCD2")
|
||||
assert scraper.is_core_title("Archery master/Uzhitz")
|
||||
assert not scraper.is_core_title("Argonautica/loc")
|
||||
|
||||
|
||||
@respx.mock
|
||||
def test_scrape_with_explicit_titles_skips_enumeration(tmp_kb) -> None:
|
||||
# Only the parse call is mocked: enumerating the whole wiki would hit an
|
||||
# unmocked allpages request and fail the test.
|
||||
respx.get(config.API_URL).mock(
|
||||
return_value=Response(
|
||||
200,
|
||||
json={"parse": {"title": "Army", "revid": 3, "categories": [], "text": "<p>x</p>"}},
|
||||
)
|
||||
)
|
||||
with httpx.Client() as client:
|
||||
result = scraper.scrape(
|
||||
titles=["Army"], client=client, throttle=_throttle(), log=lambda _: None
|
||||
)
|
||||
assert result == {"written": 1, "skipped": 0, "missing": 0}
|
||||
|
||||
|
||||
def test_main_parses_titles(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
seen: dict[str, object] = {}
|
||||
|
||||
def fake_scrape(**kwargs: object) -> dict[str, int]:
|
||||
seen.update(kwargs)
|
||||
return {"written": 0, "skipped": 0, "missing": 0}
|
||||
|
||||
monkeypatch.setattr(scraper, "scrape", fake_scrape)
|
||||
assert scraper.main(["--titles", "Alchemy/KCD2, Lockpicking ,"]) == 0
|
||||
assert seen["titles"] == ["Alchemy/KCD2", "Lockpicking"]
|
||||
181
tests/test_search.py
Normal file
181
tests/test_search.py
Normal file
@ -0,0 +1,181 @@
|
||||
"""Tests for hybrid search."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from wikikb import config, embed, indexer, search
|
||||
|
||||
_TRAITS = """---
|
||||
title: "Traits"
|
||||
url: https://x/Traits
|
||||
revid: 1
|
||||
categories: []
|
||||
---
|
||||
|
||||
# Traits
|
||||
|
||||
Genius gives plus five diplomacy. Domain limit rises with stewardship.
|
||||
"""
|
||||
|
||||
_ARMY = """---
|
||||
title: "Army"
|
||||
url: https://x/Army
|
||||
revid: 2
|
||||
categories: []
|
||||
---
|
||||
|
||||
# Army
|
||||
|
||||
Levies and knights fight battles.
|
||||
"""
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def indexed(tmp_kb, fake_model):
|
||||
(config.CORPUS_DIR / "traits.md").write_text(_TRAITS, encoding="utf-8")
|
||||
(config.CORPUS_DIR / "army.md").write_text(_ARMY, encoding="utf-8")
|
||||
indexer.build_index(
|
||||
embed_fn=lambda texts: embed.embed_passages(texts, model=fake_model),
|
||||
log=lambda _: None,
|
||||
)
|
||||
return fake_model
|
||||
|
||||
|
||||
def test_fts_match_query_ors_terms() -> None:
|
||||
assert search._fts_match_query("Domain limit") == '"domain" OR "limit"'
|
||||
assert search._fts_match_query("!!!") == ""
|
||||
|
||||
|
||||
def test_reciprocal_rank_fusion_orders_and_dedupes() -> None:
|
||||
fused = search.reciprocal_rank_fusion([[1, 2, 3], [2, 1]])
|
||||
assert set(fused) == {1, 2, 3}
|
||||
assert fused[-1] == 3 # only appears once, at rank 2 in one list
|
||||
|
||||
|
||||
def test_vector_ids_empty() -> None:
|
||||
empty = np.zeros((0, config.EMBED_DIM), dtype=np.float32)
|
||||
query = np.zeros(config.EMBED_DIM, dtype=np.float32)
|
||||
assert search._vector_ids(empty, query, 5, None) == []
|
||||
|
||||
|
||||
def test_search_ranks_relevant_page_first(indexed, fake_model) -> None:
|
||||
results = search.search("domain limit stewardship", model=fake_model)
|
||||
assert results
|
||||
assert results[0].title == "Traits"
|
||||
assert "stewardship" in results[0].text.lower()
|
||||
|
||||
|
||||
def test_search_no_fts_match_still_returns_vector_hits(indexed, fake_model) -> None:
|
||||
results = search.search("zzznonexistentxyz", model=fake_model)
|
||||
assert isinstance(results, list)
|
||||
assert results # vector retrieval still yields candidates
|
||||
|
||||
|
||||
def test_search_punctuation_query_skips_fts(indexed, fake_model) -> None:
|
||||
# punctuation-only query -> empty FTS match -> _fts_ids returns []; vector still runs
|
||||
results = search.search("!!!", model=fake_model)
|
||||
assert isinstance(results, list)
|
||||
|
||||
|
||||
def test_search_empty_index_returns_nothing(tmp_kb, fake_model) -> None:
|
||||
indexer.build_index(
|
||||
embed_fn=lambda texts: embed.embed_passages(texts, model=fake_model),
|
||||
log=lambda _: None,
|
||||
)
|
||||
assert search.search("anything", model=fake_model) == []
|
||||
|
||||
|
||||
def test_format_empty() -> None:
|
||||
assert search._format([]) == "(no results)"
|
||||
|
||||
|
||||
def test_format_truncates_long_text() -> None:
|
||||
result = search.Result(
|
||||
title="T",
|
||||
url="u",
|
||||
source="mediawiki",
|
||||
variant="kcd2",
|
||||
section="S",
|
||||
text="x" * 600,
|
||||
score=1.0,
|
||||
)
|
||||
out = search._format([result])
|
||||
assert "…" in out
|
||||
assert "[kcd2] T > S" in out
|
||||
|
||||
|
||||
def test_main_requires_query() -> None:
|
||||
assert search.main([]) == 2
|
||||
|
||||
|
||||
def test_main_runs(monkeypatch, capsys) -> None:
|
||||
monkeypatch.setattr(
|
||||
search,
|
||||
"search",
|
||||
lambda *_a, **_k: [search.Result("T", "u", "mediawiki", "kcd2", "S", "body", 1.0)],
|
||||
)
|
||||
assert search.main(["domain", "limit"]) == 0
|
||||
assert "[kcd2] T > S" in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_vector_ids_filters_before_truncating() -> None:
|
||||
# Three vectors ranked 3,2,1 by similarity; only rowid 1 (the WORST match) is in
|
||||
# scope. Truncating to one candidate first returns nothing; filtering during
|
||||
# selection returns the in-scope hit.
|
||||
vectors = np.array([[1.0], [2.0], [3.0]], dtype=np.float32)
|
||||
query = np.array([1.0], dtype=np.float32)
|
||||
assert search._vector_ids(vectors, query, 1, {1}) == [1]
|
||||
assert search._vector_ids(vectors, query, 1, None) == [3]
|
||||
|
||||
|
||||
def _variant_doc(title: str, name: str, body: str) -> str:
|
||||
return (
|
||||
f'---\ntitle: "{title}"\nurl: https://x/{title}\nrevid: 1\n'
|
||||
f"source: mediawiki\nvariant: {name}\ncategories: []\n---\n\n# {title}\n\n{body}\n"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def indexed_variants(tmp_kb, fake_model, monkeypatch):
|
||||
monkeypatch.setattr(config, "VARIANTS", (("kcd2", "KCD2"), ("kcd1", "KCD")))
|
||||
(config.CORPUS_DIR / "new.md").write_text(
|
||||
_variant_doc("Alchemy KCD2", "kcd2", "Brewing potions with herbs."), encoding="utf-8"
|
||||
)
|
||||
(config.CORPUS_DIR / "old.md").write_text(
|
||||
_variant_doc("Alchemy", "kcd1", "Brewing potions with herbs in the first game."),
|
||||
encoding="utf-8",
|
||||
)
|
||||
(config.CORPUS_DIR / "shared.md").write_text(
|
||||
_variant_doc("Moravia", "unknown", "Brewing herbs region shared by both."),
|
||||
encoding="utf-8",
|
||||
)
|
||||
indexer.build_index(
|
||||
embed_fn=lambda texts: embed.embed_passages(texts, model=fake_model), log=lambda _: None
|
||||
)
|
||||
return fake_model
|
||||
|
||||
|
||||
def test_search_variant_scope_excludes_the_other_edition(indexed_variants, fake_model) -> None:
|
||||
titles = {
|
||||
r.title
|
||||
for r in search.search("brewing potions herbs", k=10, variant="kcd2", model=fake_model)
|
||||
}
|
||||
assert "Alchemy KCD2" in titles
|
||||
assert "Moravia" in titles # unknown is always in scope
|
||||
assert "Alchemy" not in titles
|
||||
|
||||
|
||||
def test_search_any_scope_returns_every_edition(indexed_variants, fake_model) -> None:
|
||||
titles = {
|
||||
r.title
|
||||
for r in search.search("brewing potions herbs", k=10, variant="any", model=fake_model)
|
||||
}
|
||||
assert {"Alchemy KCD2", "Alchemy", "Moravia"} <= titles
|
||||
|
||||
|
||||
def test_search_results_carry_source_and_variant(indexed_variants, fake_model) -> None:
|
||||
results = search.search("brewing potions herbs", k=10, variant="any", model=fake_model)
|
||||
assert {r.source for r in results} == {"mediawiki"}
|
||||
assert {r.title: r.variant for r in results}["Alchemy"] == "kcd1"
|
||||
63
tests/test_server.py
Normal file
63
tests/test_server.py
Normal file
@ -0,0 +1,63 @@
|
||||
"""Tests for the MCP server tool functions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from wikikb import config, scraper, search, server
|
||||
|
||||
|
||||
def test_wiki_search_maps_results(monkeypatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
server,
|
||||
"search",
|
||||
lambda _q, k, variant: [search.Result("T", "u", "mediawiki", "kcd2", "S", "body", 1.0)],
|
||||
)
|
||||
assert server.wiki_search("q", k=3) == [
|
||||
{
|
||||
"title": "T",
|
||||
"source": "mediawiki",
|
||||
"variant": "kcd2",
|
||||
"section": "S",
|
||||
"url": "u",
|
||||
"text": "body",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_wiki_get_page_from_manifest(tmp_kb) -> None:
|
||||
(config.CORPUS_DIR / "Army.md").write_text("PAGE BODY", encoding="utf-8")
|
||||
scraper.save_manifest({"Army": {"filename": "Army.md"}})
|
||||
assert "PAGE BODY" in server.wiki_get_page("Army")
|
||||
|
||||
|
||||
def test_wiki_get_page_sanitize_fallback(tmp_kb) -> None:
|
||||
(config.CORPUS_DIR / "Foo_Bar.md").write_text("FB", encoding="utf-8")
|
||||
assert "FB" in server.wiki_get_page("Foo Bar") # no manifest entry -> sanitized name
|
||||
|
||||
|
||||
def test_wiki_get_page_missing(tmp_kb) -> None:
|
||||
assert "No mirrored page" in server.wiki_get_page("Nope")
|
||||
|
||||
|
||||
def test_wiki_list_pages(tmp_kb) -> None:
|
||||
scraper.save_manifest({"Army": {}, "Archery": {}, "Barony": {}})
|
||||
assert server.wiki_list_pages() == ["Archery", "Army", "Barony"]
|
||||
assert server.wiki_list_pages(prefix="ar") == ["Archery", "Army"]
|
||||
assert server.wiki_list_pages(limit=1) == ["Archery"]
|
||||
|
||||
|
||||
class _FakeServer:
|
||||
def __init__(self) -> None:
|
||||
self.registered: list[str] = []
|
||||
|
||||
def tool(self):
|
||||
def decorator(fn):
|
||||
self.registered.append(fn.__name__)
|
||||
return fn
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def test_build_server_registers_tools() -> None:
|
||||
fake = _FakeServer()
|
||||
assert server.build_server(fake) is fake
|
||||
assert fake.registered == ["wiki_search", "wiki_get_page", "wiki_list_pages"]
|
||||
190
tests/test_sources.py
Normal file
190
tests/test_sources.py
Normal file
@ -0,0 +1,190 @@
|
||||
"""Tests for the pluggable page sources and host detection."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from wikikb import config
|
||||
from wikikb.sources import (
|
||||
HTML,
|
||||
MEDIAWIKI,
|
||||
MediaWikiSource,
|
||||
SitemapHtmlSource,
|
||||
detect_source_kind,
|
||||
)
|
||||
|
||||
_SITEMAP = (
|
||||
"<urlset><url><loc>https://w.example/Common+Longsword</loc></url>"
|
||||
"<url><loc>https://w.example/Stealth</loc></url></urlset>"
|
||||
)
|
||||
_INDEX = "<sitemapindex><sitemap><loc>https://w.example/s1.xml</loc></sitemap></sitemapindex>"
|
||||
_PAGE = '<html><body><nav>menu</nav><div id="main-content"><p>Stab 122</p></div></body></html>'
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- MediaWiki
|
||||
|
||||
|
||||
def _mw(pages: list[dict]) -> MediaWikiSource:
|
||||
return MediaWikiSource(lambda _params: pages.pop(0))
|
||||
|
||||
|
||||
def test_mediawiki_iter_titles_paginates() -> None:
|
||||
source = _mw(
|
||||
[
|
||||
{
|
||||
"query": {"allpages": [{"title": "A"}, {"title": "B"}]},
|
||||
"continue": {"apcontinue": "C"},
|
||||
},
|
||||
{"query": {"allpages": [{"title": "C"}]}},
|
||||
]
|
||||
)
|
||||
assert list(source.iter_titles()) == ["A", "B", "C"]
|
||||
|
||||
|
||||
def test_mediawiki_iter_titles_respects_limit() -> None:
|
||||
source = _mw(
|
||||
[{"query": {"allpages": [{"title": "A"}, {"title": "B"}]}, "continue": {"apcontinue": "C"}}]
|
||||
)
|
||||
assert list(source.iter_titles(limit=1)) == ["A"]
|
||||
|
||||
|
||||
def test_mediawiki_fetch_page_carries_metadata() -> None:
|
||||
source = _mw(
|
||||
[
|
||||
{
|
||||
"parse": {
|
||||
"title": "Army",
|
||||
"revid": 7,
|
||||
"categories": [{"category": "War_stuff"}],
|
||||
"text": "<p>hi</p>",
|
||||
}
|
||||
}
|
||||
]
|
||||
)
|
||||
page = source.fetch_page("Army")
|
||||
assert page is not None
|
||||
assert (page.title, page.revid, page.source) == ("Army", 7, MEDIAWIKI)
|
||||
assert page.categories == ["War stuff"]
|
||||
|
||||
|
||||
def test_mediawiki_fetch_page_missing() -> None:
|
||||
assert _mw([{"error": {"code": "missingtitle"}}]).fetch_page("Nope") is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- Sitemap/HTML
|
||||
|
||||
|
||||
def _html_source(responses: dict[str, httpx.Response]) -> SitemapHtmlSource:
|
||||
return SitemapHtmlSource(
|
||||
lambda url: responses[url], sitemap_url="https://w.example/sitemap.xml"
|
||||
)
|
||||
|
||||
|
||||
def test_title_from_url_decodes_slug() -> None:
|
||||
assert SitemapHtmlSource.title_from_url("https://w.example/Common+Longsword") == (
|
||||
"Common Longsword"
|
||||
)
|
||||
assert SitemapHtmlSource.title_from_url("https://w.example/A_B/") == "A B"
|
||||
assert SitemapHtmlSource.title_from_url("https://w.example/") == ""
|
||||
|
||||
|
||||
def test_sitemap_iter_titles() -> None:
|
||||
source = _html_source({"https://w.example/sitemap.xml": httpx.Response(200, text=_SITEMAP)})
|
||||
assert list(source.iter_titles()) == ["Common Longsword", "Stealth"]
|
||||
|
||||
|
||||
def test_sitemap_iter_titles_skips_empty_slug_and_respects_limit() -> None:
|
||||
xml = "<urlset><url><loc>https://w.example/</loc></url>" + _SITEMAP
|
||||
source = _html_source({"https://w.example/sitemap.xml": httpx.Response(200, text=xml)})
|
||||
assert list(source.iter_titles(limit=1)) == ["Common Longsword"]
|
||||
|
||||
|
||||
def test_sitemap_follows_a_sitemap_index() -> None:
|
||||
source = _html_source(
|
||||
{
|
||||
"https://w.example/sitemap.xml": httpx.Response(200, text=_INDEX),
|
||||
"https://w.example/s1.xml": httpx.Response(200, text=_SITEMAP),
|
||||
}
|
||||
)
|
||||
assert list(source.iter_titles()) == ["Common Longsword", "Stealth"]
|
||||
|
||||
|
||||
def test_sitemap_fetch_page_extracts_the_content_container() -> None:
|
||||
source = _html_source(
|
||||
{
|
||||
"https://w.example/sitemap.xml": httpx.Response(200, text=_SITEMAP),
|
||||
"https://w.example/Stealth": httpx.Response(200, text=_PAGE),
|
||||
}
|
||||
)
|
||||
list(source.iter_titles()) # populates the title -> url map
|
||||
page = source.fetch_page("Stealth")
|
||||
assert page is not None
|
||||
assert page.source == HTML
|
||||
assert page.revid == 0 # no revision concept without an API
|
||||
assert "Stab 122" in page.html
|
||||
assert "menu" not in page.html # nav is outside the content container
|
||||
|
||||
|
||||
def test_sitemap_fetch_page_404_and_unextractable(monkeypatch) -> None:
|
||||
monkeypatch.setattr(config, "WIKI_HOST", "https://w.example")
|
||||
source = _html_source(
|
||||
{
|
||||
"https://w.example/Gone": httpx.Response(404),
|
||||
"https://w.example/Empty": httpx.Response(200, text="<html><body>no container</body>"),
|
||||
}
|
||||
)
|
||||
assert source.fetch_page("Gone") is None # honest 404 from the server
|
||||
assert source.fetch_page("Empty") is None # no recognisable content container
|
||||
|
||||
|
||||
def test_sitemap_fetch_page_falls_back_to_a_derived_url(monkeypatch) -> None:
|
||||
monkeypatch.setattr(config, "WIKI_HOST", "https://w.example")
|
||||
source = _html_source({"https://w.example/Never%20Listed": httpx.Response(200, text=_PAGE)})
|
||||
assert source.fetch_page("Never Listed") is not None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------- detection
|
||||
|
||||
|
||||
def test_detect_prefers_the_mediawiki_api() -> None:
|
||||
def get(url: str) -> httpx.Response:
|
||||
if url.startswith("https://w.example/api.php"):
|
||||
return httpx.Response(200, json={"query": {"general": {"generator": "MediaWiki 1.43"}}})
|
||||
return httpx.Response(404)
|
||||
|
||||
assert detect_source_kind(get, "https://w.example") == (MEDIAWIKI, "https://w.example/api.php")
|
||||
|
||||
|
||||
def test_detect_finds_the_wikipedia_style_api_path() -> None:
|
||||
def get(url: str) -> httpx.Response:
|
||||
if url.startswith("https://w.example/w/api.php"):
|
||||
return httpx.Response(200, json={"query": {"general": {"generator": "MediaWiki 1.41"}}})
|
||||
return httpx.Response(404)
|
||||
|
||||
assert detect_source_kind(get, "https://w.example")[1] == "https://w.example/w/api.php"
|
||||
|
||||
|
||||
def test_detect_falls_back_to_sitemap_when_there_is_no_api() -> None:
|
||||
def get(url: str) -> httpx.Response:
|
||||
if url.endswith("sitemap.xml"):
|
||||
return httpx.Response(200, text=_SITEMAP)
|
||||
return httpx.Response(404)
|
||||
|
||||
assert detect_source_kind(get, "https://w.example") == (HTML, "https://w.example/sitemap.xml")
|
||||
|
||||
|
||||
def test_detect_tolerates_a_non_json_or_non_mediawiki_api_response() -> None:
|
||||
# A host that serves HTML at /api.php (or some other JSON API) must not be
|
||||
# mistaken for MediaWiki; detection falls through to the sitemap.
|
||||
def get(url: str) -> httpx.Response:
|
||||
if "api.php" in url:
|
||||
return httpx.Response(200, text="<html>not an api</html>")
|
||||
return httpx.Response(200, text=_SITEMAP)
|
||||
|
||||
assert detect_source_kind(get, "https://w.example")[0] == HTML
|
||||
|
||||
|
||||
def test_detect_raises_when_neither_is_available() -> None:
|
||||
with pytest.raises(RuntimeError, match="neither a MediaWiki API nor a sitemap"):
|
||||
detect_source_kind(lambda _url: httpx.Response(404), "https://w.example")
|
||||
50
tests/test_variant.py
Normal file
50
tests/test_variant.py
Normal file
@ -0,0 +1,50 @@
|
||||
"""Tests for edition/variant tagging on wikis that cover two editions at once."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from wikikb import config, variant
|
||||
|
||||
# Ordered most-specific-first, exactly as a profile declares it.
|
||||
_KCD = (("kcd2", "KCD2"), ("kcd1", "KCD"))
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _variants(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(config, "VARIANTS", _KCD)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("title", "categories", "expected"),
|
||||
[
|
||||
("Alchemy/KCD2", [], "kcd2"), # title suffix
|
||||
("Common longsword", ["KCD2: Weapons"], "kcd2"), # bare title, category only
|
||||
("Lockpicking", ["KCD: Skills"], "kcd1"),
|
||||
# "KCD main quests" has no colon: a "KCD:"-prefix test would miss it, and a
|
||||
# naive "KCD" test would also match every KCD2 category.
|
||||
("Awakening", ["KCD main quests"], "kcd1"),
|
||||
("Weapons", ["KCD: Items", "KCD2: Weapons"], "both"),
|
||||
("Moravia", ["Countries"], "unknown"),
|
||||
],
|
||||
)
|
||||
def test_classify(title: str, categories: list[str], expected: str) -> None:
|
||||
assert variant.classify(title, categories) == expected
|
||||
|
||||
|
||||
def test_classify_without_configured_variants(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# Most wikis cover one subject; tagging is off unless a profile declares it.
|
||||
monkeypatch.setattr(config, "VARIANTS", ())
|
||||
assert variant.classify("Alchemy/KCD2", ["KCD2: Skills"]) == "unknown"
|
||||
|
||||
|
||||
def test_matches_excludes_only_the_other_edition() -> None:
|
||||
assert variant.matches("kcd2", "kcd2")
|
||||
assert variant.matches("both", "kcd2")
|
||||
assert variant.matches("unknown", "kcd2") # untagged pages stay in scope
|
||||
assert not variant.matches("kcd1", "kcd2")
|
||||
|
||||
|
||||
def test_matches_any_scope() -> None:
|
||||
for name in ("kcd1", "kcd2", "both", "unknown"):
|
||||
assert variant.matches(name, "any")
|
||||
1
wikikb/__init__.py
Normal file
1
wikikb/__init__.py
Normal file
@ -0,0 +1 @@
|
||||
"""wiki-kb: local searchable mirror of the Crusader Kings III wiki."""
|
||||
95
wikikb/chunker.py
Normal file
95
wikikb/chunker.py
Normal file
@ -0,0 +1,95 @@
|
||||
"""Split corpus markdown documents into overlapping, heading-scoped chunks."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
from . import config
|
||||
|
||||
_FRONTMATTER = re.compile(r"^---\n(.*?)\n---\n", re.DOTALL)
|
||||
_HEADING = re.compile(r"^#{1,6}\s+(.*\S)\s*$")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Chunk:
|
||||
"""A retrievable unit of text with its source page + section."""
|
||||
|
||||
title: str
|
||||
url: str
|
||||
source: str
|
||||
variant: str
|
||||
section: str
|
||||
text: str
|
||||
|
||||
|
||||
def parse_frontmatter(text: str) -> tuple[dict[str, str], str]:
|
||||
"""Return (metadata, body). Metadata values are unquoted single-line scalars."""
|
||||
match = _FRONTMATTER.match(text)
|
||||
if not match:
|
||||
return {}, text
|
||||
meta: dict[str, str] = {}
|
||||
for line in match.group(1).splitlines():
|
||||
if ":" in line:
|
||||
key, value = line.split(":", 1)
|
||||
meta[key.strip()] = value.strip().strip('"')
|
||||
return meta, text[match.end() :]
|
||||
|
||||
|
||||
def _split_sections(body: str) -> list[tuple[str, str]]:
|
||||
"""Group body lines under their most recent heading."""
|
||||
sections: list[tuple[str, str]] = []
|
||||
heading = "Overview"
|
||||
buffer: list[str] = []
|
||||
for line in body.splitlines():
|
||||
match = _HEADING.match(line)
|
||||
if match:
|
||||
if buffer:
|
||||
sections.append((heading, "\n".join(buffer).strip()))
|
||||
buffer = []
|
||||
heading = match.group(1)
|
||||
else:
|
||||
buffer.append(line)
|
||||
if buffer:
|
||||
sections.append((heading, "\n".join(buffer).strip()))
|
||||
return [(h, t) for h, t in sections if t]
|
||||
|
||||
|
||||
def _window(text: str, size: int, overlap: int) -> list[str]:
|
||||
"""Slice text into <=size windows overlapping by `overlap` chars."""
|
||||
if len(text) <= size:
|
||||
return [text]
|
||||
step = max(size - overlap, 1)
|
||||
return [text[start : start + size] for start in range(0, len(text), step)]
|
||||
|
||||
|
||||
def chunk_document(
|
||||
text: str,
|
||||
*,
|
||||
target_chars: int = config.CHUNK_TARGET_CHARS,
|
||||
overlap_chars: int = config.CHUNK_OVERLAP_CHARS,
|
||||
) -> list[Chunk]:
|
||||
"""Parse a corpus document into heading-scoped, size-bounded chunks."""
|
||||
meta, body = parse_frontmatter(text)
|
||||
title = meta.get("title", "")
|
||||
url = meta.get("url", "")
|
||||
source = meta.get("source", "mediawiki")
|
||||
page_variant = meta.get("variant", "unknown")
|
||||
chunks: list[Chunk] = []
|
||||
for heading, section_text in _split_sections(body):
|
||||
if heading.lower() in config.NOISE_SECTIONS:
|
||||
continue
|
||||
for piece in _window(section_text, target_chars, overlap_chars):
|
||||
stripped = piece.strip()
|
||||
if stripped:
|
||||
chunks.append(
|
||||
Chunk(
|
||||
title=title,
|
||||
url=url,
|
||||
source=source,
|
||||
variant=page_variant,
|
||||
section=heading,
|
||||
text=stripped,
|
||||
)
|
||||
)
|
||||
return chunks
|
||||
175
wikikb/config.py
Normal file
175
wikikb/config.py
Normal file
@ -0,0 +1,175 @@
|
||||
"""Central configuration: endpoints, paths, crawl politeness, model choice."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# --- Wiki endpoints ---------------------------------------------------------
|
||||
# Override WIKI_KB_HOST to re-point this KB at a different MediaWiki (scaffolding
|
||||
# rewrites the literal default below for a new wiki).
|
||||
WIKI_HOST = os.environ.get("WIKI_KB_HOST", "https://ck3.paradoxwikis.com").rstrip("/")
|
||||
# The API path varies by install: paradoxwikis / wiki.gg use /api.php, Wikipedia
|
||||
# uses /w/api.php. The article URL pattern varies too: /{Title} vs /wiki/{Title}.
|
||||
# Override both when re-pointing at a wiki that doesn't match the paradoxwikis layout.
|
||||
API_URL = os.environ.get("WIKI_KB_API_URL", f"{WIKI_HOST}/api.php")
|
||||
PAGE_URL_TEMPLATE = os.environ.get("WIKI_KB_PAGE_URL", f"{WIKI_HOST}/{{title}}")
|
||||
|
||||
# Namespace 0 is the main (content) namespace.
|
||||
CONTENT_NAMESPACE = 0
|
||||
|
||||
# Which page source to use: "mediawiki" (an open api.php — always preferred) or
|
||||
# "html" (sitemap.xml + content-container extraction, for wikis with no API at all).
|
||||
# scripts/detect_wiki.py probes a host and reports which applies.
|
||||
SOURCE_KIND = os.environ.get("WIKI_KB_SOURCE", "mediawiki")
|
||||
SITEMAP_URL = os.environ.get("WIKI_KB_SITEMAP", f"{WIKI_HOST}/sitemap.xml")
|
||||
|
||||
# --- Subpage policy -----------------------------------------------------------
|
||||
# Whether "A/B" titles are noise or content. This is PER WIKI and must be measured,
|
||||
# never assumed: on the CK3 wiki subpages are total-conversion mods (reject), while on
|
||||
# the KCD wiki they are the sequel's pages -- "Alchemy/KCD2" -- so rejecting them
|
||||
# deletes 441 pages of exactly the content you want (keep).
|
||||
# Counter(t.rsplit("/", 1)[-1] for t in titles if "/" in t)
|
||||
# answers it in one API pass.
|
||||
SUBPAGE_POLICY = os.environ.get("WIKI_KB_SUBPAGES", "reject")
|
||||
# Subpage tails dropped even when the policy is "keep" (e.g. "loc" localisation dumps).
|
||||
NOISE_SUBPAGE_TAILS = tuple(
|
||||
t.strip().lower() for t in os.environ.get("WIKI_KB_NOISE_TAILS", "").split(",") if t.strip()
|
||||
)
|
||||
|
||||
# --- Editions / variants ------------------------------------------------------
|
||||
# Some wikis cover two editions of one subject in a single namespace (a game and its
|
||||
# sequel, two versions of a product). Declare them as "name:PREFIX" pairs, MOST
|
||||
# SPECIFIC FIRST -- ordering is what lets "KCD" mean edition 1 without also matching
|
||||
# "KCD2". A page matches a variant if its title ends "/PREFIX" or a category starts
|
||||
# with PREFIX. Empty (the default) disables variant tagging entirely.
|
||||
VARIANTS = tuple(
|
||||
(part.split(":", 1)[0].strip(), part.split(":", 1)[1].strip())
|
||||
for part in os.environ.get("WIKI_KB_VARIANTS", "").split(",")
|
||||
if ":" in part
|
||||
)
|
||||
# Default search scope. "any" returns everything; a variant name excludes the OTHER
|
||||
# variants while always keeping pages tagged "both" or "unknown".
|
||||
DEFAULT_VARIANT = os.environ.get("WIKI_KB_DEFAULT_VARIANT", "any")
|
||||
|
||||
# Total-conversion mod pages live in ns0 as subpages ("CK3AGOT/Dragon") or under
|
||||
# these top-level prefixes. They are not base-game mechanics and pollute retrieval,
|
||||
# so the scraper skips them (any title containing "/" is also treated as a mod/subpage).
|
||||
_DEFAULT_MOD_PREFIXES = (
|
||||
"ck3agot",
|
||||
"lotr",
|
||||
"elder kings",
|
||||
"after the end",
|
||||
"the fallen eagle",
|
||||
"fallen eagle",
|
||||
"godherja",
|
||||
"princes of darkness",
|
||||
"rice",
|
||||
"realms in exile",
|
||||
"a game of thrones",
|
||||
"when the world stopped making sense",
|
||||
"more bookmarks",
|
||||
)
|
||||
# Override per-wiki with WIKI_KB_MOD_PREFIXES (comma-separated). Branchless so the
|
||||
# import line is always exercised regardless of whether the env var is set.
|
||||
MOD_TITLE_PREFIXES = tuple(
|
||||
p.strip().lower()
|
||||
for p in os.environ.get("WIKI_KB_MOD_PREFIXES", ",".join(_DEFAULT_MOD_PREFIXES)).split(",")
|
||||
if p.strip()
|
||||
)
|
||||
|
||||
# --- Politeness / crawl -----------------------------------------------------
|
||||
# Identify the bot honestly with a contact, per wiki etiquette.
|
||||
# MediaWiki API:Etiquette requires a descriptive User-Agent with a real contact.
|
||||
# Set WIKI_KB_CONTACT — the placeholder is deliberately obvious so an unconfigured
|
||||
# run is embarrassing rather than silently anonymous.
|
||||
CONTACT = os.environ.get("WIKI_KB_CONTACT", "UNCONFIGURED-set-WIKI_KB_CONTACT")
|
||||
USER_AGENT = f"wiki-kb/0.1 (personal offline RAG mirror; contact: {CONTACT})"
|
||||
|
||||
# Tool prefix for the MCP server: <namespace>_search / _get_page / _list_pages.
|
||||
MCP_NAMESPACE = os.environ.get("WIKI_KB_NAMESPACE", "wiki")
|
||||
# Minimum seconds between requests (single-flight, ~1 req/s).
|
||||
REQUEST_INTERVAL_S = 1.0
|
||||
# maxlag tells MediaWiki to reject our request if replication lag is high,
|
||||
# so we back off instead of adding load. See MediaWiki API:Etiquette.
|
||||
MAXLAG = 5
|
||||
REQUEST_TIMEOUT_S = 30.0
|
||||
MAX_RETRIES = 5
|
||||
BACKOFF_BASE_S = 2.0 # exponential: BACKOFF_BASE ** attempt
|
||||
|
||||
# --- Paths ------------------------------------------------------------------
|
||||
# WIKI_KB_ROOT isolates a KB's generated data (corpus/index/manifest) — lets one
|
||||
# codebase serve multiple wikis by pointing at different roots. Branchless default.
|
||||
_DEFAULT_ROOT = Path(__file__).resolve().parent.parent
|
||||
ROOT = Path(os.environ.get("WIKI_KB_ROOT", str(_DEFAULT_ROOT)))
|
||||
CORPUS_DIR = ROOT / "corpus"
|
||||
INDEX_DIR = ROOT / "index"
|
||||
DB_PATH = INDEX_DIR / "wiki.db"
|
||||
EMBEDDINGS_PATH = INDEX_DIR / "embeddings.npy"
|
||||
MANIFEST_PATH = ROOT / "manifest.json"
|
||||
|
||||
# --- Chunking ---------------------------------------------------------------
|
||||
CHUNK_TARGET_CHARS = 2000 # ~500 tokens
|
||||
CHUNK_OVERLAP_CHARS = 200
|
||||
# Section headings excluded from the index (the page is still mirrored in full and
|
||||
# still served by *_get_page). Wikis carry per-page boilerplate sections that are
|
||||
# both expensive to embed and harmful to retrieval — on the KCD wiki, "Other
|
||||
# languages" tables of localized UI names were 7% of all chunks and repeated every
|
||||
# page's title in eight languages. Empty by default; set per wiki after inspecting
|
||||
# the section histogram. Compared case-insensitively.
|
||||
NOISE_SECTIONS = frozenset(
|
||||
s.strip().lower() for s in os.environ.get("WIKI_KB_NOISE_SECTIONS", "").split(",") if s.strip()
|
||||
)
|
||||
|
||||
# --- Embeddings -------------------------------------------------------------
|
||||
EMBED_MODEL = "BAAI/bge-small-en-v1.5" # 384-dim, CPU-friendly via fastembed
|
||||
EMBED_DIM = 384
|
||||
EMBED_BATCH_SIZE = 32
|
||||
|
||||
|
||||
def _parse_core_spec(spec: str, total: int) -> set[int]:
|
||||
"""Parse "0-3,8" into a set of core ids clamped to [0, total)."""
|
||||
cores: set[int] = set()
|
||||
for part in spec.split(","):
|
||||
part = part.strip()
|
||||
if not part:
|
||||
continue
|
||||
if "-" in part:
|
||||
lo, hi = part.split("-", 1)
|
||||
cores.update(range(int(lo), int(hi) + 1))
|
||||
else:
|
||||
cores.add(int(part))
|
||||
return {c for c in cores if 0 <= c < total}
|
||||
|
||||
|
||||
def _allowed_cores(total: int, reserved: set[int]) -> list[int]:
|
||||
"""Cores usable by heavy jobs = all minus reserved (never empty)."""
|
||||
allowed = sorted(set(range(total)) - reserved)
|
||||
return allowed or list(range(total))
|
||||
|
||||
|
||||
# --- CPU isolation ----------------------------------------------------------
|
||||
# The bottleneck is core *contention* with the foreground game, not thread count.
|
||||
# Keep RESERVED_CORES free for the game and let heavy jobs use every other core
|
||||
# fully. Override WIKI_KB_RESERVED_CORES (e.g. "0-3" or "8,9"); set "" to use all cores.
|
||||
TOTAL_CPUS = os.cpu_count() or 4
|
||||
RESERVED_CORES = _parse_core_spec(os.environ.get("WIKI_KB_RESERVED_CORES", "0-3"), TOTAL_CPUS)
|
||||
ALLOWED_CORES = _allowed_cores(TOTAL_CPUS, RESERVED_CORES)
|
||||
# GENTLE BY DEFAULT: a running game floats across ALL cores, so heavy background
|
||||
# work drops FPS no matter which cores it "owns" — nice only softens it. Default to
|
||||
# a single-threaded, single-worker build (~1 core) pinned off RESERVED_CORES and
|
||||
# niced. When the game is CLOSED, set WIKI_KB_EMBED_PARALLEL=<n> for a fast build.
|
||||
EMBED_PARALLEL = int(os.environ.get("WIKI_KB_EMBED_PARALLEL", "1"))
|
||||
EMBED_THREADS = int(os.environ.get("WIKI_KB_EMBED_THREADS", "1"))
|
||||
# Lower scheduling priority so the game still wins on any shared core.
|
||||
NICE_INCREMENT = int(os.environ.get("WIKI_KB_NICE", "10"))
|
||||
|
||||
# --- Retrieval --------------------------------------------------------------
|
||||
RRF_K = 60 # Reciprocal Rank Fusion constant
|
||||
DEFAULT_TOP_K = 8
|
||||
CANDIDATE_POOL = 40 # per-retriever candidates before fusion
|
||||
|
||||
|
||||
def page_url(title: str) -> str:
|
||||
"""Public wiki URL for a page title."""
|
||||
return PAGE_URL_TEMPLATE.format(title=title.replace(" ", "_"))
|
||||
33
wikikb/cpu.py
Normal file
33
wikikb/cpu.py
Normal file
@ -0,0 +1,33 @@
|
||||
"""Keep heavy jobs (embedding) off reserved cores and at low priority.
|
||||
|
||||
The foreground game keeps its cores; this process is pinned to the rest and niced
|
||||
down so it yields on any shared core. Linux-only affinity is applied when present.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
|
||||
from . import config
|
||||
|
||||
|
||||
def isolate(
|
||||
*,
|
||||
allowed: Sequence[int] | None = None,
|
||||
nice_increment: int | None = None,
|
||||
) -> list[int]:
|
||||
"""Pin this process to `allowed` cores and lower its priority.
|
||||
|
||||
Returns the cores it was pinned to. No-ops the affinity call on platforms
|
||||
without ``os.sched_setaffinity`` (non-Linux).
|
||||
"""
|
||||
cores = list(config.ALLOWED_CORES if allowed is None else allowed)
|
||||
increment = config.NICE_INCREMENT if nice_increment is None else nice_increment
|
||||
|
||||
set_affinity = getattr(os, "sched_setaffinity", None)
|
||||
if set_affinity is not None:
|
||||
set_affinity(0, set(cores))
|
||||
if increment and hasattr(os, "nice"):
|
||||
os.nice(increment)
|
||||
return cores
|
||||
60
wikikb/embed.py
Normal file
60
wikikb/embed.py
Normal file
@ -0,0 +1,60 @@
|
||||
"""Local embeddings via fastembed (bge-small), L2-normalized for cosine = dot."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Iterable
|
||||
from typing import Protocol
|
||||
|
||||
import numpy as np
|
||||
|
||||
from . import config
|
||||
|
||||
|
||||
class EmbeddingModel(Protocol):
|
||||
"""Minimal fastembed-compatible surface (also satisfied by test fakes)."""
|
||||
|
||||
def embed(self, texts: Iterable[str], **kwargs: object) -> Iterable[np.ndarray]: ...
|
||||
def query_embed(self, texts: Iterable[str], **kwargs: object) -> Iterable[np.ndarray]: ...
|
||||
|
||||
|
||||
_model: EmbeddingModel | None = None
|
||||
|
||||
|
||||
def get_model() -> EmbeddingModel:
|
||||
"""Lazily construct the (heavy) fastembed model once."""
|
||||
global _model
|
||||
if _model is None: # pragma: no cover - exercised only with the real model
|
||||
from fastembed import TextEmbedding
|
||||
|
||||
# threads caps onnxruntime's intra-op parallelism (the main CPU driver);
|
||||
# CPUExecutionProvider avoids grabbing any accelerator.
|
||||
_model = TextEmbedding(
|
||||
model_name=config.EMBED_MODEL,
|
||||
threads=config.EMBED_THREADS,
|
||||
providers=["CPUExecutionProvider"],
|
||||
)
|
||||
return _model
|
||||
|
||||
|
||||
def _normalize(matrix: np.ndarray) -> np.ndarray:
|
||||
if matrix.size == 0:
|
||||
return matrix.reshape(0, config.EMBED_DIM).astype(np.float32)
|
||||
norms = np.linalg.norm(matrix, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1.0
|
||||
return (matrix / norms).astype(np.float32)
|
||||
|
||||
|
||||
def embed_passages(texts: list[str], *, model: EmbeddingModel | None = None) -> np.ndarray:
|
||||
"""Embed corpus passages -> (N, dim) normalized float32."""
|
||||
model = model or get_model()
|
||||
vectors = list(
|
||||
model.embed(texts, batch_size=config.EMBED_BATCH_SIZE, parallel=config.EMBED_PARALLEL)
|
||||
)
|
||||
return _normalize(np.array(vectors, dtype=np.float32))
|
||||
|
||||
|
||||
def embed_query(text: str, *, model: EmbeddingModel | None = None) -> np.ndarray:
|
||||
"""Embed a search query -> (dim,) normalized float32 (with bge query prefix)."""
|
||||
model = model or get_model()
|
||||
vector = np.array(next(iter(model.query_embed([text]))), dtype=np.float32)
|
||||
return _normalize(vector.reshape(1, -1))[0]
|
||||
167
wikikb/html2md.py
Normal file
167
wikikb/html2md.py
Normal file
@ -0,0 +1,167 @@
|
||||
"""Convert MediaWiki `action=parse` HTML into clean markdown with frontmatter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from bs4 import BeautifulSoup
|
||||
from markdownify import markdownify as _markdownify
|
||||
|
||||
# Chrome / navigation / reference elements we don't want in the corpus.
|
||||
_REMOVE_SELECTORS = (
|
||||
"style",
|
||||
"script",
|
||||
"link",
|
||||
".mw-editsection",
|
||||
".mw-jump-link",
|
||||
".noprint",
|
||||
".navbox",
|
||||
".toc",
|
||||
"#toc",
|
||||
".mw-empty-elt",
|
||||
".reference",
|
||||
".mw-references-wrap",
|
||||
".reflist",
|
||||
".ambox",
|
||||
".metadata",
|
||||
".pagetabber-absolute",
|
||||
".pi-image",
|
||||
)
|
||||
|
||||
_BLANK_LINES = re.compile(r"\n{3,}")
|
||||
|
||||
|
||||
def _infobox_rows(element, group: str, tab: str, rows: list[tuple[str, str, str, str]]) -> None:
|
||||
"""Collect (group, tab, label, value) from a portable-infobox subtree, in order.
|
||||
|
||||
Fandom stat blocks are an <aside> of heading/value *siblings*, not a table, so
|
||||
markdownify alone emits a stream of headings followed by orphaned numbers — the
|
||||
values survive but stop being attached to their labels. This walk restores the
|
||||
pairing before any markdown is generated.
|
||||
"""
|
||||
for child in element.find_all(recursive=False):
|
||||
classes = set(child.get("class") or [])
|
||||
if "pi-header" in classes:
|
||||
# A group heading scopes every following sibling until the next one.
|
||||
group = child.get_text(" ", strip=True)
|
||||
elif "wds-tabber" in classes:
|
||||
# Tier tabs (I/II/III/IV) are all inlined; without their label the four
|
||||
# stat blocks are indistinguishable.
|
||||
labels = [t.get_text(" ", strip=True) for t in child.select(".wds-tabs__tab-label")]
|
||||
for index, panel in enumerate(child.select(".wds-tab__content")):
|
||||
_infobox_rows(panel, group, labels[index] if index < len(labels) else "", rows)
|
||||
elif "pi-smart-group" in classes:
|
||||
# Horizontal group: labels in one row, values positionally in the next.
|
||||
labels = [h.get_text(" ", strip=True) for h in child.select(".pi-smart-data-label")]
|
||||
values = [v.get_text(" ", strip=True) for v in child.select(".pi-smart-data-value")]
|
||||
# strict=False: a hand-edited infobox with a stray label and no value
|
||||
# should degrade to the pairs that do line up, not abort the scrape.
|
||||
rows.extend(
|
||||
(group, tab, label, value) for label, value in zip(labels, values, strict=False)
|
||||
)
|
||||
elif "pi-data" in classes:
|
||||
label = child.select_one(".pi-data-label")
|
||||
value = child.select_one(".pi-data-value")
|
||||
if label is not None and value is not None:
|
||||
rows.append(
|
||||
(group, tab, label.get_text(" ", strip=True), value.get_text(" ", strip=True))
|
||||
)
|
||||
else:
|
||||
_infobox_rows(child, group, tab, rows)
|
||||
|
||||
|
||||
def _rows_to_html(rows: list[tuple[str, str, str, str]]) -> str:
|
||||
"""Render collected infobox rows as heading + table HTML, grouped in order."""
|
||||
parts: list[str] = []
|
||||
current: tuple[str, str] | None = None
|
||||
for group, tab, label, value in rows:
|
||||
if (group, tab) != current:
|
||||
if current is not None:
|
||||
parts.append("</table>")
|
||||
current = (group, tab)
|
||||
heading = "Infobox" if not group else f"Infobox: {group}"
|
||||
if tab:
|
||||
heading = f"{heading} ({tab})"
|
||||
parts.append(f"<h2>{heading}</h2>")
|
||||
parts.append("<table><tr><th>Field</th><th>Value</th></tr>")
|
||||
parts.append(f"<tr><td>{label}</td><td>{value}</td></tr>")
|
||||
if current is not None:
|
||||
parts.append("</table>")
|
||||
# The infobox emits headings, so any lead paragraph that follows it would
|
||||
# otherwise be filed under the LAST infobox group ("Infobox: Technical") —
|
||||
# mislabelling the page's defining sentence, which is the worst possible
|
||||
# chunk to misname given section text is weighted in both BM25 and the
|
||||
# embedding. Give that text its own section instead; if a real heading
|
||||
# follows immediately, this one holds no text and the chunker drops it.
|
||||
parts.append("<h2>Overview</h2>")
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def transform_infoboxes(root) -> None:
|
||||
"""Replace every portable-infobox in `root` with label/value tables, in place."""
|
||||
for aside in root.select("aside.portable-infobox"):
|
||||
rows: list[tuple[str, str, str, str]] = []
|
||||
_infobox_rows(aside, "", "", rows)
|
||||
replacement = _rows_to_html(rows)
|
||||
if replacement:
|
||||
aside.replace_with(BeautifulSoup(replacement, "html.parser"))
|
||||
else:
|
||||
# Theme/dummy infoboxes carry no data; dropping them avoids an empty
|
||||
# zero-width-space heading in the corpus.
|
||||
aside.decompose()
|
||||
|
||||
|
||||
def clean_html(html: str) -> str:
|
||||
"""Strip nav/edit/reference chrome, returning the article body HTML."""
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
root = soup.select_one(".mw-parser-output") or soup
|
||||
for selector in _REMOVE_SELECTORS:
|
||||
for element in root.select(selector):
|
||||
element.decompose()
|
||||
transform_infoboxes(root)
|
||||
return str(root)
|
||||
|
||||
|
||||
def html_to_markdown(html: str) -> str:
|
||||
"""Rendered HTML fragment -> collapsed ATX markdown (images dropped)."""
|
||||
cleaned = clean_html(html)
|
||||
markdown = _markdownify(cleaned, heading_style="ATX", strip=["img"])
|
||||
return _BLANK_LINES.sub("\n\n", markdown).strip()
|
||||
|
||||
|
||||
def _yaml_scalar(value: str) -> str:
|
||||
"""Quote a string so it is a safe single-line YAML scalar."""
|
||||
escaped = value.replace("\\", "\\\\").replace('"', '\\"')
|
||||
return f'"{escaped}"'
|
||||
|
||||
|
||||
def build_document(
|
||||
*,
|
||||
title: str,
|
||||
url: str,
|
||||
revid: int,
|
||||
categories: list[str],
|
||||
html: str,
|
||||
source: str = "mediawiki",
|
||||
variant: str = "unknown",
|
||||
) -> str:
|
||||
"""Assemble a corpus markdown document with YAML frontmatter."""
|
||||
body = html_to_markdown(html)
|
||||
category_list = ", ".join(_yaml_scalar(c) for c in categories)
|
||||
return "\n".join(
|
||||
(
|
||||
"---",
|
||||
f"title: {_yaml_scalar(title)}",
|
||||
f"url: {url}",
|
||||
f"revid: {revid}",
|
||||
f"source: {source}",
|
||||
f"variant: {variant}",
|
||||
f"categories: [{category_list}]",
|
||||
"---",
|
||||
"",
|
||||
f"# {title}",
|
||||
"",
|
||||
body,
|
||||
"",
|
||||
)
|
||||
)
|
||||
145
wikikb/indexer.py
Normal file
145
wikikb/indexer.py
Normal file
@ -0,0 +1,145 @@
|
||||
"""Build the search index: SQLite (chunks + FTS5) and a row-aligned embeddings file."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlite3
|
||||
import sys
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from . import config
|
||||
from .chunker import Chunk, chunk_document
|
||||
from .embed import embed_passages
|
||||
|
||||
_SCHEMA = """
|
||||
CREATE TABLE chunks (
|
||||
id INTEGER PRIMARY KEY,
|
||||
title TEXT NOT NULL,
|
||||
url TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
variant TEXT NOT NULL,
|
||||
section TEXT NOT NULL,
|
||||
text TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX chunks_source ON chunks(source);
|
||||
CREATE INDEX chunks_variant ON chunks(variant);
|
||||
CREATE VIRTUAL TABLE chunks_fts USING fts5(
|
||||
text, title, section,
|
||||
content='chunks', content_rowid='id'
|
||||
);
|
||||
"""
|
||||
|
||||
# Passages are stored at rowid = position+1; embeddings row `i` aligns with rowid i+1.
|
||||
EmbedFn = Callable[[list[str]], np.ndarray]
|
||||
|
||||
|
||||
def _embed_text(chunk: Chunk) -> str:
|
||||
"""Context-augmented text for embedding: page title + section + body.
|
||||
|
||||
The stored/displayed chunk text stays clean, but the embedding carries which
|
||||
page and section a chunk is from, so trait/topic queries rank the right page up
|
||||
instead of any page that merely mentions the query's generic words.
|
||||
"""
|
||||
return f"{chunk.title}. {chunk.section}. {chunk.text}"
|
||||
|
||||
|
||||
def embed_efficiently(texts: list[str], embed_fn: EmbedFn) -> np.ndarray:
|
||||
"""Embed passages, avoiding two large sources of wasted compute.
|
||||
|
||||
1. **Length bucketing.** fastembed pads every batch to its longest member, so one
|
||||
2000-char passage makes the other 31 in its batch pay that cost. Wiki chunks
|
||||
are wildly uneven (on the KCD wiki: 184-char median, 2000-char max, 53% under
|
||||
200 chars), so most of the compute was padding. Sorting by length before
|
||||
batching measured **2.7x faster** end to end.
|
||||
2. **Deduplication.** Identical passages are embedded once and the vector reused.
|
||||
Note this is usually a small win: ``_embed_text`` prefixes each passage with
|
||||
its page title, which makes per-page boilerplate distinct again (16% of raw
|
||||
chunk texts were duplicates on the KCD wiki, but only 0.1% after prefixing).
|
||||
|
||||
Vectors are returned in the caller's ORIGINAL order. That is load-bearing: the
|
||||
index's one alignment invariant is that embedding row ``i`` is chunk rowid
|
||||
``i+1``, so the sort must be undone here rather than leaking to the caller.
|
||||
"""
|
||||
if not texts:
|
||||
return embed_fn([])
|
||||
first_index: dict[str, int] = {}
|
||||
for text in texts:
|
||||
first_index.setdefault(text, len(first_index))
|
||||
unique = list(first_index)
|
||||
by_length = sorted(range(len(unique)), key=lambda i: len(unique[i]))
|
||||
sorted_vectors = embed_fn([unique[i] for i in by_length])
|
||||
unique_vectors = np.empty_like(sorted_vectors)
|
||||
unique_vectors[by_length] = sorted_vectors # undo the length sort
|
||||
return unique_vectors[[first_index[t] for t in texts]] # scatter duplicates back
|
||||
|
||||
|
||||
def collect_chunks(corpus_dir: Path) -> list[Chunk]:
|
||||
"""Chunk every markdown document in the corpus, in stable filename order."""
|
||||
chunks: list[Chunk] = []
|
||||
for doc in sorted(corpus_dir.glob("*.md")):
|
||||
chunks.extend(chunk_document(doc.read_text(encoding="utf-8")))
|
||||
return chunks
|
||||
|
||||
|
||||
def build_index(
|
||||
*,
|
||||
corpus_dir: Path | None = None,
|
||||
db_path: Path | None = None,
|
||||
embeddings_path: Path | None = None,
|
||||
embed_fn: EmbedFn = embed_passages,
|
||||
log: Callable[[str], None] = print,
|
||||
) -> int:
|
||||
"""(Re)build the index from the corpus. Returns the chunk count.
|
||||
|
||||
Paths default to the live ``config`` values, resolved at call time so tests
|
||||
can redirect them via monkeypatch.
|
||||
"""
|
||||
corpus_dir = corpus_dir or config.CORPUS_DIR
|
||||
db_path = db_path or config.DB_PATH
|
||||
embeddings_path = embeddings_path or config.EMBEDDINGS_PATH
|
||||
chunks = collect_chunks(corpus_dir)
|
||||
log(f"Chunked {len(chunks)} passages from {corpus_dir}")
|
||||
|
||||
vectors = embed_efficiently([_embed_text(c) for c in chunks], embed_fn)
|
||||
log(f"Embedded {vectors.shape[0]} passages (dim={vectors.shape[1] if vectors.size else 0})")
|
||||
|
||||
db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
db_path.unlink(missing_ok=True)
|
||||
con = sqlite3.connect(db_path)
|
||||
try:
|
||||
con.executescript(_SCHEMA)
|
||||
con.executemany(
|
||||
"INSERT INTO chunks (id, title, url, source, variant, section, text) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?, ?)",
|
||||
[
|
||||
(i + 1, c.title, c.url, c.source, c.variant, c.section, c.text)
|
||||
for i, c in enumerate(chunks)
|
||||
],
|
||||
)
|
||||
con.execute(
|
||||
"INSERT INTO chunks_fts (rowid, text, title, section) "
|
||||
"SELECT id, text, title, section FROM chunks"
|
||||
)
|
||||
con.commit()
|
||||
finally:
|
||||
con.close()
|
||||
|
||||
np.save(embeddings_path, vectors)
|
||||
log(f"Wrote {db_path} and {embeddings_path}")
|
||||
return len(chunks)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
count = build_index()
|
||||
print(f"\nIndex built: {count} chunks.")
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
from . import cpu
|
||||
|
||||
pinned = cpu.isolate()
|
||||
print(f"CPU-isolated to {len(pinned)} cores (reserved {sorted(config.RESERVED_CORES)}).")
|
||||
sys.exit(main())
|
||||
255
wikikb/scraper.py
Normal file
255
wikikb/scraper.py
Normal file
@ -0,0 +1,255 @@
|
||||
"""Mirror a wiki into `corpus/`.
|
||||
|
||||
Where pages come from is decided by a :mod:`.sources` Source — a MediaWiki API when
|
||||
one exists, a sitemap + HTML extraction when it does not. This module owns the parts
|
||||
that are the same either way: politeness (User-Agent, ~1 req/s, exponential backoff),
|
||||
noise filtering, resumability, and writing the corpus + manifest.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import sys
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
|
||||
import httpx
|
||||
|
||||
from . import config, variant
|
||||
from .html2md import build_document
|
||||
from .sources import HTML, MediaWikiSource, SitemapHtmlSource, Source
|
||||
|
||||
_UNSAFE_CHARS = re.compile(r'[/\\:*?"<>|\x00-\x1f]')
|
||||
|
||||
|
||||
class Throttle:
|
||||
"""Single-flight rate limiter maintaining a minimum request interval."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
interval: float = config.REQUEST_INTERVAL_S,
|
||||
*,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._interval = interval
|
||||
self._sleep = sleep
|
||||
self._clock = clock
|
||||
self._last = 0.0
|
||||
|
||||
def wait(self) -> None:
|
||||
now = self._clock()
|
||||
delay = self._interval - (now - self._last)
|
||||
if delay > 0:
|
||||
self._sleep(delay)
|
||||
self._last = self._clock()
|
||||
|
||||
|
||||
def is_core_title(title: str) -> bool:
|
||||
"""True for content pages worth mirroring; False for noise titles.
|
||||
|
||||
Whether a "/" in a title means noise is per-wiki and configured, never assumed —
|
||||
see ``config.SUBPAGE_POLICY``.
|
||||
"""
|
||||
if "/" in title:
|
||||
if title.rsplit("/", 1)[-1].lower() in config.NOISE_SUBPAGE_TAILS:
|
||||
return False
|
||||
if config.SUBPAGE_POLICY == "reject":
|
||||
return False
|
||||
lowered = title.lower()
|
||||
return not any(lowered.startswith(prefix) for prefix in config.MOD_TITLE_PREFIXES)
|
||||
|
||||
|
||||
def sanitize_filename(title: str) -> str:
|
||||
"""Map a page title to a readable, filesystem-safe ``.md`` filename."""
|
||||
name = _UNSAFE_CHARS.sub("_", title).replace(" ", "_").strip("._")
|
||||
return f"{name or 'untitled'}.md"
|
||||
|
||||
|
||||
def _api_get(
|
||||
client: httpx.Client,
|
||||
params: dict[str, object],
|
||||
*,
|
||||
throttle: Throttle,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
) -> dict:
|
||||
"""GET the API with maxlag, retrying on 429/503 and maxlag errors."""
|
||||
query = {
|
||||
**params,
|
||||
"format": "json",
|
||||
"formatversion": "2",
|
||||
"maxlag": config.MAXLAG,
|
||||
}
|
||||
for attempt in range(config.MAX_RETRIES):
|
||||
throttle.wait()
|
||||
response = client.get(config.API_URL, params=query)
|
||||
if response.status_code in (429, 503):
|
||||
sleep(config.BACKOFF_BASE_S**attempt)
|
||||
continue
|
||||
response.raise_for_status()
|
||||
try:
|
||||
data = response.json()
|
||||
except ValueError:
|
||||
# Wikis behind Cloudflare (Fandom, many wiki.gg installs) serve a bot
|
||||
# challenge as a 200 with an HTML body, so raise_for_status() passes it
|
||||
# straight through. Back off and retry rather than letting a challenge
|
||||
# page be converted to markdown and written into the corpus.
|
||||
sleep(config.BACKOFF_BASE_S**attempt)
|
||||
continue
|
||||
error = data.get("error")
|
||||
if error and error.get("code") == "maxlag":
|
||||
sleep(config.BACKOFF_BASE_S**attempt)
|
||||
continue
|
||||
return data
|
||||
raise RuntimeError(f"API request failed after {config.MAX_RETRIES} retries: {params}")
|
||||
|
||||
|
||||
def build_source(
|
||||
client: httpx.Client,
|
||||
*,
|
||||
throttle: Throttle,
|
||||
kind: str | None = None,
|
||||
) -> Source:
|
||||
"""Construct the page source for the configured wiki.
|
||||
|
||||
Defaults to the MediaWiki API, which is preferred whenever it exists: it gives a
|
||||
real page list, expanded templates, categories and revision ids. ``kind="html"``
|
||||
selects the sitemap fallback for wikis with no API.
|
||||
"""
|
||||
if (kind or config.SOURCE_KIND) == HTML:
|
||||
return SitemapHtmlSource(
|
||||
lambda url: _get(client, url, throttle=throttle),
|
||||
sitemap_url=config.SITEMAP_URL,
|
||||
)
|
||||
return MediaWikiSource(lambda params: _api_get(client, params, throttle=throttle))
|
||||
|
||||
|
||||
def _get(
|
||||
client: httpx.Client,
|
||||
url: str,
|
||||
*,
|
||||
throttle: Throttle,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
) -> httpx.Response:
|
||||
"""Throttled GET with the same backoff policy as the API path.
|
||||
|
||||
404 is returned to the caller rather than retried: for the sitemap source a
|
||||
missing page is a normal outcome, not a transient failure.
|
||||
"""
|
||||
for attempt in range(config.MAX_RETRIES):
|
||||
throttle.wait()
|
||||
response = client.get(url, follow_redirects=True)
|
||||
if response.status_code in (429, 503):
|
||||
sleep(config.BACKOFF_BASE_S**attempt)
|
||||
continue
|
||||
if response.status_code == 404:
|
||||
return response
|
||||
response.raise_for_status()
|
||||
return response
|
||||
raise RuntimeError(f"GET failed after {config.MAX_RETRIES} retries: {url}")
|
||||
|
||||
|
||||
def load_manifest() -> dict[str, dict]:
|
||||
if config.MANIFEST_PATH.exists():
|
||||
return json.loads(config.MANIFEST_PATH.read_text(encoding="utf-8"))
|
||||
return {}
|
||||
|
||||
|
||||
def save_manifest(manifest: dict[str, dict]) -> None:
|
||||
config.MANIFEST_PATH.write_text(
|
||||
json.dumps(manifest, indent=2, sort_keys=True, ensure_ascii=False),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
|
||||
def scrape(
|
||||
*,
|
||||
limit: int | None = None,
|
||||
refresh: bool = False,
|
||||
titles: list[str] | None = None,
|
||||
client: httpx.Client | None = None,
|
||||
throttle: Throttle | None = None,
|
||||
source: Source | None = None,
|
||||
log: Callable[[str], None] = print,
|
||||
) -> dict[str, int]:
|
||||
"""Mirror the wiki. Returns counts of {written, skipped, missing}."""
|
||||
owns_client = client is None
|
||||
client = client or httpx.Client(
|
||||
headers={"User-Agent": config.USER_AGENT},
|
||||
timeout=config.REQUEST_TIMEOUT_S,
|
||||
)
|
||||
throttle = throttle or Throttle()
|
||||
source = source or build_source(client, throttle=throttle)
|
||||
selected = (
|
||||
(t for t in titles if is_core_title(t))
|
||||
if titles is not None
|
||||
else source.iter_titles(limit=limit)
|
||||
)
|
||||
config.CORPUS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
manifest = load_manifest()
|
||||
written = skipped = missing = 0
|
||||
try:
|
||||
for title in selected:
|
||||
if titles is None and not is_core_title(title):
|
||||
continue
|
||||
filename = sanitize_filename(title)
|
||||
target = config.CORPUS_DIR / filename
|
||||
if not refresh and title in manifest and target.exists():
|
||||
skipped += 1
|
||||
continue
|
||||
page = source.fetch_page(title)
|
||||
if page is None:
|
||||
missing += 1
|
||||
log(f"MISSING {title}")
|
||||
continue
|
||||
page_variant = variant.classify(page.title, page.categories)
|
||||
document = build_document(
|
||||
variant=page_variant,
|
||||
title=page.title,
|
||||
url=page.url,
|
||||
revid=page.revid,
|
||||
categories=page.categories,
|
||||
source=page.source,
|
||||
html=page.html,
|
||||
)
|
||||
target.write_text(document, encoding="utf-8")
|
||||
manifest[page.title] = {
|
||||
"revid": page.revid,
|
||||
"filename": filename,
|
||||
"url": page.url,
|
||||
"categories": page.categories,
|
||||
"source": page.source,
|
||||
"variant": page_variant,
|
||||
}
|
||||
written += 1
|
||||
log(f"[{written:>5}] {page.title}")
|
||||
if written % 25 == 0:
|
||||
save_manifest(manifest)
|
||||
finally:
|
||||
save_manifest(manifest)
|
||||
if owns_client:
|
||||
client.close()
|
||||
return {"written": written, "skipped": skipped, "missing": missing}
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = argparse.ArgumentParser(description="Mirror the CK3 wiki into corpus/.")
|
||||
parser.add_argument("--limit", type=int, default=None, help="max pages to fetch")
|
||||
parser.add_argument("--refresh", action="store_true", help="re-fetch pages already mirrored")
|
||||
parser.add_argument(
|
||||
"--titles", default=None, help="comma-separated titles to mirror instead of the whole wiki"
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
named = [t.strip() for t in args.titles.split(",") if t.strip()] if args.titles else None
|
||||
result = scrape(limit=args.limit, refresh=args.refresh, titles=named)
|
||||
print(
|
||||
f"\nDone: {result['written']} written, "
|
||||
f"{result['skipped']} skipped, {result['missing']} missing."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
174
wikikb/search.py
Normal file
174
wikikb/search.py
Normal file
@ -0,0 +1,174 @@
|
||||
"""Hybrid retrieval: FTS5 (BM25) + embedding cosine, fused with Reciprocal Rank Fusion."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import sqlite3
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from . import config
|
||||
from . import variant as variant_mod
|
||||
from .embed import EmbeddingModel, embed_query
|
||||
|
||||
_WORD = re.compile(r"\w+", re.UNICODE)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Result:
|
||||
"""A ranked retrieval hit."""
|
||||
|
||||
title: str
|
||||
url: str
|
||||
source: str
|
||||
variant: str
|
||||
section: str
|
||||
text: str
|
||||
score: float
|
||||
|
||||
|
||||
def _fts_match_query(query: str) -> str:
|
||||
"""Build an FTS5 MATCH: OR of quoted terms so any term can hit.
|
||||
|
||||
OR (not AND) keeps recall high on natural-language queries; BM25's IDF then
|
||||
ranks chunks that match the rarer, more meaningful terms above those that only
|
||||
share common words.
|
||||
"""
|
||||
return " OR ".join(f'"{word}"' for word in _WORD.findall(query.lower()))
|
||||
|
||||
|
||||
def _allowed_ids(con: sqlite3.Connection, scope: str) -> set[int] | None:
|
||||
"""Chunk ids in scope, or None when unrestricted."""
|
||||
if scope == "any" or not config.VARIANTS:
|
||||
return None
|
||||
rows = con.execute("SELECT id, variant FROM chunks").fetchall()
|
||||
return {int(r[0]) for r in rows if variant_mod.matches(r[1], scope)}
|
||||
|
||||
|
||||
def _fts_ids(
|
||||
con: sqlite3.Connection, query: str, limit: int, allowed: set[int] | None
|
||||
) -> list[int]:
|
||||
match = _fts_match_query(query)
|
||||
if not match:
|
||||
return []
|
||||
# Weight matches in the title/section columns above body text, so a page whose
|
||||
# heading is literally about the query terms outranks one that merely mentions
|
||||
# them in passing. Weights map to the FTS columns in order: (text, title, section).
|
||||
# Filter BEFORE the LIMIT: truncating first would leave only a handful of hits
|
||||
# whenever the excluded edition dominates the lexical matches for a query.
|
||||
rows = con.execute(
|
||||
"SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? "
|
||||
"ORDER BY bm25(chunks_fts, 1.0, 5.0, 3.0)",
|
||||
(match,),
|
||||
).fetchall()
|
||||
ids = [int(r[0]) for r in rows]
|
||||
if allowed is not None:
|
||||
ids = [i for i in ids if i in allowed]
|
||||
return ids[:limit]
|
||||
|
||||
|
||||
def _vector_ids(
|
||||
vectors: np.ndarray, query_vec: np.ndarray, limit: int, allowed: set[int] | None
|
||||
) -> list[int]:
|
||||
if vectors.shape[0] == 0:
|
||||
return []
|
||||
sims = vectors @ query_vec
|
||||
out: list[int] = []
|
||||
for index in np.argsort(-sims):
|
||||
rowid = int(index) + 1 # rowid = embedding index + 1
|
||||
if allowed is None or rowid in allowed:
|
||||
out.append(rowid)
|
||||
if len(out) == limit:
|
||||
break
|
||||
return out
|
||||
|
||||
|
||||
def reciprocal_rank_fusion(rankings: list[list[int]], k: int = config.RRF_K) -> list[int]:
|
||||
"""Fuse ranked id lists; higher fused score = better."""
|
||||
scores: dict[int, float] = {}
|
||||
for ranking in rankings:
|
||||
for rank, cid in enumerate(ranking):
|
||||
scores[cid] = scores.get(cid, 0.0) + 1.0 / (k + rank + 1)
|
||||
return sorted(scores, key=lambda cid: -scores[cid])
|
||||
|
||||
|
||||
def search(
|
||||
query: str,
|
||||
*,
|
||||
k: int = config.DEFAULT_TOP_K,
|
||||
variant: str | None = None,
|
||||
db_path: Path | None = None,
|
||||
embeddings_path: Path | None = None,
|
||||
model: EmbeddingModel | None = None,
|
||||
) -> list[Result]:
|
||||
"""Run hybrid retrieval and return the top-k fused results.
|
||||
|
||||
Paths default to the live ``config`` values, resolved at call time.
|
||||
"""
|
||||
scope = variant or config.DEFAULT_VARIANT
|
||||
db_path = db_path or config.DB_PATH
|
||||
embeddings_path = embeddings_path or config.EMBEDDINGS_PATH
|
||||
con = sqlite3.connect(db_path)
|
||||
try:
|
||||
allowed = _allowed_ids(con, scope)
|
||||
fts = _fts_ids(con, query, config.CANDIDATE_POOL, allowed)
|
||||
vectors = np.load(embeddings_path)
|
||||
vec = _vector_ids(vectors, embed_query(query, model=model), config.CANDIDATE_POOL, allowed)
|
||||
fused = reciprocal_rank_fusion([fts, vec])[:k]
|
||||
if not fused:
|
||||
return []
|
||||
placeholders = ",".join("?" * len(fused))
|
||||
rows = {
|
||||
row[0]: row
|
||||
for row in con.execute(
|
||||
"SELECT id, title, url, source, variant, section, text FROM chunks "
|
||||
f"WHERE id IN ({placeholders})",
|
||||
fused,
|
||||
).fetchall()
|
||||
}
|
||||
finally:
|
||||
con.close()
|
||||
|
||||
results: list[Result] = []
|
||||
for rank, cid in enumerate(fused):
|
||||
row = rows.get(cid)
|
||||
if row is None: # pragma: no cover - defensive; ids come from the same db
|
||||
continue
|
||||
results.append(
|
||||
Result(
|
||||
title=row[1],
|
||||
url=row[2],
|
||||
source=row[3],
|
||||
variant=row[4],
|
||||
section=row[5],
|
||||
text=row[6],
|
||||
score=1.0 / (rank + 1),
|
||||
)
|
||||
)
|
||||
return results
|
||||
|
||||
|
||||
def _format(results: list[Result]) -> str:
|
||||
if not results:
|
||||
return "(no results)"
|
||||
blocks = []
|
||||
for i, r in enumerate(results, 1):
|
||||
snippet = r.text if len(r.text) <= 500 else r.text[:500] + "…"
|
||||
blocks.append(f"{i}. [{r.variant}] {r.title} > {r.section}\n {r.url}\n {snippet}")
|
||||
return "\n\n".join(blocks)
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
args = sys.argv[1:] if argv is None else argv
|
||||
if not args:
|
||||
print('usage: run.sh search "<query>"', file=sys.stderr)
|
||||
return 2
|
||||
print(_format(search(" ".join(args))))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
84
wikikb/server.py
Normal file
84
wikikb/server.py
Normal file
@ -0,0 +1,84 @@
|
||||
"""MCP server exposing the wiki knowledge base to Claude Code.
|
||||
|
||||
Tools:
|
||||
wiki_search(query, k) hybrid BM25 + embedding search over the wiki
|
||||
wiki_get_page(title) full markdown of one mirrored page
|
||||
wiki_list_pages(prefix) enumerate mirrored page titles
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from typing import Any
|
||||
|
||||
from . import config
|
||||
from .scraper import load_manifest, sanitize_filename
|
||||
from .search import search
|
||||
|
||||
|
||||
def wiki_search(
|
||||
query: str, k: int = config.DEFAULT_TOP_K, variant: str = config.DEFAULT_VARIANT
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Search the wiki knowledge base (hybrid BM25 + embeddings).
|
||||
|
||||
``variant`` scopes results when the wiki covers two editions of one subject (see
|
||||
wikikb.variant); it excludes the other edition while always keeping pages that
|
||||
apply to both or carry no edition tag. Every hit reports its source and variant,
|
||||
so a result can never be silently attributed to the wrong wiki or edition.
|
||||
|
||||
Returns ranked hits with title, source, variant, section, url and passage text.
|
||||
"""
|
||||
return [
|
||||
{
|
||||
"title": r.title,
|
||||
"source": r.source,
|
||||
"variant": r.variant,
|
||||
"section": r.section,
|
||||
"url": r.url,
|
||||
"text": r.text,
|
||||
}
|
||||
for r in search(query, k=k, variant=variant)
|
||||
]
|
||||
|
||||
|
||||
def wiki_get_page(title: str) -> str:
|
||||
"""Return the full markdown of a mirrored wiki page by exact title."""
|
||||
entry = load_manifest().get(title)
|
||||
filename = entry["filename"] if entry else sanitize_filename(title)
|
||||
path = config.CORPUS_DIR / filename
|
||||
if not path.exists():
|
||||
return f"No mirrored page found for {title!r}. Try wiki_list_pages or wiki_search."
|
||||
return path.read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def wiki_list_pages(prefix: str = "", limit: int = 100) -> list[str]:
|
||||
"""List mirrored wiki page titles, filtered by case-insensitive prefix."""
|
||||
titles = sorted(load_manifest())
|
||||
if prefix:
|
||||
lowered = prefix.lower()
|
||||
titles = [t for t in titles if t.lower().startswith(lowered)]
|
||||
return titles[:limit]
|
||||
|
||||
|
||||
def build_server(server: Any = None) -> Any:
|
||||
"""Construct (or reuse) a FastMCP server with the wiki tools registered."""
|
||||
if server is None: # pragma: no cover - real server needs the mcp runtime
|
||||
from mcp.server.fastmcp import FastMCP
|
||||
|
||||
server = FastMCP(config.MCP_NAMESPACE)
|
||||
server.tool()(wiki_search)
|
||||
server.tool()(wiki_get_page)
|
||||
server.tool()(wiki_list_pages)
|
||||
return server
|
||||
|
||||
|
||||
def main() -> int: # pragma: no cover - starts the stdio transport loop
|
||||
from . import cpu
|
||||
|
||||
cpu.isolate() # queries are light, but still keep off the reserved game cores
|
||||
build_server().run()
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
228
wikikb/sources.py
Normal file
228
wikikb/sources.py
Normal file
@ -0,0 +1,228 @@
|
||||
"""Pluggable page sources — where a wiki's pages come from.
|
||||
|
||||
Everything downstream of "fetch -> HTML" is already source-agnostic: html2md,
|
||||
chunker, embed, indexer, search and server never learn where the HTML originated.
|
||||
The wiki-engine assumption is confined to exactly two operations, which is what this
|
||||
module abstracts:
|
||||
|
||||
* **enumerate** the content pages, and
|
||||
* **fetch** one page as server-rendered HTML.
|
||||
|
||||
Two implementations ship:
|
||||
|
||||
* :class:`MediaWikiSource` — the good case. An open ``api.php`` gives a real page
|
||||
list, server-rendered HTML with templates expanded, categories, revision ids for
|
||||
fidelity diffing, and ``maxlag`` so the server can push back.
|
||||
* :class:`SitemapHtmlSource` — the fallback for wikis with no API at all (custom
|
||||
engines such as Fextralife). Enumerates ``sitemap.xml`` and extracts a content
|
||||
container from each page. Works, but gives up revision ids and ``maxlag``.
|
||||
|
||||
Use :func:`detect_source` to pick one by probing the host.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import urllib.parse
|
||||
from collections.abc import Callable, Iterator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol
|
||||
|
||||
import httpx
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
from . import config
|
||||
|
||||
MEDIAWIKI = "mediawiki"
|
||||
HTML = "html"
|
||||
|
||||
# Candidate content containers for HTML wikis, most specific first. Fextralife uses
|
||||
# #main-content; MediaWiki mirrors use .mw-parser-output; the rest are conventional.
|
||||
_CONTENT_SELECTORS = (
|
||||
"#main-content",
|
||||
"#wiki-content-block",
|
||||
".mw-parser-output",
|
||||
"#content",
|
||||
"main",
|
||||
"article",
|
||||
)
|
||||
_LOC = re.compile(r"<loc>\s*([^<\s]+)\s*</loc>", re.IGNORECASE)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PageDoc:
|
||||
"""A fetched page ready to be written to the corpus."""
|
||||
|
||||
title: str
|
||||
url: str
|
||||
html: str
|
||||
source: str
|
||||
# Revision id, or 0 for sources that have no revision concept. A real revid lets
|
||||
# the fidelity check say "the wiki was edited" rather than "our extractor broke";
|
||||
# without one, only a content diff is available.
|
||||
revid: int = 0
|
||||
categories: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
class Source(Protocol):
|
||||
"""Where pages come from. Implementations must be polite and resumable."""
|
||||
|
||||
name: str
|
||||
|
||||
def iter_titles(self, *, limit: int | None = None) -> Iterator[str]:
|
||||
"""Yield content-page titles."""
|
||||
...
|
||||
|
||||
def fetch_page(self, title: str) -> PageDoc | None:
|
||||
"""Fetch one page, or None if it is missing."""
|
||||
...
|
||||
|
||||
|
||||
class MediaWikiSource:
|
||||
"""Pages from a MediaWiki ``api.php``: page list + server-rendered HTML."""
|
||||
|
||||
name = MEDIAWIKI
|
||||
|
||||
def __init__(self, api_get: Callable[..., dict], *, api_url: str | None = None) -> None:
|
||||
self._api_get = api_get
|
||||
self._api_url = api_url or config.API_URL
|
||||
|
||||
def iter_titles(self, *, limit: int | None = None) -> Iterator[str]:
|
||||
apcontinue: str | None = None
|
||||
seen = 0
|
||||
while True:
|
||||
params: dict[str, object] = {
|
||||
"action": "query",
|
||||
"list": "allpages",
|
||||
"apnamespace": config.CONTENT_NAMESPACE,
|
||||
"apfilterredir": "nonredirects",
|
||||
"aplimit": "500",
|
||||
}
|
||||
if apcontinue is not None:
|
||||
params["apcontinue"] = apcontinue
|
||||
data = self._api_get(params)
|
||||
for page in data["query"]["allpages"]:
|
||||
yield page["title"]
|
||||
seen += 1
|
||||
if limit is not None and seen >= limit:
|
||||
return
|
||||
cont = data.get("continue")
|
||||
if not cont:
|
||||
return
|
||||
apcontinue = cont["apcontinue"]
|
||||
|
||||
def fetch_page(self, title: str) -> PageDoc | None:
|
||||
data = self._api_get(
|
||||
{
|
||||
"action": "parse",
|
||||
"page": title,
|
||||
"prop": "text|categories|revid",
|
||||
"redirects": "1",
|
||||
}
|
||||
)
|
||||
if "error" in data:
|
||||
return None
|
||||
parse = data["parse"]
|
||||
return PageDoc(
|
||||
title=parse["title"],
|
||||
url=config.page_url(parse["title"]),
|
||||
html=parse["text"],
|
||||
source=self.name,
|
||||
revid=parse["revid"],
|
||||
categories=[c["category"].replace("_", " ") for c in parse.get("categories", [])],
|
||||
)
|
||||
|
||||
|
||||
class SitemapHtmlSource:
|
||||
"""Pages from ``sitemap.xml`` + a content container in each page's HTML.
|
||||
|
||||
For wikis with no API. Two capabilities are genuinely lost and callers should not
|
||||
pretend otherwise: there is no revision id (``revid`` stays 0), and there is no
|
||||
``maxlag``, so the rate limit is purely self-imposed with no server backpressure.
|
||||
"""
|
||||
|
||||
name = HTML
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
get: Callable[[str], httpx.Response],
|
||||
*,
|
||||
sitemap_url: str | None = None,
|
||||
selectors: tuple[str, ...] = _CONTENT_SELECTORS,
|
||||
) -> None:
|
||||
self._get = get
|
||||
self._sitemap_url = sitemap_url or f"{config.WIKI_HOST}/sitemap.xml"
|
||||
self._selectors = selectors
|
||||
self._urls: dict[str, str] = {}
|
||||
|
||||
@staticmethod
|
||||
def title_from_url(url: str) -> str:
|
||||
"""Last path segment as a human title ("/Common+Longsword" -> "Common Longsword")."""
|
||||
slug = urllib.parse.urlparse(url).path.rstrip("/").rsplit("/", 1)[-1]
|
||||
return urllib.parse.unquote_plus(slug).replace("_", " ").strip()
|
||||
|
||||
def iter_titles(self, *, limit: int | None = None) -> Iterator[str]:
|
||||
text = self._get(self._sitemap_url).text
|
||||
locs = _LOC.findall(text)
|
||||
# A sitemap index points at further sitemaps rather than pages.
|
||||
if locs and all(loc.endswith((".xml", ".xml.gz")) for loc in locs):
|
||||
nested: list[str] = []
|
||||
for child in locs:
|
||||
nested.extend(_LOC.findall(self._get(child).text))
|
||||
locs = nested
|
||||
seen = 0
|
||||
for url in locs:
|
||||
title = self.title_from_url(url)
|
||||
if not title:
|
||||
continue
|
||||
self._urls.setdefault(title, url)
|
||||
yield title
|
||||
seen += 1
|
||||
if limit is not None and seen >= limit:
|
||||
return
|
||||
|
||||
def _extract(self, html: str) -> str | None:
|
||||
soup = BeautifulSoup(html, "html.parser")
|
||||
for selector in self._selectors:
|
||||
node = soup.select_one(selector)
|
||||
if node is not None:
|
||||
return str(node)
|
||||
return None
|
||||
|
||||
def fetch_page(self, title: str) -> PageDoc | None:
|
||||
url = self._urls.get(title) or f"{config.WIKI_HOST}/{urllib.parse.quote(title)}"
|
||||
response = self._get(url)
|
||||
if response.status_code == 404:
|
||||
return None
|
||||
body = self._extract(response.text)
|
||||
if body is None:
|
||||
return None
|
||||
return PageDoc(title=title, url=url, html=body, source=self.name)
|
||||
|
||||
|
||||
def detect_source_kind(get: Callable[[str], httpx.Response], host: str) -> tuple[str, str]:
|
||||
"""Probe a host and return (kind, endpoint).
|
||||
|
||||
Order matters: an open MediaWiki API beats HTML scraping every time, so try both
|
||||
conventional API paths ("/api.php" on Fandom, wiki.gg and paradoxwikis;
|
||||
"/w/api.php" on Wikipedia) before falling back to the sitemap. Raises if neither
|
||||
is available — that is a real answer, not a failure to be papered over.
|
||||
"""
|
||||
for path in ("/api.php", "/w/api.php"):
|
||||
endpoint = f"{host}{path}"
|
||||
try:
|
||||
response = get(f"{endpoint}?action=query&meta=siteinfo&format=json")
|
||||
if response.status_code == 200 and "mediawiki" in (
|
||||
response.json().get("query", {}).get("general", {}).get("generator", "").lower()
|
||||
):
|
||||
return MEDIAWIKI, endpoint
|
||||
except (httpx.HTTPError, ValueError, KeyError):
|
||||
continue
|
||||
sitemap = f"{host}/sitemap.xml"
|
||||
response = get(sitemap)
|
||||
if response.status_code == 200 and "<loc>" in response.text.lower():
|
||||
return HTML, sitemap
|
||||
raise RuntimeError(
|
||||
f"{host} exposes neither a MediaWiki API nor a sitemap.xml; "
|
||||
"it cannot be mirrored by this tool."
|
||||
)
|
||||
59
wikikb/variant.py
Normal file
59
wikikb/variant.py
Normal file
@ -0,0 +1,59 @@
|
||||
"""Tag pages with which edition of a subject they describe.
|
||||
|
||||
Some wikis cover two editions in one namespace — a game and its sequel, two versions
|
||||
of a product. The Kingdom Come: Deliverance wiki is the worked example: the original
|
||||
game's `Alchemy` and the sequel's `Alchemy/KCD2` sit side by side, and a search that
|
||||
cannot tell them apart will happily answer a question about one with the other's
|
||||
numbers.
|
||||
|
||||
Configured by ``config.VARIANTS`` as ordered ``(name, prefix)`` pairs, most specific
|
||||
first. A page matches a variant if its title ends ``/PREFIX`` or a category starts
|
||||
with ``PREFIX``, ignoring categories already claimed by an earlier, more specific
|
||||
prefix — that ordering is what lets ``KCD`` mean edition 1 without also matching
|
||||
``KCD2``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from . import config
|
||||
|
||||
BOTH = "both"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
def _matches(prefix: str, title: str, categories: list[str], more_specific: list[str]) -> bool:
|
||||
if title.endswith(f"/{prefix}"):
|
||||
return True
|
||||
return any(
|
||||
c.startswith(prefix) and not any(c.startswith(other) for other in more_specific)
|
||||
for c in categories
|
||||
)
|
||||
|
||||
|
||||
def classify(title: str, categories: list[str]) -> str:
|
||||
"""Return a variant name, "both", or "unknown".
|
||||
|
||||
``unknown`` is a real answer, not a failure: on the KCD wiki ~9% of pages carry no
|
||||
edition-specific category yet remain relevant to both ("Moravia", "Codex").
|
||||
Callers must treat it as in scope for every variant filter.
|
||||
"""
|
||||
hits = [
|
||||
name
|
||||
for index, (name, prefix) in enumerate(config.VARIANTS)
|
||||
if _matches(prefix, title, categories, [p for _, p in config.VARIANTS[:index]])
|
||||
]
|
||||
if len(hits) > 1:
|
||||
return BOTH
|
||||
return hits[0] if hits else UNKNOWN
|
||||
|
||||
|
||||
def matches(variant: str, scope: str) -> bool:
|
||||
"""True if a page tagged ``variant`` is in scope for filter ``scope``.
|
||||
|
||||
``scope="any"`` returns everything. Otherwise the scope EXCLUDES the other
|
||||
editions rather than filtering TO one — dropping ``both``/``unknown`` pages would
|
||||
be a recall hole, not a precision win.
|
||||
"""
|
||||
if scope == "any":
|
||||
return True
|
||||
return variant in (scope, BOTH, UNKNOWN)
|
||||
Loading…
Reference in New Issue
Block a user