YoVinchen
f4ba5d5342
Some checks failed
verify / verify (push) Has been cancelled
The Stage 1 foundation now proves provider contracts with mocks and validates PostgreSQL/pgvector ingestion, approval binding, retrieval, reranking, and idempotency using only synthetic data. Live Bailian validation remains gated on rotating the exposed key. Constraint: The key shown in chat is compromised and cannot be used or committed Rejected: Mark Stage 1 complete from mock and offline results | real three-model smoke is still required Confidence: high Scope-risk: moderate Reversibility: clean Directive: Do not enable real-data ingestion until Stage 3 cloud approval and outbound manifest controls are enforced end to end Tested: make verify; 41 pytest tests; strict mypy; Ruff; Compose config; pinned image build; empty-volume migration; role denial; two idempotent 20-vector seeds; database restart persistence Not-tested: Live Bailian calls require a newly rotated key; React product UI is not implemented
667 lines
24 KiB
Python
667 lines
24 KiB
Python
"""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.bailian import BailianEmbeddingAdapter, BailianRerankerAdapter
|
|
from app.adapters.fake import FakeEmbeddingProvider, FakeReranker, lexical_features
|
|
from app.core.config import Settings
|
|
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"
|
|
IDENTITY_NAMESPACE = uuid.UUID("eef85571-1f64-4a09-86d7-53fd329c3eb2")
|
|
KNOWLEDGE_BASE_ID = uuid.uuid5(IDENTITY_NAMESPACE, "synthetic-demo-knowledge-base")
|
|
ACCESS_SCOPE_ID = uuid.uuid5(IDENTITY_NAMESPACE, "synthetic-demo-public-scope")
|
|
|
|
|
|
@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 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
|
|
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 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:
|
|
endpoint_identity = "local-fake"
|
|
model = "fake-feature-hash-v1"
|
|
api_mode = "deterministic-offline"
|
|
if mode == "bailian":
|
|
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[tuple[float, ...], ...], str]:
|
|
vectors: list[tuple[float, ...]] = []
|
|
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
|
|
vectors.extend(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[tuple[float, ...]],
|
|
*,
|
|
profile_hash: str,
|
|
embedding_model: str,
|
|
) -> 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_id = uuid.uuid5(IDENTITY_NAMESPACE, f"document:{document.source_id}")
|
|
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,
|
|
embedding_model=embedding_model,
|
|
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]) -> dict[str, int]:
|
|
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.knowledge_bases (id, name, description)
|
|
VALUES (%s, %s, %s)
|
|
ON CONFLICT (id) DO UPDATE
|
|
SET name = EXCLUDED.name,
|
|
description = EXCLUDED.description,
|
|
updated_at = now()
|
|
""",
|
|
(KNOWLEDGE_BASE_ID, "虚构地质 PoC 知识库", "仅含公开的合成验证文本"),
|
|
)
|
|
connection.execute(
|
|
"""
|
|
INSERT INTO rag.access_scopes (id, knowledge_base_id, name)
|
|
VALUES (%s, %s, %s)
|
|
ON CONFLICT (id) DO NOTHING
|
|
""",
|
|
(ACCESS_SCOPE_ID, KNOWLEDGE_BASE_ID, "synthetic-demo"),
|
|
)
|
|
|
|
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,
|
|
KNOWLEDGE_BASE_ID,
|
|
ACCESS_SCOPE_ID,
|
|
item.raw_sha256,
|
|
f"{item.source_id}.json",
|
|
f"synthetic/{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,
|
|
KNOWLEDGE_BASE_ID,
|
|
item.document_id,
|
|
item.version_id,
|
|
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(
|
|
"""
|
|
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
|
|
""",
|
|
(KNOWLEDGE_BASE_ID,),
|
|
).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, ...],
|
|
) -> 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 id, metadata, embedding_text,
|
|
1 - (embedding <=> %s) AS vector_score
|
|
FROM rag.chunks
|
|
WHERE searchable
|
|
AND knowledge_base_id = %s
|
|
AND access_scope_id = %s
|
|
ORDER BY embedding <=> %s
|
|
LIMIT %s
|
|
""",
|
|
(
|
|
Vector(list(query_vector)),
|
|
KNOWLEDGE_BASE_ID,
|
|
ACCESS_SCOPE_ID,
|
|
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,
|
|
) -> 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])
|
|
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()
|
|
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_embedder: BailianEmbeddingAdapter | None = None
|
|
cloud_reranker: BailianRerankerAdapter | 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":
|
|
api_key = settings.bailian_api_key()
|
|
cloud_embedder = BailianEmbeddingAdapter(
|
|
api_key=api_key,
|
|
base_url=settings.bailian_openai_base_url,
|
|
model=settings.embedding_model,
|
|
dimensions=settings.embedding_dimension,
|
|
timeout_seconds=settings.model_timeout_seconds,
|
|
max_retries=settings.model_max_retries,
|
|
)
|
|
cloud_reranker = BailianRerankerAdapter(
|
|
api_key=api_key,
|
|
base_url=settings.bailian_rerank_base_url,
|
|
model=settings.rerank_model,
|
|
timeout_seconds=settings.model_timeout_seconds,
|
|
max_retries=settings.model_max_retries,
|
|
)
|
|
embedder = cloud_embedder
|
|
reranker = cloud_reranker
|
|
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,
|
|
)
|
|
counts = write_chunks(settings, prepared)
|
|
metrics = await evaluate_queries(settings, queries, embedder, reranker)
|
|
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_embedder is not None:
|
|
await cloud_embedder.aclose()
|
|
if cloud_reranker is not None:
|
|
await cloud_reranker.aclose()
|
|
|
|
|
|
def main() -> None:
|
|
raise SystemExit(asyncio.run(async_main()))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|