Make the first RAG slice executable without risking production data
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:
2026-07-12 15:41:58 +08:00
parent ec1acb36b5
commit f4ba5d5342
61 changed files with 6886 additions and 20 deletions

3
backend/app/__init__.py Normal file
View File

@@ -0,0 +1,3 @@
"""Geological RAG backend package."""
__version__ = "0.1.0"

View File

@@ -1 +0,0 @@

View File

@@ -0,0 +1,11 @@
"""Alibaba Cloud Model Studio provider adapters."""
from app.adapters.bailian.chat import BailianChatAdapter
from app.adapters.bailian.embedding import BailianEmbeddingAdapter
from app.adapters.bailian.rerank import BailianRerankerAdapter
__all__ = [
"BailianChatAdapter",
"BailianEmbeddingAdapter",
"BailianRerankerAdapter",
]

View File

@@ -0,0 +1,382 @@
"""Shared HTTP and validation machinery for Alibaba Cloud Model Studio."""
from __future__ import annotations
import asyncio
import re
import secrets
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import UTC, datetime
from email.utils import parsedate_to_datetime
from time import perf_counter
from typing import Any, Self
from urllib.parse import urlsplit
import httpx
from app.ports.model_providers import (
ModelProviderError,
ProviderErrorKind,
ProviderUsage,
TokenCounter,
)
_SAFE_IDENTIFIER = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.:/-]{0,127}$")
_BAILIAN_BEIJING_HOST = re.compile(
r"^[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.cn-beijing\.maas\.aliyuncs\.com$"
)
def conservative_utf8_token_count(text: str) -> int:
"""Return a dependency-free conservative upper bound for byte-BPE tokens.
Production may inject the provider tokenizer through ``token_counter``. A
UTF-8 byte count is deliberately conservative and, unlike a word count,
does not undercount Chinese text or byte-fallback tokens.
"""
return len(text.encode("utf-8"))
def count_tokens(
counter: TokenCounter,
text: str,
*,
operation: str,
) -> int:
try:
count = counter(text)
except Exception:
raise sanitized_error(
operation=operation,
kind=ProviderErrorKind.INVALID_REQUEST,
provider_code="token_count_failed",
) from None
if isinstance(count, bool) or not isinstance(count, int) or count < 0:
raise sanitized_error(
operation=operation,
kind=ProviderErrorKind.INVALID_REQUEST,
provider_code="invalid_token_count",
)
return count
def sanitized_error(
*,
operation: str,
kind: ProviderErrorKind,
status_code: int | None = None,
provider_code: str | None = None,
request_id: str | None = None,
retryable: bool = False,
) -> ModelProviderError:
return ModelProviderError(
operation=operation,
kind=kind,
status_code=status_code,
provider_code=provider_code,
request_id=request_id,
retryable=retryable,
)
def invalid_request(operation: str, code: str) -> ModelProviderError:
return sanitized_error(
operation=operation,
kind=ProviderErrorKind.INVALID_REQUEST,
provider_code=code,
)
def invalid_response(operation: str, code: str) -> ModelProviderError:
return sanitized_error(
operation=operation,
kind=ProviderErrorKind.INVALID_RESPONSE,
provider_code=code,
)
def parse_usage(value: Any) -> ProviderUsage:
if not isinstance(value, Mapping):
return ProviderUsage()
return ProviderUsage(
input_tokens=_first_nonnegative_int(
value.get("prompt_tokens"),
value.get("input_tokens"),
),
output_tokens=_first_nonnegative_int(
value.get("completion_tokens"),
value.get("output_tokens"),
),
total_tokens=_nonnegative_int(value.get("total_tokens")),
)
def response_model(
body: Mapping[str, Any],
requested_model: str,
*,
sensitive_values: Sequence[str] = (),
) -> str:
return safe_identifier(body.get("model"), sensitive_values=sensitive_values) or requested_model
def safe_identifier(
value: Any,
*,
sensitive_values: Sequence[str] = (),
) -> str | None:
if not isinstance(value, str) or not _SAFE_IDENTIFIER.fullmatch(value):
return None
if any(value in sensitive or sensitive in value for sensitive in sensitive_values if sensitive):
return None
return value
def extract_request_id(
body: Mapping[str, Any],
*,
sensitive_values: Sequence[str] = (),
) -> str | None:
return safe_identifier(
body.get("id") or body.get("request_id"),
sensitive_values=sensitive_values,
)
def _nonnegative_int(value: Any) -> int | None:
if isinstance(value, bool) or not isinstance(value, int) or value < 0:
return None
return int(value)
def _first_nonnegative_int(*values: Any) -> int | None:
for value in values:
parsed = _nonnegative_int(value)
if parsed is not None:
return parsed
return None
@dataclass(frozen=True, slots=True)
class JsonResponse:
body: Mapping[str, Any]
elapsed_ms: float
class BailianHttpAdapter:
"""Minimal async HTTP client with sanitized error translation."""
def __init__(
self,
*,
api_key: str,
base_url: str,
expected_path: str,
http_client: httpx.AsyncClient | None,
timeout_seconds: float,
max_retries: int = 0,
retry_base_seconds: float = 0.5,
) -> None:
if not api_key or api_key != api_key.strip():
raise invalid_request("configuration", "invalid_api_key_value")
if timeout_seconds <= 0:
raise invalid_request("configuration", "invalid_timeout")
if isinstance(max_retries, bool) or not isinstance(max_retries, int) or max_retries < 0:
raise invalid_request("configuration", "invalid_max_retries")
if retry_base_seconds < 0:
raise invalid_request("configuration", "invalid_retry_base")
self._base_url = _validated_base_url(base_url, expected_path)
self._api_key = api_key
self._owns_client = http_client is None
self._client = http_client or httpx.AsyncClient(
timeout=httpx.Timeout(timeout_seconds),
follow_redirects=False,
)
self._max_retries = max_retries
self._retry_base_seconds = retry_base_seconds
async def aclose(self) -> None:
if self._owns_client:
await self._client.aclose()
async def __aenter__(self) -> Self:
return self
async def __aexit__(self, *_: object) -> None:
await self.aclose()
def _url(self, path: str) -> str:
return f"{self._base_url}/{path.lstrip('/')}"
def _headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self._api_key}",
"Content-Type": "application/json",
}
async def _post_json(
self,
*,
operation: str,
path: str,
payload: Mapping[str, Any],
sensitive_values: Sequence[str],
) -> JsonResponse:
started = perf_counter()
for attempt in range(self._max_retries + 1):
try:
response = await self._client.post(
self._url(path),
headers=self._headers(),
json=payload,
)
except httpx.TimeoutException:
error = sanitized_error(
operation=operation,
kind=ProviderErrorKind.TIMEOUT,
provider_code="request_timeout",
retryable=True,
)
if await self._maybe_retry(error, attempt=attempt, response=None):
continue
raise error from None
except httpx.HTTPError:
error = sanitized_error(
operation=operation,
kind=ProviderErrorKind.TRANSPORT,
provider_code="transport_error",
retryable=True,
)
if await self._maybe_retry(error, attempt=attempt, response=None):
continue
raise error from None
if response.status_code >= 400:
try:
self._raise_http_error(
operation=operation,
response=response,
sensitive_values=(*sensitive_values, self._api_key),
)
except ModelProviderError as error:
if await self._maybe_retry(error, attempt=attempt, response=response):
continue
raise
try:
body = response.json()
except ValueError:
raise invalid_response(operation, "invalid_json") from None
if not isinstance(body, Mapping):
raise invalid_response(operation, "invalid_json_object")
return JsonResponse(body=body, elapsed_ms=(perf_counter() - started) * 1000)
raise AssertionError("bounded provider retry loop exhausted unexpectedly")
async def _maybe_retry(
self,
error: ModelProviderError,
*,
attempt: int,
response: httpx.Response | None,
) -> bool:
if not error.retryable or attempt >= self._max_retries:
return False
await asyncio.sleep(self._retry_delay(attempt=attempt, response=response))
return True
def _retry_delay(self, *, attempt: int, response: httpx.Response | None) -> float:
if response is not None:
retry_after = response.headers.get("Retry-After")
if retry_after:
parsed_delay = _parse_retry_after(retry_after)
if parsed_delay is not None:
return min(max(parsed_delay, 0.0), 30.0)
base_delay = min(self._retry_base_seconds * (2**attempt), 30.0)
if base_delay == 0:
return 0.0
jitter = base_delay * (float(secrets.randbelow(251)) / 1000.0)
return float(min(base_delay + jitter, 30.0))
def _raise_http_error(
self,
*,
operation: str,
response: httpx.Response,
sensitive_values: Sequence[str],
) -> None:
body: Mapping[str, Any] = {}
try:
decoded = response.json()
if isinstance(decoded, Mapping):
body = decoded
except (ValueError, httpx.ResponseNotRead):
pass
nested_error = body.get("error")
error_body = nested_error if isinstance(nested_error, Mapping) else body
code = safe_identifier(
error_body.get("code"),
sensitive_values=sensitive_values,
)
request_id = extract_request_id(body, sensitive_values=sensitive_values)
kind, retryable = _kind_for_status(response.status_code)
raise sanitized_error(
operation=operation,
kind=kind,
status_code=response.status_code,
provider_code=code,
request_id=request_id,
retryable=retryable,
) from None
def _validated_base_url(base_url: str, expected_path: str) -> str:
if not base_url or base_url != base_url.strip():
raise invalid_request("configuration", "invalid_base_url")
parsed = urlsplit(base_url)
normalized_path = parsed.path.rstrip("/")
if (
parsed.scheme != "https"
or not parsed.hostname
or not _BAILIAN_BEIJING_HOST.fullmatch(parsed.hostname)
or parsed.username is not None
or parsed.password is not None
or parsed.query
or parsed.fragment
or normalized_path != expected_path
):
raise invalid_request("configuration", "invalid_base_url")
return base_url.rstrip("/")
def _kind_for_status(status_code: int) -> tuple[ProviderErrorKind, bool]:
if status_code == 400:
return ProviderErrorKind.INVALID_REQUEST, False
if status_code == 401:
return ProviderErrorKind.AUTHENTICATION, False
if status_code == 403:
return ProviderErrorKind.PERMISSION_DENIED, False
if status_code == 404:
return ProviderErrorKind.NOT_FOUND, False
if status_code == 408:
return ProviderErrorKind.TIMEOUT, True
if status_code == 429:
return ProviderErrorKind.RATE_LIMITED, True
return ProviderErrorKind.UPSTREAM, status_code >= 500
def _parse_retry_after(value: str) -> float | None:
try:
return float(value)
except ValueError:
try:
retry_at = parsedate_to_datetime(value)
except (TypeError, ValueError, OverflowError):
return None
if retry_at.tzinfo is None:
retry_at = retry_at.replace(tzinfo=UTC)
return (retry_at - datetime.now(UTC)).total_seconds()

