RAG ค้นเจอแล้ว อ่านได้ไหม? ตรวจสิทธิ์เอกสารก่อนส่งให้โมเดล

"""Offline mock: authorize RAG candidate IDs before loading document text."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Mapping, Sequence


MAX_CANDIDATES = 50
MAX_ID_LENGTH = 128
MAX_CONTEXT_FIELD_LENGTH = 128
MAX_QUERY_LENGTH = 2048


@dataclass(frozen=True)
class AuthContext:
    """Trusted server context, assumed established outside this example."""

    principal_id: str
    tenant_id: str


@dataclass(frozen=True)
class DocumentMetadata:
    """Metadata-only catalog entry; document text is stored separately."""

    document_id: str
    tenant_id: str | None


@dataclass(frozen=True)
class Document:
    document_id: str
    tenant_id: str
    text: str


@dataclass(frozen=True)
class PolicySnapshot:
    """Server-owned snapshot keyed by (tenant_id, principal_id)."""

    allowed_by_principal: Mapping[tuple[str, str], frozenset[str]]


class TextLoaderSpy:
    def __init__(self, text_by_id: Mapping[str, str]) -> None:
        self._text_by_id = text_by_id
        self.hydrated_ids: list[str] = []

    def load(self, document_id: str, tenant_id: str) -> Document:
        text = self._text_by_id[document_id]
        self.hydrated_ids.append(document_id)
        return Document(document_id, tenant_id, text)


class RerankerSpy:
    def __init__(self) -> None:
        self.input_ids: list[str] = []

    def rerank(self, query: str, documents: Sequence[Document]) -> tuple[Document, ...]:
        # A deterministic placeholder: preserve candidate order; no relevance claim.
        self.input_ids.extend(document.document_id for document in documents)
        return tuple(documents)


def _valid_context(auth: AuthContext) -> bool:
    return bool(
        isinstance(auth, AuthContext)
        and isinstance(auth.principal_id, str)
        and auth.principal_id.strip()
        and len(auth.principal_id) <= MAX_CONTEXT_FIELD_LENGTH
        and isinstance(auth.tenant_id, str)
        and auth.tenant_id.strip()
        and len(auth.tenant_id) <= MAX_CONTEXT_FIELD_LENGTH
    )


def authorize_ids(
    candidate_ids: Sequence[str],
    *,
    auth: AuthContext,
    catalog: Mapping[str, DocumentMetadata],
    policy: PolicySnapshot,
) -> tuple[str, ...]:
    """Return only IDs explicitly allowed by context, metadata, and policy."""
    if not _valid_context(auth) or not isinstance(candidate_ids, Sequence):
        return ()
    if isinstance(candidate_ids, (str, bytes)) or len(candidate_ids) > MAX_CANDIDATES:
        return ()
    if (
        not isinstance(policy, PolicySnapshot)
        or not isinstance(policy.allowed_by_principal, Mapping)
        or not isinstance(catalog, Mapping)
    ):
        return ()

    grants = policy.allowed_by_principal.get((auth.tenant_id, auth.principal_id))
    if not isinstance(grants, frozenset) or not grants:
        return ()

    allowed: list[str] = []
    seen: set[str] = set()
    for document_id in candidate_ids:
        if (
            not isinstance(document_id, str)
            or not document_id
            or len(document_id) > MAX_ID_LENGTH
            or document_id in seen
        ):
            continue
        seen.add(document_id)
        metadata = catalog.get(document_id)
        if not isinstance(metadata, DocumentMetadata):
            continue
        if metadata.document_id != document_id:
            continue
        if (
            not isinstance(metadata.tenant_id, str)
            or not metadata.tenant_id.strip()
            or len(metadata.tenant_id) > MAX_CONTEXT_FIELD_LENGTH
            or metadata.tenant_id != auth.tenant_id
        ):
            continue
        if document_id not in grants:
            continue
        allowed.append(document_id)
    return tuple(allowed)


def run_query(
    query: str,
    candidate_ids: Sequence[str],
    *,
    auth: AuthContext,
    catalog: Mapping[str, DocumentMetadata],
    policy: PolicySnapshot,
    loader: TextLoaderSpy,
    reranker: RerankerSpy,
) -> tuple[Document, ...]:
    """Authorize before hydration, then rerank only the authorized subset."""
    if (
        not isinstance(query, str)
        or not query.strip()
        or len(query) > MAX_QUERY_LENGTH
    ):
        return ()
    authorized_ids = authorize_ids(
        candidate_ids, auth=auth, catalog=catalog, policy=policy
    )
    documents = tuple(
        loader.load(document_id, catalog[document_id].tenant_id or "")
        for document_id in authorized_ids
    )
    return reranker.rerank(query, documents)


def demo() -> None:
    # Metadata catalog is separate from text: constructing it reads no document text.
    catalog = {
        "a-public": DocumentMetadata("a-public", "tenant-a"),
        "a-denied": DocumentMetadata("a-denied", "tenant-a"),
        "b-secret": DocumentMetadata("b-secret", "tenant-b"),
        "unclassified": DocumentMetadata("unclassified", None),
    }
    text_by_id = {
        "a-public": "Synthetic tenant A document.",
        "a-denied": "Synthetic denied A document.",
        "b-secret": "Synthetic tenant B secret.",
        "unclassified": "Synthetic document without tenant classification.",
    }
    policy = PolicySnapshot({("tenant-a", "principal-a"): frozenset({"a-public"})})
    loader = TextLoaderSpy(text_by_id)
    reranker = RerankerSpy()
    result = run_query(
        "synthetic query",
        ("a-public", "a-denied", "b-secret", "unclassified", "unknown-id"),
        auth=AuthContext("principal-a", "tenant-a"),
        catalog=catalog,
        policy=policy,
        loader=loader,
        reranker=reranker,
    )
    print("result_ids:", [document.document_id for document in result])
    print("hydrated_ids:", loader.hydrated_ids)
    print("reranker_input_ids:", reranker.input_ids)


if __name__ == "__main__":
    demo()

"""Auth boundary specifications. Authored; NOT RUN."""

import unittest

from auth_boundary import (
    AuthContext,
    DocumentMetadata,
    PolicySnapshot,
    RerankerSpy,
    TextLoaderSpy,
    authorize_ids,
    run_query,
)


class AuthorizationBoundaryTests(unittest.TestCase):
    def test_absent_or_malformed_policy_does_not_hydrate(self) -> None:
        for policy in (None, PolicySnapshot(None), PolicySnapshot({})):  # type: ignore[arg-type]
            loader = TextLoaderSpy(self.texts)
            reranker = RerankerSpy()
            result = run_query(
                "q", ("a-allowed",), auth=self.auth, catalog=self.catalog,
                policy=policy, loader=loader, reranker=reranker,  # type: ignore[arg-type]
            )
            self.assertEqual(result, ())
            self.assertEqual(loader.hydrated_ids, [])
            self.assertEqual(reranker.input_ids, [])

    def setUp(self) -> None:
        self.catalog = {
            "a-allowed": DocumentMetadata("a-allowed", "tenant-a"),
            "a-denied": DocumentMetadata("a-denied", "tenant-a"),
            "b-secret": DocumentMetadata("b-secret", "tenant-b"),
            "unclassified": DocumentMetadata("unclassified", None),
        }
        self.texts = {key: f"synthetic {key}" for key in self.catalog}
        self.policy = PolicySnapshot(
            {("tenant-a", "principal-a"): frozenset({"a-allowed"})}
        )
        self.auth = AuthContext("principal-a", "tenant-a")

    def test_only_explicitly_allowed_same_tenant_id_is_hydrated_and_reranked(self) -> None:
        loader = TextLoaderSpy(self.texts)
        reranker = RerankerSpy()
        result = run_query(
            "q",
            ("a-allowed", "a-denied", "b-secret", "unclassified", "missing"),
            auth=self.auth,
            catalog=self.catalog,
            policy=self.policy,
            loader=loader,
            reranker=reranker,
        )
        self.assertEqual([doc.document_id for doc in result], ["a-allowed"])
        self.assertEqual(loader.hydrated_ids, ["a-allowed"])
        self.assertEqual(reranker.input_ids, ["a-allowed"])
        for forbidden in ("a-denied", "b-secret", "unclassified", "missing"):
            self.assertNotIn(forbidden, loader.hydrated_ids)
            self.assertNotIn(forbidden, reranker.input_ids)

    def test_missing_principal_policy_denies_all(self) -> None:
        self.assertEqual(
            authorize_ids(
                ("a-allowed",),
                auth=AuthContext("unlisted-principal", "tenant-a"),
                catalog=self.catalog,
                policy=self.policy,
            ),
            (),
        )

    def test_empty_or_invalid_context_denies_all(self) -> None:
        for auth in (
            AuthContext("", "tenant-a"),
            AuthContext("principal-a", " "),
            None,
            "not-an-auth-context",
            AuthContext("p" * 129, "tenant-a"),
        ):
            self.assertEqual(
                authorize_ids(
                    ("a-allowed",),
                    auth=auth,
                    catalog=self.catalog,
                    policy=self.policy,
                ),
                (),
            )

    def test_oversized_candidate_batch_and_malformed_id_are_denied(self) -> None:
        self.assertEqual(
            authorize_ids(
                tuple("a-allowed" for _ in range(51)),
                auth=self.auth,
                catalog=self.catalog,
                policy=self.policy,
            ),
            (),
        )
        self.assertEqual(
            authorize_ids(
                (None, "x" * 129, "a-allowed"),  # type: ignore[arg-type]
                auth=self.auth,
                catalog=self.catalog,
                policy=self.policy,
            ),
            ("a-allowed",),
        )

    def test_wrong_tenant_is_denied_even_if_policy_mentions_id(self) -> None:
        policy = PolicySnapshot(
            {("tenant-a", "principal-a"): frozenset({"b-secret"})}
        )
        self.assertEqual(
            authorize_ids(
                ("b-secret",), auth=self.auth, catalog=self.catalog, policy=policy
            ),
            (),
        )

    def test_catalog_key_metadata_id_mismatch_is_denied(self) -> None:
        altered = dict(self.catalog)
        altered["a-allowed"] = DocumentMetadata("different-id", "tenant-a")
        self.assertEqual(
            authorize_ids(
                ("a-allowed",), auth=self.auth, catalog=altered, policy=self.policy
            ),
            (),
        )

    def test_duplicates_unknown_ids_and_query_bounds_do_not_reach_stages(self) -> None:
        loader = TextLoaderSpy(self.texts)
        reranker = RerankerSpy()
        result = run_query(
            "q",
            ("a-allowed", "a-allowed", "missing"),
            auth=self.auth,
            catalog=self.catalog,
            policy=self.policy,
            loader=loader,
            reranker=reranker,
        )
        self.assertEqual([doc.document_id for doc in result], ["a-allowed"])
        self.assertEqual(loader.hydrated_ids, ["a-allowed"])
        loader.hydrated_ids.clear()
        reranker.input_ids.clear()
        run_query(
            "q" * 2049,
            ("a-allowed",),
            auth=self.auth,
            catalog=self.catalog,
            policy=self.policy,
            loader=loader,
            reranker=reranker,
        )
        self.assertEqual(loader.hydrated_ids, [])
        self.assertEqual(reranker.input_ids, [])

    def test_replacement_policy_applies_to_next_request_snapshot(self) -> None:
        loader = TextLoaderSpy(self.texts)
        reranker = RerankerSpy()
        revoked = PolicySnapshot({("tenant-a", "principal-a"): frozenset()})
        result = run_query(
            "q",
            ("a-allowed",),
            auth=self.auth,
            catalog=self.catalog,
            policy=revoked,
            loader=loader,
            reranker=reranker,
        )
        self.assertEqual(result, ())
        self.assertEqual(loader.hydrated_ids, [])
        self.assertEqual(reranker.input_ids, [])

    def test_client_claims_do_not_mutate_trusted_context_or_policy(self) -> None:
        # The request payload is deliberately not an input to run_query.
        client_payload = {"tenant_id": "tenant-b", "grants": ["b-secret"]}
        trusted_auth = AuthContext("principal-a", "tenant-a")
        self.assertEqual(client_payload["tenant_id"], "tenant-b")
        self.assertEqual(trusted_auth.tenant_id, "tenant-a")
        self.assertNotIn(
            "b-secret", self.policy.allowed_by_principal[("tenant-a", "principal-a")]
        )
        loader = TextLoaderSpy(self.texts)
        reranker = RerankerSpy()
        run_query(
            "q",
            ("b-secret",),
            auth=trusted_auth,
            catalog=self.catalog,
            policy=self.policy,
            loader=loader,
            reranker=reranker,
        )
        self.assertEqual(loader.hydrated_ids, [])
        self.assertEqual(reranker.input_ids, [])


if __name__ == "__main__":
    unittest.main()

# Offline RAG authorization boundary example

This Python 3.12+ example demonstrates authorization at the document boundary. The server supplies a trusted `AuthContext` and server-owned policy. Retrieval supplies candidate IDs. The code checks each candidate against metadata and policy before loading text, then sends only authorized documents to the reranker.

## Runtime and setup

- Python 3.12 or newer
- Python standard library only
- No environment variables, installation, API keys, database, vector service, network access, or GPU
- All fixtures are synthetic and held in memory

## Run the mock

From this directory, run:

```text
python code/auth_boundary.py
```

Expected output (illustrative; not observed by running the code):

```text
result_ids: ['a-public']
hydrated_ids: ['a-public']
reranker_input_ids: ['a-public']
```

## Run the authored checks

```text
python -m unittest discover -s code -p "test_*.py"
```

The checks are authored but **NOT RUN**. Do not interpret the expected output or the test code as a test result.

## Scope and limits

This mock assumes `AuthContext`, the tenant-scoped policy snapshot, and metadata catalog have already been established by trusted server code; it does not authenticate users or verify that these inputs are authentic. Client request fields do not establish identity or grants. The catalog is metadata-only and is separate from the text loader, so denied text is not fetched by this pipeline. A real vector/search provider might expose snippets before returning IDs and must be designed to prevent that. A replacement policy snapshot applies to the next request in this synchronous illustration; concurrent revocation races are out of scope. The example does not validate a real provider, establish retrieval quality, or make performance claims. It reads no files or external data and has no paid service cost. **Mock** only; no **Sandbox** or **Live** integration is supplied.