Day 35: Metadata, Citation, Permission-aware RAG

source document
  -> parser
  -> chunker
  -> metadata enricher
  -> embedding
  -> vector/BM25 index
  -> retriever
  -> context builder
  -> LLM
  -> answer + citation

{
  "chunk_id": "company_a:policy_001:v2026-01:00012",
  "document_id": "policy_001",
  "document_title": "Employee Leave Policy",
  "text": "Nhân viên full-time có 12 ngày nghỉ phép năm...",
  "metadata": {
    "tenant_id": "company_a",
    "source_type": "pdf",
    "source_uri": "s3://private-kb/company_a/hr/policy.pdf",
    "page_start": 12,
    "page_end": 13,
    "section_path": ["HR", "Leave Policy", "Annual Leave"],
    "document_version": "2026-01",
    "chunk_index": 12,
    "chunking_version": "markdown-pdf-v3",
    "embedding_model": "BAAI/bge-m3",
    "index_version": "rag-index-2026-05-09",
    "visibility": "restricted",
    "acl_roles": ["hr", "manager"],
    "acl_groups": [],
    "acl_users": [],
    "acl_version": "acl-2026-05-09T10:00:00Z",
    "acl_updated_at": "2026-05-09T10:00:00Z",
    "text_hash": "sha256:...",
    "deleted_at": null,
    "created_at": "2026-05-09T09:00:00Z",
    "updated_at": "2026-05-09T10:00:00Z"
  }
}

user auth context
  -> normalize tenant, roles, groups, attributes
  -> pre-filter tenant + ACL trong search query
  -> retrieve candidate chunks
  -> post-filter lại ở backend
  -> build context + source_map
  -> call LLM
  -> validate citation + permission lần cuối

tenant_id = current_user.tenant_id
AND deleted_at IS NULL
AND (
  acl_roles intersects current_user.roles
  OR acl_groups intersects current_user.groups
  OR acl_users contains current_user.user_id
)

retrieved visible chunks
  -> backend gán source_id S1/S2/S3
  -> build context có [S1], [S2]
  -> LLM chỉ được cite source_id đã cấp
  -> validator kiểm tra citation
  -> renderer tạo citation link qua proxy/signed URL

{
  "source_id": "S1",
  "chunk_id": "company_a:policy_001:v2026-01:00012",
  "document_id": "policy_001",
  "title": "Employee Leave Policy",
  "page_start": 12,
  "page_end": 13,
  "section": "HR > Leave Policy > Annual Leave",
  "document_version": "2026-01",
  "access_url": "/api/rag/sources/S1?trace_id=trace_123",
  "expires_at": "2026-05-09T10:15:00Z"
}

client click citation
  -> GET /api/rag/traces/{trace_id}/sources/{source_id}
  -> backend load source_map
  -> re-check user permission
  -> generate signed URL or proxy page snippet
  -> audit source_opened event

source delete event
  -> mark document/chunks as tombstone: deleted_at != null
  -> invalidate query/result cache
  -> remove chunks from vector index and BM25 index
  -> remove or revoke generated signed URLs
  -> write deletion audit event
  -> run verification job: deleted chunks not retrievable

{
  "trace_id": "trace_123",
  "timestamp": "2026-05-09T10:02:03Z",
  "user_id": "u123",
  "tenant_id": "company_a",
  "query_hash": "sha256:...",
  "query_redacted": "tôi còn bao nhiêu ngày nghỉ phép",
  "permission_snapshot": {
    "roles": ["employee"],
    "groups": ["engineering"],
    "acl_version": "acl-2026-05-09T10:00:00Z"
  },
  "retrieval": {
    "index_version": "rag-index-2026-05-09",
    "filters": ["tenant", "acl", "deleted_at"],
    "retrieved_chunk_ids": ["c1", "c2"],
    "post_filtered_chunk_ids": []
  },
  "context_source_ids": ["S1", "S2"],
  "citation_ids": ["S1"],
  "latency_ms": {
    "embedding": 45,
    "retrieval": 38,
    "rerank": 120,
    "llm": 1400
  }
}

Client
  -> API Gateway/AuthN
  -> RAG API
      -> Authz/Directory service
      -> Query normalizer
      -> Retriever
          -> Vector DB pre-filter tenant/ACL/deleted_at
          -> BM25/hybrid search optional
      -> Backend post-filter
      -> Reranker optional
      -> Context builder + source_map
      -> LLM
      -> Citation validator
      -> Citation renderer via proxy/signed URL
      -> Audit logger

{
  "document_id": "policy_001",
  "tenant_id": "company_a",
  "title": "Employee Leave Policy",
  "source_type": "pdf",
  "source_uri": "s3://private-kb/company_a/hr/policy.pdf",
  "document_version": "2026-01",
  "status": "active",
  "acl_policy_id": "acl_policy_789",
  "acl_version": "acl-2026-05-09T10:00:00Z",
  "created_at": "2026-05-01T00:00:00Z",
  "updated_at": "2026-05-09T10:00:00Z",
  "deleted_at": null
}

