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