“ydy0615”
5a26dfde2a
后端变更: - 新增 risk_config.py: 风险配置数据类，支持环境变量驱动 - 新增 risk_control.py: 风险控制控制器，管理并发和预算 - 新增 session_store.py: 匿名会话存储，基于 cookie 的 session ID - 新增 audit_store.py: API 审计日志存储，记录请求和 LLM 调用 - 新增 captcha_api.py: 验证码 API，用于验证用户操作真实性 - 新增 llm_policy.py: LLM 策略配置，管理 completion/pro/vision 模型 - main.py: 集成 middleware、risk/audit/session 模块 (+467/-7) - job_handlers.py: LLM 执行流程重构，新增 risk/audit 集成 (+207/-4) - llm.py: 异步客户端封装，新增 max_output_tokens 参数 (+78/-1) - job_system.py: stream_events 逻辑优化，支持心跳检测 (+12/-4) - pro_completions.py: SSE heartbeat 机制，防止连接超时 (+14/-4) - prompt.py: _normalize_preferences 支持 Mapping 类型 (+13/-0) - tts_asr.py: asyncio loop 初始化，router export (+10/-0) 前端变更: - src/components/CaptchaComponent.vue: 新增验证码组件 (NEW) - src/utils/cookie_policy.js: Cookie 策略工具 (NEW) - SettingsPanel.vue: 集成验证码组件，新增安全设置部分 (+59/-0) - MilkdownEditor.vue: 移除硬编码 API_KEY，新增 credentials (+32/-10) - ProBlockCrepe.vue: 样式简化，移除渐变动画 (+18/-4) - proBlockPlugin.ts: 重构 schema/serializer 引用方式，通过 Ctx 管理 (+40/-10) - api.js: 新增 credentials，重构 headers 条件逻辑 (+50/-14) - config.js: API 基址改为 https://api.imageteach.tech:8002 (+8/-4) - convert.js, docsApi.js, i18n.js: 新增 credentials 和验证码 i18n (+54/-12) - proAccept.js: 重构正则和转义处理，修复捕获组索引 (+14/-4) 配置和基础设施: - docker-compose.yml: 新增端口映射 8001:8001 (+2/-0) - docker/nginx.conf: 改为 307 redirect，优化代理配置 (+8/-6) - vite.config.js: 移除 proxy 配置，直接调用远程 API (+8/-4) - .env.example: 新增 VITE_API_BASE_URL, VITE_API_KEY (+3/-1) - backend/.env.example: 大量 RISK_*, SESSION_*, CORS_* 配置 (+54/-0) - pytest.ini: 扩展 coverage 范围到整个 backend，移除 fail_under (+3/-2) - .coveragerc: 移除 fail_under = 90 (+0/-1) - .gitignore: 新增 docker-data/ (+3/-0) - package.json: 新增 vue3-captcha 依赖 (+3/-1) - AGENTS.md, README.md: 更新 Docker 部署和前端网络约定 (+20/-5) - public/sw.js: Service Worker cache 版本从 v1 升级到 v2 (+0/-1) 测试变更: - test_main_endpoints.py: 新增 session/risk/audit reset，新增测试用例 (+63/-4) - test_main_cancel.py: 新增 reset 调用 (+6/-0) - test_pro_completions.py: 新增 preferences 序列化和测试 (+23/-0) 总计: 45 个文件变更，+1009/-280 行
355 lines
14 KiB
Python
355 lines
14 KiB
Python
import asyncio
|
|
import hashlib
|
|
import math
|
|
import os
|
|
import time
|
|
from dataclasses import dataclass
|
|
from datetime import date
|
|
from typing import Any
|
|
|
|
from risk_config import RiskConfig
|
|
|
|
try: # pragma: no cover
|
|
from redis import asyncio as redis_asyncio
|
|
except Exception: # pragma: no cover
|
|
redis_asyncio = None
|
|
|
|
|
|
def _now_ms() -> int:
|
|
return int(time.time() * 1000)
|
|
|
|
|
|
def _utc_day() -> str:
|
|
return date.today().isoformat()
|
|
|
|
|
|
def stable_hash(value: str) -> str:
|
|
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def estimate_tokens(text: str) -> int:
|
|
if not text:
|
|
return 0
|
|
ascii_chars = sum(1 for ch in text if ord(ch) < 128)
|
|
non_ascii = len(text) - ascii_chars
|
|
ascii_tokens = math.ceil(ascii_chars / 4)
|
|
non_ascii_tokens = math.ceil(non_ascii * 1.5)
|
|
return max(ascii_tokens + non_ascii_tokens, 1)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RiskIdentity:
|
|
request_id: str
|
|
session_hash: str
|
|
ip_hash: str
|
|
route: str
|
|
method: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RiskDecision:
|
|
allowed: bool
|
|
status_code: int = 200
|
|
reason: str = ""
|
|
error_code: str = ""
|
|
retry_after_seconds: int = 0
|
|
delay_ms: int = 0
|
|
|
|
|
|
class RiskRejected(RuntimeError):
|
|
def __init__(self, decision: RiskDecision) -> None:
|
|
super().__init__(decision.reason or decision.error_code or "request rejected")
|
|
self.decision = decision
|
|
|
|
|
|
class BaseRiskBackend:
|
|
async def incr_window(self, key: str, ttl_seconds: int) -> int:
|
|
raise NotImplementedError
|
|
|
|
async def get_float(self, key: str) -> float:
|
|
raise NotImplementedError
|
|
|
|
async def add_float(self, key: str, value: float, ttl_seconds: int) -> float:
|
|
raise NotImplementedError
|
|
|
|
async def get_int(self, key: str) -> int:
|
|
raise NotImplementedError
|
|
|
|
async def set_int(self, key: str, value: int, ttl_seconds: int) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def set_float(self, key: str, value: float, ttl_seconds: int) -> None:
|
|
raise NotImplementedError
|
|
|
|
async def acquire_lock(self, key: str, ttl_seconds: int) -> bool:
|
|
raise NotImplementedError
|
|
|
|
async def release_lock(self, key: str) -> None:
|
|
raise NotImplementedError
|
|
|
|
|
|
class InMemoryRiskBackend(BaseRiskBackend):
|
|
def __init__(self) -> None:
|
|
self.values: dict[str, tuple[float, float]] = {}
|
|
self.locks: dict[str, float] = {}
|
|
self.guard = asyncio.Lock()
|
|
|
|
def _purge(self) -> None:
|
|
now = time.time()
|
|
for key, (_, expires_at) in list(self.values.items()):
|
|
if expires_at and expires_at <= now:
|
|
self.values.pop(key, None)
|
|
for key, expires_at in list(self.locks.items()):
|
|
if expires_at <= now:
|
|
self.locks.pop(key, None)
|
|
|
|
async def incr_window(self, key: str, ttl_seconds: int) -> int:
|
|
async with self.guard:
|
|
self._purge()
|
|
value, _ = self.values.get(key, (0.0, 0.0))
|
|
next_value = int(value) + 1
|
|
self.values[key] = (float(next_value), time.time() + ttl_seconds)
|
|
return next_value
|
|
|
|
async def get_float(self, key: str) -> float:
|
|
async with self.guard:
|
|
self._purge()
|
|
return float(self.values.get(key, (0.0, 0.0))[0])
|
|
|
|
async def add_float(self, key: str, value: float, ttl_seconds: int) -> float:
|
|
async with self.guard:
|
|
self._purge()
|
|
current, _ = self.values.get(key, (0.0, 0.0))
|
|
next_value = current + value
|
|
self.values[key] = (next_value, time.time() + ttl_seconds)
|
|
return next_value
|
|
|
|
async def get_int(self, key: str) -> int:
|
|
return int(await self.get_float(key))
|
|
|
|
async def set_int(self, key: str, value: int, ttl_seconds: int) -> None:
|
|
async with self.guard:
|
|
self._purge()
|
|
self.values[key] = (float(value), time.time() + ttl_seconds)
|
|
|
|
async def set_float(self, key: str, value: float, ttl_seconds: int) -> None:
|
|
async with self.guard:
|
|
self._purge()
|
|
self.values[key] = (float(value), time.time() + ttl_seconds)
|
|
|
|
async def acquire_lock(self, key: str, ttl_seconds: int) -> bool:
|
|
async with self.guard:
|
|
self._purge()
|
|
if key in self.locks:
|
|
return False
|
|
self.locks[key] = time.time() + ttl_seconds
|
|
return True
|
|
|
|
async def release_lock(self, key: str) -> None:
|
|
async with self.guard:
|
|
self.locks.pop(key, None)
|
|
|
|
|
|
class RedisRiskBackend(BaseRiskBackend):
|
|
def __init__(self, redis_url: str) -> None:
|
|
if redis_asyncio is None:
|
|
raise RuntimeError("redis package is not installed")
|
|
self.redis = redis_asyncio.from_url(redis_url, encoding="utf-8", decode_responses=True)
|
|
|
|
async def incr_window(self, key: str, ttl_seconds: int) -> int:
|
|
value = await self.redis.incr(key)
|
|
if value == 1:
|
|
await self.redis.expire(key, ttl_seconds)
|
|
return int(value)
|
|
|
|
async def get_float(self, key: str) -> float:
|
|
value = await self.redis.get(key)
|
|
if value is None:
|
|
return 0.0
|
|
return float(value)
|
|
|
|
async def add_float(self, key: str, value: float, ttl_seconds: int) -> float:
|
|
current = await self.get_float(key)
|
|
next_value = current + value
|
|
await self.redis.set(key, next_value, ex=ttl_seconds)
|
|
return next_value
|
|
|
|
async def get_int(self, key: str) -> int:
|
|
value = await self.redis.get(key)
|
|
if value is None:
|
|
return 0
|
|
return int(value)
|
|
|
|
async def set_int(self, key: str, value: int, ttl_seconds: int) -> None:
|
|
await self.redis.set(key, value, ex=ttl_seconds)
|
|
|
|
async def set_float(self, key: str, value: float, ttl_seconds: int) -> None:
|
|
await self.redis.set(key, value, ex=ttl_seconds)
|
|
|
|
async def acquire_lock(self, key: str, ttl_seconds: int) -> bool:
|
|
return bool(await self.redis.set(key, "1", ex=ttl_seconds, nx=True))
|
|
|
|
async def release_lock(self, key: str) -> None:
|
|
await self.redis.delete(key)
|
|
|
|
|
|
class RiskController:
|
|
def __init__(self, config: RiskConfig) -> None:
|
|
self.config = config
|
|
self.prefix = "llmtext:risk"
|
|
redis_url = os.getenv("REDIS_URL", "").strip()
|
|
if redis_url and redis_asyncio is not None:
|
|
self.backend: BaseRiskBackend = RedisRiskBackend(redis_url)
|
|
else:
|
|
self.backend = InMemoryRiskBackend()
|
|
|
|
def _api_key(self, identity: RiskIdentity, scope: str) -> str:
|
|
return f"{self.prefix}:api:{scope}:{identity.session_hash}:{identity.ip_hash}"
|
|
|
|
def _llm_key(self, identity: RiskIdentity, scope: str) -> str:
|
|
return f"{self.prefix}:llm:{scope}:{identity.session_hash}:{identity.ip_hash}"
|
|
|
|
def _budget_key(self, scope: str, scope_hash: str, current_day: str) -> str:
|
|
return f"{self.prefix}:budget:{scope}:{scope_hash}:{current_day}"
|
|
|
|
def _lock_key(self, scope: str, scope_hash: str) -> str:
|
|
return f"{self.prefix}:lock:{scope}:{scope_hash}"
|
|
|
|
def _circuit_key(self, scope: str) -> str:
|
|
return f"{self.prefix}:circuit:{scope}"
|
|
|
|
def _failure_key(self, model: str) -> str:
|
|
return f"{self.prefix}:failure:{model}"
|
|
|
|
async def check_api(self, identity: RiskIdentity, *, scope: str = "default") -> RiskDecision:
|
|
key = self._api_key(identity, scope)
|
|
count = await self.backend.incr_window(key, self.config.api_window_seconds)
|
|
if count > self.config.api_hard_limit_per_window:
|
|
return RiskDecision(
|
|
allowed=False,
|
|
status_code=429,
|
|
reason="请求过于频繁，请稍后再试",
|
|
error_code="api_rate_limited",
|
|
retry_after_seconds=self.config.api_window_seconds,
|
|
)
|
|
if count > self.config.api_soft_limit_per_window:
|
|
overflow = count - self.config.api_soft_limit_per_window
|
|
delay_ms = min(self.config.delay_cap_ms, overflow * self.config.delay_step_ms)
|
|
return RiskDecision(allowed=True, delay_ms=delay_ms)
|
|
return RiskDecision(allowed=True)
|
|
|
|
async def check_llm(
|
|
self,
|
|
identity: RiskIdentity,
|
|
*,
|
|
scope: str,
|
|
estimated_cost: float,
|
|
) -> RiskDecision:
|
|
global_circuit = await self.backend.get_int(self._circuit_key("global"))
|
|
model_circuit = await self.backend.get_int(self._circuit_key(scope))
|
|
if global_circuit > 0 or model_circuit > 0:
|
|
return RiskDecision(
|
|
allowed=False,
|
|
status_code=503,
|
|
reason="当前推理服务繁忙，请稍后再试",
|
|
error_code="llm_circuit_open",
|
|
retry_after_seconds=self.config.model_circuit_ttl_seconds,
|
|
)
|
|
if estimated_cost > self.config.single_request_max_cost_usd:
|
|
return RiskDecision(
|
|
allowed=False,
|
|
status_code=429,
|
|
reason="单次请求成本过高，已被拒绝",
|
|
error_code="llm_cost_too_high",
|
|
)
|
|
day = _utc_day()
|
|
session_budget = await self.backend.get_float(self._budget_key("session", identity.session_hash, day))
|
|
ip_budget = await self.backend.get_float(self._budget_key("ip", identity.ip_hash, day))
|
|
global_budget = await self.backend.get_float(self._budget_key("global", "global", day))
|
|
if session_budget + estimated_cost > self.config.daily_budget_session_usd:
|
|
return RiskDecision(False, 429, "当前匿名会话今日额度已用尽", "session_budget_exhausted", 3600)
|
|
if ip_budget + estimated_cost > self.config.daily_budget_ip_usd:
|
|
return RiskDecision(False, 429, "当前网络环境今日额度已用尽", "ip_budget_exhausted", 3600)
|
|
if global_budget + estimated_cost > self.config.daily_budget_global_usd:
|
|
await self.backend.set_int(self._circuit_key("global"), 1, self.config.model_circuit_ttl_seconds)
|
|
return RiskDecision(False, 503, "今日全局推理预算已耗尽", "global_budget_exhausted", 3600)
|
|
key = self._llm_key(identity, scope)
|
|
count = await self.backend.incr_window(key, self.config.llm_window_seconds)
|
|
if count > self.config.llm_hard_limit_per_window:
|
|
return RiskDecision(False, 429, "推理请求过于频繁，请稍后重试", "llm_rate_limited", self.config.llm_window_seconds)
|
|
if count > self.config.llm_soft_limit_per_window:
|
|
overflow = count - self.config.llm_soft_limit_per_window
|
|
delay_ms = min(self.config.delay_cap_ms, overflow * self.config.delay_step_ms)
|
|
return RiskDecision(True, delay_ms=delay_ms)
|
|
return RiskDecision(True)
|
|
|
|
async def reserve_budget(self, identity: RiskIdentity, estimated_cost: float) -> None:
|
|
day = _utc_day()
|
|
ttl_seconds = 60 * 60 * 24
|
|
await self.backend.add_float(self._budget_key("session", identity.session_hash, day), estimated_cost, ttl_seconds)
|
|
await self.backend.add_float(self._budget_key("ip", identity.ip_hash, day), estimated_cost, ttl_seconds)
|
|
await self.backend.add_float(self._budget_key("global", "global", day), estimated_cost, ttl_seconds)
|
|
|
|
async def acquire_execution_slot(self, identity: RiskIdentity, *, model: str) -> list[str]:
|
|
ttl_seconds = 60 * 15
|
|
keys = [
|
|
self._lock_key("session", f"{identity.session_hash}:{identity.request_id}"),
|
|
self._lock_key("global", identity.request_id),
|
|
]
|
|
session_running = await self.backend.get_int(self._lock_key("session-count", identity.session_hash))
|
|
global_running = await self.backend.get_int(self._lock_key("global-count", "global"))
|
|
if session_running >= self.config.session_concurrency_limit:
|
|
raise RiskRejected(
|
|
RiskDecision(False, 429, "当前会话并发推理过多，请稍后重试", "session_concurrency_limited", 30)
|
|
)
|
|
if global_running >= self.config.global_concurrency_limit:
|
|
raise RiskRejected(
|
|
RiskDecision(False, 503, "当前全局推理负载过高，请稍后重试", "global_concurrency_limited", 30)
|
|
)
|
|
acquired: list[str] = []
|
|
for key in keys:
|
|
ok = await self.backend.acquire_lock(key, ttl_seconds)
|
|
if not ok:
|
|
for acquired_key in acquired:
|
|
await self.backend.release_lock(acquired_key)
|
|
raise RiskRejected(
|
|
RiskDecision(False, 429, "当前请求正在执行，请勿重复提交", "duplicate_request", 10)
|
|
)
|
|
acquired.append(key)
|
|
await self.backend.add_float(self._lock_key("session-count", identity.session_hash), 1.0, ttl_seconds)
|
|
await self.backend.add_float(self._lock_key("global-count", "global"), 1.0, ttl_seconds)
|
|
return acquired
|
|
|
|
async def release_execution_slot(self, identity: RiskIdentity, lock_keys: list[str], *, model: str) -> None:
|
|
for key in lock_keys:
|
|
await self.backend.release_lock(key)
|
|
session_count_key = self._lock_key("session-count", identity.session_hash)
|
|
global_count_key = self._lock_key("global-count", "global")
|
|
session_count = max(0.0, await self.backend.get_float(session_count_key) - 1.0)
|
|
global_count = max(0.0, await self.backend.get_float(global_count_key) - 1.0)
|
|
await self.backend.set_float(session_count_key, session_count, 60 * 15)
|
|
await self.backend.set_float(global_count_key, global_count, 60 * 15)
|
|
|
|
async def record_model_result(self, *, model: str, success: bool) -> None:
|
|
if success:
|
|
await self.backend.set_int(self._failure_key(model), 0, self.config.model_circuit_ttl_seconds)
|
|
return
|
|
failures = await self.backend.incr_window(self._failure_key(model), self.config.model_circuit_ttl_seconds)
|
|
if failures >= self.config.model_circuit_breaker_failures:
|
|
await self.backend.set_int(self._circuit_key(model), 1, self.config.model_circuit_ttl_seconds)
|
|
|
|
|
|
_risk_controller: RiskController | None = None
|
|
|
|
|
|
def get_risk_controller(config: RiskConfig) -> RiskController:
|
|
global _risk_controller
|
|
if _risk_controller is None:
|
|
_risk_controller = RiskController(config)
|
|
return _risk_controller
|
|
|
|
|
|
def reset_risk_controller() -> None:
|
|
global _risk_controller
|
|
_risk_controller = None
|