Make the first RAG slice executable without risking production data
Some checks failed
verify / verify (push) Has been cancelled
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
This commit is contained in:
666
backend/app/tools/seed_demo.py
Normal file
666
backend/app/tools/seed_demo.py
Normal file
@@ -0,0 +1,666 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user