View File

@@ -0,0 +1,326 @@
"""Alibaba Cloud Model Studio OpenAI-compatible chat adapter."""
from __future__ import annotations
import json
from collections.abc import AsyncIterator, Mapping, Sequence
from time import perf_counter
from typing import Any
import httpx
from app.adapters.bailian._base import (
BailianHttpAdapter,
extract_request_id,
invalid_request,
invalid_response,
parse_usage,
response_model,
sanitized_error,
)
from app.ports.model_providers import (
ChatCompletionResult,
ChatMessage,
ChatStreamEvent,
ModelProviderError,
ProviderErrorKind,
ProviderUsage,
)
_ALLOWED_ROLES = frozenset({"system", "user", "assistant"})
class BailianChatAdapter(BailianHttpAdapter):
"""Call chat completions with thinking and web search forcibly disabled."""
def __init__(
self,
*,
api_key: str,
base_url: str,
model: str = "deepseek-v4-flash",
http_client: httpx.AsyncClient | None = None,
timeout_seconds: float = 60.0,
max_retries: int = 0,
retry_base_seconds: float = 0.5,
) -> None:
if not model or model != model.strip():
raise invalid_request("chat.configuration", "invalid_model")
super().__init__(
api_key=api_key,
base_url=base_url,
expected_path="/compatible-mode/v1",
http_client=http_client,
timeout_seconds=timeout_seconds,
max_retries=max_retries,
retry_base_seconds=retry_base_seconds,
)
self._model = model
async def complete(
self,
messages: Sequence[ChatMessage],
*,
max_tokens: int,
) -> ChatCompletionResult:
operation = "chat.complete"
validated_messages = self._validate_messages(
messages,
max_tokens=max_tokens,
operation=operation,
)
payload = self._payload(validated_messages, max_tokens=max_tokens, stream=False)
sensitive_values = tuple(message.content for message in validated_messages)
response = await self._post_json(
operation=operation,
path="chat/completions",
payload=payload,
sensitive_values=sensitive_values,
)
content, finish_reason = self._parse_completion(
response.body,
operation=operation,
)
return ChatCompletionResult(
content=content,
finish_reason=finish_reason,
model=response_model(
response.body,
self._model,
sensitive_values=(*sensitive_values, self._api_key),
),
request_id=extract_request_id(
response.body,
sensitive_values=(*sensitive_values, self._api_key),
),
usage=parse_usage(response.body.get("usage")),
elapsed_ms=response.elapsed_ms,
)
async def stream(
self,
messages: Sequence[ChatMessage],
*,
max_tokens: int,
) -> AsyncIterator[ChatStreamEvent]:
operation = "chat.stream"
validated_messages = self._validate_messages(
messages,
max_tokens=max_tokens,
operation=operation,
)
payload = self._payload(validated_messages, max_tokens=max_tokens, stream=True)
sensitive_values = tuple(message.content for message in validated_messages)
started = perf_counter()
for attempt in range(self._max_retries + 1):
emitted = False
try:
async with self._client.stream(
"POST",
self._url("chat/completions"),
headers=self._headers(),
json=payload,
) as response:
if response.status_code >= 400:
await response.aread()
try:
self._raise_http_error(
operation=operation,
response=response,
sensitive_values=(*sensitive_values, self._api_key),
)
except ModelProviderError as error:
if await self._maybe_retry(
error,
attempt=attempt,
response=response,
):
continue
raise
async for line in response.aiter_lines():
if not line or line.startswith(":"):
continue
if not line.startswith("data:"):
raise invalid_response(operation, "invalid_sse_event")
raw_data = line[5:].strip()
if raw_data == "[DONE]":
return
event_body = self._decode_stream_event(raw_data, operation=operation)
event = self._parse_stream_event(
event_body,
operation=operation,
sensitive_values=sensitive_values,
elapsed_ms=(perf_counter() - started) * 1000,
)
emitted = True
yield event
return
except ModelProviderError as error:
if not emitted and await self._maybe_retry(
error,
attempt=attempt,
response=None,
):
continue
raise
except httpx.TimeoutException:
timeout_error = sanitized_error(
operation=operation,
kind=ProviderErrorKind.TIMEOUT,
provider_code="request_timeout",
retryable=True,
)
if not emitted and await self._maybe_retry(
timeout_error,
attempt=attempt,
response=None,
):
continue
raise timeout_error from None
except httpx.HTTPError:
transport_error = sanitized_error(
operation=operation,
kind=ProviderErrorKind.TRANSPORT,
provider_code="transport_error",
retryable=True,
)
if not emitted and await self._maybe_retry(
transport_error,
attempt=attempt,
response=None,
):
continue
raise transport_error from None
raise AssertionError("bounded chat retry loop exhausted unexpectedly")
def _validate_messages(
self,
messages: object,
*,
max_tokens: int,
operation: str,
) -> tuple[ChatMessage, ...]:
if isinstance(messages, (str, bytes)) or not isinstance(messages, Sequence):
raise invalid_request(operation, "invalid_message_collection")
validated = tuple(messages)
if not validated:
raise invalid_request(operation, "empty_messages")
if isinstance(max_tokens, bool) or not isinstance(max_tokens, int) or max_tokens <= 0:
raise invalid_request(operation, "invalid_max_tokens")
for message in validated:
if (
not isinstance(message, ChatMessage)
or message.role not in _ALLOWED_ROLES
or not isinstance(message.content, str)
or not message.content
):
raise invalid_request(operation, "invalid_message")
return validated
def _payload(
self,
messages: tuple[ChatMessage, ...],
*,
max_tokens: int,
stream: bool,
) -> dict[str, Any]:
payload: dict[str, Any] = {
"model": self._model,
"messages": [
{"role": message.role, "content": message.content} for message in messages
],
"max_tokens": max_tokens,
"stream": stream,
"enable_thinking": False,
"enable_search": False,
}
if stream:
payload["stream_options"] = {"include_usage": True}
return payload
def _parse_completion(
self,
body: Mapping[str, Any],
*,
operation: str,
) -> tuple[str, str | None]:
choices = body.get("choices")
if not isinstance(choices, list) or len(choices) != 1:
raise invalid_response(operation, "invalid_choices")
choice = choices[0]
if not isinstance(choice, Mapping):
raise invalid_response(operation, "invalid_choice")
message = choice.get("message")
if not isinstance(message, Mapping):
raise invalid_response(operation, "invalid_message")
content = message.get("content")
if not isinstance(content, str):
raise invalid_response(operation, "invalid_content")
finish_reason = choice.get("finish_reason")
if finish_reason is not None and not isinstance(finish_reason, str):
raise invalid_response(operation, "invalid_finish_reason")
return content, finish_reason
def _decode_stream_event(
self,
raw_data: str,
*,
operation: str,
) -> Mapping[str, Any]:
try:
decoded = json.loads(raw_data)
except (TypeError, ValueError):
raise invalid_response(operation, "invalid_stream_json") from None
if not isinstance(decoded, Mapping):
raise invalid_response(operation, "invalid_stream_object")
return decoded
def _parse_stream_event(
self,
body: Mapping[str, Any],
*,
operation: str,
sensitive_values: tuple[str, ...],
elapsed_ms: float,
) -> ChatStreamEvent:
choices = body.get("choices")
delta = ""
finish_reason: str | None = None
if choices not in (None, []):
if not isinstance(choices, list) or len(choices) != 1:
raise invalid_response(operation, "invalid_stream_choices")
choice = choices[0]
if not isinstance(choice, Mapping):
raise invalid_response(operation, "invalid_stream_choice")
raw_delta = choice.get("delta")
if not isinstance(raw_delta, Mapping):
raise invalid_response(operation, "invalid_stream_delta")
content = raw_delta.get("content", "")
if not isinstance(content, str):
raise invalid_response(operation, "invalid_stream_content")
delta = content
raw_finish_reason = choice.get("finish_reason")
if raw_finish_reason is not None and not isinstance(raw_finish_reason, str):
raise invalid_response(operation, "invalid_stream_finish_reason")
finish_reason = raw_finish_reason
usage = parse_usage(body.get("usage"))
if choices in (None, []) and usage == ProviderUsage():
raise invalid_response(operation, "empty_stream_event")
return ChatStreamEvent(
delta=delta,
finish_reason=finish_reason,
model=response_model(
body,
self._model,
sensitive_values=(*sensitive_values, self._api_key),
),
request_id=extract_request_id(
body,
sensitive_values=(*sensitive_values, self._api_key),
),
usage=usage,
elapsed_ms=elapsed_ms,
)

