brain-of-reese/tests/unit/test_embed_question_length.py

"""Unit: the chat question-embed prefix budget (phase 114, TODO L6; LOCKED A1)
and the too-large embed error mapping (task 02; LOCKED A3).

The embed step of ``POST /api/chat`` embeds at most
``settings.embed_question_max_chars`` (default 1200 — the chunker's
``HARD_MAX_CHARS`` budget: worst-case ~1.4 chars/token, so it stays
under the endpoint's ~1024-token per-request input cap) of the
question; the FULL question still reaches the LLM prompt. A question
at or under the budget embeds byte-identically to the pre-phase path.

Error mapping (task 02): a single text over the endpoint's input cap
is a DETERMINISTIC size failure (``EmbeddingInputTooLargeError``, the
real ``_post_embeddings`` → ``_TooLarge`` branch) — the chat endpoint
settles it with the accurate "question too long" terminal frame + the
reachability-fine hint, ONE attempt, NO retry frame (locked A3). An
embed failure without the too-large signature keeps the phase-67
reachability path byte-identically (retry frames + the old copy).

The endpoint-level tests drive ``POST /api/chat`` with the LLM (a
recording fake that captures every ``embed_one`` input and the
messages of each request, or a real ``LLMClient`` on a canned-failure
transport for the error-mapping tests), the retriever, and the DB
session all faked (the ``test_chat_gate.py`` wiring), so the whole
embed → retrieve → prompt contract runs without a stack.

Frontend pins (task 02, source-assertion house style): the SSE
error frame's optional ``hint`` threads through the stream state
machine to ``showErrorBanner`` (shown in place of the default
reachability hint); the phase-111 Retry button rides the same
turn-error path.
"""
from __future__ import annotations

import asyncio
import json
import uuid
from collections.abc import Iterator
from datetime import UTC, datetime
from pathlib import Path
from types import SimpleNamespace
from typing import TYPE_CHECKING, Any

import pytest
from fastapi.testclient import TestClient
from pydantic import ValidationError

from app.api import chat as chat_api
from app.config import Settings
from app.main import app as fastapi_app
from app.models import Document, KbOverview
from app.rag.chunker import HARD_MAX_CHARS
from app.rag.llm import (
    EmbeddingError,
    EmbeddingInputTooLargeError,
    LLMClient,
    StreamPiece,
)
from app.rag.retriever import RetrievedChunk
from tests.conftest import ADMIN_PASSWORD

if TYPE_CHECKING:
    from app.rag.scaffolding import ScaffoldingFilter

ANSWER = "Here is what your notes say about that."


def _question(n: int) -> str:
    """A deterministic *n*-char question with a distinct head and tail.

    Repeated, index-marked sentences cut at exactly *n* chars — the
    4,000-char case is the composer's schema clamp
    (``ChatRequest.message`` ``max_length=4000``), the L6 repro.
    """
    sentence = "How is my homelab kubernetes cluster configured for long-running batch jobs? "
    parts: list[str] = []
    total = 0
    i = 0
    while total < n:
        part = f"[{i}] " + sentence
        parts.append(part)
        total += len(part)
        i += 1
    return "".join(parts)[:n]


# ---------- the setting (default + validator) ----------


def test_default_budget_matches_the_chunker_hard_cap() -> None:
    """LOCKED A1: the default is the chunker's ``HARD_MAX_CHARS`` budget."""
    assert Settings(_env_file=None).embed_question_max_chars == 1200  # pyright: ignore[reportCallIssue]
    assert Settings.model_fields["embed_question_max_chars"].default == HARD_MAX_CHARS


@pytest.mark.parametrize("bad", [0, -1, -1200])
def test_budget_rejects_zero_and_negative(bad: int) -> None:
    """``0``/negative would embed an empty/absent prefix — a typo that
    must fail loudly at startup (the ``agent_max_rounds`` pattern)."""
    with pytest.raises(ValidationError, match="embed_question_max_chars must be > 0"):
        Settings(_env_file=None, embed_question_max_chars=bad)  # pyright: ignore[reportCallIssue]


