mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
refactor: 调度器迁移至独立模块,消除 services->api 反向依赖
- 将调度器相关模块从 src/services/cache/ 迁移到 src/services/scheduling/ - 下沉类型定义到 core 层: AccessRestrictions, ProviderAuthInfo, ParsedChunk/StreamStats, 视频工具函数 - 提取 thinking_cache 签名缓存到 core/api_format/conversion/ - 提取 provider 认证逻辑到 services/provider/auth - 提取遥测记录到 services/usage/telemetry - 提取 models 列表缓存到 services/cache/model_list_cache - 更新所有引用方的 import 路径及相关测试
This commit is contained in:
113
src/core/access_restrictions.py
Normal file
113
src/core/access_restrictions.py
Normal file
@@ -0,0 +1,113 @@
|
||||
"""
|
||||
API Key / User 访问限制数据类型。
|
||||
|
||||
从 api/base/models_service.py 下沉到 core 层,
|
||||
消除 services→api 的反向依赖。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
from src.core.logger import logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from src.models.database import ApiKey, User
|
||||
|
||||
|
||||
def _safe_normalize_signature(value: str) -> str:
|
||||
"""归一化 endpoint signature,解析失败时原样返回(小写)。"""
|
||||
try:
|
||||
return normalize_signature_key(value)
|
||||
except ValueError:
|
||||
logger.warning("[AccessRestrictions] 无法归一化 API 格式 '{}', 原样使用小写形式", value)
|
||||
return value.strip().lower()
|
||||
|
||||
|
||||
@dataclass
|
||||
class AccessRestrictions:
|
||||
"""API Key 或 User 的访问限制"""
|
||||
|
||||
allowed_providers: list[str] | None = None # 允许的 Provider ID 列表
|
||||
allowed_models: list[str] | None = None # 允许的模型名称列表
|
||||
allowed_api_formats: list[str] | None = None # 允许的 API 格式列表
|
||||
|
||||
@classmethod
|
||||
def from_api_key_and_user(cls, api_key: ApiKey | None, user: User | None) -> AccessRestrictions:
|
||||
"""
|
||||
从 API Key 和 User 合并访问限制
|
||||
|
||||
限制逻辑:
|
||||
- API Key 的限制优先于 User 的限制
|
||||
- 如果 API Key 有限制,使用 API Key 的限制
|
||||
- 如果 API Key 无限制但 User 有限制,使用 User 的限制
|
||||
- 两者都无限制则返回空限制
|
||||
"""
|
||||
allowed_providers: list[str] | None = None
|
||||
allowed_models: list[str] | None = None
|
||||
allowed_api_formats: list[str] | None = None
|
||||
|
||||
# 优先使用 API Key 的限制
|
||||
if api_key:
|
||||
if api_key.allowed_providers is not None:
|
||||
allowed_providers = api_key.allowed_providers
|
||||
if api_key.allowed_models is not None:
|
||||
allowed_models = api_key.allowed_models
|
||||
if api_key.allowed_api_formats is not None:
|
||||
allowed_api_formats = api_key.allowed_api_formats
|
||||
|
||||
# 如果 API Key 没有限制,检查 User 的限制
|
||||
if user:
|
||||
if allowed_providers is None and user.allowed_providers is not None:
|
||||
allowed_providers = user.allowed_providers
|
||||
if allowed_models is None and user.allowed_models is not None:
|
||||
allowed_models = user.allowed_models
|
||||
if allowed_api_formats is None and user.allowed_api_formats is not None:
|
||||
allowed_api_formats = user.allowed_api_formats
|
||||
|
||||
return cls(
|
||||
allowed_providers=allowed_providers,
|
||||
allowed_models=allowed_models,
|
||||
allowed_api_formats=allowed_api_formats,
|
||||
)
|
||||
|
||||
def is_api_format_allowed(self, api_format: str) -> bool:
|
||||
"""
|
||||
检查 API 格式是否被允许
|
||||
|
||||
Args:
|
||||
api_format: endpoint signature(如 "openai:chat")
|
||||
|
||||
Returns:
|
||||
True 如果格式被允许,False 否则
|
||||
"""
|
||||
if self.allowed_api_formats is None:
|
||||
return True
|
||||
target = _safe_normalize_signature(api_format)
|
||||
allowed = {_safe_normalize_signature(f) for f in self.allowed_api_formats if f}
|
||||
return target in allowed
|
||||
|
||||
def is_model_allowed(self, model_id: str, provider_id: str) -> bool:
|
||||
"""
|
||||
检查模型是否被允许访问
|
||||
|
||||
Args:
|
||||
model_id: 模型 ID
|
||||
provider_id: Provider ID
|
||||
|
||||
Returns:
|
||||
True 如果模型被允许,False 否则
|
||||
"""
|
||||
# 检查 Provider 限制
|
||||
if self.allowed_providers is not None:
|
||||
if provider_id not in self.allowed_providers:
|
||||
return False
|
||||
|
||||
# 检查模型限制
|
||||
if self.allowed_models is not None:
|
||||
if model_id not in self.allowed_models:
|
||||
return False
|
||||
|
||||
return True
|
||||
@@ -1425,7 +1425,7 @@ class GeminiNormalizer(FormatNormalizer):
|
||||
)
|
||||
|
||||
try:
|
||||
from src.services.provider.adapters.antigravity.signature_cache import (
|
||||
from src.core.api_format.conversion.thinking_cache import (
|
||||
signature_cache,
|
||||
)
|
||||
|
||||
@@ -1565,7 +1565,7 @@ class GeminiNormalizer(FormatNormalizer):
|
||||
from src.core.api_format.conversion.constants import DUMMY_THOUGHT_SIGNATURE
|
||||
|
||||
try:
|
||||
from src.services.provider.adapters.antigravity.signature_cache import signature_cache
|
||||
from src.core.api_format.conversion.thinking_cache import signature_cache
|
||||
|
||||
cached_or_dummy = signature_cache.get_or_dummy(model, text_val)
|
||||
|
||||
|
||||
237
src/core/api_format/conversion/thinking_cache.py
Normal file
237
src/core/api_format/conversion/thinking_cache.py
Normal file
@@ -0,0 +1,237 @@
|
||||
"""Thinking block signature cache (triple-layer).
|
||||
|
||||
三层缓存设计:
|
||||
Layer 1: tool_use_id -> thoughtSignature (工具调用签名恢复)
|
||||
Layer 2: signature -> model_family (跨模型兼容校验)
|
||||
Layer 3: session_id -> latest signature (会话级签名追踪 + rewind 检测)
|
||||
|
||||
同时保留原有的 model:text -> signature 兼容层。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import threading
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from src.core.api_format.conversion.constants import DUMMY_THOUGHT_SIGNATURE
|
||||
|
||||
# 签名最小长度阈值
|
||||
MIN_SIGNATURE_LENGTH = 50
|
||||
|
||||
# TTL: 2 小时
|
||||
_SIGNATURE_TTL_SECONDS = 2 * 60 * 60
|
||||
|
||||
# 各层缓存上限
|
||||
_TOOL_CACHE_LIMIT = 500
|
||||
_FAMILY_CACHE_LIMIT = 200
|
||||
_SESSION_CACHE_LIMIT = 1000
|
||||
_TEXT_CACHE_LIMIT = 1000
|
||||
|
||||
|
||||
class _CacheEntry:
|
||||
"""带时间戳的缓存条目,支持 TTL 过期。"""
|
||||
|
||||
__slots__ = ("data", "created_at")
|
||||
|
||||
def __init__(self, data: Any) -> None:
|
||||
self.data = data
|
||||
self.created_at: float = time.monotonic()
|
||||
|
||||
def is_expired(self, now: float | None = None) -> bool:
|
||||
return ((now or time.monotonic()) - self.created_at) > _SIGNATURE_TTL_SECONDS
|
||||
|
||||
|
||||
class _SessionEntry:
|
||||
"""Session 层缓存数据,包含消息计数用于 rewind 检测。"""
|
||||
|
||||
__slots__ = ("signature", "message_count")
|
||||
|
||||
def __init__(self, signature: str, message_count: int) -> None:
|
||||
self.signature = signature
|
||||
self.message_count = message_count
|
||||
|
||||
|
||||
class ThinkingSignatureCache:
|
||||
"""Triple-layer thinking signature cache.
|
||||
|
||||
Layer 1 (tool): tool_use_id -> thoughtSignature
|
||||
当客户端(如 OpenCode) 在 tool_result 中丢弃了 signature 时用于恢复。
|
||||
|
||||
Layer 2 (family): signature -> model_family
|
||||
防止跨模型签名污染(Claude 签名不能用在 Gemini 上)。
|
||||
|
||||
Layer 3 (session): session_id -> latest signature + message_count
|
||||
会话级追踪,支持 rewind 检测(用户删除消息后不会注入来自"未来"的签名)。
|
||||
|
||||
Legacy (text): SHA256(model + text) -> signature
|
||||
向后兼容的 get_or_dummy() 接口。
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._tool_sigs: dict[str, _CacheEntry] = {}
|
||||
self._families: dict[str, _CacheEntry] = {}
|
||||
self._sessions: dict[str, _CacheEntry] = {}
|
||||
self._text_sigs: dict[str, _CacheEntry] = {}
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# ===== Layer 1: Tool Use ID -> Signature =====
|
||||
|
||||
def cache_tool_signature(self, tool_use_id: str, signature: str) -> None:
|
||||
"""缓存工具调用对应的 thinking signature。"""
|
||||
if len(signature) < MIN_SIGNATURE_LENGTH:
|
||||
return
|
||||
with self._lock:
|
||||
self._tool_sigs[tool_use_id] = _CacheEntry(signature)
|
||||
if len(self._tool_sigs) > _TOOL_CACHE_LIMIT:
|
||||
self._prune(self._tool_sigs, limit=_TOOL_CACHE_LIMIT)
|
||||
|
||||
def get_tool_signature(self, tool_use_id: str) -> str | None:
|
||||
"""查找工具调用对应的 signature。"""
|
||||
with self._lock:
|
||||
entry = self._tool_sigs.get(tool_use_id)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry.is_expired():
|
||||
self._tool_sigs.pop(tool_use_id, None)
|
||||
return None
|
||||
return entry.data
|
||||
|
||||
# ===== Layer 2: Signature -> Model Family =====
|
||||
|
||||
def cache_thinking_family(self, signature: str, family: str) -> None:
|
||||
"""记录 signature 所属的模型家族。"""
|
||||
if len(signature) < MIN_SIGNATURE_LENGTH:
|
||||
return
|
||||
with self._lock:
|
||||
self._families[signature] = _CacheEntry(family)
|
||||
if len(self._families) > _FAMILY_CACHE_LIMIT:
|
||||
self._prune(self._families, limit=_FAMILY_CACHE_LIMIT)
|
||||
|
||||
def get_signature_family(self, signature: str) -> str | None:
|
||||
"""查找 signature 所属的模型家族。"""
|
||||
with self._lock:
|
||||
entry = self._families.get(signature)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry.is_expired():
|
||||
self._families.pop(signature, None)
|
||||
return None
|
||||
return entry.data
|
||||
|
||||
# ===== Layer 3: Session ID -> Latest Signature =====
|
||||
|
||||
def cache_session_signature(
|
||||
self, session_id: str, signature: str, message_count: int = 0
|
||||
) -> None:
|
||||
"""存储会话的最新 thinking signature。
|
||||
|
||||
Rewind 检测:当 message_count 小于已缓存值时,说明用户删除了消息,
|
||||
强制更新签名以避免注入来自"未来"的签名。
|
||||
"""
|
||||
if len(signature) < MIN_SIGNATURE_LENGTH:
|
||||
return
|
||||
|
||||
with self._lock:
|
||||
existing = self._sessions.get(session_id)
|
||||
should_store = True
|
||||
|
||||
if existing and not existing.is_expired():
|
||||
entry: _SessionEntry = existing.data
|
||||
if message_count < entry.message_count:
|
||||
# Rewind detected: 用户删除了消息,强制更新
|
||||
pass
|
||||
elif message_count == entry.message_count:
|
||||
# 同一轮消息:仅当新签名更长(更完整)时才替换
|
||||
should_store = len(signature) > len(entry.signature)
|
||||
# else: 正常递增,更新
|
||||
|
||||
if should_store:
|
||||
self._sessions[session_id] = _CacheEntry(_SessionEntry(signature, message_count))
|
||||
if len(self._sessions) > _SESSION_CACHE_LIMIT:
|
||||
self._prune(self._sessions, limit=_SESSION_CACHE_LIMIT)
|
||||
|
||||
def get_session_signature(self, session_id: str) -> str | None:
|
||||
"""获取会话的最新 thinking signature。"""
|
||||
with self._lock:
|
||||
entry = self._sessions.get(session_id)
|
||||
if entry is None:
|
||||
return None
|
||||
if entry.is_expired():
|
||||
self._sessions.pop(session_id, None)
|
||||
return None
|
||||
return entry.data.signature
|
||||
|
||||
# ===== Legacy: model:text -> signature(向后兼容) =====
|
||||
|
||||
def get_or_dummy(self, model: str, thinking_text: str) -> str | None:
|
||||
"""Legacy: 根据 model + thinking_text 查找 signature。
|
||||
|
||||
Gemini 模型在未命中时返回 DUMMY_THOUGHT_SIGNATURE(跳过验证)。
|
||||
"""
|
||||
key = self._text_key(model, thinking_text)
|
||||
with self._lock:
|
||||
entry = self._text_sigs.get(key)
|
||||
if entry is not None:
|
||||
if entry.is_expired():
|
||||
self._text_sigs.pop(key, None)
|
||||
else:
|
||||
return entry.data
|
||||
if str(model).startswith("gemini-"):
|
||||
return DUMMY_THOUGHT_SIGNATURE
|
||||
return None
|
||||
|
||||
def cache(self, model: str, thinking_text: str, signature: str) -> None:
|
||||
"""Legacy: 缓存 model + thinking_text -> signature。"""
|
||||
if len(signature) < MIN_SIGNATURE_LENGTH:
|
||||
return
|
||||
|
||||
key = self._text_key(model, thinking_text)
|
||||
with self._lock:
|
||||
if key in self._text_sigs:
|
||||
self._text_sigs[key] = _CacheEntry(signature)
|
||||
return
|
||||
|
||||
if len(self._text_sigs) >= _TEXT_CACHE_LIMIT:
|
||||
# FIFO 淘汰 1/4
|
||||
evict_n = max(1, _TEXT_CACHE_LIMIT // 4)
|
||||
for k in list(self._text_sigs.keys())[:evict_n]:
|
||||
self._text_sigs.pop(k, None)
|
||||
|
||||
self._text_sigs[key] = _CacheEntry(signature)
|
||||
|
||||
# ===== Utilities =====
|
||||
|
||||
@staticmethod
|
||||
def _text_key(model: str, thinking_text: str) -> str:
|
||||
content = f"{model}\x00{thinking_text}"
|
||||
return hashlib.sha256(content.encode("utf-8")).hexdigest()[:32]
|
||||
|
||||
@staticmethod
|
||||
def _prune(d: dict[str, _CacheEntry], *, limit: int | None = None) -> None:
|
||||
"""Remove expired entries and optionally enforce a size limit."""
|
||||
now = time.monotonic()
|
||||
expired = [k for k, v in d.items() if v.is_expired(now)]
|
||||
for k in expired:
|
||||
d.pop(k, None)
|
||||
|
||||
if limit is None or len(d) <= limit:
|
||||
return
|
||||
|
||||
excess = len(d) - limit
|
||||
for k, _entry in sorted(d.items(), key=lambda kv: kv[1].created_at)[:excess]:
|
||||
d.pop(k, None)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""清空所有缓存层(用于测试或手动重置)。"""
|
||||
with self._lock:
|
||||
self._tool_sigs.clear()
|
||||
self._families.clear()
|
||||
self._sessions.clear()
|
||||
self._text_sigs.clear()
|
||||
|
||||
|
||||
signature_cache = ThinkingSignatureCache()
|
||||
|
||||
__all__ = ["ThinkingSignatureCache", "signature_cache", "MIN_SIGNATURE_LENGTH"]
|
||||
@@ -8,7 +8,7 @@ from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import os
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import TYPE_CHECKING, Any, Protocol
|
||||
|
||||
from src.core.logger import logger
|
||||
from src.core.modules.base import (
|
||||
@@ -22,6 +22,17 @@ if TYPE_CHECKING:
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
|
||||
class ConfigBackend(Protocol):
|
||||
"""模块配置读写后端协议。
|
||||
|
||||
通过 ``ModuleRegistry.set_config_backend()`` 在应用启动时注入实现,
|
||||
使 core 层无需在运行时 import services 层。
|
||||
"""
|
||||
|
||||
def get_config(self, db: Any, key: str, default: Any = None) -> Any: ...
|
||||
def set_config(self, db: Any, key: str, value: Any, description: Any = None) -> Any: ...
|
||||
|
||||
|
||||
class ModuleRegistry:
|
||||
"""
|
||||
模块注册中心 - 单例模式
|
||||
@@ -34,11 +45,17 @@ class ModuleRegistry:
|
||||
"""
|
||||
|
||||
_instance: ModuleRegistry | None = None
|
||||
_config_backend: ConfigBackend | None = None
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._modules: dict[str, ModuleDefinition] = {}
|
||||
self._initialized: set[str] = set()
|
||||
|
||||
@classmethod
|
||||
def set_config_backend(cls, backend: ConfigBackend) -> None:
|
||||
"""注入配置读写后端,消除 core→services 的运行时依赖"""
|
||||
cls._config_backend = backend
|
||||
|
||||
@classmethod
|
||||
def get_instance(cls) -> ModuleRegistry:
|
||||
"""获取单例实例"""
|
||||
@@ -50,6 +67,7 @@ class ModuleRegistry:
|
||||
def reset_instance(cls) -> None:
|
||||
"""重置单例(仅用于测试)"""
|
||||
cls._instance = None
|
||||
cls._config_backend = None
|
||||
|
||||
def register(self, module: ModuleDefinition) -> None:
|
||||
"""
|
||||
@@ -123,6 +141,15 @@ class ModuleRegistry:
|
||||
|
||||
# ========== 启用状态检查(运行级)==========
|
||||
|
||||
def _get_config_backend(self) -> ConfigBackend:
|
||||
"""获取配置后端(优先使用已注入的,兜底 lazy import)"""
|
||||
if self._config_backend is not None:
|
||||
return self._config_backend
|
||||
# 兜底: 未注入时使用 lazy import(向后兼容独立脚本/测试场景)
|
||||
from src.services.system.config import SystemConfigService # noqa: lazy fallback
|
||||
|
||||
return SystemConfigService # type: ignore[return-value]
|
||||
|
||||
def is_enabled(self, name: str, db: Session) -> bool:
|
||||
"""
|
||||
检查模块是否运行启用(数据库配置)
|
||||
@@ -131,10 +158,8 @@ class ModuleRegistry:
|
||||
name: 模块名称
|
||||
db: 数据库会话
|
||||
"""
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
config_key = f"module.{name}.enabled"
|
||||
value = SystemConfigService.get_config(db, config_key, default=False)
|
||||
value = self._get_config_backend().get_config(db, config_key, default=False)
|
||||
return bool(value)
|
||||
|
||||
def set_enabled(self, name: str, enabled: bool, db: Session) -> None:
|
||||
@@ -146,15 +171,13 @@ class ModuleRegistry:
|
||||
enabled: 是否启用
|
||||
db: 数据库会话
|
||||
"""
|
||||
from src.services.system.config import SystemConfigService
|
||||
|
||||
if name not in self._modules:
|
||||
raise ValueError(f"Module [{name}] not registered")
|
||||
|
||||
config_key = f"module.{name}.enabled"
|
||||
module = self._modules[name]
|
||||
description = f"模块 [{module.metadata.display_name}] 启用状态"
|
||||
SystemConfigService.set_config(db, config_key, enabled, description)
|
||||
self._get_config_backend().set_config(db, config_key, enabled, description)
|
||||
|
||||
# ========== 激活状态检查 ==========
|
||||
|
||||
|
||||
25
src/core/provider_auth_types.py
Normal file
25
src/core/provider_auth_types.py
Normal file
@@ -0,0 +1,25 @@
|
||||
"""
|
||||
Provider 认证相关的数据类型。
|
||||
|
||||
从 api/handlers/base/request_builder.py 下沉到 core 层,
|
||||
消除 services→api 的反向依赖。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProviderAuthInfo:
|
||||
"""Provider 认证信息(用于 Service Account 等异步认证场景)"""
|
||||
|
||||
auth_header: str
|
||||
auth_value: str
|
||||
# 解密后的认证配置(用于 URL 构建等场景,避免重复解密)
|
||||
decrypted_auth_config: dict[str, Any] | None = None
|
||||
|
||||
def as_tuple(self) -> tuple[str, str]:
|
||||
"""返回 (auth_header, auth_value) 元组"""
|
||||
return (self.auth_header, self.auth_value)
|
||||
@@ -10,7 +10,6 @@ import jwt
|
||||
from src.clients.http_client import HTTPClientPool
|
||||
from src.core.logger import logger
|
||||
from src.core.provider_types import ProviderType
|
||||
from src.services.proxy_node.resolver import build_proxy_url
|
||||
|
||||
_ANTHROPIC_TOKEN_URL = "https://console.anthropic.com/v1/oauth/token"
|
||||
_GOOGLE_USERINFO_URL = "https://www.googleapis.com/oauth2/v1/userinfo?alt=json"
|
||||
@@ -22,6 +21,8 @@ def _coerce_proxy_url(proxy_config: dict[str, Any] | None) -> str | None:
|
||||
try:
|
||||
if not proxy_config.get("enabled", True):
|
||||
return None
|
||||
from src.services.proxy_node.resolver import build_proxy_url # lazy: core→services
|
||||
|
||||
return build_proxy_url(proxy_config)
|
||||
except Exception:
|
||||
return None
|
||||
@@ -363,11 +364,10 @@ async def enrich_auth_config(
|
||||
"""Enrich auth_config with non-secret metadata (email/account_id).
|
||||
|
||||
各 provider 的 enrichment 逻辑通过 register_auth_enricher 注册。
|
||||
注意: ensure_providers_bootstrapped() 在应用启动时(main.py lifespan)已显式调用。
|
||||
"""
|
||||
from src.core.provider_types import normalize_provider_type
|
||||
from src.services.provider.envelope import ensure_providers_bootstrapped
|
||||
|
||||
ensure_providers_bootstrapped()
|
||||
pt = normalize_provider_type(provider_type)
|
||||
enricher = _auth_enrichers.get(pt)
|
||||
if enricher:
|
||||
|
||||
@@ -15,9 +15,9 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
|
||||
from src.core.provider_templates.types import ProviderType
|
||||
from src.services.provider.adapters.antigravity.constants import (
|
||||
PROD_BASE_URL as ANTIGRAVITY_PROD_URL,
|
||||
)
|
||||
|
||||
# Antigravity 生产环境 URL(唯一定义点,services 层通过 re-export 引用)
|
||||
ANTIGRAVITY_PROD_URL = "https://cloudcode-pa.googleapis.com"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
|
||||
224
src/core/stream_types.py
Normal file
224
src/core/stream_types.py
Normal file
@@ -0,0 +1,224 @@
|
||||
"""
|
||||
响应解析器基类与流式统计类型。
|
||||
|
||||
从 api/handlers/base/response_parser.py 下沉到 core 层,
|
||||
消除 services→api 的反向依赖。同时提供 parser 注册表,
|
||||
允许 API 层注册具体实现,services 层通过 format_id 获取实例。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedChunk:
|
||||
"""解析后的流式数据块"""
|
||||
|
||||
# 原始数据
|
||||
raw_line: str
|
||||
event_type: str | None = None
|
||||
data: dict[str, Any] | None = None
|
||||
|
||||
# 提取的内容
|
||||
text_delta: str = ""
|
||||
is_done: bool = False
|
||||
is_error: bool = False
|
||||
error_message: str | None = None
|
||||
|
||||
# 使用量信息(通常在最后一个 chunk 中)
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
|
||||
# 响应 ID
|
||||
response_id: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class StreamStats:
|
||||
"""流式响应统计信息"""
|
||||
|
||||
# 计数
|
||||
chunk_count: int = 0
|
||||
data_count: int = 0
|
||||
|
||||
# Token 使用量
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
|
||||
# 内容
|
||||
collected_text: str = ""
|
||||
response_id: str | None = None
|
||||
|
||||
# 状态
|
||||
has_completion: bool = False
|
||||
status_code: int = 200
|
||||
error_message: str | None = None
|
||||
|
||||
# Provider 信息
|
||||
provider_name: str | None = None
|
||||
endpoint_id: str | None = None
|
||||
key_id: str | None = None
|
||||
|
||||
# 响应头和完整响应
|
||||
response_headers: dict[str, str] = field(default_factory=dict)
|
||||
final_response: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class ParsedResponse:
|
||||
"""解析后的非流式响应"""
|
||||
|
||||
# 原始响应
|
||||
raw_response: dict[str, Any]
|
||||
status_code: int
|
||||
|
||||
# 提取的内容
|
||||
text_content: str = ""
|
||||
response_id: str | None = None
|
||||
|
||||
# 使用量
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_creation_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
|
||||
# 错误信息
|
||||
is_error: bool = False
|
||||
error_type: str | None = None
|
||||
error_message: str | None = None
|
||||
# 从响应体解析出的嵌套状态码(当 HTTP 200 但响应体含错误时使用)
|
||||
embedded_status_code: int | None = None
|
||||
|
||||
|
||||
class ResponseParser(ABC):
|
||||
"""
|
||||
响应解析器基类
|
||||
|
||||
定义统一的接口来解析不同 API 格式的响应。
|
||||
子类需要实现具体的解析逻辑。
|
||||
"""
|
||||
|
||||
# 解析器名称(用于日志)
|
||||
name: str = "base"
|
||||
|
||||
# 支持的 API 格式
|
||||
api_format: str = "UNKNOWN"
|
||||
|
||||
@abstractmethod
|
||||
def parse_sse_line(self, line: str, stats: StreamStats) -> ParsedChunk | None:
|
||||
"""
|
||||
解析单行 SSE 数据
|
||||
|
||||
Args:
|
||||
line: SSE 行数据
|
||||
stats: 流统计对象(会被更新)
|
||||
|
||||
Returns:
|
||||
解析后的数据块,如果行不包含有效数据则返回 None
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def parse_response(self, response: dict[str, Any], status_code: int) -> ParsedResponse:
|
||||
"""
|
||||
解析非流式响应
|
||||
|
||||
Args:
|
||||
response: 响应 JSON
|
||||
status_code: HTTP 状态码
|
||||
|
||||
Returns:
|
||||
解析后的响应对象
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def extract_usage_from_response(self, response: dict[str, Any]) -> dict[str, int]:
|
||||
"""
|
||||
从响应中提取 token 使用量
|
||||
|
||||
Args:
|
||||
response: 响应 JSON
|
||||
|
||||
Returns:
|
||||
包含 input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens 的字典
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def extract_text_content(self, response: dict[str, Any]) -> str:
|
||||
"""
|
||||
从响应中提取文本内容
|
||||
|
||||
Args:
|
||||
response: 响应 JSON
|
||||
|
||||
Returns:
|
||||
提取的文本内容
|
||||
"""
|
||||
pass
|
||||
|
||||
def is_error_response(self, response: dict[str, Any]) -> bool:
|
||||
"""
|
||||
判断响应是否为错误响应
|
||||
|
||||
Args:
|
||||
response: 响应 JSON
|
||||
|
||||
Returns:
|
||||
是否为错误响应
|
||||
"""
|
||||
return "error" in response
|
||||
|
||||
def create_stats(self) -> StreamStats:
|
||||
"""创建新的流统计对象"""
|
||||
return StreamStats()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Parser 注册表 -- API 层注册具体实现,services 层通过 format_id 获取
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_PARSER_REGISTRY: dict[str, type[ResponseParser]] = {}
|
||||
|
||||
|
||||
def register_parser(format_id: str, parser_class: type[ResponseParser]) -> None:
|
||||
"""注册一个格式对应的 ResponseParser 实现"""
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
|
||||
normalized = normalize_signature_key(format_id)
|
||||
_PARSER_REGISTRY[normalized] = parser_class
|
||||
|
||||
|
||||
def get_parser_for_format(format_id: str) -> ResponseParser:
|
||||
"""
|
||||
根据格式 ID 获取 ResponseParser 实例
|
||||
|
||||
Args:
|
||||
format_id: endpoint signature,如 "claude:chat", "openai:cli"
|
||||
|
||||
Returns:
|
||||
ResponseParser 实例
|
||||
|
||||
Raises:
|
||||
KeyError: 格式不存在
|
||||
"""
|
||||
from src.core.api_format.signature import normalize_signature_key
|
||||
|
||||
if not _PARSER_REGISTRY:
|
||||
raise KeyError(
|
||||
f"Parser registry is empty when looking up '{format_id}'. "
|
||||
"Ensure parsers are registered at startup (import src.api.handlers.base.parsers)."
|
||||
)
|
||||
|
||||
normalized = normalize_signature_key(format_id)
|
||||
if normalized not in _PARSER_REGISTRY:
|
||||
raise KeyError(f"Unknown format: {normalized}")
|
||||
return _PARSER_REGISTRY[normalized]()
|
||||
@@ -10,6 +10,7 @@ from __future__ import annotations
|
||||
import json
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
import jwt
|
||||
@@ -102,13 +103,17 @@ class VertexAuthService:
|
||||
}
|
||||
return jwt.encode(payload, self.private_key, algorithm="RS256")
|
||||
|
||||
async def get_access_token(self) -> str:
|
||||
async def get_access_token(self, *, httpx_client_kwargs: dict[str, Any] | None = None) -> str:
|
||||
"""
|
||||
获取 Access Token(带 LRU 缓存)
|
||||
|
||||
如果缓存中有有效的 Token(距离过期超过 60 秒),直接返回。
|
||||
否则重新获取 Token。缓存采用 LRU 策略,超过 100 个条目时淘汰最旧的。
|
||||
|
||||
Args:
|
||||
httpx_client_kwargs: 传给 httpx.AsyncClient 的额外参数(如代理配置)。
|
||||
调用者(services 层)负责构建,core 层不关心代理细节。
|
||||
|
||||
Returns:
|
||||
Access Token 字符串
|
||||
|
||||
@@ -129,10 +134,10 @@ class VertexAuthService:
|
||||
try:
|
||||
signed_jwt = self._create_jwt()
|
||||
|
||||
# 使用系统默认代理(Vertex AI token endpoint 是外部服务)
|
||||
from src.services.proxy_node.resolver import build_proxy_client_kwargs
|
||||
|
||||
async with httpx.AsyncClient(**build_proxy_client_kwargs(timeout=30)) as client:
|
||||
client_kwargs = (
|
||||
httpx_client_kwargs if httpx_client_kwargs is not None else {"timeout": 30}
|
||||
)
|
||||
async with httpx.AsyncClient(**client_kwargs) as client:
|
||||
resp = await client.post(
|
||||
self.TOKEN_URL,
|
||||
data={
|
||||
@@ -186,16 +191,21 @@ class VertexAuthService:
|
||||
cls._token_cache.clear()
|
||||
|
||||
|
||||
async def get_vertex_access_token(service_account_json: str) -> tuple[str, str]:
|
||||
async def get_vertex_access_token(
|
||||
service_account_json: str,
|
||||
*,
|
||||
httpx_client_kwargs: dict[str, Any] | None = None,
|
||||
) -> tuple[str, str]:
|
||||
"""
|
||||
便捷函数:获取 Vertex AI Access Token 和 Project ID
|
||||
|
||||
Args:
|
||||
service_account_json: Service Account JSON 字符串
|
||||
httpx_client_kwargs: 传给 httpx.AsyncClient 的额外参数(如代理配置)
|
||||
|
||||
Returns:
|
||||
(access_token, project_id) 元组
|
||||
"""
|
||||
service = VertexAuthService(service_account_json)
|
||||
token = await service.get_access_token()
|
||||
token = await service.get_access_token(httpx_client_kwargs=httpx_client_kwargs)
|
||||
return token, service.project_id
|
||||
|
||||
77
src/core/video_utils.py
Normal file
77
src/core/video_utils.py
Normal file
@@ -0,0 +1,77 @@
|
||||
"""
|
||||
视频/图像相关的纯工具函数。
|
||||
|
||||
从 api/handlers 层下沉到 core 层,消除 services→api 的反向依赖。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
# 敏感信息匹配正则(预编译提升性能)
|
||||
_SENSITIVE_PATTERN = re.compile(
|
||||
r"(api[_-]?key|token|bearer|authorization)[=:\s]+\S+",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def sanitize_error_message(message: str, max_length: int = 200) -> str:
|
||||
"""
|
||||
移除错误消息中可能包含的敏感信息
|
||||
|
||||
Args:
|
||||
message: 原始错误消息
|
||||
max_length: 最大长度,默认 200
|
||||
|
||||
Returns:
|
||||
脱敏后的消息
|
||||
"""
|
||||
if not message:
|
||||
return "Request failed"
|
||||
# 先脱敏再截断,确保敏感信息不会因截断位置而泄露
|
||||
sanitized = _SENSITIVE_PATTERN.sub("[REDACTED]", message)
|
||||
return sanitized[:max_length]
|
||||
|
||||
|
||||
def extract_short_id_from_operation(operation_id: str) -> str:
|
||||
"""
|
||||
从 operation ID 中提取短 ID
|
||||
|
||||
我们对外暴露的 operation name 格式是:
|
||||
- models/{model}/operations/{short_id}
|
||||
|
||||
此函数提取最后一部分作为 short_id,用于在数据库中查找任务。
|
||||
|
||||
Args:
|
||||
operation_id: 原始 operation ID(如 "models/veo-3.1/operations/abc123")
|
||||
|
||||
Returns:
|
||||
short_id(如 "abc123")
|
||||
"""
|
||||
# 格式: models/{model}/operations/{short_id}
|
||||
# 或者直接是 short_id
|
||||
if "/" in operation_id:
|
||||
# 提取最后一部分
|
||||
return operation_id.rsplit("/", 1)[-1]
|
||||
return operation_id
|
||||
|
||||
|
||||
def normalize_gemini_operation_id(operation_id: str) -> str:
|
||||
"""
|
||||
规范化 Gemini operation ID(保留用于向后兼容)
|
||||
|
||||
Args:
|
||||
operation_id: 原始 operation ID
|
||||
|
||||
Returns:
|
||||
规范化后的 operation ID(原样返回)
|
||||
"""
|
||||
return operation_id
|
||||
|
||||
|
||||
def is_image_gen_model(model: str | None) -> bool:
|
||||
"""判断是否为图像生成模型(模式匹配,覆盖 gemini-*-image / imagen-* 系列)"""
|
||||
if not model:
|
||||
return False
|
||||
m = model.lower()
|
||||
return "image" in m and ("gemini" in m or "imagen" in m)
|
||||
Reference in New Issue
Block a user