View File

@@ -0,0 +1,177 @@
"""Alibaba Cloud Model Studio OpenAI-compatible embedding adapter."""
from __future__ import annotations
import math
from collections.abc import Mapping, Sequence
from typing import Any
import httpx
from app.adapters.bailian._base import (
BailianHttpAdapter,
conservative_utf8_token_count,
count_tokens,
extract_request_id,
invalid_request,
invalid_response,
parse_usage,
response_model,
)
from app.ports.model_providers import EmbeddingResult, TokenCounter
EMBEDDING_MAX_BATCH_SIZE = 10
EMBEDDING_MAX_TOKENS_PER_TEXT = 8_192
EMBEDDING_MAX_TOKENS_PER_REQUEST = 33_000
EMBEDDING_DIMENSIONS = 1_024
class BailianEmbeddingAdapter(BailianHttpAdapter):
"""Call ``compatible-mode/v1/embeddings`` with strict input/output checks."""
def __init__(
self,
*,
api_key: str,
base_url: str,
model: str = "text-embedding-v4",
dimensions: int = EMBEDDING_DIMENSIONS,
token_counter: TokenCounter = conservative_utf8_token_count,
http_client: httpx.AsyncClient | None = None,
timeout_seconds: float = 30.0,
max_retries: int = 0,
retry_base_seconds: float = 0.5,
) -> None:
if not model or model != model.strip():
raise invalid_request("embedding.configuration", "invalid_model")
if isinstance(dimensions, bool) or dimensions != EMBEDDING_DIMENSIONS:
raise invalid_request("embedding.configuration", "unsupported_dimensions")
if not callable(token_counter):
raise invalid_request("embedding.configuration", "invalid_token_counter")
super().__init__(
api_key=api_key,
base_url=base_url,
expected_path="/compatible-mode/v1",
http_client=http_client,
timeout_seconds=timeout_seconds,
max_retries=max_retries,
retry_base_seconds=retry_base_seconds,
)
self._model = model
self._dimensions = dimensions
self._token_counter = token_counter
async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult:
operation = "embedding.create"
validated_texts = self._validate_texts(texts, operation=operation)
payload = {
"model": self._model,
"input": list(validated_texts),
"dimensions": self._dimensions,
"encoding_format": "float",
}
response = await self._post_json(
operation=operation,
path="embeddings",
payload=payload,
sensitive_values=validated_texts,
)
vectors = self._parse_vectors(
response.body,
expected_count=len(validated_texts),
operation=operation,
)
return EmbeddingResult(
vectors=vectors,
model=response_model(
response.body,
self._model,
sensitive_values=(*validated_texts, self._api_key),
),
request_id=extract_request_id(
response.body,
sensitive_values=(*validated_texts, self._api_key),
),
usage=parse_usage(response.body.get("usage")),
elapsed_ms=response.elapsed_ms,
)
async def embed_query(self, text: str) -> EmbeddingResult:
return await self.embed_documents((text,))
def _validate_texts(
self,
texts: Sequence[str],
*,
operation: str,
) -> tuple[str, ...]:
if isinstance(texts, (str, bytes)) or not isinstance(texts, Sequence):
raise invalid_request(operation, "invalid_input_collection")
validated = tuple(texts)
if not validated:
raise invalid_request(operation, "empty_input")
if len(validated) > EMBEDDING_MAX_BATCH_SIZE:
raise invalid_request(operation, "batch_size_exceeded")
total_tokens = 0
for text in validated:
if not isinstance(text, str) or not text:
raise invalid_request(operation, "invalid_input_text")
token_count = count_tokens(
self._token_counter,
text,
operation=operation,
)
if token_count > EMBEDDING_MAX_TOKENS_PER_TEXT:
raise invalid_request(operation, "text_token_limit_exceeded")
total_tokens += token_count
if total_tokens > EMBEDDING_MAX_TOKENS_PER_REQUEST:
raise invalid_request(operation, "request_token_limit_exceeded")
return validated
def _parse_vectors(
self,
body: Mapping[str, Any],
*,
expected_count: int,
operation: str,
) -> tuple[tuple[float, ...], ...]:
data = body.get("data")
if not isinstance(data, list) or len(data) != expected_count:
raise invalid_response(operation, "invalid_embedding_count")
by_index: list[tuple[float, ...] | None] = [None] * expected_count
for item in data:
if not isinstance(item, Mapping):
raise invalid_response(operation, "invalid_embedding_item")
index = item.get("index")
if (
isinstance(index, bool)
or not isinstance(index, int)
or not 0 <= index < expected_count
or by_index[index] is not None
):
raise invalid_response(operation, "invalid_embedding_index")
raw_vector = item.get("embedding")
if not isinstance(raw_vector, list) or len(raw_vector) != self._dimensions:
raise invalid_response(operation, "invalid_embedding_dimensions")
vector: list[float] = []
for component in raw_vector:
if isinstance(component, bool) or not isinstance(component, (int, float)):
raise invalid_response(operation, "invalid_embedding_component")
normalized = float(component)
if not math.isfinite(normalized):
raise invalid_response(operation, "invalid_embedding_component")
vector.append(normalized)
norm = math.hypot(*vector)
if not math.isfinite(norm) or norm <= 0:
raise invalid_response(operation, "invalid_embedding_norm")
by_index[index] = tuple(vector)
if any(vector is None for vector in by_index):
raise invalid_response(operation, "missing_embedding_index")
return tuple(vector for vector in by_index if vector is not None)