{
  "chunk_id": "company_a:policy_001:v2026-01:00012",
  "document_id": "policy_001",
  "tenant_id": "company_a",
  "text": "Nhân viên full-time có 12 ngày nghỉ phép năm...",
  "metadata": {
    "document_title": "Employee Leave Policy",
    "source_type": "pdf",
    "source_uri": "s3://private-kb/company_a/hr/policy.pdf",
    "page_start": 12,
    "page_end": 13,
    "section_path": ["HR", "Leave Policy", "Annual Leave"],
    "document_version": "2026-01",
    "chunk_index": 12,
    "chunking_version": "pdf-layout-v3",
    "embedding_model": "BAAI/bge-m3",
    "embedding_dimension": 1024,
    "index_version": "rag-index-2026-05-09",
    "visibility": "restricted",
    "acl_roles": ["hr", "manager"],
    "acl_groups": [],
    "acl_users": [],
    "acl_version": "acl-2026-05-09T10:00:00Z",
    "text_hash": "sha256:...",
    "deleted_at": null
  }
}

{
  "source_id": "S1",
  "chunk_id": "company_a:policy_001:v2026-01:00012",
  "document_id": "policy_001",
  "tenant_id": "company_a",
  "title": "Employee Leave Policy",
  "page_start": 12,
  "page_end": 13,
  "section_path": ["HR", "Leave Policy", "Annual Leave"],
  "document_version": "2026-01",
  "raw_source_uri": "s3://private-kb/company_a/hr/policy.pdf"
}

from __future__ import annotations

from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
import hashlib
import re
from typing import Any


PRIVATE_URI_PREFIXES = ("s3://", "gs://", "file://", "/mnt/", "/var/")


@dataclass(frozen=True)
class UserContext:
    user_id: str
    tenant_id: str
    roles: frozenset[str]
    groups: frozenset[str]
    attributes: dict[str, str]


@dataclass(frozen=True)
class Chunk:
    chunk_id: str
    document_id: str
    text: str
    metadata: dict[str, Any]
    score: float = 0.0


@dataclass(frozen=True)
class SourceRef:
    source_id: str
    chunk_id: str
    document_id: str
    title: str
    page_start: int | None
    page_end: int | None
    section_path: list[str]
    document_version: str
    raw_source_uri: str


def hash_for_log(value: str) -> str:
    return "sha256:" + hashlib.sha256(value.encode("utf-8")).hexdigest()


def _has_any(left: frozenset[str], right: list[str] | tuple[str, ...]) -> bool:
    return bool(left.intersection(right))


def can_read(user: UserContext, chunk: Chunk) -> bool:
    meta = chunk.metadata

    if meta.get("tenant_id") != user.tenant_id:
        return False
    if meta.get("deleted_at"):
        return False

    visibility = meta.get("visibility", "restricted")
    if visibility == "public_to_tenant":
        return True
    if visibility != "restricted":
        return False

    roles = meta.get("acl_roles") or []
    groups = meta.get("acl_groups") or []
    users = meta.get("acl_users") or []

    if not roles and not groups and not users:
        return False

    return (
        _has_any(user.roles, roles)
        or _has_any(user.groups, groups)
        or user.user_id in users
    )


def build_vector_filter(user: UserContext) -> dict[str, Any]:
    return {
        "must": [
            {"key": "tenant_id", "match": user.tenant_id},
            {"key": "deleted_at", "is_null": True},
        ],
        "should": [
            {"key": "visibility", "match": "public_to_tenant"},
            {"key": "acl_roles", "intersects": sorted(user.roles)},
            {"key": "acl_groups", "intersects": sorted(user.groups)},
            {"key": "acl_users", "contains": user.user_id},
        ],
        "minimum_should_match": 1,
    }


def post_filter_chunks(user: UserContext, candidates: list[Chunk]) -> list[Chunk]:
    return [chunk for chunk in candidates if can_read(user, chunk)]