@pytest.mark.parametrize("good", [1, 500, 10_000])
def test_budget_accepts_positive_values(good: int) -> None:
    """A model with a smaller/larger cap is env-tunable, no code change."""
    settings = Settings(_env_file=None, embed_question_max_chars=good)  # pyright: ignore[reportCallIssue]
    assert settings.embed_question_max_chars == good


# ---------- endpoint-level (fake LLM + fake retriever + fake session) ----------


class _RecordingLLM:
    """Records every ``embed_one`` input and the messages of each request.

    Streams a canned answer and never emits tool calls, so a grounded
    turn through the agent loop ends after the single (tools-offered)
    request. Mirrors the ``test_chat_gate.py`` fake LLM.
    """

    def __init__(self, answer: str = ANSWER) -> None:
        self.settings = Settings(_env_file=None)  # pyright: ignore[reportCallIssue]
        self.embedded: list[str] = []
        self.answer = answer
        self.seen: list[list[dict[str, str]]] = []

    async def embed_one(self, text: str) -> list[float]:
        self.embedded.append(text)
        return [0.0] * 768

    async def chat_stream(
        self,
        messages: list[dict[str, str]],
        tools: list[dict[str, Any]] | None = None,
        scaffolding: ScaffoldingFilter | None = None,  # phase 71 pass-through
    ):
        self.seen.append(messages)
        for i in range(0, len(self.answer), 12):
            yield StreamPiece("content", self.answer[i : i + 12])


class _FakeSteeringResult:
    """Empty steering-note result (no stored notes in these tests)."""

    def all(self) -> list[Any]:
        return []


class _FakeSession:
    """Stands in for the DB session (the ``test_chat_gate.py`` fake)."""

    def __init__(self) -> None:
        self.added: list[Any] = []
        self.commits = 0

    def __enter__(self) -> _FakeSession:
        return self

    def __exit__(self, *args: Any) -> None:
        pass

    def add(self, obj: Any) -> None:
        self.added.append(obj)

    def commit(self) -> None:
        self.commits += 1

    def scalars(self, _stmt: Any) -> _FakeSteeringResult:
        return _FakeSteeringResult()

    def get(self, model: Any, pk: Any) -> Any:
        if model is KbOverview:
            return None
        return None


def _doc(title: str, content: str) -> Document:
    return Document(
        id=uuid.uuid4(),
        source="Homelab",
        path=f"{title.lower().replace(' ', '-')}.md",
        full_path="/tmp/doc.md",
        title=title,
        content=content,
        content_hash="0" * 64,
        created_at=datetime(2024, 6, 15, 12, 0, 0, tzinfo=UTC),
    )


def _chunk(doc: Document, score: float) -> RetrievedChunk:
    return RetrievedChunk(
        chunk_id=uuid.uuid4(),
        position=0,
        content=doc.content[:32],
        score=score,
        document=doc,
        cosine=score,
        fts_hit=False,
        is_summary=False,
    )


def _fake_retriever(chunks: list[RetrievedChunk]) -> Any:
    def retrieve(_db: Any, _question: str, _vec: list[float]) -> list[RetrievedChunk]:
        return chunks

    return retrieve


@pytest.fixture(autouse=True)
def _admin_signed_in(client: TestClient) -> None:
    """``POST /api/chat`` is user-gated — the endpoint-level tests run
    as the signed-in ADMIN (the ``test_chat_gate.py`` pattern)."""
    r = client.post("/api/login", json={"password": ADMIN_PASSWORD})
    assert r.status_code == 204, f"admin login failed: {r.status_code} {r.text}"


@pytest.fixture()
def embed_env(monkeypatch: pytest.MonkeyPatch) -> Iterator[tuple[_FakeSession, _RecordingLLM]]:
    """``POST /api/chat`` with retriever, session, and LLM all faked.

    The code defaults apply (``embed_question_max_chars=1200``); the
    gate threshold is pinned low so a 0.9-chunk goes grounded and a
    0.29-chunk deflects, regardless of any local ``.env``.
    """
    monkeypatch.setattr(chat_api, "db_available", lambda: True)
    session = _FakeSession()
    llm = _RecordingLLM()
    monkeypatch.setattr(chat_api, "SessionLocal", lambda: session)
    monkeypatch.setitem(fastapi_app.dependency_overrides, chat_api.get_llm, lambda: llm)
    monkeypatch.setattr(
        chat_api,
        "get_settings",
        lambda: Settings(_env_file=None, relevance_threshold=0.30),  # pyright: ignore[reportCallIssue]
    )
    yield session, llm
    fastapi_app.dependency_overrides.clear()