View File

@@ -0,0 +1,208 @@
"""Alibaba Cloud Model Studio compatible rerank adapter."""
from __future__ import annotations
import math
from collections.abc import Mapping, Sequence
from typing import Any
import httpx
from app.adapters.bailian._base import (
BailianHttpAdapter,
conservative_utf8_token_count,
count_tokens,
extract_request_id,
invalid_request,
invalid_response,
parse_usage,
response_model,
)
from app.ports.model_providers import RankedItem, RerankResult, TokenCounter
RERANK_MAX_DOCUMENTS = 500
RERANK_MAX_TOKENS_PER_TEXT = 4_000
RERANK_MAX_TOKENS_PER_REQUEST = 120_000
class BailianRerankerAdapter(BailianHttpAdapter):
"""Call ``compatible-api/v1/reranks`` and map indices locally."""
def __init__(
self,
*,
api_key: str,
base_url: str,
model: str = "qwen3-rerank",
token_counter: TokenCounter = conservative_utf8_token_count,
http_client: httpx.AsyncClient | None = None,
timeout_seconds: float = 30.0,
max_retries: int = 0,
retry_base_seconds: float = 0.5,
) -> None:
if not model or model != model.strip():
raise invalid_request("rerank.configuration", "invalid_model")
if not callable(token_counter):
raise invalid_request("rerank.configuration", "invalid_token_counter")
super().__init__(
api_key=api_key,
base_url=base_url,
expected_path="/compatible-api/v1",
http_client=http_client,
timeout_seconds=timeout_seconds,
max_retries=max_retries,
retry_base_seconds=retry_base_seconds,
)
self._model = model
self._token_counter = token_counter
async def rerank(
self,
query: str,
documents: Sequence[str],
*,
top_n: int,
instruct: str | None = None,
) -> RerankResult:
operation = "rerank.create"
validated_documents = self._validate_request(
query=query,
documents=documents,
top_n=top_n,
instruct=instruct,
operation=operation,
)
payload: dict[str, Any] = {
"model": self._model,
"query": query,
"documents": list(validated_documents),
"top_n": top_n,
}
if instruct is not None:
payload["instruct"] = instruct
sensitive_values = (
query,
*validated_documents,
*((instruct,) if instruct is not None else ()),
)
response = await self._post_json(
operation=operation,
path="reranks",
payload=payload,
sensitive_values=sensitive_values,
)
items = self._parse_results(
response.body,
documents=validated_documents,
top_n=top_n,
operation=operation,
)
return RerankResult(
items=items,
model=response_model(
response.body,
self._model,
sensitive_values=(*sensitive_values, self._api_key),
),
request_id=extract_request_id(
response.body,
sensitive_values=(*sensitive_values, self._api_key),
),
usage=parse_usage(response.body.get("usage")),
elapsed_ms=response.elapsed_ms,
)
def _validate_request(
self,
*,
query: str,
documents: Sequence[str],
top_n: int,
instruct: str | None,
operation: str,
) -> tuple[str, ...]:
if not isinstance(query, str) or not query:
raise invalid_request(operation, "invalid_query")
if isinstance(documents, (str, bytes)) or not isinstance(documents, Sequence):
raise invalid_request(operation, "invalid_document_collection")
validated_documents = tuple(documents)
if not validated_documents:
raise invalid_request(operation, "empty_documents")
if len(validated_documents) > RERANK_MAX_DOCUMENTS:
raise invalid_request(operation, "document_count_exceeded")
if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n <= 0:
raise invalid_request(operation, "invalid_top_n")
if instruct is not None and (not isinstance(instruct, str) or not instruct):
raise invalid_request(operation, "invalid_instruct")
query_tokens = count_tokens(
self._token_counter,
query,
operation=operation,
)
if query_tokens > RERANK_MAX_TOKENS_PER_TEXT:
raise invalid_request(operation, "query_token_limit_exceeded")
document_tokens_total = 0
for document in validated_documents:
if not isinstance(document, str) or not document:
raise invalid_request(operation, "invalid_document")
document_tokens = count_tokens(
self._token_counter,
document,
operation=operation,
)
if document_tokens > RERANK_MAX_TOKENS_PER_TEXT:
raise invalid_request(operation, "document_token_limit_exceeded")
document_tokens_total += document_tokens
request_tokens = query_tokens * len(validated_documents) + document_tokens_total
if request_tokens > RERANK_MAX_TOKENS_PER_REQUEST:
raise invalid_request(operation, "request_token_limit_exceeded")
return validated_documents
def _parse_results(
self,
body: Mapping[str, Any],
*,
documents: tuple[str, ...],
top_n: int,
operation: str,
) -> tuple[RankedItem, ...]:
results = body.get("results")
if not isinstance(results, list) or len(results) > min(top_n, len(documents)):
raise invalid_response(operation, "invalid_rerank_count")
seen_indices: set[int] = set()
parsed: list[RankedItem] = []
previous_score = math.inf
for item in results:
if not isinstance(item, Mapping):
raise invalid_response(operation, "invalid_rerank_item")
index = item.get("index")
if (
isinstance(index, bool)
or not isinstance(index, int)
or not 0 <= index < len(documents)
or index in seen_indices
):
raise invalid_response(operation, "invalid_rerank_index")
raw_score = item.get("relevance_score")
if isinstance(raw_score, bool) or not isinstance(raw_score, (int, float)):
raise invalid_response(operation, "invalid_rerank_score")
score = float(raw_score)
if not math.isfinite(score) or not 0 <= score <= 1 or score > previous_score:
raise invalid_response(operation, "invalid_rerank_score")
seen_indices.add(index)
previous_score = score
parsed.append(
RankedItem(
index=index,
relevance_score=score,
document=documents[index],
)
)
return tuple(parsed)