def build_context(
    user: UserContext,
    candidates: list[Chunk],
    *,
    max_chars: int = 6000,
) -> tuple[str, dict[str, SourceRef]]:
    visible_chunks = post_filter_chunks(user, candidates)
    source_map: dict[str, SourceRef] = {}
    context_blocks: list[str] = []
    used_chars = 0

    for idx, chunk in enumerate(visible_chunks, start=1):
        meta = chunk.metadata
        source_id = f"S{idx}"
        section_path = list(meta.get("section_path") or [])
        section_text = " > ".join(section_path) if section_path else "Unknown section"
        title = meta.get("document_title") or chunk.document_id

        block = (
            f"[{source_id}]\n"
            f"Title: {title}\n"
            f"Version: {meta.get('document_version', 'unknown')}\n"
            f"Page: {meta.get('page_start')}-{meta.get('page_end')}\n"
            f"Section: {section_text}\n"
            f"Text: {chunk.text}\n"
        )
        if used_chars + len(block) > max_chars:
            break

        source_map[source_id] = SourceRef(
            source_id=source_id,
            chunk_id=chunk.chunk_id,
            document_id=chunk.document_id,
            title=title,
            page_start=meta.get("page_start"),
            page_end=meta.get("page_end"),
            section_path=section_path,
            document_version=meta.get("document_version", "unknown"),
            raw_source_uri=meta.get("source_uri", ""),
        )
        context_blocks.append(block)
        used_chars += len(block)

    return "\n---\n".join(context_blocks), source_map


def validate_citations(
    answer: str,
    source_map: dict[str, SourceRef],
    user: UserContext,
    chunk_lookup: dict[str, Chunk],
) -> list[str]:
    errors: list[str] = []
    cited_ids = set(re.findall(r"\[(S\d+)\]", answer))

    for source_id in sorted(cited_ids):
        source = source_map.get(source_id)
        if source is None:
            errors.append(f"invalid_citation:{source_id}")
            continue

        chunk = chunk_lookup.get(source.chunk_id)
        if chunk is None:
            errors.append(f"missing_chunk:{source_id}")
            continue
        if not can_read(user, chunk):
            errors.append(f"not_visible:{source_id}")

    if answer.strip() and not cited_ids:
        errors.append("missing_citation")

    return errors


def render_access_url(source: SourceRef, *, trace_id: str) -> dict[str, str]:
    if source.raw_source_uri.startswith(PRIVATE_URI_PREFIXES):
        access_url = f"/api/rag/traces/{trace_id}/sources/{source.source_id}"
    else:
        access_url = f"/api/rag/traces/{trace_id}/sources/{source.source_id}"

    expires_at = datetime.now(UTC) + timedelta(minutes=10)
    return {
        "source_id": source.source_id,
        "title": source.title,
        "page": str(source.page_start or ""),
        "section": " > ".join(source.section_path),
        "document_version": source.document_version,
        "access_url": access_url,
        "expires_at": expires_at.isoformat(),
    }


def build_audit_event(
    *,
    trace_id: str,
    user: UserContext,
    query: str,
    vector_filter: dict[str, Any],
    retrieved: list[Chunk],
    visible: list[Chunk],
    source_map: dict[str, SourceRef],
    citation_errors: list[str],
) -> dict[str, Any]:
    return {
        "trace_id": trace_id,
        "timestamp": datetime.now(UTC).isoformat(),
        "user_id": user.user_id,
        "tenant_id": user.tenant_id,
        "query_hash": hash_for_log(query),
        "query_redacted": query[:300],
        "permission_snapshot": {
            "roles": sorted(user.roles),
            "groups": sorted(user.groups),
        },
        "vector_filter": vector_filter,
        "retrieved_chunk_ids": [chunk.chunk_id for chunk in retrieved],
        "visible_chunk_ids": [chunk.chunk_id for chunk in visible],
        "context_source_ids": list(source_map.keys()),
        "citation_errors": citation_errors,
    }

def sample_chunks() -> list[Chunk]:
    return [
        Chunk(
            chunk_id="company_a:policy:v1:001",
            document_id="policy",
            text="Nhân viên full-time có 12 ngày nghỉ phép năm.",
            metadata={
                "tenant_id": "company_a",
                "document_title": "Leave Policy",
                "source_uri": "s3://private/company_a/policy.pdf",
                "page_start": 1,
                "page_end": 1,
                "section_path": ["HR", "Leave"],
                "document_version": "v1",
                "visibility": "public_to_tenant",
                "deleted_at": None,
            },
        ),
        Chunk(
            chunk_id="company_a:salary:v1:001",
            document_id="salary",
            text="Bảng lương chỉ dành cho HR.",
            metadata={
                "tenant_id": "company_a",
                "document_title": "Salary Policy",
                "source_uri": "s3://private/company_a/salary.pdf",
                "page_start": 2,
                "page_end": 2,
                "section_path": ["HR", "Compensation"],
                "document_version": "v1",
                "visibility": "restricted",
                "acl_roles": ["hr"],
                "acl_groups": [],
                "acl_users": [],
                "deleted_at": None,
            },
        ),
        Chunk(
            chunk_id="company_b:policy:v1:001",
            document_id="policy_b",
            text="Policy của tenant B.",
            metadata={
                "tenant_id": "company_b",
                "document_title": "Tenant B Policy",
                "visibility": "public_to_tenant",
                "deleted_at": None,
            },
        ),
        Chunk(
            chunk_id="company_a:old:v1:001",
            document_id="old",
            text="Document đã xóa.",
            metadata={
                "tenant_id": "company_a",
                "document_title": "Old Policy",
                "visibility": "public_to_tenant",
                "deleted_at": "2026-05-09T10:00:00Z",
            },
        ),
    ]


