yovinchen/RAG
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Isolate cloud model access before enabling product RAG workflows
Some checks failed
verify / verify (push) Has been cancelled
Details

Browse Source 
The API and ingestion tools now use a fixed internal model gateway while
governed profiles, embedding cache assignments, traceable citations, and
stable API errors establish the boundaries required by later workflows.

Constraint: The current Alibaba Cloud workspace rejects all three live model calls with authentication failures
Rejected: Give the API or seed tools the Bailian key and direct egress | combines database access, cloud credentials, and public network access
Rejected: Mix offline and Bailian vectors in one demo namespace | makes profile activation and retrieval ambiguous
Confidence: high
Scope-risk: moderate
Reversibility: clean
Directive: Keep Bailian credentials and egress exclusive to model-gateway and create a new immutable profile hash for any embedding identity change
Tested: make verify; 121 backend tests; 14 frontend tests; fresh and populated Alembic upgrade-downgrade-upgrade; two idempotent offline seeds; Docker health and HTTP retrieval; isolated provider smoke
Not-tested: Successful live Bailian responses because the supplied workspace credential currently fails authentication
This commit is contained in:
YoVinchen
2026-07-13 04:09:06 +08:00
parent 99b7df64ea
commit 75592af33a
28 changed files with 3932 additions and 254 deletions
Download Patch File Download Diff File Expand all files Collapse all files

539
backend/app/adapters/model_gateway.py Normal file
View File