View File

@@ -0,0 +1,154 @@
"""Deterministic local providers for tests and offline Docker verification only."""
from __future__ import annotations
import hashlib
import math
import re
import time
from collections.abc import Sequence
from app.ports.model_providers import (
EmbeddingResult,
ModelProviderError,
ProviderErrorKind,
ProviderUsage,
RankedItem,
RerankResult,
)
_TOKEN_PATTERN = re.compile(r"[\u3400-\u9fff]+|[a-zA-Z0-9_.%+-]+")
def invalid_input(operation: str, code: str) -> ModelProviderError:
return ModelProviderError(
operation=operation,
kind=ProviderErrorKind.INVALID_REQUEST,
provider_code=code,
)
def lexical_features(text: str) -> tuple[str, ...]:
"""Create stable character n-grams/words without pretending to be a tokenizer."""
features: list[str] = []
for token in _TOKEN_PATTERN.findall(text.lower()):
if "\u3400" <= token[0] <= "\u9fff":
features.extend(token)
features.extend(token[index : index + 2] for index in range(len(token) - 1))
else:
features.append(token)
return tuple(features)
class FakeEmbeddingProvider:
"""Feature-hashing embedding used to validate plumbing without cloud calls."""
def __init__(self, dimension: int = 1024) -> None:
if dimension < 1:
raise ValueError("dimension must be positive")
self.dimension = dimension
def _vector(self, text: str) -> tuple[float, ...]:
vector = [0.0] * self.dimension
for feature in lexical_features(text):
digest = hashlib.sha256(feature.encode("utf-8")).digest()
index = int.from_bytes(digest[:4], "big") % self.dimension
sign = 1.0 if digest[4] & 1 else -1.0
vector[index] += sign
norm = math.sqrt(sum(value * value for value in vector))
if norm == 0:
vector[0] = 1.0
norm = 1.0
return tuple(value / norm for value in vector)
async def _embed(self, texts: Sequence[str]) -> EmbeddingResult:
if isinstance(texts, (str, bytes)) or not isinstance(texts, Sequence):
raise invalid_input("fake.embedding", "invalid_input_collection")
if not texts:
raise invalid_input("fake.embedding", "empty_input")
if len(texts) > 10:
raise invalid_input("fake.embedding", "batch_size_exceeded")
token_counts = []
for text in texts:
if not isinstance(text, str) or not text:
raise invalid_input("fake.embedding", "invalid_input_text")
token_count = len(text.encode("utf-8"))
if token_count > 8_192:
raise invalid_input("fake.embedding", "text_token_limit_exceeded")
token_counts.append(token_count)
if sum(token_counts) > 33_000:
raise invalid_input("fake.embedding", "request_token_limit_exceeded")
started = time.perf_counter()
vectors = tuple(self._vector(text) for text in texts)
return EmbeddingResult(
vectors=vectors,
model="fake-feature-hash-v1",
request_id=None,
usage=ProviderUsage(input_tokens=sum(len(lexical_features(text)) for text in texts)),
elapsed_ms=(time.perf_counter() - started) * 1000,
)
async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult:
return await self._embed(texts)
async def embed_query(self, text: str) -> EmbeddingResult:
return await self._embed([text])
class FakeReranker:
"""Lexical-overlap reranker for deterministic offline flow tests."""
async def rerank(
self,
query: str,
documents: Sequence[str],
*,
top_n: int,
instruct: str | None = None,
) -> RerankResult:
del instruct
if not isinstance(query, str) or not query:
raise invalid_input("fake.rerank", "invalid_query")
if isinstance(documents, (str, bytes)) or not isinstance(documents, Sequence):
raise invalid_input("fake.rerank", "invalid_document_collection")
if not documents:
raise invalid_input("fake.rerank", "empty_documents")
if len(documents) > 500:
raise invalid_input("fake.rerank", "document_count_exceeded")
if isinstance(top_n, bool) or not isinstance(top_n, int) or top_n <= 0:
raise invalid_input("fake.rerank", "invalid_top_n")
query_tokens = len(query.encode("utf-8"))
if query_tokens > 4_000:
raise invalid_input("fake.rerank", "query_token_limit_exceeded")
document_tokens = []
for document in documents:
if not isinstance(document, str) or not document:
raise invalid_input("fake.rerank", "invalid_document")
count = len(document.encode("utf-8"))
if count > 4_000:
raise invalid_input("fake.rerank", "document_token_limit_exceeded")
document_tokens.append(count)
if query_tokens * len(documents) + sum(document_tokens) > 120_000:
raise invalid_input("fake.rerank", "request_token_limit_exceeded")
started = time.perf_counter()
query_features = set(lexical_features(query))
ranked: list[RankedItem] = []
for index, document in enumerate(documents):
document_features = set(lexical_features(document))
union = query_features | document_features
score = len(query_features & document_features) / len(union) if union else 0.0
ranked.append(RankedItem(index=index, relevance_score=score, document=document))
ranked.sort(key=lambda item: (-item.relevance_score, item.index))
return RerankResult(
items=tuple(ranked[:top_n]),
model="fake-lexical-rerank-v1",
request_id=None,
usage=ProviderUsage(
input_tokens=len(query_features)
+ sum(len(lexical_features(document)) for document in documents)
),
elapsed_ms=(time.perf_counter() - started) * 1000,
)

