app/tests/test_personalised_rag.py

"""Personalised tenant-library RAG (6-step private style flow)."""

from __future__ import annotations

import pytest
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.db.models import Document, IngestStatus
from app.services.personalised_rag import (
    kb_style_fallback_allowed,
    resolve_retrieval_doc_allowlist,
    tenant_library_document_ids,
)


@pytest.mark.asyncio
async def test_tenant_library_excludes_kb_and_pending(
    test_db: AsyncSession,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(settings, "knowledge_base_tenant_id", "__rics_kb__")
    test_db.add(
        Document(
            id="doc-a",
            tenant_id="firm-1",
            filename="a.pdf",
            file_path="/tmp/a.pdf",
            status=IngestStatus.complete,
        )
    )
    test_db.add(
        Document(
            id="doc-b",
            tenant_id="firm-1",
            filename="b.pdf",
            file_path="/tmp/b.pdf",
            status=IngestStatus.pending,
        )
    )
    await test_db.commit()
    ids = await tenant_library_document_ids(test_db, "firm-1")
    assert ids == frozenset({"doc-a"})


@pytest.mark.asyncio
async def test_personalised_allowlist_uses_full_library(
    test_db: AsyncSession,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(settings, "personalised_style_rag_enabled", True)
    test_db.add(
        Document(
            id="lib-1",
            tenant_id="firm-2",
            filename="r1.docx",
            file_path="/tmp/r1.docx",
            status=IngestStatus.complete,
        )
    )
    test_db.add(
        Document(
            id="lib-2",
            tenant_id="firm-2",
            filename="r2.docx",
            file_path="/tmp/r2.docx",
            status=IngestStatus.complete,
        )
    )
    await test_db.commit()
    allowed = await resolve_retrieval_doc_allowlist(
        test_db,
        "firm-2",
        primary_document_id=None,
        reference_document_ids=[],
        runtime_doc_ids=[],
        strict_uploaded_only=False,
    )
    assert allowed == frozenset({"lib-1", "lib-2"})


@pytest.mark.asyncio
async def test_strict_mode_only_attached_docs(
    test_db: AsyncSession,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setattr(settings, "personalised_style_rag_enabled", True)
    test_db.add(
        Document(
            id="lib-1",
            tenant_id="firm-3",
            filename="r1.docx",
            file_path="/tmp/r1.docx",
            status=IngestStatus.complete,
        )
    )
    await test_db.commit()
    allowed = await resolve_retrieval_doc_allowlist(
        test_db,
        "firm-3",
        primary_document_id="only-this",
        reference_document_ids=[],
        runtime_doc_ids=["runtime-sec"],
        strict_uploaded_only=True,
    )
    assert allowed == frozenset({"only-this", "runtime-sec"})


def test_kb_fallback_blocked_when_library_exists(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(settings, "personalised_style_rag_enabled", True)
    monkeypatch.setattr(settings, "knowledge_base_enabled", True)
    assert kb_style_fallback_allowed("firm-x", has_personal_library=True) is False
    assert kb_style_fallback_allowed("firm-x", has_personal_library=False) is True