def _ask(client: TestClient, message: str) -> list[dict[str, Any]]:
    with client.stream("POST", "/api/chat", json={"message": message}) as r:
        assert r.status_code == 200
        frames: list[dict[str, Any]] = []
        buf = ""
        for part in r.iter_text():
            buf += part
            while "\n\n" in buf:
                frame, buf = buf.split("\n\n", 1)
                frame = frame.strip()
                if frame.startswith("data:"):
                    frames.append(json.loads(frame.removeprefix("data:").strip()))
    assert buf.strip() == ""
    return frames


def test_long_question_embeds_exactly_the_bounded_prefix(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """L6 repro: the 4,000-char (composer-clamp) question embeds ONLY the
    1200-char prefix — one embed call, exactly the head, and the turn
    completes (no error frame)."""
    _session, llm = embed_env
    question = _question(4_000)
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(_doc("T", "C"), 0.29)]))

    frames = _ask(client, question)

    # The code default (1200), derived from the field so this never drifts.
    budget = Settings.model_fields["embed_question_max_chars"].default
    assert llm.embedded == [question[:budget]]
    assert llm.embedded[0] != question  # it really was cut
    assert all(f["type"] != "error" for f in frames)
    assert frames[-1]["type"] == "done"


def test_long_question_full_text_reaches_llm_prompt(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Truncation is the embed step ONLY: the deflected turn's request
    carries the FULL 4,000-char question as the user message."""
    _session, llm = embed_env
    question = _question(4_000)
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(_doc("T", "C"), 0.29)]))

    _frames = _ask(client, question)

    assert len(llm.seen) == 1
    assert llm.seen[0][-1] == {"role": "user", "content": question}
    assert len(llm.seen[0][-1]["content"]) == 4_000


def test_grounded_turn_agent_request_carries_full_question(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Same contract on the grounded (agent-loop) branch: the single
    tools-offered request carries the FULL question."""
    _session, llm = embed_env
    question = _question(4_000)
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(_doc("T", "C"), 0.90)]))

    frames = _ask(client, question)

    assert frames[-1]["deflected"] is False
    assert len(llm.seen) == 1
    assert llm.seen[0][-1] == {"role": "user", "content": question}
    assert llm.embedded == [question[:1200]]  # the prefix, not the full text


@pytest.mark.parametrize("n", [100, 900])
def test_short_question_embeds_byte_identically(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
    n: int,
) -> None:
    """A question under the budget embeds the WHOLE question — the
    pre-phase call, byte for byte (one call, the exact string)."""
    _session, llm = embed_env
    question = _question(n)
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(_doc("T", "C"), 0.29)]))

    _frames = _ask(client, question)

    assert llm.embedded == [question]


def test_question_at_exactly_the_budget_embeds_whole(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """The budget is an INCLUSIVE cap (``[:budget]``): a question exactly
    1200 chars long embeds in full — no char lost at the boundary."""
    _session, llm = embed_env
    question = _question(1_200)
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(_doc("T", "C"), 0.29)]))

    _frames = _ask(client, question)

    assert llm.embedded == [question]


