RAG/backend/app/tools/seed_demo.py

"""Idempotently embed, store, retrieve, and rerank the public synthetic corpus."""

from __future__ import annotations

import asyncio
import hashlib
import json
import os
import sys
import uuid
from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path
from typing import Any, cast
from urllib.parse import urlsplit

import psycopg
from pgvector import Vector
from pgvector.psycopg import register_vector
from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from app.adapters.fake import FakeEmbeddingProvider, FakeReranker, lexical_features
from app.adapters.model_gateway import ModelGatewayAdapter
from app.core.config import Settings
from app.core.demo_identity import (
    ACCESS_SCOPE_ID,
    IDENTITY_NAMESPACE,
    KNOWLEDGE_BASE_ID,
    offline_embedding_profile_hash,
)
from app.core.secrets import SecretFileError
from app.ports.model_providers import EmbeddingProvider, ModelProviderError, Reranker

PROJECT_ROOT = Path(__file__).resolve().parents[3]
DEFAULT_SAMPLE_ROOT = PROJECT_ROOT / "data" / "samples" / "public"


@dataclass(frozen=True, slots=True)
class DemoDocument:
    source_id: str
    title: str
    content: str
    region: str
    mineral: str
    page_no: int
    cloud_policy_id: str


@dataclass(frozen=True, slots=True)
class DemoQuery:
    qid: str
    query: str
    expected_doc_ids: tuple[str, ...]
    answerable: bool


@dataclass(frozen=True, slots=True)
class DemoNamespace:
    mode: str
    knowledge_base_id: uuid.UUID
    access_scope_id: uuid.UUID
    scope_name: str
    knowledge_base_name: str
    storage_prefix: str


OFFLINE_NAMESPACE = DemoNamespace(
    mode="fake",
    knowledge_base_id=KNOWLEDGE_BASE_ID,
    access_scope_id=ACCESS_SCOPE_ID,
    scope_name="synthetic-demo",
    knowledge_base_name="虚构地质 PoC 知识库（离线）",
    storage_prefix="synthetic/offline",
)
BAILIAN_NAMESPACE = DemoNamespace(
    mode="bailian",
    knowledge_base_id=uuid.uuid5(IDENTITY_NAMESPACE, "synthetic-bailian-knowledge-base"),
    access_scope_id=uuid.uuid5(IDENTITY_NAMESPACE, "synthetic-bailian-public-scope"),
    scope_name="synthetic-bailian-validation",
    knowledge_base_name="虚构地质 PoC 知识库（百炼验证）",
    storage_prefix="synthetic/bailian",
)


@dataclass(frozen=True, slots=True)
class EmbeddedVector:
    vector: tuple[float, ...]
    request_id: str | None
    usage: dict[str, int | None]
    elapsed_ms: int


@dataclass(frozen=True, slots=True)
class PreparedChunk:
    source_id: str
    document_id: uuid.UUID
    version_id: uuid.UUID
    chunk_id: uuid.UUID
    raw_sha256: str
    cloud_text: str
    cloud_text_sha256: str
    embedding_prefix: str
    embedding_text: str
    embedding_text_sha256: str
    outbound_manifest_sha256: str
    embedding_profile_hash: str
    vector: tuple[float, ...]
    embedding_model: str
    provider_request_id: str | None
    embedding_usage: dict[str, int | None]
    embedding_elapsed_ms: int
    title: str
    region: str
    mineral: str
    page_no: int
    cloud_policy_id: str


class SeedContractError(ValueError):
    def __init__(self, code: str) -> None:
        self.code = code
        super().__init__(code)


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


def namespace_for_mode(mode: str) -> DemoNamespace:
    if mode == "fake":
        return OFFLINE_NAMESPACE
    if mode == "bailian":
        return BAILIAN_NAMESPACE
    raise SeedContractError("invalid_provider_mode")


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    if not path.is_file():
        raise SeedContractError("fixture_missing")
    records: list[dict[str, Any]] = []
    for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), start=1):
        try:
            value = json.loads(line)
        except json.JSONDecodeError as exc:
            raise SeedContractError(f"invalid_jsonl_line_{line_number}") from exc
        if not isinstance(value, dict):
            raise SeedContractError(f"jsonl_object_required_line_{line_number}")
        records.append(value)
    return records