@@ -0,0 +1,539 @@
"""Typed client for the fixed internal model-gateway trust boundary."""
from __future__ import annotations
import json
import math
from collections.abc import AsyncIterator, Mapping, Sequence
from typing import Any, Literal, Self
from urllib.parse import urlsplit
import httpx
from app.adapters.bailian._base import (
extract_request_id,
invalid_request,
invalid_response,
parse_usage,
response_model,
safe_identifier,
sanitized_error,
)
from app.core.config import Settings
from app.core.secrets import read_secret_file
from app.ports.model_providers import (
ChatCompletionResult,
ChatMessage,
ChatStreamEvent,
EmbeddingResult,
ModelProviderError,
ProviderErrorKind,
ProviderUsage,
RankedItem,
RerankResult,
)
_GATEWAY_HOST = "model-gateway"
_GATEWAY_PORT = 8000
_DIMENSION = 1024
_ALLOWED_ROLES = frozenset({"system", "user", "assistant"})
class ModelGatewayAdapter:
"""Expose internal gateway calls through the provider-neutral model ports."""
def __init__(
self,
*,
token: str,
caller: Literal["api", "worker"],
base_url: str = "http://model-gateway:8000",
embedding_model: str = "text-embedding-v4",
rerank_model: str = "qwen3-rerank",
chat_model: str = "deepseek-v4-flash",
http_client: httpx.AsyncClient | None = None,
timeout_seconds: float = 120.0,
) -> None:
if not token or token != token.strip():
raise invalid_request("model_gateway.configuration", "invalid_token")
if caller not in ("api", "worker"):
raise invalid_request("model_gateway.configuration", "invalid_caller")
self._base_url = self._validate_base_url(base_url)
self._token = token
self._caller = caller
self._embedding_model = embedding_model
self._rerank_model = rerank_model
self._chat_model = chat_model
self._owns_client = http_client is None
self._client = http_client or httpx.AsyncClient(
timeout=httpx.Timeout(timeout_seconds),
follow_redirects=False,
trust_env=False,
)
@classmethod
def from_settings(
cls,
settings: Settings,
*,
http_client: httpx.AsyncClient | None = None,
) -> Self:
return cls(
token=read_secret_file(settings.model_gateway_token_file),
caller=settings.model_gateway_caller,
base_url=settings.model_gateway_base_url,
embedding_model=settings.embedding_model,
rerank_model=settings.rerank_model,
chat_model=settings.llm_model,
http_client=http_client,
timeout_seconds=settings.model_gateway_timeout_seconds,
)
async def __aenter__(self) -> Self:
return self
async def __aexit__(self, *_: object) -> None:
await self.aclose()
async def aclose(self) -> None:
if self._owns_client:
await self._client.aclose()
async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult:
return await self._embed(texts, input_type="document")
async def embed_query(self, text: str) -> EmbeddingResult:
return await self._embed((text,), input_type="query")
async def _embed(
self,
texts: Sequence[str],
*,
input_type: Literal["document", "query"],
) -> EmbeddingResult:
operation = f"model_gateway.embedding.{input_type}"
validated = self._texts(texts, operation=operation, maximum=10)
if input_type == "query" and len(validated) != 1:
raise invalid_request(operation, "query_requires_one_text")
body = await self._post_json(
operation=operation,
path="embeddings",
payload={"texts": list(validated), "input_type": input_type},
)
vectors = self._vectors(body, expected=len(validated), operation=operation)
return EmbeddingResult(
vectors=vectors,
model=response_model(body, self._embedding_model, sensitive_values=validated),
request_id=extract_request_id(body, sensitive_values=validated),
usage=parse_usage(body.get("usage")),
elapsed_ms=self._elapsed(body, operation=operation),
)
async def rerank(
self,
query: str,
documents: Sequence[str],
*,
top_n: int,
instruct: str | None = None,
) -> RerankResult:
operation = "model_gateway.rerank"
if not isinstance(query, str) or not query:
raise invalid_request(operation, "invalid_query")
validated = self._texts(documents, operation=operation, maximum=500)
if (
isinstance(top_n, bool)
or not isinstance(top_n, int)
or not 1 <= top_n <= len(validated)
):
raise invalid_request(operation, "invalid_top_n")
payload: dict[str, Any] = {
"query": query,
"documents": list(validated),
"top_n": top_n,
}
if instruct is not None:
if not isinstance(instruct, str) or not instruct:
raise invalid_request(operation, "invalid_instruct")
payload["instruct"] = instruct
body = await self._post_json(operation=operation, path="rerank", payload=payload)
raw_items = body.get("items")
if not isinstance(raw_items, list) or len(raw_items) > top_n:
raise invalid_response(operation, "invalid_items")
items: list[RankedItem] = []
seen: set[int] = set()
for raw_item in raw_items:
if not isinstance(raw_item, Mapping):
raise invalid_response(operation, "invalid_item")
index = raw_item.get("index")
score = raw_item.get("relevance_score")
document = raw_item.get("document")
if (
isinstance(index, bool)
or not isinstance(index, int)
or not 0 <= index < len(validated)
or index in seen
or isinstance(score, bool)
or not isinstance(score, (int, float))
or not math.isfinite(float(score))
or document != validated[index]
):
raise invalid_response(operation, "invalid_item")
seen.add(index)
items.append(
RankedItem(index=index, relevance_score=float(score), document=validated[index])
)
return RerankResult(
items=tuple(items),
model=response_model(body, self._rerank_model, sensitive_values=(query, *validated)),
request_id=extract_request_id(body, sensitive_values=(query, *validated)),
usage=parse_usage(body.get("usage")),
elapsed_ms=self._elapsed(body, operation=operation),
)
async def complete(
self,
messages: Sequence[ChatMessage],
*,
max_tokens: int,
) -> ChatCompletionResult:
operation = "model_gateway.chat.complete"
validated = self._messages(messages, max_tokens=max_tokens, operation=operation)
body = await self._post_json(
operation=operation,
path="chat/completions",
payload={
"messages": [
{"role": message.role, "content": message.content} for message in validated
],
"max_tokens": max_tokens,
},
)
content = body.get("content")
finish_reason = body.get("finish_reason")
if not isinstance(content, str) or (
finish_reason is not None and not isinstance(finish_reason, str)
):
raise invalid_response(operation, "invalid_completion")
sensitive = tuple(message.content for message in validated)
return ChatCompletionResult(
content=content,
finish_reason=finish_reason,
model=response_model(body, self._chat_model, sensitive_values=sensitive),
request_id=extract_request_id(body, sensitive_values=sensitive),
usage=parse_usage(body.get("usage")),
elapsed_ms=self._elapsed(body, operation=operation),
)
async def stream(
self,
messages: Sequence[ChatMessage],
*,
max_tokens: int,
) -> AsyncIterator[ChatStreamEvent]:
operation = "model_gateway.chat.stream"
validated = self._messages(messages, max_tokens=max_tokens, operation=operation)
payload = {
"messages": [
{"role": message.role, "content": message.content} for message in validated
],
"max_tokens": max_tokens,
}
try:
async with self._client.stream(
"POST",
self._url("chat/stream"),
headers=self._headers(),
json=payload,
) as response:
if response.status_code >= 400:
await response.aread()
self._raise_http_error(operation=operation, response=response)
event_name: str | None = None
complete_seen = False
async for line in response.aiter_lines():
if not line:
event_name = None
continue
if line.startswith(":"):
continue
if line.startswith("event:"):
event_name = line[6:].strip()
continue
if not line.startswith("data:") or event_name is None:
raise invalid_response(operation, "invalid_sse_event")
body = self._json_object(line[5:].strip(), operation=operation)
if event_name == "error":
self._raise_stream_error(body, operation=operation)
if event_name not in {"delta", "complete"}:
raise invalid_response(operation, "unsupported_sse_event")
if complete_seen:
raise invalid_response(operation, "event_after_complete")
terminal = event_name == "complete"
if terminal:
complete_seen = True
yield self._stream_event(
body,
operation=operation,
terminal=terminal,
)
if not complete_seen:
raise invalid_response(operation, "missing_complete_event")
except ModelProviderError:
raise
except httpx.TimeoutException:
raise sanitized_error(
operation=operation,
kind=ProviderErrorKind.TIMEOUT,
provider_code="request_timeout",
retryable=True,
) from None
except httpx.HTTPError:
raise sanitized_error(
operation=operation,
kind=ProviderErrorKind.TRANSPORT,
provider_code="transport_error",
retryable=True,
) from None
def _headers(self) -> dict[str, str]:
return {
"Authorization": f"Bearer {self._token}",
"Content-Type": "application/json",
"X-RAG-Caller": self._caller,
}
def _url(self, path: str) -> str:
return f"{self._base_url}/internal/v1/{path.lstrip('/')}"
async def _post_json(
self,
*,
operation: str,
path: str,
payload: Mapping[str, Any],
) -> Mapping[str, Any]:
try:
response = await self._client.post(
self._url(path), headers=self._headers(), json=payload
)
except httpx.TimeoutException:
raise sanitized_error(
operation=operation,
kind=ProviderErrorKind.TIMEOUT,
provider_code="request_timeout",
retryable=True,
) from None
except httpx.HTTPError:
raise sanitized_error(
operation=operation,
kind=ProviderErrorKind.TRANSPORT,
provider_code="transport_error",
retryable=True,
) from None
if response.status_code >= 400:
self._raise_http_error(operation=operation, response=response)
return self._json_object(response.text, operation=operation)
def _raise_http_error(self, *, operation: str, response: httpx.Response) -> None:
status = response.status_code
if status == 400 or status == 422:
kind, retryable = ProviderErrorKind.INVALID_REQUEST, False
elif status == 401:
kind, retryable = ProviderErrorKind.AUTHENTICATION, False
elif status == 403:
kind, retryable = ProviderErrorKind.PERMISSION_DENIED, False
elif status == 404:
kind, retryable = ProviderErrorKind.NOT_FOUND, False
elif status == 408 or status == 504:
kind, retryable = ProviderErrorKind.TIMEOUT, True
elif status == 429:
kind, retryable = ProviderErrorKind.RATE_LIMITED, True
else:
kind, retryable = ProviderErrorKind.UPSTREAM, status >= 500
# The trusted gateway exposes only a fixed provider-neutral error object.
# Preserve that category for diagnostics while discarding all other body data.
request_id: str | None = None
try:
decoded = response.json()
except (ValueError, httpx.ResponseNotRead):
decoded = None
if isinstance(decoded, Mapping):
raw_error = decoded.get("error")
if isinstance(raw_error, Mapping):
try:
kind = ProviderErrorKind(str(raw_error.get("kind")))
except ValueError:
pass
retryable = raw_error.get("retryable") is True
request_id = extract_request_id(
raw_error,
sensitive_values=(self._token,),
)
raise sanitized_error(
operation=operation,
kind=kind,
status_code=status,
provider_code="model_gateway_rejected",
request_id=request_id,
retryable=retryable,
)
def _raise_stream_error(self, body: Mapping[str, Any], *, operation: str) -> None:
raw_kind = body.get("kind")
try:
kind = ProviderErrorKind(str(raw_kind))
except ValueError:
kind = ProviderErrorKind.UPSTREAM
retryable = body.get("retryable") is True
request_id = body.get("request_id")
raise sanitized_error(
operation=operation,
kind=kind,
provider_code="model_gateway_stream_error",
request_id=request_id if isinstance(request_id, str) else None,
retryable=retryable,
)
@staticmethod
def _validate_base_url(value: str) -> str:
normalized = value.rstrip("/")
parsed = urlsplit(normalized)
if (
parsed.scheme != "http"
or parsed.hostname != _GATEWAY_HOST
or parsed.port != _GATEWAY_PORT
or parsed.path not in ("", "/")
or parsed.username is not None
or parsed.password is not None
or parsed.query
or parsed.fragment
):
raise invalid_request("model_gateway.configuration", "invalid_base_url")
return normalized
@staticmethod
def _texts(
values: Sequence[str],
*,
operation: str,
maximum: int,
) -> tuple[str, ...]:
if isinstance(values, (str, bytes)) or not isinstance(values, Sequence):
raise invalid_request(operation, "invalid_text_collection")
validated = tuple(values)
if not validated or len(validated) > maximum:
raise invalid_request(operation, "invalid_text_count")
if any(not isinstance(value, str) or not value for value in validated):
raise invalid_request(operation, "invalid_text")
return validated
@staticmethod
def _messages(
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 or isinstance(max_tokens, bool) or not isinstance(max_tokens, int):
raise invalid_request(operation, "invalid_messages")
if not 1 <= max_tokens <= 8192:
raise invalid_request(operation, "invalid_max_tokens")
if any(
not isinstance(message, ChatMessage)
or message.role not in _ALLOWED_ROLES
or not message.content
for message in validated
):
raise invalid_request(operation, "invalid_message")
return validated
@staticmethod
def _json_object(value: str, *, operation: str) -> Mapping[str, Any]:
try:
body = json.loads(value)
except (TypeError, ValueError):
raise invalid_response(operation, "invalid_json") from None
if not isinstance(body, Mapping):
raise invalid_response(operation, "invalid_json_object")
return body
@staticmethod
def _elapsed(body: Mapping[str, Any], *, operation: str) -> float:
value = body.get("elapsed_ms")
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise invalid_response(operation, "invalid_elapsed_ms")
elapsed = float(value)
if not math.isfinite(elapsed) or elapsed < 0:
raise invalid_response(operation, "invalid_elapsed_ms")
return elapsed
@staticmethod
def _vectors(
body: Mapping[str, Any],
*,
expected: int,
operation: str,
) -> tuple[tuple[float, ...], ...]:
raw_vectors = body.get("vectors")
if not isinstance(raw_vectors, list) or len(raw_vectors) != expected:
raise invalid_response(operation, "invalid_embedding_count")
vectors: list[tuple[float, ...]] = []
for raw_vector in raw_vectors:
if not isinstance(raw_vector, list) or len(raw_vector) != _DIMENSION:
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")
number = float(component)
if not math.isfinite(number):
raise invalid_response(operation, "invalid_embedding_component")
vector.append(number)
if math.hypot(*vector) <= 0:
raise invalid_response(operation, "invalid_embedding_norm")
vectors.append(tuple(vector))
return tuple(vectors)
def _stream_event(
self,
body: Mapping[str, Any],
*,
operation: str,
terminal: bool,
) -> ChatStreamEvent:
delta = body.get("delta", "")
finish_reason = body.get("finish_reason")
if not isinstance(delta, str) or (
finish_reason is not None and not isinstance(finish_reason, str)
):
raise invalid_response(operation, "invalid_stream_event")
raw_model = body.get("model")
model = safe_identifier(raw_model, sensitive_values=(self._token,))
if model is None:
raise invalid_response(operation, "invalid_stream_model")
request_id = extract_request_id(body, sensitive_values=(self._token,))
if body.get("request_id") is not None and request_id is None:
raise invalid_response(operation, "invalid_stream_request_id")
if terminal and "elapsed_ms" not in body:
raise invalid_response(operation, "missing_elapsed_ms")
elapsed = body.get("elapsed_ms", 0.0)
if isinstance(elapsed, bool) or not isinstance(elapsed, (int, float)):
raise invalid_response(operation, "invalid_elapsed_ms")
normalized_elapsed = float(elapsed)
if not math.isfinite(normalized_elapsed) or normalized_elapsed < 0:
raise invalid_response(operation, "invalid_elapsed_ms")
if terminal and not isinstance(body.get("usage"), Mapping):
raise invalid_response(operation, "missing_usage")
return ChatStreamEvent(
delta=delta,
finish_reason=finish_reason,
model=model,
request_id=request_id,
usage=parse_usage(body.get("usage")) if "usage" in body else ProviderUsage(),
elapsed_ms=normalized_elapsed,
)

23
backend/app/core/config.py
View File

@@ -35,6 +35,11 @@ class Settings(BaseSettings):
upload_root: Path = Path("/data/uploads")
max_upload_mb: int = Field(default=100, ge=1, le=2048)
model_gateway_base_url: str = "http://model-gateway:8000"
model_gateway_token_file: Path = Path("/run/secrets/model_gateway_api_token")
model_gateway_caller: Literal["api", "worker"] = "api"
model_gateway_timeout_seconds: float = Field(default=120, gt=0, le=600)
bailian_openai_base_url: str = (
"https://<workspace-id>.cn-beijing.maas.aliyuncs.com/compatible-mode/v1"
)
@@ -71,6 +76,24 @@ class Settings(BaseSettings):
def normalize_base_url(cls, value: str) -> str:
return value.rstrip("/")
@field_validator("model_gateway_base_url")
@classmethod
def validate_model_gateway_base_url(cls, value: str) -> str:
normalized = value.rstrip("/")
parsed = urlsplit(normalized)
if (
parsed.scheme != "http"
or parsed.hostname != "model-gateway"
or parsed.port != 8000
or parsed.path not in ("", "/")
or parsed.username is not None
or parsed.password is not None
or parsed.query
or parsed.fragment
):
raise ValueError("MODEL_GATEWAY_BASE_URL must be the fixed internal service URL")
return normalized
@field_validator("embedding_dimension", mode="before")
@classmethod
def parse_embedding_dimension(cls, value: object) -> int:

62
backend/app/core/problems.py Normal file
View File

@@ -0,0 +1,62 @@
"""Stable RFC 9457-style problem responses for public API failures."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from fastapi import Request
from fastapi.responses import JSONResponse
PROBLEM_MEDIA_TYPE = "application/problem+json"
PROBLEM_BASE = "https://geological-rag.local/problems"
@dataclass(frozen=True, slots=True)
class ApiProblem(Exception):
"""An intentionally public and sanitized application failure."""
status: int
code: str
title: str
detail: str
def __post_init__(self) -> None:
Exception.__init__(self, self.code)
def problem_payload(
*,
status: int,
code: str,
title: str,
detail: str,
trace_id: str,
field_errors: list[dict[str, Any]] | None = None,
) -> dict[str, Any]:
"""Build the only public error envelope used by formal product routes."""
return {
"type": f"{PROBLEM_BASE}/{code.lower().replace('_', '-')}",
"title": title,
"status": status,
"code": code,
"detail": detail,
"trace_id": trace_id,
"field_errors": field_errors or [],
}
def api_problem_handler(request: Request, exc: ApiProblem) -> JSONResponse:
trace_id = str(getattr(request.state, "trace_id", "unavailable"))
return JSONResponse(
status_code=exc.status,
media_type=PROBLEM_MEDIA_TYPE,
content=problem_payload(
status=exc.status,
code=exc.code,
title=exc.title,
detail=exc.detail,
trace_id=trace_id,
),
)

30
backend/app/core/request_context.py Normal file
View File

@@ -0,0 +1,30 @@
"""Per-request trace context with a bounded, non-secret public identifier."""
from __future__ import annotations
import uuid
from collections.abc import Awaitable, Callable
from fastapi import Request, Response
REQUEST_ID_HEADER = "x-request-id"
type CallNext = Callable[[Request], Awaitable[Response]]
def _request_id(value: str | None) -> str:
if value is not None:
try:
return str(uuid.UUID(value))
except (ValueError, AttributeError):
pass
return str(uuid.uuid4())
async def trace_request(request: Request, call_next: CallNext) -> Response:
"""Attach a UUID trace ID and return it without trusting arbitrary input."""
trace_id = _request_id(request.headers.get(REQUEST_ID_HEADER))
request.state.trace_id = trace_id
response = await call_next(request)
response.headers[REQUEST_ID_HEADER] = trace_id
return response

74
backend/app/main.py
View File

@@ -1,4 +1,4 @@
"""FastAPI entrypoint with dependency-free liveness and database readiness probes."""
"""FastAPI application factory and production entrypoint."""
from typing import Any
@@ -9,26 +9,17 @@ from fastapi import FastAPI, Response, status
from app import __version__
from app.api.v1 import demo_router
from app.core.config import get_settings
from app.core.problems import ApiProblem, api_problem_handler
from app.core.request_context import trace_request
from app.core.secrets import SecretFileError
app = FastAPI(title="Geological RAG API", version=__version__)
app.include_router(demo_router)
type HealthPayload = dict[str, str | dict[str, str]]
@app.get("/health/live", tags=["health"])
@app.get("/api/v1/health/live", tags=["health"])
def live() -> dict[str, str]:
return {"status": "ok", "version": __version__}
@app.get(
"/health/ready",
tags=["health"],
responses={status.HTTP_503_SERVICE_UNAVAILABLE: {"description": "Database unavailable"}},
)
@app.get("/api/v1/health/ready", tags=["health"])
def ready(response: Response) -> HealthPayload:
settings = get_settings()
try:
@@ -48,7 +39,6 @@ def ready(response: Response) -> HealthPayload:
return {"status": "ready", "checks": {"database": "ok"}}
@app.get("/api/v1/meta", tags=["meta"])
def meta() -> dict[str, Any]:
settings = get_settings()
return {
@@ -63,6 +53,64 @@ def meta() -> dict[str, Any]:
}
def create_app() -> FastAPI:
"""Create the API without opening a database or loading model credentials."""
api = FastAPI(
title="Geological RAG API",
version=__version__,
openapi_tags=[
{"name": "health", "description": "Process and database health probes."},
{"name": "meta", "description": "Safe runtime capability metadata."},
{"name": "offline-demo", "description": "Synthetic offline validation only."},
],
)
api.middleware("http")(trace_request)
api.add_exception_handler(ApiProblem, api_problem_handler) # type: ignore[arg-type]
api.include_router(demo_router)
api.add_api_route(
"/health/live",
live,
methods=["GET"],
tags=["health"],
include_in_schema=False,
)
api.add_api_route(
"/api/v1/health/live",
live,
methods=["GET"],
tags=["health"],
operation_id="getLiveness",
)
api.add_api_route(
"/health/ready",
ready,
methods=["GET"],
tags=["health"],
include_in_schema=False,
responses={status.HTTP_503_SERVICE_UNAVAILABLE: {"description": "Database unavailable"}},
)
api.add_api_route(
"/api/v1/health/ready",
ready,
methods=["GET"],
tags=["health"],
operation_id="getReadiness",
responses={status.HTTP_503_SERVICE_UNAVAILABLE: {"description": "Database unavailable"}},
)
api.add_api_route(
"/api/v1/meta",
meta,
methods=["GET"],
tags=["meta"],
operation_id="getRuntimeMetadata",
)
return api
app = create_app()
if __name__ == "__main__":
# Compose publishes this listener only on the host loopback interface.
uvicorn.run("app.main:app", host="0.0.0.0", port=8000) # noqa: S104

755
backend/app/model_gateway.py Normal file
View File

@@ -0,0 +1,755 @@
"""Credential-isolated internal gateway for all cloud model capabilities."""
import asyncio
import json
import os
import secrets
from collections.abc import AsyncIterator, Callable, Mapping
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated, Literal, Protocol, Self, runtime_checkable
from fastapi import Depends, FastAPI, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel, ConfigDict, Field, model_validator
from starlette.types import Receive, Scope, Send
from app import __version__
from app.adapters.bailian import (
BailianChatAdapter,
BailianEmbeddingAdapter,
BailianRerankerAdapter,
)
from app.core.config import Settings
from app.core.secrets import SecretFileError, read_secret_file
from app.ports.model_providers import (
ChatMessage,
ChatProvider,
ChatStreamEvent,
EmbeddingProvider,
ModelProviderError,
ProviderErrorKind,
ProviderUsage,
Reranker,
)
Caller = Literal["api", "worker"]
InputType = Literal["document", "query"]
Role = Literal["system", "user", "assistant"]
SettingsFactory = Callable[[], Settings]
CALLERS: tuple[Caller, Caller] = ("api", "worker")
DEFAULT_ALLOWED_TOKEN_FILES = (
"/run/secrets/model_gateway_api_token,/run/secrets/model_gateway_worker_token" # noqa: S105
)
MAX_CHAT_MESSAGES = 100
MAX_CHAT_CONTENT_CHARS = 100_000
MAX_CHAT_OUTPUT_TOKENS = 8_192
ALLOWED_FINISH_REASONS = frozenset(
{"stop", "length", "content_filter", "tool_calls", "function_call"}
)
class _StrictModel(BaseModel):
model_config = ConfigDict(extra="forbid")
class EmbeddingRequest(_StrictModel):
texts: list[Annotated[str, Field(min_length=1, max_length=8_192)]] = Field(
min_length=1,
max_length=10,
)
input_type: InputType
@model_validator(mode="after")
def validate_query_count(self) -> Self:
if self.input_type == "query" and len(self.texts) != 1:
raise ValueError("query embedding accepts exactly one text")
return self
class RerankRequest(_StrictModel):
query: str = Field(min_length=1, max_length=4_000)
documents: list[Annotated[str, Field(min_length=1, max_length=4_000)]] = Field(
min_length=1,
max_length=500,
)
top_n: int = Field(ge=1, le=500)
instruct: str | None = Field(default=None, min_length=1, max_length=4_000)
@model_validator(mode="after")
def validate_top_n(self) -> Self:
if self.top_n > len(self.documents):
raise ValueError("top_n must not exceed document count")
return self
class ChatMessageRequest(_StrictModel):
role: Role
content: str = Field(min_length=1, max_length=MAX_CHAT_CONTENT_CHARS)
class ChatRequest(_StrictModel):
messages: list[ChatMessageRequest] = Field(min_length=1, max_length=MAX_CHAT_MESSAGES)
max_tokens: int = Field(default=1_024, ge=1, le=MAX_CHAT_OUTPUT_TOKENS)
class UsageResponse(_StrictModel):
input_tokens: int | None
output_tokens: int | None
total_tokens: int | None
class EmbeddingResponse(_StrictModel):
vectors: list[list[float]]
model: str
request_id: str | None
usage: UsageResponse
elapsed_ms: float
class RankedItemResponse(_StrictModel):
index: int
relevance_score: float
document: str
class RerankResponse(_StrictModel):
items: list[RankedItemResponse]
model: str
request_id: str | None
usage: UsageResponse
elapsed_ms: float
class ChatResponse(_StrictModel):
content: str
finish_reason: str | None
model: str
request_id: str | None
usage: UsageResponse
elapsed_ms: float
class _UnauthorizedError(RuntimeError):
pass
class _ForbiddenError(RuntimeError):
pass
class _UnavailableError(RuntimeError):
pass
class _RestartRequiredError(RuntimeError):
pass
@runtime_checkable
class _SupportsAclose(Protocol):
async def aclose(self) -> None: ...
class _ClosingStreamingResponse(StreamingResponse):
"""Close the response iterator on completion, cancellation, or send failure."""
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
try:
await super().__call__(scope, receive, send)
finally:
await _close_stream(self.body_iterator)
class _Runtime:
def __init__(self) -> None:
self.embedding: EmbeddingProvider | None = None
self.reranker: Reranker | None = None
self.chat: ChatProvider | None = None
self.semaphore: asyncio.Semaphore | None = None
self.allowed_tokens: dict[Caller, str] = {}
self.local_configuration_check: Callable[[], bool] = lambda: False
self.restart_required = False
@property
def available(self) -> bool:
return all(
provider is not None
for provider in (self.embedding, self.reranker, self.chat, self.semaphore)
) and set(self.allowed_tokens) == {"api", "worker"}
def invalidate_for_restart(self) -> None:
"""Atomically stop accepting work after mounted configuration changes."""
self.restart_required = True
self.embedding = None
self.reranker = None
self.chat = None
self.semaphore = None
self.allowed_tokens = {}
def _usage_response(usage: ProviderUsage) -> UsageResponse:
return UsageResponse(
input_tokens=usage.input_tokens,
output_tokens=usage.output_tokens,
total_tokens=usage.total_tokens,
)
def _merge_usage(current: ProviderUsage, update: ProviderUsage) -> ProviderUsage:
return ProviderUsage(
input_tokens=(
update.input_tokens if update.input_tokens is not None else current.input_tokens
),
output_tokens=(
update.output_tokens if update.output_tokens is not None else current.output_tokens
),
total_tokens=(
update.total_tokens if update.total_tokens is not None else current.total_tokens
),
)
def _boundary_error(operation: str) -> ModelProviderError:
return ModelProviderError(
operation=operation,
kind=ProviderErrorKind.INVALID_RESPONSE,
provider_code="gateway_boundary_error",
)
def _normalize_allowed_tokens(allowed_tokens: Mapping[str, str]) -> dict[Caller, str]:
if set(allowed_tokens) != {"api", "worker"}:
raise ValueError("allowed_tokens must define api and worker identities")
normalized: dict[Caller, str] = {}
for caller in CALLERS:
token = allowed_tokens.get(caller)
if (
not isinstance(token, str)
or not token
or token != token.strip()
or "\n" in token
or "\r" in token
or len(token) > 4_096
):
raise ValueError("allowed token is invalid")
normalized[caller] = token
if secrets.compare_digest(normalized["api"], normalized["worker"]):
raise ValueError("api and worker tokens must be different")
return normalized
def _load_allowed_tokens_from_files() -> dict[Caller, str]:
raw = os.environ.get("MODEL_GATEWAY_ALLOWED_TOKEN_FILES", DEFAULT_ALLOWED_TOKEN_FILES)
entries = raw.split(",")
if len(entries) != 2 or any(not entry.strip() for entry in entries):
raise ValueError("MODEL_GATEWAY_ALLOWED_TOKEN_FILES must define two files")
paths: dict[str, str] = {}
if all("=" not in entry for entry in entries):
paths = {caller: entry.strip() for caller, entry in zip(CALLERS, entries, strict=True)}
else:
for entry in entries:
caller, separator, path = entry.partition("=")
if not separator or caller.strip() in paths:
raise ValueError("invalid model gateway token file mapping")
paths[caller.strip()] = path.strip()
if set(paths) != {"api", "worker"}:
raise ValueError("model gateway token files must map api and worker")
return _normalize_allowed_tokens(
{
"api": read_secret_file(Path(paths["api"])),
"worker": read_secret_file(Path(paths["worker"])),
}
)
def _provider_status(error: ModelProviderError) -> int:
return {
ProviderErrorKind.INVALID_REQUEST: status.HTTP_422_UNPROCESSABLE_CONTENT,
ProviderErrorKind.RATE_LIMITED: status.HTTP_429_TOO_MANY_REQUESTS,
ProviderErrorKind.TIMEOUT: status.HTTP_504_GATEWAY_TIMEOUT,
ProviderErrorKind.TRANSPORT: status.HTTP_503_SERVICE_UNAVAILABLE,
}.get(error.kind, status.HTTP_502_BAD_GATEWAY)
def _error_payload(
kind: str,
*,
retryable: bool = False,
request_id: str | None = None,
) -> dict[str, dict[str, str | bool | None]]:
return {
"error": {
"kind": kind,
"retryable": retryable,
"request_id": request_id,
}
}
def _sse(event: str, payload: Mapping[str, object]) -> bytes:
data = json.dumps(payload, ensure_ascii=False, separators=(",", ":"), sort_keys=True)
return f"event: {event}\ndata: {data}\n\n".encode()
async def _close_stream(stream: object | None) -> None:
if isinstance(stream, _SupportsAclose):
try:
await stream.aclose()
except Exception:
# Closing is best-effort and must never expose provider exception text.
return
def create_model_gateway_app(
*,
embedding_provider: EmbeddingProvider | None = None,
reranker: Reranker | None = None,
chat_provider: ChatProvider | None = None,
allowed_tokens: Mapping[str, str] | None = None,
max_concurrency: int | None = None,
settings_factory: SettingsFactory = Settings,
) -> FastAPI:
"""Create the internal model gateway with injectable hermetic providers."""
injected = (embedding_provider, reranker, chat_provider)
if any(provider is not None for provider in injected) != all(
provider is not None for provider in injected
):
raise ValueError("all three providers must be injected together")
providers_are_injected = all(provider is not None for provider in injected)
if providers_are_injected != (allowed_tokens is not None):
raise ValueError("injected providers and allowed_tokens must be supplied together")
if max_concurrency is not None and (
isinstance(max_concurrency, bool) or max_concurrency < 1 or max_concurrency > 100
):
raise ValueError("max_concurrency must be between 1 and 100")
runtime = _Runtime()
owned_adapters: list[BailianEmbeddingAdapter | BailianRerankerAdapter | BailianChatAdapter] = []
@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
if providers_are_injected:
assert embedding_provider is not None
assert reranker is not None
assert chat_provider is not None
assert allowed_tokens is not None
runtime.embedding = embedding_provider
runtime.reranker = reranker
runtime.chat = chat_provider
runtime.allowed_tokens = _normalize_allowed_tokens(allowed_tokens)
runtime.semaphore = asyncio.Semaphore(max_concurrency or 4)
runtime.local_configuration_check = lambda: True
runtime.restart_required = False
else:
try:
settings = settings_factory()
api_key = settings.bailian_api_key()
loaded_tokens = _load_allowed_tokens_from_files()
embedding_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,
)
owned_adapters.append(embedding_adapter)
rerank_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,
)
owned_adapters.append(rerank_adapter)
chat_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,
)
owned_adapters.append(chat_adapter)
runtime.embedding = embedding_adapter
runtime.reranker = rerank_adapter
runtime.chat = chat_adapter
runtime.allowed_tokens = loaded_tokens
runtime.semaphore = asyncio.Semaphore(
max_concurrency or settings.model_max_concurrency
)
runtime.restart_required = False
def check_local_configuration() -> bool:
try:
current_settings = settings_factory()
current_api_key = current_settings.bailian_api_key()
current_tokens = _load_allowed_tokens_from_files()
same_key = secrets.compare_digest(current_api_key, api_key)
same_tokens = all(
secrets.compare_digest(current_tokens[caller], loaded_tokens[caller])
for caller in CALLERS
)
same_provider_contract = (
current_settings.bailian_openai_base_url
== settings.bailian_openai_base_url
and current_settings.bailian_rerank_base_url
== settings.bailian_rerank_base_url
and current_settings.embedding_model == settings.embedding_model
and current_settings.embedding_dimension == settings.embedding_dimension
and current_settings.rerank_model == settings.rerank_model
and current_settings.llm_model == settings.llm_model
)
return same_key and same_tokens and same_provider_contract
except (OSError, SecretFileError, ValueError, ModelProviderError):
return False
runtime.local_configuration_check = check_local_configuration
except (OSError, SecretFileError, ValueError, ModelProviderError):
runtime.local_configuration_check = lambda: False
try:
yield
finally:
runtime.embedding = None
runtime.reranker = None
runtime.chat = None
runtime.semaphore = None
runtime.allowed_tokens = {}
runtime.local_configuration_check = lambda: False
runtime.restart_required = False
for adapter in reversed(owned_adapters):
await adapter.aclose()
owned_adapters.clear()
gateway = FastAPI(
title="Geological RAG Model Gateway",
version=__version__,
docs_url=None,
redoc_url=None,
openapi_url=None,
lifespan=lifespan,
)
@gateway.exception_handler(RequestValidationError)
async def request_validation_error(
_: Request,
__: RequestValidationError,
) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_422_UNPROCESSABLE_CONTENT,
content=_error_payload("invalid_request"),
)
@gateway.exception_handler(ModelProviderError)
async def model_provider_error(_: Request, error: ModelProviderError) -> JSONResponse:
return JSONResponse(
status_code=_provider_status(error),
content=_error_payload(
error.kind.value,
retryable=error.retryable,
request_id=error.request_id,
),
)
@gateway.exception_handler(_UnauthorizedError)
async def unauthorized_error(_: Request, __: _UnauthorizedError) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_401_UNAUTHORIZED,
content=_error_payload("unauthorized"),
headers={"WWW-Authenticate": "Bearer"},
)
@gateway.exception_handler(_ForbiddenError)
async def forbidden_error(_: Request, __: _ForbiddenError) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_403_FORBIDDEN,
content=_error_payload("forbidden"),
)
@gateway.exception_handler(_UnavailableError)
async def unavailable_error(_: Request, __: _UnavailableError) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
content=_error_payload("unavailable", retryable=True),
)
@gateway.exception_handler(_RestartRequiredError)
async def restart_required_error(_: Request, __: _RestartRequiredError) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
content=_error_payload("restart_required"),
)
@gateway.exception_handler(Exception)
async def unexpected_error(_: Request, __: Exception) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
content=_error_payload("internal_error"),
)
def ensure_current_configuration() -> None:
if runtime.restart_required:
raise _RestartRequiredError
if not runtime.available:
raise _UnavailableError
if not runtime.local_configuration_check():
runtime.invalidate_for_restart()
raise _RestartRequiredError
def require_runtime() -> tuple[
EmbeddingProvider,
Reranker,
ChatProvider,
asyncio.Semaphore,
]:
ensure_current_configuration()
assert runtime.embedding is not None
assert runtime.reranker is not None
assert runtime.chat is not None
assert runtime.semaphore is not None
return runtime.embedding, runtime.reranker, runtime.chat, runtime.semaphore
async def authorize(request: Request) -> Caller:
ensure_current_configuration()
authorization = request.headers.get("authorization", "")
scheme, separator, credential = authorization.partition(" ")
caller_value = request.headers.get("x-rag-caller", "")
if (
not separator
or scheme.lower() != "bearer"
or not credential
or len(credential) > 4_096
or caller_value not in {"api", "worker"}
):
raise _UnauthorizedError
matched_identity: Caller | None = None
for identity, allowed_token in runtime.allowed_tokens.items():
if secrets.compare_digest(credential, allowed_token):
matched_identity = identity
if matched_identity is None or matched_identity != caller_value:
raise _UnauthorizedError
return matched_identity
@gateway.get("/health/live", include_in_schema=False)
async def live() -> dict[str, str]:
return {"status": "ok", "version": __version__}
@gateway.get("/health/ready", include_in_schema=False)
async def ready() -> JSONResponse:
if runtime.restart_required:
configuration_status = "restart_required"
elif runtime.available and runtime.local_configuration_check():
return JSONResponse(
status_code=status.HTTP_200_OK,
content={"status": "ready", "checks": {"configuration": "ok"}},
)
elif runtime.available:
runtime.invalidate_for_restart()
configuration_status = "restart_required"
else:
configuration_status = "unavailable"
return JSONResponse(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
content={
"status": "not_ready",
"checks": {"configuration": configuration_status},
},
)
@gateway.post("/internal/v1/embeddings", response_model=EmbeddingResponse)
async def embeddings(
payload: EmbeddingRequest,
caller: Annotated[Caller, Depends(authorize)],
) -> EmbeddingResponse:
embedding, _, _, semaphore = require_runtime()
if payload.input_type == "document" and caller != "worker":
raise _ForbiddenError
try:
async with semaphore:
if payload.input_type == "query":
result = await embedding.embed_query(payload.texts[0])
else:
result = await embedding.embed_documents(payload.texts)
return EmbeddingResponse(
vectors=[list(vector) for vector in result.vectors],
model=result.model,
request_id=result.request_id,
usage=_usage_response(result.usage),
elapsed_ms=result.elapsed_ms,
)
except ModelProviderError:
raise
except Exception:
raise _boundary_error("model_gateway.embedding") from None
@gateway.post("/internal/v1/rerank", response_model=RerankResponse)
async def rerank(
payload: RerankRequest,
_: Annotated[Caller, Depends(authorize)],
) -> RerankResponse:
embedding_provider_unused, reranker_provider, chat_provider_unused, semaphore = (
require_runtime()
)
del embedding_provider_unused, chat_provider_unused
try:
async with semaphore:
result = await reranker_provider.rerank(
payload.query,
payload.documents,
top_n=payload.top_n,
instruct=payload.instruct,
)
return RerankResponse(
items=[
RankedItemResponse(
index=item.index,
relevance_score=item.relevance_score,
document=item.document,
)
for item in result.items
],
model=result.model,
request_id=result.request_id,
usage=_usage_response(result.usage),
elapsed_ms=result.elapsed_ms,
)
except ModelProviderError:
raise
except Exception:
raise _boundary_error("model_gateway.rerank") from None
def chat_messages(payload: ChatRequest) -> list[ChatMessage]:
return [
ChatMessage(role=message.role, content=message.content) for message in payload.messages
]
@gateway.post("/internal/v1/chat/completions", response_model=ChatResponse)
async def chat_completion(
payload: ChatRequest,
_: Annotated[Caller, Depends(authorize)],
) -> ChatResponse:
embedding_provider_unused, reranker_unused, chat, semaphore = require_runtime()
del embedding_provider_unused, reranker_unused
try:
async with semaphore:
result = await chat.complete(chat_messages(payload), max_tokens=payload.max_tokens)
return ChatResponse(
content=result.content,
finish_reason=result.finish_reason,
model=result.model,
request_id=result.request_id,
usage=_usage_response(result.usage),
elapsed_ms=result.elapsed_ms,
)
except ModelProviderError:
raise
except Exception:
raise _boundary_error("model_gateway.chat_complete") from None
@gateway.post("/internal/v1/chat/stream")
async def chat_stream(
payload: ChatRequest,
_: Annotated[Caller, Depends(authorize)],
) -> StreamingResponse:
embedding_provider_unused, reranker_unused, chat, semaphore = require_runtime()
del embedding_provider_unused, reranker_unused
async def event_stream() -> AsyncIterator[bytes]:
events: AsyncIterator[ChatStreamEvent] | None = None
finish_reason: str | None = None
model: str | None = None
request_id: str | None = None
usage = ProviderUsage()
elapsed_ms = 0.0
try:
events = chat.stream(chat_messages(payload), max_tokens=payload.max_tokens)
async with semaphore:
async for event in events:
if event.finish_reason is not None:
if event.finish_reason not in ALLOWED_FINISH_REASONS:
raise ModelProviderError(
operation="chat.stream",
kind=ProviderErrorKind.INVALID_RESPONSE,
provider_code="invalid_finish_reason",
)
finish_reason = event.finish_reason
if event.model:
model = event.model
if event.request_id:
request_id = event.request_id
usage = _merge_usage(usage, event.usage)
elapsed_ms = max(elapsed_ms, event.elapsed_ms)
if event.delta:
yield _sse(
"delta",
{
"delta": event.delta,
"finish_reason": (
event.finish_reason
if event.finish_reason in ALLOWED_FINISH_REASONS
else None
),
"model": event.model or model,
"request_id": event.request_id or request_id,
},
)
if finish_reason is None or model is None:
raise ModelProviderError(
operation="chat.stream",
kind=ProviderErrorKind.INVALID_RESPONSE,
provider_code="missing_terminal_event",
)
yield _sse(
"complete",
{
"finish_reason": finish_reason,
"model": model,
"request_id": request_id,
"usage": _usage_response(usage).model_dump(),
"elapsed_ms": elapsed_ms,
},
)
except ModelProviderError as error:
yield _sse(
"error",
{
"kind": error.kind.value,
"retryable": error.retryable,
"request_id": error.request_id,
},
)
except Exception:
boundary_failure = _boundary_error("model_gateway.chat_stream")
yield _sse(
"error",
{
"kind": boundary_failure.kind.value,
"retryable": boundary_failure.retryable,
"request_id": boundary_failure.request_id,
},
)
finally:
await _close_stream(events)
return _ClosingStreamingResponse(
event_stream(),
media_type="text/event-stream",
headers={
"Cache-Control": "no-store",
"X-Accel-Buffering": "no",
},
)
return gateway
app = create_model_gateway_app()

156
backend/app/tools/provider_smoke.py
View File

@@ -9,11 +9,7 @@ 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.adapters.model_gateway import ModelGatewayAdapter
from app.core.config import Settings
from app.core.secrets import SecretFileError
from app.ports.model_providers import ChatMessage, ModelProviderError
@@ -30,89 +26,59 @@ class ProbeResult:
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,
async def probe_embedding(settings: Settings, adapter: ModelGatewayAdapter) -> ProbeResult:
# API identity probes query embedding. Document embedding remains worker-only.
result = await adapter.embed_query("用于能力探测的虚构地质问题。")
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,
)
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,
async def probe_rerank(_: Settings, adapter: ModelGatewayAdapter) -> ProbeResult:
result = await adapter.rerank(
"哪段文本提到了斑岩铜矿？",
["虚构斑岩铜矿具有钾化带。", "虚构煤层采用测井曲线对比。"],
top_n=1,
)
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,
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,
)
async def probe_chat(_: Settings, adapter: ModelGatewayAdapter) -> ProbeResult:
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()
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,
)
def failed_probe(capability: str, error: BaseException) -> ProbeResult:
@@ -133,12 +99,12 @@ def failed_probe(capability: str, error: BaseException) -> ProbeResult:
async def run_probe(
capability: str,
operation: Callable[[Settings, str], Awaitable[ProbeResult]],
operation: Callable[[Settings, ModelGatewayAdapter], Awaitable[ProbeResult]],
settings: Settings,
api_key: str,
adapter: ModelGatewayAdapter,
) -> ProbeResult:
try:
return await operation(settings, api_key)
return await operation(settings, adapter)
except Exception as exc: # The output is deliberately reduced to a safe category.
return failed_probe(capability, exc)
@@ -148,17 +114,10 @@ def write_json_line(payload: dict[str, Any]) -> None:
async def async_main() -> int:
adapter: ModelGatewayAdapter | None = None
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()
adapter = ModelGatewayAdapter.from_settings(settings)
except (SecretFileError, ValueError):
write_json_line(
{
@@ -174,12 +133,15 @@ async def async_main() -> int:
("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
try:
results = []
for capability, operation in probes:
result = await run_probe(capability, operation, settings, adapter)
results.append(result)
write_json_line(asdict(result))
return 0 if all(result.status == "ok" for result in results) else 1
finally:
await adapter.aclose()
def main() -> None:

307
backend/app/tools/seed_demo.py
View File

@@ -20,8 +20,8 @@ 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.adapters.model_gateway import ModelGatewayAdapter
from app.core.config import Settings
from app.core.demo_identity import (
ACCESS_SCOPE_ID,
@@ -55,6 +55,42 @@ class DemoQuery:
answerable: bool
@dataclass(frozen=True, slots=True)
class DemoNamespace:
mode: str
knowledge_base_id: uuid.UUID
access_scope_id: uuid.UUID
scope_name: str
knowledge_base_name: str
storage_prefix: str
OFFLINE_NAMESPACE = DemoNamespace(
mode="fake",
knowledge_base_id=KNOWLEDGE_BASE_ID,
access_scope_id=ACCESS_SCOPE_ID,
scope_name="synthetic-demo",
knowledge_base_name="虚构地质 PoC 知识库（离线）",
storage_prefix="synthetic/offline",
)
BAILIAN_NAMESPACE = DemoNamespace(
mode="bailian",
knowledge_base_id=uuid.uuid5(IDENTITY_NAMESPACE, "synthetic-bailian-knowledge-base"),
access_scope_id=uuid.uuid5(IDENTITY_NAMESPACE, "synthetic-bailian-public-scope"),
scope_name="synthetic-bailian-validation",
knowledge_base_name="虚构地质 PoC 知识库（百炼验证）",
storage_prefix="synthetic/bailian",
)
@dataclass(frozen=True, slots=True)
class EmbeddedVector:
vector: tuple[float, ...]
request_id: str | None
usage: dict[str, int | None]
elapsed_ms: int
@dataclass(frozen=True, slots=True)
class PreparedChunk:
source_id: str
@@ -71,6 +107,9 @@ class PreparedChunk:
embedding_profile_hash: str
vector: tuple[float, ...]
embedding_model: str
provider_request_id: str | None
embedding_usage: dict[str, int | None]
embedding_elapsed_ms: int
title: str
region: str
mineral: str
@@ -88,6 +127,14 @@ def sha256_text(value: str) -> str:
return hashlib.sha256(value.encode("utf-8")).hexdigest()
def namespace_for_mode(mode: str) -> DemoNamespace:
if mode == "fake":
return OFFLINE_NAMESPACE
if mode == "bailian":
return BAILIAN_NAMESPACE
raise SeedContractError("invalid_provider_mode")
def load_jsonl(path: Path) -> list[dict[str, Any]]:
if not path.is_file():
raise SeedContractError("fixture_missing")
@@ -142,8 +189,10 @@ def load_queries(path: Path) -> list[DemoQuery]:
def embedding_profile_hash(settings: Settings, mode: str) -> str:
if mode != "bailian":
if mode == "fake":
return offline_embedding_profile_hash(settings.embedding_dimension)
if mode != "bailian":
raise SeedContractError("invalid_provider_mode")
endpoint_identity = sha256_text(urlsplit(settings.bailian_openai_base_url).hostname or "")
model = settings.embedding_model
@@ -164,15 +213,30 @@ def embedding_profile_hash(settings: Settings, mode: str) -> str:
async def embed_in_batches(
provider: EmbeddingProvider,
texts: Sequence[str],
) -> tuple[tuple[tuple[float, ...], ...], str]:
vectors: list[tuple[float, ...]] = []
) -> tuple[tuple[EmbeddedVector, ...], str]:
vectors: list[EmbeddedVector] = []
resolved_model: str | None = None
for offset in range(0, len(texts), 10):
result = await provider.embed_documents(texts[offset : offset + 10])
if resolved_model is not None and result.model != resolved_model:
raise SeedContractError("embedding_model_changed_between_batches")
resolved_model = result.model
vectors.extend(result.vectors)
if len(result.vectors) != len(texts[offset : offset + 10]):
raise SeedContractError("embedding_batch_count_mismatch")
usage = {
"input_tokens": result.usage.input_tokens,
"output_tokens": result.usage.output_tokens,
"total_tokens": result.usage.total_tokens,
}
vectors.extend(
EmbeddedVector(
vector=vector,
request_id=result.request_id,
usage=usage,
elapsed_ms=max(0, round(result.elapsed_ms)),
)
for vector in result.vectors
)
if len(vectors) != len(texts) or resolved_model is None:
raise SeedContractError("embedding_result_count_mismatch")
return tuple(vectors), resolved_model
@@ -180,10 +244,11 @@ async def embed_in_batches(
def prepare_chunks(
documents: Sequence[DemoDocument],
vectors: Sequence[tuple[float, ...]],
vectors: Sequence[EmbeddedVector],
*,
profile_hash: str,
embedding_model: str,
namespace: DemoNamespace = OFFLINE_NAMESPACE,
) -> list[PreparedChunk]:
prepared = []
for document, vector in zip(documents, vectors, strict=True):
@@ -200,7 +265,12 @@ def prepare_chunks(
separators=(",", ":"),
)
raw_hash = sha256_text(raw_payload)
document_id = uuid.uuid5(IDENTITY_NAMESPACE, f"document:{document.source_id}")
document_identity = (
f"document:{document.source_id}"
if namespace.mode == "fake"
else f"document:{namespace.mode}:{document.source_id}"
)
document_id = uuid.uuid5(IDENTITY_NAMESPACE, document_identity)
version_id = uuid.uuid5(
IDENTITY_NAMESPACE,
f"version:{document.source_id}:{raw_hash}:{profile_hash}",
@@ -237,8 +307,11 @@ def prepare_chunks(
embedding_text_sha256=embedding_hash,
outbound_manifest_sha256=sha256_text(manifest_payload),
embedding_profile_hash=profile_hash,
vector=vector,
vector=vector.vector,
embedding_model=embedding_model,
provider_request_id=vector.request_id,
embedding_usage=vector.usage,
embedding_elapsed_ms=vector.elapsed_ms,
title=document.title,
region=document.region,
mineral=document.mineral,
@@ -255,20 +328,96 @@ def database_dsn(settings: Settings) -> str:
)
def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[str, int]:
def write_chunks(
settings: Settings,
chunks: Sequence[PreparedChunk],
*,
namespace: DemoNamespace,
) -> dict[str, int]:
if not chunks:
raise SeedContractError("chunks_empty")
profile_hashes = {item.embedding_profile_hash for item in chunks}
resolved_models = {item.embedding_model for item in chunks}
if len(profile_hashes) != 1 or len(resolved_models) != 1:
raise SeedContractError("mixed_embedding_profiles")
profile_hash = next(iter(profile_hashes))
resolved_model = next(iter(resolved_models))
if namespace.mode == "fake":
provider = "local-synthetic"
api_mode = "deterministic-offline"
endpoint_identity_hash = sha256_text("local-fake")
else:
provider = "aliyun-bailian"
api_mode = "model-gateway/openai-compatible"
endpoint_identity_hash = sha256_text(
urlsplit(settings.bailian_openai_base_url).hostname or ""
)
with psycopg.connect(database_dsn(settings), row_factory=dict_row) as connection:
register_vector(connection)
connection.execute("SELECT pg_advisory_xact_lock(724202607120001)")
connection.execute(
"""
INSERT INTO rag.knowledge_bases (id, name, description)
VALUES (%s, %s, %s)
INSERT INTO rag.model_profiles (
profile_hash, alias, kind, provider, model, api_mode, dimension,
endpoint_identity_hash, config_snapshot, synthetic, enabled
) VALUES (
%s, %s, 'embedding', %s, %s, %s, 1024, %s, %s, %s, true
)
ON CONFLICT (profile_hash) DO NOTHING
""",
(
profile_hash,
f"{namespace.mode}-embedding-{profile_hash[:12]}",
provider,
resolved_model,
api_mode,
endpoint_identity_hash,
Jsonb(
{
"dimension": settings.embedding_dimension,
"requested_model": settings.embedding_model,
"source": "synthetic-seed-v1",
}
),
namespace.mode == "fake",
),
)
registered_profile = connection.execute(
"""
SELECT kind, provider, model, api_mode, dimension, endpoint_identity_hash
FROM rag.model_profiles
WHERE profile_hash = %s
""",
(profile_hash,),
).fetchone()
if registered_profile is None or (
registered_profile["kind"] != "embedding"
or registered_profile["provider"] != provider
or registered_profile["model"] != resolved_model
or registered_profile["api_mode"] != api_mode
or registered_profile["dimension"] != settings.embedding_dimension
or registered_profile["endpoint_identity_hash"] != endpoint_identity_hash
):
raise SeedContractError("embedding_profile_collision")
connection.execute(
"""
INSERT INTO rag.knowledge_bases (
id, name, description, active_embedding_profile_hash
)
VALUES (%s, %s, %s, %s)
ON CONFLICT (id) DO UPDATE
SET name = EXCLUDED.name,
description = EXCLUDED.description,
active_embedding_profile_hash = EXCLUDED.active_embedding_profile_hash,
updated_at = now()
""",
(KNOWLEDGE_BASE_ID, "虚构地质 PoC 知识库", "仅含公开的合成验证文本"),
(
namespace.knowledge_base_id,
namespace.knowledge_base_name,
"仅含公开的合成验证文本",
profile_hash,
),
)
connection.execute(
"""
@@ -276,7 +425,11 @@ def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[st
VALUES (%s, %s, %s)
ON CONFLICT (id) DO NOTHING
""",
(ACCESS_SCOPE_ID, KNOWLEDGE_BASE_ID, "synthetic-demo"),
(
namespace.access_scope_id,
namespace.knowledge_base_id,
namespace.scope_name,
),
)
for item in chunks:
@@ -295,11 +448,11 @@ def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[st
""",
(
item.document_id,
KNOWLEDGE_BASE_ID,
ACCESS_SCOPE_ID,
namespace.knowledge_base_id,
namespace.access_scope_id,
item.raw_sha256,
f"{item.source_id}.json",
f"synthetic/{item.source_id}",
f"{namespace.storage_prefix}/{item.source_id}",
),
)
connection.execute(
@@ -384,10 +537,10 @@ def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[st
""",
(
item.chunk_id,
KNOWLEDGE_BASE_ID,
namespace.knowledge_base_id,
item.document_id,
item.version_id,
ACCESS_SCOPE_ID,
namespace.access_scope_id,
item.cloud_text,
item.cloud_text,
item.cloud_text_sha256,
@@ -425,6 +578,41 @@ def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[st
""",
(item.version_id,),
)
connection.execute(
"""
INSERT INTO rag.embedding_cache (
profile_hash, embedding_text_sha256, embedding, resolved_model,
provider_request_id, usage, elapsed_ms
) VALUES (%s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (profile_hash, embedding_text_sha256) DO NOTHING
""",
(
item.embedding_profile_hash,
item.embedding_text_sha256,
Vector(list(item.vector)),
item.embedding_model,
item.provider_request_id,
Jsonb(item.embedding_usage),
item.embedding_elapsed_ms,
),
)
connection.execute(
"""
INSERT INTO rag.chunk_embedding_assignments (
chunk_id, profile_hash, embedding_text_sha256,
cache_profile_hash, cache_embedding_text_sha256,
status, completed_at
) VALUES (%s, %s, %s, %s, %s, 'READY', now())
ON CONFLICT (chunk_id, profile_hash) DO NOTHING
""",
(
item.chunk_id,
item.embedding_profile_hash,
item.embedding_text_sha256,
item.embedding_profile_hash,
item.embedding_text_sha256,
),
)
connection.execute(
"""
UPDATE rag.chunks
@@ -459,8 +647,9 @@ def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[st
count(*) FILTER (WHERE searchable)::integer AS searchable
FROM rag.chunks
WHERE knowledge_base_id = %s
AND embedding_profile_hash = %s
""",
(KNOWLEDGE_BASE_ID,),
(namespace.knowledge_base_id, profile_hash),
).fetchone()
if counts is None:
raise SeedContractError("database_count_missing")
@@ -470,6 +659,9 @@ def write_chunks(settings: Settings, chunks: Sequence[PreparedChunk]) -> dict[st
def retrieve(
settings: Settings,
query_vector: tuple[float, ...],
*,
namespace: DemoNamespace,
profile_hash: str,
) -> list[dict[str, Any]]:
with psycopg.connect(database_dsn(settings), row_factory=dict_row) as connection:
register_vector(connection)
@@ -477,19 +669,25 @@ def retrieve(
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
SELECT chunk.id, chunk.metadata, chunk.embedding_text,
1 - (chunk.embedding <=> %s) AS vector_score
FROM rag.chunks AS chunk
JOIN rag.knowledge_bases AS knowledge_base
ON knowledge_base.id = chunk.knowledge_base_id
AND knowledge_base.active_embedding_profile_hash = %s
WHERE chunk.searchable
AND chunk.knowledge_base_id = %s
AND chunk.access_scope_id = %s
AND chunk.embedding_profile_hash = %s
ORDER BY chunk.embedding <=> %s
LIMIT %s
""",
(
Vector(list(query_vector)),
KNOWLEDGE_BASE_ID,
ACCESS_SCOPE_ID,
profile_hash,
namespace.knowledge_base_id,
namespace.access_scope_id,
profile_hash,
Vector(list(query_vector)),
settings.vector_top_k,
),
@@ -502,12 +700,20 @@ async def evaluate_queries(
queries: Sequence[DemoQuery],
embedder: EmbeddingProvider,
reranker: Reranker,
*,
namespace: DemoNamespace,
profile_hash: str,
) -> dict[str, float | int]:
hits = 0
answerable = 0
for query in queries:
query_result = await embedder.embed_query(query.query)
candidates = retrieve(settings, query_result.vectors[0])
candidates = retrieve(
settings,
query_result.vectors[0],
namespace=namespace,
profile_hash=profile_hash,
)
if not candidates:
continue
reranked = await reranker.rerank(
@@ -552,14 +758,14 @@ async def async_main() -> int:
return 2
settings = Settings()
namespace = namespace_for_mode(mode)
documents_path = Path(
os.getenv("DEMO_DOCUMENTS_PATH", str(DEFAULT_SAMPLE_ROOT / "demo_documents.jsonl"))
)
queries_path = Path(
os.getenv("DEMO_QUERIES_PATH", str(DEFAULT_SAMPLE_ROOT / "demo_queries.jsonl"))
)
cloud_embedder: BailianEmbeddingAdapter | None = None
cloud_reranker: BailianRerankerAdapter | None = None
cloud_gateway: ModelGatewayAdapter | None = None
try:
documents = load_documents(documents_path)
queries = load_queries(queries_path)
@@ -567,24 +773,9 @@ async def async_main() -> int:
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
cloud_gateway = ModelGatewayAdapter.from_settings(settings)
embedder = cloud_gateway
reranker = cloud_gateway
else:
embedder = FakeEmbeddingProvider(settings.embedding_dimension)
reranker = FakeReranker()
@@ -599,9 +790,17 @@ async def async_main() -> int:
vectors,
profile_hash=profile_hash,
embedding_model=resolved_model,
namespace=namespace,
)
counts = write_chunks(settings, prepared, namespace=namespace)
metrics = await evaluate_queries(
settings,
queries,
embedder,
reranker,
namespace=namespace,
profile_hash=profile_hash,
)
counts = write_chunks(settings, prepared)
metrics = await evaluate_queries(settings, queries, embedder, reranker)
output_summary(
{
"counts": counts,
@@ -654,10 +853,8 @@ async def async_main() -> int:
)
return 1
finally:
if cloud_embedder is not None:
await cloud_embedder.aclose()
if cloud_reranker is not None:
await cloud_reranker.aclose()
if cloud_gateway is not None:
await cloud_gateway.aclose()
def main() -> None: