mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-09 12:40:20 +08:00
feat: 引入 Rust executor/gateway sidecar 及 Python 侧双后端适配
- 新增 Rust workspace crates: aether-contracts, aether-executor, aether-gateway - aether-executor: 支持 Unix Socket/TCP 双传输模式,处理同步/流式上游请求 - aether-gateway: 作为本地主入口代理,集成 /api/internal/gateway/resolve 认证预解析 - Python 侧新增 ExecutionPlan 契约和 RustExecutorClient,各 handler 支持 executor_backend=rust 时将可序列化请求转发给 Rust executor 执行 - 重构 dev.sh 支持 executor/gateway 进程编排与生命周期管理 - 新增 internal gateway 路由,提供 resolve/passthrough 端点 - handler 层(chat/cli/video/endpoint_checker 等)全面适配 Rust executor 回退逻辑 - pipeline 层支持 trusted auth context 跳过重复认证 - 新增 Rust CI workflow 及对应测试用例
This commit is contained in:
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
from enum import Enum
|
||||
@@ -43,6 +44,12 @@ QUIET_POLLING_PATHS: set[str] = {
|
||||
"/api/wallet/today-cost",
|
||||
}
|
||||
|
||||
TRUSTED_GATEWAY_HEADER = "x-aether-gateway"
|
||||
TRUSTED_AUTH_USER_ID_HEADER = "x-aether-auth-user-id"
|
||||
TRUSTED_AUTH_API_KEY_ID_HEADER = "x-aether-auth-api-key-id"
|
||||
TRUSTED_AUTH_BALANCE_HEADER = "x-aether-auth-balance-remaining"
|
||||
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER = "x-aether-auth-access-allowed"
|
||||
|
||||
|
||||
class ApiRequestPipeline:
|
||||
"""负责统一执行认证、余额校验、上下文构建等通用逻辑的管道。"""
|
||||
@@ -349,6 +356,10 @@ class ApiRequestPipeline:
|
||||
async def _authenticate_client(
|
||||
self, request: Request, db: Session, adapter: ApiAdapter, **_kw: object
|
||||
) -> tuple[User, ApiKey]:
|
||||
trusted_auth = self._try_trusted_gateway_auth(request, db)
|
||||
if trusted_auth is not None:
|
||||
return trusted_auth
|
||||
|
||||
client_api_key = adapter.extract_api_key(request)
|
||||
if not client_api_key:
|
||||
raise HTTPException(status_code=401, detail="请提供API密钥")
|
||||
@@ -393,6 +404,86 @@ class ApiRequestPipeline:
|
||||
|
||||
return db_user, db_api_key
|
||||
|
||||
def _try_trusted_gateway_auth(
|
||||
self,
|
||||
request: Request,
|
||||
db: Session,
|
||||
) -> tuple[User, ApiKey] | None:
|
||||
if not self._is_trusted_gateway_request(request):
|
||||
return None
|
||||
|
||||
user_id = str(request.headers.get(TRUSTED_AUTH_USER_ID_HEADER) or "").strip()
|
||||
api_key_id = str(request.headers.get(TRUSTED_AUTH_API_KEY_ID_HEADER) or "").strip()
|
||||
if not user_id or not api_key_id:
|
||||
return None
|
||||
|
||||
db_user = db.query(User).filter(User.id == user_id).first()
|
||||
db_api_key = db.query(ApiKey).filter(ApiKey.id == api_key_id).first()
|
||||
if not db_user or not db_api_key:
|
||||
return None
|
||||
if not db_user.is_active or db_user.is_deleted:
|
||||
return None
|
||||
if not db_api_key.is_active:
|
||||
return None
|
||||
if db_api_key.is_locked and not db_api_key.is_standalone:
|
||||
raise HTTPException(status_code=403, detail="该密钥已被管理员锁定,请联系管理员")
|
||||
if db_api_key.user_id != db_user.id:
|
||||
return None
|
||||
if db_api_key.expires_at:
|
||||
expires_at = db_api_key.expires_at
|
||||
if expires_at.tzinfo is None:
|
||||
expires_at = expires_at.replace(tzinfo=timezone.utc)
|
||||
if expires_at < datetime.now(timezone.utc):
|
||||
return None
|
||||
|
||||
balance_remaining = self._parse_trusted_balance_remaining(
|
||||
request.headers.get(TRUSTED_AUTH_BALANCE_HEADER)
|
||||
)
|
||||
access_allowed = self._parse_trusted_bool_header(
|
||||
request.headers.get(TRUSTED_AUTH_ACCESS_ALLOWED_HEADER),
|
||||
default=True,
|
||||
)
|
||||
|
||||
request.state.user_id = db_user.id
|
||||
request.state.api_key_id = db_api_key.id
|
||||
request.state.prefetched_balance_remaining = balance_remaining
|
||||
|
||||
if not access_allowed:
|
||||
raise BalanceInsufficientException(balance_type="USD", remaining=balance_remaining)
|
||||
|
||||
return db_user, db_api_key
|
||||
|
||||
@staticmethod
|
||||
def _is_trusted_gateway_request(request: Request) -> bool:
|
||||
gateway_marker = str(request.headers.get(TRUSTED_GATEWAY_HEADER) or "").strip().lower()
|
||||
if not gateway_marker.startswith("rust-phase3"):
|
||||
return False
|
||||
|
||||
host = request.client.host if request.client else ""
|
||||
try:
|
||||
return ipaddress.ip_address(host).is_loopback
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
@staticmethod
|
||||
def _parse_trusted_balance_remaining(value: str | None) -> float | None:
|
||||
raw = str(value or "").strip()
|
||||
if not raw:
|
||||
return None
|
||||
try:
|
||||
return float(raw)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _parse_trusted_bool_header(value: str | None, *, default: bool) -> bool:
|
||||
raw = str(value or "").strip().lower()
|
||||
if raw in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if raw in {"0", "false", "no", "off"}:
|
||||
return False
|
||||
return default
|
||||
|
||||
async def _try_token_prefix_auth(
|
||||
self, token: str, request: Request, db: Session
|
||||
) -> tuple[User, Any] | None:
|
||||
|
||||
@@ -92,6 +92,17 @@ from src.services.provider.transport import (
|
||||
build_provider_url,
|
||||
)
|
||||
from src.services.provider.upstream_headers import build_upstream_extra_headers
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
build_execution_plan_body,
|
||||
is_remote_contract_eligible,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
from src.services.scheduling.aware_scheduler import ProviderCandidate
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.task.request_state import MutableRequestBodyState
|
||||
@@ -958,7 +969,10 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
|
||||
|
||||
# 解析有效代理(Key 级别优先于 Provider 级别)
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_proxy_label,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_effective_proxy,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
@@ -971,6 +985,7 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
|
||||
f" [{self.request_id}] 发送流式请求: Provider={provider.name}, "
|
||||
f"模型={ctx.model} -> {mapped_model or '无映射'}, 代理={proxy_label}"
|
||||
)
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
|
||||
# If upstream is forced to non-stream mode, we execute a sync request and then
|
||||
# simulate streaming to the client (sync -> stream bridge).
|
||||
@@ -978,11 +993,9 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_post_kwargs_async,
|
||||
resolve_delegate_config_async,
|
||||
)
|
||||
|
||||
request_timeout_sync = provider.request_timeout or config.http_request_timeout
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
proxy_config=effective_proxy,
|
||||
@@ -1179,19 +1192,142 @@ class ChatHandlerBase(BaseMessageHandler, ABC):
|
||||
|
||||
return _streamified()
|
||||
|
||||
if config.executor_backend == "rust":
|
||||
effective_proxy_for_contract = effective_proxy
|
||||
if not effective_proxy_for_contract or not effective_proxy_for_contract.get(
|
||||
"enabled", True
|
||||
):
|
||||
effective_proxy_for_contract = await get_system_proxy_config_async()
|
||||
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy_for_contract and not is_tunnel_delegate:
|
||||
proxy_url = await build_proxy_url_async(effective_proxy_for_contract)
|
||||
|
||||
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
ctx.proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if is_tunnel_delegate else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if is_tunnel_delegate
|
||||
else None
|
||||
),
|
||||
)
|
||||
rust_plan = ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=str(
|
||||
getattr(candidate, "request_candidate_id", "")
|
||||
or getattr(candidate, "id", "")
|
||||
or ""
|
||||
)
|
||||
or None,
|
||||
provider_name=str(provider.name),
|
||||
provider_id=str(provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(provider_headers),
|
||||
body=build_execution_plan_body(
|
||||
provider_payload,
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
),
|
||||
stream=True,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
model_name=str(ctx.model or ""),
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
content_encoding=client_content_encoding,
|
||||
proxy=proxy_snapshot,
|
||||
tls_profile=tls_profile,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=int(config.http_connect_timeout * 1000),
|
||||
read_ms=int(config.http_read_timeout * 1000),
|
||||
write_ms=int(config.http_write_timeout * 1000),
|
||||
pool_ms=int(config.http_pool_timeout * 1000),
|
||||
total_ms=None,
|
||||
),
|
||||
)
|
||||
|
||||
if is_remote_contract_eligible(rust_plan):
|
||||
try:
|
||||
rust_stream = await RustExecutorClient().execute_stream(rust_plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[{}] Rust executor stream 不可用,回退 Python 执行: {}",
|
||||
self.request_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
ctx.status_code = rust_stream.status_code
|
||||
ctx.response_headers = dict(rust_stream.headers)
|
||||
ctx.set_proxy_timing(ctx.response_headers)
|
||||
if envelope:
|
||||
envelope.on_http_status(
|
||||
base_url=ctx.selected_base_url,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
if ctx.status_code >= 400:
|
||||
error_chunks: list[bytes] = []
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
if chunk:
|
||||
error_chunks.append(chunk)
|
||||
if sum(len(item) for item in error_chunks) >= 4000:
|
||||
break
|
||||
error_body = b"".join(error_chunks)[:4000].decode(
|
||||
"utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
request = httpx.Request("POST", url, headers=provider_headers)
|
||||
response = httpx.Response(
|
||||
ctx.status_code,
|
||||
request=request,
|
||||
headers=rust_stream.headers,
|
||||
content=b"".join(error_chunks),
|
||||
)
|
||||
error = httpx.HTTPStatusError(
|
||||
f"Upstream status error: {ctx.status_code}",
|
||||
request=request,
|
||||
response=response,
|
||||
)
|
||||
error.upstream_response = error_body # type: ignore[attr-defined]
|
||||
raise error
|
||||
|
||||
prefetched_chunks = await stream_processor.prefetch_and_check_error(
|
||||
rust_stream.byte_iterator,
|
||||
provider,
|
||||
endpoint,
|
||||
ctx,
|
||||
max_prefetch_lines=config.stream_prefetch_lines,
|
||||
)
|
||||
except Exception:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
raise
|
||||
|
||||
return stream_processor.create_response_stream(
|
||||
ctx,
|
||||
rust_stream.byte_iterator,
|
||||
rust_stream.response_ctx,
|
||||
prefetched_chunks,
|
||||
start_time=self.start_time,
|
||||
)
|
||||
|
||||
# 流式请求使用 stream_first_byte_timeout 作为首字节超时
|
||||
# 优先使用 Provider 配置,否则使用全局配置
|
||||
request_timeout = provider.stream_first_byte_timeout or config.stream_first_byte_timeout
|
||||
request_timeout = (
|
||||
getattr(provider, "stream_first_byte_timeout", None) or config.stream_first_byte_timeout
|
||||
)
|
||||
|
||||
# 获取 HTTP 客户端(支持代理配置,Key 级别优先于 Provider 级别)
|
||||
# 使用连接池复用客户端,避免每次流式请求都新建 TCP/TLS 连接
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_stream_kwargs_async,
|
||||
resolve_delegate_config_async,
|
||||
)
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
proxy_config=effective_proxy,
|
||||
|
||||
@@ -35,6 +35,7 @@ from src.api.handlers.base.utils import (
|
||||
resolve_client_accept_encoding,
|
||||
resolve_client_content_encoding,
|
||||
)
|
||||
from src.config.settings import config
|
||||
from src.core.error_utils import extract_client_error_message
|
||||
from src.core.exceptions import (
|
||||
EmbeddedErrorException,
|
||||
@@ -46,6 +47,17 @@ from src.core.exceptions import (
|
||||
UpstreamClientException,
|
||||
)
|
||||
from src.core.logger import logger
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
PreparedExecutionPlan,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
from src.services.task.request_state import MutableRequestBodyState
|
||||
|
||||
if TYPE_CHECKING:
|
||||
@@ -442,6 +454,39 @@ class ChatSyncExecutor:
|
||||
client_content_encoding: str | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""单次同步请求(原 sync_request_func 内嵌函数)"""
|
||||
prepared_plan = await self._build_sync_execution_plan(
|
||||
provider,
|
||||
endpoint,
|
||||
key,
|
||||
candidate,
|
||||
model=model,
|
||||
api_format=api_format,
|
||||
original_headers=original_headers,
|
||||
request_state=request_state,
|
||||
query_params=query_params,
|
||||
client_content_encoding=client_content_encoding,
|
||||
)
|
||||
return await self._execute_sync_plan(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=provider,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def _build_sync_execution_plan(
|
||||
self,
|
||||
provider: Provider,
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
candidate: ProviderCandidate,
|
||||
*,
|
||||
model: str,
|
||||
api_format: Any,
|
||||
original_headers: dict[str, Any],
|
||||
request_state: MutableRequestBodyState,
|
||||
query_params: dict[str, str] | None = None,
|
||||
client_content_encoding: str | None = None,
|
||||
) -> PreparedExecutionPlan:
|
||||
"""构建可序列化的执行计划,并保留本地执行所需的运行时上下文。"""
|
||||
handler = self._handler
|
||||
ctx = self._ctx
|
||||
|
||||
@@ -518,7 +563,9 @@ class ChatSyncExecutor:
|
||||
|
||||
# 解析有效代理(Key 级别优先于 Provider 级别)
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_proxy_label,
|
||||
get_system_proxy_config_async,
|
||||
resolve_effective_proxy,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
@@ -536,13 +583,8 @@ class ChatSyncExecutor:
|
||||
)
|
||||
logger.debug(f" [{handler.request_id}] 请求URL: {redact_url_for_log(url)}")
|
||||
|
||||
# 获取复用的 HTTP 客户端(支持代理配置,Key 级别优先于 Provider 级别)
|
||||
# 注意:使用 get_proxy_client 复用连接池,不再每次创建新客户端
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.config.settings import config
|
||||
# 解析 delegate 配置,用于后续本地执行或交给 Rust executor。
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_post_kwargs_async,
|
||||
build_stream_kwargs_async,
|
||||
resolve_delegate_config_async,
|
||||
)
|
||||
|
||||
@@ -551,49 +593,354 @@ class ChatSyncExecutor:
|
||||
request_timeout = provider.request_timeout or config.http_request_timeout
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(_effective_proxy)
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
effective_proxy_for_contract = _effective_proxy
|
||||
if not effective_proxy_for_contract or not effective_proxy_for_contract.get(
|
||||
"enabled", True
|
||||
):
|
||||
effective_proxy_for_contract = await get_system_proxy_config_async()
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy_for_contract and not is_tunnel_delegate:
|
||||
proxy_url = await build_proxy_url_async(effective_proxy_for_contract)
|
||||
|
||||
return PreparedExecutionPlan(
|
||||
contract=ExecutionPlan(
|
||||
request_id=str(handler.request_id or ""),
|
||||
candidate_id=str(
|
||||
getattr(candidate, "request_candidate_id", "")
|
||||
or getattr(candidate, "id", "")
|
||||
or ""
|
||||
)
|
||||
or None,
|
||||
provider_name=str(provider.name),
|
||||
provider_id=str(provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(provider_hdrs),
|
||||
body=build_execution_plan_body(
|
||||
provider_payload,
|
||||
content_type=str(provider_hdrs.get("content-type") or "").strip() or None,
|
||||
),
|
||||
stream=upstream_is_stream,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
model_name=str(model or ""),
|
||||
content_type=str(provider_hdrs.get("content-type") or "").strip() or None,
|
||||
content_encoding=client_content_encoding,
|
||||
proxy=ExecutionProxySnapshot.from_proxy_info(
|
||||
ctx.sync_proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if is_tunnel_delegate else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if is_tunnel_delegate
|
||||
else None
|
||||
),
|
||||
),
|
||||
tls_profile=tls_profile,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=int(config.http_connect_timeout * 1000),
|
||||
read_ms=int(config.http_read_timeout * 1000),
|
||||
write_ms=int(config.http_write_timeout * 1000),
|
||||
pool_ms=int(config.http_pool_timeout * 1000),
|
||||
total_ms=int(request_timeout * 1000),
|
||||
),
|
||||
),
|
||||
payload=provider_payload,
|
||||
headers=dict(provider_hdrs),
|
||||
upstream_is_stream=upstream_is_stream,
|
||||
needs_conversion=needs_conversion,
|
||||
provider_type=provider_type,
|
||||
request_timeout=request_timeout,
|
||||
delegate_config=delegate_cfg,
|
||||
proxy_config=_effective_proxy,
|
||||
tls_profile=tls_profile,
|
||||
envelope=envelope,
|
||||
selected_base_url=selected_base_url_cached,
|
||||
client_content_encoding=client_content_encoding,
|
||||
proxy_info=ctx.sync_proxy_info,
|
||||
)
|
||||
|
||||
async def _execute_sync_plan(
|
||||
self,
|
||||
*,
|
||||
prepared_plan: PreparedExecutionPlan,
|
||||
provider: Provider,
|
||||
model: str,
|
||||
) -> dict[str, Any]:
|
||||
handler = self._handler
|
||||
ctx = self._ctx
|
||||
|
||||
if config.executor_backend == "rust" and prepared_plan.remote_eligible:
|
||||
try:
|
||||
rust_result = await RustExecutorClient().execute_sync_json(prepared_plan.contract)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[{}] Rust executor 不可用,回退 Python 执行: {}",
|
||||
handler.request_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
result = await self._finalize_rust_sync_result(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=provider,
|
||||
model=model,
|
||||
status_code=rust_result.status_code,
|
||||
response_headers=rust_result.headers,
|
||||
response_json=rust_result.response_json,
|
||||
provider_response_json=rust_result.provider_response_json,
|
||||
response_body_bytes=rust_result.response_body_bytes,
|
||||
)
|
||||
logger.debug(
|
||||
"[{}] sync chat 请求由 Rust executor 执行完成",
|
||||
handler.request_id,
|
||||
)
|
||||
return result
|
||||
|
||||
return await self._execute_sync_plan_locally(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=provider,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def _finalize_rust_sync_result(
|
||||
self,
|
||||
*,
|
||||
prepared_plan: PreparedExecutionPlan,
|
||||
provider: Provider,
|
||||
model: str,
|
||||
status_code: int,
|
||||
response_headers: dict[str, str],
|
||||
response_json: dict[str, Any] | None,
|
||||
provider_response_json: dict[str, Any] | None = None,
|
||||
response_body_bytes: bytes | None = None,
|
||||
) -> dict[str, Any]:
|
||||
ctx = self._ctx
|
||||
|
||||
synthetic_content = response_body_bytes
|
||||
if synthetic_content is None:
|
||||
synthetic_content = json.dumps(
|
||||
response_json or {},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
|
||||
request = httpx.Request(
|
||||
prepared_plan.contract.method,
|
||||
prepared_plan.contract.url,
|
||||
headers=prepared_plan.headers,
|
||||
)
|
||||
synthetic_response = httpx.Response(
|
||||
status_code,
|
||||
request=request,
|
||||
headers=response_headers,
|
||||
content=synthetic_content,
|
||||
)
|
||||
|
||||
ctx.status_code = status_code
|
||||
ctx.response_headers = dict(synthetic_response.headers)
|
||||
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
|
||||
|
||||
if prepared_plan.envelope:
|
||||
prepared_plan.envelope.on_http_status(
|
||||
base_url=prepared_plan.selected_base_url,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
|
||||
if status_code >= 400:
|
||||
error = httpx.HTTPStatusError(
|
||||
f"Upstream status error: {status_code}",
|
||||
request=request,
|
||||
response=synthetic_response,
|
||||
)
|
||||
error_body = ""
|
||||
try:
|
||||
if prepared_plan.envelope and hasattr(prepared_plan.envelope, "extract_error_text"):
|
||||
error_body = await prepared_plan.envelope.extract_error_text(synthetic_response)
|
||||
else:
|
||||
error_body = synthetic_response.text[:4000] if synthetic_response.text else ""
|
||||
except Exception:
|
||||
error_body = synthetic_response.text[:4000] if synthetic_response.text else ""
|
||||
error.upstream_response = error_body[:4000] # type: ignore[attr-defined]
|
||||
raise error
|
||||
|
||||
if prepared_plan.upstream_is_stream:
|
||||
if response_body_bytes is None:
|
||||
raise RustExecutorClientError("Rust executor stream result must contain body bytes")
|
||||
return await self._finalize_rust_stream_sync_result(
|
||||
prepared_plan=prepared_plan,
|
||||
provider=provider,
|
||||
model=model,
|
||||
response_body_bytes=response_body_bytes,
|
||||
)
|
||||
|
||||
if response_json is None:
|
||||
raise RustExecutorClientError("Rust executor sync result must contain response_json")
|
||||
|
||||
ctx.response_json = dict(response_json)
|
||||
if provider_response_json is not None:
|
||||
ctx.provider_response_json = dict(provider_response_json)
|
||||
|
||||
if prepared_plan.envelope:
|
||||
ctx.response_json = prepared_plan.envelope.unwrap_response(ctx.response_json)
|
||||
prepared_plan.envelope.postprocess_unwrapped_response(
|
||||
model=model,
|
||||
data=ctx.response_json,
|
||||
)
|
||||
|
||||
if isinstance(ctx.response_json, dict):
|
||||
parser = get_parser_for_format(ctx.provider_api_format_for_error or "")
|
||||
if parser.is_error_response(ctx.response_json):
|
||||
parsed = parser.parse_response(ctx.response_json, status_code)
|
||||
raise EmbeddedErrorException(
|
||||
provider_name=str(provider.name),
|
||||
error_code=parsed.embedded_status_code,
|
||||
error_message=parsed.error_message,
|
||||
error_status=parsed.error_type,
|
||||
)
|
||||
|
||||
if ctx.needs_conversion_for_error and isinstance(ctx.response_json, dict):
|
||||
ctx.provider_response_json = ctx.response_json.copy()
|
||||
registry = get_format_converter_registry()
|
||||
ctx.response_json = registry.convert_response(
|
||||
ctx.response_json,
|
||||
ctx.provider_api_format_for_error or "",
|
||||
ctx.client_api_format_for_error or "",
|
||||
requested_model=model,
|
||||
)
|
||||
|
||||
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
|
||||
|
||||
async def _finalize_rust_stream_sync_result(
|
||||
self,
|
||||
*,
|
||||
prepared_plan: PreparedExecutionPlan,
|
||||
provider: Provider,
|
||||
model: str,
|
||||
response_body_bytes: bytes,
|
||||
) -> dict[str, Any]:
|
||||
async def _iter_body() -> Any:
|
||||
yield response_body_bytes
|
||||
|
||||
return await self._aggregate_upstream_stream_response(
|
||||
byte_iter=_iter_body(),
|
||||
prepared_plan=prepared_plan,
|
||||
provider=provider,
|
||||
model=model,
|
||||
)
|
||||
|
||||
async def _aggregate_upstream_stream_response(
|
||||
self,
|
||||
*,
|
||||
byte_iter: Any,
|
||||
prepared_plan: PreparedExecutionPlan,
|
||||
provider: Provider,
|
||||
model: str,
|
||||
) -> dict[str, Any]:
|
||||
ctx = self._ctx
|
||||
|
||||
provider_parser = (
|
||||
get_parser_for_format(ctx.provider_api_format_for_error)
|
||||
if ctx.provider_api_format_for_error
|
||||
else None
|
||||
)
|
||||
|
||||
if (
|
||||
prepared_plan.provider_type == "kiro"
|
||||
and prepared_plan.envelope
|
||||
and prepared_plan.envelope.force_stream_rewrite()
|
||||
):
|
||||
from src.services.provider.adapters.kiro.eventstream_rewriter import (
|
||||
apply_kiro_stream_rewrite,
|
||||
)
|
||||
|
||||
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
|
||||
|
||||
from src.api.handlers.base.upstream_stream_bridge import (
|
||||
aggregate_upstream_stream_to_internal_response,
|
||||
)
|
||||
|
||||
internal_resp = await aggregate_upstream_stream_to_internal_response(
|
||||
byte_iter,
|
||||
provider_api_format=ctx.provider_api_format_for_error or "",
|
||||
provider_name=str(provider.name),
|
||||
model=str(model or ""),
|
||||
request_id=str(self._handler.request_id or ""),
|
||||
envelope=prepared_plan.envelope,
|
||||
provider_parser=provider_parser,
|
||||
)
|
||||
|
||||
registry = get_format_converter_registry()
|
||||
tgt_norm = (
|
||||
registry.get_normalizer(ctx.client_api_format_for_error)
|
||||
if ctx.client_api_format_for_error
|
||||
else None
|
||||
)
|
||||
if tgt_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {ctx.client_api_format_for_error}")
|
||||
|
||||
ctx.response_json = tgt_norm.response_from_internal(
|
||||
internal_resp,
|
||||
requested_model=model,
|
||||
)
|
||||
ctx.response_json = ctx.response_json if isinstance(ctx.response_json, dict) else {}
|
||||
return ctx.response_json
|
||||
|
||||
async def _execute_sync_plan_locally(
|
||||
self,
|
||||
*,
|
||||
prepared_plan: PreparedExecutionPlan,
|
||||
provider: Provider,
|
||||
model: str,
|
||||
) -> dict[str, Any]:
|
||||
handler = self._handler
|
||||
ctx = self._ctx
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_post_kwargs_async,
|
||||
build_stream_kwargs_async,
|
||||
)
|
||||
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
prepared_plan.delegate_config,
|
||||
proxy_config=prepared_plan.proxy_config,
|
||||
tls_profile=prepared_plan.contract.tls_profile,
|
||||
)
|
||||
|
||||
# 注意:不使用 async with,因为复用的客户端不应该被关闭
|
||||
# 超时通过 timeout 参数控制
|
||||
resp: httpx.Response | None = None
|
||||
if not upstream_is_stream:
|
||||
if not prepared_plan.upstream_is_stream:
|
||||
try:
|
||||
_pkw = await build_post_kwargs_async(
|
||||
delegate_cfg,
|
||||
url=url,
|
||||
headers=provider_hdrs,
|
||||
payload=provider_payload,
|
||||
timeout=request_timeout,
|
||||
client_content_encoding=client_content_encoding,
|
||||
prepared_plan.delegate_config,
|
||||
url=prepared_plan.contract.url,
|
||||
headers=prepared_plan.headers,
|
||||
payload=prepared_plan.payload,
|
||||
timeout=prepared_plan.request_timeout,
|
||||
client_content_encoding=prepared_plan.client_content_encoding,
|
||||
)
|
||||
resp = await http_client.post(**_pkw)
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
|
||||
if envelope:
|
||||
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
|
||||
if selected_base_url_cached:
|
||||
if prepared_plan.envelope:
|
||||
prepared_plan.envelope.on_connection_error(
|
||||
base_url=prepared_plan.selected_base_url,
|
||||
exc=e,
|
||||
)
|
||||
if prepared_plan.selected_base_url:
|
||||
logger.warning(
|
||||
f"[{envelope.name}] Connection error: "
|
||||
f"{selected_base_url_cached} ({e})"
|
||||
f"[{prepared_plan.envelope.name}] Connection error: "
|
||||
f"{prepared_plan.selected_base_url} ({e})"
|
||||
)
|
||||
raise
|
||||
else:
|
||||
# Forced upstream streaming: aggregate SSE to a sync JSON response.
|
||||
provider_parser = (
|
||||
get_parser_for_format(provider_api_format) if provider_api_format else None
|
||||
)
|
||||
|
||||
try:
|
||||
_stream_args = await build_stream_kwargs_async(
|
||||
delegate_cfg,
|
||||
url=url,
|
||||
headers=provider_hdrs,
|
||||
payload=provider_payload,
|
||||
timeout=request_timeout,
|
||||
client_content_encoding=client_content_encoding,
|
||||
prepared_plan.delegate_config,
|
||||
url=prepared_plan.contract.url,
|
||||
headers=prepared_plan.headers,
|
||||
payload=prepared_plan.payload,
|
||||
timeout=prepared_plan.request_timeout,
|
||||
client_content_encoding=prepared_plan.client_content_encoding,
|
||||
)
|
||||
async with http_client.stream(**_stream_args) as stream_resp:
|
||||
resp = stream_resp
|
||||
@@ -602,58 +949,31 @@ class ChatSyncExecutor:
|
||||
ctx.response_headers = dict(stream_resp.headers)
|
||||
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
|
||||
|
||||
if envelope:
|
||||
envelope.on_http_status(
|
||||
base_url=selected_base_url_cached,
|
||||
if prepared_plan.envelope:
|
||||
prepared_plan.envelope.on_http_status(
|
||||
base_url=prepared_plan.selected_base_url,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
|
||||
stream_resp.raise_for_status()
|
||||
|
||||
byte_iter = stream_resp.aiter_bytes()
|
||||
if provider_type == "kiro" and envelope and envelope.force_stream_rewrite():
|
||||
from src.services.provider.adapters.kiro.eventstream_rewriter import (
|
||||
apply_kiro_stream_rewrite,
|
||||
)
|
||||
|
||||
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
|
||||
|
||||
from src.api.handlers.base.upstream_stream_bridge import (
|
||||
aggregate_upstream_stream_to_internal_response,
|
||||
)
|
||||
|
||||
internal_resp = await aggregate_upstream_stream_to_internal_response(
|
||||
byte_iter,
|
||||
provider_api_format=provider_api_format,
|
||||
provider_name=str(provider.name),
|
||||
model=str(model or ""),
|
||||
request_id=str(handler.request_id or ""),
|
||||
envelope=envelope,
|
||||
provider_parser=provider_parser,
|
||||
)
|
||||
|
||||
registry = get_format_converter_registry()
|
||||
tgt_norm = (
|
||||
registry.get_normalizer(client_api_format) if client_api_format else None
|
||||
)
|
||||
if tgt_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
|
||||
|
||||
ctx.response_json = tgt_norm.response_from_internal(
|
||||
internal_resp,
|
||||
requested_model=model,
|
||||
)
|
||||
ctx.response_json = (
|
||||
ctx.response_json if isinstance(ctx.response_json, dict) else {}
|
||||
ctx.response_json = await self._aggregate_upstream_stream_response(
|
||||
byte_iter=stream_resp.aiter_bytes(),
|
||||
prepared_plan=prepared_plan,
|
||||
provider=provider,
|
||||
model=model,
|
||||
)
|
||||
|
||||
except (httpx.ConnectError, httpx.ConnectTimeout, httpx.TimeoutException) as e:
|
||||
if envelope:
|
||||
envelope.on_connection_error(base_url=selected_base_url_cached, exc=e)
|
||||
if selected_base_url_cached:
|
||||
if prepared_plan.envelope:
|
||||
prepared_plan.envelope.on_connection_error(
|
||||
base_url=prepared_plan.selected_base_url,
|
||||
exc=e,
|
||||
)
|
||||
if prepared_plan.selected_base_url:
|
||||
logger.warning(
|
||||
f"[{envelope.name}] Connection error: "
|
||||
f"{selected_base_url_cached} ({e})"
|
||||
f"[{prepared_plan.envelope.name}] Connection error: "
|
||||
f"{prepared_plan.selected_base_url} ({e})"
|
||||
)
|
||||
raise
|
||||
|
||||
@@ -661,35 +981,32 @@ class ChatSyncExecutor:
|
||||
ctx.response_headers = dict(resp.headers)
|
||||
extract_proxy_timing(ctx.sync_proxy_info, ctx.response_headers)
|
||||
|
||||
if envelope:
|
||||
envelope.on_http_status(base_url=selected_base_url_cached, status_code=ctx.status_code)
|
||||
if prepared_plan.envelope:
|
||||
prepared_plan.envelope.on_http_status(
|
||||
base_url=prepared_plan.selected_base_url,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
|
||||
# Forced upstream streaming already built response_json via aggregator.
|
||||
if upstream_is_stream:
|
||||
if prepared_plan.upstream_is_stream:
|
||||
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
|
||||
|
||||
# 统一使用 HTTPStatusError,让 TaskService/error_classifier 负责分类
|
||||
# (客户端错误/兼容性错误/限流等)
|
||||
try:
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
error_body = ""
|
||||
try:
|
||||
if envelope and hasattr(envelope, "extract_error_text"):
|
||||
error_body = await envelope.extract_error_text(resp)
|
||||
if prepared_plan.envelope and hasattr(prepared_plan.envelope, "extract_error_text"):
|
||||
error_body = await prepared_plan.envelope.extract_error_text(resp)
|
||||
else:
|
||||
error_body = resp.text[:4000] if resp.text else ""
|
||||
except Exception:
|
||||
error_body = ""
|
||||
# 供 ErrorClassifier 优先读取
|
||||
e.upstream_response = error_body # type: ignore[attr-defined]
|
||||
raise
|
||||
|
||||
# 安全解析 JSON 响应,处理可能的编码错误
|
||||
try:
|
||||
ctx.response_json = resp.json()
|
||||
except (UnicodeDecodeError, json.JSONDecodeError) as e:
|
||||
# 获取原始响应内容用于调试(存入 upstream_response)
|
||||
raw_content = ""
|
||||
try:
|
||||
raw_content = resp.text[:500] if resp.text else "(empty)"
|
||||
@@ -699,7 +1016,6 @@ class ChatSyncExecutor:
|
||||
except Exception:
|
||||
raw_content = "(unable to read)"
|
||||
logger.error(f"[{handler.request_id}] 无法解析响应 JSON: {e}, 原始内容: {raw_content}")
|
||||
# 判断错误类型,生成友好的客户端错误消息(不暴露提供商信息)
|
||||
if raw_content == "(empty)" or not raw_content.strip():
|
||||
client_message = "上游服务返回了空响应"
|
||||
elif raw_content.strip().startswith(("<", "<!doctype", "<!DOCTYPE")):
|
||||
@@ -713,13 +1029,15 @@ class ChatSyncExecutor:
|
||||
upstream_response=raw_content,
|
||||
)
|
||||
|
||||
if envelope:
|
||||
ctx.response_json = envelope.unwrap_response(ctx.response_json)
|
||||
envelope.postprocess_unwrapped_response(model=model, data=ctx.response_json)
|
||||
if prepared_plan.envelope:
|
||||
ctx.response_json = prepared_plan.envelope.unwrap_response(ctx.response_json)
|
||||
prepared_plan.envelope.postprocess_unwrapped_response(
|
||||
model=model,
|
||||
data=ctx.response_json,
|
||||
)
|
||||
|
||||
# 检查响应体中的嵌套错误(HTTP 200 但响应体包含错误)
|
||||
if isinstance(ctx.response_json, dict):
|
||||
parser = get_parser_for_format(provider_api_format)
|
||||
parser = get_parser_for_format(ctx.provider_api_format_for_error or "")
|
||||
if parser.is_error_response(ctx.response_json):
|
||||
parsed = parser.parse_response(ctx.response_json, 200)
|
||||
logger.warning(
|
||||
@@ -736,15 +1054,14 @@ class ChatSyncExecutor:
|
||||
error_status=parsed.error_type,
|
||||
)
|
||||
|
||||
# 跨格式:响应转换回 client_format(失败触发 failover)
|
||||
if needs_conversion and isinstance(ctx.response_json, dict):
|
||||
if ctx.needs_conversion_for_error and isinstance(ctx.response_json, dict):
|
||||
ctx.provider_response_json = ctx.response_json.copy()
|
||||
registry = get_format_converter_registry()
|
||||
ctx.response_json = registry.convert_response(
|
||||
ctx.response_json,
|
||||
provider_api_format,
|
||||
client_api_format,
|
||||
requested_model=model, # 使用用户请求的原始模型名
|
||||
ctx.provider_api_format_for_error or "",
|
||||
ctx.client_api_format_for_error or "",
|
||||
requested_model=model,
|
||||
)
|
||||
|
||||
return ctx.response_json if isinstance(ctx.response_json, dict) else {}
|
||||
|
||||
@@ -38,6 +38,17 @@ from src.core.exceptions import (
|
||||
)
|
||||
from src.core.logger import logger
|
||||
from src.services.provider.behavior import get_provider_behavior
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
build_execution_plan_body,
|
||||
is_remote_contract_eligible,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
from src.services.scheduling.aware_scheduler import ProviderCandidate
|
||||
from src.services.system.config import SystemConfigService
|
||||
from src.services.task.request_state import MutableRequestBodyState
|
||||
@@ -53,6 +64,67 @@ if TYPE_CHECKING:
|
||||
class CliStreamMixin:
|
||||
"""流式处理核心方法的 Mixin"""
|
||||
|
||||
def _streamify_sync_response(
|
||||
self: CliHandlerProtocol,
|
||||
*,
|
||||
ctx: StreamContext,
|
||||
response_json: dict[str, Any],
|
||||
client_api_format: str,
|
||||
provider_api_format: str,
|
||||
) -> AsyncGenerator[bytes]:
|
||||
registry = get_format_converter_registry()
|
||||
src_norm = registry.get_normalizer(provider_api_format) if provider_api_format else None
|
||||
if src_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {provider_api_format}")
|
||||
|
||||
internal_resp = src_norm.response_to_internal(
|
||||
response_json if isinstance(response_json, dict) else {}
|
||||
)
|
||||
internal_resp.model = str(ctx.model or internal_resp.model or "")
|
||||
if internal_resp.id:
|
||||
ctx.response_id = internal_resp.id
|
||||
|
||||
if internal_resp.usage:
|
||||
ctx.input_tokens = int(internal_resp.usage.input_tokens or 0)
|
||||
ctx.output_tokens = int(internal_resp.usage.output_tokens or 0)
|
||||
ctx.cached_tokens = int(internal_resp.usage.cache_read_tokens or 0)
|
||||
ctx.cache_creation_tokens = int(internal_resp.usage.cache_write_tokens or 0)
|
||||
|
||||
from src.core.api_format.conversion.stream_state import StreamState
|
||||
|
||||
tgt_norm = registry.get_normalizer(client_api_format) if client_api_format else None
|
||||
if tgt_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
|
||||
|
||||
state = StreamState(
|
||||
model=str(ctx.model or ""),
|
||||
message_id=str(ctx.response_id or ctx.request_id or self.request_id or ""),
|
||||
)
|
||||
output_state = {"first_yield": True, "streaming_updated": False}
|
||||
|
||||
async def _streamified() -> AsyncGenerator[bytes]:
|
||||
for ev in iter_internal_response_as_stream_events(internal_resp):
|
||||
converted_events = tgt_norm.stream_event_from_internal(ev, state)
|
||||
if not converted_events:
|
||||
continue
|
||||
self._record_converted_chunks(ctx, converted_events)
|
||||
for sse_line in _format_converted_events_to_sse(
|
||||
converted_events, client_api_format
|
||||
):
|
||||
if not sse_line:
|
||||
continue
|
||||
ctx.chunk_count += 1
|
||||
self._mark_first_output(ctx, output_state)
|
||||
yield (sse_line + "\n").encode("utf-8")
|
||||
|
||||
if str(client_api_format or "").strip().lower() == "openai:chat":
|
||||
ctx.chunk_count += 1
|
||||
self._mark_first_output(ctx, output_state)
|
||||
yield b"data: [DONE]\n\n"
|
||||
ctx.has_completion = True
|
||||
|
||||
return _streamified()
|
||||
|
||||
async def process_stream(
|
||||
self: CliHandlerProtocol,
|
||||
original_request_body: dict[str, Any],
|
||||
@@ -341,24 +413,156 @@ class CliStreamMixin:
|
||||
ctx.selected_base_url = upstream_request.selected_base_url
|
||||
|
||||
# 解析有效代理(Key 级别优先于 Provider 级别)
|
||||
from src.services.proxy_node.resolver import build_proxy_url_async as _bpua
|
||||
from src.services.proxy_node.resolver import get_proxy_label as _gpl
|
||||
from src.services.proxy_node.resolver import resolve_delegate_config_async as _rda
|
||||
from src.services.proxy_node.resolver import resolve_effective_proxy as _rep
|
||||
from src.services.proxy_node.resolver import resolve_proxy_info_async as _rpi_async
|
||||
|
||||
effective_proxy = _rep(provider.proxy, getattr(key, "proxy", None))
|
||||
ctx.proxy_info = await _rpi_async(effective_proxy)
|
||||
delegate_cfg = await _rda(effective_proxy)
|
||||
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy and not is_tunnel_delegate:
|
||||
proxy_url = await _bpua(effective_proxy)
|
||||
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
ctx.proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if is_tunnel_delegate else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if is_tunnel_delegate
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
# If upstream is forced to non-stream mode, we execute a sync request and then
|
||||
# simulate streaming to the client (sync -> stream bridge).
|
||||
if not upstream_is_stream:
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_post_kwargs_async,
|
||||
resolve_delegate_config_async,
|
||||
)
|
||||
from src.services.proxy_node.resolver import build_post_kwargs_async
|
||||
|
||||
request_timeout_sync = provider.request_timeout or config.http_request_timeout
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
|
||||
rust_plan = ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=str(
|
||||
getattr(candidate, "request_candidate_id", "")
|
||||
or getattr(candidate, "id", "")
|
||||
or ""
|
||||
)
|
||||
or None,
|
||||
provider_name=str(provider.name),
|
||||
provider_id=str(provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(provider_headers),
|
||||
body=build_execution_plan_body(
|
||||
provider_payload,
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
),
|
||||
stream=False,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
model_name=str(ctx.model or ""),
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
content_encoding=client_content_encoding,
|
||||
proxy=proxy_snapshot,
|
||||
tls_profile=envelope_tls_profile,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=int(config.http_connect_timeout * 1000),
|
||||
read_ms=int(config.http_read_timeout * 1000),
|
||||
write_ms=int(config.http_write_timeout * 1000),
|
||||
pool_ms=int(config.http_pool_timeout * 1000),
|
||||
total_ms=int(request_timeout_sync * 1000),
|
||||
),
|
||||
)
|
||||
|
||||
if config.executor_backend == "rust" and is_remote_contract_eligible(rust_plan):
|
||||
try:
|
||||
rust_result = await RustExecutorClient().execute_sync_json(rust_plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[{}] CLI Rust executor(sync->stream) 不可用,回退 Python 执行: {}",
|
||||
self.request_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
ctx.status_code = rust_result.status_code
|
||||
ctx.response_headers = dict(rust_result.headers)
|
||||
ctx.set_proxy_timing(ctx.response_headers)
|
||||
if envelope:
|
||||
envelope.on_http_status(
|
||||
base_url=ctx.selected_base_url,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
|
||||
request = httpx.Request("POST", url, headers=provider_headers)
|
||||
synthetic_content = rust_result.response_body_bytes
|
||||
if synthetic_content is None:
|
||||
synthetic_content = json.dumps(
|
||||
rust_result.response_json or {},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
synthetic_response = httpx.Response(
|
||||
ctx.status_code,
|
||||
request=request,
|
||||
headers=ctx.response_headers,
|
||||
content=synthetic_content,
|
||||
)
|
||||
|
||||
if ctx.status_code >= 400:
|
||||
error = httpx.HTTPStatusError(
|
||||
f"Upstream status error: {ctx.status_code}",
|
||||
request=request,
|
||||
response=synthetic_response,
|
||||
)
|
||||
error_body = ""
|
||||
try:
|
||||
if envelope and hasattr(envelope, "extract_error_text"):
|
||||
error_body = await envelope.extract_error_text(synthetic_response)
|
||||
else:
|
||||
error_body = (
|
||||
synthetic_response.text[:4000]
|
||||
if synthetic_response.text
|
||||
else ""
|
||||
)
|
||||
except Exception:
|
||||
error_body = (
|
||||
synthetic_response.text[:4000] if synthetic_response.text else ""
|
||||
)
|
||||
error.upstream_response = error_body # type: ignore[attr-defined]
|
||||
raise error
|
||||
|
||||
response_json = rust_result.response_json or {}
|
||||
if envelope:
|
||||
response_json = envelope.unwrap_response(response_json)
|
||||
envelope.postprocess_unwrapped_response(model=ctx.model, data=response_json)
|
||||
|
||||
if isinstance(response_json, dict) and provider_api_format:
|
||||
parser = get_parser_for_format(provider_api_format)
|
||||
if parser.is_error_response(response_json):
|
||||
parsed = parser.parse_response(response_json, 200)
|
||||
raise EmbeddedErrorException(
|
||||
provider_name=str(provider.name),
|
||||
error_code=parsed.embedded_status_code,
|
||||
error_message=parsed.error_message,
|
||||
error_status=parsed.error_type,
|
||||
)
|
||||
|
||||
if isinstance(response_json, dict):
|
||||
ctx.response_metadata = self._extract_response_metadata(response_json)
|
||||
|
||||
return self._streamify_sync_response(
|
||||
ctx=ctx,
|
||||
response_json=response_json if isinstance(response_json, dict) else {},
|
||||
client_api_format=client_api_format,
|
||||
provider_api_format=provider_api_format,
|
||||
)
|
||||
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
proxy_config=effective_proxy,
|
||||
@@ -488,60 +692,12 @@ class CliStreamMixin:
|
||||
if isinstance(response_json, dict):
|
||||
ctx.response_metadata = self._extract_response_metadata(response_json)
|
||||
|
||||
# Convert sync JSON -> InternalResponse, then InternalResponse -> client stream events.
|
||||
registry = get_format_converter_registry()
|
||||
src_norm = registry.get_normalizer(provider_api_format) if provider_api_format else None
|
||||
if src_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {provider_api_format}")
|
||||
|
||||
internal_resp = src_norm.response_to_internal(
|
||||
response_json if isinstance(response_json, dict) else {}
|
||||
return self._streamify_sync_response(
|
||||
ctx=ctx,
|
||||
response_json=response_json if isinstance(response_json, dict) else {},
|
||||
client_api_format=client_api_format,
|
||||
provider_api_format=provider_api_format,
|
||||
)
|
||||
internal_resp.model = str(ctx.model or internal_resp.model or "")
|
||||
if internal_resp.id:
|
||||
ctx.response_id = internal_resp.id
|
||||
|
||||
if internal_resp.usage:
|
||||
ctx.input_tokens = int(internal_resp.usage.input_tokens or 0)
|
||||
ctx.output_tokens = int(internal_resp.usage.output_tokens or 0)
|
||||
ctx.cached_tokens = int(internal_resp.usage.cache_read_tokens or 0)
|
||||
ctx.cache_creation_tokens = int(internal_resp.usage.cache_write_tokens or 0)
|
||||
|
||||
from src.core.api_format.conversion.stream_state import StreamState
|
||||
|
||||
tgt_norm = registry.get_normalizer(client_api_format) if client_api_format else None
|
||||
if tgt_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
|
||||
|
||||
state = StreamState(
|
||||
model=str(ctx.model or ""),
|
||||
message_id=str(ctx.response_id or ctx.request_id or self.request_id or ""),
|
||||
)
|
||||
output_state = {"first_yield": True, "streaming_updated": False}
|
||||
|
||||
async def _streamified() -> AsyncGenerator[bytes]:
|
||||
for ev in iter_internal_response_as_stream_events(internal_resp):
|
||||
converted_events = tgt_norm.stream_event_from_internal(ev, state)
|
||||
if not converted_events:
|
||||
continue
|
||||
self._record_converted_chunks(ctx, converted_events)
|
||||
for sse_line in _format_converted_events_to_sse(
|
||||
converted_events, client_api_format
|
||||
):
|
||||
if not sse_line:
|
||||
continue
|
||||
ctx.chunk_count += 1
|
||||
self._mark_first_output(ctx, output_state)
|
||||
yield (sse_line + "\n").encode("utf-8")
|
||||
|
||||
# OpenAI chat clients expect a final [DONE] marker.
|
||||
if str(client_api_format or "").strip().lower() == "openai:chat":
|
||||
ctx.chunk_count += 1
|
||||
self._mark_first_output(ctx, output_state)
|
||||
yield b"data: [DONE]\n\n"
|
||||
ctx.has_completion = True
|
||||
|
||||
return _streamified()
|
||||
|
||||
# 流式请求使用 stream_first_byte_timeout 作为首字节超时
|
||||
# 优先使用 Provider 配置,否则使用全局配置
|
||||
@@ -560,12 +716,106 @@ class CliStreamMixin:
|
||||
# 获取 HTTP 客户端(支持代理配置,Key 级别优先于 Provider 级别)
|
||||
# 使用连接池复用客户端,避免每次流式请求都新建 TCP/TLS 连接
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_stream_kwargs_async,
|
||||
resolve_delegate_config_async,
|
||||
from src.services.proxy_node.resolver import build_stream_kwargs_async
|
||||
|
||||
rust_plan = ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=str(
|
||||
getattr(candidate, "request_candidate_id", "") or getattr(candidate, "id", "") or ""
|
||||
)
|
||||
or None,
|
||||
provider_name=str(provider.name),
|
||||
provider_id=str(provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(provider_headers),
|
||||
body=build_execution_plan_body(
|
||||
provider_payload,
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
),
|
||||
stream=True,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
model_name=str(ctx.model or ""),
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
content_encoding=client_content_encoding,
|
||||
proxy=proxy_snapshot,
|
||||
tls_profile=envelope_tls_profile,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=int(config.http_connect_timeout * 1000),
|
||||
read_ms=int(config.http_read_timeout * 1000),
|
||||
write_ms=int(config.http_write_timeout * 1000),
|
||||
pool_ms=int(config.http_pool_timeout * 1000),
|
||||
total_ms=None,
|
||||
),
|
||||
)
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
if config.executor_backend == "rust" and is_remote_contract_eligible(rust_plan):
|
||||
try:
|
||||
rust_stream = await RustExecutorClient().execute_stream(rust_plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[{}] CLI Rust executor stream 不可用,回退 Python 执行: {}",
|
||||
self.request_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
ctx.status_code = rust_stream.status_code
|
||||
ctx.response_headers = dict(rust_stream.headers)
|
||||
ctx.set_proxy_timing(ctx.response_headers)
|
||||
|
||||
if envelope:
|
||||
envelope.on_http_status(
|
||||
base_url=ctx.selected_base_url,
|
||||
status_code=ctx.status_code,
|
||||
)
|
||||
|
||||
try:
|
||||
if ctx.status_code >= 400:
|
||||
error_chunks: list[bytes] = []
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
if chunk:
|
||||
error_chunks.append(chunk)
|
||||
if sum(len(item) for item in error_chunks) >= 4000:
|
||||
break
|
||||
error_body = b"".join(error_chunks)[:4000].decode(
|
||||
"utf-8",
|
||||
errors="replace",
|
||||
)
|
||||
request = httpx.Request("POST", url, headers=provider_headers)
|
||||
response = httpx.Response(
|
||||
ctx.status_code,
|
||||
request=request,
|
||||
headers=rust_stream.headers,
|
||||
content=b"".join(error_chunks),
|
||||
)
|
||||
error = httpx.HTTPStatusError(
|
||||
f"Upstream status error: {ctx.status_code}",
|
||||
request=request,
|
||||
response=response,
|
||||
)
|
||||
error.upstream_response = error_body # type: ignore[attr-defined]
|
||||
raise error
|
||||
|
||||
prefetched_chunks = await self._prefetch_and_check_embedded_error(
|
||||
rust_stream.byte_iterator,
|
||||
provider,
|
||||
endpoint,
|
||||
ctx,
|
||||
)
|
||||
except Exception:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
raise
|
||||
|
||||
return self._create_response_stream_with_prefetch(
|
||||
ctx,
|
||||
rust_stream.byte_iterator,
|
||||
rust_stream.response_ctx,
|
||||
prefetched_chunks,
|
||||
)
|
||||
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
proxy_config=effective_proxy,
|
||||
|
||||
@@ -31,6 +31,17 @@ from src.core.exceptions import (
|
||||
ThinkingSignatureException,
|
||||
)
|
||||
from src.core.logger import logger
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
build_execution_plan_body,
|
||||
is_remote_contract_eligible,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
from src.services.scheduling.aware_scheduler import ProviderCandidate
|
||||
from src.services.task.request_state import MutableRequestBodyState
|
||||
|
||||
@@ -42,6 +53,54 @@ if TYPE_CHECKING:
|
||||
class CliSyncMixin:
|
||||
"""同步处理相关方法的 Mixin"""
|
||||
|
||||
async def _aggregate_upstream_stream_sync_response(
|
||||
self: CliHandlerProtocol,
|
||||
*,
|
||||
body_bytes: bytes,
|
||||
provider_api_format: str,
|
||||
client_api_format: str,
|
||||
provider_name: str,
|
||||
provider_type: str,
|
||||
model: str,
|
||||
request_id: str,
|
||||
envelope: Any,
|
||||
) -> dict[str, Any]:
|
||||
registry = get_format_converter_registry()
|
||||
provider_parser = (
|
||||
get_parser_for_format(provider_api_format) if provider_api_format else None
|
||||
)
|
||||
|
||||
async def _byte_iter() -> Any:
|
||||
yield body_bytes
|
||||
|
||||
byte_iter = _byte_iter()
|
||||
if provider_type == "kiro" and envelope and envelope.force_stream_rewrite():
|
||||
from src.services.provider.adapters.kiro.eventstream_rewriter import (
|
||||
apply_kiro_stream_rewrite,
|
||||
)
|
||||
|
||||
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
|
||||
|
||||
internal_resp = await aggregate_upstream_stream_to_internal_response(
|
||||
byte_iter,
|
||||
provider_api_format=provider_api_format,
|
||||
provider_name=provider_name,
|
||||
model=model,
|
||||
request_id=request_id,
|
||||
envelope=envelope,
|
||||
provider_parser=provider_parser,
|
||||
)
|
||||
|
||||
tgt_norm = registry.get_normalizer(client_api_format) if client_api_format else None
|
||||
if tgt_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
|
||||
|
||||
response_json = tgt_norm.response_from_internal(
|
||||
internal_resp,
|
||||
requested_model=model,
|
||||
)
|
||||
return response_json if isinstance(response_json, dict) else {}
|
||||
|
||||
async def process_sync(
|
||||
self: CliHandlerProtocol,
|
||||
original_request_body: dict[str, Any],
|
||||
@@ -179,6 +238,7 @@ class CliSyncMixin:
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_post_kwargs_async,
|
||||
build_proxy_url_async,
|
||||
build_stream_kwargs_async,
|
||||
resolve_delegate_config_async,
|
||||
)
|
||||
@@ -188,6 +248,141 @@ class CliSyncMixin:
|
||||
request_timeout = provider.request_timeout or config.http_request_timeout
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(_effective_proxy)
|
||||
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
proxy_url: str | None = None
|
||||
if _effective_proxy and not is_tunnel_delegate:
|
||||
proxy_url = await build_proxy_url_async(_effective_proxy)
|
||||
|
||||
rust_plan = ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=str(
|
||||
getattr(candidate, "request_candidate_id", "")
|
||||
or getattr(candidate, "id", "")
|
||||
or ""
|
||||
)
|
||||
or None,
|
||||
provider_name=str(provider.name),
|
||||
provider_id=str(provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(key.id),
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(provider_headers),
|
||||
body=build_execution_plan_body(
|
||||
provider_payload,
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
),
|
||||
stream=upstream_is_stream,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
model_name=str(model or ""),
|
||||
content_type=str(provider_headers.get("content-type") or "").strip() or None,
|
||||
content_encoding=effective_client_content_encoding,
|
||||
proxy=ExecutionProxySnapshot.from_proxy_info(
|
||||
sync_proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if is_tunnel_delegate else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if is_tunnel_delegate
|
||||
else None
|
||||
),
|
||||
),
|
||||
tls_profile=envelope_tls_profile,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=int(config.http_connect_timeout * 1000),
|
||||
read_ms=int(config.http_read_timeout * 1000),
|
||||
write_ms=int(config.http_write_timeout * 1000),
|
||||
pool_ms=int(config.http_pool_timeout * 1000),
|
||||
total_ms=int(request_timeout * 1000),
|
||||
),
|
||||
)
|
||||
|
||||
if config.executor_backend == "rust" and is_remote_contract_eligible(rust_plan):
|
||||
try:
|
||||
rust_result = await RustExecutorClient().execute_sync_json(rust_plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[{}] CLI Rust executor 不可用,回退 Python 执行: {}",
|
||||
self.request_id,
|
||||
exc,
|
||||
)
|
||||
else:
|
||||
status_code = rust_result.status_code
|
||||
response_headers = dict(rust_result.headers)
|
||||
extract_proxy_timing(sync_proxy_info, response_headers)
|
||||
|
||||
if envelope:
|
||||
envelope.on_http_status(
|
||||
base_url=selected_base_url_cached,
|
||||
status_code=status_code,
|
||||
)
|
||||
|
||||
request = httpx.Request("POST", url, headers=provider_headers)
|
||||
synthetic_content = rust_result.response_body_bytes
|
||||
if synthetic_content is None:
|
||||
synthetic_content = json.dumps(
|
||||
rust_result.response_json or {},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
synthetic_response = httpx.Response(
|
||||
status_code,
|
||||
request=request,
|
||||
headers=response_headers,
|
||||
content=synthetic_content,
|
||||
)
|
||||
|
||||
if status_code >= 400:
|
||||
error = httpx.HTTPStatusError(
|
||||
f"Upstream status error: {status_code}",
|
||||
request=request,
|
||||
response=synthetic_response,
|
||||
)
|
||||
error_body = ""
|
||||
try:
|
||||
if envelope and hasattr(envelope, "extract_error_text"):
|
||||
error_body = await envelope.extract_error_text(synthetic_response)
|
||||
else:
|
||||
error_body = (
|
||||
synthetic_response.text[:4000]
|
||||
if synthetic_response.text
|
||||
else ""
|
||||
)
|
||||
except Exception:
|
||||
error_body = (
|
||||
synthetic_response.text[:4000] if synthetic_response.text else ""
|
||||
)
|
||||
error.upstream_response = error_body[:4000] # type: ignore[attr-defined]
|
||||
raise error
|
||||
|
||||
if upstream_is_stream:
|
||||
if rust_result.response_body_bytes is None:
|
||||
raise RustExecutorClientError(
|
||||
"Rust executor stream sync result must contain body bytes"
|
||||
)
|
||||
response_json = await self._aggregate_upstream_stream_sync_response(
|
||||
body_bytes=rust_result.response_body_bytes,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
provider_name=str(provider.name),
|
||||
provider_type=str(getattr(provider, "provider_type", "") or "").lower(),
|
||||
model=str(model or ""),
|
||||
request_id=str(self.request_id or ""),
|
||||
envelope=envelope,
|
||||
)
|
||||
response_metadata_result = self._extract_response_metadata(
|
||||
response_json or {}
|
||||
)
|
||||
return response_json if isinstance(response_json, dict) else {}
|
||||
|
||||
response_json = rust_result.response_json or {}
|
||||
if envelope:
|
||||
response_json = envelope.unwrap_response(response_json)
|
||||
envelope.postprocess_unwrapped_response(model=model, data=response_json)
|
||||
|
||||
response_metadata_result = self._extract_response_metadata(response_json)
|
||||
return response_json if isinstance(response_json, dict) else {}
|
||||
|
||||
http_client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
proxy_config=_effective_proxy,
|
||||
@@ -218,11 +413,6 @@ class CliSyncMixin:
|
||||
raise
|
||||
else:
|
||||
# Forced upstream streaming: aggregate SSE to a sync JSON response.
|
||||
registry = get_format_converter_registry()
|
||||
provider_parser = (
|
||||
get_parser_for_format(provider_api_format) if provider_api_format else None
|
||||
)
|
||||
|
||||
try:
|
||||
_stream_args = await build_stream_kwargs_async(
|
||||
delegate_cfg,
|
||||
@@ -246,41 +436,16 @@ class CliSyncMixin:
|
||||
)
|
||||
|
||||
stream_resp.raise_for_status()
|
||||
|
||||
byte_iter = stream_resp.aiter_bytes()
|
||||
_provider_type = str(getattr(provider, "provider_type", "") or "").lower()
|
||||
if (
|
||||
_provider_type == "kiro"
|
||||
and envelope
|
||||
and envelope.force_stream_rewrite()
|
||||
):
|
||||
from src.services.provider.adapters.kiro.eventstream_rewriter import (
|
||||
apply_kiro_stream_rewrite,
|
||||
)
|
||||
|
||||
byte_iter = apply_kiro_stream_rewrite(byte_iter, model=str(model or ""))
|
||||
|
||||
internal_resp = await aggregate_upstream_stream_to_internal_response(
|
||||
byte_iter,
|
||||
response_body = await stream_resp.aread()
|
||||
response_json = await self._aggregate_upstream_stream_sync_response(
|
||||
body_bytes=response_body,
|
||||
provider_api_format=provider_api_format,
|
||||
client_api_format=client_api_format,
|
||||
provider_name=str(provider.name),
|
||||
provider_type=str(getattr(provider, "provider_type", "") or "").lower(),
|
||||
model=str(model or ""),
|
||||
request_id=str(self.request_id or ""),
|
||||
envelope=envelope,
|
||||
provider_parser=provider_parser,
|
||||
)
|
||||
|
||||
tgt_norm = (
|
||||
registry.get_normalizer(client_api_format)
|
||||
if client_api_format
|
||||
else None
|
||||
)
|
||||
if tgt_norm is None:
|
||||
raise RuntimeError(f"未注册 Normalizer: {client_api_format}")
|
||||
|
||||
response_json = tgt_norm.response_from_internal(
|
||||
internal_resp,
|
||||
requested_model=model,
|
||||
)
|
||||
response_json = response_json if isinstance(response_json, dict) else {}
|
||||
|
||||
|
||||
@@ -17,6 +17,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import codecs
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
@@ -26,6 +27,7 @@ from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.api_format import (
|
||||
CORE_REDACT_HEADERS,
|
||||
merge_headers_with_protection,
|
||||
@@ -655,6 +657,16 @@ class HttpRequestExecutor:
|
||||
is_stream = True
|
||||
|
||||
try:
|
||||
rust_result = await self._execute_via_rust(
|
||||
request=request,
|
||||
request_id=request_id,
|
||||
is_stream=is_stream,
|
||||
start_time=start_time,
|
||||
effective_timeout=effective_timeout,
|
||||
)
|
||||
if rust_result is not None:
|
||||
return rust_result
|
||||
|
||||
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
||||
|
||||
# 统一通过 build_proxy_client_kwargs 构建(支持 tunnel 模式 + 普通代理 + 系统默认回退)
|
||||
@@ -758,112 +770,350 @@ class HttpRequestExecutor:
|
||||
async with client.stream(
|
||||
"POST", request.url, json=request.json_body, headers=request.headers
|
||||
) as response:
|
||||
headers = dict(response.headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
# 读取上限 16KB 用于 debug response_body;
|
||||
# error 字段仅取前 500 字符,避免日志/展示过长
|
||||
error_body = ""
|
||||
async for chunk in response.aiter_text():
|
||||
error_body += chunk
|
||||
if len(error_body) > 16384:
|
||||
break
|
||||
logger.debug(
|
||||
"[{}] check_endpoint | stream error | {}",
|
||||
request.api_format,
|
||||
error_body[:500],
|
||||
)
|
||||
return {
|
||||
"error": f"HTTP {response.status_code}: {error_body[:500]}",
|
||||
"status_code": response.status_code,
|
||||
"headers": headers,
|
||||
"response_body": error_body,
|
||||
}
|
||||
|
||||
# 收集 SSE 事件(兼容多种 API 格式)
|
||||
final_response: dict[str, Any] = {}
|
||||
collected_text = ""
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line or not line.startswith("data:"):
|
||||
continue
|
||||
|
||||
data_str = line[5:].strip()
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
event = json.loads(data_str)
|
||||
if "/v1internal:" in (request.url or ""):
|
||||
event = self._unwrap_gemini_cli_response_wrapper(event)
|
||||
event_type = event.get("type", "")
|
||||
|
||||
# OpenAI Responses API 事件
|
||||
if event_type == "response.output_text.delta":
|
||||
delta = event.get("delta", "")
|
||||
if isinstance(delta, str):
|
||||
collected_text += delta
|
||||
elif event_type == "response.completed":
|
||||
final_response = event.get("response", {})
|
||||
break
|
||||
|
||||
# OpenAI Chat Completions 格式
|
||||
elif "choices" in event:
|
||||
for choice in event.get("choices", []):
|
||||
delta = choice.get("delta", {})
|
||||
content = delta.get("content")
|
||||
if content:
|
||||
collected_text += content
|
||||
if choice.get("finish_reason"):
|
||||
final_response = event
|
||||
break
|
||||
|
||||
# Claude Messages API 格式
|
||||
elif event_type == "content_block_delta":
|
||||
delta = event.get("delta", {})
|
||||
text = delta.get("text", "")
|
||||
if text:
|
||||
collected_text += text
|
||||
elif event_type == "message_stop":
|
||||
break
|
||||
|
||||
# Gemini SSE 格式
|
||||
elif "candidates" in event:
|
||||
for candidate in event.get("candidates", []):
|
||||
content = candidate.get("content", {})
|
||||
for part in content.get("parts", []):
|
||||
text = part.get("text", "")
|
||||
if text:
|
||||
collected_text += text
|
||||
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# 如果没有收到最终响应事件,构建一个基本响应
|
||||
if not final_response:
|
||||
final_response = {
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": collected_text}],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
"[{}] check_endpoint | stream completed | text_length={}",
|
||||
request.api_format,
|
||||
len(collected_text),
|
||||
return await self._consume_stream_lines(
|
||||
line_iter=response.aiter_lines(),
|
||||
request=request,
|
||||
status_code=response.status_code,
|
||||
headers=dict(response.headers),
|
||||
error_text_iter=response.aiter_text(),
|
||||
)
|
||||
|
||||
return {"final_response": final_response, "headers": headers}
|
||||
|
||||
except Exception as e:
|
||||
logger.warning("[{}] check_endpoint | stream error | {}", request.api_format, e)
|
||||
return {"error": str(e), "status_code": 500, "headers": {}, "response_body": None}
|
||||
|
||||
async def _execute_via_rust(
|
||||
self,
|
||||
*,
|
||||
request: EndpointCheckRequest,
|
||||
request_id: str,
|
||||
is_stream: bool,
|
||||
start_time: float,
|
||||
effective_timeout: float,
|
||||
) -> EndpointCheckResult | None:
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
proxy_snapshot = await self._build_rust_proxy_snapshot(request.proxy_config)
|
||||
plan = ExecutionPlan(
|
||||
request_id=f"endpoint-check-{request_id}",
|
||||
candidate_id=None,
|
||||
provider_name=str(request.provider_name or request.api_format or ""),
|
||||
provider_id=str(request.provider_id or ""),
|
||||
endpoint_id="",
|
||||
key_id=str(request.api_key_id or ""),
|
||||
method="POST",
|
||||
url=request.url,
|
||||
headers=dict(request.headers),
|
||||
body=build_execution_plan_body(request.json_body, content_type="application/json"),
|
||||
stream=is_stream,
|
||||
provider_api_format=str(request.api_format or ""),
|
||||
client_api_format=str(request.api_format or ""),
|
||||
model_name=str(request.model_name or ""),
|
||||
content_type="application/json",
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=min(int(effective_timeout * 1000), 30_000),
|
||||
read_ms=(None if is_stream else int(effective_timeout * 1000)),
|
||||
write_ms=int(effective_timeout * 1000),
|
||||
pool_ms=min(int(effective_timeout * 1000), 30_000),
|
||||
total_ms=(None if is_stream else int(effective_timeout * 1000)),
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
if is_stream:
|
||||
rust_stream = await RustExecutorClient().execute_stream(plan)
|
||||
try:
|
||||
if rust_stream.status_code >= 400:
|
||||
error_bytes = await self._read_limited_bytes(rust_stream.byte_iterator)
|
||||
return await ErrorHandler.handle_error(
|
||||
httpx.HTTPStatusError(
|
||||
message=f"HTTP {rust_stream.status_code}",
|
||||
request=httpx.Request("POST", request.url, headers=request.headers),
|
||||
response=httpx.Response(
|
||||
rust_stream.status_code,
|
||||
request=httpx.Request(
|
||||
"POST", request.url, headers=request.headers
|
||||
),
|
||||
headers=rust_stream.headers,
|
||||
content=error_bytes,
|
||||
),
|
||||
),
|
||||
request,
|
||||
)
|
||||
|
||||
stream_result = await self._consume_stream_lines(
|
||||
line_iter=self._aiter_lines_from_bytes(rust_stream.byte_iterator),
|
||||
request=request,
|
||||
status_code=rust_stream.status_code,
|
||||
headers=dict(rust_stream.headers),
|
||||
)
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
response_time_ms = int((time.time() - start_time) * 1000)
|
||||
if stream_result.get("error"):
|
||||
return EndpointCheckResult(
|
||||
status_code=stream_result.get("status_code", 500),
|
||||
headers=stream_result.get("headers", {}),
|
||||
response_time_ms=response_time_ms,
|
||||
request_id=request_id,
|
||||
response_data=None,
|
||||
error_message=stream_result.get("error"),
|
||||
raw_response_body=stream_result.get("response_body"),
|
||||
)
|
||||
return EndpointCheckResult(
|
||||
status_code=rust_stream.status_code,
|
||||
headers=stream_result.get("headers", {}),
|
||||
response_time_ms=response_time_ms,
|
||||
request_id=request_id,
|
||||
response_data=stream_result.get("final_response"),
|
||||
raw_response_body=stream_result.get("final_response"),
|
||||
)
|
||||
|
||||
result = await RustExecutorClient().execute_sync_json(plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError, ValueError) as exc:
|
||||
logger.warning(
|
||||
"[{}] endpoint check rust fallback | provider={} model={} error={}",
|
||||
request.api_format,
|
||||
request.provider_name,
|
||||
request.model_name,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
response = self._build_httpx_response_from_rust(
|
||||
method="POST",
|
||||
url=request.url,
|
||||
request_headers=request.headers,
|
||||
status_code=result.status_code,
|
||||
headers=result.headers,
|
||||
response_json=result.response_json,
|
||||
response_body_bytes=result.response_body_bytes,
|
||||
)
|
||||
if result.status_code >= 400:
|
||||
return await ErrorHandler.handle_error(
|
||||
httpx.HTTPStatusError(
|
||||
message=f"HTTP {result.status_code}",
|
||||
request=response.request,
|
||||
response=response,
|
||||
),
|
||||
request,
|
||||
)
|
||||
|
||||
response_time_ms = int((time.time() - start_time) * 1000)
|
||||
response_data: dict[str, Any] | None
|
||||
if isinstance(result.response_json, dict):
|
||||
response_data = result.response_json
|
||||
else:
|
||||
try:
|
||||
parsed_body = response.json()
|
||||
response_data = parsed_body if isinstance(parsed_body, dict) else None
|
||||
except Exception:
|
||||
response_data = None
|
||||
|
||||
return EndpointCheckResult(
|
||||
status_code=result.status_code,
|
||||
headers=dict(result.headers),
|
||||
response_time_ms=response_time_ms,
|
||||
request_id=request_id,
|
||||
response_data=response_data,
|
||||
raw_response_body=response_data if response_data is not None else response.text,
|
||||
)
|
||||
|
||||
async def _build_rust_proxy_snapshot(
|
||||
self,
|
||||
proxy_config: dict[str, Any] | None,
|
||||
) -> Any:
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
|
||||
effective_proxy = proxy_config
|
||||
if not effective_proxy or not effective_proxy.get("enabled", True):
|
||||
effective_proxy = await get_system_proxy_config_async()
|
||||
if not effective_proxy:
|
||||
return None
|
||||
|
||||
try:
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
|
||||
proxy_url = await build_proxy_url_async(effective_proxy)
|
||||
proxy_info = await resolve_proxy_info_async(effective_proxy)
|
||||
return ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("endpoint check proxy snapshot build failed: {}", exc)
|
||||
return None
|
||||
|
||||
async def _consume_stream_lines(
|
||||
self,
|
||||
*,
|
||||
line_iter: Any,
|
||||
request: EndpointCheckRequest,
|
||||
status_code: int,
|
||||
headers: dict[str, str],
|
||||
error_text_iter: Any | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if status_code != 200:
|
||||
error_body = ""
|
||||
if error_text_iter is not None:
|
||||
async for chunk in error_text_iter:
|
||||
error_body += chunk
|
||||
if len(error_body) > 16384:
|
||||
break
|
||||
logger.debug(
|
||||
"[{}] check_endpoint | stream error | {}",
|
||||
request.api_format,
|
||||
error_body[:500],
|
||||
)
|
||||
return {
|
||||
"error": f"HTTP {status_code}: {error_body[:500]}",
|
||||
"status_code": status_code,
|
||||
"headers": headers,
|
||||
"response_body": error_body,
|
||||
}
|
||||
|
||||
final_response: dict[str, Any] = {}
|
||||
collected_text = ""
|
||||
|
||||
async for line in line_iter:
|
||||
if not line or not line.startswith("data:"):
|
||||
continue
|
||||
|
||||
data_str = line[5:].strip()
|
||||
if data_str == "[DONE]":
|
||||
break
|
||||
|
||||
try:
|
||||
event = json.loads(data_str)
|
||||
if "/v1internal:" in (request.url or ""):
|
||||
event = self._unwrap_gemini_cli_response_wrapper(event)
|
||||
event_type = event.get("type", "")
|
||||
|
||||
if event_type == "response.output_text.delta":
|
||||
delta = event.get("delta", "")
|
||||
if isinstance(delta, str):
|
||||
collected_text += delta
|
||||
elif event_type == "response.completed":
|
||||
final_response = event.get("response", {})
|
||||
break
|
||||
elif "choices" in event:
|
||||
for choice in event.get("choices", []):
|
||||
delta = choice.get("delta", {})
|
||||
content = delta.get("content")
|
||||
if content:
|
||||
collected_text += content
|
||||
if choice.get("finish_reason"):
|
||||
final_response = event
|
||||
break
|
||||
elif event_type == "content_block_delta":
|
||||
delta = event.get("delta", {})
|
||||
text = delta.get("text", "")
|
||||
if text:
|
||||
collected_text += text
|
||||
elif event_type == "message_stop":
|
||||
break
|
||||
elif "candidates" in event:
|
||||
for candidate in event.get("candidates", []):
|
||||
content = candidate.get("content", {})
|
||||
for part in content.get("parts", []):
|
||||
text = part.get("text", "")
|
||||
if text:
|
||||
collected_text += text
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
if not final_response:
|
||||
final_response = {
|
||||
"status": "completed",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": collected_text}],
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
logger.debug(
|
||||
"[{}] check_endpoint | stream completed | text_length={}",
|
||||
request.api_format,
|
||||
len(collected_text),
|
||||
)
|
||||
|
||||
return {"final_response": final_response, "headers": headers}
|
||||
|
||||
async def _aiter_lines_from_bytes(self, byte_iter: Any) -> Any:
|
||||
decoder = codecs.getincrementaldecoder("utf-8")()
|
||||
buffer = ""
|
||||
async for chunk in byte_iter:
|
||||
buffer += decoder.decode(chunk)
|
||||
while "\n" in buffer:
|
||||
line, buffer = buffer.split("\n", 1)
|
||||
yield line.rstrip("\r")
|
||||
buffer += decoder.decode(b"", final=True)
|
||||
if buffer:
|
||||
yield buffer.rstrip("\r")
|
||||
|
||||
async def _read_limited_bytes(self, byte_iter: Any, limit: int = 16_384) -> bytes:
|
||||
chunks: list[bytes] = []
|
||||
total = 0
|
||||
async for chunk in byte_iter:
|
||||
if not chunk:
|
||||
continue
|
||||
chunks.append(chunk)
|
||||
total += len(chunk)
|
||||
if total >= limit:
|
||||
break
|
||||
return b"".join(chunks)
|
||||
|
||||
@staticmethod
|
||||
def _build_httpx_response_from_rust(
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
request_headers: dict[str, str],
|
||||
status_code: int,
|
||||
headers: dict[str, str],
|
||||
response_json: Any,
|
||||
response_body_bytes: bytes | None,
|
||||
) -> httpx.Response:
|
||||
if response_json is not None:
|
||||
content = json.dumps(response_json, ensure_ascii=False).encode("utf-8")
|
||||
else:
|
||||
content = response_body_bytes or b""
|
||||
return httpx.Response(
|
||||
status_code=status_code,
|
||||
request=httpx.Request(method, url, headers=request_headers),
|
||||
headers=headers,
|
||||
content=content,
|
||||
)
|
||||
|
||||
|
||||
class UsageCalculator:
|
||||
"""用量计算器 - 专门负责Token计数和费用计算"""
|
||||
|
||||
@@ -6,10 +6,12 @@ Video Handler 基类
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException, Request
|
||||
from fastapi.responses import JSONResponse, Response, StreamingResponse
|
||||
from sqlalchemy.orm import Session
|
||||
@@ -28,8 +30,6 @@ from src.services.billing.rule_service import BillingRuleLookupResult
|
||||
from src.services.scheduling.aware_scheduler import ProviderCandidate
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import httpx
|
||||
|
||||
from src.services.candidate.submit import SubmitOutcome
|
||||
|
||||
|
||||
@@ -167,6 +167,109 @@ class VideoHandlerBase(ABC):
|
||||
content={"error": fallback_payload},
|
||||
)
|
||||
|
||||
async def _try_rust_sync_http_response(
|
||||
self,
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: Any = None,
|
||||
provider_name: str | None = None,
|
||||
provider_id: str | None = None,
|
||||
endpoint_id: str | None = None,
|
||||
key_id: str | None = None,
|
||||
provider_api_format: str | None = None,
|
||||
client_api_format: str | None = None,
|
||||
model_name: str | None = None,
|
||||
content_type: str | None = None,
|
||||
content_encoding: str | None = None,
|
||||
proxy: Any = None,
|
||||
tls_profile: str | None = None,
|
||||
request_timeout_ms: int = 300_000,
|
||||
connect_timeout_ms: int = 30_000,
|
||||
pool_timeout_ms: int = 30_000,
|
||||
log_label: str = "VideoRequest",
|
||||
) -> httpx.Response | None:
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
request_headers = dict(headers)
|
||||
if (
|
||||
body is not None
|
||||
and content_type
|
||||
and not any(str(key).lower() == "content-type" for key in request_headers)
|
||||
):
|
||||
request_headers["content-type"] = content_type
|
||||
|
||||
plan = ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=None,
|
||||
provider_name=str(provider_name or ""),
|
||||
provider_id=str(provider_id or ""),
|
||||
endpoint_id=str(endpoint_id or ""),
|
||||
key_id=str(key_id or ""),
|
||||
method=str(method or "POST").upper(),
|
||||
url=url,
|
||||
headers=request_headers,
|
||||
body=build_execution_plan_body(body, content_type=content_type),
|
||||
stream=False,
|
||||
provider_api_format=str(provider_api_format or self.FORMAT_ID),
|
||||
client_api_format=str(client_api_format or self.FORMAT_ID),
|
||||
model_name=str(model_name or ""),
|
||||
content_type=content_type,
|
||||
content_encoding=content_encoding,
|
||||
proxy=proxy,
|
||||
tls_profile=tls_profile,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=connect_timeout_ms,
|
||||
read_ms=request_timeout_ms,
|
||||
write_ms=request_timeout_ms,
|
||||
pool_ms=pool_timeout_ms,
|
||||
total_ms=request_timeout_ms,
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
rust_result = await RustExecutorClient().execute_sync_json(plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[{}] Rust executor unavailable request_id={} method={} url={}: {}",
|
||||
log_label,
|
||||
self.request_id,
|
||||
method,
|
||||
url,
|
||||
sanitize_error_message(str(exc)),
|
||||
)
|
||||
return None
|
||||
|
||||
response_headers = dict(rust_result.headers)
|
||||
if rust_result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(rust_result.response_json, ensure_ascii=False).encode(
|
||||
"utf-8"
|
||||
)
|
||||
elif rust_result.response_body_bytes is not None:
|
||||
response_body = rust_result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
return httpx.Response(
|
||||
status_code=rust_result.status_code,
|
||||
request=httpx.Request(str(method or "POST").upper(), url, headers=request_headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
def _format_error_payload(self, error: dict[str, Any], status_code: int) -> dict[str, Any]:
|
||||
"""
|
||||
格式化错误负载,子类可重写以匹配特定 API 格式
|
||||
|
||||
@@ -183,6 +183,25 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
original_body=original_request_body,
|
||||
)
|
||||
|
||||
rust_response = await self._try_rust_sync_http_response(
|
||||
method="POST",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=converted_body,
|
||||
provider_name=str(candidate.provider.name),
|
||||
provider_id=str(candidate.provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(_key.id),
|
||||
provider_api_format=provider_format,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=internal_request.model,
|
||||
content_type=str(headers.get("content-type") or "").strip()
|
||||
or "application/json",
|
||||
log_label="GeminiVideoCreate",
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
return await client.post(upstream_url, headers=headers, json=converted_body)
|
||||
else:
|
||||
@@ -206,6 +225,25 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
body=request_body,
|
||||
original_body=original_request_body,
|
||||
)
|
||||
rust_response = await self._try_rust_sync_http_response(
|
||||
method="POST",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=request_body,
|
||||
provider_name=str(candidate.provider.name),
|
||||
provider_id=str(candidate.provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(_key.id),
|
||||
provider_api_format=provider_format,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=internal_request.model,
|
||||
content_type=str(headers.get("content-type") or "").strip()
|
||||
or "application/json",
|
||||
log_label="GeminiVideoCreate",
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
return await client.post(upstream_url, headers=headers, json=request_body)
|
||||
|
||||
@@ -470,6 +508,17 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
|
||||
# 代理下载而非直接重定向,避免暴露上游存储 URL
|
||||
# 使用 httpx 支持重定向(Gemini 视频 URL 会重定向到实际存储位置)
|
||||
rust_response = await self._try_rust_download_stream(
|
||||
url=task.video_url,
|
||||
headers=download_headers,
|
||||
task_id=str(task.id),
|
||||
endpoint=endpoint,
|
||||
key=key,
|
||||
model_name=str(getattr(task, "model", "") or "") or None,
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
import httpx
|
||||
|
||||
try:
|
||||
@@ -516,6 +565,125 @@ class GeminiVeoHandler(VideoHandlerBase):
|
||||
media_type=response.headers.get("content-type", "video/mp4"),
|
||||
)
|
||||
|
||||
async def _try_rust_download_stream(
|
||||
self,
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
task_id: str,
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
model_name: str | None = None,
|
||||
) -> Response | StreamingResponse | None:
|
||||
import httpx
|
||||
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_effective_proxy,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
effective_proxy = resolve_effective_proxy(
|
||||
resolve_provider_proxy(endpoint=endpoint, key=key),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
if not effective_proxy or not effective_proxy.get("enabled", True):
|
||||
effective_proxy = await get_system_proxy_config_async()
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
|
||||
proxy_url = await build_proxy_url_async(effective_proxy)
|
||||
|
||||
proxy_info = await resolve_proxy_info_async(effective_proxy)
|
||||
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
rust_stream = await RustExecutorClient().execute_stream(
|
||||
ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=None,
|
||||
provider_name="gemini",
|
||||
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or ""),
|
||||
key_id=str(getattr(key, "id", "") or ""),
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=dict(headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=True,
|
||||
provider_api_format=self.FORMAT_ID,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=str(model_name or ""),
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=300_000,
|
||||
write_ms=300_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=None,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (RustExecutorClientError, httpx.HTTPError, ValueError) as exc:
|
||||
logger.warning(
|
||||
"[VideoDownload] Rust executor unavailable task={} url={}: {}",
|
||||
task_id,
|
||||
url,
|
||||
sanitize_error_message(str(exc)),
|
||||
)
|
||||
return None
|
||||
|
||||
safe_headers = {
|
||||
k: v for k, v in rust_stream.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS
|
||||
}
|
||||
|
||||
if rust_stream.status_code >= 400:
|
||||
try:
|
||||
async for _ in rust_stream.byte_iterator:
|
||||
pass
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
raise HTTPException(status_code=rust_stream.status_code, detail="Upstream error")
|
||||
|
||||
async def _iter_bytes() -> AsyncIterator[bytes]:
|
||||
try:
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
_iter_bytes(),
|
||||
status_code=rust_stream.status_code,
|
||||
headers=safe_headers,
|
||||
media_type=safe_headers.get("content-type", "video/mp4"),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@@ -43,6 +43,15 @@ from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ApiKey, ProviderAPIKey, ProviderEndpoint, User, VideoTask
|
||||
from src.services.billing.rule_service import BillingRuleLookupResult, BillingRuleService
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
from src.services.scheduling.aware_scheduler import ProviderCandidate
|
||||
from src.services.usage.service import UsageService
|
||||
|
||||
@@ -180,6 +189,25 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
original_body=original_request_body,
|
||||
)
|
||||
|
||||
rust_response = await self._try_rust_sync_http_response(
|
||||
method="POST",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=converted_body,
|
||||
provider_name=str(candidate.provider.name),
|
||||
provider_id=str(candidate.provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(_provider_key.id),
|
||||
provider_api_format=provider_format,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=internal_request.model,
|
||||
content_type=str(headers.get("content-type") or "").strip()
|
||||
or "application/json",
|
||||
log_label="OpenAIVideoCreate",
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
return await client.post(upstream_url, headers=headers, json=converted_body)
|
||||
else:
|
||||
@@ -199,6 +227,25 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
body=request_body,
|
||||
original_body=original_request_body,
|
||||
)
|
||||
rust_response = await self._try_rust_sync_http_response(
|
||||
method="POST",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=request_body,
|
||||
provider_name=str(candidate.provider.name),
|
||||
provider_id=str(candidate.provider.id),
|
||||
endpoint_id=str(endpoint.id),
|
||||
key_id=str(_provider_key.id),
|
||||
provider_api_format=provider_format,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=internal_request.model,
|
||||
content_type=str(headers.get("content-type") or "").strip()
|
||||
or "application/json",
|
||||
log_label="OpenAIVideoCreate",
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
return await client.post(upstream_url, headers=headers, json=request_body)
|
||||
|
||||
@@ -478,8 +525,23 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
)
|
||||
headers = self._build_upstream_headers(original_headers, upstream_key, endpoint)
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.delete(upstream_url, headers=headers)
|
||||
response = await self._try_rust_sync_http_response(
|
||||
method="DELETE",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=None,
|
||||
provider_name="openai",
|
||||
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or ""),
|
||||
key_id=str(getattr(key, "id", "") or ""),
|
||||
provider_api_format=self.FORMAT_ID,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=str(getattr(task, "model", "") or ""),
|
||||
log_label="OpenAIVideoDelete",
|
||||
)
|
||||
if response is None:
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.delete(upstream_url, headers=headers)
|
||||
if response.status_code >= 400 and response.status_code != 404:
|
||||
# 404 表示上游已删除,不算错误
|
||||
return self._build_error_response(response)
|
||||
@@ -549,8 +611,24 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
original_body=original_request_body,
|
||||
)
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.post(upstream_url, headers=headers, json=request_body)
|
||||
response = await self._try_rust_sync_http_response(
|
||||
method="POST",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=request_body,
|
||||
provider_name="openai",
|
||||
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or ""),
|
||||
key_id=str(getattr(key, "id", "") or ""),
|
||||
provider_api_format=self.FORMAT_ID,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=str(getattr(original_task, "model", "") or ""),
|
||||
content_type=str(headers.get("content-type") or "").strip() or "application/json",
|
||||
log_label="OpenAIVideoRemix",
|
||||
)
|
||||
if response is None:
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.post(upstream_url, headers=headers, json=request_body)
|
||||
|
||||
if response.status_code >= 400:
|
||||
return self._build_error_response(response)
|
||||
@@ -669,6 +747,16 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
task_id,
|
||||
task.video_url,
|
||||
)
|
||||
rust_response = await self._try_rust_download_stream(
|
||||
url=task.video_url,
|
||||
headers={},
|
||||
task_id=task_id,
|
||||
model_name=str(getattr(task, "model", "") or "") or None,
|
||||
default_media_type="video/mp4",
|
||||
default_error_message="Video not available",
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
return await self._proxy_direct_url(task.video_url, task_id)
|
||||
|
||||
if not task.external_task_id:
|
||||
@@ -697,6 +785,20 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
task_id,
|
||||
task.external_task_id,
|
||||
)
|
||||
rust_response = await self._try_rust_download_stream(
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
task_id=task_id,
|
||||
provider_id=str(getattr(endpoint, "provider_id", "") or "") or None,
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or "") or None,
|
||||
key_id=str(getattr(key, "id", "") or "") or None,
|
||||
model_name=str(getattr(task, "model", "") or "") or None,
|
||||
default_media_type="application/octet-stream",
|
||||
default_error_message="Upstream connection failed",
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
try:
|
||||
# 使用 httpx 的 stream 方法并正确管理上下文
|
||||
# 视频下载可能较大,设置 5 分钟超时
|
||||
@@ -811,6 +913,111 @@ class OpenAIVideoHandler(VideoHandlerBase):
|
||||
media_type=response.headers.get("content-type", "video/mp4"),
|
||||
)
|
||||
|
||||
async def _try_rust_download_stream(
|
||||
self,
|
||||
*,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
task_id: str,
|
||||
provider_id: str | None = None,
|
||||
endpoint_id: str | None = None,
|
||||
key_id: str | None = None,
|
||||
model_name: str | None = None,
|
||||
default_media_type: str,
|
||||
default_error_message: str,
|
||||
) -> Response | StreamingResponse | None:
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
plan = ExecutionPlan(
|
||||
request_id=str(self.request_id or ""),
|
||||
candidate_id=None,
|
||||
provider_name="openai",
|
||||
provider_id=str(provider_id or ""),
|
||||
endpoint_id=str(endpoint_id or ""),
|
||||
key_id=str(key_id or ""),
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=dict(headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=True,
|
||||
provider_api_format=self.FORMAT_ID,
|
||||
client_api_format=self.FORMAT_ID,
|
||||
model_name=str(model_name or "") or "",
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=300_000,
|
||||
write_ms=300_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=None,
|
||||
),
|
||||
)
|
||||
|
||||
try:
|
||||
rust_stream = await RustExecutorClient().execute_stream(plan)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[VideoDownload] Rust executor unavailable task={} url={}: {}",
|
||||
task_id,
|
||||
url,
|
||||
sanitize_error_message(str(exc)),
|
||||
)
|
||||
return None
|
||||
|
||||
safe_headers = {
|
||||
k: v for k, v in rust_stream.headers.items() if k.lower() not in HOP_BY_HOP_HEADERS
|
||||
}
|
||||
|
||||
if rust_stream.status_code >= 400:
|
||||
error_chunks: list[bytes] = []
|
||||
try:
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
if chunk:
|
||||
error_chunks.append(chunk)
|
||||
if sum(len(item) for item in error_chunks) >= 16_384:
|
||||
break
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
error_body = b"".join(error_chunks)[:16_384]
|
||||
content_type = str(safe_headers.get("content-type") or "").lower()
|
||||
if "application/json" in content_type:
|
||||
try:
|
||||
data = json.loads(error_body)
|
||||
if isinstance(data, dict) and isinstance(data.get("error"), dict):
|
||||
if "message" in data["error"]:
|
||||
data["error"]["message"] = sanitize_error_message(
|
||||
str(data["error"]["message"])
|
||||
)
|
||||
return JSONResponse(status_code=rust_stream.status_code, content=data)
|
||||
except json.JSONDecodeError:
|
||||
pass
|
||||
|
||||
message = sanitize_error_message(error_body.decode(errors="ignore"))
|
||||
return JSONResponse(
|
||||
status_code=rust_stream.status_code,
|
||||
content={
|
||||
"error": {
|
||||
"type": "upstream_error",
|
||||
"message": message or default_error_message,
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
async def _iter_bytes() -> AsyncIterator[bytes]:
|
||||
try:
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
_iter_bytes(),
|
||||
status_code=rust_stream.status_code,
|
||||
headers=safe_headers,
|
||||
media_type=safe_headers.get("content-type", default_media_type),
|
||||
)
|
||||
|
||||
def _build_upstream_url(self, base_url: str | None, suffix: str | None = None) -> str:
|
||||
base = (base_url or self.DEFAULT_BASE_URL).rstrip("/")
|
||||
if base.endswith("/v1"):
|
||||
|
||||
@@ -1,3 +1,10 @@
|
||||
from .hub import router
|
||||
from fastapi import APIRouter
|
||||
|
||||
from .gateway import router as gateway_router
|
||||
from .hub import router as hub_router
|
||||
|
||||
router = APIRouter()
|
||||
router.include_router(hub_router)
|
||||
router.include_router(gateway_router)
|
||||
|
||||
__all__ = ["router"]
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
|
||||
from fastapi import HTTPException, Request
|
||||
|
||||
|
||||
def ensure_loopback(request: Request) -> None:
|
||||
host = request.client.host if request.client else ""
|
||||
try:
|
||||
if not ipaddress.ip_address(host).is_loopback:
|
||||
raise ValueError(host)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=403, detail="loopback access only") from exc
|
||||
File diff suppressed because it is too large
Load Diff
+4
-12
@@ -1,7 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ipaddress
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request
|
||||
@@ -9,6 +8,8 @@ from pydantic import BaseModel, Field
|
||||
|
||||
from src.services.proxy_node.service import ProxyNodeService, build_heartbeat_ack
|
||||
|
||||
from .common import ensure_loopback
|
||||
|
||||
router = APIRouter(prefix="/api/internal/hub", tags=["Internal - Hub"], include_in_schema=False)
|
||||
|
||||
|
||||
@@ -31,18 +32,9 @@ class HubNodeStatusRequest(BaseModel):
|
||||
conn_count: int = Field(0, ge=0)
|
||||
|
||||
|
||||
def _ensure_loopback(request: Request) -> None:
|
||||
host = request.client.host if request.client else ""
|
||||
try:
|
||||
if not ipaddress.ip_address(host).is_loopback:
|
||||
raise ValueError(host)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=403, detail="loopback access only") from exc
|
||||
|
||||
|
||||
@router.post("/heartbeat")
|
||||
async def hub_heartbeat(request: Request, payload: HubHeartbeatRequest) -> dict[str, Any]:
|
||||
_ensure_loopback(request)
|
||||
ensure_loopback(request)
|
||||
|
||||
def _sync_apply() -> dict[str, Any]:
|
||||
from src.database import create_session
|
||||
@@ -74,7 +66,7 @@ async def hub_heartbeat(request: Request, payload: HubHeartbeatRequest) -> dict[
|
||||
|
||||
@router.post("/node-status")
|
||||
async def hub_node_status(request: Request, payload: HubNodeStatusRequest) -> dict[str, Any]:
|
||||
_ensure_loopback(request)
|
||||
ensure_loopback(request)
|
||||
|
||||
def _sync_apply() -> dict[str, Any]:
|
||||
from src.database import create_session
|
||||
|
||||
+442
-105
@@ -21,12 +21,15 @@ https://ai.google.dev/api/files
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
from typing import Any
|
||||
from urllib.parse import urlencode
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import APIRouter, HTTPException, Request, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
@@ -37,6 +40,7 @@ from src.database import create_session
|
||||
from src.models.database import ApiKey, GlobalModel, Model, Provider, ProviderEndpoint, User
|
||||
from src.services.auth.service import AuthService
|
||||
from src.services.gemini_files_mapping import delete_file_key_mapping, store_file_key_mapping
|
||||
from src.services.provider.provider_context import resolve_provider_proxy
|
||||
from src.services.provider.transport import redact_url_for_log
|
||||
from src.services.scheduling.aware_scheduler import CacheAwareScheduler, ProviderCandidate
|
||||
from src.services.usage.service import UsageService
|
||||
@@ -50,6 +54,13 @@ class UpstreamContext:
|
||||
base_url: str
|
||||
file_key_id: str
|
||||
user_id: str
|
||||
provider_id: str
|
||||
endpoint_id: str
|
||||
provider_proxy: dict[str, Any] | None = None
|
||||
key_proxy: dict[str, Any] | None = None
|
||||
proxy_config: dict[str, Any] | None = None
|
||||
delegate_config: dict[str, Any] | None = None
|
||||
proxy_snapshot: Any = None
|
||||
|
||||
|
||||
router = APIRouter(tags=["Gemini Files API"])
|
||||
@@ -252,7 +263,7 @@ async def _select_provider_candidate(
|
||||
async def _resolve_upstream_context(
|
||||
request: Request,
|
||||
db: Session,
|
||||
) -> tuple[str, str, str, str]:
|
||||
) -> UpstreamContext:
|
||||
"""
|
||||
解析上游 Key 与 Base URL(需要外部提供 db session)
|
||||
|
||||
@@ -263,7 +274,7 @@ async def _resolve_upstream_context(
|
||||
db: 数据库会话
|
||||
|
||||
Returns:
|
||||
(upstream_key, base_url, key_id, user_id)
|
||||
UpstreamContext
|
||||
"""
|
||||
|
||||
client_key = _extract_gemini_api_key(request)
|
||||
@@ -341,7 +352,53 @@ async def _resolve_upstream_context(
|
||||
)
|
||||
|
||||
base_url = candidate.endpoint.base_url or GEMINI_FILES_BASE_URL
|
||||
return upstream_key, base_url, str(candidate.key.id), str(user.id)
|
||||
return UpstreamContext(
|
||||
upstream_key,
|
||||
base_url,
|
||||
str(candidate.key.id),
|
||||
str(user.id),
|
||||
str(candidate.provider.id),
|
||||
str(candidate.endpoint.id),
|
||||
provider_proxy=resolve_provider_proxy(endpoint=candidate.endpoint, key=candidate.key),
|
||||
key_proxy=(
|
||||
candidate.key.proxy if isinstance(getattr(candidate.key, "proxy", None), dict) else None
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
async def _enrich_upstream_context_proxy(ctx: UpstreamContext) -> UpstreamContext:
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_effective_proxy,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
|
||||
effective_proxy = resolve_effective_proxy(ctx.provider_proxy, ctx.key_proxy)
|
||||
if not effective_proxy or not effective_proxy.get("enabled", True):
|
||||
effective_proxy = await get_system_proxy_config_async()
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
is_tunnel_delegate = bool(delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy and not is_tunnel_delegate:
|
||||
proxy_url = await build_proxy_url_async(effective_proxy)
|
||||
|
||||
proxy_info = await resolve_proxy_info_async(effective_proxy)
|
||||
ctx.proxy_config = effective_proxy
|
||||
ctx.delegate_config = delegate_cfg
|
||||
ctx.proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if is_tunnel_delegate else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None if is_tunnel_delegate else None
|
||||
),
|
||||
)
|
||||
return ctx
|
||||
|
||||
|
||||
async def _resolve_upstream_context_standalone(request: Request) -> UpstreamContext:
|
||||
@@ -357,13 +414,170 @@ async def _resolve_upstream_context_standalone(request: Request) -> UpstreamCont
|
||||
UpstreamContext: 包含所有必要信息的上下文对象
|
||||
"""
|
||||
with create_session() as db:
|
||||
upstream_key, base_url, file_key_id, user_id = await _resolve_upstream_context(request, db)
|
||||
return UpstreamContext(
|
||||
upstream_key=upstream_key,
|
||||
base_url=base_url,
|
||||
file_key_id=file_key_id,
|
||||
user_id=user_id,
|
||||
ctx = await _resolve_upstream_context(request, db)
|
||||
return await _enrich_upstream_context_proxy(ctx)
|
||||
|
||||
|
||||
def _build_response_headers(headers: dict[str, str]) -> dict[str, str]:
|
||||
response_headers = {}
|
||||
hop_by_hop = {"connection", "keep-alive", "transfer-encoding", "upgrade"}
|
||||
for name, value in headers.items():
|
||||
if name.lower() not in hop_by_hop:
|
||||
response_headers[name] = value
|
||||
return response_headers
|
||||
|
||||
|
||||
async def _maybe_store_file_mapping_from_payload(
|
||||
*,
|
||||
status_code: int,
|
||||
headers: dict[str, str],
|
||||
content_bytes: bytes,
|
||||
file_key_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
) -> None:
|
||||
if (
|
||||
not file_key_id
|
||||
or status_code >= 300
|
||||
or not headers.get("content-type", "").startswith("application/json")
|
||||
):
|
||||
return
|
||||
|
||||
try:
|
||||
payload = json.loads(content_bytes)
|
||||
file_name = None
|
||||
file_obj = None
|
||||
|
||||
if isinstance(payload, dict):
|
||||
file_name = payload.get("name")
|
||||
file_obj = payload
|
||||
|
||||
if not file_name and isinstance(payload.get("file"), dict):
|
||||
file_name = payload["file"].get("name")
|
||||
file_obj = payload["file"]
|
||||
|
||||
if file_name and file_obj:
|
||||
display_name = file_obj.get("displayName") or file_obj.get("display_name")
|
||||
mime_type = file_obj.get("mimeType") or file_obj.get("mime_type")
|
||||
await store_file_key_mapping(
|
||||
file_name,
|
||||
file_key_id,
|
||||
user_id=user_id,
|
||||
display_name=display_name,
|
||||
mime_type=mime_type,
|
||||
)
|
||||
logger.debug(f"Gemini file→key 映射已存储: {file_name} → key_id={file_key_id}")
|
||||
|
||||
files_list = payload.get("files")
|
||||
if isinstance(files_list, list):
|
||||
mapped_count = 0
|
||||
for item in files_list:
|
||||
if isinstance(item, dict) and item.get("name"):
|
||||
item_display_name = item.get("displayName") or item.get("display_name")
|
||||
item_mime_type = item.get("mimeType") or item.get("mime_type")
|
||||
await store_file_key_mapping(
|
||||
item["name"],
|
||||
file_key_id,
|
||||
user_id=user_id,
|
||||
display_name=item_display_name,
|
||||
mime_type=item_mime_type,
|
||||
)
|
||||
mapped_count += 1
|
||||
if mapped_count > 0:
|
||||
logger.debug(
|
||||
"Gemini list_files 批量映射已存储: {} 个文件 → key_id={}",
|
||||
mapped_count,
|
||||
file_key_id,
|
||||
)
|
||||
except (ValueError, KeyError) as e:
|
||||
logger.debug("Failed to store Gemini file mapping: {}", e)
|
||||
|
||||
|
||||
async def _try_rust_sync_proxy_request(
|
||||
method: str,
|
||||
upstream_url: str,
|
||||
headers: dict[str, str],
|
||||
*,
|
||||
content: bytes | None = None,
|
||||
json_body: dict[str, Any] | None = None,
|
||||
file_key_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
provider_id: str = "",
|
||||
endpoint_id: str = "",
|
||||
proxy: Any = None,
|
||||
) -> Response | None:
|
||||
from src.config.settings import config
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
request_headers = dict(headers)
|
||||
content_type = str(request_headers.get("content-type") or "").strip() or None
|
||||
request_body = json_body if json_body is not None else content
|
||||
|
||||
try:
|
||||
result = await RustExecutorClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=f"gemini-files-{uuid4().hex}",
|
||||
candidate_id=None,
|
||||
provider_name="gemini",
|
||||
provider_id=provider_id,
|
||||
endpoint_id=endpoint_id,
|
||||
key_id=str(file_key_id or ""),
|
||||
method=method.upper(),
|
||||
url=upstream_url,
|
||||
headers=request_headers,
|
||||
body=build_execution_plan_body(request_body, content_type=content_type),
|
||||
stream=False,
|
||||
provider_api_format="gemini:files",
|
||||
client_api_format="gemini:files",
|
||||
model_name="gemini-files",
|
||||
proxy=proxy,
|
||||
content_type=content_type,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=300_000,
|
||||
write_ms=300_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=300_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (RustExecutorClientError, HTTPException, json.JSONDecodeError, ValueError) as exc:
|
||||
logger.warning("Gemini Files Rust proxy unavailable: {}", redact_url_for_log(str(exc)))
|
||||
return None
|
||||
|
||||
response_headers = _build_response_headers(dict(result.headers))
|
||||
if result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
|
||||
elif result.response_body_bytes is not None:
|
||||
response_body = result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
await _maybe_store_file_mapping_from_payload(
|
||||
status_code=result.status_code,
|
||||
headers=response_headers,
|
||||
content_bytes=response_body,
|
||||
file_key_id=file_key_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=response_body,
|
||||
status_code=result.status_code,
|
||||
headers=response_headers,
|
||||
media_type=response_headers.get("content-type", "application/json"),
|
||||
)
|
||||
|
||||
|
||||
async def _proxy_request(
|
||||
@@ -374,6 +588,11 @@ async def _proxy_request(
|
||||
json_body: dict[str, Any] | None = None,
|
||||
file_key_id: str | None = None,
|
||||
user_id: str | None = None,
|
||||
provider_id: str = "",
|
||||
endpoint_id: str = "",
|
||||
proxy: Any = None,
|
||||
proxy_config: dict[str, Any] | None = None,
|
||||
delegate_config: dict[str, Any] | None = None,
|
||||
) -> Response:
|
||||
"""
|
||||
代理请求到上游 Gemini API
|
||||
@@ -390,7 +609,25 @@ async def _proxy_request(
|
||||
Returns:
|
||||
FastAPI Response 对象
|
||||
"""
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
rust_response = await _try_rust_sync_proxy_request(
|
||||
method,
|
||||
upstream_url,
|
||||
headers,
|
||||
content=content,
|
||||
json_body=json_body,
|
||||
file_key_id=file_key_id,
|
||||
user_id=user_id,
|
||||
provider_id=provider_id,
|
||||
endpoint_id=endpoint_id,
|
||||
proxy=proxy,
|
||||
)
|
||||
if rust_response is not None:
|
||||
return rust_response
|
||||
|
||||
client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_config,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
|
||||
try:
|
||||
if method.upper() == "GET":
|
||||
@@ -408,73 +645,14 @@ async def _proxy_request(
|
||||
raise HTTPException(status_code=405, detail="Method not allowed")
|
||||
|
||||
# 构建响应头(排除 hop-by-hop 头部)
|
||||
response_headers = {}
|
||||
hop_by_hop = {"connection", "keep-alive", "transfer-encoding", "upgrade"}
|
||||
for name, value in response.headers.items():
|
||||
if name.lower() not in hop_by_hop:
|
||||
response_headers[name] = value
|
||||
|
||||
if (
|
||||
file_key_id
|
||||
and response.status_code < 300
|
||||
and response.headers.get("content-type", "").startswith("application/json")
|
||||
):
|
||||
try:
|
||||
payload = response.json()
|
||||
file_name = None
|
||||
file_obj = None
|
||||
|
||||
if isinstance(payload, dict):
|
||||
# 单文件上传响应
|
||||
file_name = payload.get("name")
|
||||
file_obj = payload
|
||||
|
||||
# 嵌套格式:{"file": {...}}
|
||||
if not file_name and isinstance(payload.get("file"), dict):
|
||||
file_name = payload["file"].get("name")
|
||||
file_obj = payload["file"]
|
||||
|
||||
if file_name and file_obj:
|
||||
display_name = file_obj.get("displayName") or file_obj.get("display_name")
|
||||
mime_type = file_obj.get("mimeType") or file_obj.get("mime_type")
|
||||
await store_file_key_mapping(
|
||||
file_name,
|
||||
file_key_id,
|
||||
user_id=user_id,
|
||||
display_name=display_name,
|
||||
mime_type=mime_type,
|
||||
)
|
||||
logger.debug(
|
||||
f"Gemini file→key 映射已存储: {file_name} → key_id={file_key_id}"
|
||||
)
|
||||
|
||||
# 为 list_files 响应中的所有文件建立映射
|
||||
# 这是正确的:Gemini API 按 Key 隔离文件,返回的文件必然属于当前 Key
|
||||
files_list = payload.get("files")
|
||||
if isinstance(files_list, list):
|
||||
mapped_count = 0
|
||||
for item in files_list:
|
||||
if isinstance(item, dict) and item.get("name"):
|
||||
item_display_name = item.get("displayName") or item.get(
|
||||
"display_name"
|
||||
)
|
||||
item_mime_type = item.get("mimeType") or item.get("mime_type")
|
||||
await store_file_key_mapping(
|
||||
item["name"],
|
||||
file_key_id,
|
||||
user_id=user_id,
|
||||
display_name=item_display_name,
|
||||
mime_type=item_mime_type,
|
||||
)
|
||||
mapped_count += 1
|
||||
if mapped_count > 0:
|
||||
logger.debug(
|
||||
"Gemini list_files 批量映射已存储: {} 个文件 → key_id={}",
|
||||
mapped_count,
|
||||
file_key_id,
|
||||
)
|
||||
except (ValueError, KeyError) as e:
|
||||
logger.debug("Failed to store Gemini file mapping: {}", e)
|
||||
response_headers = _build_response_headers(dict(response.headers))
|
||||
await _maybe_store_file_mapping_from_payload(
|
||||
status_code=response.status_code,
|
||||
headers=response_headers,
|
||||
content_bytes=response.content,
|
||||
file_key_id=file_key_id,
|
||||
user_id=user_id,
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=response.content,
|
||||
@@ -560,6 +738,11 @@ async def upload_file(
|
||||
content=body,
|
||||
file_key_id=ctx.file_key_id,
|
||||
user_id=ctx.user_id,
|
||||
provider_id=ctx.provider_id,
|
||||
endpoint_id=ctx.endpoint_id,
|
||||
proxy=ctx.proxy_snapshot,
|
||||
proxy_config=ctx.proxy_config,
|
||||
delegate_config=ctx.delegate_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -624,7 +807,16 @@ async def list_files(
|
||||
logger.debug("Gemini Files list proxy: GET {}", redact_url_for_log(upstream_url))
|
||||
|
||||
return await _proxy_request(
|
||||
"GET", upstream_url, headers, file_key_id=ctx.file_key_id, user_id=ctx.user_id
|
||||
"GET",
|
||||
upstream_url,
|
||||
headers,
|
||||
file_key_id=ctx.file_key_id,
|
||||
user_id=ctx.user_id,
|
||||
provider_id=ctx.provider_id,
|
||||
endpoint_id=ctx.endpoint_id,
|
||||
proxy=ctx.proxy_snapshot,
|
||||
proxy_config=ctx.proxy_config,
|
||||
delegate_config=ctx.delegate_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -716,9 +908,19 @@ async def download_file(
|
||||
|
||||
优化:HTTP 下载期间不持有数据库连接
|
||||
"""
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
from fastapi.responses import JSONResponse, Response
|
||||
|
||||
from src.config.settings import config
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
ExecutionProxySnapshot,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
# ========== 阶段 1:数据库操作(短暂持有连接)==========
|
||||
client_key = _extract_gemini_api_key(request)
|
||||
@@ -730,6 +932,8 @@ async def download_file(
|
||||
},
|
||||
)
|
||||
|
||||
regular_file_ctx: UpstreamContext | None = None
|
||||
|
||||
# 在数据库会话内完成所有查询
|
||||
with create_session() as db:
|
||||
auth_result = AuthService.authenticate_api_key(db, client_key)
|
||||
@@ -766,12 +970,13 @@ async def download_file(
|
||||
},
|
||||
)
|
||||
upstream_url = video_url
|
||||
file_key_id = ""
|
||||
provider_id = ""
|
||||
endpoint_id = ""
|
||||
else:
|
||||
# 普通文件下载:透传到 Gemini
|
||||
try:
|
||||
upstream_key, base_url, _file_key_id, _user_id = await _resolve_upstream_context(
|
||||
request, db
|
||||
)
|
||||
regular_file_ctx = await _resolve_upstream_context(request, db)
|
||||
except HTTPException:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
@@ -783,26 +988,145 @@ async def download_file(
|
||||
}
|
||||
},
|
||||
)
|
||||
file_name = f"files/{file_id}" if not file_id.startswith("files/") else file_id
|
||||
upstream_url = _build_upstream_url(
|
||||
base_url,
|
||||
f"/v1beta/{file_name}:download",
|
||||
dict(request.query_params),
|
||||
)
|
||||
|
||||
if regular_file_ctx is not None:
|
||||
ctx = await _enrich_upstream_context_proxy(regular_file_ctx)
|
||||
upstream_key = ctx.upstream_key
|
||||
file_key_id = ctx.file_key_id
|
||||
provider_id = ctx.provider_id
|
||||
endpoint_id = ctx.endpoint_id
|
||||
file_name = f"files/{file_id}" if not file_id.startswith("files/") else file_id
|
||||
upstream_url = _build_upstream_url(
|
||||
ctx.base_url,
|
||||
f"/v1beta/{file_name}:download",
|
||||
dict(request.query_params),
|
||||
)
|
||||
|
||||
# ========== 阶段 2:HTTP 下载(不持有数据库连接)==========
|
||||
headers = _build_upstream_headers(dict(request.headers), upstream_key)
|
||||
|
||||
logger.debug("Gemini Files download proxy: GET {}", redact_url_for_log(upstream_url))
|
||||
|
||||
proxy_snapshot = None
|
||||
proxy_config = None
|
||||
delegate_config = None
|
||||
if regular_file_ctx is not None:
|
||||
proxy_snapshot = ctx.proxy_snapshot
|
||||
proxy_config = ctx.proxy_config
|
||||
delegate_config = ctx.delegate_config
|
||||
|
||||
if config.executor_backend == "rust":
|
||||
try:
|
||||
if proxy_snapshot is None and file_id.startswith("aev_"):
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
|
||||
system_proxy = await get_system_proxy_config_async()
|
||||
delegate_cfg = await resolve_delegate_config_async(system_proxy)
|
||||
proxy_config = system_proxy
|
||||
delegate_config = delegate_cfg
|
||||
proxy_url: str | None = None
|
||||
if system_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
|
||||
proxy_url = await build_proxy_url_async(system_proxy)
|
||||
proxy_info = await resolve_proxy_info_async(system_proxy)
|
||||
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Gemini Files download proxy snapshot build failed: {}", exc)
|
||||
|
||||
try:
|
||||
rust_stream = await RustExecutorClient().execute_stream(
|
||||
ExecutionPlan(
|
||||
request_id=f"gemini-files-download-{uuid4().hex}",
|
||||
candidate_id=None,
|
||||
provider_name="gemini",
|
||||
provider_id=provider_id,
|
||||
endpoint_id=endpoint_id,
|
||||
key_id=str(file_key_id or ""),
|
||||
method="GET",
|
||||
url=upstream_url,
|
||||
headers=headers,
|
||||
body=ExecutionPlanBody(),
|
||||
stream=True,
|
||||
provider_api_format="gemini:files",
|
||||
client_api_format="gemini:files",
|
||||
model_name="gemini-files",
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=300_000,
|
||||
write_ms=300_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=None,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (RustExecutorClientError, ValueError) as exc:
|
||||
logger.warning("Gemini Files Rust download unavailable: {}", exc)
|
||||
else:
|
||||
safe_headers = _build_response_headers(dict(rust_stream.headers))
|
||||
if rust_stream.status_code >= 400:
|
||||
error_chunks: list[bytes] = []
|
||||
try:
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
if chunk:
|
||||
error_chunks.append(chunk)
|
||||
if sum(len(item) for item in error_chunks) >= 16_384:
|
||||
break
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
raw_error = b"".join(error_chunks)
|
||||
if safe_headers.get("content-type", "").startswith("application/json"):
|
||||
try:
|
||||
return JSONResponse(
|
||||
content=json.loads(raw_error.decode("utf-8")),
|
||||
status_code=rust_stream.status_code,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
return JSONResponse(
|
||||
content={"error": raw_error.decode("utf-8", errors="replace")},
|
||||
status_code=rust_stream.status_code,
|
||||
)
|
||||
|
||||
async def _iter_bytes() -> AsyncIterator[bytes]:
|
||||
try:
|
||||
async for chunk in rust_stream.byte_iterator:
|
||||
yield chunk
|
||||
finally:
|
||||
await rust_stream.response_ctx.__aexit__(None, None, None)
|
||||
|
||||
return StreamingResponse(
|
||||
_iter_bytes(),
|
||||
status_code=rust_stream.status_code,
|
||||
headers=safe_headers,
|
||||
media_type=safe_headers.get("content-type", "application/octet-stream"),
|
||||
)
|
||||
|
||||
# 使用 follow_redirects=True 跟随重定向(Gemini 文件下载会重定向)
|
||||
try:
|
||||
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
||||
|
||||
async with httpx.AsyncClient(
|
||||
**build_proxy_client_kwargs(timeout=httpx.Timeout(300.0), follow_redirects=True)
|
||||
) as client:
|
||||
response = await client.get(upstream_url, headers=headers)
|
||||
client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_config,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
response = await client.get(
|
||||
upstream_url,
|
||||
headers=headers,
|
||||
timeout=300.0,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error("Gemini Files download failed: {}", exc)
|
||||
raise HTTPException(status_code=502, detail="Failed to download file")
|
||||
@@ -818,15 +1142,10 @@ async def download_file(
|
||||
content = {"error": response.text}
|
||||
return JSONResponse(content=content, status_code=response.status_code)
|
||||
|
||||
# 返回文件内容
|
||||
return Response(
|
||||
content=response.content,
|
||||
status_code=response.status_code,
|
||||
headers={
|
||||
k: v
|
||||
for k, v in response.headers.items()
|
||||
if k.lower() not in {"transfer-encoding", "connection", "keep-alive"}
|
||||
},
|
||||
headers=_build_response_headers(dict(response.headers)),
|
||||
media_type=response.headers.get("content-type", "application/octet-stream"),
|
||||
)
|
||||
|
||||
@@ -887,7 +1206,16 @@ async def get_file(
|
||||
logger.debug("Gemini Files get proxy: GET {}", redact_url_for_log(upstream_url))
|
||||
|
||||
return await _proxy_request(
|
||||
"GET", upstream_url, headers, file_key_id=ctx.file_key_id, user_id=ctx.user_id
|
||||
"GET",
|
||||
upstream_url,
|
||||
headers,
|
||||
file_key_id=ctx.file_key_id,
|
||||
user_id=ctx.user_id,
|
||||
provider_id=ctx.provider_id,
|
||||
endpoint_id=ctx.endpoint_id,
|
||||
proxy=ctx.proxy_snapshot,
|
||||
proxy_config=ctx.proxy_config,
|
||||
delegate_config=ctx.delegate_config,
|
||||
)
|
||||
|
||||
|
||||
@@ -933,7 +1261,16 @@ async def delete_file(
|
||||
|
||||
logger.debug("Gemini Files delete proxy: DELETE {}", redact_url_for_log(upstream_url))
|
||||
|
||||
response = await _proxy_request("DELETE", upstream_url, headers)
|
||||
response = await _proxy_request(
|
||||
"DELETE",
|
||||
upstream_url,
|
||||
headers,
|
||||
provider_id=ctx.provider_id,
|
||||
endpoint_id=ctx.endpoint_id,
|
||||
proxy=ctx.proxy_snapshot,
|
||||
proxy_config=ctx.proxy_config,
|
||||
delegate_config=ctx.delegate_config,
|
||||
)
|
||||
if response.status_code < 300:
|
||||
await delete_file_key_mapping(file_name)
|
||||
else:
|
||||
|
||||
@@ -19,13 +19,15 @@ from src.api.handlers.base.request_builder import (
|
||||
build_test_request_body,
|
||||
get_provider_auth,
|
||||
)
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.clients.redis_client import get_redis_client
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.database import get_db
|
||||
from src.database.database import get_pool_status
|
||||
from src.models.database import GlobalModel, Model, Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.provider.provider_context import resolve_provider_proxy
|
||||
from src.services.provider.transport import build_provider_url
|
||||
from src.utils.ssl_utils import get_ssl_context
|
||||
|
||||
router = APIRouter(tags=["System Catalog"])
|
||||
|
||||
@@ -98,6 +100,132 @@ def _select_provider(db: Session, provider_name: str | None) -> Provider | None:
|
||||
return query.order_by(Provider.provider_priority.asc()).first()
|
||||
|
||||
|
||||
async def _build_test_connection_transport_context(
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
) -> tuple[dict[str, Any] | None, dict[str, Any] | None, Any]:
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_effective_proxy,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
|
||||
try:
|
||||
effective_proxy = resolve_effective_proxy(
|
||||
resolve_provider_proxy(endpoint=endpoint, key=key),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
if not effective_proxy or not effective_proxy.get("enabled", True):
|
||||
effective_proxy = await get_system_proxy_config_async()
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
|
||||
proxy_url = await build_proxy_url_async(effective_proxy)
|
||||
|
||||
proxy_info = await resolve_proxy_info_async(effective_proxy)
|
||||
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
return effective_proxy, delegate_cfg, proxy_snapshot
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Failed to build test-connection transport context endpoint={} key={}: {}",
|
||||
getattr(endpoint, "id", None),
|
||||
getattr(key, "id", None),
|
||||
exc,
|
||||
)
|
||||
return None, None, None
|
||||
|
||||
|
||||
async def _try_rust_test_connection_response(
|
||||
*,
|
||||
request_id: str,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
body: dict[str, Any],
|
||||
provider_name: str,
|
||||
provider_id: str | None,
|
||||
endpoint_id: str | None,
|
||||
key_id: str | None,
|
||||
api_format: str,
|
||||
model_name: str,
|
||||
proxy_snapshot: Any,
|
||||
) -> httpx.Response | None:
|
||||
import json
|
||||
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await RustExecutorClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=request_id,
|
||||
candidate_id=None,
|
||||
provider_name=provider_name,
|
||||
provider_id=str(provider_id or ""),
|
||||
endpoint_id=str(endpoint_id or ""),
|
||||
key_id=str(key_id or ""),
|
||||
method="POST",
|
||||
url=url,
|
||||
headers=dict(headers),
|
||||
body=build_execution_plan_body(body, content_type="application/json"),
|
||||
stream=False,
|
||||
provider_api_format=api_format,
|
||||
client_api_format=api_format,
|
||||
model_name=model_name,
|
||||
content_type="application/json",
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=30_000,
|
||||
write_ms=30_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=30_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning("Rust test-connection fallback url={}: {}", url, exc)
|
||||
return None
|
||||
|
||||
response_headers = dict(result.headers)
|
||||
if result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
|
||||
elif result.response_body_bytes is not None:
|
||||
response_body = result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
return httpx.Response(
|
||||
status_code=result.status_code,
|
||||
request=httpx.Request("POST", url, headers=headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
|
||||
# ============== 端点 ==============
|
||||
|
||||
|
||||
@@ -375,11 +503,33 @@ async def test_connection(
|
||||
key=key,
|
||||
decrypted_auth_config=auth_info.decrypted_auth_config if auth_info else None,
|
||||
)
|
||||
proxy_config, delegate_cfg, proxy_snapshot = await _build_test_connection_transport_context(
|
||||
endpoint,
|
||||
key,
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0, verify=get_ssl_context()) as client:
|
||||
resp = await _try_rust_test_connection_response(
|
||||
request_id=f"test-connection:{selected_provider.id}:{model}",
|
||||
url=url,
|
||||
headers=provider_headers,
|
||||
body=provider_payload,
|
||||
provider_name=selected_provider.name,
|
||||
provider_id=getattr(selected_provider, "id", None),
|
||||
endpoint_id=getattr(endpoint, "id", None),
|
||||
key_id=getattr(key, "id", None),
|
||||
api_format=format_value,
|
||||
model_name=model,
|
||||
proxy_snapshot=proxy_snapshot,
|
||||
)
|
||||
if resp is None:
|
||||
client = await HTTPClientPool.get_upstream_client(
|
||||
delegate_cfg,
|
||||
proxy_config=proxy_config,
|
||||
)
|
||||
resp = await client.post(url, json=provider_payload, headers=provider_headers)
|
||||
resp.raise_for_status()
|
||||
response = resp.json()
|
||||
|
||||
resp.raise_for_status()
|
||||
response = resp.json()
|
||||
|
||||
return {
|
||||
"status": "success",
|
||||
|
||||
@@ -192,6 +192,20 @@ class Config:
|
||||
# HTTP_REQUEST_TIMEOUT: 非流式请求整体超时(秒),默认 300 秒
|
||||
self.http_request_timeout = float(os.getenv("HTTP_REQUEST_TIMEOUT", "300.0"))
|
||||
|
||||
# 内部 executor 配置
|
||||
# EXECUTOR_BACKEND:
|
||||
# - rust: 当前 pioneer 分支默认,优先将可序列化的执行计划转发给 Rust executor
|
||||
# - python: 显式回退到现有 httpx 路径
|
||||
self.executor_backend = os.getenv("EXECUTOR_BACKEND", "rust").strip().lower()
|
||||
self.executor_transport = os.getenv("EXECUTOR_TRANSPORT", "unix_socket").strip().lower()
|
||||
self.executor_socket_path = os.getenv(
|
||||
"EXECUTOR_SOCKET_PATH", "/tmp/aether-executor.sock"
|
||||
).strip()
|
||||
self.executor_base_url = os.getenv("EXECUTOR_BASE_URL", "http://127.0.0.1:5219").strip()
|
||||
self.executor_request_timeout = float(
|
||||
os.getenv("EXECUTOR_REQUEST_TIMEOUT", str(self.http_request_timeout))
|
||||
)
|
||||
|
||||
# HTTP 连接池配置
|
||||
# HTTP_MAX_CONNECTIONS: 最大连接数,影响并发能力
|
||||
# - 每个连接占用一个 socket,过多会耗尽系统资源
|
||||
|
||||
@@ -13,6 +13,7 @@ import httpx
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.core.crypto import crypto_service
|
||||
from src.core.logger import logger
|
||||
from src.models.database import Provider, ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.provider.auth import get_provider_auth
|
||||
from src.services.provider.oauth_token import looks_like_token_invalidated
|
||||
@@ -70,6 +71,138 @@ def _extract_error_message_from_response(response: httpx.Response) -> str:
|
||||
return text[:300] if text else ""
|
||||
|
||||
|
||||
async def _build_codex_proxy_snapshot(
|
||||
effective_proxy: dict[str, Any] | None,
|
||||
) -> Any:
|
||||
from importlib import import_module
|
||||
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
|
||||
resolver_module = import_module("src.services.proxy_node.resolver")
|
||||
build_proxy_url_async = getattr(resolver_module, "build_proxy_url_async", None)
|
||||
get_system_proxy_config_async = getattr(resolver_module, "get_system_proxy_config_async", None)
|
||||
resolve_delegate_config_async = getattr(resolver_module, "resolve_delegate_config_async", None)
|
||||
resolve_proxy_info_async = getattr(resolver_module, "resolve_proxy_info_async", None)
|
||||
|
||||
proxy_config = effective_proxy
|
||||
if (not proxy_config or not proxy_config.get("enabled", True)) and callable(
|
||||
get_system_proxy_config_async
|
||||
):
|
||||
proxy_config = await get_system_proxy_config_async()
|
||||
if not proxy_config:
|
||||
return None
|
||||
|
||||
try:
|
||||
delegate_cfg = (
|
||||
await resolve_delegate_config_async(proxy_config)
|
||||
if callable(resolve_delegate_config_async)
|
||||
else None
|
||||
)
|
||||
proxy_url: str | None = None
|
||||
if (
|
||||
proxy_config
|
||||
and not (delegate_cfg and delegate_cfg.get("tunnel"))
|
||||
and callable(build_proxy_url_async)
|
||||
):
|
||||
proxy_url = await build_proxy_url_async(proxy_config)
|
||||
proxy_info = (
|
||||
await resolve_proxy_info_async(proxy_config)
|
||||
if callable(resolve_proxy_info_async)
|
||||
else {"url": proxy_url}
|
||||
)
|
||||
return ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning("Codex quota proxy snapshot build failed: {}", exc)
|
||||
return None
|
||||
|
||||
|
||||
async def _try_rust_codex_quota_response(
|
||||
*,
|
||||
key: ProviderAPIKey,
|
||||
provider: Provider,
|
||||
endpoint: ProviderEndpoint,
|
||||
url: str,
|
||||
headers: dict[str, str],
|
||||
proxy_snapshot: Any,
|
||||
) -> httpx.Response | None:
|
||||
from src.config import config
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await RustExecutorClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=f"codex-quota:{key.id}",
|
||||
candidate_id=None,
|
||||
provider_name="codex",
|
||||
provider_id=str(getattr(provider, "id", "") or ""),
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or ""),
|
||||
key_id=str(getattr(key, "id", "") or ""),
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=dict(headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=False,
|
||||
provider_api_format="openai:cli",
|
||||
client_api_format="openai:cli",
|
||||
model_name="codex-wham-usage",
|
||||
proxy=proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=30_000,
|
||||
write_ms=30_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=30_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (
|
||||
RustExecutorClientError,
|
||||
httpx.HTTPError,
|
||||
json.JSONDecodeError,
|
||||
ValueError,
|
||||
AttributeError,
|
||||
TypeError,
|
||||
) as exc:
|
||||
logger.warning("Codex quota Rust fallback key_id={} url={}: {}", key.id, url, exc)
|
||||
return None
|
||||
|
||||
response_headers = dict(result.headers)
|
||||
if result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
|
||||
elif result.response_body_bytes is not None:
|
||||
response_body = result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
return httpx.Response(
|
||||
status_code=result.status_code,
|
||||
request=httpx.Request("GET", url, headers=headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
|
||||
def _looks_like_account_deactivated(message: str | None) -> bool:
|
||||
lowered = str(message or "").strip().lower()
|
||||
return "account has been deactivated" in lowered or "account deactivated" in lowered
|
||||
@@ -212,12 +345,22 @@ async def refresh_codex_key_quota(
|
||||
getattr(provider, "proxy", None),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
proxy_snapshot = await _build_codex_proxy_snapshot(effective_proxy)
|
||||
|
||||
# 使用 wham/usage API 获取限额信息
|
||||
async with httpx.AsyncClient(
|
||||
**build_proxy_client_kwargs(effective_proxy, timeout=30.0)
|
||||
) as client:
|
||||
response = await client.get(codex_wham_usage_url, headers=headers)
|
||||
response = await _try_rust_codex_quota_response(
|
||||
key=key,
|
||||
provider=provider,
|
||||
endpoint=endpoint,
|
||||
url=codex_wham_usage_url,
|
||||
headers=headers,
|
||||
proxy_snapshot=proxy_snapshot,
|
||||
)
|
||||
if response is None:
|
||||
async with httpx.AsyncClient(
|
||||
**build_proxy_client_kwargs(effective_proxy, timeout=30.0)
|
||||
) as client:
|
||||
response = await client.get(codex_wham_usage_url, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
status_code = int(response.status_code)
|
||||
|
||||
@@ -145,6 +145,123 @@ class ProviderOpsService:
|
||||
finally:
|
||||
self.db.expire_on_commit = original_expire_on_commit
|
||||
|
||||
@staticmethod
|
||||
def _build_ops_proxy_snapshot(
|
||||
proxy: Any,
|
||||
tunnel_node_id: str | None,
|
||||
) -> Any:
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
|
||||
if tunnel_node_id:
|
||||
return ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode="tunnel",
|
||||
node_id=str(tunnel_node_id).strip() or None,
|
||||
)
|
||||
|
||||
if proxy is None:
|
||||
return None
|
||||
|
||||
proxy_url = getattr(proxy, "url", None)
|
||||
if proxy_url is not None:
|
||||
proxy_url = str(proxy_url).strip()
|
||||
else:
|
||||
proxy_url = str(proxy).strip()
|
||||
if not proxy_url:
|
||||
return None
|
||||
|
||||
mode = proxy_url.split("://", 1)[0].strip().lower() or None
|
||||
return ExecutionProxySnapshot(
|
||||
enabled=True,
|
||||
mode=mode,
|
||||
url=proxy_url,
|
||||
)
|
||||
|
||||
async def _try_rust_verify_response(
|
||||
self,
|
||||
*,
|
||||
architecture_id: str,
|
||||
verify_endpoint: str,
|
||||
headers: dict[str, str],
|
||||
proxy: Any,
|
||||
tunnel_node_id: str | None,
|
||||
) -> Any:
|
||||
import json
|
||||
|
||||
import httpx
|
||||
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await RustExecutorClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=f"provider-ops-verify:{architecture_id}",
|
||||
candidate_id=None,
|
||||
provider_name=architecture_id,
|
||||
provider_id="",
|
||||
endpoint_id="",
|
||||
key_id="",
|
||||
method="GET",
|
||||
url=verify_endpoint,
|
||||
headers=dict(headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=False,
|
||||
provider_api_format="provider_ops:verify",
|
||||
client_api_format="provider_ops:verify",
|
||||
model_name="verify-auth",
|
||||
proxy=self._build_ops_proxy_snapshot(proxy, tunnel_node_id),
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=30_000,
|
||||
write_ms=30_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=30_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (
|
||||
RustExecutorClientError,
|
||||
httpx.HTTPError,
|
||||
json.JSONDecodeError,
|
||||
ValueError,
|
||||
AttributeError,
|
||||
TypeError,
|
||||
) as exc:
|
||||
logger.warning(
|
||||
"Provider ops Rust verify fallback architecture={} endpoint={}: {}",
|
||||
architecture_id,
|
||||
verify_endpoint,
|
||||
exc,
|
||||
)
|
||||
return None
|
||||
|
||||
response_headers = dict(result.headers)
|
||||
if result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
|
||||
elif result.response_body_bytes is not None:
|
||||
response_body = result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
return httpx.Response(
|
||||
status_code=result.status_code,
|
||||
request=httpx.Request("GET", verify_endpoint, headers=headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
# ==================== 配置管理 ====================
|
||||
|
||||
def get_config(self, provider_id: str) -> ProviderOpsConfig | None:
|
||||
@@ -1046,64 +1163,74 @@ class ProviderOpsService:
|
||||
|
||||
proxy, tunnel_node_id = await resolve_ops_proxy_config_async(config)
|
||||
|
||||
# 构建 httpx client 参数
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"timeout": 30.0,
|
||||
"verify": get_ssl_context(),
|
||||
}
|
||||
if tunnel_node_id:
|
||||
from src.services.proxy_node.tunnel_transport import create_tunnel_transport
|
||||
response = await self._try_rust_verify_response(
|
||||
architecture_id=architecture_id,
|
||||
verify_endpoint=verify_endpoint,
|
||||
headers=headers,
|
||||
proxy=proxy,
|
||||
tunnel_node_id=tunnel_node_id,
|
||||
)
|
||||
if response is None:
|
||||
# 构建 httpx client 参数
|
||||
client_kwargs: dict[str, Any] = {
|
||||
"timeout": 30.0,
|
||||
"verify": get_ssl_context(),
|
||||
}
|
||||
if tunnel_node_id:
|
||||
from src.services.proxy_node.tunnel_transport import create_tunnel_transport
|
||||
|
||||
client_kwargs["transport"] = create_tunnel_transport(tunnel_node_id, timeout=30.0)
|
||||
logger.debug("使用 tunnel 代理: node_id={}", tunnel_node_id)
|
||||
elif proxy:
|
||||
client_kwargs["proxy"] = proxy
|
||||
logger.debug("使用代理: {}", proxy)
|
||||
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
response = await client.get(verify_endpoint, headers=headers)
|
||||
|
||||
logger.debug(
|
||||
"验证响应: status={}, content_type={}",
|
||||
response.status_code,
|
||||
response.headers.get("content-type"),
|
||||
)
|
||||
|
||||
# 尝试解析 JSON
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception:
|
||||
data = {}
|
||||
|
||||
# 将预处理获取的额外数据合并到响应中
|
||||
if "_combined_data" in merged_config:
|
||||
data["_combined_data"] = merged_config["_combined_data"]
|
||||
elif "_balance_data" in merged_config:
|
||||
data["_balance_data"] = merged_config["_balance_data"]
|
||||
|
||||
# 使用架构的方法解析响应
|
||||
result = architecture.parse_verify_response(response.status_code, data)
|
||||
result_dict = result.to_dict()
|
||||
|
||||
# 将凭据更新信息附加到响应,供前端同步更新表单
|
||||
# 过滤掉内部缓存字段(以 _ 开头),前端不需要这些
|
||||
frontend_creds = {k: v for k, v in updated_creds.items() if not k.startswith("_")}
|
||||
if frontend_creds:
|
||||
result_dict["updated_credentials"] = frontend_creds
|
||||
|
||||
# 验证成功且有 provider_id 时,缓存余额
|
||||
if result.success and provider_id and result.quota is not None:
|
||||
# 从架构的默认配置获取 quota_divisor
|
||||
balance_config = architecture.default_action_configs.get(
|
||||
ProviderActionType.QUERY_BALANCE, {}
|
||||
client_kwargs["transport"] = create_tunnel_transport(
|
||||
tunnel_node_id, timeout=30.0
|
||||
)
|
||||
quota_divisor = balance_config.get("quota_divisor", 1)
|
||||
# 转换为美元值后缓存
|
||||
quota_usd = result.quota / quota_divisor
|
||||
# 传入 extra 信息(如窗口限额)
|
||||
await self._cache_balance_from_verify(provider_id, quota_usd, result.extra)
|
||||
logger.debug("使用 tunnel 代理: node_id={}", tunnel_node_id)
|
||||
elif proxy:
|
||||
client_kwargs["proxy"] = proxy
|
||||
logger.debug("使用代理: {}", proxy)
|
||||
|
||||
return result_dict
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
response = await client.get(verify_endpoint, headers=headers)
|
||||
|
||||
logger.debug(
|
||||
"验证响应: status={}, content_type={}",
|
||||
response.status_code,
|
||||
response.headers.get("content-type"),
|
||||
)
|
||||
|
||||
# 尝试解析 JSON
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception:
|
||||
data = {}
|
||||
|
||||
# 将预处理获取的额外数据合并到响应中
|
||||
if "_combined_data" in merged_config:
|
||||
data["_combined_data"] = merged_config["_combined_data"]
|
||||
elif "_balance_data" in merged_config:
|
||||
data["_balance_data"] = merged_config["_balance_data"]
|
||||
|
||||
# 使用架构的方法解析响应
|
||||
result = architecture.parse_verify_response(response.status_code, data)
|
||||
result_dict = result.to_dict()
|
||||
|
||||
# 将凭据更新信息附加到响应,供前端同步更新表单
|
||||
# 过滤掉内部缓存字段(以 _ 开头),前端不需要这些
|
||||
frontend_creds = {k: v for k, v in updated_creds.items() if not k.startswith("_")}
|
||||
if frontend_creds:
|
||||
result_dict["updated_credentials"] = frontend_creds
|
||||
|
||||
# 验证成功且有 provider_id 时,缓存余额
|
||||
if result.success and provider_id and result.quota is not None:
|
||||
# 从架构的默认配置获取 quota_divisor
|
||||
balance_config = architecture.default_action_configs.get(
|
||||
ProviderActionType.QUERY_BALANCE, {}
|
||||
)
|
||||
quota_divisor = balance_config.get("quota_divisor", 1)
|
||||
# 转换为美元值后缓存
|
||||
quota_usd = result.quota / quota_divisor
|
||||
# 传入 extra 信息(如窗口限额)
|
||||
await self._cache_balance_from_verify(provider_id, quota_usd, result.extra)
|
||||
|
||||
return result_dict
|
||||
|
||||
except ValueError as e:
|
||||
return {"success": False, "message": str(e)}
|
||||
|
||||
@@ -0,0 +1,183 @@
|
||||
"""
|
||||
执行计划契约
|
||||
|
||||
用于在 Python 控制面和未来的 Rust executor 之间传递稳定的请求执行信息。
|
||||
当前阶段先服务于非流式 chat 路径的计划构建与本地执行拆分。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
def _drop_none(value: Any) -> Any:
|
||||
"""递归移除 None 字段,便于序列化为紧凑 payload。"""
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
key: _drop_none(item)
|
||||
for key, item in value.items()
|
||||
if item is not None
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [_drop_none(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ExecutionPlanTimeouts:
|
||||
connect_ms: int | None = None
|
||||
read_ms: int | None = None
|
||||
write_ms: int | None = None
|
||||
pool_ms: int | None = None
|
||||
total_ms: int | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ExecutionPlanBody:
|
||||
json_body: Any = None
|
||||
body_bytes_b64: str | None = None
|
||||
body_ref: str | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ExecutionProxySnapshot:
|
||||
enabled: bool
|
||||
mode: str | None = None
|
||||
node_id: str | None = None
|
||||
label: str | None = None
|
||||
url: str | None = None
|
||||
extra: dict[str, Any] | None = None
|
||||
|
||||
@classmethod
|
||||
def from_proxy_info(
|
||||
cls,
|
||||
proxy_info: dict[str, Any] | None,
|
||||
*,
|
||||
proxy_url: str | None = None,
|
||||
mode_override: str | None = None,
|
||||
node_id_override: str | None = None,
|
||||
extra: dict[str, Any] | None = None,
|
||||
) -> ExecutionProxySnapshot | None:
|
||||
if not proxy_info and not proxy_url:
|
||||
return None
|
||||
mode = mode_override or str(proxy_info.get("type") or proxy_info.get("mode") or "").strip() or None
|
||||
if not mode and proxy_url:
|
||||
mode = proxy_url.split("://", 1)[0].strip().lower() or None
|
||||
return cls(
|
||||
enabled=True,
|
||||
mode=mode,
|
||||
node_id=node_id_override
|
||||
or str((proxy_info or {}).get("node_id") or "").strip()
|
||||
or None,
|
||||
label=str((proxy_info or {}).get("label") or "").strip() or None,
|
||||
url=str(proxy_url or (proxy_info or {}).get("url") or "").strip() or None,
|
||||
extra=extra or None,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ExecutionPlan:
|
||||
request_id: str
|
||||
candidate_id: str | None
|
||||
provider_name: str
|
||||
provider_id: str
|
||||
endpoint_id: str
|
||||
key_id: str
|
||||
method: str
|
||||
url: str
|
||||
headers: dict[str, str]
|
||||
body: ExecutionPlanBody
|
||||
stream: bool
|
||||
provider_api_format: str
|
||||
client_api_format: str
|
||||
model_name: str
|
||||
content_type: str | None = None
|
||||
content_encoding: str | None = None
|
||||
proxy: ExecutionProxySnapshot | None = None
|
||||
tls_profile: str | None = None
|
||||
timeouts: ExecutionPlanTimeouts | None = None
|
||||
|
||||
def to_payload(self) -> dict[str, Any]:
|
||||
return _drop_none(asdict(self))
|
||||
|
||||
|
||||
def is_remote_proxy_supported(proxy: ExecutionProxySnapshot | None) -> bool:
|
||||
proxy_mode = str((proxy.mode if proxy else "") or "").strip().lower()
|
||||
return proxy is None or (
|
||||
str(proxy.url or "").strip() != ""
|
||||
and proxy_mode not in {"tunnel"}
|
||||
) or (
|
||||
proxy_mode == "tunnel"
|
||||
and str(proxy.node_id or "").strip() != ""
|
||||
)
|
||||
|
||||
|
||||
def is_remote_contract_eligible(contract: ExecutionPlan) -> bool:
|
||||
content_encoding = str(contract.content_encoding or "").strip().lower()
|
||||
has_json_body = contract.body.json_body is not None
|
||||
has_raw_body = bool(str(contract.body.body_bytes_b64 or "").strip())
|
||||
has_body = has_json_body or has_raw_body
|
||||
return (
|
||||
((not has_body and content_encoding == "") or has_body)
|
||||
and (not has_json_body or content_encoding in {"", "gzip"} or has_raw_body)
|
||||
and is_remote_proxy_supported(contract.proxy)
|
||||
)
|
||||
|
||||
|
||||
def build_execution_plan_body(
|
||||
payload: Any,
|
||||
*,
|
||||
content_type: str | None = None,
|
||||
) -> ExecutionPlanBody:
|
||||
normalized_content_type = str(content_type or "").strip().lower()
|
||||
|
||||
if isinstance(payload, dict):
|
||||
return ExecutionPlanBody(json_body=payload)
|
||||
|
||||
if isinstance(payload, list) and "json" in normalized_content_type:
|
||||
return ExecutionPlanBody(json_body=payload)
|
||||
|
||||
if isinstance(payload, (bytes, bytearray, memoryview)):
|
||||
return ExecutionPlanBody(
|
||||
body_bytes_b64=base64.b64encode(bytes(payload)).decode("ascii")
|
||||
)
|
||||
|
||||
if isinstance(payload, str):
|
||||
return ExecutionPlanBody(
|
||||
body_bytes_b64=base64.b64encode(payload.encode("utf-8")).decode("ascii")
|
||||
)
|
||||
|
||||
if payload is None:
|
||||
return ExecutionPlanBody()
|
||||
|
||||
return ExecutionPlanBody(
|
||||
body_bytes_b64=base64.b64encode(
|
||||
json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
|
||||
).decode("ascii")
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PreparedExecutionPlan:
|
||||
"""本地执行所需的运行时上下文;`contract` 是可序列化的稳定边界。"""
|
||||
|
||||
contract: ExecutionPlan
|
||||
payload: dict[str, Any]
|
||||
headers: dict[str, str]
|
||||
upstream_is_stream: bool
|
||||
needs_conversion: bool
|
||||
provider_type: str
|
||||
request_timeout: float
|
||||
delegate_config: dict[str, Any] | None = None
|
||||
proxy_config: dict[str, Any] | None = None
|
||||
envelope: Any = None
|
||||
selected_base_url: str | None = None
|
||||
client_content_encoding: str | None = None
|
||||
proxy_info: dict[str, Any] | None = None
|
||||
|
||||
@property
|
||||
def remote_eligible(self) -> bool:
|
||||
return is_remote_contract_eligible(self.contract)
|
||||
@@ -0,0 +1,245 @@
|
||||
"""
|
||||
Rust executor 客户端
|
||||
|
||||
当前实现支持两类同步请求:
|
||||
- 普通 JSON 响应
|
||||
- 上游流式但客户端同步聚合的请求(返回原始流字节给 Python 后处理)
|
||||
|
||||
整体仍然是 best-effort:
|
||||
- 当 backend 未就绪或连接失败时,由调用方决定是否回退到 Python 路径
|
||||
- 传输协议优先支持 Unix Socket + HTTP,其次是 TCP + HTTP
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from src.config.settings import config
|
||||
from src.services.request.executor_plan import ExecutionPlan
|
||||
|
||||
|
||||
class RustExecutorClientError(RuntimeError):
|
||||
"""Rust executor 客户端错误。"""
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class RustExecutorSyncResult:
|
||||
status_code: int
|
||||
response_json: Any = None
|
||||
headers: dict[str, str] = field(default_factory=dict)
|
||||
provider_response_json: Any = None
|
||||
response_body_bytes: bytes | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class RustExecutorStreamResult:
|
||||
status_code: int
|
||||
headers: dict[str, str]
|
||||
byte_iterator: AsyncIterator[bytes]
|
||||
response_ctx: Any
|
||||
|
||||
|
||||
class _RustExecutorManagedStreamContext:
|
||||
def __init__(self, client: httpx.AsyncClient, response_ctx: Any) -> None:
|
||||
self._client = client
|
||||
self._response_ctx = response_ctx
|
||||
self._closed = False
|
||||
|
||||
async def __aexit__(self, exc_type: object, exc: object, tb: object) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
try:
|
||||
await self._response_ctx.__aexit__(exc_type, exc, tb)
|
||||
finally:
|
||||
await self._client.aclose()
|
||||
|
||||
|
||||
class RustExecutorClient:
|
||||
"""Python 控制面访问 Rust executor 的轻量客户端。"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
transport: str | None = None,
|
||||
base_url: str | None = None,
|
||||
socket_path: str | None = None,
|
||||
request_timeout: float | None = None,
|
||||
) -> None:
|
||||
self.transport = (transport or config.executor_transport).strip().lower()
|
||||
self.base_url = (base_url or config.executor_base_url).strip()
|
||||
self.socket_path = (socket_path or config.executor_socket_path).strip()
|
||||
self.request_timeout = (
|
||||
request_timeout if request_timeout is not None else config.executor_request_timeout
|
||||
)
|
||||
|
||||
def _build_client(self, *, streaming: bool = False) -> httpx.AsyncClient:
|
||||
if streaming:
|
||||
timeout = httpx.Timeout(
|
||||
connect=min(self.request_timeout, 30.0),
|
||||
read=None,
|
||||
write=self.request_timeout,
|
||||
pool=self.request_timeout,
|
||||
)
|
||||
else:
|
||||
timeout = httpx.Timeout(self.request_timeout)
|
||||
if self.transport == "unix_socket":
|
||||
if not self.socket_path:
|
||||
raise RustExecutorClientError("EXECUTOR_SOCKET_PATH is required for unix_socket")
|
||||
transport = httpx.AsyncHTTPTransport(uds=self.socket_path, retries=0)
|
||||
return httpx.AsyncClient(
|
||||
transport=transport,
|
||||
base_url=self.base_url,
|
||||
timeout=timeout,
|
||||
)
|
||||
|
||||
if self.transport != "tcp":
|
||||
raise RustExecutorClientError(f"Unsupported executor transport: {self.transport}")
|
||||
|
||||
return httpx.AsyncClient(
|
||||
base_url=self.base_url,
|
||||
timeout=timeout,
|
||||
transport=httpx.AsyncHTTPTransport(retries=0),
|
||||
)
|
||||
|
||||
async def execute_sync_json(self, plan: ExecutionPlan) -> RustExecutorSyncResult:
|
||||
async with self._build_client() as client:
|
||||
response = await client.post(
|
||||
"/v1/execute/sync",
|
||||
json=plan.to_payload(),
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
|
||||
status_code = int(payload.get("status_code") or 200)
|
||||
headers = payload.get("headers") or {}
|
||||
if not isinstance(headers, dict):
|
||||
raise RustExecutorClientError("Executor response headers must be an object")
|
||||
|
||||
response_json = payload.get("response_json")
|
||||
provider_response_json = payload.get("provider_response_json")
|
||||
body_payload = payload.get("body")
|
||||
body_bytes_b64 = None
|
||||
if response_json is None and isinstance(body_payload, dict):
|
||||
response_json = body_payload.get("json_body")
|
||||
body_bytes_b64 = body_payload.get("body_bytes_b64")
|
||||
|
||||
response_body_bytes: bytes | None = None
|
||||
if body_bytes_b64 is not None:
|
||||
if not isinstance(body_bytes_b64, str):
|
||||
raise RustExecutorClientError("Executor body_bytes_b64 must be a string")
|
||||
try:
|
||||
response_body_bytes = base64.b64decode(body_bytes_b64)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise RustExecutorClientError(
|
||||
"Executor body_bytes_b64 must be valid base64"
|
||||
) from exc
|
||||
|
||||
return RustExecutorSyncResult(
|
||||
status_code=status_code,
|
||||
response_json=response_json,
|
||||
headers={str(k): str(v) for k, v in headers.items()},
|
||||
provider_response_json=provider_response_json,
|
||||
response_body_bytes=response_body_bytes,
|
||||
)
|
||||
|
||||
async def execute_stream(self, plan: ExecutionPlan) -> RustExecutorStreamResult:
|
||||
client = self._build_client(streaming=True)
|
||||
response_ctx = client.stream(
|
||||
"POST",
|
||||
"/v1/execute/stream",
|
||||
json=plan.to_payload(),
|
||||
)
|
||||
try:
|
||||
response = await response_ctx.__aenter__()
|
||||
response.raise_for_status()
|
||||
line_iter = response.aiter_lines()
|
||||
headers_frame = await self._read_first_stream_frame(line_iter)
|
||||
payload = headers_frame.get("payload")
|
||||
if not isinstance(payload, dict) or payload.get("kind") != "headers":
|
||||
raise RustExecutorClientError("Executor stream must start with headers frame")
|
||||
|
||||
status_code = int(payload.get("status_code") or 200)
|
||||
headers = payload.get("headers") or {}
|
||||
if not isinstance(headers, dict):
|
||||
raise RustExecutorClientError("Executor stream headers must be an object")
|
||||
|
||||
async def _byte_iter() -> AsyncIterator[bytes]:
|
||||
async for line in line_iter:
|
||||
if not line:
|
||||
continue
|
||||
frame = self._decode_stream_frame(line)
|
||||
frame_payload = frame["payload"]
|
||||
kind = str(frame_payload.get("kind") or "").strip().lower()
|
||||
if kind == "data":
|
||||
chunk_b64 = frame_payload.get("chunk_b64")
|
||||
if isinstance(chunk_b64, str):
|
||||
if chunk_b64:
|
||||
try:
|
||||
yield base64.b64decode(chunk_b64)
|
||||
except Exception as exc: # noqa: BLE001
|
||||
raise RustExecutorClientError(
|
||||
"Executor stream chunk_b64 must be valid base64"
|
||||
) from exc
|
||||
continue
|
||||
|
||||
text = frame_payload.get("text")
|
||||
if isinstance(text, str):
|
||||
if text:
|
||||
yield text.encode("utf-8")
|
||||
continue
|
||||
|
||||
if kind == "error":
|
||||
error = frame_payload.get("error") or {}
|
||||
message = str(error.get("message") or "executor stream error")
|
||||
raise httpx.ReadError(message)
|
||||
|
||||
if kind == "telemetry":
|
||||
continue
|
||||
|
||||
if kind == "eof":
|
||||
break
|
||||
|
||||
raise RustExecutorClientError(f"Unexpected executor stream frame kind: {kind}")
|
||||
|
||||
return RustExecutorStreamResult(
|
||||
status_code=status_code,
|
||||
headers={str(k): str(v) for k, v in headers.items()},
|
||||
byte_iterator=_byte_iter(),
|
||||
response_ctx=_RustExecutorManagedStreamContext(client, response_ctx),
|
||||
)
|
||||
except Exception:
|
||||
try:
|
||||
await response_ctx.__aexit__(None, None, None)
|
||||
finally:
|
||||
await client.aclose()
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _decode_stream_frame(line: str) -> dict[str, Any]:
|
||||
try:
|
||||
frame = json.loads(line)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise RustExecutorClientError("Executor stream frame must be valid JSON") from exc
|
||||
if not isinstance(frame, dict):
|
||||
raise RustExecutorClientError("Executor stream frame must be an object")
|
||||
payload = frame.get("payload")
|
||||
if not isinstance(payload, dict):
|
||||
raise RustExecutorClientError("Executor stream frame payload must be an object")
|
||||
return frame
|
||||
|
||||
async def _read_first_stream_frame(
|
||||
self,
|
||||
line_iter: AsyncIterator[str],
|
||||
) -> dict[str, Any]:
|
||||
async for line in line_iter:
|
||||
if not line:
|
||||
continue
|
||||
return self._decode_stream_frame(line)
|
||||
raise RustExecutorClientError("Executor stream ended before headers frame")
|
||||
@@ -4,6 +4,7 @@ from typing import Any
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from src.config.settings import config
|
||||
from src.core.logger import logger
|
||||
from src.models.database import ProviderAPIKey, ProviderEndpoint
|
||||
from src.services.usage.service import UsageService
|
||||
@@ -29,8 +30,10 @@ class VideoTaskCancelService:
|
||||
- None on success
|
||||
- upstream httpx.Response when upstream returns an error (status >= 400)
|
||||
"""
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import httpx
|
||||
from fastapi import HTTPException
|
||||
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
@@ -94,12 +97,100 @@ class VideoTaskCancelService:
|
||||
header_rules=getattr(endpoint, "header_rules", None),
|
||||
)
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
async def _try_rust_cancel_response(
|
||||
*,
|
||||
method: str,
|
||||
url: str,
|
||||
request_headers: dict[str, str],
|
||||
body: Any,
|
||||
content_type: str | None = None,
|
||||
) -> httpx.Response | None:
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanTimeouts,
|
||||
build_execution_plan_body,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
final_headers = dict(request_headers)
|
||||
if (
|
||||
body is not None
|
||||
and content_type
|
||||
and not any(str(key).lower() == "content-type" for key in final_headers)
|
||||
):
|
||||
final_headers["content-type"] = content_type
|
||||
|
||||
try:
|
||||
result = await RustExecutorClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=str(getattr(task, "request_id", "") or task_id),
|
||||
candidate_id=None,
|
||||
provider_name=provider_format_norm.split(":", 1)[0],
|
||||
provider_id=str(getattr(endpoint, "provider_id", "") or ""),
|
||||
endpoint_id=str(getattr(endpoint, "id", "") or ""),
|
||||
key_id=str(getattr(key, "id", "") or ""),
|
||||
method=method,
|
||||
url=url,
|
||||
headers=final_headers,
|
||||
body=build_execution_plan_body(body, content_type=content_type),
|
||||
stream=False,
|
||||
provider_api_format=provider_format,
|
||||
client_api_format=provider_format,
|
||||
model_name=str(getattr(task, "model", "") or ""),
|
||||
content_type=content_type,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=300_000,
|
||||
write_ms=300_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=300_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError) as exc:
|
||||
logger.warning(
|
||||
"[VideoCancel] Rust executor unavailable task={} method={} url={}: {}",
|
||||
getattr(task, "id", task_id),
|
||||
method,
|
||||
url,
|
||||
str(exc),
|
||||
)
|
||||
return None
|
||||
|
||||
response_headers = dict(result.headers)
|
||||
if result.response_json is not None:
|
||||
response_headers.setdefault("content-type", "application/json")
|
||||
response_body = json.dumps(result.response_json, ensure_ascii=False).encode("utf-8")
|
||||
elif result.response_body_bytes is not None:
|
||||
response_body = result.response_body_bytes
|
||||
else:
|
||||
response_body = b""
|
||||
|
||||
return httpx.Response(
|
||||
status_code=result.status_code,
|
||||
request=httpx.Request(method, url, headers=final_headers),
|
||||
headers=response_headers,
|
||||
content=response_body,
|
||||
)
|
||||
|
||||
if provider_format_norm.startswith("openai:"):
|
||||
upstream_url = build_provider_url(endpoint, is_stream=False, key=key)
|
||||
upstream_url = f"{upstream_url.rstrip('/')}/{str(external_task_id).lstrip('/')}"
|
||||
response = await client.delete(upstream_url, headers=headers)
|
||||
response = await _try_rust_cancel_response(
|
||||
method="DELETE",
|
||||
url=upstream_url,
|
||||
request_headers=headers,
|
||||
body=None,
|
||||
)
|
||||
if response is None:
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.delete(upstream_url, headers=headers)
|
||||
if response.status_code >= 400:
|
||||
return response
|
||||
|
||||
@@ -125,7 +216,16 @@ class VideoTaskCancelService:
|
||||
headers.pop("x-goog-api-key", None)
|
||||
headers[auth_info.auth_header] = auth_info.auth_value
|
||||
|
||||
response = await client.post(upstream_url, headers=headers, json={})
|
||||
response = await _try_rust_cancel_response(
|
||||
method="POST",
|
||||
url=upstream_url,
|
||||
request_headers=headers,
|
||||
body={},
|
||||
content_type="application/json",
|
||||
)
|
||||
if response is None:
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.post(upstream_url, headers=headers, json={})
|
||||
if response.status_code >= 400:
|
||||
return response
|
||||
|
||||
|
||||
@@ -37,6 +37,7 @@ from src.core.video_utils import (
|
||||
from src.database import create_session
|
||||
from src.models.database import ProviderAPIKey, ProviderEndpoint, VideoTask
|
||||
from src.services.provider.auth import get_provider_auth
|
||||
from src.services.provider.provider_context import resolve_provider_proxy
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
@@ -55,6 +56,9 @@ class VideoPollContext:
|
||||
poll_interval_seconds: int
|
||||
max_poll_count: int
|
||||
current_status: str
|
||||
proxy_config: dict[str, Any] | None = None
|
||||
delegate_config: dict[str, Any] | None = None
|
||||
proxy_snapshot: Any = None
|
||||
|
||||
|
||||
# 永久性错误指示词(用于降级判断,不应重试)
|
||||
@@ -200,6 +204,10 @@ class VideoTaskPollerAdapter:
|
||||
else:
|
||||
auth_info = None
|
||||
headers = self._build_headers(provider_format, upstream_key, endpoint, auth_info)
|
||||
proxy_config, delegate_config, proxy_snapshot = await self._build_transport_context(
|
||||
endpoint=endpoint,
|
||||
key=key,
|
||||
)
|
||||
|
||||
return VideoPollContext(
|
||||
task_id=task.id,
|
||||
@@ -213,6 +221,9 @@ class VideoTaskPollerAdapter:
|
||||
poll_interval_seconds=task.poll_interval_seconds,
|
||||
max_poll_count=task.max_poll_count,
|
||||
current_status=task.status,
|
||||
proxy_config=proxy_config,
|
||||
delegate_config=delegate_config,
|
||||
proxy_snapshot=proxy_snapshot,
|
||||
)
|
||||
|
||||
async def poll_task_http(self, ctx: VideoPollContext) -> InternalVideoPollResult:
|
||||
@@ -351,13 +362,18 @@ class VideoTaskPollerAdapter:
|
||||
"""使用上下文进行 OpenAI 轮询(不需要数据库)"""
|
||||
url = self._build_openai_url(ctx.base_url, ctx.external_task_id)
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.get(url, headers=ctx.headers)
|
||||
if response.status_code >= 400:
|
||||
error_message = self._extract_error_message(response.text, response.status_code)
|
||||
raise PollHTTPError(response.status_code, error_message)
|
||||
payload = await self._try_rust_poll_payload(ctx=ctx, url=url)
|
||||
if payload is None:
|
||||
client = await HTTPClientPool.get_upstream_client(
|
||||
ctx.delegate_config,
|
||||
proxy_config=ctx.proxy_config,
|
||||
)
|
||||
response = await client.get(url, headers=ctx.headers)
|
||||
if response.status_code >= 400:
|
||||
error_message = self._extract_error_message(response.text, response.status_code)
|
||||
raise PollHTTPError(response.status_code, error_message)
|
||||
payload = response.json()
|
||||
|
||||
payload = response.json()
|
||||
return self._openai_normalizer.video_poll_to_internal(payload)
|
||||
|
||||
async def _poll_gemini_with_context(self, ctx: VideoPollContext) -> InternalVideoPollResult:
|
||||
@@ -372,21 +388,111 @@ class VideoTaskPollerAdapter:
|
||||
url,
|
||||
)
|
||||
|
||||
client = await HTTPClientPool.get_default_client_async()
|
||||
response = await client.get(url, headers=ctx.headers)
|
||||
if response.status_code >= 400:
|
||||
logger.warning(
|
||||
"[VideoPoller] Gemini poll failed: task={} status={} response={}",
|
||||
ctx.task_id,
|
||||
response.status_code,
|
||||
response.text[:500] if response.text else "(empty)",
|
||||
payload = await self._try_rust_poll_payload(ctx=ctx, url=url)
|
||||
if payload is None:
|
||||
client = await HTTPClientPool.get_upstream_client(
|
||||
ctx.delegate_config,
|
||||
proxy_config=ctx.proxy_config,
|
||||
)
|
||||
error_message = self._extract_error_message(response.text, response.status_code)
|
||||
raise PollHTTPError(response.status_code, error_message)
|
||||
response = await client.get(url, headers=ctx.headers)
|
||||
if response.status_code >= 400:
|
||||
logger.warning(
|
||||
"[VideoPoller] Gemini poll failed: task={} status={} response={}",
|
||||
ctx.task_id,
|
||||
response.status_code,
|
||||
response.text[:500] if response.text else "(empty)",
|
||||
)
|
||||
error_message = self._extract_error_message(response.text, response.status_code)
|
||||
raise PollHTTPError(response.status_code, error_message)
|
||||
payload = response.json()
|
||||
|
||||
payload = response.json()
|
||||
return self._gemini_normalizer.video_poll_to_internal(payload)
|
||||
|
||||
async def _try_rust_poll_payload(
|
||||
self,
|
||||
*,
|
||||
ctx: VideoPollContext,
|
||||
url: str,
|
||||
) -> dict[str, Any] | None:
|
||||
import httpx
|
||||
|
||||
from src.services.request.executor_plan import (
|
||||
ExecutionPlan,
|
||||
ExecutionPlanBody,
|
||||
ExecutionPlanTimeouts,
|
||||
)
|
||||
from src.services.request.rust_executor_client import (
|
||||
RustExecutorClient,
|
||||
RustExecutorClientError,
|
||||
)
|
||||
|
||||
if config.executor_backend != "rust":
|
||||
return None
|
||||
|
||||
try:
|
||||
result = await RustExecutorClient().execute_sync_json(
|
||||
ExecutionPlan(
|
||||
request_id=f"video-poll-{ctx.task_id}",
|
||||
candidate_id=None,
|
||||
provider_name=ctx.provider_api_format.split(":", 1)[0],
|
||||
provider_id="",
|
||||
endpoint_id="",
|
||||
key_id="",
|
||||
method="GET",
|
||||
url=url,
|
||||
headers=dict(ctx.headers),
|
||||
body=ExecutionPlanBody(),
|
||||
stream=False,
|
||||
provider_api_format=ctx.provider_api_format,
|
||||
client_api_format=ctx.provider_api_format,
|
||||
model_name="video-poll",
|
||||
proxy=ctx.proxy_snapshot,
|
||||
timeouts=ExecutionPlanTimeouts(
|
||||
connect_ms=30_000,
|
||||
read_ms=300_000,
|
||||
write_ms=300_000,
|
||||
pool_ms=30_000,
|
||||
total_ms=300_000,
|
||||
),
|
||||
)
|
||||
)
|
||||
except (RustExecutorClientError, httpx.HTTPError, json.JSONDecodeError, ValueError) as exc:
|
||||
logger.warning(
|
||||
"[VideoPoller] Rust poll fallback task={} url={} error={}",
|
||||
ctx.task_id,
|
||||
url,
|
||||
sanitize_error_message(str(exc)),
|
||||
)
|
||||
return None
|
||||
|
||||
if result.status_code >= 400:
|
||||
response_text = ""
|
||||
if result.response_json is not None:
|
||||
response_text = json.dumps(result.response_json, ensure_ascii=False)
|
||||
elif result.response_body_bytes is not None:
|
||||
response_text = result.response_body_bytes.decode("utf-8", errors="replace")
|
||||
raise PollHTTPError(
|
||||
result.status_code,
|
||||
self._extract_error_message(response_text, result.status_code),
|
||||
)
|
||||
|
||||
if isinstance(result.response_json, dict):
|
||||
return result.response_json
|
||||
|
||||
if result.response_body_bytes is not None:
|
||||
try:
|
||||
payload = json.loads(result.response_body_bytes.decode("utf-8"))
|
||||
if isinstance(payload, dict):
|
||||
return payload
|
||||
except Exception:
|
||||
logger.warning(
|
||||
"[VideoPoller] Rust poll returned non-json body task={} url={}",
|
||||
ctx.task_id,
|
||||
url,
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
# ==================== 旧版方法(保留兼容性)====================
|
||||
|
||||
async def poll_single_task(
|
||||
@@ -591,6 +697,55 @@ class VideoTaskPollerAdapter:
|
||||
headers[auth_info.auth_header] = auth_info.auth_value
|
||||
return headers
|
||||
|
||||
async def _build_transport_context(
|
||||
self,
|
||||
*,
|
||||
endpoint: ProviderEndpoint,
|
||||
key: ProviderAPIKey,
|
||||
) -> tuple[dict[str, Any] | None, dict[str, Any] | None, Any]:
|
||||
from src.services.proxy_node.resolver import (
|
||||
build_proxy_url_async,
|
||||
get_system_proxy_config_async,
|
||||
resolve_delegate_config_async,
|
||||
resolve_effective_proxy,
|
||||
resolve_proxy_info_async,
|
||||
)
|
||||
from src.services.request.executor_plan import ExecutionProxySnapshot
|
||||
|
||||
try:
|
||||
effective_proxy = resolve_effective_proxy(
|
||||
resolve_provider_proxy(endpoint=endpoint, key=key),
|
||||
getattr(key, "proxy", None),
|
||||
)
|
||||
if not effective_proxy or not effective_proxy.get("enabled", True):
|
||||
effective_proxy = await get_system_proxy_config_async()
|
||||
|
||||
delegate_cfg = await resolve_delegate_config_async(effective_proxy)
|
||||
proxy_url: str | None = None
|
||||
if effective_proxy and not (delegate_cfg and delegate_cfg.get("tunnel")):
|
||||
proxy_url = await build_proxy_url_async(effective_proxy)
|
||||
|
||||
proxy_info = await resolve_proxy_info_async(effective_proxy)
|
||||
proxy_snapshot = ExecutionProxySnapshot.from_proxy_info(
|
||||
proxy_info,
|
||||
proxy_url=proxy_url,
|
||||
mode_override="tunnel" if delegate_cfg and delegate_cfg.get("tunnel") else None,
|
||||
node_id_override=(
|
||||
str(delegate_cfg.get("node_id") or "").strip() or None
|
||||
if delegate_cfg and delegate_cfg.get("tunnel")
|
||||
else None
|
||||
),
|
||||
)
|
||||
return effective_proxy, delegate_cfg, proxy_snapshot
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[VideoPoller] Failed to build transport context endpoint={} key={}: {}",
|
||||
getattr(endpoint, "id", None),
|
||||
getattr(key, "id", None),
|
||||
sanitize_error_message(str(exc)),
|
||||
)
|
||||
return None, None, None
|
||||
|
||||
def _get_endpoint(self, db: Session, endpoint_id: str) -> ProviderEndpoint:
|
||||
endpoint = db.query(ProviderEndpoint).filter(ProviderEndpoint.id == endpoint_id).first()
|
||||
if not endpoint:
|
||||
|
||||
Reference in New Issue
Block a user