app/services/personalised_rag.py

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