Spaces:
Sleeping
Sleeping
File size: 3,725 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 | """Private per-tenant style RAG (upload → extract → sanitise → index → retrieve → generate).
Maps to the product flow:
1. User uploads previous reports
2. Ingest extracts text
3. document_sanitiser redacts PII before embedding
4. Chunks live in the tenant-scoped vector index only
5. Generation retrieves from that tenant's library (not shared KB)
6. OpenAI receives retrieved style examples as prompt context (not fine-tuning)
"""
from __future__ import annotations
import logging
from typing import TYPE_CHECKING
from sqlalchemy import select
from app.config import settings
from app.db.database import get_session_factory
from app.db.models import Document, IngestStatus
if TYPE_CHECKING:
from sqlalchemy.ext.asyncio import AsyncSession
logger = logging.getLogger(__name__)
async def tenant_library_document_ids(
db: AsyncSession,
tenant_id: str,
) -> frozenset[str]:
"""UUIDs of ingested reference reports for this tenant (excludes KB tenant)."""
tid = (tenant_id or "").strip()
if not tid or tid == settings.knowledge_base_tenant_id:
return frozenset()
result = await db.execute(
select(Document.id).where(
Document.tenant_id == tid,
Document.status == IngestStatus.complete,
)
)
return frozenset(str(row[0]) for row in result.all())
async def tenant_has_personal_library(tenant_id: str) -> bool:
"""True when the tenant has at least one successfully indexed upload."""
factory = get_session_factory()
async with factory() as db:
return bool(await tenant_library_document_ids(db, tenant_id))
async def resolve_retrieval_doc_allowlist(
db: AsyncSession | None,
tenant_id: str,
*,
primary_document_id: str | None,
reference_document_ids: list[str] | None,
runtime_doc_ids: list[str],
strict_uploaded_only: bool,
) -> frozenset[str] | None:
"""Doc IDs that may supply RAG snippets for this generation call.
* **Personalised RAG (default):** all completed uploads for ``tenant_id``,
plus runtime section index virtual docs. ``primary_document_id`` /
``reference_document_ids`` only affect ranking (via ``retrieve_for_report``),
not membership.
* **strict_uploaded_only:** only docs attached to this report + runtime.
* **Legacy (personalised off):** same as strict — attached docs only.
"""
runtime = {str(d) for d in (runtime_doc_ids or []) if str(d).strip()}
attached: set[str] = set(runtime)
if primary_document_id:
attached.add(str(primary_document_id))
attached.update(str(d) for d in (reference_document_ids or []) if str(d).strip())
if strict_uploaded_only or not settings.personalised_style_rag_enabled:
return frozenset(attached) if attached else None
if db is not None:
library = await tenant_library_document_ids(db, tenant_id)
else:
factory = get_session_factory()
async with factory() as session:
library = await tenant_library_document_ids(session, tenant_id)
if not library:
return None
allow = set(library) | runtime
logger.debug(
"Personalised RAG allowlist tenant=%s docs=%d runtime=%d attached_for_rank=%d",
tenant_id,
len(library),
len(runtime),
len(attached) - len(runtime),
)
return frozenset(allow)
def kb_style_fallback_allowed(tenant_id: str, *, has_personal_library: bool) -> bool:
"""KB reference style is only used before the tenant has their own library."""
if not settings.personalised_style_rag_enabled:
return True
if has_personal_library:
return False
return bool(settings.knowledge_base_enabled)
|