def load_documents(path: Path) -> list[DemoDocument]:
    documents = []
    for value in load_jsonl(path):
        if value.get("source_type") != "synthetic":
            raise SeedContractError("non_synthetic_document")
        if value.get("review_state") != "LOCAL_PARSED_PENDING_CLOUD_REVIEW":
            raise SeedContractError("invalid_initial_review_state")
        documents.append(
            DemoDocument(
                source_id=str(value["doc_id"]),
                title=str(value["title"]),
                content=str(value["content"]),
                region=str(value["region"]),
                mineral=str(value["mineral"]),
                page_no=int(value["page_no"]),
                cloud_policy_id=str(value["cloud_policy_id"]),
            )
        )
    if len(documents) != 20 or len({item.source_id for item in documents}) != 20:
        raise SeedContractError("expected_twenty_unique_documents")
    return documents


def load_queries(path: Path) -> list[DemoQuery]:
    queries = [
        DemoQuery(
            qid=str(value["qid"]),
            query=str(value["query"]),
            expected_doc_ids=tuple(str(item) for item in value["expected_doc_ids"]),
            answerable=bool(value["answerable"]),
        )
        for value in load_jsonl(path)
    ]
    if not queries:
        raise SeedContractError("query_set_empty")
    return queries


def embedding_profile_hash(settings: Settings, mode: str) -> str:
    if mode == "fake":
        return offline_embedding_profile_hash(settings.embedding_dimension)
    if mode != "bailian":
        raise SeedContractError("invalid_provider_mode")

    endpoint_identity = sha256_text(urlsplit(settings.bailian_openai_base_url).hostname or "")
    model = settings.embedding_model
    api_mode = "openai-compatible"
    profile = {
        "api_mode": api_mode,
        "dimension": settings.embedding_dimension,
        "endpoint_identity_hash": endpoint_identity,
        "model": model,
        "normalization": "provider-default",
        "profile_version": 1,
    }
    return sha256_text(
        json.dumps(profile, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    )


async def embed_in_batches(
    provider: EmbeddingProvider,
    texts: Sequence[str],
) -> tuple[tuple[EmbeddedVector, ...], str]:
    vectors: list[EmbeddedVector] = []
    resolved_model: str | None = None
    for offset in range(0, len(texts), 10):
        result = await provider.embed_documents(texts[offset : offset + 10])
        if resolved_model is not None and result.model != resolved_model:
            raise SeedContractError("embedding_model_changed_between_batches")
        resolved_model = result.model
        if len(result.vectors) != len(texts[offset : offset + 10]):
            raise SeedContractError("embedding_batch_count_mismatch")
        usage = {
            "input_tokens": result.usage.input_tokens,
            "output_tokens": result.usage.output_tokens,
            "total_tokens": result.usage.total_tokens,
        }
        vectors.extend(
            EmbeddedVector(
                vector=vector,
                request_id=result.request_id,
                usage=usage,
                elapsed_ms=max(0, round(result.elapsed_ms)),
            )
            for vector in result.vectors
        )
    if len(vectors) != len(texts) or resolved_model is None:
        raise SeedContractError("embedding_result_count_mismatch")
    return tuple(vectors), resolved_model


def prepare_chunks(
    documents: Sequence[DemoDocument],
    vectors: Sequence[EmbeddedVector],
    *,
    profile_hash: str,
    embedding_model: str,
    namespace: DemoNamespace = OFFLINE_NAMESPACE,
) -> list[PreparedChunk]:
    prepared = []
    for document, vector in zip(documents, vectors, strict=True):
        raw_payload = json.dumps(
            {
                "content": document.content,
                "mineral": document.mineral,
                "page_no": document.page_no,
                "region": document.region,
                "title": document.title,
            },
            ensure_ascii=False,
            sort_keys=True,
            separators=(",", ":"),
        )
        raw_hash = sha256_text(raw_payload)
        document_identity = (
            f"document:{document.source_id}"
            if namespace.mode == "fake"
            else f"document:{namespace.mode}:{document.source_id}"
        )
        document_id = uuid.uuid5(IDENTITY_NAMESPACE, document_identity)
        version_id = uuid.uuid5(
            IDENTITY_NAMESPACE,
            f"version:{document.source_id}:{raw_hash}:{profile_hash}",
        )
        chunk_id = uuid.uuid5(IDENTITY_NAMESPACE, f"chunk:{version_id}:0")
        prefix = (
            f"标题：{document.title}\n地区：{document.region}\n矿种：{document.mineral}\n正文："
        )
        cloud_hash = sha256_text(document.content)
        embedding_text = prefix + document.content
        embedding_hash = sha256_text(embedding_text)
        manifest_payload = json.dumps(
            [
                {
                    "cloud_text_sha256": cloud_hash,
                    "embedding_text_sha256": embedding_hash,
                    "ordinal": 0,
                }
            ],
            sort_keys=True,
            separators=(",", ":"),
        )
        prepared.append(
            PreparedChunk(
                source_id=document.source_id,
                document_id=document_id,
                version_id=version_id,
                chunk_id=chunk_id,
                raw_sha256=raw_hash,
                cloud_text=document.content,
                cloud_text_sha256=cloud_hash,
                embedding_prefix=prefix,
                embedding_text=embedding_text,
                embedding_text_sha256=embedding_hash,
                outbound_manifest_sha256=sha256_text(manifest_payload),
                embedding_profile_hash=profile_hash,
                vector=vector.vector,
                embedding_model=embedding_model,
                provider_request_id=vector.request_id,
                embedding_usage=vector.usage,
                embedding_elapsed_ms=vector.elapsed_ms,
                title=document.title,
                region=document.region,
                mineral=document.mineral,
                page_no=document.page_no,
                cloud_policy_id=document.cloud_policy_id,
            )
        )
    return prepared


def database_dsn(settings: Settings) -> str:
    return (
        settings.database_url().set(drivername="postgresql").render_as_string(hide_password=False)
    )


def write_chunks(
    settings: Settings,
    chunks: Sequence[PreparedChunk],
    *,
    namespace: DemoNamespace,
) -> dict[str, int]:
    if not chunks:
        raise SeedContractError("chunks_empty")
    profile_hashes = {item.embedding_profile_hash for item in chunks}
    resolved_models = {item.embedding_model for item in chunks}
    if len(profile_hashes) != 1 or len(resolved_models) != 1:
        raise SeedContractError("mixed_embedding_profiles")
    profile_hash = next(iter(profile_hashes))
    resolved_model = next(iter(resolved_models))
    if namespace.mode == "fake":
        provider = "local-synthetic"
        api_mode = "deterministic-offline"
        endpoint_identity_hash = sha256_text("local-fake")
    else:
        provider = "aliyun-bailian"
        api_mode = "model-gateway/openai-compatible"
        endpoint_identity_hash = sha256_text(
            urlsplit(settings.bailian_openai_base_url).hostname or ""
        )

    with psycopg.connect(database_dsn(settings), row_factory=dict_row) as connection:
        register_vector(connection)
        connection.execute("SELECT pg_advisory_xact_lock(724202607120001)")
        connection.execute(
            """
            INSERT INTO rag.model_profiles (
              profile_hash, alias, kind, provider, model, api_mode, dimension,
              endpoint_identity_hash, config_snapshot, synthetic, enabled
            ) VALUES (
              %s, %s, 'embedding', %s, %s, %s, 1024, %s, %s, %s, true
            )
            ON CONFLICT (profile_hash) DO NOTHING
            """,
            (
                profile_hash,
                f"{namespace.mode}-embedding-{profile_hash[:12]}",
                provider,
                resolved_model,
                api_mode,
                endpoint_identity_hash,
                Jsonb(
                    {
                        "dimension": settings.embedding_dimension,
                        "requested_model": settings.embedding_model,
                        "source": "synthetic-seed-v1",
                    }
                ),
                namespace.mode == "fake",
            ),
        )
        registered_profile = connection.execute(
            """
            SELECT kind, provider, model, api_mode, dimension, endpoint_identity_hash
            FROM rag.model_profiles
            WHERE profile_hash = %s
            """,
            (profile_hash,),
        ).fetchone()
        if registered_profile is None or (
            registered_profile["kind"] != "embedding"
            or registered_profile["provider"] != provider
            or registered_profile["model"] != resolved_model
            or registered_profile["api_mode"] != api_mode
            or registered_profile["dimension"] != settings.embedding_dimension
            or registered_profile["endpoint_identity_hash"] != endpoint_identity_hash
        ):
            raise SeedContractError("embedding_profile_collision")
        connection.execute(
            """
            INSERT INTO rag.knowledge_bases (
              id, name, description, active_embedding_profile_hash
            )
            VALUES (%s, %s, %s, %s)
            ON CONFLICT (id) DO UPDATE
              SET name = EXCLUDED.name,
                  description = EXCLUDED.description,
                  active_embedding_profile_hash = EXCLUDED.active_embedding_profile_hash,
                  updated_at = now()
            """,
            (
                namespace.knowledge_base_id,
                namespace.knowledge_base_name,
                "仅含公开的合成验证文本",
                profile_hash,
            ),
        )
        connection.execute(
            """
            INSERT INTO rag.access_scopes (id, knowledge_base_id, name)
            VALUES (%s, %s, %s)
            ON CONFLICT (id) DO NOTHING
            """,
            (
                namespace.access_scope_id,
                namespace.knowledge_base_id,
                namespace.scope_name,
            ),
        )

        for item in chunks:
            connection.execute(
                """
                INSERT INTO rag.documents (
                  id, knowledge_base_id, access_scope_id, raw_sha256,
                  filename, storage_key, mime_type, status
                ) VALUES (%s, %s, %s, %s, %s, %s, 'application/json',
                          'LOCAL_PARSED_PENDING_CLOUD_REVIEW')
                ON CONFLICT (id) DO UPDATE
                  SET filename = EXCLUDED.filename,
                      storage_key = EXCLUDED.storage_key,
                      mime_type = EXCLUDED.mime_type,
                      updated_at = now()
                """,
                (
                    item.document_id,
                    namespace.knowledge_base_id,
                    namespace.access_scope_id,
                    item.raw_sha256,
                    f"{item.source_id}.json",
                    f"{namespace.storage_prefix}/{item.source_id}",
                ),
            )
            connection.execute(
                """
                INSERT INTO rag.document_versions (
                  id, document_id, parser_profile_hash, ocr_profile_hash,
                  normalization_profile_hash, chunk_profile_hash, cloud_policy_id,
                  outbound_manifest_sha256, review_state, embedding_profile_hash,
                  status, expected_chunk_count
                ) VALUES (
                  %s, %s, %s, NULL, %s, %s, %s, %s,
                  'LOCAL_PARSED_PENDING_CLOUD_REVIEW', %s, 'PROCESSING', 1
                )
                ON CONFLICT (id) DO NOTHING
                """,
                (
                    item.version_id,
                    item.document_id,
                    sha256_text("synthetic-jsonl-parser-v1"),
                    sha256_text("identity-normalization-v1"),
                    sha256_text("one-record-one-chunk-v1"),
                    item.cloud_policy_id,
                    item.outbound_manifest_sha256,
                    item.embedding_profile_hash,
                ),
            )
            connection.execute(
                """
                INSERT INTO rag.outbound_manifest_items (
                  document_version_id, ordinal, outbound_manifest_sha256,
                  cloud_text_sha256, embedding_text_sha256
                )
                SELECT %s, 0, %s, %s, %s
                WHERE NOT EXISTS (
                  SELECT 1
                  FROM rag.outbound_manifest_items
                  WHERE document_version_id = %s AND ordinal = 0
                )
                ON CONFLICT (document_version_id, ordinal) DO NOTHING
                """,
                (
                    item.version_id,
                    item.outbound_manifest_sha256,
                    item.cloud_text_sha256,
                    item.embedding_text_sha256,
                    item.version_id,
                ),
            )
            connection.execute(
                """
                INSERT INTO rag.chunks (
                  id, knowledge_base_id, document_id, document_version_id,
                  access_scope_id, ordinal, display_text, cloud_text,
                  cloud_text_sha256, embedding_prefix, embedding_text,
                  embedding_text_sha256, embedded_text_sha256,
                  embedding_profile_hash, outbound_manifest_sha256, token_count,
                  page_start, page_end, section_path, metadata, embedding_model,
                  embedding_dimension, embedding, approval_status, index_status,
                  searchable
                ) VALUES (
                  %s, %s, %s, %s, %s, 0, %s, %s, %s, %s, %s, %s, %s,
                  %s, %s, %s, %s, %s, %s, %s, %s, 1024, %s,
                  'LOCAL_PARSED_PENDING_CLOUD_REVIEW', 'READY', false
                )
                ON CONFLICT (id) DO UPDATE SET
                  display_text = EXCLUDED.display_text,
                  cloud_text = EXCLUDED.cloud_text,
                  cloud_text_sha256 = EXCLUDED.cloud_text_sha256,
                  embedding_prefix = EXCLUDED.embedding_prefix,
                  embedding_text = EXCLUDED.embedding_text,
                  embedding_text_sha256 = EXCLUDED.embedding_text_sha256,
                  embedded_text_sha256 = EXCLUDED.embedded_text_sha256,
                  embedding_profile_hash = EXCLUDED.embedding_profile_hash,
                  outbound_manifest_sha256 = EXCLUDED.outbound_manifest_sha256,
                  token_count = EXCLUDED.token_count,
                  page_start = EXCLUDED.page_start,
                  page_end = EXCLUDED.page_end,
                  section_path = EXCLUDED.section_path,
                  metadata = EXCLUDED.metadata,
                  embedding_model = EXCLUDED.embedding_model,
                  updated_at = now()
                """,
                (
                    item.chunk_id,
                    namespace.knowledge_base_id,
                    item.document_id,
                    item.version_id,
                    namespace.access_scope_id,
                    item.cloud_text,
                    item.cloud_text,
                    item.cloud_text_sha256,
                    item.embedding_prefix,
                    item.embedding_text,
                    item.embedding_text_sha256,
                    item.embedding_text_sha256,
                    item.embedding_profile_hash,
                    item.outbound_manifest_sha256,
                    max(1, len(lexical_features(item.embedding_text))),
                    item.page_no,
                    item.page_no,
                    Jsonb([item.title]),
                    Jsonb(
                        {
                            "mineral": item.mineral,
                            "region": item.region,
                            "source_doc_id": item.source_id,
                            "source_type": "synthetic",
                        }
                    ),
                    item.embedding_model,
                    Vector(list(item.vector)),
                ),
            )
            connection.execute(
                """
                UPDATE rag.document_versions
                   SET review_state = 'CLOUD_APPROVED',
                       cloud_approved_at = COALESCE(cloud_approved_at, now()),
                       cloud_approved_by = 'seed-demo:synthetic-policy',
                       status = 'READY',
                       completed_at = COALESCE(completed_at, now())
                 WHERE id = %s
                """,
                (item.version_id,),
            )
            connection.execute(
                """
                INSERT INTO rag.embedding_cache (
                  profile_hash, embedding_text_sha256, embedding, resolved_model,
                  provider_request_id, usage, elapsed_ms
                ) VALUES (%s, %s, %s, %s, %s, %s, %s)
                ON CONFLICT (profile_hash, embedding_text_sha256) DO NOTHING
                """,
                (
                    item.embedding_profile_hash,
                    item.embedding_text_sha256,
                    Vector(list(item.vector)),
                    item.embedding_model,
                    item.provider_request_id,
                    Jsonb(item.embedding_usage),
                    item.embedding_elapsed_ms,
                ),
            )
            connection.execute(
                """
                INSERT INTO rag.chunk_embedding_assignments (
                  chunk_id, profile_hash, embedding_text_sha256,
                  cache_profile_hash, cache_embedding_text_sha256,
                  status, completed_at
                ) VALUES (%s, %s, %s, %s, %s, 'READY', now())
                ON CONFLICT (chunk_id, profile_hash) DO NOTHING
                """,
                (
                    item.chunk_id,
                    item.embedding_profile_hash,
                    item.embedding_text_sha256,
                    item.embedding_profile_hash,
                    item.embedding_text_sha256,
                ),
            )
            connection.execute(
                """
                UPDATE rag.chunks
                   SET approval_status = 'CLOUD_APPROVED',
                       index_status = 'READY',
                       updated_at = now()
                 WHERE id = %s
                """,
                (item.chunk_id,),
            )
            connection.execute(
                """
                UPDATE rag.documents
                   SET active_version_id = %s,
                       raw_sha256 = %s,
                       status = 'READY',
                       updated_at = now()
                 WHERE id = %s
                """,
                (item.version_id, item.raw_sha256, item.document_id),
            )
            connection.execute(
                "UPDATE rag.chunks SET searchable = true, updated_at = now() WHERE id = %s",
                (item.chunk_id,),
            )

        counts = connection.execute(
            """
            SELECT
              count(*)::integer AS chunks,
              count(*) FILTER (WHERE embedding IS NOT NULL)::integer AS vectors,
              count(*) FILTER (WHERE searchable)::integer AS searchable
            FROM rag.chunks
            WHERE knowledge_base_id = %s
              AND embedding_profile_hash = %s
            """,
            (namespace.knowledge_base_id, profile_hash),
        ).fetchone()
        if counts is None:
            raise SeedContractError("database_count_missing")
        return {key: int(counts[key]) for key in ("chunks", "vectors", "searchable")}


def retrieve(
    settings: Settings,
    query_vector: tuple[float, ...],
    *,
    namespace: DemoNamespace,
    profile_hash: str,
) -> list[dict[str, Any]]:
    with psycopg.connect(database_dsn(settings), row_factory=dict_row) as connection:
        register_vector(connection)
        connection.execute("SET LOCAL hnsw.iterative_scan = strict_order")
        connection.execute("SET LOCAL hnsw.ef_search = 100")
        rows = connection.execute(
            """
            SELECT chunk.id, chunk.metadata, chunk.embedding_text,
                   1 - (chunk.embedding <=> %s) AS vector_score
            FROM rag.chunks AS chunk
            JOIN rag.knowledge_bases AS knowledge_base
              ON knowledge_base.id = chunk.knowledge_base_id
             AND knowledge_base.active_embedding_profile_hash = %s
            WHERE chunk.searchable
              AND chunk.knowledge_base_id = %s
              AND chunk.access_scope_id = %s
              AND chunk.embedding_profile_hash = %s
            ORDER BY chunk.embedding <=> %s
            LIMIT %s
            """,
            (
                Vector(list(query_vector)),
                profile_hash,
                namespace.knowledge_base_id,
                namespace.access_scope_id,
                profile_hash,
                Vector(list(query_vector)),
                settings.vector_top_k,
            ),
        ).fetchall()
        return [dict(row) for row in rows]


async def evaluate_queries(
    settings: Settings,
    queries: Sequence[DemoQuery],
    embedder: EmbeddingProvider,
    reranker: Reranker,
    *,
    namespace: DemoNamespace,
    profile_hash: str,
) -> dict[str, float | int]:
    hits = 0
    answerable = 0
    for query in queries:
        query_result = await embedder.embed_query(query.query)
        candidates = retrieve(
            settings,
            query_result.vectors[0],
            namespace=namespace,
            profile_hash=profile_hash,
        )
        if not candidates:
            continue
        reranked = await reranker.rerank(
            query.query,
            [cast(str, item["embedding_text"]) for item in candidates],
            top_n=min(settings.rerank_top_n, len(candidates)),
        )
        result_doc_ids = [
            cast(dict[str, Any], candidates[item.index]["metadata"])["source_doc_id"]
            for item in reranked.items[:3]
        ]
        if query.answerable:
            answerable += 1
            if set(query.expected_doc_ids) & set(result_doc_ids):
                hits += 1
    return {
        "answerable_queries": answerable,
        "hit_at_3": hits,
        "hit_rate_at_3": round(hits / answerable, 4) if answerable else 0.0,
    }


def output_summary(payload: dict[str, Any]) -> None:
    sys.stdout.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")


def safe_failure_site(error: BaseException) -> str:
    traceback = error.__traceback__
    selected: str | None = None
    while traceback is not None:
        filename = Path(traceback.tb_frame.f_code.co_filename).name
        if filename == "seed_demo.py":
            selected = f"{filename}:{traceback.tb_frame.f_code.co_name}:{traceback.tb_lineno}"
        traceback = traceback.tb_next
    return selected or "external_dependency"


async def async_main() -> int:
    mode = os.getenv("DEMO_PROVIDER_MODE", "fake").strip().lower()
    if mode not in {"fake", "bailian"}:
        output_summary({"status": "failed", "error_kind": "invalid_provider_mode"})
        return 2

    settings = Settings()
    namespace = namespace_for_mode(mode)
    documents_path = Path(
        os.getenv("DEMO_DOCUMENTS_PATH", str(DEFAULT_SAMPLE_ROOT / "demo_documents.jsonl"))
    )
    queries_path = Path(
        os.getenv("DEMO_QUERIES_PATH", str(DEFAULT_SAMPLE_ROOT / "demo_queries.jsonl"))
    )
    cloud_gateway: ModelGatewayAdapter | None = None
    try:
        documents = load_documents(documents_path)
        queries = load_queries(queries_path)
        profile_hash = embedding_profile_hash(settings, mode)
        embedder: EmbeddingProvider
        reranker: Reranker
        if mode == "bailian":
            cloud_gateway = ModelGatewayAdapter.from_settings(settings)
            embedder = cloud_gateway
            reranker = cloud_gateway
        else:
            embedder = FakeEmbeddingProvider(settings.embedding_dimension)
            reranker = FakeReranker()

        texts = [
            f"标题：{item.title}\n地区：{item.region}\n矿种：{item.mineral}\n正文：{item.content}"
            for item in documents
        ]
        vectors, resolved_model = await embed_in_batches(embedder, texts)
        prepared = prepare_chunks(
            documents,
            vectors,
            profile_hash=profile_hash,
            embedding_model=resolved_model,
            namespace=namespace,
        )
        counts = write_chunks(settings, prepared, namespace=namespace)
        metrics = await evaluate_queries(
            settings,
            queries,
            embedder,
            reranker,
            namespace=namespace,
            profile_hash=profile_hash,
        )
        output_summary(
            {
                "counts": counts,
                "embedding_model": resolved_model,
                "metrics": metrics,
                "provider_mode": mode,
                "status": "ok",
            }
        )
        return 0
    except ModelProviderError as exc:
        output_summary(
            {
                "status": "failed",
                "error_kind": f"model_provider_{exc.kind.value}",
                "status_code": exc.status_code,
            }
        )
        return 1
    except psycopg.Error as exc:
        constraint_name = exc.diag.constraint_name
        output_summary(
            {
                "status": "failed",
                "error_kind": "database_error",
                "sqlstate": exc.sqlstate,
                "failure_site": safe_failure_site(exc),
                "constraint": constraint_name
                if constraint_name and constraint_name.replace("_", "").isalnum()
                else None,
            }
        )
        return 1
    except SecretFileError:
        output_summary({"status": "failed", "error_kind": "secret_configuration"})
        return 1
    except OSError:
        output_summary({"status": "failed", "error_kind": "fixture_io_error"})
        return 1
    except SeedContractError as exc:
        output_summary({"status": "failed", "error_kind": "seed_contract_error", "code": exc.code})
        return 1
    except ValueError as exc:
        output_summary(
            {
                "status": "failed",
                "error_kind": "fixture_or_contract_error",
                "failure_site": safe_failure_site(exc),
            }
        )
        return 1
    finally:
        if cloud_gateway is not None:
            await cloud_gateway.aclose()


def main() -> None:
    raise SystemExit(asyncio.run(async_main()))


if __name__ == "__main__":
    main()