View File

@@ -1 +0,0 @@

View File

@@ -0,0 +1 @@
"""Configuration and security primitives."""

128
backend/app/core/config.py Normal file
View File

@@ -0,0 +1,128 @@
"""Typed runtime configuration loaded from non-secret environment values."""
from functools import lru_cache
from pathlib import Path
from typing import Literal, Self
from urllib.parse import urlsplit
from pydantic import Field, field_validator, model_validator
from pydantic_settings import BaseSettings, SettingsConfigDict
from sqlalchemy import URL
from app.core.secrets import read_secret_file
class Settings(BaseSettings):
"""Application settings; secret values stay in mounted files."""
model_config = SettingsConfigDict(
env_file=".env",
env_file_encoding="utf-8",
case_sensitive=False,
extra="ignore",
)
app_env: Literal["development", "test", "production"] = "development"
app_name: str = "geological-rag"
app_secret_key_file: Path = Path("/run/secrets/app_secret_key")
postgres_host: str = "db"
postgres_port: int = Field(default=5432, ge=1, le=65535)
postgres_db: str = "geological_rag"
postgres_user: str = "geological_rag_app"
postgres_password_file: Path = Path("/run/secrets/postgres_app_password")
upload_root: Path = Path("/data/uploads")
max_upload_mb: int = Field(default=100, ge=1, le=2048)
bailian_openai_base_url: str = (
"https://<workspace-id>.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"
)
bailian_native_base_url: str = "https://<workspace-id>.cn-beijing.maas.aliyuncs.com/api/v1"
bailian_rerank_base_url: str = (
"https://<workspace-id>.cn-beijing.maas.aliyuncs.com/compatible-api/v1"
)
dashscope_api_key_file: Path = Path("/run/secrets/bailian_api_key")
embedding_model: str = "text-embedding-v4"
embedding_dimension: Literal[1024] = 1024
rerank_model: str = "qwen3-rerank"
llm_model: str = "deepseek-v4-flash"
chunk_target_tokens: int = Field(default=512, ge=120, le=800)
chunk_max_tokens: int = Field(default=800, ge=120, le=4000)
chunk_overlap_tokens: int = Field(default=64, ge=0, le=256)
vector_top_k: int = Field(default=50, ge=1, le=500)
rerank_top_n: int = Field(default=10, ge=1, le=500)
context_top_n: int = Field(default=8, ge=1, le=50)
max_context_tokens: int = Field(default=10_000, ge=1, le=100_000)
model_timeout_seconds: float = Field(default=90, gt=0, le=600)
model_max_retries: int = Field(default=3, ge=0, le=10)
model_max_concurrency: int = Field(default=4, ge=1, le=100)
worker_capabilities: str = "document_parse,embedding,rerank,evaluation"
@field_validator(
"bailian_openai_base_url",
"bailian_native_base_url",
"bailian_rerank_base_url",
)
@classmethod
def normalize_base_url(cls, value: str) -> str:
return value.rstrip("/")
@model_validator(mode="after")
def validate_rag_limits(self) -> Self:
if self.chunk_overlap_tokens >= self.chunk_target_tokens:
raise ValueError("CHUNK_OVERLAP_TOKENS must be smaller than CHUNK_TARGET_TOKENS")
if self.chunk_target_tokens > self.chunk_max_tokens:
raise ValueError("CHUNK_TARGET_TOKENS must not exceed CHUNK_MAX_TOKENS")
if self.context_top_n > self.rerank_top_n:
raise ValueError("CONTEXT_TOP_N must not exceed RERANK_TOP_N")
if self.rerank_top_n > self.vector_top_k:
raise ValueError("RERANK_TOP_N must not exceed VECTOR_TOP_K")
return self
def bailian_api_key(self) -> str:
self.validate_live_bailian_endpoints()
return read_secret_file(self.dashscope_api_key_file)
def validate_live_bailian_endpoints(self) -> str:
endpoints = {
self.bailian_openai_base_url: "/compatible-mode/v1",
self.bailian_rerank_base_url: "/compatible-api/v1",
}
hosts: set[str] = set()
for endpoint, expected_path in endpoints.items():
parsed = urlsplit(endpoint)
host = parsed.hostname or ""
if (
parsed.scheme != "https"
or parsed.path.rstrip("/") != expected_path
or parsed.username is not None
or parsed.password is not None
or parsed.query
or parsed.fragment
or "<" in endpoint
or not host.endswith(".cn-beijing.maas.aliyuncs.com")
):
raise ValueError("Bailian endpoint is not an approved Beijing MaaS URL")
hosts.add(host)
if len(hosts) != 1:
raise ValueError("Bailian endpoints must use the same workspace host")
return hosts.pop()
def database_url(self) -> URL:
return URL.create(
drivername="postgresql+psycopg",
username=self.postgres_user,
password=read_secret_file(self.postgres_password_file),
host=self.postgres_host,
port=self.postgres_port,
database=self.postgres_db,
)
@lru_cache(maxsize=1)
def get_settings() -> Settings:
return Settings()

View File

@@ -0,0 +1,24 @@
"""Secret-file helpers that never include secret values in errors or reprs."""
from pathlib import Path
class SecretFileError(RuntimeError):
"""Raised when a configured secret file cannot be read safely."""
def read_secret_file(path: Path) -> str:
"""Read one Docker/Kubernetes secret without exposing its content."""
try:
if not path.is_file():
raise SecretFileError(f"Secret file is missing or not a file: {path}")
value = path.read_text(encoding="utf-8").strip()
except OSError as exc:
raise SecretFileError(f"Secret file cannot be read: {path}") from exc
if not value:
raise SecretFileError(f"Secret file is empty: {path}")
if "\n" in value or "\r" in value:
raise SecretFileError(f"Secret file must contain exactly one logical line: {path}")
return value

56
backend/app/main.py Normal file
View File

