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:
3
backend/app/__init__.py
Normal file
3
backend/app/__init__.py
Normal file
@@ -0,0 +1,3 @@
|
||||
"""Geological RAG backend package."""
|
||||
|
||||
__version__ = "0.1.0"
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
11
backend/app/adapters/bailian/__init__.py
Normal file
11
backend/app/adapters/bailian/__init__.py
Normal 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",
|
||||
]
|
||||
382
backend/app/adapters/bailian/_base.py
Normal file
382
backend/app/adapters/bailian/_base.py
Normal 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()
|
||||
326
backend/app/adapters/bailian/chat.py
Normal file
326
backend/app/adapters/bailian/chat.py
Normal 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,
|
||||
)
|
||||
177
backend/app/adapters/bailian/embedding.py
Normal file
177
backend/app/adapters/bailian/embedding.py
Normal 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)
|
||||
208
backend/app/adapters/bailian/rerank.py
Normal file
208
backend/app/adapters/bailian/rerank.py
Normal 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)
|
||||
154
backend/app/adapters/fake.py
Normal file
154
backend/app/adapters/fake.py
Normal 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,
|
||||
)
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
1
backend/app/core/__init__.py
Normal file
1
backend/app/core/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Configuration and security primitives."""
|
||||
128
backend/app/core/config.py
Normal file
128
backend/app/core/config.py
Normal 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()
|
||||
24
backend/app/core/secrets.py
Normal file
24
backend/app/core/secrets.py
Normal 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
56
backend/app/main.py
Normal 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
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
17
backend/app/persistence/__init__.py
Normal file
17
backend/app/persistence/__init__.py
Normal 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",
|
||||
]
|
||||
129
backend/app/persistence/job_queue_sql.py
Normal file
129
backend/app/persistence/job_queue_sql.py
Normal 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.*
|
||||
"""
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
151
backend/app/ports/model_providers.py
Normal file
151
backend/app/ports/model_providers.py
Normal 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]: ...
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
1
backend/app/tools/__init__.py
Normal file
1
backend/app/tools/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Operational command entrypoints shipped in the backend image."""
|
||||
190
backend/app/tools/provider_smoke.py
Normal file
190
backend/app/tools/provider_smoke.py
Normal 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()
|
||||
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