def test_employee_cannot_read_hr_or_other_tenant_or_deleted():
    user = UserContext(
        user_id="u1",
        tenant_id="company_a",
        roles=frozenset({"employee"}),
        groups=frozenset({"engineering"}),
        attributes={},
    )

    visible = post_filter_chunks(user, sample_chunks())

    assert [chunk.chunk_id for chunk in visible] == ["company_a:policy:v1:001"]


def test_hr_can_read_restricted_salary_policy():
    user = UserContext(
        user_id="u2",
        tenant_id="company_a",
        roles=frozenset({"hr"}),
        groups=frozenset(),
        attributes={},
    )

    visible = post_filter_chunks(user, sample_chunks())

    assert "company_a:salary:v1:001" in {chunk.chunk_id for chunk in visible}


def test_citation_validator_rejects_unknown_source():
    user = UserContext("u1", "company_a", frozenset({"employee"}), frozenset(), {})
    chunks = sample_chunks()
    context, source_map = build_context(user, chunks)
    chunk_lookup = {chunk.chunk_id: chunk for chunk in chunks}

    errors = validate_citations(
        "Bạn có 12 ngày nghỉ phép năm [S1]. Thông tin khác [S99].",
        source_map,
        user,
        chunk_lookup,
    )

    assert "invalid_citation:S99" in errors

def answer_question(user: UserContext, query: str) -> dict:
    trace_id = create_trace_id()
    vector_filter = build_vector_filter(user)

    candidates = vector_db.search(
        query_embedding=embed(query),
        top_k=80,
        filter=vector_filter,
    )

    visible = post_filter_chunks(user, candidates)
    reranked = rerank(query, visible)[:8]
    context, source_map = build_context(user, reranked, max_chars=6000)

    llm_answer = llm.generate(
        system=(
            "Trả lời chỉ dựa trên context. "
            "Mỗi claim factual phải cite bằng [S1], [S2]. "
            "Không tự tạo citation ngoài source đã có."
        ),
        user=f"Question: {query}\n\nContext:\n{context}",
    )

    chunk_lookup = {chunk.chunk_id: chunk for chunk in reranked}
    errors = validate_citations(llm_answer, source_map, user, chunk_lookup)
    if errors:
        raise CitationValidationError(errors)

    citations = [
        render_access_url(source, trace_id=trace_id)
        for source in source_map.values()
        if f"[{source.source_id}]" in llm_answer
    ]

    audit_logger.write_async(
        build_audit_event(
            trace_id=trace_id,
            user=user,
            query=query,
            vector_filter=vector_filter,
            retrieved=candidates,
            visible=visible,
            source_map=source_map,
            citation_errors=[],
        )
    )

    return {
        "trace_id": trace_id,
        "answer": llm_answer,
        "citations": citations,
    }

def can_read(user: UserContext, chunk: Chunk) -> bool:
    ...

def build_vector_filter(user: UserContext) -> dict:
    ...

def post_filter_chunks(user: UserContext, chunks: list[Chunk]) -> list[Chunk]:
    ...

def build_context(user: UserContext, chunks: list[Chunk], max_chars: int) -> tuple[str, dict]:
    ...

[S1]
Title: Employee Leave Policy
Version: 2026-01
Page: 12-13
Section: HR > Leave Policy > Annual Leave
Text: Nhân viên full-time có 12 ngày nghỉ phép năm.

def validate_citations(answer: str, source_map: dict, user: UserContext) -> list[str]:
    ...

"Bạn có 12 ngày nghỉ phép năm [S1]. Mức lương theo bảng nội bộ [S99]."

GET /api/rag/traces/{trace_id}/sources/{source_id}

delete event
  -> mark tombstone
  -> invalidate cache
  -> remove vector chunks
  -> remove BM25 entries
  -> revoke signed URL if possible
  -> write audit
  -> verify not retrievable

{
  "trace_id": "trace_123",
  "timestamp": "2026-05-09T10:02:03Z",
  "user_id": "u123",
  "tenant_id": "company_a",
  "query_hash": "sha256:...",
  "permission_snapshot": {},
  "filters": {},
  "retrieved_chunk_ids": [],
  "visible_chunk_ids": [],
  "context_source_ids": [],
  "citation_ids": [],
  "citation_errors": [],
  "index_version": "rag-index-2026-05-09"
}

Dùng được trong production không?
Nếu có thì cần điều kiện gì?
Nếu chưa đủ thì đang thiếu gì?