cluster-knaster/core/tests/conftest.py

from __future__ import annotations

import textwrap
from collections.abc import Callable
from pathlib import Path

import pytest

from knaster.embeddings import HashingEmbedder
from knaster.llm import ChatResult
from knaster.registry import Registry
from knaster.skills import SkillLibrary
from knaster.store import Store

REGISTRY = {
    "slots": {
        "tier0-strong": {
            "tier": 0, "device": "external", "engine": "openrouter", "litellm_model": "tier0-strong",
            "base_model": "strong-x", "max_context": 128000, "effective_context": 100000,
        },
        "tier1-general": {
            "tier": 1, "engine": "vllm", "node": "gpu-main", "litellm_model": "tier1-general",
            "base_model": "base-14b", "max_context": 8192, "effective_context": 2000,
            "supports_lora": True, "fallbacks": ["tier1-flex"],
        },
        "tier1-flex": {
            "tier": 1, "engine": "vllm", "node": "gpu-flex", "litellm_model": "tier1-flex",
            "base_model": "base-14b", "max_context": 32768, "effective_context": 16000,
            "supports_lora": True,
        },
        "verifier-small": {
            "tier": 1, "engine": "llamacpp", "node": "gpu-small", "litellm_model": "verifier-small",
            "base_model": "small-4b", "max_context": 8192,
        },
        "embed-default": {
            "tier": 1, "engine": "llamacpp", "node": "gpu-embed", "litellm_model": "embed-default",
            "base_model": "embed-0.6b", "max_context": 8192, "kind": "embedding",
        },
    },
    "adapters": [
        {"name": "rechnungen-v1", "base_model": "base-14b", "cluster": "rechnungen", "version": 1},
    ],
}

SKILLS = {
    "rechnung-extrahieren": """\
        ---
        name: rechnung-extrahieren
        description: Rechnung auslesen und Rechnungsnummer Betrag Datum als JSON extrahieren
        metadata:
          knaster:
            status: active
            stage: 1
            cluster: rechnungen
            threshold: 0.3
            triggers: ["Rechnungsdaten aus einer Rechnung extrahieren"]
            max_output_tokens: 300
            output:
              format: json
              schema:
                type: object
                required: [nummer, betrag]
                properties:
                  nummer: {type: string}
                  betrag: {type: number}
        ---
        # Rechnung extrahieren
        Lies Nummer und Betrag aus.

        ## Beispiele
        Eingabe: Rechnung 7 über 10 EUR -> {"nummer": "7", "betrag": 10}
        """,
    "ticket-kategorisieren": """\
        ---
        name: ticket-kategorisieren
        description: Support Ticket einer Kategorie zuordnen Drucker Netzwerk Passwort
        metadata:
          knaster:
            status: active
            stage: 3
            execution: classifier
            classifier: ticket-kat
            threshold: 0.3
            min_confidence: 0.6
        ---
        Ordne das Ticket einer Kategorie zu: drucker, netzwerk, passwort.
        """,
    "text-gross": """\
        ---
        name: text-gross
        description: Text in Großbuchstaben umwandeln uppercase
        metadata:
          knaster:
            status: active
            stage: 2
            execution: program
            program: echo-upper
            threshold: 0.3
        ---
        Wandle den Text in Großbuchstaben um.
        """,
    "entwurf": """\
        ---
        name: entwurf
        description: Gedicht über Rechnung schreiben
        metadata:
          knaster:
            status: draft
        ---
        Entwurf.
        """,
    "vertrag-pruefen": """\
        ---
        name: vertrag-pruefen
        description: Vertrag prüfen Kündigungsfrist Laufzeit Klauseln
        metadata:
          knaster:
            status: demoted
            stage: 1
            threshold: 0.3
            verifier:
              - {type: contains, value: "Kündigungsfrist"}
        ---
        Prüfe den Vertrag.
        """,
}


@pytest.fixture
def store() -> Store:
    return Store("sqlite://")


@pytest.fixture
def registry(store: Store) -> Registry:
    return Registry.from_dict(REGISTRY, store)


@pytest.fixture
def skills_dir(tmp_path: Path) -> Path:
    for name, text in SKILLS.items():
        d = tmp_path / "skills" / name
        d.mkdir(parents=True)
        (d / "SKILL.md").write_text(textwrap.dedent(text), encoding="utf-8")
    return tmp_path / "skills"


@pytest.fixture
def library(skills_dir: Path, store: Store) -> SkillLibrary:
    return SkillLibrary.from_dir(skills_dir, store)


@pytest.fixture
def embedder() -> HashingEmbedder:
    return HashingEmbedder()


class FakeChat:
    """Antworten pro Modellname; ein Handler darf ChatResult liefern oder eine Exception werfen."""

    def __init__(self, handlers: dict[str, Callable[[list[dict[str, str]]], ChatResult | str]] | None = None):
        self.handlers = handlers or {}
        self.calls: list[tuple[str, list[dict[str, str]]]] = []

    def chat(self, model, messages, *, max_tokens=None, temperature=None):
        self.calls.append((model, messages))
        handler = self.handlers.get(model)
        if handler is None:
            return ChatResult(content="ok", finish_reason="stop")
        result = handler(messages)
        if isinstance(result, str):
            return ChatResult(content=result, finish_reason="stop")
        return result


@pytest.fixture
def fake_chat() -> FakeChat:
    return FakeChat()


class Clock:
    def __init__(self, t: float | None = None):
        import time

        self.t = time.time() if t is None else t

    def __call__(self) -> float:
        return self.t

    def advance(self, seconds: float) -> None:
        self.t += seconds


@pytest.fixture
def clock() -> Clock:
    return Clock()


@pytest.fixture
def make_system(tmp_path, skills_dir, fake_chat, clock):
    import yaml

    from knaster.embeddings import HashingEmbedder
    from knaster.healer.drivers import DryRunDriver
    from knaster.lifecycle import LoggingJobRunner
    from knaster.settings import REPO_ROOT, Settings
    from knaster.system import build_system

    def make(**overrides):
        site = tmp_path / "site"
        site.mkdir(exist_ok=True)
        (site / "registry.yaml").write_text(yaml.safe_dump(REGISTRY))
        settings = Settings(db_url="sqlite://", site_dir=site, skills_dir=skills_dir,
                            config_dir=overrides.pop("config_dir", REPO_ROOT / "config"),
                            api_token=overrides.pop("api_token", "secret"))
        kw = dict(embedder=HashingEmbedder(), chat=fake_chat, jobs=LoggingJobRunner(), notifier=DryRunDriver(),
                  clock=clock)
        kw.update(overrides)
        return build_system(settings, **kw)

    return make