@@ -0,0 +1,56 @@
"""Minimal FastAPI entrypoint; product endpoints are added in stage 2."""
from typing import Any
import psycopg
import uvicorn
from fastapi import FastAPI, HTTPException, status
from app import __version__
from app.core.config import get_settings
from app.core.secrets import SecretFileError
app = FastAPI(title="Geological RAG API", version=__version__)
@app.get("/api/v1/health/live", tags=["health"])
def live() -> dict[str, str]:
return {"status": "ok", "version": __version__}
@app.get("/api/v1/health/ready", tags=["health"])
def ready() -> dict[str, str]:
settings = get_settings()
try:
dsn = settings.database_url().set(drivername="postgresql")
with psycopg.connect(
dsn.render_as_string(hide_password=False),
connect_timeout=2,
) as connection:
connection.execute("SELECT 1")
except (OSError, SecretFileError, psycopg.Error) as exc:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="database unavailable",
) from exc
return {"status": "ready"}
@app.get("/api/v1/meta", tags=["meta"])
def meta() -> dict[str, Any]:
settings = get_settings()
return {
"name": settings.app_name,
"environment": settings.app_env,
"version": __version__,
"models": {
"embedding": settings.embedding_model,
"rerank": settings.rerank_model,
"generation": settings.llm_model,
},
}
if __name__ == "__main__":
# Container ingress is controlled by Compose; only the edge proxy publishes a port.
uvicorn.run("app.main:app", host="0.0.0.0", port=8000) # noqa: S104

View File

@@ -1 +0,0 @@

View File

@@ -0,0 +1,17 @@
"""Persistence-level contracts shared by workers and integration tests."""
from .job_queue_sql import (
CLAIM_JOB_SQL,
COMPLETE_JOB_SQL,
FAIL_OR_RETRY_JOB_SQL,
HEARTBEAT_JOB_SQL,
REAP_EXPIRED_JOBS_SQL,
)
__all__ = [
"CLAIM_JOB_SQL",
"COMPLETE_JOB_SQL",
"FAIL_OR_RETRY_JOB_SQL",
"HEARTBEAT_JOB_SQL",
"REAP_EXPIRED_JOBS_SQL",
]

View File

@@ -0,0 +1,129 @@
"""Fenced PostgreSQL job-queue statements.
Callers must execute each statement in a transaction and treat an empty
``RETURNING`` result as loss of the lease. No external model call should hold
the claim transaction open.
"""
CLAIM_JOB_SQL = """
WITH candidate AS (
SELECT job.id
FROM rag.background_jobs AS job
WHERE job.attempt < job.max_attempts
AND job.required_capability = ANY(CAST(:worker_capabilities AS text[]))
AND (
(job.status = 'QUEUED' AND job.run_after <= now())
OR (job.status = 'RUNNING' AND job.lease_until < now())
)
ORDER BY job.priority DESC, job.run_after, job.created_at
FOR UPDATE SKIP LOCKED
LIMIT 1
)
UPDATE rag.background_jobs AS job
SET status = 'RUNNING',
lease_owner = :worker_id,
lease_token = gen_random_uuid(),
lease_until = now() + make_interval(secs => CAST(:lease_seconds AS integer)),
attempt = job.attempt + 1,
updated_at = now(),
finished_at = NULL
FROM candidate
WHERE job.id = candidate.id
RETURNING job.*
"""
HEARTBEAT_JOB_SQL = """
UPDATE rag.background_jobs AS job
SET lease_until = now() + make_interval(secs => CAST(:lease_seconds AS integer)),
updated_at = now()
WHERE job.id = :job_id
AND job.status = 'RUNNING'
AND job.lease_owner = :worker_id
AND job.lease_token = :lease_token
RETURNING job.id, job.lease_until
"""
COMPLETE_JOB_SQL = """
UPDATE rag.background_jobs AS job
SET status = 'SUCCEEDED',
progress = 100,
lease_owner = NULL,
lease_token = NULL,
lease_until = NULL,
finished_at = now(),
updated_at = now()
WHERE job.id = :job_id
AND job.status = 'RUNNING'
AND job.lease_owner = :worker_id
AND job.lease_token = :lease_token
RETURNING job.*
"""
FAIL_OR_RETRY_JOB_SQL = """
UPDATE rag.background_jobs AS job
SET status = CASE
WHEN job.attempt < job.max_attempts THEN 'QUEUED'
ELSE 'FAILED'
END,
run_after = CASE
WHEN job.attempt < job.max_attempts
THEN now() + make_interval(
secs => GREATEST(CAST(:retry_delay_seconds AS integer), 0)
)
ELSE job.run_after
END,
last_error_code = :error_code,
last_error_message = left(:error_message, 2000),
lease_owner = NULL,
lease_token = NULL,
lease_until = NULL,
finished_at = CASE
WHEN job.attempt < job.max_attempts THEN NULL
ELSE now()
END,
updated_at = now()
WHERE job.id = :job_id
AND job.status = 'RUNNING'
AND job.lease_owner = :worker_id
AND job.lease_token = :lease_token
RETURNING job.*
"""
REAP_EXPIRED_JOBS_SQL = """
WITH maintenance_lock AS (
SELECT pg_try_advisory_xact_lock(CAST(:lock_key AS bigint)) AS acquired
),
expired AS (
SELECT job.id
FROM rag.background_jobs AS job
CROSS JOIN maintenance_lock
WHERE maintenance_lock.acquired
AND job.status = 'RUNNING'
AND job.lease_until < now()
ORDER BY job.lease_until, job.created_at
FOR UPDATE OF job SKIP LOCKED
LIMIT CAST(:batch_size AS integer)
)
UPDATE rag.background_jobs AS job
SET status = CASE
WHEN job.attempt < job.max_attempts THEN 'QUEUED'
ELSE 'FAILED'
END,
run_after = CASE
WHEN job.attempt < job.max_attempts THEN now()
ELSE job.run_after
END,
last_error_code = 'WORKER_LEASE_EXPIRED',
last_error_message = 'Worker lease expired before a fenced terminal update.',
lease_owner = NULL,
lease_token = NULL,
lease_until = NULL,
finished_at = CASE
WHEN job.attempt < job.max_attempts THEN NULL
ELSE now()
END,
updated_at = now()
FROM expired
WHERE job.id = expired.id
RETURNING job.*
"""

View File

@@ -1 +0,0 @@

View File

