pablan/backend/tests/test_refine_grounding.py

"""Retrieval-aware section refinement: a refined section may draw on what the
company has already documented, and only on documents the author is allowed to
read.

Uses deterministic fake embeddings (fake_embed): identical text lands at
distance 0, unrelated text near-orthogonal. That is enough to prove the wiring,
the permission boundary and the gates; whether grounding helps the writing is
measured against the real model in tests/evals.
"""

import uuid

import pytest
from httpx import AsyncClient
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.api.authoring.grounding import for_section as grounding_for
from app.auth.passwords import hash_password
from app.authoring.prompts import render_refine_prompt
from app.models import (
    Chunk,
    Department,
    Document,
    DocumentStatus,
    DocumentVisibility,
    User,
    UserRole,
)
from app.rag.indexing import embedding_text, reindex_document
from tests.fake_openai import FakeOpenAI

pytestmark = pytest.mark.usefixtures("fake_embed")

# A sentence long enough to clear the grounding minimum, reused as both the
# stored document and the query so the fake embedding matches exactly.
COFFEE = "Die Kaffeemaschine wird jeden Freitag gründlich entkalkt und gereinigt."


async def _user(db: AsyncSession, email: str, department_id: uuid.UUID) -> User:
    user = User(
        email=email,
        name=email.split("@")[0],
        role=UserRole.member,
        password_hash=hash_password("secret123"),
        department_id=department_id,
    )
    db.add(user)
    await db.flush()
    return user


async def _published(
    db: AsyncSession,
    *,
    title: str,
    content: str,
    author: User,
    visibility: DocumentVisibility = DocumentVisibility.public,
) -> Document:
    document = Document(
        title=title,
        status=DocumentStatus.published,
        visibility=visibility,
        content_md=content,
        author_id=author.id,
        department_id=author.department_id,
    )
    db.add(document)
    await db.flush()
    await reindex_document(db, document)
    return document


async def _chunk_text(db: AsyncSession, document: Document) -> str:
    """A chunk exactly as it was embedded — heading path and all, so a search
    for it lands at distance 0 (see indexing.embedding_text)."""
    row = (
        await db.execute(
            select(Chunk.content, Chunk.meta).where(Chunk.document_id == document.id)
        )
    ).first()
    return embedding_text(row.meta["heading_path"], row.content)


# --- The prompt only offers grounding as a reference, never as a fact source.


def test_prompt_omits_grounding_when_there_is_none() -> None:
    messages = render_refine_prompt(
        "## Pflege\n\nNotizen", prefix="", suffix="", persona=None, hint=None
    )
    assert "Related knowledge" not in messages[-1]["content"]


def test_prompt_appends_grounding_after_the_section() -> None:
    messages = render_refine_prompt(
        "## Pflege\n\nNotizen",
        prefix="",
        suffix="",
        persona=None,
        hint=None,
        knowledge=['From "Kaffeemaschine": entkalken.'],
    )
    turn = messages[-1]["content"]
    assert "Related knowledge" in turn
    assert 'From "Kaffeemaschine"' in turn
    # The section to refine still leads; grounding trails it.
    assert turn.index("Refine only this section") < turn.index("Related knowledge")


# --- grounding.for_section: finds related knowledge, excludes self, respects the gates.


async def test_grounding_surfaces_a_related_document(db: AsyncSession) -> None:
    engineering = Department(name="Engineering")
    db.add(engineering)
    await db.flush()
    author = await _user(db, "pablo@test.dev", engineering.id)
    document = await _published(
        db, title="Kaffeemaschine", content=f"## Pflege\n\n{COFFEE}", author=author
    )
    await db.commit()

    query = await _chunk_text(db, document)
    references = await grounding_for(db, query, author, document_id=uuid.uuid4())
    assert references, "an exact-text match should be grounded"
    assert references[0].title == "Kaffeemaschine"


async def test_grounding_never_includes_the_document_being_edited(
    db: AsyncSession,
) -> None:
    engineering = Department(name="Engineering")
    db.add(engineering)
    await db.flush()
    author = await _user(db, "pablo@test.dev", engineering.id)
    document = await _published(
        db, title="Kaffeemaschine", content=f"## Pflege\n\n{COFFEE}", author=author
    )
    await db.commit()

    query = await _chunk_text(db, document)
    references = await grounding_for(db, query, author, document_id=document.id)
    assert references == []


async def test_grounding_skips_a_section_that_is_still_just_a_heading(
    db: AsyncSession,
) -> None:
    engineering = Department(name="Engineering")
    db.add(engineering)
    await db.flush()
    author = await _user(db, "pablo@test.dev", engineering.id)
    await _published(
        db, title="Kaffeemaschine", content=f"## Pflege\n\n{COFFEE}", author=author
    )
    await db.commit()

    # Heading plus a few words is below the minimum, so nothing is searched.
    references = await grounding_for(
        db, "## Pflege\n\nnoch nichts", author, uuid.uuid4()
    )
    assert references == []


async def test_grounding_cannot_reach_a_document_the_author_may_not_read(
    db: AsyncSession,
) -> None:
    engineering = Department(name="Engineering")
    sales = Department(name="Sales")
    db.add_all([engineering, sales])
    await db.flush()
    pablo = await _user(db, "pablo@test.dev", engineering.id)
    max_user = await _user(db, "max@test.dev", sales.id)
    secret = await _published(
        db,
        title="Preisliste",
        content=f"## Preise\n\n{COFFEE}",
        author=max_user,
        visibility=DocumentVisibility.restricted,
    )
    await db.commit()

    query = await _chunk_text(db, secret)
    # Max authored it, so he is grounded on it; Pablo has no access, so he is not.
    assert await grounding_for(db, query, max_user, uuid.uuid4())
    assert await grounding_for(db, query, pablo, uuid.uuid4()) == []


# --- The endpoint wires grounding into the streamed prompt.


async def test_refine_endpoint_passes_grounding_to_the_model(
    client: AsyncClient, db: AsyncSession, fake_llm: FakeOpenAI
) -> None:
    engineering = Department(name="Engineering")
    db.add(engineering)
    await db.flush()
    author = await _user(db, "pablo@test.dev", engineering.id)
    reference = await _published(
        db, title="Kaffeemaschine", content=f"## Pflege\n\n{COFFEE}", author=author
    )
    draft = Document(
        title="Entwurf",
        status=DocumentStatus.draft,
        visibility=DocumentVisibility.public,
        content_md="## Pflege\n\nStichpunkte",
        author_id=author.id,
        department_id=engineering.id,
    )
    db.add(draft)
    await db.commit()

    query = await _chunk_text(db, reference)
    response = await client.post(
        "/api/auth/login", json={"email": "pablo@test.dev", "password": "secret123"}
    )
    assert response.status_code == 200

    refined = await client.post(
        f"/api/documents/{draft.id}/refine",
        json={"content_md": query, "cursor_line": 1},
    )
    assert refined.status_code == 200
    prompt = fake_llm.requests[-1]["messages"][-1]["content"]
    assert "Kaffeemaschine" in prompt