Spaces:
Sleeping
Sleeping
File size: 3,328 Bytes
be9fd4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """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
|