def test_budget_is_env_tunable_via_settings(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """LOCKED A1: the budget is a setting — a deployment with a smaller-cap
    model lowers it via ``BOR_EMBED_QUESTION_MAX_CHARS`` (here: 500) and
    the prefix follows, no code change."""
    _session, llm = embed_env
    monkeypatch.setattr(
        chat_api,
        "get_settings",
        lambda: Settings(
            _env_file=None,  # pyright: ignore[reportCallIssue]
            relevance_threshold=0.30,
            embed_question_max_chars=500,
        ),
    )
    question = _question(4_000)
    monkeypatch.setattr(chat_api, "retrieve", _fake_retriever([_chunk(_doc("T", "C"), 0.29)]))

    _frames = _ask(client, question)

    assert llm.embedded == [question[:500]]
    assert llm.seen[0][-1] == {"role": "user", "content": question}


# ---------- error mapping (task 02, LOCKED A3) ----------

#: The aipi/litellm signature of a too-large input (the L6 repro body,
#: truncated the way the client sees it — the client keys off "too
#: large" in the body).
_TOO_LARGE_BODY = (
    "input (903 tokens) is too large to process. increase the physical "
    "batch size (current batch size: 512)"
)


class _CannedHttpResponse:
    """One canned transport reply (status + text body, no JSON)."""

    def __init__(self, status_code: int, text: str) -> None:
        self.status_code = status_code
        self.text = text


class _CannedHttp:
    """Stands in for the httpx transport the openai client owns.

    Every POST returns the same canned failure and records the request
    body — the embed attempt counter.
    """

    def __init__(self, status_code: int, text: str) -> None:
        self.status_code = status_code
        self.text = text
        self.posts: list[dict[str, Any]] = []

    async def post(
        self, url: str, *, json: dict[str, Any], headers: dict[str, str] | None = None
    ) -> _CannedHttpResponse:
        self.posts.append(json)
        return _CannedHttpResponse(self.status_code, self.text)


def _canned_embed_llm(
    settings: Settings, status_code: int, text: str
) -> tuple[LLMClient, _CannedHttp]:
    """A REAL ``LLMClient`` whose transport is a canned failure.

    ``embed_one`` runs the real ``_post_embeddings`` → ``_embed_batch``
    path — the ``_TooLarge`` branch fires for real. The chat stream is
    never reached: the embed step settles the turn first.
    """
    llm = LLMClient(settings)
    http = _CannedHttp(status_code, text)
    llm._client = SimpleNamespace(_client=http)  # pyright: ignore[reportAttributeAccessIssue]
    return llm, http


def test_single_oversized_text_raises_the_too_large_subclass() -> None:
    """The single-text ``_TooLarge`` branch of ``LLMClient._embed_batch``
    raises ``EmbeddingInputTooLargeError`` — still an
    ``EmbeddingError`` (the importer path is a drop-in) with the
    byte-identical import-oriented message."""
    settings = Settings(_env_file=None)  # pyright: ignore[reportCallIssue]
    llm, http = _canned_embed_llm(settings, 500, _TOO_LARGE_BODY)
    with pytest.raises(EmbeddingInputTooLargeError, match="token cap") as exc:
        asyncio.run(llm.embed_one("x" * 3000))
    assert isinstance(exc.value, EmbeddingError)
    assert str(exc.value) == (
        "a single 3000-char chunk exceeded the endpoint's per-request input "
        "token cap — lower BOR_CHUNK_TARGET_CHARS and re-import"
    )
    assert len(http.posts) == 1
    assert llm.embed_batches == 0


def test_too_large_embed_maps_to_terminal_too_long_frame(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """L6 acceptance: the litellm "too large to process" 500 maps to the
    ACCURATE terminal error — exactly ONE error frame with the precise
    detail + the reachability-fine hint, NO retry frame, ONE embed
    attempt (locked A3: deterministic — never retried), no "couldn't
    reach" copy."""
    _session, _recording = embed_env
    settings = Settings(
        _env_file=None,  # pyright: ignore[reportCallIssue]
        relevance_threshold=0.30,
        llm_retry_delay=0.01,  # keep the (unused here) budget cheap
    )
    monkeypatch.setattr(chat_api, "get_settings", lambda: settings)
    llm, http = _canned_embed_llm(settings, 500, _TOO_LARGE_BODY)
    monkeypatch.setitem(fastapi_app.dependency_overrides, chat_api.get_llm, lambda: llm)

    frames = _ask(client, "How is my homelab kubernetes cluster configured?")

    assert len(http.posts) == 1  # ONE attempt — no restart (locked A3)
    assert [f["type"] for f in frames] == ["error"]  # terminal: no retry, no done
    frame = frames[0]
    assert frame["detail"] == "Question too long — trim it and re-ask."
    assert frame["hint"] == (
        "The app reached the embedding model fine — only the question length is the problem."
    )


def test_transport_embed_failure_keeps_the_legacy_retry_path(
    client: TestClient,
    embed_env: tuple[_FakeSession, _RecordingLLM],
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    """Regression pin: an embed 500 WITHOUT the too-large signature
    keeps the phase-67 reachability behavior byte-identical — one
    retry frame per restart (attempts 2–4 of 4), then the OLD
    "couldn't reach" copy (``hint`` null) after the full attempt
    budget."""
    _session, _recording = embed_env
    settings = Settings(
        _env_file=None,  # pyright: ignore[reportCallIssue]
        relevance_threshold=0.30,
        llm_retry_delay=0.01,  # 3 restarts × 0.01 s — the shape is what is pinned
    )
    monkeypatch.setattr(chat_api, "get_settings", lambda: settings)
    llm, http = _canned_embed_llm(settings, 500, "internal server error")
    monkeypatch.setitem(fastapi_app.dependency_overrides, chat_api.get_llm, lambda: llm)

    frames = _ask(client, "How is my homelab kubernetes cluster configured?")

    assert len(http.posts) == 4  # the full budget was spent (it retried)
    assert [f["type"] for f in frames] == ["retry", "retry", "retry", "error"]
    for i, frame in enumerate(frames[:3], start=2):
        assert frame == {"type": "retry", "attempt": i, "max_attempts": 4}
    error = frames[-1]
    assert error["detail"] == (
        "I couldn't reach the embedding model — please try again."
    )
    assert error["hint"] is None  # the additive field is null, never a too-long hint


# ---------- frontend hint support (task 02, source-assertion house style) ----------

APP_JS = Path(__file__).resolve().parents[2] / "frontend" / "assets" / "app.js"


def _js() -> str:
    return APP_JS.read_text(encoding="utf-8")


def test_sse_error_branch_threads_the_frame_hint() -> None:
    """The stream state machine's error branch carries the frame's
    optional ``hint`` through the throw (``err.hint``) — the banner
    shows it in place of the default reachability hint."""
    js = _js()
    idx = js.find('ev.type === "error"')
    assert idx != -1, "the readSSE handler must branch on error frames"
    end = js.find("throw err;", idx)
    assert end != -1, "the error branch throws the detail as the turn error"
    branch = js[idx:end]
    assert "err.hint = ev.hint;" in branch, (
        "the frame's optional hint must ride the thrown error"
    )


def test_catch_passes_the_thrown_hint_to_set_ui_state() -> None:
    """The turn catch passes the thrown error's hint to ``setUiState``
    (the third argument) — the banner call gets it via the opts merge.
    Non-Error throws pass no opts (the default hint applies)."""
    js = _js()
    idx = js.find("setUiState(UI_STATE.error, detail,")
    assert idx != -1, "the turn-error landing must flow through setUiState"
    call = js[idx : js.find(");", idx)]
    assert "{ hint: err.hint }" in call, "the thrown hint must reach setUiState"


def test_set_ui_state_merges_opts_into_the_banner_call() -> None:
    """``setUiState``'s error transition keeps the phase-111
    ``{ retryable: true }`` (the banner Retry button) AND merges the
    turn opts (the hint) into the ``showErrorBanner`` call."""
    js = _js()
    idx = js.find("export function setUiState")
    assert idx != -1
    body = js[idx : js.find("\n}\n", idx)]
    assert "opts = {}" in body, "the opts parameter carries the turn hint"
    assert "showErrorBanner(errorDetail, { retryable: true, ...opts });" in body


def test_show_error_banner_honors_opts_hint() -> None:
    """``showErrorBanner`` shows ``opts.hint`` in place of the default
    ``ERROR_HINT`` — in BOTH the with-detail and the detail-less forms
    (``??`` falls back on null/undefined, so hint-less frames keep the
    old copy byte-identically)."""
    js = _js()
    idx = js.find("function showErrorBanner")
    assert idx != -1
    body = js[idx : js.find("\n}\n", idx)]
    assert body.count("opts.hint ?? ERROR_HINT") == 2, (
        "the hint fallback must cover the detail and no-detail forms"
    )