nisse/nisse-dev/tests/test_expert.py

import io
from pathlib import Path
from typing import Any

import pytest

from nisse_core.log.canonical import (
    AssistantText,
    CanonicalEvent,
    DelegationCompleted,
    DelegationStarted,
    StoredEvent,
)
from nisse_core.models.llm import LLMError, RoleMap
from nisse_core.oplog import OpLog
from nisse_core.state.store import Transcript
from nisse_dev.expert import (
    DEFAULT_BUDGET,
    BudgetExhausted,
    ExpertAnswer,
    ExpertDesk,
    parse_answer,
)
from nisse_dev.state import StateStore

ROLES = RoleMap(
    driver="haiku",
    helper="haiku",
    expert="claude-haiku-4-5",
)

GOOD = '{"answer": "Use SQLite; it is enough here.", "confidence": "high", "caveats": "Assumes one writer."}'


class ScriptedClient:
    def __init__(self, replies: list[str]) -> None:
        self._replies = replies
        self.calls: list[dict[str, Any]] = []

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]:
        self.calls.append(
            {"events": events, "tools": tools, "model": model, "system": system}
        )
        if not self._replies:
            return [AssistantText(text="nothing more", model=model)]
        return [AssistantText(text=self._replies.pop(0), model=model)]


@pytest.fixture
def store(tmp_path: Path) -> StateStore:
    return StateStore(tmp_path / "state.db", oplog=OpLog("nisse-dev", io.StringIO()))


ITEM = "t"  # the work item the consults happen in


def log(store: StateStore, item: str = ITEM) -> Transcript:
    return store.transcript(item, "expert")


class Bound:
    """The desk as the executor uses it: within one work item."""

    def __init__(self, desk: ExpertDesk, store: StateStore, item: str) -> None:
        self.desk, self.store, self.item = desk, store, item

    def ask(self, question: str, **kw: Any) -> Any:
        return self.desk.ask(log(self.store, self.item), question, **kw)


def desk(client: Any, store: StateStore, item: str = ITEM, **kw: Any) -> Bound:
    return Bound(ExpertDesk(client, ROLES, **kw), store, item)


# --- parsing (the contract is prompt-level, so parsing is the guard) ------


def test_parses_a_clean_object() -> None:
    parsed = parse_answer(GOOD)
    assert parsed == ExpertAnswer(
        answer="Use SQLite; it is enough here.",
        confidence="high",
        caveats="Assumes one writer.",
    )


def test_parses_a_fenced_block_because_models_wrap_json() -> None:
    assert parse_answer(f"Here you go:\n```json\n{GOOD}\n```") is not None


def test_missing_or_bad_fields_are_a_failed_turn() -> None:
    assert parse_answer('{"answer": "", "confidence": "high"}') is None
    assert parse_answer('{"answer": "x", "confidence": "certain"}') is None
    assert parse_answer('{"confidence": "high"}') is None
    assert parse_answer("no json at all") is None
    assert parse_answer('["not", "an", "object"]') is None


def test_absent_caveats_is_tolerated() -> None:
    parsed = parse_answer('{"answer": "yes", "confidence": "low"}')
    assert parsed is not None
    assert parsed.caveats == ""


# --- the desk ------------------------------------------------------------


def test_ask_returns_a_validated_answer(store: StateStore) -> None:
    client = ScriptedClient([GOOD])
    answer = desk(client, store).ask("SQLite or Postgres?")
    assert answer.confidence == "high"
    assert "SQLite" in answer.render()


def test_consult_is_stateless_and_uses_the_consultant_role(
    store: StateStore,
) -> None:
    log(store).append(AssistantText(text="secret repo context", model="m"))
    client = ScriptedClient([GOOD])
    desk(client, store).ask("Is this API deprecated?")

    call = client.calls[0]
    # L9: no repo access, no session state — exactly the question, and none
    # of the log the caller was working from.
    assert [e.event.text for e in call["events"]] == ["Is this API deprecated?"]
    assert all("secret repo context" not in e.event.text for e in call["events"])
    assert call["tools"] is None
    assert call["model"] == "claude-haiku-4-5"


def test_unparseable_reply_is_retried_once_then_raises(store: StateStore) -> None:
    client = ScriptedClient(["waffle", "more waffle"])
    with pytest.raises(LLMError, match="no usable answer"):
        desk(client, store).ask("q")
    assert len(client.calls) == 2
    assert "not a single valid JSON object" in client.calls[1]["system"]


def test_retry_succeeds_after_a_bad_first_reply(store: StateStore) -> None:
    client = ScriptedClient(["waffle", GOOD])
    assert desk(client, store).ask("q").confidence == "high"


def test_consults_are_recorded_as_composite_events(store: StateStore) -> None:
    # L9 wants a consult/delegation eval set; reusing the pair gets consult
    # evaluation for free.
    desk(ScriptedClient([GOOD]), store).ask("q")
    events = [e.event for e in log(store).read()]
    started = next(e for e in events if isinstance(e, DelegationStarted))
    completed = next(e for e in events if isinstance(e, DelegationCompleted))
    assert started.role == "expert"
    assert started.brief == "q"
    assert completed.exit_status == "answered"


def test_a_failed_consult_is_recorded_too(store: StateStore) -> None:
    with pytest.raises(LLMError):
        desk(ScriptedClient(["waffle", "waffle"]), store).ask("q")
    completed = [
        e.event for e in log(store).read() if isinstance(e.event, DelegationCompleted)
    ]
    assert completed[0].exit_status == "failed"


def test_budget_is_enforced_across_calls(store: StateStore) -> None:
    client = ScriptedClient([GOOD] * 10)
    d = desk(client, store, budget=2)
    d.ask("one")
    d.ask("two")
    with pytest.raises(BudgetExhausted, match="budget of 2"):
        d.ask("three")
    assert len(client.calls) == 2  # the third never reached the model


def test_budget_survives_a_restart(tmp_path: Path) -> None:
    # Counted from the canonical log, not from memory, so a crash cannot
    # hand a circling executor a fresh allowance.
    db = tmp_path / "state.db"
    with StateStore(db, oplog=OpLog("d", io.StringIO())) as store:
        desk(ScriptedClient([GOOD]), store, budget=1).ask("one")
    with (
        StateStore(db, oplog=OpLog("d", io.StringIO())) as store,
        pytest.raises(BudgetExhausted),
    ):
        desk(ScriptedClient([GOOD]), store, budget=1).ask("two")


def test_budget_counts_only_after_the_task_marker(store: StateStore) -> None:
    d = desk(ScriptedClient([GOOD] * 5), store, budget=1)
    d.ask("earlier task")
    marker = log(store).read()[-1].id
    # A new task starts its allowance afresh from its own marker.
    assert d.ask("new task", since_event_id=marker).confidence == "high"


def test_the_budget_is_per_work_item(store: StateStore) -> None:
    # Consults in one task never use up another's allowance (L9).
    desk(ScriptedClient([GOOD]), store, item="a", budget=1).ask("one")
    assert desk(ScriptedClient([GOOD]), store, item="b", budget=1).ask("two")
    with pytest.raises(BudgetExhausted):
        desk(ScriptedClient([GOOD]), store, item="a", budget=1).ask("three")


def test_default_budget_is_small() -> None:
    assert DEFAULT_BUDGET <= 5  # L9: do not let this become a second subsystem