"""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