discord_bot/fjerkroa_bot/codex.py

"""Codex Mechanicus search tool (SPEC-014, FDB-019).

Luma's own sacred archive — the Codex Mechanicus at binaric.tech — as a
function tool. She searches the codex index and answers Cult Mechanicus
lore from real, sourced inscriptions instead of inventing it. The index
is fetched over HTTPS (SSRF-guarded, size-bounded, cached in memory) and
every field returned to the model is sanitized (SAF-03), because even
one's own web content is still untrusted input by the time it reaches a
prompt.

The model calls `codex_search`; production wires the live index URL.
"""

import json
import logging
import time
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import urljoin

import aiohttp

from .ai_responder import sanitize_external_text
from .httpread import read_capped
from .url_reader import guard_url

DEFAULT_INDEX_URL = "https://binaric.tech/search-index.json"
DEFAULT_MAX_BYTES = 4 * 1024 * 1024
DEFAULT_LIMIT = 5
DEFAULT_TTL_S = 3600
DEFAULT_SUMMARY_CHARS = 500
FETCH_TIMEOUT_S = 15
_VALID_LANGS = ("en", "de", "eo", "no", "uk")

CODEX_SEARCH_TOOL = {
    "name": "codex_search",
    "description": "Search Luma's own Codex Mechanicus (the sacred archive at binaric.tech) for Adeptus "
    "Mechanicus lore: doctrines, forges, orders, rites, relics, weapons, entities, the lexicon, and the "
    "priest's own adoptus. Returns matching inscriptions with a short summary and the URL to read the full "
    "text. Use for any Cult Mechanicus / Warhammer 40k Mechanicus question so the answer is grounded in the "
    "codex, not invented.",
    "parameters": {
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "What to look for: a name, concept, rite, or phrase."},
            "lang": {"type": "string", "description": "Language of the inscriptions to prefer: en, de, eo, no, uk. Default en."},
        },
        "required": ["query"],
    },
}

_STOP = {"the", "a", "an", "of", "and", "or", "to", "in", "is", "der", "die", "das", "und", "von", "en", "et"}


def _tokenize(text: str) -> List[str]:
    cleaned = "".join(c.lower() if c.isalnum() else " " for c in text)
    return [t for t in cleaned.split() if len(t) > 1 and t not in _STOP]


def _score(item: Dict[str, Any], terms: List[str]) -> int:
    """Weight a hit by field: title beats summary beats body (CDX-03)."""
    title = str(item.get("title") or "").lower()
    summary = str(item.get("summary") or "").lower()
    body = str(item.get("body") or "").lower()
    score = 0
    for term in terms:
        score += 8 if term in title else 0
        score += 3 if term in summary else 0
        score += 1 if term in body else 0
    return score


def _rank(items: List[Dict[str, Any]], terms: List[str], lang: str) -> List[Dict[str, Any]]:
    """Score items in the given language; fall back to all languages if empty (CDX-04)."""

    def scored(only_lang: Optional[str]) -> List[Any]:
        out = []
        for item in items:
            if only_lang and f"/{only_lang}/" not in str(item.get("url") or ""):
                continue
            hit = _score(item, terms)
            if hit > 0:
                out.append((hit, item))
        out.sort(key=lambda pair: pair[0], reverse=True)
        return out

    ranked = scored(lang) or scored(None)
    return [item for _, item in ranked]


class CodexSearch:
    def __init__(self, config_getter: Callable[[], Dict[str, Any]]) -> None:
        self._config = config_getter
        self._cache: Optional[List[Dict[str, Any]]] = None
        self._fetched_at = 0.0

    def enabled(self) -> bool:
        return bool(self._config().get("enable-codex", False))

    def _index_url(self) -> str:
        return str(self._config().get("codex-index-url", DEFAULT_INDEX_URL))

    async def _load_index(self) -> List[Dict[str, Any]]:
        """Fetch + cache the codex index, SSRF-guarded and size-bounded (CDX-02)."""
        ttl = float(self._config().get("codex-cache-ttl", DEFAULT_TTL_S))
        if self._cache is not None and (time.monotonic() - self._fetched_at) < ttl:
            return self._cache
        url = self._index_url()
        reason = guard_url(url)
        if reason:
            raise ValueError(reason)
        max_bytes = int(self._config().get("codex-max-bytes", DEFAULT_MAX_BYTES))
        timeout = aiohttp.ClientTimeout(total=FETCH_TIMEOUT_S)
        async with aiohttp.ClientSession(timeout=timeout, headers={"User-Agent": "FjerkroaBot-codex/1.0"}) as session:
            async with session.get(url) as response:
                response.raise_for_status()
                raw = await read_capped(response, max_bytes)
        data = json.loads(raw.decode("utf-8", "ignore"))
        items = data.get("items", []) if isinstance(data, dict) else []
        self._cache = [i for i in items if isinstance(i, dict)]
        self._fetched_at = time.monotonic()
        return self._cache

    async def search(self, query: str, lang: str = "en", limit: int = DEFAULT_LIMIT) -> Dict[str, Any]:
        """Return sanitized top matches, or an error dict — never raise (CDX-05)."""
        try:
            items = await self._load_index()
        except Exception as err:
            logging.warning(f"codex: index load failed: {err!r}")
            return {"error": f"codex unavailable: {err}"}
        terms = _tokenize(query)
        if not terms:
            return {"query": query, "results": []}
        pick = (lang or "en").lower()
        if pick not in _VALID_LANGS:
            pick = "en"
        summary_chars = int(self._config().get("codex-summary-chars", DEFAULT_SUMMARY_CHARS))
        results = []
        for item in _rank(items, terms, pick)[: max(1, limit)]:
            results.append(
                {
                    "title": sanitize_external_text(str(item.get("title") or ""), 200),
                    "summary": sanitize_external_text(str(item.get("summary") or ""), summary_chars),
                    "collection": str(item.get("collection") or ""),
                    "url": urljoin(self._index_url(), str(item.get("url") or "")),
                }
            )
        return {"query": query, "lang": pick, "results": results}