@@ -0,0 +1,151 @@
"""Dependency-inversion ports for external model providers.
The domain and service layers depend on these provider-neutral value objects and
protocols. Vendor SDK/HTTP response objects must not cross this boundary.
"""
from __future__ import annotations
from collections.abc import AsyncIterator, Callable, Sequence
from dataclasses import dataclass
from enum import StrEnum
from typing import Protocol
TokenCounter = Callable[[str], int]
class ProviderErrorKind(StrEnum):
"""Stable, provider-neutral error categories."""
INVALID_REQUEST = "invalid_request"
AUTHENTICATION = "authentication"
PERMISSION_DENIED = "permission_denied"
NOT_FOUND = "not_found"
RATE_LIMITED = "rate_limited"
TIMEOUT = "timeout"
TRANSPORT = "transport"
UPSTREAM = "upstream"
INVALID_RESPONSE = "invalid_response"
class ModelProviderError(RuntimeError):
"""Sanitized provider failure safe for structured application handling.
The exception intentionally never stores request payloads, response bodies,
HTTP request objects, headers, or credentials.
"""
def __init__(
self,
*,
operation: str,
kind: ProviderErrorKind,
status_code: int | None = None,
provider_code: str | None = None,
request_id: str | None = None,
retryable: bool = False,
) -> None:
self.operation = operation
self.kind = kind
self.status_code = status_code
self.provider_code = provider_code
self.request_id = request_id
self.retryable = retryable
details = [f"kind={kind.value}"]
if status_code is not None:
details.append(f"status={status_code}")
if provider_code is not None:
details.append(f"code={provider_code}")
super().__init__(f"model provider {operation} failed ({', '.join(details)})")
@dataclass(frozen=True, slots=True)
class ProviderUsage:
input_tokens: int | None = None
output_tokens: int | None = None
total_tokens: int | None = None
@dataclass(frozen=True, slots=True)
class EmbeddingResult:
vectors: tuple[tuple[float, ...], ...]
model: str
request_id: str | None
usage: ProviderUsage
elapsed_ms: float
@dataclass(frozen=True, slots=True)
class RankedItem:
index: int
relevance_score: float
document: str
@dataclass(frozen=True, slots=True)
class RerankResult:
items: tuple[RankedItem, ...]
model: str
request_id: str | None
usage: ProviderUsage
elapsed_ms: float
@dataclass(frozen=True, slots=True)
class ChatMessage:
role: str
content: str
@dataclass(frozen=True, slots=True)
class ChatCompletionResult:
content: str
finish_reason: str | None
model: str
request_id: str | None
usage: ProviderUsage
elapsed_ms: float
@dataclass(frozen=True, slots=True)
class ChatStreamEvent:
delta: str
finish_reason: str | None
model: str
request_id: str | None
usage: ProviderUsage
elapsed_ms: float
class EmbeddingProvider(Protocol):
async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult: ...
async def embed_query(self, text: str) -> EmbeddingResult: ...
class Reranker(Protocol):
async def rerank(
self,
query: str,
documents: Sequence[str],
*,
top_n: int,
instruct: str | None = None,
) -> RerankResult: ...
class ChatProvider(Protocol):
async def complete(
self,
messages: Sequence[ChatMessage],
*,
max_tokens: int,
) -> ChatCompletionResult: ...
def stream(
self,
messages: Sequence[ChatMessage],
*,
max_tokens: int,
) -> AsyncIterator[ChatStreamEvent]: ...

View File

@@ -1 +0,0 @@

View File

@@ -0,0 +1 @@
"""Operational command entrypoints shipped in the backend image."""

View File

@@ -0,0 +1,190 @@
"""Minimal live probes for the three configured Bailian capabilities."""
from __future__ import annotations
import asyncio
import json
import sys
from collections.abc import Awaitable, Callable
from dataclasses import asdict, dataclass
from typing import Any
from app.adapters.bailian import (
BailianChatAdapter,
BailianEmbeddingAdapter,
BailianRerankerAdapter,
)
from app.core.config import Settings
from app.core.secrets import SecretFileError
from app.ports.model_providers import ChatMessage, ModelProviderError
@dataclass(frozen=True, slots=True)
class ProbeResult:
capability: str
status: str
model: str | None = None
elapsed_ms: float | None = None
request_id: str | None = None
error_kind: str | None = None
status_code: int | None = None
async def probe_embedding(settings: Settings, api_key: str) -> ProbeResult:
adapter = 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,
)
try:
result = await adapter.embed_documents(["用于能力探测的虚构地质文本。"])
if len(result.vectors) != 1 or len(result.vectors[0]) != settings.embedding_dimension:
raise RuntimeError("embedding contract mismatch")
return ProbeResult(
capability="embedding",
status="ok",
model=result.model,
elapsed_ms=round(result.elapsed_ms, 2),
request_id=result.request_id,
)
finally:
await adapter.aclose()
async def probe_rerank(settings: Settings, api_key: str) -> ProbeResult:
adapter = 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,
)
try:
result = await adapter.rerank(
"哪段文本提到了斑岩铜矿?",
["虚构斑岩铜矿具有钾化带。", "虚构煤层采用测井曲线对比。"],
top_n=1,
)
if len(result.items) != 1 or result.items[0].index not in (0, 1):
raise RuntimeError("rerank contract mismatch")
return ProbeResult(
capability="rerank",
status="ok",
model=result.model,
elapsed_ms=round(result.elapsed_ms, 2),
request_id=result.request_id,
)
finally:
await adapter.aclose()
async def probe_chat(settings: Settings, api_key: str) -> ProbeResult:
adapter = BailianChatAdapter(
api_key=api_key,
base_url=settings.bailian_openai_base_url,
model=settings.llm_model,
timeout_seconds=settings.model_timeout_seconds,
max_retries=settings.model_max_retries,
)
model: str | None = None
request_id: str | None = None
elapsed_ms = 0.0
content_seen = False
try:
async for event in adapter.stream(
[ChatMessage(role="user", content="只回复:能力正常")],
max_tokens=16,
):
model = event.model
request_id = event.request_id or request_id
elapsed_ms = max(elapsed_ms, event.elapsed_ms)
content_seen = content_seen or bool(event.delta)
if not content_seen:
raise RuntimeError("chat stream contained no text")
return ProbeResult(
capability="chat",
status="ok",
model=model,
elapsed_ms=round(elapsed_ms, 2),
request_id=request_id,
)
finally:
await adapter.aclose()
def failed_probe(capability: str, error: BaseException) -> ProbeResult:
if isinstance(error, ModelProviderError):
return ProbeResult(
capability=capability,
status="failed",
request_id=error.request_id,
error_kind=error.kind.value,
status_code=error.status_code,
)
return ProbeResult(
capability=capability,
status="failed",
error_kind="internal_contract_error",
)
async def run_probe(
capability: str,
operation: Callable[[Settings, str], Awaitable[ProbeResult]],
settings: Settings,
api_key: str,
) -> ProbeResult:
try:
return await operation(settings, api_key)
except Exception as exc: # The output is deliberately reduced to a safe category.
return failed_probe(capability, exc)
def write_json_line(payload: dict[str, Any]) -> None:
sys.stdout.write(json.dumps(payload, ensure_ascii=False, sort_keys=True) + "\n")
async def async_main() -> int:
try:
settings = Settings()
if any(
"<workspace-id>" in url
for url in (
settings.bailian_openai_base_url,
settings.bailian_rerank_base_url,
)
):
raise ValueError("workspace endpoint placeholders are not runnable")
api_key = settings.bailian_api_key()
except (SecretFileError, ValueError):
write_json_line(
{
"capability": "configuration",
"status": "failed",
"error_kind": "invalid_local_configuration",
}
)
return 2
probes = (
("embedding", probe_embedding),
("rerank", probe_rerank),
("chat", probe_chat),
)
results = []
for capability, operation in probes:
result = await run_probe(capability, operation, settings, api_key)
results.append(result)
write_json_line(asdict(result))
return 0 if all(result.status == "ok" for result in results) else 1
def main() -> None:
raise SystemExit(asyncio.run(async_main()))
if __name__ == "__main__":
main()

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