Initial project import
This commit is contained in:
XM
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,253 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import secrets
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from threading import Lock
|
||||
from typing import Literal
|
||||
|
||||
from services.config import config
|
||||
from services.storage.base import StorageBackend
|
||||
|
||||
AuthRole = Literal["admin", "user"]
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _hash_key(value: str) -> str:
|
||||
return hashlib.sha256(value.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
class AuthService:
|
||||
def __init__(self, storage: StorageBackend):
|
||||
self.storage = storage
|
||||
self._lock = Lock()
|
||||
self._items = self._load()
|
||||
self._last_used_flush_at: dict[str, datetime] = {}
|
||||
|
||||
@staticmethod
|
||||
def _clean(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
@staticmethod
|
||||
def _default_name(role: object) -> str:
|
||||
return "管理员密钥" if str(role or "").strip().lower() == "admin" else "普通用户"
|
||||
|
||||
def _normalize_item(self, raw: object) -> dict[str, object] | None:
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
role = self._clean(raw.get("role")).lower()
|
||||
if role not in {"admin", "user"}:
|
||||
return None
|
||||
key_hash = self._clean(raw.get("key_hash"))
|
||||
if not key_hash:
|
||||
return None
|
||||
item_id = self._clean(raw.get("id")) or uuid.uuid4().hex[:12]
|
||||
name = self._clean(raw.get("name")) or self._default_name(role)
|
||||
created_at = self._clean(raw.get("created_at")) or _now_iso()
|
||||
last_used_at = self._clean(raw.get("last_used_at")) or None
|
||||
return {
|
||||
"id": item_id,
|
||||
"name": name,
|
||||
"role": role,
|
||||
"key_hash": key_hash,
|
||||
"enabled": bool(raw.get("enabled", True)),
|
||||
"created_at": created_at,
|
||||
"last_used_at": last_used_at,
|
||||
}
|
||||
|
||||
def _load(self) -> list[dict[str, object]]:
|
||||
try:
|
||||
items = self.storage.load_auth_keys()
|
||||
except Exception:
|
||||
return []
|
||||
if not isinstance(items, list):
|
||||
return []
|
||||
return [normalized for item in items if (normalized := self._normalize_item(item)) is not None]
|
||||
|
||||
def _save(self) -> None:
|
||||
self.storage.save_auth_keys(self._items)
|
||||
|
||||
def _reload_locked(self) -> None:
|
||||
self._items = self._load()
|
||||
|
||||
@staticmethod
|
||||
def _public_item(item: dict[str, object]) -> dict[str, object]:
|
||||
return {
|
||||
"id": item.get("id"),
|
||||
"name": item.get("name"),
|
||||
"role": item.get("role"),
|
||||
"enabled": bool(item.get("enabled", True)),
|
||||
"created_at": item.get("created_at"),
|
||||
"last_used_at": item.get("last_used_at"),
|
||||
}
|
||||
|
||||
def list_keys(self, role: AuthRole | None = None) -> list[dict[str, object]]:
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
items = [item for item in self._items if role is None or item.get("role") == role]
|
||||
return [self._public_item(item) for item in items]
|
||||
|
||||
def _has_key_hash_locked(self, key_hash: str, *, exclude_id: str = "") -> bool:
|
||||
for item in self._items:
|
||||
item_id = self._clean(item.get("id"))
|
||||
if exclude_id and item_id == exclude_id:
|
||||
continue
|
||||
stored_hash = self._clean(item.get("key_hash"))
|
||||
if stored_hash and hmac.compare_digest(stored_hash, key_hash):
|
||||
return True
|
||||
return False
|
||||
|
||||
def _build_key_hash_locked(self, raw_key: str, *, exclude_id: str = "") -> str:
|
||||
candidate = self._clean(raw_key)
|
||||
if not candidate:
|
||||
raise ValueError("请输入新的专用密钥")
|
||||
admin_key = self._clean(config.auth_key)
|
||||
if admin_key and hmac.compare_digest(candidate, admin_key):
|
||||
raise ValueError("这个密钥和管理员密钥冲突了，请换一个新的密钥")
|
||||
key_hash = _hash_key(candidate)
|
||||
if self._has_key_hash_locked(key_hash, exclude_id=exclude_id):
|
||||
raise ValueError("这个专用密钥已经存在，请换一个新的密钥")
|
||||
return key_hash
|
||||
|
||||
def _has_name_locked(self, name: str, *, role: AuthRole | None = None, exclude_id: str = "") -> bool:
|
||||
candidate = self._clean(name)
|
||||
if not candidate:
|
||||
return False
|
||||
for item in self._items:
|
||||
item_id = self._clean(item.get("id"))
|
||||
if exclude_id and item_id == exclude_id:
|
||||
continue
|
||||
if role is not None and item.get("role") != role:
|
||||
continue
|
||||
if self._clean(item.get("name")) == candidate:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _build_default_name_locked(self, role: AuthRole, *, exclude_id: str = "") -> str:
|
||||
base_name = self._default_name(role)
|
||||
if not self._has_name_locked(base_name, role=role, exclude_id=exclude_id):
|
||||
return base_name
|
||||
suffix = 2
|
||||
while True:
|
||||
candidate = f"{base_name} {suffix}"
|
||||
if not self._has_name_locked(candidate, role=role, exclude_id=exclude_id):
|
||||
return candidate
|
||||
suffix += 1
|
||||
|
||||
def _build_name_locked(self, name: str, *, role: AuthRole, exclude_id: str = "") -> str:
|
||||
candidate = self._clean(name)
|
||||
if not candidate:
|
||||
return self._build_default_name_locked(role, exclude_id=exclude_id)
|
||||
if self._has_name_locked(candidate, role=role, exclude_id=exclude_id):
|
||||
raise ValueError("这个名称已经在使用中了，换一个更容易区分的名称吧")
|
||||
return candidate
|
||||
|
||||
def create_key(self, *, role: AuthRole, name: str = "") -> tuple[dict[str, object], str]:
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
normalized_name = self._build_name_locked(name, role=role)
|
||||
while True:
|
||||
raw_key = f"sk-{secrets.token_urlsafe(24)}"
|
||||
try:
|
||||
key_hash = self._build_key_hash_locked(raw_key)
|
||||
break
|
||||
except ValueError:
|
||||
continue
|
||||
item = {
|
||||
"id": uuid.uuid4().hex[:12],
|
||||
"name": normalized_name,
|
||||
"role": role,
|
||||
"key_hash": key_hash,
|
||||
"enabled": True,
|
||||
"created_at": _now_iso(),
|
||||
"last_used_at": None,
|
||||
}
|
||||
self._items.append(item)
|
||||
self._save()
|
||||
return self._public_item(item), raw_key
|
||||
|
||||
def update_key(
|
||||
self,
|
||||
key_id: str,
|
||||
updates: dict[str, object],
|
||||
*,
|
||||
role: AuthRole | None = None,
|
||||
) -> dict[str, object] | None:
|
||||
normalized_id = self._clean(key_id)
|
||||
if not normalized_id:
|
||||
return None
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
for index, item in enumerate(self._items):
|
||||
if item.get("id") != normalized_id:
|
||||
continue
|
||||
if role is not None and item.get("role") != role:
|
||||
return None
|
||||
next_item = dict(item)
|
||||
next_role = "admin" if str(next_item.get("role") or "").strip().lower() == "admin" else "user"
|
||||
if "name" in updates and updates.get("name") is not None:
|
||||
next_item["name"] = self._build_name_locked(
|
||||
str(updates.get("name") or ""),
|
||||
role=next_role,
|
||||
exclude_id=normalized_id,
|
||||
)
|
||||
if "enabled" in updates and updates.get("enabled") is not None:
|
||||
next_item["enabled"] = bool(updates.get("enabled"))
|
||||
if "key" in updates and updates.get("key") is not None:
|
||||
next_item["key_hash"] = self._build_key_hash_locked(str(updates.get("key") or ""), exclude_id=normalized_id)
|
||||
self._items[index] = next_item
|
||||
self._save()
|
||||
return self._public_item(next_item)
|
||||
return None
|
||||
|
||||
def delete_key(self, key_id: str, *, role: AuthRole | None = None) -> bool:
|
||||
normalized_id = self._clean(key_id)
|
||||
if not normalized_id:
|
||||
return False
|
||||
with self._lock:
|
||||
self._reload_locked()
|
||||
before = len(self._items)
|
||||
self._items = [
|
||||
item
|
||||
for item in self._items
|
||||
if not (item.get("id") == normalized_id and (role is None or item.get("role") == role))
|
||||
]
|
||||
if len(self._items) == before:
|
||||
return False
|
||||
self._save()
|
||||
return True
|
||||
|
||||
def authenticate(self, raw_key: str) -> dict[str, object] | None:
|
||||
candidate = self._clean(raw_key)
|
||||
if not candidate:
|
||||
return None
|
||||
candidate_hash = _hash_key(candidate)
|
||||
with self._lock:
|
||||
for index, item in enumerate(self._items):
|
||||
if not bool(item.get("enabled", True)):
|
||||
continue
|
||||
stored_hash = self._clean(item.get("key_hash"))
|
||||
if not stored_hash or not hmac.compare_digest(stored_hash, candidate_hash):
|
||||
continue
|
||||
next_item = dict(item)
|
||||
now = datetime.now(timezone.utc)
|
||||
next_item["last_used_at"] = now.isoformat()
|
||||
self._items[index] = next_item
|
||||
item_id = self._clean(next_item.get("id"))
|
||||
last_flush_at = self._last_used_flush_at.get(item_id)
|
||||
if last_flush_at is None or (now - last_flush_at).total_seconds() >= 60:
|
||||
try:
|
||||
self._save()
|
||||
self._last_used_flush_at[item_id] = now
|
||||
except Exception:
|
||||
pass
|
||||
return self._public_item(next_item)
|
||||
return None
|
||||
|
||||
|
||||
auth_service = AuthService(config.get_storage_backend())
|
||||
@@ -0,0 +1,671 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import tarfile
|
||||
import threading
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from urllib.parse import quote, urlencode
|
||||
|
||||
from curl_cffi import requests
|
||||
|
||||
from services.config import BASE_DIR, CONFIG_FILE, DATA_DIR, config, load_backup_state, save_backup_state
|
||||
from services.image_storage_service import IMAGE_INDEX_FILE
|
||||
from services.image_tags_service import TAGS_FILE
|
||||
|
||||
|
||||
def _utc_now() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
def _iso_now() -> str:
|
||||
return _utc_now().replace(microsecond=0).isoformat().replace("+00:00", "Z")
|
||||
|
||||
|
||||
def _clean(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def _sha256_hex(value: bytes) -> str:
|
||||
return hashlib.sha256(value).hexdigest()
|
||||
|
||||
|
||||
def _is_backup_object(key: object) -> bool:
|
||||
name = _clean(key).rsplit("/", 1)[-1]
|
||||
return name.startswith("backup-") and (name.endswith(".tar.gz") or name.endswith(".tar.gz.enc"))
|
||||
|
||||
|
||||
def _hmac_sha256(key: bytes, message: str) -> bytes:
|
||||
return hmac.new(key, message.encode("utf-8"), hashlib.sha256).digest()
|
||||
|
||||
|
||||
def _openssl_encrypt(data: bytes, passphrase: str) -> bytes:
|
||||
env = dict(os.environ)
|
||||
env["CHATGPT2API_BACKUP_PASSPHRASE"] = passphrase
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"openssl",
|
||||
"enc",
|
||||
"-aes-256-cbc",
|
||||
"-pbkdf2",
|
||||
"-salt",
|
||||
"-md",
|
||||
"sha256",
|
||||
"-pass",
|
||||
"env:CHATGPT2API_BACKUP_PASSPHRASE",
|
||||
],
|
||||
input=data,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
check=True,
|
||||
env=env,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise BackupError("当前环境缺少 openssl，无法执行加密备份") from exc
|
||||
except subprocess.CalledProcessError as exc:
|
||||
detail = (exc.stderr or b"").decode("utf-8", errors="replace").strip()
|
||||
raise BackupError(f"加密备份失败：{detail or 'openssl 执行失败'}") from exc
|
||||
return result.stdout
|
||||
|
||||
|
||||
def _openssl_decrypt(data: bytes, passphrase: str) -> bytes:
|
||||
env = dict(os.environ)
|
||||
env["CHATGPT2API_BACKUP_PASSPHRASE"] = passphrase
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[
|
||||
"openssl",
|
||||
"enc",
|
||||
"-d",
|
||||
"-aes-256-cbc",
|
||||
"-pbkdf2",
|
||||
"-md",
|
||||
"sha256",
|
||||
"-pass",
|
||||
"env:CHATGPT2API_BACKUP_PASSPHRASE",
|
||||
],
|
||||
input=data,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
check=True,
|
||||
env=env,
|
||||
)
|
||||
except FileNotFoundError as exc:
|
||||
raise BackupError("当前环境缺少 openssl，无法解密备份内容") from exc
|
||||
except subprocess.CalledProcessError as exc:
|
||||
detail = (exc.stderr or b"").decode("utf-8", errors="replace").strip()
|
||||
raise BackupError(f"解密备份失败：{detail or 'openssl 执行失败'}") from exc
|
||||
return result.stdout
|
||||
|
||||
|
||||
def _guess_content_type(name: str) -> str:
|
||||
if name.endswith(".json"):
|
||||
return "application/json"
|
||||
if name.endswith(".jsonl"):
|
||||
return "application/x-ndjson"
|
||||
if name.endswith(".tar.gz"):
|
||||
return "application/gzip"
|
||||
if name.endswith(".gz"):
|
||||
return "application/gzip"
|
||||
return "application/octet-stream"
|
||||
|
||||
|
||||
def _json_bytes(value: object) -> bytes:
|
||||
return json.dumps(value, ensure_ascii=False, indent=2).encode("utf-8")
|
||||
|
||||
|
||||
def _count_items(value: object) -> int:
|
||||
if isinstance(value, list):
|
||||
return len(value)
|
||||
if isinstance(value, dict):
|
||||
return len(value)
|
||||
return 0
|
||||
|
||||
|
||||
class BackupError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class CloudflareR2Client:
|
||||
def __init__(self, settings: dict[str, object]) -> None:
|
||||
self.account_id = _clean(settings.get("account_id"))
|
||||
self.access_key_id = _clean(settings.get("access_key_id"))
|
||||
self.secret_access_key = _clean(settings.get("secret_access_key"))
|
||||
self.bucket = _clean(settings.get("bucket"))
|
||||
self.prefix = _clean(settings.get("prefix")) or "backups"
|
||||
self.session = requests.Session(impersonate="chrome", verify=True)
|
||||
|
||||
def validate(self) -> None:
|
||||
missing = []
|
||||
if not self.account_id:
|
||||
missing.append("Account ID")
|
||||
if not self.access_key_id:
|
||||
missing.append("Access Key ID")
|
||||
if not self.secret_access_key:
|
||||
missing.append("Secret Access Key")
|
||||
if not self.bucket:
|
||||
missing.append("Bucket")
|
||||
if missing:
|
||||
raise BackupError(f"R2 配置不完整：缺少 {'、'.join(missing)}")
|
||||
|
||||
@property
|
||||
def endpoint(self) -> str:
|
||||
return f"https://{self.account_id}.r2.cloudflarestorage.com"
|
||||
|
||||
def _aws_v4_headers(
|
||||
self,
|
||||
method: str,
|
||||
path: str,
|
||||
*,
|
||||
query: dict[str, str] | None = None,
|
||||
body: bytes = b"",
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> tuple[str, dict[str, str]]:
|
||||
now = _utc_now()
|
||||
amz_date = now.strftime("%Y%m%dT%H%M%SZ")
|
||||
date_stamp = now.strftime("%Y%m%d")
|
||||
encoded_query = urlencode(sorted((query or {}).items()))
|
||||
payload_hash = _sha256_hex(body)
|
||||
host = f"{self.account_id}.r2.cloudflarestorage.com"
|
||||
headers = {
|
||||
"host": host,
|
||||
"x-amz-content-sha256": payload_hash,
|
||||
"x-amz-date": amz_date,
|
||||
}
|
||||
if extra_headers:
|
||||
for key, value in extra_headers.items():
|
||||
headers[key.lower()] = value.strip()
|
||||
sorted_items = sorted((key.lower(), " ".join(str(value).strip().split())) for key, value in headers.items())
|
||||
canonical_headers = "".join(f"{key}:{value}\n" for key, value in sorted_items)
|
||||
signed_headers = ";".join(key for key, _ in sorted_items)
|
||||
canonical_request = "\n".join([
|
||||
method.upper(),
|
||||
path,
|
||||
encoded_query,
|
||||
canonical_headers,
|
||||
signed_headers,
|
||||
payload_hash,
|
||||
])
|
||||
credential_scope = f"{date_stamp}/auto/s3/aws4_request"
|
||||
string_to_sign = "\n".join([
|
||||
"AWS4-HMAC-SHA256",
|
||||
amz_date,
|
||||
credential_scope,
|
||||
_sha256_hex(canonical_request.encode("utf-8")),
|
||||
])
|
||||
k_date = _hmac_sha256(("AWS4" + self.secret_access_key).encode("utf-8"), date_stamp)
|
||||
k_region = hmac.new(k_date, b"auto", hashlib.sha256).digest()
|
||||
k_service = hmac.new(k_region, b"s3", hashlib.sha256).digest()
|
||||
k_signing = hmac.new(k_service, b"aws4_request", hashlib.sha256).digest()
|
||||
signature = hmac.new(k_signing, string_to_sign.encode("utf-8"), hashlib.sha256).hexdigest()
|
||||
authorization = (
|
||||
"AWS4-HMAC-SHA256 "
|
||||
f"Credential={self.access_key_id}/{credential_scope}, "
|
||||
f"SignedHeaders={signed_headers}, "
|
||||
f"Signature={signature}"
|
||||
)
|
||||
request_headers = {key: value for key, value in headers.items()}
|
||||
request_headers["authorization"] = authorization
|
||||
return encoded_query, request_headers
|
||||
|
||||
def _request(
|
||||
self,
|
||||
method: str,
|
||||
key: str = "",
|
||||
*,
|
||||
query: dict[str, str] | None = None,
|
||||
body: bytes = b"",
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
timeout: float = 60.0,
|
||||
):
|
||||
object_path = f"/{self.bucket}"
|
||||
if key:
|
||||
object_path += f"/{quote(key.lstrip('/'), safe='/')}"
|
||||
encoded_query, headers = self._aws_v4_headers(method, object_path, query=query, body=body, extra_headers=extra_headers)
|
||||
url = f"{self.endpoint}{object_path}"
|
||||
if encoded_query:
|
||||
url += f"?{encoded_query}"
|
||||
response = self.session.request(method.upper(), url, headers=headers, data=body, timeout=timeout)
|
||||
return response
|
||||
|
||||
def test_connection(self) -> dict[str, object]:
|
||||
self.validate()
|
||||
response = self._request("GET", query={"list-type": "2", "max-keys": "1"}, timeout=30.0)
|
||||
if response.status_code >= 400:
|
||||
raise BackupError(f"连接 R2 失败：HTTP {response.status_code}")
|
||||
return {"ok": True, "status": int(response.status_code)}
|
||||
|
||||
def upload_bytes(self, key: str, payload: bytes, *, content_type: str, metadata: dict[str, str] | None = None) -> dict[str, object]:
|
||||
headers = {"content-type": content_type}
|
||||
if metadata:
|
||||
for item_key, item_value in metadata.items():
|
||||
headers[f"x-amz-meta-{item_key}"] = str(item_value)
|
||||
response = self._request("PUT", key, body=payload, extra_headers=headers)
|
||||
if response.status_code >= 400:
|
||||
raise BackupError(f"上传备份失败：HTTP {response.status_code}")
|
||||
return {"key": key, "etag": str(response.headers.get("etag") or "").strip('"')}
|
||||
|
||||
def delete_object(self, key: str) -> None:
|
||||
response = self._request("DELETE", key, timeout=30.0)
|
||||
if response.status_code >= 400 and response.status_code != 404:
|
||||
raise BackupError(f"删除备份失败：HTTP {response.status_code}")
|
||||
|
||||
def download_bytes(self, key: str) -> bytes:
|
||||
response = self._request("GET", key, timeout=60.0)
|
||||
if response.status_code >= 400:
|
||||
raise BackupError(f"读取备份失败：HTTP {response.status_code}")
|
||||
return bytes(response.content or b"")
|
||||
|
||||
def list_objects(self) -> list[dict[str, object]]:
|
||||
items: list[dict[str, object]] = []
|
||||
continuation = ""
|
||||
while True:
|
||||
query = {"list-type": "2", "prefix": f"{self.prefix.rstrip('/')}/", "max-keys": "1000"}
|
||||
if continuation:
|
||||
query["continuation-token"] = continuation
|
||||
response = self._request("GET", query=query, timeout=30.0)
|
||||
if response.status_code >= 400:
|
||||
raise BackupError(f"获取备份列表失败：HTTP {response.status_code}")
|
||||
text = response.text
|
||||
for block in text.split("<Contents>")[1:]:
|
||||
key = _clean(block.split("<Key>", 1)[1].split("</Key>", 1)[0]) if "<Key>" in block else ""
|
||||
if not key:
|
||||
continue
|
||||
size_text = _clean(block.split("<Size>", 1)[1].split("</Size>", 1)[0]) if "<Size>" in block else "0"
|
||||
updated = _clean(block.split("<LastModified>", 1)[1].split("</LastModified>", 1)[0]) if "<LastModified>" in block else ""
|
||||
items.append({
|
||||
"key": key,
|
||||
"size": int(size_text or 0),
|
||||
"updated_at": updated,
|
||||
})
|
||||
truncated = "<IsTruncated>true</IsTruncated>" in text
|
||||
if not truncated:
|
||||
break
|
||||
if "<NextContinuationToken>" not in text:
|
||||
break
|
||||
continuation = _clean(text.split("<NextContinuationToken>", 1)[1].split("</NextContinuationToken>", 1)[0])
|
||||
if not continuation:
|
||||
break
|
||||
items.sort(key=lambda item: str(item.get("updated_at") or ""), reverse=True)
|
||||
return items
|
||||
|
||||
def close(self) -> None:
|
||||
self.session.close()
|
||||
|
||||
|
||||
class BackupService:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.RLock()
|
||||
self._stop_event = threading.Event()
|
||||
self._thread: threading.Thread | None = None
|
||||
self._running = False
|
||||
|
||||
def start(self) -> None:
|
||||
with self._lock:
|
||||
if self._thread and self._thread.is_alive():
|
||||
return
|
||||
self._stop_event.clear()
|
||||
self._thread = threading.Thread(target=self._run, daemon=True, name="r2-backup-scheduler")
|
||||
self._thread.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
with self._lock:
|
||||
self._stop_event.set()
|
||||
thread = self._thread
|
||||
self._thread = None
|
||||
if thread and thread.is_alive():
|
||||
thread.join(timeout=2)
|
||||
|
||||
def _run(self) -> None:
|
||||
while not self._stop_event.is_set():
|
||||
try:
|
||||
self.run_scheduled_backup_if_needed()
|
||||
except Exception:
|
||||
pass
|
||||
self._stop_event.wait(30)
|
||||
|
||||
def run_scheduled_backup_if_needed(self) -> None:
|
||||
settings = config.get_backup_settings()
|
||||
if not settings.get("enabled"):
|
||||
return
|
||||
state = self.get_status()
|
||||
if state.get("running"):
|
||||
return
|
||||
interval_minutes = int(settings.get("interval_minutes") or 360)
|
||||
last_finished_raw = _clean(state.get("last_finished_at"))
|
||||
if last_finished_raw:
|
||||
try:
|
||||
last_finished = datetime.fromisoformat(last_finished_raw.replace("Z", "+00:00"))
|
||||
elapsed = (_utc_now() - last_finished.astimezone(UTC)).total_seconds()
|
||||
if elapsed < interval_minutes * 60:
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
self.run_backup(trigger="schedule")
|
||||
|
||||
def get_status(self) -> dict[str, object]:
|
||||
return {
|
||||
**load_backup_state(),
|
||||
"running": self._running,
|
||||
}
|
||||
|
||||
def is_configured(self) -> bool:
|
||||
settings = config.get_backup_settings()
|
||||
return all([
|
||||
_clean(settings.get("account_id")),
|
||||
_clean(settings.get("access_key_id")),
|
||||
_clean(settings.get("secret_access_key")),
|
||||
_clean(settings.get("bucket")),
|
||||
])
|
||||
|
||||
def get_settings(self) -> dict[str, object]:
|
||||
settings = dict(config.get_backup_settings())
|
||||
settings["secret_access_key"] = "********" if _clean(settings.get("secret_access_key")) else ""
|
||||
settings["passphrase"] = "********" if _clean(settings.get("passphrase")) else ""
|
||||
return settings
|
||||
|
||||
def update_settings(self, payload: dict[str, object]) -> dict[str, object]:
|
||||
current = config.get_backup_settings()
|
||||
merged = dict(current)
|
||||
merged.update(dict(payload or {}))
|
||||
if "include" in payload and isinstance(payload.get("include"), dict):
|
||||
include = dict(current.get("include") or {})
|
||||
include.update(payload.get("include") or {})
|
||||
merged["include"] = include
|
||||
if payload.get("secret_access_key") == "********":
|
||||
merged["secret_access_key"] = current.get("secret_access_key")
|
||||
if payload.get("passphrase") == "********":
|
||||
merged["passphrase"] = current.get("passphrase")
|
||||
updated = config.update({"backup": merged})
|
||||
return dict(updated.get("backup") or {})
|
||||
|
||||
def test_connection(self) -> dict[str, object]:
|
||||
client = CloudflareR2Client(config.get_backup_settings())
|
||||
try:
|
||||
return client.test_connection()
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
def list_backups(self) -> list[dict[str, object]]:
|
||||
if not self.is_configured():
|
||||
return []
|
||||
client = CloudflareR2Client(config.get_backup_settings())
|
||||
try:
|
||||
items = client.list_objects()
|
||||
finally:
|
||||
client.close()
|
||||
parsed: list[dict[str, object]] = []
|
||||
for item in items:
|
||||
key = _clean(item.get("key"))
|
||||
name = key.rsplit("/", 1)[-1]
|
||||
encrypted = name.endswith(".enc")
|
||||
parsed.append({
|
||||
"key": key,
|
||||
"name": name,
|
||||
"size": int(item.get("size") or 0),
|
||||
"updated_at": item.get("updated_at"),
|
||||
"encrypted": encrypted,
|
||||
})
|
||||
return parsed
|
||||
|
||||
def delete_backup(self, key: str) -> None:
|
||||
candidate = _clean(key)
|
||||
if not candidate:
|
||||
raise BackupError("备份对象 key 不能为空")
|
||||
client = CloudflareR2Client(config.get_backup_settings())
|
||||
try:
|
||||
client.delete_object(candidate)
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
def download_backup(self, key: str) -> dict[str, object]:
|
||||
candidate = _clean(key)
|
||||
if not candidate:
|
||||
raise BackupError("备份对象 key 不能为空")
|
||||
client = CloudflareR2Client(config.get_backup_settings())
|
||||
try:
|
||||
payload = client.download_bytes(candidate)
|
||||
finally:
|
||||
client.close()
|
||||
name = candidate.rsplit("/", 1)[-1] or "backup.bin"
|
||||
if candidate.endswith(".enc"):
|
||||
passphrase = _clean(config.get_backup_settings().get("passphrase"))
|
||||
if not passphrase:
|
||||
raise BackupError("当前未配置加密口令，无法下载并解密已加密备份")
|
||||
payload = _openssl_decrypt(payload, passphrase)
|
||||
if name.endswith(".enc"):
|
||||
name = name[:-4] or "backup.tar.gz"
|
||||
return {
|
||||
"key": candidate,
|
||||
"name": name,
|
||||
"content_type": _guess_content_type(name),
|
||||
"payload": payload,
|
||||
"size": len(payload),
|
||||
}
|
||||
|
||||
def get_backup_detail(self, key: str) -> dict[str, object]:
|
||||
candidate = _clean(key)
|
||||
if not candidate:
|
||||
raise BackupError("备份对象 key 不能为空")
|
||||
client = CloudflareR2Client(config.get_backup_settings())
|
||||
try:
|
||||
payload = client.download_bytes(candidate)
|
||||
finally:
|
||||
client.close()
|
||||
detail = self._decode_backup_payload(candidate, payload)
|
||||
detail["key"] = candidate
|
||||
detail["name"] = candidate.rsplit("/", 1)[-1]
|
||||
detail["encrypted"] = candidate.endswith(".enc")
|
||||
return detail
|
||||
|
||||
def run_backup(self, *, trigger: str = "manual") -> dict[str, object]:
|
||||
with self._lock:
|
||||
current = self.get_status()
|
||||
if self._running:
|
||||
raise BackupError("当前已有备份任务正在执行")
|
||||
started_at = _iso_now()
|
||||
self._running = True
|
||||
save_backup_state({
|
||||
"last_started_at": started_at,
|
||||
"last_finished_at": current.get("last_finished_at"),
|
||||
"last_status": "idle",
|
||||
"last_error": None,
|
||||
"last_object_key": current.get("last_object_key"),
|
||||
})
|
||||
try:
|
||||
result = self._run_backup_once(trigger=trigger)
|
||||
save_backup_state({
|
||||
"last_started_at": started_at,
|
||||
"last_finished_at": _iso_now(),
|
||||
"last_status": "success",
|
||||
"last_error": None,
|
||||
"last_object_key": result["key"],
|
||||
})
|
||||
return result
|
||||
except Exception as exc:
|
||||
save_backup_state({
|
||||
"last_started_at": started_at,
|
||||
"last_finished_at": _iso_now(),
|
||||
"last_status": "error",
|
||||
"last_error": str(exc) or exc.__class__.__name__,
|
||||
"last_object_key": current.get("last_object_key"),
|
||||
})
|
||||
raise
|
||||
finally:
|
||||
self._running = False
|
||||
|
||||
def _run_backup_once(self, *, trigger: str) -> dict[str, object]:
|
||||
settings = config.get_backup_settings()
|
||||
client = CloudflareR2Client(settings)
|
||||
client.validate()
|
||||
payload_raw = self._build_backup_archive(settings, trigger=trigger)
|
||||
encrypted = bool(settings.get("encrypt"))
|
||||
if encrypted:
|
||||
passphrase = _clean(settings.get("passphrase"))
|
||||
if not passphrase:
|
||||
raise BackupError("已启用备份加密，但未设置加密口令")
|
||||
payload = _openssl_encrypt(payload_raw, passphrase)
|
||||
suffix = ".tar.gz.enc"
|
||||
else:
|
||||
payload = payload_raw
|
||||
suffix = ".tar.gz"
|
||||
timestamp = _utc_now().strftime("%Y%m%dT%H%M%SZ")
|
||||
random_tag = f"{random.randint(0, 0xFFFF):04x}"
|
||||
object_key = f"{client.prefix.rstrip('/')}/backup-{timestamp}-{random_tag}{suffix}"
|
||||
metadata = {
|
||||
"created-at": _iso_now(),
|
||||
"encrypted": "true" if encrypted else "false",
|
||||
"trigger": trigger,
|
||||
}
|
||||
try:
|
||||
result = client.upload_bytes(object_key, payload, content_type="application/octet-stream", metadata=metadata)
|
||||
self._apply_rotation(client, int(settings.get("rotation_keep") or 0))
|
||||
return {
|
||||
"key": result["key"],
|
||||
"size": len(payload),
|
||||
"encrypted": encrypted,
|
||||
}
|
||||
finally:
|
||||
client.close()
|
||||
|
||||
def _decode_backup_payload(self, key: str, payload: bytes) -> dict[str, object]:
|
||||
decoded = payload
|
||||
if key.endswith(".enc"):
|
||||
passphrase = _clean(config.get_backup_settings().get("passphrase"))
|
||||
if not passphrase:
|
||||
raise BackupError("当前未配置加密口令，无法查看已加密备份")
|
||||
decoded = _openssl_decrypt(decoded, passphrase)
|
||||
return self._decode_archive_detail(decoded)
|
||||
|
||||
def _apply_rotation(self, client: CloudflareR2Client, keep: int) -> None:
|
||||
if keep <= 0:
|
||||
return
|
||||
items = [item for item in client.list_objects() if _is_backup_object(item.get("key"))]
|
||||
if len(items) <= keep:
|
||||
return
|
||||
for item in items[keep:]:
|
||||
key = _clean(item.get("key"))
|
||||
if key:
|
||||
client.delete_object(key)
|
||||
|
||||
def _decode_archive_detail(self, payload: bytes) -> dict[str, object]:
|
||||
files: list[dict[str, object]] = []
|
||||
snapshots: list[dict[str, object]] = []
|
||||
metadata: dict[str, object] = {}
|
||||
try:
|
||||
with tarfile.open(fileobj=io.BytesIO(payload), mode="r:gz") as archive:
|
||||
members = [member for member in archive.getmembers() if member.isfile()]
|
||||
for member in members:
|
||||
extracted = archive.extractfile(member)
|
||||
if extracted is None:
|
||||
continue
|
||||
raw = extracted.read()
|
||||
name = member.name
|
||||
if name == "backup-metadata.json":
|
||||
try:
|
||||
parsed = json.loads(raw.decode("utf-8"))
|
||||
if isinstance(parsed, dict):
|
||||
metadata = parsed
|
||||
except Exception:
|
||||
metadata = {}
|
||||
continue
|
||||
if name.startswith("snapshots/") and name.endswith(".json"):
|
||||
count = 0
|
||||
try:
|
||||
parsed_snapshot = json.loads(raw.decode("utf-8"))
|
||||
count = _count_items(parsed_snapshot)
|
||||
except Exception:
|
||||
count = 0
|
||||
snapshots.append({
|
||||
"name": name.removeprefix("snapshots/").removesuffix(".json"),
|
||||
"count": count,
|
||||
})
|
||||
continue
|
||||
files.append({
|
||||
"name": name,
|
||||
"exists": True,
|
||||
"content_type": _guess_content_type(name),
|
||||
"size": len(raw),
|
||||
"sha256": _sha256_hex(raw),
|
||||
})
|
||||
except tarfile.TarError as exc:
|
||||
raise BackupError("解析备份压缩包失败，备份可能已损坏") from exc
|
||||
files.sort(key=lambda item: str(item.get("name") or ""))
|
||||
snapshots.sort(key=lambda item: str(item.get("name") or ""))
|
||||
return {
|
||||
"created_at": metadata.get("created_at"),
|
||||
"trigger": metadata.get("trigger"),
|
||||
"app_version": metadata.get("app_version"),
|
||||
"storage_backend": metadata.get("storage_backend"),
|
||||
"files": files,
|
||||
"snapshots": snapshots,
|
||||
}
|
||||
|
||||
def _build_backup_archive(self, settings: dict[str, object], *, trigger: str) -> bytes:
|
||||
include = settings.get("include") if isinstance(settings.get("include"), dict) else {}
|
||||
metadata = {
|
||||
"version": 2,
|
||||
"created_at": _iso_now(),
|
||||
"trigger": trigger,
|
||||
"app_version": config.app_version,
|
||||
"storage_backend": config.get_storage_backend().get_backend_info(),
|
||||
}
|
||||
buffer = io.BytesIO()
|
||||
with tarfile.open(fileobj=buffer, mode="w:gz") as archive:
|
||||
self._add_bytes_to_archive(archive, "backup-metadata.json", _json_bytes(metadata))
|
||||
if include.get("config"):
|
||||
self._add_file_to_archive(archive, CONFIG_FILE, "config.json")
|
||||
if include.get("cpa"):
|
||||
self._add_file_to_archive(archive, DATA_DIR / "cpa_config.json", "data/cpa_config.json")
|
||||
if include.get("sub2api"):
|
||||
self._add_file_to_archive(archive, DATA_DIR / "sub2api_config.json", "data/sub2api_config.json")
|
||||
if include.get("logs"):
|
||||
self._add_file_to_archive(archive, DATA_DIR / "logs.jsonl", "data/logs.jsonl")
|
||||
if include.get("image_tasks"):
|
||||
self._add_file_to_archive(archive, DATA_DIR / "image_tasks.json", "data/image_tasks.json")
|
||||
self._add_file_to_archive(archive, IMAGE_INDEX_FILE, "data/image_index.json")
|
||||
if include.get("accounts_snapshot"):
|
||||
self._add_bytes_to_archive(
|
||||
archive,
|
||||
"snapshots/accounts.json",
|
||||
_json_bytes(config.get_storage_backend().load_accounts()),
|
||||
)
|
||||
if include.get("auth_keys_snapshot"):
|
||||
self._add_bytes_to_archive(
|
||||
archive,
|
||||
"snapshots/auth_keys.json",
|
||||
_json_bytes(config.get_storage_backend().load_auth_keys()),
|
||||
)
|
||||
if include.get("images"):
|
||||
self._add_file_to_archive(archive, TAGS_FILE, "data/image_tags.json")
|
||||
self._add_directory_to_archive(archive, config.images_dir, "data/images")
|
||||
return buffer.getvalue()
|
||||
|
||||
def _add_bytes_to_archive(self, archive: tarfile.TarFile, name: str, payload: bytes) -> None:
|
||||
info = tarfile.TarInfo(name=name)
|
||||
info.size = len(payload)
|
||||
info.mtime = int(_utc_now().timestamp())
|
||||
archive.addfile(info, io.BytesIO(payload))
|
||||
|
||||
def _add_file_to_archive(self, archive: tarfile.TarFile, source: Path, arcname: str) -> None:
|
||||
if not source.exists() or not source.is_file():
|
||||
return
|
||||
archive.add(source, arcname=arcname)
|
||||
|
||||
def _add_directory_to_archive(self, archive: tarfile.TarFile, source_dir: Path, arcname_root: str) -> None:
|
||||
if not source_dir.exists() or not source_dir.is_dir():
|
||||
return
|
||||
for path in sorted(source_dir.rglob("*")):
|
||||
if path.is_file():
|
||||
relative = path.relative_to(source_dir).as_posix()
|
||||
archive.add(path, arcname=f"{arcname_root}/{relative}")
|
||||
|
||||
|
||||
backup_service = BackupService()
|
||||
@@ -0,0 +1,667 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from dataclasses import dataclass
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import time
|
||||
|
||||
from services.storage.base import StorageBackend
|
||||
|
||||
BASE_DIR = Path(__file__).resolve().parents[1]
|
||||
DATA_DIR = BASE_DIR / "data"
|
||||
CONFIG_FILE = BASE_DIR / "config.json"
|
||||
VERSION_FILE = BASE_DIR / "VERSION"
|
||||
BACKUP_STATE_FILE = DATA_DIR / "backup_state.json"
|
||||
|
||||
DEFAULT_BACKUP_INCLUDE = {
|
||||
"config": True,
|
||||
"cpa": True,
|
||||
"sub2api": True,
|
||||
"logs": True,
|
||||
"image_tasks": True,
|
||||
"accounts_snapshot": True,
|
||||
"auth_keys_snapshot": True,
|
||||
"images": False,
|
||||
}
|
||||
|
||||
DEFAULT_IMAGE_STORAGE = {
|
||||
"enabled": False,
|
||||
"mode": "local",
|
||||
"webdav_url": "",
|
||||
"webdav_username": "",
|
||||
"webdav_password": "",
|
||||
"webdav_root_path": "chatgpt2api/images",
|
||||
"public_base_url": "",
|
||||
}
|
||||
|
||||
DEFAULT_CHAT_COMPLETION_CACHE = {
|
||||
"enabled": True,
|
||||
"ttl_seconds": 60,
|
||||
"max_entries": 256,
|
||||
"dedupe_inflight": True,
|
||||
"stream_cache": True,
|
||||
"normalize_messages": True,
|
||||
"drop_adjacent_duplicates": True,
|
||||
"drop_assistant_history": False,
|
||||
}
|
||||
|
||||
DEFAULT_PROXY_RUNTIME_USER_AGENT = (
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
"Chrome/145.0.0.0 Safari/537.36"
|
||||
)
|
||||
|
||||
DEFAULT_PROXY_RUNTIME = {
|
||||
"enabled": False,
|
||||
"egress_mode": "direct",
|
||||
"proxy_url": "",
|
||||
"resource_proxy_url": "",
|
||||
"skip_ssl_verify": False,
|
||||
"reset_session_status_codes": [403],
|
||||
"clearance": {
|
||||
"enabled": False,
|
||||
"mode": "none",
|
||||
"cf_cookies": "",
|
||||
"cf_clearance": "",
|
||||
"user_agent": DEFAULT_PROXY_RUNTIME_USER_AGENT,
|
||||
"browser": "chrome",
|
||||
"flaresolverr_url": "",
|
||||
"timeout_sec": 60,
|
||||
"refresh_interval": 3600,
|
||||
"warm_up_on_start": False,
|
||||
},
|
||||
}
|
||||
|
||||
DEFAULT_THIRD_PARTY_APPS = {
|
||||
"infinite_canvas": {
|
||||
"enabled": False,
|
||||
"url": "https://canvas.best",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _normalize_bool(value: object, default: bool = False) -> bool:
|
||||
if isinstance(value, str):
|
||||
lowered = value.strip().lower()
|
||||
if lowered in {"1", "true", "yes", "on"}:
|
||||
return True
|
||||
if lowered in {"0", "false", "no", "off"}:
|
||||
return False
|
||||
return default
|
||||
if value is None:
|
||||
return default
|
||||
return bool(value)
|
||||
|
||||
|
||||
def _normalize_positive_int(value: object, default: int, minimum: int = 0) -> int:
|
||||
try:
|
||||
normalized = int(value)
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
normalized = default
|
||||
return max(minimum, normalized)
|
||||
|
||||
|
||||
def _normalize_backup_include(value: object) -> dict[str, bool]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
normalized = dict(DEFAULT_BACKUP_INCLUDE)
|
||||
for key in normalized:
|
||||
normalized[key] = _normalize_bool(source.get(key), normalized[key])
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_backup_settings(value: object) -> dict[str, object]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
return {
|
||||
"enabled": _normalize_bool(source.get("enabled"), False),
|
||||
"provider": "cloudflare_r2",
|
||||
"account_id": str(source.get("account_id") or "").strip(),
|
||||
"access_key_id": str(source.get("access_key_id") or "").strip(),
|
||||
"secret_access_key": str(source.get("secret_access_key") or "").strip(),
|
||||
"bucket": str(source.get("bucket") or "").strip(),
|
||||
"prefix": str(source.get("prefix") or "backups").strip().strip("/") or "backups",
|
||||
"interval_minutes": _normalize_positive_int(source.get("interval_minutes"), 360, 1),
|
||||
"rotation_keep": _normalize_positive_int(source.get("rotation_keep"), 10, 0),
|
||||
"encrypt": _normalize_bool(source.get("encrypt"), False),
|
||||
"passphrase": str(source.get("passphrase") or "").strip(),
|
||||
"include": _normalize_backup_include(source.get("include")),
|
||||
}
|
||||
|
||||
|
||||
def _normalize_backup_state(value: object) -> dict[str, object]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
return {
|
||||
"last_started_at": str(source.get("last_started_at") or "").strip() or None,
|
||||
"last_finished_at": str(source.get("last_finished_at") or "").strip() or None,
|
||||
"last_status": str(source.get("last_status") or "idle").strip() or "idle",
|
||||
"last_error": str(source.get("last_error") or "").strip() or None,
|
||||
"last_object_key": str(source.get("last_object_key") or "").strip() or None,
|
||||
}
|
||||
|
||||
|
||||
def _normalize_image_storage_settings(value: object) -> dict[str, object]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
mode = str(source.get("mode") or "local").strip().lower()
|
||||
if mode not in {"local", "webdav", "both"}:
|
||||
mode = "local"
|
||||
enabled = _normalize_bool(source.get("enabled"), False)
|
||||
if not enabled:
|
||||
mode = "local"
|
||||
root_path = str(source.get("webdav_root_path") or DEFAULT_IMAGE_STORAGE["webdav_root_path"]).strip().strip("/")
|
||||
return {
|
||||
"enabled": enabled,
|
||||
"mode": mode,
|
||||
"webdav_url": str(source.get("webdav_url") or "").strip().rstrip("/"),
|
||||
"webdav_username": str(source.get("webdav_username") or "").strip(),
|
||||
"webdav_password": str(source.get("webdav_password") or "").strip(),
|
||||
"webdav_root_path": root_path or str(DEFAULT_IMAGE_STORAGE["webdav_root_path"]),
|
||||
"public_base_url": str(source.get("public_base_url") or "").strip().rstrip("/"),
|
||||
}
|
||||
|
||||
|
||||
def _normalize_chat_completion_cache_settings(value: object) -> dict[str, object]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
return {
|
||||
"enabled": _normalize_bool(source.get("enabled"), DEFAULT_CHAT_COMPLETION_CACHE["enabled"]),
|
||||
"ttl_seconds": _normalize_positive_int(
|
||||
source.get("ttl_seconds"),
|
||||
int(DEFAULT_CHAT_COMPLETION_CACHE["ttl_seconds"]),
|
||||
0,
|
||||
),
|
||||
"max_entries": _normalize_positive_int(
|
||||
source.get("max_entries"),
|
||||
int(DEFAULT_CHAT_COMPLETION_CACHE["max_entries"]),
|
||||
1,
|
||||
),
|
||||
"dedupe_inflight": _normalize_bool(
|
||||
source.get("dedupe_inflight"),
|
||||
bool(DEFAULT_CHAT_COMPLETION_CACHE["dedupe_inflight"]),
|
||||
),
|
||||
"stream_cache": _normalize_bool(
|
||||
source.get("stream_cache"),
|
||||
bool(DEFAULT_CHAT_COMPLETION_CACHE["stream_cache"]),
|
||||
),
|
||||
"normalize_messages": _normalize_bool(
|
||||
source.get("normalize_messages"),
|
||||
bool(DEFAULT_CHAT_COMPLETION_CACHE["normalize_messages"]),
|
||||
),
|
||||
"drop_adjacent_duplicates": _normalize_bool(
|
||||
source.get("drop_adjacent_duplicates"),
|
||||
bool(DEFAULT_CHAT_COMPLETION_CACHE["drop_adjacent_duplicates"]),
|
||||
),
|
||||
"drop_assistant_history": _normalize_bool(
|
||||
source.get("drop_assistant_history"),
|
||||
bool(DEFAULT_CHAT_COMPLETION_CACHE["drop_assistant_history"]),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _normalize_status_codes(value: object) -> list[int]:
|
||||
items = value if isinstance(value, list) else DEFAULT_PROXY_RUNTIME["reset_session_status_codes"]
|
||||
normalized: list[int] = []
|
||||
for item in items:
|
||||
if isinstance(item, bool):
|
||||
continue
|
||||
try:
|
||||
status = int(item)
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
continue
|
||||
if 100 <= status <= 599 and status not in normalized:
|
||||
normalized.append(status)
|
||||
if not normalized:
|
||||
return list(DEFAULT_PROXY_RUNTIME["reset_session_status_codes"])
|
||||
return normalized
|
||||
|
||||
|
||||
def _normalize_proxy_runtime_settings(value: object) -> dict[str, object]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
default_clearance = DEFAULT_PROXY_RUNTIME["clearance"]
|
||||
clearance_source = source.get("clearance") if isinstance(source.get("clearance"), dict) else {}
|
||||
|
||||
egress_mode = str(source.get("egress_mode") or DEFAULT_PROXY_RUNTIME["egress_mode"]).strip().lower()
|
||||
if egress_mode not in {"direct", "single_proxy"}:
|
||||
egress_mode = str(DEFAULT_PROXY_RUNTIME["egress_mode"])
|
||||
|
||||
clearance_mode = str(clearance_source.get("mode") or default_clearance["mode"]).strip().lower()
|
||||
if clearance_mode not in {"none", "manual", "flaresolverr"}:
|
||||
clearance_mode = str(default_clearance["mode"])
|
||||
|
||||
user_agent = str(clearance_source.get("user_agent") or default_clearance["user_agent"]).strip()
|
||||
browser = str(clearance_source.get("browser") or default_clearance["browser"]).strip()
|
||||
|
||||
existing_clearance_cookies = str(source.get("_existing_cf_cookies") or "").strip()
|
||||
existing_cf_clearance = str(source.get("_existing_cf_clearance") or "").strip()
|
||||
cf_cookies = str(clearance_source.get("cf_cookies") or "").strip()
|
||||
cf_clearance = str(clearance_source.get("cf_clearance") or "").strip()
|
||||
if not cf_cookies and _normalize_bool(clearance_source.get("has_cf_cookies"), False):
|
||||
cf_cookies = existing_clearance_cookies
|
||||
if not cf_clearance and _normalize_bool(clearance_source.get("has_cf_clearance"), False):
|
||||
cf_clearance = existing_cf_clearance
|
||||
|
||||
return {
|
||||
"enabled": _normalize_bool(source.get("enabled"), bool(DEFAULT_PROXY_RUNTIME["enabled"])),
|
||||
"egress_mode": egress_mode,
|
||||
"proxy_url": str(source.get("proxy_url") or "").strip(),
|
||||
"resource_proxy_url": str(source.get("resource_proxy_url") or "").strip(),
|
||||
"skip_ssl_verify": _normalize_bool(
|
||||
source.get("skip_ssl_verify"),
|
||||
bool(DEFAULT_PROXY_RUNTIME["skip_ssl_verify"]),
|
||||
),
|
||||
"reset_session_status_codes": _normalize_status_codes(source.get("reset_session_status_codes")),
|
||||
"clearance": {
|
||||
"enabled": _normalize_bool(clearance_source.get("enabled"), bool(default_clearance["enabled"])),
|
||||
"mode": clearance_mode,
|
||||
"cf_cookies": cf_cookies,
|
||||
"cf_clearance": cf_clearance,
|
||||
"user_agent": user_agent or str(default_clearance["user_agent"]),
|
||||
"browser": browser or str(default_clearance["browser"]),
|
||||
"flaresolverr_url": str(clearance_source.get("flaresolverr_url") or "").strip(),
|
||||
"timeout_sec": _normalize_positive_int(
|
||||
clearance_source.get("timeout_sec"),
|
||||
int(default_clearance["timeout_sec"]),
|
||||
1,
|
||||
),
|
||||
"refresh_interval": _normalize_positive_int(
|
||||
clearance_source.get("refresh_interval"),
|
||||
int(default_clearance["refresh_interval"]),
|
||||
60,
|
||||
),
|
||||
"warm_up_on_start": _normalize_bool(
|
||||
clearance_source.get("warm_up_on_start"),
|
||||
bool(default_clearance["warm_up_on_start"]),
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _normalize_third_party_apps_settings(value: object) -> dict[str, object]:
|
||||
source = value if isinstance(value, dict) else {}
|
||||
canvas_source = source.get("infinite_canvas") if isinstance(source.get("infinite_canvas"), dict) else {}
|
||||
return {
|
||||
"infinite_canvas": {
|
||||
"enabled": _normalize_bool(canvas_source.get("enabled"), False),
|
||||
"url": str(canvas_source.get("url") or DEFAULT_THIRD_PARTY_APPS["infinite_canvas"]["url"]).strip(),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _validate_image_storage_settings(settings: dict[str, object]) -> None:
|
||||
if not _normalize_bool(settings.get("enabled"), False):
|
||||
return
|
||||
if not str(settings.get("webdav_url") or "").strip():
|
||||
raise ValueError("启用 WebDAV 图片存储后必须填写 WebDAV URL")
|
||||
if not str(settings.get("webdav_password") or "").strip():
|
||||
raise ValueError("启用 WebDAV 图片存储后必须填写 WebDAV 密码")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LoadedSettings:
|
||||
auth_key: str
|
||||
refresh_account_interval_minute: int
|
||||
|
||||
|
||||
def _normalize_auth_key(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def _is_invalid_auth_key(value: object) -> bool:
|
||||
return _normalize_auth_key(value) == ""
|
||||
|
||||
|
||||
def _read_json_object(path: Path, *, name: str) -> dict[str, object]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
if path.is_dir():
|
||||
print(
|
||||
f"Warning: {name} at '{path}' is a directory, ignoring it and falling back to other configuration sources.",
|
||||
file=sys.stderr,
|
||||
)
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def _load_settings() -> LoadedSettings:
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
raw_config = _read_json_object(CONFIG_FILE, name="config.json")
|
||||
auth_key = _normalize_auth_key(os.getenv("CHATGPT2API_AUTH_KEY") or raw_config.get("auth-key"))
|
||||
if _is_invalid_auth_key(auth_key):
|
||||
raise ValueError(
|
||||
"❌ auth-key 未设置！\n"
|
||||
"请在环境变量 CHATGPT2API_AUTH_KEY 中设置，或者在 config.json 中填写 auth-key。"
|
||||
)
|
||||
|
||||
try:
|
||||
refresh_interval = int(raw_config.get("refresh_account_interval_minute", 5))
|
||||
except (TypeError, ValueError):
|
||||
refresh_interval = 5
|
||||
|
||||
return LoadedSettings(
|
||||
auth_key=auth_key,
|
||||
refresh_account_interval_minute=refresh_interval,
|
||||
)
|
||||
|
||||
|
||||
class ConfigStore:
|
||||
def __init__(self, path: Path):
|
||||
self.path = path
|
||||
DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
self.data = self._load()
|
||||
self._storage_backend: StorageBackend | None = None
|
||||
if _is_invalid_auth_key(self.auth_key):
|
||||
raise ValueError(
|
||||
"❌ auth-key 未设置！\n"
|
||||
"请按以下任意一种方式解决：\n"
|
||||
"1. 在 Render 的 Environment 变量中添加：\n"
|
||||
" CHATGPT2API_AUTH_KEY = your_real_auth_key\n"
|
||||
"2. 或者在 config.json 中填写：\n"
|
||||
' "auth-key": "your_real_auth_key"'
|
||||
)
|
||||
|
||||
def _load(self) -> dict[str, object]:
|
||||
return _read_json_object(self.path, name="config.json")
|
||||
|
||||
def _save(self) -> None:
|
||||
self.path.write_text(json.dumps(self.data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
@property
|
||||
def auth_key(self) -> str:
|
||||
return _normalize_auth_key(os.getenv("CHATGPT2API_AUTH_KEY") or self.data.get("auth-key"))
|
||||
|
||||
@property
|
||||
def accounts_file(self) -> Path:
|
||||
return DATA_DIR / "accounts.json"
|
||||
|
||||
@property
|
||||
def refresh_account_interval_minute(self) -> int:
|
||||
try:
|
||||
return int(self.data.get("refresh_account_interval_minute", 5))
|
||||
except (TypeError, ValueError):
|
||||
return 5
|
||||
|
||||
@property
|
||||
def image_retention_days(self) -> int:
|
||||
try:
|
||||
return max(1, int(self.data.get("image_retention_days", 30)))
|
||||
except (TypeError, ValueError):
|
||||
return 30
|
||||
|
||||
@property
|
||||
def image_poll_timeout_secs(self) -> int:
|
||||
try:
|
||||
return max(1, int(self.data.get("image_poll_timeout_secs", 120)))
|
||||
except (TypeError, ValueError):
|
||||
return 120
|
||||
|
||||
@property
|
||||
def image_poll_interval_secs(self) -> float:
|
||||
try:
|
||||
return max(0.5, float(self.data.get("image_poll_interval_secs", 10.0)))
|
||||
except (TypeError, ValueError):
|
||||
return 10.0
|
||||
|
||||
@property
|
||||
def image_poll_initial_wait_secs(self) -> float:
|
||||
"""Image generation upstream takes ~30s; polling immediately wastes requests
|
||||
and trips a transient 429. Default 10s gives the conversation document time
|
||||
to commit before the first poll."""
|
||||
try:
|
||||
return max(0.0, float(self.data.get("image_poll_initial_wait_secs", 10.0)))
|
||||
except (TypeError, ValueError):
|
||||
return 10.0
|
||||
|
||||
@property
|
||||
def image_account_concurrency(self) -> int:
|
||||
try:
|
||||
return max(1, int(self.data.get("image_account_concurrency", 3)))
|
||||
except (TypeError, ValueError):
|
||||
return 3
|
||||
|
||||
@property
|
||||
def image_parallel_generation(self) -> bool:
|
||||
value = self.data.get("image_parallel_generation", True)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def image_settle_enabled(self) -> bool:
|
||||
"""图片二次确认机制：找到 file_ids 后等待一段时间再次确认。"""
|
||||
value = self.data.get("image_settle_enabled", True)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def image_check_before_hit_enabled(self) -> bool:
|
||||
"""先check再hit：通过轮询确认 file_ids 存在后再返回，而非仅依赖 SSE 事件。"""
|
||||
value = self.data.get("image_check_before_hit_enabled", True)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def image_remove_conversation_after_result(self) -> bool:
|
||||
"""出图成功后异步隐藏 ChatGPT 本地对话记录。"""
|
||||
value = self.data.get("image_remove_conversation_after_result", False)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def image_remove_conversation_always(self) -> bool:
|
||||
"""无论是否出图，画图请求结束后都异步隐藏 ChatGPT 本地对话记录。"""
|
||||
return _normalize_bool(self.data.get("image_remove_conversation_always"), False)
|
||||
|
||||
@property
|
||||
def image_settle_secs(self) -> float:
|
||||
"""二次确认等待时间（秒）。"""
|
||||
try:
|
||||
return max(0.5, float(self.data.get("image_settle_secs", 2.0)))
|
||||
except (TypeError, ValueError):
|
||||
return 2.0
|
||||
|
||||
@property
|
||||
def auto_remove_invalid_accounts(self) -> bool:
|
||||
value = self.data.get("auto_remove_invalid_accounts", False)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def auto_remove_rate_limited_accounts(self) -> bool:
|
||||
value = self.data.get("auto_remove_rate_limited_accounts", False)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def auto_relogin_after_refresh(self) -> bool:
|
||||
value = self.data.get("auto_relogin_after_refresh", False)
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
@property
|
||||
def log_levels(self) -> list[str]:
|
||||
levels = self.data.get("log_levels")
|
||||
if not isinstance(levels, list):
|
||||
return []
|
||||
allowed = {"debug", "info", "warning", "error"}
|
||||
return [level for item in levels if (level := str(item or "").strip().lower()) in allowed]
|
||||
|
||||
@property
|
||||
def sensitive_words(self) -> list[str]:
|
||||
words = self.data.get("sensitive_words")
|
||||
return [word for item in words if (word := str(item or "").strip())] if isinstance(words, list) else []
|
||||
|
||||
@property
|
||||
def ai_review(self) -> dict[str, object]:
|
||||
value = self.data.get("ai_review")
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
||||
@property
|
||||
def global_system_prompt(self) -> str:
|
||||
return str(self.data.get("global_system_prompt") or "").strip()
|
||||
|
||||
@property
|
||||
def default_upstream_model_name(self) -> str:
|
||||
return str(self.data.get("default_upstream_model_name") or "gpt-5-5").strip()
|
||||
|
||||
@property
|
||||
def default_thinking_effort(self) -> str:
|
||||
value = str(self.data.get("default_thinking_effort") or "auto").strip().lower()
|
||||
return value if value in {"auto", "standard", "extended", "max"} else "auto"
|
||||
|
||||
@property
|
||||
def images_dir(self) -> Path:
|
||||
path = DATA_DIR / "images"
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
@property
|
||||
def image_thumbnails_dir(self) -> Path:
|
||||
path = DATA_DIR / "image_thumbnails"
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
def cleanup_old_images(self) -> int:
|
||||
cutoff = time.time() - self.image_retention_days * 86400
|
||||
removed = 0
|
||||
for path in self.images_dir.rglob("*"):
|
||||
if path.is_file() and path.stat().st_mtime < cutoff:
|
||||
path.unlink()
|
||||
removed += 1
|
||||
for path in sorted((p for p in self.images_dir.rglob("*") if p.is_dir()), key=lambda p: len(p.parts), reverse=True):
|
||||
try:
|
||||
path.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
return removed
|
||||
|
||||
@property
|
||||
def base_url(self) -> str:
|
||||
return str(
|
||||
os.getenv("CHATGPT2API_BASE_URL")
|
||||
or self.data.get("base_url")
|
||||
or ""
|
||||
).strip().rstrip("/")
|
||||
|
||||
@property
|
||||
def app_version(self) -> str:
|
||||
try:
|
||||
value = VERSION_FILE.read_text(encoding="utf-8").strip()
|
||||
except FileNotFoundError:
|
||||
return "0.0.0"
|
||||
return value or "0.0.0"
|
||||
|
||||
def get(self) -> dict[str, object]:
|
||||
data = dict(self.data)
|
||||
data["refresh_account_interval_minute"] = self.refresh_account_interval_minute
|
||||
data["image_retention_days"] = self.image_retention_days
|
||||
data["image_poll_timeout_secs"] = self.image_poll_timeout_secs
|
||||
data["image_poll_interval_secs"] = self.image_poll_interval_secs
|
||||
data["image_poll_initial_wait_secs"] = self.image_poll_initial_wait_secs
|
||||
data["image_account_concurrency"] = self.image_account_concurrency
|
||||
data["image_parallel_generation"] = self.image_parallel_generation
|
||||
data["image_remove_conversation_after_result"] = self.image_remove_conversation_after_result
|
||||
data["image_remove_conversation_always"] = self.image_remove_conversation_always
|
||||
data["auto_remove_invalid_accounts"] = self.auto_remove_invalid_accounts
|
||||
data["auto_remove_rate_limited_accounts"] = self.auto_remove_rate_limited_accounts
|
||||
data["auto_relogin_after_refresh"] = self.auto_relogin_after_refresh
|
||||
data["log_levels"] = self.log_levels
|
||||
data["sensitive_words"] = self.sensitive_words
|
||||
data["ai_review"] = self.ai_review
|
||||
data["global_system_prompt"] = self.global_system_prompt
|
||||
data["default_upstream_model_name"] = self.default_upstream_model_name
|
||||
data["default_thinking_effort"] = self.default_thinking_effort
|
||||
data["backup"] = self.get_backup_settings()
|
||||
data["image_storage"] = self.get_image_storage_settings()
|
||||
data["chat_completion_cache"] = self.get_chat_completion_cache_settings()
|
||||
data["proxy_runtime"] = self.get_public_proxy_runtime_settings()
|
||||
data["third_party_apps"] = self.get_third_party_apps_settings()
|
||||
data.pop("auth-key", None)
|
||||
return data
|
||||
|
||||
def get_proxy_settings(self) -> str:
|
||||
return str(self.data.get("proxy") or "").strip()
|
||||
|
||||
def get_proxy_runtime_settings(self) -> dict[str, object]:
|
||||
return _normalize_proxy_runtime_settings(self.data.get("proxy_runtime"))
|
||||
|
||||
def get_public_proxy_runtime_settings(self) -> dict[str, object]:
|
||||
runtime = copy.deepcopy(self.get_proxy_runtime_settings())
|
||||
clearance = runtime.get("clearance") if isinstance(runtime.get("clearance"), dict) else {}
|
||||
if isinstance(clearance, dict):
|
||||
cf_cookies = str(clearance.get("cf_cookies") or "").strip()
|
||||
cf_clearance = str(clearance.get("cf_clearance") or "").strip()
|
||||
clearance["cf_cookies"] = ""
|
||||
clearance["cf_clearance"] = ""
|
||||
clearance["has_cf_cookies"] = bool(cf_cookies)
|
||||
clearance["has_cf_clearance"] = bool(cf_clearance)
|
||||
return runtime
|
||||
|
||||
def get_third_party_apps_settings(self) -> dict[str, object]:
|
||||
return _normalize_third_party_apps_settings(self.data.get("third_party_apps"))
|
||||
|
||||
def update(self, data: dict[str, object]) -> dict[str, object]:
|
||||
next_data = dict(self.data)
|
||||
next_data.update(dict(data or {}))
|
||||
if "backup" in next_data:
|
||||
next_data["backup"] = _normalize_backup_settings(next_data.get("backup"))
|
||||
if "image_storage" in next_data:
|
||||
next_data["image_storage"] = _normalize_image_storage_settings(next_data.get("image_storage"))
|
||||
_validate_image_storage_settings(next_data["image_storage"])
|
||||
if "chat_completion_cache" in next_data:
|
||||
next_data["chat_completion_cache"] = _normalize_chat_completion_cache_settings(
|
||||
next_data.get("chat_completion_cache")
|
||||
)
|
||||
if "third_party_apps" in next_data:
|
||||
next_data["third_party_apps"] = _normalize_third_party_apps_settings(next_data.get("third_party_apps"))
|
||||
if "proxy_runtime" in next_data:
|
||||
incoming_runtime = next_data.get("proxy_runtime")
|
||||
if isinstance(incoming_runtime, dict):
|
||||
previous_clearance = self.get_proxy_runtime_settings().get("clearance")
|
||||
if isinstance(previous_clearance, dict):
|
||||
incoming_runtime = dict(incoming_runtime)
|
||||
incoming_runtime["_existing_cf_cookies"] = previous_clearance.get("cf_cookies")
|
||||
incoming_runtime["_existing_cf_clearance"] = previous_clearance.get("cf_clearance")
|
||||
next_data["proxy_runtime"] = _normalize_proxy_runtime_settings(incoming_runtime)
|
||||
next_data.pop("backup_state", None)
|
||||
self.data = next_data
|
||||
self._save()
|
||||
return self.get()
|
||||
|
||||
def get_backup_settings(self) -> dict[str, object]:
|
||||
return _normalize_backup_settings(self.data.get("backup"))
|
||||
|
||||
def get_image_storage_settings(self) -> dict[str, object]:
|
||||
return _normalize_image_storage_settings(self.data.get("image_storage"))
|
||||
|
||||
def get_chat_completion_cache_settings(self) -> dict[str, object]:
|
||||
return _normalize_chat_completion_cache_settings(self.data.get("chat_completion_cache"))
|
||||
|
||||
def get_storage_backend(self) -> StorageBackend:
|
||||
"""获取存储后端实例（单例）"""
|
||||
if self._storage_backend is None:
|
||||
from services.storage.factory import create_storage_backend
|
||||
self._storage_backend = create_storage_backend(DATA_DIR)
|
||||
return self._storage_backend
|
||||
|
||||
|
||||
def load_backup_state() -> dict[str, object]:
|
||||
return _normalize_backup_state(_read_json_object(BACKUP_STATE_FILE, name="backup_state.json"))
|
||||
|
||||
|
||||
def save_backup_state(state: dict[str, object]) -> dict[str, object]:
|
||||
normalized = _normalize_backup_state(state)
|
||||
BACKUP_STATE_FILE.write_text(json.dumps(normalized, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
return normalized
|
||||
|
||||
|
||||
config = ConfigStore(CONFIG_FILE)
|
||||
@@ -0,0 +1,242 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
|
||||
from curl_cffi import requests
|
||||
from fastapi import HTTPException
|
||||
|
||||
from services.config import config
|
||||
from services.proxy_service import proxy_settings
|
||||
from utils.log import logger
|
||||
|
||||
DEFAULT_REVIEW_PROMPT = "判断用户请求是否允许。只回答 ALLOW 或 REJECT。"
|
||||
|
||||
# Strip base64 image data URIs before review: a text-only review model can't
|
||||
# analyze image bytes, and a single inlined image easily blows past the token
|
||||
# budget of the upstream review service.
|
||||
_BASE64_DATA_URI = re.compile(r"data:[\w/.+;-]+;base64,[A-Za-z0-9+/=]+")
|
||||
|
||||
# Cap aligned to the upstream review service's max context. If text still
|
||||
# exceeds the cap after base64 stripping, keep equal head/tail halves so both
|
||||
# the system prompt and the most recent user message survive.
|
||||
_MAX_REVIEW_TEXT_LEN = 100_000
|
||||
_TRUNCATION_MARKER = "\n…[truncated]…\n"
|
||||
|
||||
|
||||
def _text(value: object) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if isinstance(value, list):
|
||||
return "\n".join(_text(item) for item in value)
|
||||
if isinstance(value, dict):
|
||||
return "\n".join(_text(value.get(key)) for key in ("text", "input_text", "content", "input", "instructions", "system", "prompt"))
|
||||
return ""
|
||||
|
||||
|
||||
def request_text(*values: object) -> str:
|
||||
return "\n".join(part for value in values if (part := _text(value).strip()))
|
||||
|
||||
|
||||
def request_shape(*values: object) -> dict[str, int]:
|
||||
"""Return a safe structural summary without logging prompts or image bytes."""
|
||||
stats = {
|
||||
"response_message_items": 0,
|
||||
"input_image_parts": 0,
|
||||
"image_url_parts": 0,
|
||||
"image_parts": 0,
|
||||
"data_url_images": 0,
|
||||
"remote_image_urls": 0,
|
||||
"literal_image_placeholders": 0,
|
||||
}
|
||||
|
||||
def walk(value: object, key: str = "") -> None:
|
||||
if isinstance(value, str):
|
||||
text = value.strip()
|
||||
lower = text.lower()
|
||||
if "<image>" in lower:
|
||||
stats["literal_image_placeholders"] += 1
|
||||
if lower.startswith("data:image/"):
|
||||
stats["data_url_images"] += 1
|
||||
elif key in {"image_url", "url"} and lower.startswith(("http://", "https://")):
|
||||
stats["remote_image_urls"] += 1
|
||||
return
|
||||
if isinstance(value, list):
|
||||
for item in value:
|
||||
walk(item, key)
|
||||
return
|
||||
if not isinstance(value, dict):
|
||||
return
|
||||
item_type = str(value.get("type") or "").strip()
|
||||
if item_type == "message":
|
||||
stats["response_message_items"] += 1
|
||||
elif item_type == "input_image":
|
||||
stats["input_image_parts"] += 1
|
||||
elif item_type == "image_url":
|
||||
stats["image_url_parts"] += 1
|
||||
elif item_type == "image":
|
||||
stats["image_parts"] += 1
|
||||
for child_key, child in value.items():
|
||||
walk(child, str(child_key))
|
||||
|
||||
for value in values:
|
||||
walk(value)
|
||||
return {key: value for key, value in stats.items() if value}
|
||||
|
||||
|
||||
def _sanitize_for_review(text: str) -> tuple[str, dict[str, int]]:
|
||||
"""Strip base64 data URIs and truncate to the review-service context limit.
|
||||
|
||||
Returns (sanitized_text, stats) where stats carries base64_blocks_stripped
|
||||
and truncated_chars so callers can emit structured logs.
|
||||
"""
|
||||
sanitized, base64_blocks_stripped = _BASE64_DATA_URI.subn("[image]", text)
|
||||
truncated_chars = 0
|
||||
if len(sanitized) > _MAX_REVIEW_TEXT_LEN:
|
||||
# Reserve marker space so the result stays within the cap.
|
||||
half = (_MAX_REVIEW_TEXT_LEN - len(_TRUNCATION_MARKER)) // 2
|
||||
truncated_chars = len(sanitized) - 2 * half
|
||||
sanitized = sanitized[:half] + _TRUNCATION_MARKER + sanitized[-half:]
|
||||
stats = {
|
||||
"base64_blocks_stripped": base64_blocks_stripped,
|
||||
"truncated_chars": truncated_chars,
|
||||
}
|
||||
return sanitized, stats
|
||||
|
||||
|
||||
def _extract_review_decision(data: object) -> str | None:
|
||||
"""Defensively pull the decision text out of the review service response.
|
||||
|
||||
Returns None when the response shape doesn't match the OpenAI chat-completion
|
||||
contract (e.g. {"error": ...} with no choices). The caller treats None as
|
||||
"undecided" and applies the configured fail-open policy.
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return None
|
||||
choices = data.get("choices")
|
||||
if not isinstance(choices, list) or not choices:
|
||||
return None
|
||||
first = choices[0]
|
||||
if not isinstance(first, dict):
|
||||
return None
|
||||
message = first.get("message")
|
||||
if not isinstance(message, dict):
|
||||
return None
|
||||
content = message.get("content")
|
||||
if content is None:
|
||||
return None
|
||||
return str(content).strip().lower()
|
||||
|
||||
|
||||
def _is_allow_decision(decision: str) -> bool:
|
||||
return decision.startswith(("allow", "pass", "true", "yes", "通过", "允许", "安全"))
|
||||
|
||||
|
||||
def _is_reject_decision(decision: str) -> bool:
|
||||
return decision.startswith(("reject", "deny", "block", "false", "no", "拒绝", "不允许", "违规", "禁止"))
|
||||
|
||||
|
||||
def _resolve_fail_open(review: dict) -> bool:
|
||||
"""Resolve fail_open from review config. Defaults to True."""
|
||||
value = review.get("fail_open")
|
||||
if value is None:
|
||||
return True
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, str):
|
||||
return value.strip().lower() in {"1", "true", "yes", "on"}
|
||||
return bool(value)
|
||||
|
||||
|
||||
def check_request(text: str) -> None:
|
||||
text = str(text or "")
|
||||
if not text.strip():
|
||||
return
|
||||
# Local sensitive-word match runs on the raw text (cheap, no network).
|
||||
for word in config.sensitive_words:
|
||||
if word in text:
|
||||
raise HTTPException(status_code=400, detail={"error": "检测到敏感词，拒绝本次任务"})
|
||||
review = config.ai_review
|
||||
if not review.get("enabled"):
|
||||
return
|
||||
base_url = str(review.get("base_url") or "").strip().rstrip("/")
|
||||
api_key = str(review.get("api_key") or "").strip()
|
||||
model = str(review.get("model") or "").strip()
|
||||
if not base_url or not api_key or not model:
|
||||
raise HTTPException(status_code=400, detail={"error": "ai review config is incomplete"})
|
||||
|
||||
fail_open = _resolve_fail_open(review)
|
||||
|
||||
review_text, sanitize_stats = _sanitize_for_review(text)
|
||||
if sanitize_stats["base64_blocks_stripped"] or sanitize_stats["truncated_chars"]:
|
||||
logger.info({
|
||||
"event": "ai_review_text_sanitized",
|
||||
"original_text_len": len(text),
|
||||
"review_text_len": len(review_text),
|
||||
**sanitize_stats,
|
||||
})
|
||||
prompt = str(review.get("prompt") or DEFAULT_REVIEW_PROMPT).strip()
|
||||
content = f"{prompt}\n\n用户请求:\n{review_text}\n\n只回答 ALLOW 或 REJECT。"
|
||||
|
||||
# fail_open=True (default): on upstream failure or ambiguous reply, let the
|
||||
# request through. The review is a soft safety net; one missed review is
|
||||
# preferable to a 5xx storm when the review service is flaky. Set
|
||||
# config.ai_review.fail_open=false for strict-compliance deployments.
|
||||
def _on_failure(event_payload: dict) -> None:
|
||||
logger.warning(event_payload)
|
||||
if not fail_open:
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail={"error": "AI 审核服务暂时不可用，请稍后重试"},
|
||||
)
|
||||
|
||||
try:
|
||||
response = requests.post(
|
||||
f"{base_url}/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"},
|
||||
json={"model": model, "messages": [{"role": "user", "content": content}], "temperature": 0},
|
||||
timeout=60,
|
||||
**proxy_settings.build_session_kwargs(),
|
||||
)
|
||||
except Exception as exc:
|
||||
_on_failure({
|
||||
"event": "ai_review_request_failed",
|
||||
"error": str(exc),
|
||||
"error_type": exc.__class__.__name__,
|
||||
"review_text_len": len(review_text),
|
||||
"original_text_len": len(text),
|
||||
})
|
||||
return
|
||||
|
||||
try:
|
||||
data = response.json()
|
||||
except Exception as exc:
|
||||
_on_failure({
|
||||
"event": "ai_review_response_not_json",
|
||||
"status_code": response.status_code,
|
||||
"body_preview": str(response.text or "")[:200],
|
||||
"error": str(exc),
|
||||
})
|
||||
return
|
||||
|
||||
decision = _extract_review_decision(data)
|
||||
if decision is None:
|
||||
_on_failure({
|
||||
"event": "ai_review_malformed_response",
|
||||
"status_code": response.status_code,
|
||||
"body_preview": str(data)[:300],
|
||||
"review_text_len": len(review_text),
|
||||
"original_text_len": len(text),
|
||||
})
|
||||
return
|
||||
|
||||
if _is_allow_decision(decision):
|
||||
return
|
||||
if _is_reject_decision(decision):
|
||||
raise HTTPException(status_code=400, detail={"error": "AI 审核未通过，拒绝本次任务"})
|
||||
# Ambiguous decisions (e.g. "MAYBE", empty content) fall back to fail-open policy.
|
||||
_on_failure({
|
||||
"event": "ai_review_ambiguous_decision",
|
||||
"decision": decision[:100],
|
||||
"review_text_len": len(review_text),
|
||||
})
|
||||
return
|
||||
@@ -0,0 +1,315 @@
|
||||
"""CLIProxyAPI integration for browsing remote auth files and importing selected tokens."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
|
||||
from curl_cffi.requests import Session
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.config import DATA_DIR
|
||||
from services.proxy_service import proxy_settings
|
||||
|
||||
|
||||
CPA_CONFIG_FILE = DATA_DIR / "cpa_config.json"
|
||||
|
||||
|
||||
def _new_id() -> str:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _normalize_import_job(raw: object, *, fail_unfinished: bool) -> dict | None:
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
status = str(raw.get("status") or "failed").strip() or "failed"
|
||||
if fail_unfinished and status in {"pending", "running"}:
|
||||
status = "failed"
|
||||
return {
|
||||
"job_id": str(raw.get("job_id") or uuid.uuid4().hex).strip(),
|
||||
"status": status,
|
||||
"created_at": str(raw.get("created_at") or _now_iso()).strip() or _now_iso(),
|
||||
"updated_at": str(raw.get("updated_at") or raw.get("created_at") or _now_iso()).strip() or _now_iso(),
|
||||
"total": int(raw.get("total") or 0),
|
||||
"completed": int(raw.get("completed") or 0),
|
||||
"added": int(raw.get("added") or 0),
|
||||
"skipped": int(raw.get("skipped") or 0),
|
||||
"refreshed": int(raw.get("refreshed") or 0),
|
||||
"failed": int(raw.get("failed") or 0),
|
||||
"errors": raw.get("errors") if isinstance(raw.get("errors"), list) else [],
|
||||
}
|
||||
|
||||
|
||||
def _normalize_pool(raw: dict) -> dict:
|
||||
return {
|
||||
"id": str(raw.get("id") or _new_id()).strip(),
|
||||
"name": str(raw.get("name") or "").strip(),
|
||||
"base_url": str(raw.get("base_url") or "").strip(),
|
||||
"secret_key": str(raw.get("secret_key") or "").strip(),
|
||||
"import_job": _normalize_import_job(raw.get("import_job"), fail_unfinished=True),
|
||||
}
|
||||
|
||||
|
||||
def _management_headers(secret_key: str) -> dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {secret_key}",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
|
||||
class CPAConfig:
|
||||
def __init__(self, store_file: Path):
|
||||
self._store_file = store_file
|
||||
self._lock = Lock()
|
||||
self._pools: list[dict] = self._load()
|
||||
|
||||
def _load(self) -> list[dict]:
|
||||
if not self._store_file.exists():
|
||||
return []
|
||||
try:
|
||||
raw = json.loads(self._store_file.read_text(encoding="utf-8"))
|
||||
if isinstance(raw, dict) and "base_url" in raw:
|
||||
pool = _normalize_pool(raw)
|
||||
return [pool] if pool["base_url"] else []
|
||||
if isinstance(raw, list):
|
||||
return [_normalize_pool(item) for item in raw if isinstance(item, dict)]
|
||||
except Exception:
|
||||
pass
|
||||
return []
|
||||
|
||||
def _save(self) -> None:
|
||||
self._store_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._store_file.write_text(json.dumps(self._pools, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
def list_pools(self) -> list[dict]:
|
||||
with self._lock:
|
||||
return [dict(pool) for pool in self._pools]
|
||||
|
||||
def get_pool(self, pool_id: str) -> dict | None:
|
||||
with self._lock:
|
||||
for pool in self._pools:
|
||||
if pool["id"] == pool_id:
|
||||
return dict(pool)
|
||||
return None
|
||||
|
||||
def add_pool(self, name: str, base_url: str, secret_key: str) -> dict:
|
||||
pool = _normalize_pool({"id": _new_id(), "name": name, "base_url": base_url, "secret_key": secret_key})
|
||||
with self._lock:
|
||||
self._pools.append(pool)
|
||||
self._save()
|
||||
return dict(pool)
|
||||
|
||||
def update_pool(self, pool_id: str, updates: dict) -> dict | None:
|
||||
with self._lock:
|
||||
for index, pool in enumerate(self._pools):
|
||||
if pool["id"] != pool_id:
|
||||
continue
|
||||
merged = {**pool, **{key: value for key, value in updates.items() if value is not None}, "id": pool_id}
|
||||
self._pools[index] = _normalize_pool(merged)
|
||||
self._save()
|
||||
return dict(self._pools[index])
|
||||
return None
|
||||
|
||||
def delete_pool(self, pool_id: str) -> bool:
|
||||
with self._lock:
|
||||
before = len(self._pools)
|
||||
self._pools = [pool for pool in self._pools if pool["id"] != pool_id]
|
||||
if len(self._pools) < before:
|
||||
self._save()
|
||||
return True
|
||||
return False
|
||||
|
||||
def set_import_job(self, pool_id: str, import_job: dict | None) -> dict | None:
|
||||
with self._lock:
|
||||
for index, pool in enumerate(self._pools):
|
||||
if pool["id"] != pool_id:
|
||||
continue
|
||||
next_pool = dict(pool)
|
||||
next_pool["import_job"] = _normalize_import_job(import_job, fail_unfinished=False)
|
||||
self._pools[index] = next_pool
|
||||
self._save()
|
||||
return dict(next_pool)
|
||||
return None
|
||||
|
||||
def get_import_job(self, pool_id: str) -> dict | None:
|
||||
with self._lock:
|
||||
for pool in self._pools:
|
||||
if pool["id"] == pool_id:
|
||||
job = pool.get("import_job")
|
||||
return dict(job) if isinstance(job, dict) else None
|
||||
return None
|
||||
|
||||
|
||||
def list_remote_files(pool: dict) -> list[dict]:
|
||||
base_url = str(pool.get("base_url") or "").strip()
|
||||
secret_key = str(pool.get("secret_key") or "").strip()
|
||||
if not base_url or not secret_key:
|
||||
return []
|
||||
|
||||
url = f"{base_url.rstrip('/')}/v0/management/auth-files"
|
||||
session = Session(**proxy_settings.build_session_kwargs(verify=True))
|
||||
try:
|
||||
response = session.get(url, headers=_management_headers(secret_key), timeout=30)
|
||||
if not response.ok:
|
||||
raise RuntimeError(f"remote list failed: HTTP {response.status_code}")
|
||||
payload = response.json()
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
files = payload.get("files") if isinstance(payload, dict) else None
|
||||
if not isinstance(files, list):
|
||||
raise RuntimeError("remote list payload is invalid")
|
||||
|
||||
items: list[dict] = []
|
||||
for item in files:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = str(item.get("name") or "").strip()
|
||||
email = str(item.get("email") or item.get("account") or "").strip()
|
||||
if not name:
|
||||
continue
|
||||
items.append({"name": name, "email": email})
|
||||
return items
|
||||
|
||||
|
||||
def fetch_remote_access_token(pool: dict, file_name: str) -> tuple[str | None, str | None]:
|
||||
base_url = str(pool.get("base_url") or "").strip()
|
||||
secret_key = str(pool.get("secret_key") or "").strip()
|
||||
file_name = str(file_name or "").strip()
|
||||
if not base_url or not secret_key or not file_name:
|
||||
return None, "invalid request"
|
||||
|
||||
url = f"{base_url.rstrip('/')}/v0/management/auth-files/download"
|
||||
session = Session(**proxy_settings.build_session_kwargs(verify=True))
|
||||
try:
|
||||
response = session.get(url, headers=_management_headers(secret_key), params={"name": file_name}, timeout=30)
|
||||
if not response.ok:
|
||||
return None, f"HTTP {response.status_code}"
|
||||
payload = response.json()
|
||||
except Exception as exc:
|
||||
return None, str(exc)
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
if not isinstance(payload, dict):
|
||||
return None, "invalid payload"
|
||||
|
||||
access_token = str(payload.get("access_token") or "").strip()
|
||||
if not access_token:
|
||||
return None, "missing access_token"
|
||||
return access_token, None
|
||||
|
||||
|
||||
class CPAImportService:
|
||||
def __init__(self, cpa_config: CPAConfig):
|
||||
self._config = cpa_config
|
||||
|
||||
def start_import(self, pool: dict, selected_files: list[str]) -> dict:
|
||||
names = [str(name or "").strip() for name in selected_files if str(name or "").strip()]
|
||||
if not names:
|
||||
raise ValueError("selected files is required")
|
||||
|
||||
pool_id = str(pool.get("id") or "").strip()
|
||||
job = {
|
||||
"job_id": uuid.uuid4().hex,
|
||||
"status": "pending",
|
||||
"created_at": _now_iso(),
|
||||
"updated_at": _now_iso(),
|
||||
"total": len(names),
|
||||
"completed": 0,
|
||||
"added": 0,
|
||||
"skipped": 0,
|
||||
"refreshed": 0,
|
||||
"failed": 0,
|
||||
"errors": [],
|
||||
}
|
||||
saved_pool = self._config.set_import_job(pool_id, job)
|
||||
if saved_pool is None:
|
||||
raise ValueError("pool not found")
|
||||
|
||||
thread = threading.Thread(
|
||||
target=self._run_import,
|
||||
args=(pool_id, pool, names),
|
||||
name=f"cpa-import-{pool_id}",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
return dict(saved_pool.get("import_job") or job)
|
||||
|
||||
def _update_job(self, pool_id: str, **updates) -> dict | None:
|
||||
current = self._config.get_import_job(pool_id)
|
||||
if current is None:
|
||||
return None
|
||||
next_job = {**current, **updates, "updated_at": _now_iso()}
|
||||
pool = self._config.set_import_job(pool_id, next_job)
|
||||
if pool is None:
|
||||
return None
|
||||
job = pool.get("import_job")
|
||||
return dict(job) if isinstance(job, dict) else None
|
||||
|
||||
def _append_error(self, pool_id: str, file_name: str, message: str) -> None:
|
||||
current = self._config.get_import_job(pool_id)
|
||||
if current is None:
|
||||
return
|
||||
errors = list(current.get("errors") or [])
|
||||
errors.append({"name": file_name, "error": message})
|
||||
self._update_job(pool_id, errors=errors, failed=len(errors))
|
||||
|
||||
def _run_import(self, pool_id: str, pool: dict, names: list[str]) -> None:
|
||||
self._update_job(pool_id, status="running")
|
||||
|
||||
tokens: list[str] = []
|
||||
max_workers = min(16, max(1, len(names)))
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
future_map = {executor.submit(fetch_remote_access_token, pool, name): name for name in names}
|
||||
for future in as_completed(future_map):
|
||||
file_name = future_map[future]
|
||||
try:
|
||||
token, error = future.result()
|
||||
except Exception as exc:
|
||||
token, error = None, str(exc)
|
||||
|
||||
if token:
|
||||
tokens.append(token)
|
||||
else:
|
||||
self._append_error(pool_id, file_name, error or "unknown error")
|
||||
|
||||
current = self._config.get_import_job(pool_id) or {}
|
||||
failed = len(current.get("errors") or [])
|
||||
self._update_job(pool_id, completed=int(current.get("completed") or 0) + 1, failed=failed)
|
||||
|
||||
if not tokens:
|
||||
current = self._config.get_import_job(pool_id) or {}
|
||||
self._update_job(
|
||||
pool_id,
|
||||
status="failed",
|
||||
completed=int(current.get("total") or 0),
|
||||
failed=len(current.get("errors") or []),
|
||||
)
|
||||
return
|
||||
|
||||
add_result = account_service.add_accounts(tokens, source_type="codex")
|
||||
refresh_result = account_service.refresh_accounts(tokens)
|
||||
current = self._config.get_import_job(pool_id) or {}
|
||||
self._update_job(
|
||||
pool_id,
|
||||
status="completed",
|
||||
completed=len(names),
|
||||
added=int(add_result.get("added") or 0),
|
||||
skipped=int(add_result.get("skipped") or 0),
|
||||
refreshed=int(refresh_result.get("refreshed") or 0),
|
||||
failed=len(current.get("errors") or []),
|
||||
)
|
||||
|
||||
|
||||
cpa_config = CPAConfig(CPA_CONFIG_FILE)
|
||||
cpa_import_service = CPAImportService(cpa_config)
|
||||
@@ -0,0 +1,260 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from urllib.parse import quote
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.config import DATA_DIR
|
||||
from services.content_filter import request_text
|
||||
from services.log_service import LOG_TYPE_CALL, log_service
|
||||
from services.openai_backend_api import EDITABLE_FILE_MODEL, OpenAIBackendAPI
|
||||
from utils.helper import new_uuid
|
||||
|
||||
TASK_STATUS_QUEUED = "queued"
|
||||
TASK_STATUS_RUNNING = "running"
|
||||
TASK_STATUS_SUCCESS = "success"
|
||||
TASK_STATUS_ERROR = "error"
|
||||
UNFINISHED_STATUSES = {TASK_STATUS_QUEUED, TASK_STATUS_RUNNING}
|
||||
EDITABLE_FILE_PLAN_TYPES = ("Plus", "Team", "Pro", "Enterprise")
|
||||
EDITABLE_FILE_ROOT = DATA_DIR / "files"
|
||||
EDITABLE_FILE_TASKS_PATH = DATA_DIR / "editable_file_tasks.json"
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
def _clean(value: object, default: str = "") -> str:
|
||||
return str(value or default).strip()
|
||||
|
||||
|
||||
def _owner_id(identity: dict[str, object]) -> str:
|
||||
return _clean(identity.get("id")) or "anonymous"
|
||||
|
||||
|
||||
def _task_key(owner_id: str, task_id: str) -> str:
|
||||
return f"{owner_id}:{task_id}"
|
||||
|
||||
|
||||
def _elapsed_seconds(task: dict[str, Any]) -> int:
|
||||
start = float(task.get("started_ts") or task.get("created_ts") or 0)
|
||||
end = float(task.get("ended_ts") or time.time())
|
||||
return max(0, int(end - start)) if start else 0
|
||||
|
||||
|
||||
def _file_url(path: Path, base_url: str) -> str:
|
||||
rel = path.resolve().relative_to(EDITABLE_FILE_ROOT.resolve()).as_posix()
|
||||
prefix = str(base_url or "").strip().rstrip("/")
|
||||
return f"{prefix}/files/{quote(rel, safe='/')}" if prefix else f"/files/{quote(rel, safe='/')}"
|
||||
|
||||
|
||||
def _editable_access_token() -> str:
|
||||
accounts = [
|
||||
item for item in account_service.list_accounts()
|
||||
if _clean(item.get("access_token"))
|
||||
and item.get("status") not in {"禁用", "异常"}
|
||||
and account_service._account_matches_any_plan_type(item, EDITABLE_FILE_PLAN_TYPES)
|
||||
]
|
||||
if not accounts:
|
||||
raise RuntimeError("no available plus/team/pro account")
|
||||
accounts.sort(key=lambda item: _clean(item.get("last_used_at")))
|
||||
token = _clean(accounts[0].get("access_token"))
|
||||
return account_service.refresh_access_token(token, event="editable_file_task") or token
|
||||
|
||||
|
||||
def _public_task(task: dict[str, Any]) -> dict[str, Any]:
|
||||
item = {
|
||||
"id": task.get("id"),
|
||||
"taskId": task.get("id"),
|
||||
"status": task.get("status"),
|
||||
"kind": task.get("kind"),
|
||||
"created_at": task.get("created_at"),
|
||||
"updated_at": task.get("updated_at"),
|
||||
"elapsed_seconds": _elapsed_seconds(task),
|
||||
}
|
||||
for key in ("result", "error"):
|
||||
if task.get(key):
|
||||
item[key] = task[key]
|
||||
return item
|
||||
|
||||
|
||||
class EditableFileTaskService:
|
||||
def __init__(self, path: Path = EDITABLE_FILE_TASKS_PATH) -> None:
|
||||
self.path = path
|
||||
self._lock = threading.RLock()
|
||||
self._tasks: dict[str, dict[str, Any]] = {}
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._lock:
|
||||
self._tasks = self._load_locked()
|
||||
if self._recover_unfinished_locked():
|
||||
self._save_locked()
|
||||
|
||||
def submit_ppt(self, identity: dict[str, object], *, client_task_id: str = "", prompt: str = "", base64_images: list[str] | None = None, base_url: str = "") -> dict[str, Any]:
|
||||
return self._submit(identity, client_task_id=client_task_id, kind="ppt", prompt=prompt, base64_images=base64_images or [], base_url=base_url)
|
||||
|
||||
def submit_psd(self, identity: dict[str, object], *, client_task_id: str = "", prompt: str = "", base64_images: list[str] | None = None, base_url: str = "") -> dict[str, Any]:
|
||||
return self._submit(identity, client_task_id=client_task_id, kind="psd", prompt=prompt, base64_images=base64_images or [], base_url=base_url)
|
||||
|
||||
def list_tasks(self, identity: dict[str, object], task_ids: list[str]) -> dict[str, Any]:
|
||||
owner = _owner_id(identity)
|
||||
requested = [_clean(item) for item in task_ids if _clean(item)]
|
||||
with self._lock:
|
||||
if requested:
|
||||
items = [task for task_id in requested if (task := self._tasks.get(_task_key(owner, task_id)))]
|
||||
return {"items": [_public_task(item) for item in items], "missing_ids": [task_id for task_id in requested if _task_key(owner, task_id) not in self._tasks]}
|
||||
items = [task for task in self._tasks.values() if task.get("owner_id") == owner]
|
||||
items.sort(key=lambda item: str(item.get("updated_at") or ""), reverse=True)
|
||||
return {"items": [_public_task(item) for item in items], "missing_ids": []}
|
||||
|
||||
def _submit(self, identity: dict[str, object], *, client_task_id: str, kind: str, prompt: str, base64_images: list[str], base_url: str) -> dict[str, Any]:
|
||||
task_id = _clean(client_task_id) or new_uuid()
|
||||
owner = _owner_id(identity)
|
||||
key = _task_key(owner, task_id)
|
||||
now = _now_iso()
|
||||
with self._lock:
|
||||
if key in self._tasks:
|
||||
return _public_task(self._tasks[key])
|
||||
ts = time.time()
|
||||
self._tasks[key] = {"id": task_id, "owner_id": owner, "status": TASK_STATUS_QUEUED, "kind": kind, "model": EDITABLE_FILE_MODEL, "created_at": now, "updated_at": now, "created_ts": ts, "updated_ts": ts}
|
||||
task = dict(self._tasks[key])
|
||||
self._save_locked()
|
||||
threading.Thread(target=self._run_task, args=(key, kind, prompt, base64_images, dict(identity), base_url), name=f"{kind}-file-task-{task_id[:16]}", daemon=True).start()
|
||||
return _public_task(task)
|
||||
|
||||
def _run_task(self, key: str, kind: str, prompt: str, base64_images: list[str], identity: dict[str, object], base_url: str) -> None:
|
||||
started = time.time()
|
||||
token = ""
|
||||
account_email = ""
|
||||
self._update_task(key, status=TASK_STATUS_RUNNING, error="", started_ts=started)
|
||||
backend = None
|
||||
try:
|
||||
if kind == "psd" and not base64_images:
|
||||
raise ValueError("base64_images is empty")
|
||||
token = _editable_access_token()
|
||||
account = account_service.get_account(token) or {}
|
||||
account_email = _clean(account.get("email"))
|
||||
backend = OpenAIBackendAPI(token)
|
||||
output_dir = EDITABLE_FILE_ROOT / kind / key.rsplit(":", 1)[-1]
|
||||
result = backend.export_psd_zip(base64_images, prompt, output_dir) if kind == "psd" else backend.export_ppt_zip(base64_images, prompt, output_dir)
|
||||
account_service.mark_text_used(token)
|
||||
data = {"conversation_id": result.conversation_id, "primary_url": _file_url(result.primary_path, base_url), "zip_url": _file_url(result.zip_path, base_url)}
|
||||
self._update_task(key, status=TASK_STATUS_SUCCESS, result=data, account_email=account_email, error="", ended_ts=time.time())
|
||||
self._log_call(identity, kind, started, request_text(prompt), account_email=account_email, result=data)
|
||||
except Exception as exc:
|
||||
error = str(exc) or "editable file task failed"
|
||||
self._update_task(key, status=TASK_STATUS_ERROR, error=error, account_email=account_email, ended_ts=time.time())
|
||||
self._log_call(identity, kind, started, request_text(prompt), status="failed", error=error, account_email=account_email)
|
||||
finally:
|
||||
if backend is not None:
|
||||
backend.close()
|
||||
|
||||
def public_file_path(self, relative_path: str) -> Path:
|
||||
raw = str(relative_path or "").replace("\\", "/").lstrip("/")
|
||||
path = (EDITABLE_FILE_ROOT / raw).resolve()
|
||||
path.relative_to(EDITABLE_FILE_ROOT.resolve())
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(raw)
|
||||
return path
|
||||
|
||||
def _update_task(self, key: str, **updates: Any) -> None:
|
||||
with self._lock:
|
||||
task = self._tasks.get(key)
|
||||
if task is None:
|
||||
return
|
||||
task.update(updates)
|
||||
task["updated_at"] = _now_iso()
|
||||
task["updated_ts"] = time.time()
|
||||
self._save_locked()
|
||||
|
||||
def _load_locked(self) -> dict[str, dict[str, Any]]:
|
||||
if not self.path.exists():
|
||||
return {}
|
||||
try:
|
||||
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
tasks: dict[str, dict[str, Any]] = {}
|
||||
for item in (raw.get("tasks") if isinstance(raw, dict) else raw) or []:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
task_id = _clean(item.get("id"))
|
||||
owner = _clean(item.get("owner_id"))
|
||||
if not task_id or not owner:
|
||||
continue
|
||||
task = {
|
||||
"id": task_id,
|
||||
"owner_id": owner,
|
||||
"status": _clean(item.get("status"), TASK_STATUS_ERROR),
|
||||
"kind": "psd" if item.get("kind") == "psd" else "ppt",
|
||||
"created_at": _clean(item.get("created_at"), _now_iso()),
|
||||
"updated_at": _clean(item.get("updated_at"), _clean(item.get("created_at"), _now_iso())),
|
||||
"created_ts": float(item.get("created_ts") or 0),
|
||||
"updated_ts": float(item.get("updated_ts") or 0),
|
||||
}
|
||||
for field in ("result", "error", "started_ts", "ended_ts"):
|
||||
if item.get(field):
|
||||
task[field] = item[field]
|
||||
tasks[_task_key(owner, task_id)] = task
|
||||
return tasks
|
||||
|
||||
def _save_locked(self) -> None:
|
||||
items = sorted(self._tasks.values(), key=lambda item: str(item.get("updated_at") or ""), reverse=True)
|
||||
tmp_path = self.path.with_suffix(self.path.suffix + ".tmp")
|
||||
tmp_path.write_text(json.dumps({"tasks": items}, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
tmp_path.replace(self.path)
|
||||
|
||||
def _recover_unfinished_locked(self) -> bool:
|
||||
changed = False
|
||||
for task in self._tasks.values():
|
||||
if task.get("status") in UNFINISHED_STATUSES:
|
||||
task["status"] = TASK_STATUS_ERROR
|
||||
task["error"] = "服务已重启，未完成的任务已中断"
|
||||
task["ended_ts"] = time.time()
|
||||
task["updated_at"] = _now_iso()
|
||||
task["updated_ts"] = time.time()
|
||||
changed = True
|
||||
return changed
|
||||
|
||||
def _log_call(
|
||||
self,
|
||||
identity: dict[str, object],
|
||||
kind: str,
|
||||
started: float,
|
||||
request_preview: str,
|
||||
*,
|
||||
status: str = "success",
|
||||
error: str = "",
|
||||
account_email: str = "",
|
||||
result: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
detail = {
|
||||
"key_id": identity.get("id"),
|
||||
"key_name": identity.get("name"),
|
||||
"role": identity.get("role"),
|
||||
"endpoint": f"/v1/{kind}/generations",
|
||||
"model": EDITABLE_FILE_MODEL,
|
||||
"started_at": datetime.fromtimestamp(started).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"ended_at": _now_iso(),
|
||||
"duration_ms": int((time.time() - started) * 1000),
|
||||
"status": status,
|
||||
}
|
||||
if request_preview:
|
||||
detail["request_text"] = request_preview
|
||||
if account_email:
|
||||
detail["account_email"] = account_email
|
||||
if error:
|
||||
detail["error"] = error
|
||||
if result:
|
||||
detail["result"] = result
|
||||
try:
|
||||
log_service.add(LOG_TYPE_CALL, f"{kind.upper()}生成任务{'失败' if status == 'failed' else '完成'}", detail)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
editable_file_task_service = EditableFileTaskService()
|
||||
@@ -0,0 +1,376 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import shutil
|
||||
import threading
|
||||
import time
|
||||
import zipfile
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import HTTPException
|
||||
from fastapi.responses import FileResponse, Response
|
||||
from PIL import Image, ImageOps
|
||||
|
||||
from services.config import config
|
||||
from services.image_storage_service import image_storage_service
|
||||
from services.image_tags_service import load_tags, remove_tags
|
||||
from utils.log import logger
|
||||
|
||||
THUMBNAIL_SIZE = (320, 320)
|
||||
|
||||
|
||||
def _cleanup_empty_dirs(root: Path) -> None:
|
||||
for path in sorted((p for p in root.rglob("*") if p.is_dir()), key=lambda p: len(p.parts), reverse=True):
|
||||
try:
|
||||
path.rmdir()
|
||||
except OSError:
|
||||
pass
|
||||
|
||||
|
||||
def _safe_relative_path(path: str) -> str:
|
||||
value = str(path or "").strip().replace("\\", "/").lstrip("/")
|
||||
if not value:
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
parts = Path(value).parts
|
||||
if any(part in {"", ".", ".."} for part in parts):
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
return Path(*parts).as_posix()
|
||||
|
||||
|
||||
def _safe_image_path(relative_path: str) -> Path:
|
||||
rel = _safe_relative_path(relative_path)
|
||||
root = config.images_dir.resolve()
|
||||
path = (root / rel).resolve()
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=404, detail="image not found") from exc
|
||||
if not path.is_file():
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
return path
|
||||
|
||||
|
||||
def get_image_response(relative_path: str) -> FileResponse | Response:
|
||||
headers = {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "*",
|
||||
}
|
||||
if image_storage_service.has_local(relative_path):
|
||||
return FileResponse(_safe_image_path(relative_path), headers=headers)
|
||||
return Response(content=image_storage_service.get_bytes(relative_path), media_type="image/png", headers=headers)
|
||||
|
||||
|
||||
def _thumbnail_path(relative_path: str) -> Path:
|
||||
rel = _safe_relative_path(relative_path)
|
||||
return config.image_thumbnails_dir / f"{rel}.png"
|
||||
|
||||
|
||||
def thumbnail_url(base_url: str, relative_path: str) -> str:
|
||||
return f"{base_url.rstrip('/')}/image-thumbnails/{_safe_relative_path(relative_path)}"
|
||||
|
||||
|
||||
def _image_dimensions(path: Path) -> tuple[int, int] | None:
|
||||
try:
|
||||
with Image.open(path) as image:
|
||||
return image.size
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def ensure_thumbnail(relative_path: str) -> Path:
|
||||
target = _thumbnail_path(relative_path)
|
||||
source_mtime = 0.0
|
||||
source: Path | None = None
|
||||
if image_storage_service.has_local(relative_path):
|
||||
source = _safe_image_path(relative_path)
|
||||
source_mtime = source.stat().st_mtime
|
||||
if target.exists() and (not source_mtime or target.stat().st_mtime >= source_mtime):
|
||||
return target
|
||||
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
image_source = source if source is not None else io.BytesIO(image_storage_service.get_bytes(relative_path))
|
||||
with Image.open(image_source) as image:
|
||||
image = ImageOps.exif_transpose(image)
|
||||
if image.mode not in {"RGB", "RGBA"}:
|
||||
image = image.convert("RGBA" if "A" in image.getbands() else "RGB")
|
||||
image.thumbnail(THUMBNAIL_SIZE, Image.Resampling.LANCZOS)
|
||||
image.save(target, format="PNG", optimize=True)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise HTTPException(status_code=422, detail="failed to create thumbnail") from exc
|
||||
return target
|
||||
|
||||
|
||||
def get_thumbnail_response(relative_path: str) -> FileResponse:
|
||||
headers = {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "*",
|
||||
}
|
||||
return FileResponse(ensure_thumbnail(relative_path), headers=headers)
|
||||
|
||||
|
||||
def get_image_download_response(relative_path: str) -> FileResponse:
|
||||
cors_headers = {
|
||||
"Access-Control-Allow-Origin": "*",
|
||||
"Access-Control-Allow-Methods": "GET, OPTIONS",
|
||||
"Access-Control-Allow-Headers": "*",
|
||||
}
|
||||
if image_storage_service.has_local(relative_path):
|
||||
path = _safe_image_path(relative_path)
|
||||
headers = {**cors_headers, "Content-Disposition": f'attachment; filename="{path.name}"'}
|
||||
return FileResponse(path, filename=path.name, headers=headers)
|
||||
rel = _safe_relative_path(relative_path)
|
||||
headers = {
|
||||
**cors_headers,
|
||||
"Content-Disposition": f'attachment; filename="{Path(rel).name}"',
|
||||
}
|
||||
return Response(
|
||||
content=image_storage_service.get_bytes(rel),
|
||||
media_type="image/png",
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
def cleanup_image_thumbnails() -> int:
|
||||
thumbnails_root = config.image_thumbnails_dir
|
||||
removed = 0
|
||||
for path in thumbnails_root.rglob("*"):
|
||||
if not path.is_file():
|
||||
continue
|
||||
rel = path.relative_to(thumbnails_root).as_posix()
|
||||
if not rel.endswith(".png") or not image_storage_service.exists(rel[:-4]):
|
||||
path.unlink()
|
||||
removed += 1
|
||||
_cleanup_empty_dirs(thumbnails_root)
|
||||
return removed
|
||||
|
||||
def list_images(base_url: str, start_date: str = "", end_date: str = "") -> dict[str, object]:
|
||||
config.cleanup_old_images()
|
||||
cleanup_image_thumbnails()
|
||||
all_tags = load_tags()
|
||||
items = [
|
||||
{
|
||||
**item,
|
||||
"url": str(item.get("url") or f"{base_url.rstrip('/')}/images/{item['path']}"),
|
||||
"thumbnail_url": thumbnail_url(base_url, str(item["path"])),
|
||||
"tags": all_tags.get(str(item["path"]), []),
|
||||
}
|
||||
for item in image_storage_service.list_items(base_url, start_date, end_date)
|
||||
]
|
||||
groups: dict[str, list[dict[str, object]]] = {}
|
||||
for item in items:
|
||||
groups.setdefault(str(item["date"]), []).append(item)
|
||||
return {"items": items, "groups": [{"date": key, "items": value} for key, value in groups.items()]}
|
||||
|
||||
|
||||
def delete_images(paths: list[str] | None = None, start_date: str = "", end_date: str = "", all_matching: bool = False) -> dict[str, int]:
|
||||
root = config.images_dir.resolve()
|
||||
targets = [
|
||||
str(item["path"])
|
||||
for item in image_storage_service.list_items("", start_date=start_date, end_date=end_date)
|
||||
] if all_matching else (paths or [])
|
||||
removed = 0
|
||||
for item in targets:
|
||||
path = (root / item).resolve()
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError:
|
||||
continue
|
||||
if image_storage_service.delete(item):
|
||||
removed += 1
|
||||
for thumbnail in (_thumbnail_path(item), config.image_thumbnails_dir / _safe_relative_path(item)):
|
||||
if thumbnail.is_file():
|
||||
thumbnail.unlink()
|
||||
remove_tags(item)
|
||||
_cleanup_empty_dirs(root)
|
||||
_cleanup_empty_dirs(config.image_thumbnails_dir)
|
||||
return {"removed": removed}
|
||||
|
||||
|
||||
def download_images_zip(paths: list[str]) -> io.BytesIO:
|
||||
root = config.images_dir.resolve()
|
||||
buf = io.BytesIO()
|
||||
added = 0
|
||||
used_names: set[str] = set()
|
||||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for item in paths:
|
||||
rel = _safe_relative_path(item)
|
||||
path = (root / rel).resolve()
|
||||
payload: bytes | None = None
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError:
|
||||
continue
|
||||
if path.is_file():
|
||||
payload = path.read_bytes()
|
||||
else:
|
||||
try:
|
||||
payload = image_storage_service.get_bytes(rel)
|
||||
except Exception:
|
||||
continue
|
||||
name = path.name
|
||||
if name in used_names:
|
||||
stem = path.stem
|
||||
suffix = path.suffix
|
||||
counter = 2
|
||||
while f"{stem}_{counter}{suffix}" in used_names:
|
||||
counter += 1
|
||||
name = f"{stem}_{counter}{suffix}"
|
||||
used_names.add(name)
|
||||
zf.writestr(name, payload)
|
||||
added += 1
|
||||
if added == 0:
|
||||
raise HTTPException(status_code=404, detail="no images found")
|
||||
buf.seek(0)
|
||||
return buf
|
||||
def storage_stats() -> dict:
|
||||
import shutil
|
||||
usage = shutil.disk_usage(config.images_dir)
|
||||
total_mb = usage.total // (1024 * 1024)
|
||||
used_mb = usage.used // (1024 * 1024)
|
||||
free_mb = usage.free // (1024 * 1024)
|
||||
|
||||
image_count = 0
|
||||
image_size = 0
|
||||
for p in config.images_dir.rglob("*"):
|
||||
if p.is_file():
|
||||
image_count += 1
|
||||
image_size += p.stat().st_size
|
||||
|
||||
return {
|
||||
"disk_total_mb": total_mb,
|
||||
"disk_used_mb": used_mb,
|
||||
"disk_free_mb": free_mb,
|
||||
"image_count": image_count,
|
||||
"image_size_mb": image_size // (1024 * 1024),
|
||||
"image_size_bytes": image_size,
|
||||
}
|
||||
|
||||
|
||||
def compress_images(quality: int = 60) -> dict:
|
||||
"""重新压缩所有图片，返回节省的空间"""
|
||||
saved = 0
|
||||
count = 0
|
||||
for p in sorted(config.images_dir.rglob("*.png")):
|
||||
if not p.is_file():
|
||||
continue
|
||||
try:
|
||||
orig = p.stat().st_size
|
||||
with Image.open(p) as img:
|
||||
img = ImageOps.exif_transpose(img)
|
||||
img.save(str(p) + ".tmp", format="PNG", optimize=True)
|
||||
new_size = Path(str(p) + ".tmp").stat().st_size
|
||||
if new_size < orig:
|
||||
Path(str(p) + ".tmp").replace(p)
|
||||
saved += orig - new_size
|
||||
count += 1
|
||||
else:
|
||||
Path(str(p) + ".tmp").unlink()
|
||||
except Exception:
|
||||
pass
|
||||
return {"compressed": count, "saved_bytes": saved, "saved_mb": saved // (1024 * 1024)}
|
||||
|
||||
|
||||
def delete_to_target(target_free_mb: int, dry_run: bool = False) -> dict:
|
||||
"""删除最旧的图片直到剩余空间达到 target_free_mb"""
|
||||
import shutil
|
||||
usage = shutil.disk_usage(config.images_dir)
|
||||
current_free = usage.free // (1024 * 1024)
|
||||
if current_free >= target_free_mb and not dry_run:
|
||||
return {"removed": 0, "current_free_mb": current_free, "target_free_mb": target_free_mb, "done": True}
|
||||
|
||||
files = sorted(
|
||||
(p for p in config.images_dir.rglob("*.png") if p.is_file()),
|
||||
key=lambda p: p.stat().st_mtime,
|
||||
)
|
||||
removed = 0
|
||||
freed = 0
|
||||
for p in files:
|
||||
if current_free + freed // (1024 * 1024) >= target_free_mb:
|
||||
break
|
||||
size = p.stat().st_size
|
||||
if not dry_run:
|
||||
rel = p.relative_to(config.images_dir).as_posix()
|
||||
for tp in (_thumbnail_path(rel), config.image_thumbnails_dir / _safe_relative_path(rel)):
|
||||
if tp.is_file():
|
||||
tp.unlink()
|
||||
remove_tags(rel)
|
||||
p.unlink()
|
||||
freed += size
|
||||
removed += 1
|
||||
|
||||
if not dry_run:
|
||||
_cleanup_empty_dirs(config.images_dir)
|
||||
_cleanup_empty_dirs(config.image_thumbnails_dir)
|
||||
|
||||
return {
|
||||
"removed": removed,
|
||||
"freed_mb": freed // (1024 * 1024),
|
||||
"target_free_mb": target_free_mb,
|
||||
"current_free_mb": current_free + (freed // (1024 * 1024)),
|
||||
"done": (current_free + freed // (1024 * 1024)) >= target_free_mb,
|
||||
"dry_run": dry_run,
|
||||
}
|
||||
|
||||
|
||||
def download_images_zip(paths: list[str]) -> io.BytesIO:
|
||||
root = config.images_dir.resolve()
|
||||
buf = io.BytesIO()
|
||||
added = 0
|
||||
used_names: set[str] = set()
|
||||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||||
for item in paths:
|
||||
rel = _safe_relative_path(item)
|
||||
path = (root / rel).resolve()
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError:
|
||||
continue
|
||||
if not path.is_file():
|
||||
continue
|
||||
name = path.name
|
||||
if name in used_names:
|
||||
stem = path.stem
|
||||
suffix = path.suffix
|
||||
counter = 2
|
||||
while f"{stem}_{counter}{suffix}" in used_names:
|
||||
counter += 1
|
||||
name = f"{stem}_{counter}{suffix}"
|
||||
used_names.add(name)
|
||||
zf.write(path, name)
|
||||
added += 1
|
||||
if added == 0:
|
||||
raise HTTPException(status_code=404, detail="no images found")
|
||||
buf.seek(0)
|
||||
return buf
|
||||
|
||||
|
||||
def _auto_cleanup_worker(stop_event: threading.Event) -> None:
|
||||
"""后台线程：每30分钟检查存储，空间低于阈值自动清理最旧图片"""
|
||||
import shutil
|
||||
min_free_mb = getattr(config, "image_min_free_mb", None)
|
||||
if min_free_mb is None:
|
||||
min_free_mb = 500
|
||||
|
||||
while not stop_event.wait(1800): # 每30分钟
|
||||
try:
|
||||
config.cleanup_old_images()
|
||||
cleanup_image_thumbnails()
|
||||
usage = shutil.disk_usage(config.images_dir)
|
||||
free_mb = usage.free // (1024 * 1024)
|
||||
if free_mb < min_free_mb:
|
||||
logger.info({"event": "image_auto_cleanup", "free_mb": free_mb, "min_free_mb": min_free_mb})
|
||||
result = delete_to_target(min_free_mb)
|
||||
logger.info({"event": "image_auto_cleanup_done", **result})
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def start_image_cleanup_scheduler(stop_event: threading.Event) -> threading.Thread:
|
||||
t = threading.Thread(target=_auto_cleanup_worker, args=(stop_event,), daemon=True, name="image-cleanup")
|
||||
t.start()
|
||||
return t
|
||||
@@ -0,0 +1,405 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import time
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
from curl_cffi import requests
|
||||
from fastapi import HTTPException
|
||||
from PIL import Image
|
||||
|
||||
from services.config import DATA_DIR, config
|
||||
|
||||
IMAGE_INDEX_FILE = DATA_DIR / "image_index.json"
|
||||
IMAGE_INDEX_LOCK = Lock()
|
||||
IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"}
|
||||
|
||||
|
||||
class ImageStorageError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StoredImage:
|
||||
rel: str
|
||||
url: str
|
||||
storage: str
|
||||
size: int
|
||||
|
||||
|
||||
def _clean(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
def _safe_relative_path(path: str) -> str:
|
||||
value = str(path or "").strip().replace("\\", "/").lstrip("/")
|
||||
if not value:
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
parts = Path(value).parts
|
||||
if any(part in {"", ".", ".."} for part in parts):
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
return Path(*parts).as_posix()
|
||||
|
||||
|
||||
def _image_dimensions(payload: bytes) -> tuple[int, int] | None:
|
||||
try:
|
||||
with Image.open(io.BytesIO(payload)) as image:
|
||||
return image.size
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _is_image_rel(path: str) -> bool:
|
||||
try:
|
||||
safe_rel = _safe_relative_path(path)
|
||||
except HTTPException:
|
||||
return False
|
||||
return Path(safe_rel).suffix.lower() in IMAGE_EXTENSIONS
|
||||
|
||||
|
||||
def _local_image_path(relative_path: str) -> Path:
|
||||
rel = _safe_relative_path(relative_path)
|
||||
root = config.images_dir.resolve()
|
||||
path = (root / rel).resolve()
|
||||
try:
|
||||
path.relative_to(root)
|
||||
except ValueError as exc:
|
||||
raise HTTPException(status_code=404, detail="image not found") from exc
|
||||
return path
|
||||
|
||||
|
||||
def _read_json_object(path: Path) -> dict[str, object]:
|
||||
if not path.exists():
|
||||
return {}
|
||||
try:
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def _write_json_object(path: Path, data: dict[str, object]) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
tmp_path = path.with_suffix(path.suffix + ".tmp")
|
||||
tmp_path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
tmp_path.replace(path)
|
||||
|
||||
|
||||
class WebDAVClient:
|
||||
def __init__(self, settings: dict[str, object]):
|
||||
self.url = _clean(settings.get("webdav_url")).rstrip("/")
|
||||
self.username = _clean(settings.get("webdav_username"))
|
||||
self.password = _clean(settings.get("webdav_password"))
|
||||
self.root_path = _clean(settings.get("webdav_root_path")).strip("/")
|
||||
self.session = requests.Session()
|
||||
|
||||
def _auth_kwargs(self) -> dict[str, object]:
|
||||
return {"auth": (self.username, self.password)} if self.username or self.password else {}
|
||||
|
||||
def _request(self, method: str, url: str, **kwargs):
|
||||
response = self.session.request(method, url, timeout=30, **self._auth_kwargs(), **kwargs)
|
||||
if response.status_code >= 400 and not (method == "MKCOL" and response.status_code in {405}):
|
||||
raise ImageStorageError(f"WebDAV {method} failed: HTTP {response.status_code}")
|
||||
return response
|
||||
|
||||
def remote_url(self, rel: str = "") -> str:
|
||||
parts = [part for part in [self.root_path, _safe_relative_path(rel) if rel else ""] if part]
|
||||
encoded = "/".join(quote(part, safe="") for item in parts for part in item.split("/") if part)
|
||||
return f"{self.url}/{encoded}" if encoded else self.url
|
||||
|
||||
def ensure_dirs(self, rel: str) -> None:
|
||||
parts = [part for part in [self.root_path, Path(_safe_relative_path(rel)).parent.as_posix()] if part and part != "."]
|
||||
current = self.url
|
||||
for item in "/".join(parts).split("/"):
|
||||
if not item:
|
||||
continue
|
||||
current = f"{current}/{quote(item, safe='')}"
|
||||
response = self.session.request("MKCOL", current, timeout=30, **self._auth_kwargs())
|
||||
if response.status_code in {201, 405}:
|
||||
continue
|
||||
if response.status_code >= 400:
|
||||
raise ImageStorageError(f"WebDAV MKCOL failed: HTTP {response.status_code}")
|
||||
|
||||
def put(self, rel: str, payload: bytes, content_type: str = "image/png") -> str:
|
||||
self.ensure_dirs(rel)
|
||||
url = self.remote_url(rel)
|
||||
self._request("PUT", url, data=payload, headers={"Content-Type": content_type})
|
||||
return url
|
||||
|
||||
def get(self, rel: str) -> bytes:
|
||||
response = self._request("GET", self.remote_url(rel))
|
||||
return bytes(response.content)
|
||||
|
||||
def delete(self, rel: str) -> bool:
|
||||
response = self.session.request("DELETE", self.remote_url(rel), timeout=30, **self._auth_kwargs())
|
||||
if response.status_code in {200, 202, 204, 404}:
|
||||
return response.status_code != 404
|
||||
raise ImageStorageError(f"WebDAV DELETE failed: HTTP {response.status_code}")
|
||||
|
||||
def test(self) -> dict[str, object]:
|
||||
if not self.url:
|
||||
return {"ok": False, "status": 0, "error": "WebDAV URL is required"}
|
||||
if urlparse(self.url).scheme not in {"http", "https"}:
|
||||
return {"ok": False, "status": 0, "error": "invalid WebDAV URL"}
|
||||
test_rel = ".chatgpt2api_webdav_test.txt"
|
||||
try:
|
||||
self.put(test_rel, b"chatgpt2api webdav test\n", content_type="text/plain")
|
||||
self.delete(test_rel)
|
||||
return {"ok": True, "status": 200, "error": None}
|
||||
except ImageStorageError as exc:
|
||||
return {"ok": False, "status": 0, "error": str(exc)}
|
||||
except Exception as exc:
|
||||
return {"ok": False, "status": 0, "error": str(exc) or exc.__class__.__name__}
|
||||
finally:
|
||||
self.session.close()
|
||||
|
||||
|
||||
class ImageStorageService:
|
||||
def __init__(self, index_file: Path = IMAGE_INDEX_FILE):
|
||||
self.index_file = index_file
|
||||
self._index_lock = IMAGE_INDEX_LOCK
|
||||
|
||||
def settings(self) -> dict[str, object]:
|
||||
return config.get_image_storage_settings()
|
||||
|
||||
def mode(self) -> str:
|
||||
return _clean(self.settings().get("mode")) or "local"
|
||||
|
||||
def _load_index(self) -> dict[str, dict[str, object]]:
|
||||
raw = _read_json_object(self.index_file)
|
||||
items = raw.get("items")
|
||||
if not isinstance(items, dict):
|
||||
return {}
|
||||
return {str(key): value for key, value in items.items() if isinstance(value, dict)}
|
||||
|
||||
def _load_clean_index(self) -> dict[str, dict[str, object]]:
|
||||
items = self._load_index()
|
||||
return {rel: item for rel, item in items.items() if _is_image_rel(rel)}
|
||||
|
||||
def _save_index(self, items: dict[str, dict[str, object]]) -> None:
|
||||
_write_json_object(self.index_file, {"items": items})
|
||||
|
||||
def _public_url(self, rel: str, base_url: str | None = None) -> str:
|
||||
settings = self.settings()
|
||||
public_base_url = _clean(settings.get("public_base_url"))
|
||||
if public_base_url:
|
||||
return f"{public_base_url.rstrip('/')}/{_safe_relative_path(rel)}"
|
||||
return f"{(base_url or config.base_url).rstrip('/')}/images/{_safe_relative_path(rel)}"
|
||||
|
||||
def make_relative_path(self, image_data: bytes) -> str:
|
||||
file_hash = hashlib.md5(image_data).hexdigest()
|
||||
filename = f"{int(time.time())}_{file_hash}.png"
|
||||
relative_dir = Path(time.strftime("%Y"), time.strftime("%m"), time.strftime("%d"))
|
||||
return f"{relative_dir.as_posix()}/{filename}"
|
||||
|
||||
def save(self, image_data: bytes, base_url: str | None = None) -> StoredImage:
|
||||
config.cleanup_old_images()
|
||||
rel = self.make_relative_path(image_data)
|
||||
mode = self.mode()
|
||||
if mode not in {"local", "webdav", "both"}:
|
||||
mode = "local"
|
||||
stored_local = False
|
||||
stored_webdav = False
|
||||
remote_url = ""
|
||||
|
||||
if mode in {"local", "both"}:
|
||||
path = _local_image_path(rel)
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_bytes(image_data)
|
||||
stored_local = True
|
||||
|
||||
if mode in {"webdav", "both"}:
|
||||
remote_url = WebDAVClient(self.settings()).put(rel, image_data)
|
||||
stored_webdav = True
|
||||
|
||||
dimensions = _image_dimensions(image_data)
|
||||
item = {
|
||||
"rel": rel,
|
||||
"path": rel,
|
||||
"name": Path(rel).name,
|
||||
"date": "-".join(rel.split("/")[:3]),
|
||||
"size": len(image_data),
|
||||
"created_at": _now_iso(),
|
||||
"storage": "both" if stored_local and stored_webdav else ("webdav" if stored_webdav else "local"),
|
||||
"local": stored_local,
|
||||
"webdav": stored_webdav,
|
||||
"remote_url": remote_url,
|
||||
}
|
||||
if dimensions:
|
||||
item["width"], item["height"] = dimensions
|
||||
with self._index_lock:
|
||||
items = self._load_clean_index()
|
||||
items[rel] = item
|
||||
self._save_index(items)
|
||||
return StoredImage(rel=rel, url=self._public_url(rel, base_url), storage=str(item["storage"]), size=len(image_data))
|
||||
|
||||
def get_bytes(self, rel: str) -> bytes:
|
||||
safe_rel = _safe_relative_path(rel)
|
||||
if not _is_image_rel(safe_rel):
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
path = _local_image_path(safe_rel)
|
||||
if path.is_file():
|
||||
return path.read_bytes()
|
||||
item = self._load_clean_index().get(safe_rel, {})
|
||||
if item.get("webdav"):
|
||||
return WebDAVClient(self.settings()).get(safe_rel)
|
||||
raise HTTPException(status_code=404, detail="image not found")
|
||||
|
||||
def exists(self, rel: str) -> bool:
|
||||
safe_rel = _safe_relative_path(rel)
|
||||
if not _is_image_rel(safe_rel):
|
||||
return False
|
||||
if _local_image_path(safe_rel).is_file():
|
||||
return True
|
||||
item = self._load_clean_index().get(safe_rel, {})
|
||||
return bool(item.get("webdav"))
|
||||
|
||||
def has_local(self, rel: str) -> bool:
|
||||
safe_rel = _safe_relative_path(rel)
|
||||
return _is_image_rel(safe_rel) and _local_image_path(safe_rel).is_file()
|
||||
|
||||
def list_items(self, base_url: str, start_date: str = "", end_date: str = "") -> list[dict[str, object]]:
|
||||
with self._index_lock:
|
||||
indexed = self._load_clean_index()
|
||||
root = config.images_dir
|
||||
changed = False
|
||||
for path in root.rglob("*"):
|
||||
if not path.is_file() or not _is_image_rel(path.name):
|
||||
continue
|
||||
rel = path.relative_to(root).as_posix()
|
||||
if rel in indexed:
|
||||
continue
|
||||
dimensions = None
|
||||
try:
|
||||
dimensions = _image_dimensions(path.read_bytes())
|
||||
except Exception:
|
||||
dimensions = None
|
||||
indexed[rel] = {
|
||||
"rel": rel,
|
||||
"path": rel,
|
||||
"name": path.name,
|
||||
"date": "-".join(rel.split("/")[:3]) if len(rel.split("/")) >= 4 else datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d"),
|
||||
"size": path.stat().st_size,
|
||||
"created_at": datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"storage": "local",
|
||||
"local": True,
|
||||
"webdav": False,
|
||||
**({"width": dimensions[0], "height": dimensions[1]} if dimensions else {}),
|
||||
}
|
||||
changed = True
|
||||
|
||||
items: list[dict[str, object]] = []
|
||||
for rel, item in list(indexed.items()):
|
||||
if not _is_image_rel(rel):
|
||||
indexed.pop(rel, None)
|
||||
changed = True
|
||||
continue
|
||||
local = _local_image_path(rel).is_file()
|
||||
webdav = bool(item.get("webdav"))
|
||||
if not local and not webdav:
|
||||
indexed.pop(rel, None)
|
||||
changed = True
|
||||
continue
|
||||
storage = "both" if local and webdav else ("webdav" if webdav else "local")
|
||||
if item.get("local") != local or item.get("storage") != storage:
|
||||
item = {
|
||||
**item,
|
||||
"local": local,
|
||||
"storage": storage,
|
||||
}
|
||||
indexed[rel] = item
|
||||
changed = True
|
||||
day = str(item.get("date") or "")
|
||||
if start_date and day < start_date:
|
||||
continue
|
||||
if end_date and day > end_date:
|
||||
continue
|
||||
items.append({
|
||||
**item,
|
||||
"rel": rel,
|
||||
"path": rel,
|
||||
"url": self._public_url(rel, base_url),
|
||||
})
|
||||
if changed:
|
||||
self._save_index(indexed)
|
||||
items.sort(key=lambda item: str(item.get("created_at") or ""), reverse=True)
|
||||
return items
|
||||
|
||||
def delete(self, rel: str) -> bool:
|
||||
safe_rel = _safe_relative_path(rel)
|
||||
removed = False
|
||||
path = _local_image_path(safe_rel)
|
||||
if path.is_file():
|
||||
path.unlink()
|
||||
removed = True
|
||||
with self._index_lock:
|
||||
items = self._load_clean_index()
|
||||
item = items.get(safe_rel, {})
|
||||
if item.get("webdav"):
|
||||
try:
|
||||
removed = WebDAVClient(self.settings()).delete(safe_rel) or removed
|
||||
except ImageStorageError:
|
||||
if not removed:
|
||||
raise
|
||||
if safe_rel in items:
|
||||
items.pop(safe_rel, None)
|
||||
self._save_index(items)
|
||||
return removed
|
||||
|
||||
def sync_all(self) -> dict[str, int]:
|
||||
settings = self.settings()
|
||||
if self.mode() not in {"webdav", "both"}:
|
||||
raise ImageStorageError("WebDAV 图片存储未启用")
|
||||
uploaded = 0
|
||||
skipped = 0
|
||||
failed = 0
|
||||
with self._index_lock:
|
||||
items = self._load_clean_index()
|
||||
client = WebDAVClient(settings)
|
||||
for path in sorted(config.images_dir.rglob("*")):
|
||||
if not path.is_file() or not _is_image_rel(path.name):
|
||||
continue
|
||||
rel = path.relative_to(config.images_dir).as_posix()
|
||||
item = items.get(rel, {})
|
||||
if item.get("webdav"):
|
||||
skipped += 1
|
||||
continue
|
||||
try:
|
||||
payload = path.read_bytes()
|
||||
remote_url = client.put(rel, payload)
|
||||
dimensions = _image_dimensions(payload)
|
||||
items[rel] = {
|
||||
**item,
|
||||
"rel": rel,
|
||||
"path": rel,
|
||||
"name": path.name,
|
||||
"date": "-".join(rel.split("/")[:3]) if len(rel.split("/")) >= 4 else datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d"),
|
||||
"size": len(payload),
|
||||
"created_at": str(item.get("created_at") or datetime.fromtimestamp(path.stat().st_mtime).strftime("%Y-%m-%d %H:%M:%S")),
|
||||
"storage": "both",
|
||||
"local": True,
|
||||
"webdav": True,
|
||||
"remote_url": remote_url,
|
||||
**({"width": dimensions[0], "height": dimensions[1]} if dimensions else {}),
|
||||
}
|
||||
uploaded += 1
|
||||
except Exception:
|
||||
failed += 1
|
||||
self._save_index(items)
|
||||
return {"uploaded": uploaded, "skipped": skipped, "failed": failed}
|
||||
|
||||
def test_webdav(self) -> dict[str, object]:
|
||||
return WebDAVClient(self.settings()).test()
|
||||
|
||||
|
||||
image_storage_service = ImageStorageService()
|
||||
@@ -0,0 +1,76 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from services.config import DATA_DIR
|
||||
|
||||
TAGS_FILE = DATA_DIR / "image_tags.json"
|
||||
|
||||
|
||||
def _ensure_file() -> None:
|
||||
TAGS_FILE.parent.mkdir(parents=True, exist_ok=True)
|
||||
if not TAGS_FILE.exists():
|
||||
TAGS_FILE.write_text("{}", encoding="utf-8")
|
||||
|
||||
|
||||
def load_tags() -> dict[str, list[str]]:
|
||||
_ensure_file()
|
||||
try:
|
||||
data = json.loads(TAGS_FILE.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
return data if isinstance(data, dict) else {}
|
||||
|
||||
|
||||
def save_tags(data: dict[str, list[str]]) -> None:
|
||||
_ensure_file()
|
||||
TAGS_FILE.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
def get_tags(image_rel: str) -> list[str]:
|
||||
return load_tags().get(image_rel, [])
|
||||
|
||||
|
||||
def set_tags(image_rel: str, tags: list[str]) -> list[str]:
|
||||
data = load_tags()
|
||||
cleaned = list(dict.fromkeys(t.strip() for t in tags if t.strip()))
|
||||
if cleaned:
|
||||
data[image_rel] = cleaned
|
||||
else:
|
||||
data.pop(image_rel, None)
|
||||
save_tags(data)
|
||||
return cleaned
|
||||
|
||||
|
||||
def remove_tags(image_rel: str) -> None:
|
||||
data = load_tags()
|
||||
if data.pop(image_rel, None) is not None:
|
||||
save_tags(data)
|
||||
|
||||
|
||||
def delete_tag(tag: str) -> int:
|
||||
"""从所有图片中删除指定标签，返回受影响的图片数。"""
|
||||
data = load_tags()
|
||||
count = 0
|
||||
for rel in list(data):
|
||||
if tag in data[rel]:
|
||||
data[rel] = [t for t in data[rel] if t != tag]
|
||||
if not data[rel]:
|
||||
del data[rel]
|
||||
count += 1
|
||||
if count > 0:
|
||||
save_tags(data)
|
||||
return count
|
||||
|
||||
|
||||
def get_all_tags() -> list[str]:
|
||||
data = load_tags()
|
||||
seen: set[str] = set()
|
||||
result: list[str] = []
|
||||
for tags in data.values():
|
||||
for t in tags:
|
||||
if t not in seen:
|
||||
seen.add(t)
|
||||
result.append(t)
|
||||
return result
|
||||
@@ -0,0 +1,576 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from services.config import DATA_DIR, config
|
||||
from services.content_filter import request_text
|
||||
from services.log_service import LOG_TYPE_CALL, log_service
|
||||
from services.protocol import openai_v1_image_edit, openai_v1_image_generations
|
||||
from utils.log import logger
|
||||
|
||||
TASK_STATUS_QUEUED = "queued"
|
||||
TASK_STATUS_RUNNING = "running"
|
||||
TASK_STATUS_SUCCESS = "success"
|
||||
TASK_STATUS_ERROR = "error"
|
||||
TERMINAL_STATUSES = {TASK_STATUS_SUCCESS, TASK_STATUS_ERROR}
|
||||
UNFINISHED_STATUSES = {TASK_STATUS_QUEUED, TASK_STATUS_RUNNING}
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||||
|
||||
|
||||
def _timestamp(value: object) -> float:
|
||||
if not isinstance(value, str) or not value.strip():
|
||||
return 0.0
|
||||
for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%dT%H:%M:%S.%f", "%Y-%m-%dT%H:%M:%S"):
|
||||
try:
|
||||
return datetime.strptime(value[:26], fmt).timestamp()
|
||||
except ValueError:
|
||||
continue
|
||||
try:
|
||||
return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp()
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
|
||||
def _clean(value: object, default: str = "") -> str:
|
||||
return str(value or default).strip()
|
||||
|
||||
|
||||
def _owner_id(identity: dict[str, object]) -> str:
|
||||
return _clean(identity.get("id")) or "anonymous"
|
||||
|
||||
|
||||
def _task_key(owner_id: str, task_id: str) -> str:
|
||||
return f"{owner_id}:{task_id}"
|
||||
|
||||
|
||||
def _collect_image_urls(data: list[Any]) -> list[str]:
|
||||
urls: list[str] = []
|
||||
for item in data:
|
||||
if isinstance(item, dict):
|
||||
url = item.get("url")
|
||||
if isinstance(url, str) and url:
|
||||
urls.append(url)
|
||||
return urls
|
||||
|
||||
|
||||
def _public_task(task: dict[str, Any]) -> dict[str, Any]:
|
||||
item = {
|
||||
"id": task.get("id"),
|
||||
"status": task.get("status"),
|
||||
"mode": task.get("mode"),
|
||||
"model": task.get("model"),
|
||||
"size": task.get("size"),
|
||||
"quality": task.get("quality"),
|
||||
"created_at": task.get("created_at"),
|
||||
"updated_at": task.get("updated_at"),
|
||||
}
|
||||
if task.get("conversation_id"):
|
||||
item["conversation_id"] = task.get("conversation_id")
|
||||
if task.get("data") is not None:
|
||||
item["data"] = task.get("data")
|
||||
if task.get("usage") is not None:
|
||||
item["usage"] = task.get("usage")
|
||||
if task.get("error"):
|
||||
item["error"] = task.get("error")
|
||||
if task.get("progress"):
|
||||
item["progress"] = task.get("progress")
|
||||
if task.get("duration_ms") is not None:
|
||||
item["duration_ms"] = task.get("duration_ms")
|
||||
if task.get("stage_durations_ms"):
|
||||
item["stage_durations_ms"] = dict(task["stage_durations_ms"])
|
||||
if task.get("status") in (TASK_STATUS_RUNNING, TASK_STATUS_QUEUED):
|
||||
if task.get("status") == TASK_STATUS_RUNNING:
|
||||
# Include account selection, upload and the upstream response wait.
|
||||
base_ts = task.get("started_ts")
|
||||
else:
|
||||
# QUEUED 状态从 created_ts 开始计时（排队等待中）
|
||||
base_ts = task.get("created_ts") or task.get("updated_ts")
|
||||
if base_ts:
|
||||
item["elapsed_secs"] = round(time.time() - base_ts, 1)
|
||||
return item
|
||||
|
||||
|
||||
class ImageTaskService:
|
||||
def __init__(
|
||||
self,
|
||||
path: Path,
|
||||
*,
|
||||
generation_handler: Callable[[dict[str, Any]], dict[str, Any]] = openai_v1_image_generations.handle,
|
||||
edit_handler: Callable[[dict[str, Any]], dict[str, Any]] = openai_v1_image_edit.handle,
|
||||
retention_days_getter: Callable[[], int] | None = None,
|
||||
):
|
||||
self.path = path
|
||||
self.generation_handler = generation_handler
|
||||
self.edit_handler = edit_handler
|
||||
self.retention_days_getter = retention_days_getter or (lambda: config.image_retention_days)
|
||||
self._lock = threading.RLock()
|
||||
self._tasks: dict[str, dict[str, Any]] = {}
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
with self._lock:
|
||||
self._tasks = self._load_locked()
|
||||
changed = self._recover_unfinished_locked()
|
||||
changed = self._cleanup_locked() or changed
|
||||
if changed:
|
||||
self._save_locked()
|
||||
|
||||
def submit_generation(
|
||||
self,
|
||||
identity: dict[str, object],
|
||||
*,
|
||||
client_task_id: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
size: str | None,
|
||||
quality: str = "auto",
|
||||
base_url: str = "",
|
||||
) -> dict[str, Any]:
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"n": 1,
|
||||
"size": size,
|
||||
"quality": quality,
|
||||
"response_format": "url",
|
||||
"base_url": base_url,
|
||||
}
|
||||
return self._submit(identity, client_task_id=client_task_id, mode="generate", payload=payload)
|
||||
|
||||
def submit_edit(
|
||||
self,
|
||||
identity: dict[str, object],
|
||||
*,
|
||||
client_task_id: str,
|
||||
prompt: str,
|
||||
model: str,
|
||||
size: str | None,
|
||||
quality: str = "auto",
|
||||
base_url: str = "",
|
||||
images: list[tuple[bytes, str, str]] | None = None,
|
||||
masks: list[tuple[bytes, str, str]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"images": images or [],
|
||||
"mask": masks or [],
|
||||
"model": model,
|
||||
"n": 1,
|
||||
"size": size,
|
||||
"quality": quality,
|
||||
"response_format": "url",
|
||||
"base_url": base_url,
|
||||
}
|
||||
return self._submit(identity, client_task_id=client_task_id, mode="edit", payload=payload)
|
||||
|
||||
def list_tasks(self, identity: dict[str, object], task_ids: list[str]) -> dict[str, Any]:
|
||||
owner = _owner_id(identity)
|
||||
requested_ids = [_clean(task_id) for task_id in task_ids if _clean(task_id)]
|
||||
with self._lock:
|
||||
if self._cleanup_locked():
|
||||
self._save_locked()
|
||||
items = []
|
||||
missing_ids = []
|
||||
for task_id in requested_ids:
|
||||
task = self._tasks.get(_task_key(owner, task_id))
|
||||
if task is None:
|
||||
missing_ids.append(task_id)
|
||||
else:
|
||||
items.append(_public_task(task))
|
||||
if not requested_ids:
|
||||
items = [
|
||||
_public_task(task)
|
||||
for task in self._tasks.values()
|
||||
if task.get("owner_id") == owner
|
||||
]
|
||||
items.sort(key=lambda item: str(item.get("updated_at") or ""), reverse=True)
|
||||
missing_ids = []
|
||||
return {"items": items, "missing_ids": missing_ids}
|
||||
|
||||
def _submit(
|
||||
self,
|
||||
identity: dict[str, object],
|
||||
*,
|
||||
client_task_id: str,
|
||||
mode: str,
|
||||
payload: dict[str, Any],
|
||||
) -> dict[str, Any]:
|
||||
task_id = _clean(client_task_id)
|
||||
if not task_id:
|
||||
raise ValueError("client_task_id is required")
|
||||
owner = _owner_id(identity)
|
||||
key = _task_key(owner, task_id)
|
||||
now = _now_iso()
|
||||
should_start = False
|
||||
with self._lock:
|
||||
cleaned = self._cleanup_locked()
|
||||
task = self._tasks.get(key)
|
||||
if task is not None:
|
||||
if cleaned:
|
||||
self._save_locked()
|
||||
return _public_task(task)
|
||||
task = {
|
||||
"id": task_id,
|
||||
"owner_id": owner,
|
||||
"status": TASK_STATUS_QUEUED,
|
||||
"mode": mode,
|
||||
"model": _clean(payload.get("model"), "gpt-image-2"),
|
||||
"size": _clean(payload.get("size")),
|
||||
"quality": _clean(payload.get("quality"), "auto"),
|
||||
"created_at": now,
|
||||
"updated_at": now,
|
||||
"created_ts": time.time(),
|
||||
}
|
||||
self._tasks[key] = task
|
||||
self._save_locked()
|
||||
should_start = True
|
||||
|
||||
if should_start:
|
||||
thread = threading.Thread(
|
||||
target=self._run_task,
|
||||
args=(key, mode, payload, dict(identity), _clean(payload.get("model"), "gpt-image-2")),
|
||||
name=f"image-task-{task_id[:16]}",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
return _public_task(task)
|
||||
|
||||
def _run_task(
|
||||
self,
|
||||
key: str,
|
||||
mode: str,
|
||||
payload: dict[str, Any],
|
||||
identity: dict[str, object],
|
||||
model: str,
|
||||
) -> None:
|
||||
started = time.time()
|
||||
self._update_task(key, status=TASK_STATUS_RUNNING, error="", started_ts=started)
|
||||
stage_started = time.monotonic()
|
||||
current_stage = "getting_account"
|
||||
stage_durations: dict[str, int] = {}
|
||||
timing_lock = threading.Lock()
|
||||
|
||||
def advance_stage(step: str | None) -> dict[str, int]:
|
||||
nonlocal stage_started, current_stage
|
||||
with timing_lock:
|
||||
now = time.monotonic()
|
||||
stage_durations[current_stage] = stage_durations.get(current_stage, 0) + max(0, int((now - stage_started) * 1000))
|
||||
stage_started = now
|
||||
if step:
|
||||
current_stage = step
|
||||
return dict(stage_durations)
|
||||
|
||||
# 创建进度回调，每个步骤完成后更新任务状态
|
||||
def progress_callback(step: str) -> None:
|
||||
self._update_task(key, progress=step, stage_durations_ms=advance_stage(step))
|
||||
# 将进度回调添加到 payload 中（handler 会提取并传递给 ConversationRequest）
|
||||
payload_with_progress = {**payload, "progress_callback": progress_callback}
|
||||
try:
|
||||
handler = self.edit_handler if mode == "edit" else self.generation_handler
|
||||
result = handler(payload_with_progress)
|
||||
if not isinstance(result, dict):
|
||||
raise RuntimeError("image task returned streaming result unexpectedly")
|
||||
data = result.get("data")
|
||||
account_email = _clean(result.get("_account_email") or result.get("account_email"))
|
||||
if not isinstance(data, list) or not data:
|
||||
upstream = _clean(result.get("message"))
|
||||
if upstream:
|
||||
message = upstream
|
||||
else:
|
||||
message = "号池中没有可用账号或所有账号均被限流，请检查号池状态（账号额度、是否被封禁、是否到达生图上限）"
|
||||
error = RuntimeError(message)
|
||||
if account_email:
|
||||
setattr(error, "account_email", account_email)
|
||||
raise error
|
||||
usage = result.get("usage")
|
||||
duration_ms = int((time.time() - started) * 1000)
|
||||
timings = advance_stage(None)
|
||||
self._update_task(key, status=TASK_STATUS_SUCCESS, data=data, usage=usage, error="", duration_ms=duration_ms, stage_durations_ms=timings)
|
||||
logger.info({"event": "image_task_timings", "task_id": payload.get("client_task_id"), "duration_ms": duration_ms, "stages_ms": timings})
|
||||
self._log_call(
|
||||
identity,
|
||||
mode,
|
||||
model,
|
||||
started,
|
||||
"调用完成",
|
||||
request_preview=request_text(payload.get("prompt")),
|
||||
urls=_collect_image_urls(data),
|
||||
account_email=account_email,
|
||||
)
|
||||
except Exception as exc:
|
||||
error_message = str(exc) or "image task failed"
|
||||
account_email = _clean(getattr(exc, "account_email", ""))
|
||||
conversation_id = _clean(getattr(exc, "conversation_id", ""))
|
||||
duration_ms = int((time.time() - started) * 1000)
|
||||
timings = advance_stage(None)
|
||||
self._update_task(key, status=TASK_STATUS_ERROR, error=error_message, data=[],
|
||||
duration_ms=duration_ms, stage_durations_ms=timings,
|
||||
**({"conversation_id": conversation_id} if conversation_id else {}))
|
||||
logger.info({"event": "image_task_timings", "task_id": payload.get("client_task_id"), "duration_ms": duration_ms, "stages_ms": timings})
|
||||
self._log_call(
|
||||
identity,
|
||||
mode,
|
||||
model,
|
||||
started,
|
||||
"调用失败",
|
||||
request_preview=request_text(payload.get("prompt")),
|
||||
status="failed",
|
||||
error=error_message,
|
||||
account_email=account_email,
|
||||
)
|
||||
|
||||
def _log_call(
|
||||
self,
|
||||
identity: dict[str, object],
|
||||
mode: str,
|
||||
model: str,
|
||||
started: float,
|
||||
suffix: str,
|
||||
*,
|
||||
request_preview: str = "",
|
||||
status: str = "success",
|
||||
error: str = "",
|
||||
urls: list[str] | None = None,
|
||||
account_email: str = "",
|
||||
) -> None:
|
||||
endpoint = "/v1/images/edits" if mode == "edit" else "/v1/images/generations"
|
||||
summary_prefix = "图生图" if mode == "edit" else "文生图"
|
||||
detail = {
|
||||
"key_id": identity.get("id"),
|
||||
"key_name": identity.get("name"),
|
||||
"role": identity.get("role"),
|
||||
"endpoint": endpoint,
|
||||
"model": model,
|
||||
"started_at": datetime.fromtimestamp(started).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"ended_at": _now_iso(),
|
||||
"duration_ms": int((time.time() - started) * 1000),
|
||||
"status": status,
|
||||
}
|
||||
if request_preview:
|
||||
detail["request_text"] = request_preview
|
||||
if error:
|
||||
detail["error"] = error
|
||||
if account_email:
|
||||
detail["account_email"] = account_email
|
||||
if urls:
|
||||
detail["urls"] = list(dict.fromkeys(urls))
|
||||
try:
|
||||
log_service.add(LOG_TYPE_CALL, f"{summary_prefix}{suffix}", detail)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _update_task(self, key: str, **updates: Any) -> None:
|
||||
with self._lock:
|
||||
task = self._tasks.get(key)
|
||||
if task is None:
|
||||
return
|
||||
task.update(updates)
|
||||
task["updated_at"] = _now_iso()
|
||||
task["updated_ts"] = time.time()
|
||||
self._save_locked()
|
||||
|
||||
def _load_locked(self) -> dict[str, dict[str, Any]]:
|
||||
if not self.path.exists():
|
||||
return {}
|
||||
try:
|
||||
raw = json.loads(self.path.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
raw_items = raw.get("tasks") if isinstance(raw, dict) else raw
|
||||
if not isinstance(raw_items, list):
|
||||
return {}
|
||||
tasks: dict[str, dict[str, Any]] = {}
|
||||
for item in raw_items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
task_id = _clean(item.get("id"))
|
||||
owner = _clean(item.get("owner_id"))
|
||||
if not task_id or not owner:
|
||||
continue
|
||||
status = _clean(item.get("status"))
|
||||
if status not in {TASK_STATUS_QUEUED, TASK_STATUS_RUNNING, TASK_STATUS_SUCCESS, TASK_STATUS_ERROR}:
|
||||
status = TASK_STATUS_ERROR
|
||||
task = {
|
||||
"id": task_id,
|
||||
"owner_id": owner,
|
||||
"status": status,
|
||||
"mode": "edit" if item.get("mode") == "edit" else "generate",
|
||||
"model": _clean(item.get("model"), "gpt-image-2"),
|
||||
"size": _clean(item.get("size")),
|
||||
"quality": _clean(item.get("quality"), "auto"),
|
||||
"created_at": _clean(item.get("created_at"), _now_iso()),
|
||||
"updated_at": _clean(item.get("updated_at"), _clean(item.get("created_at"), _now_iso())),
|
||||
"created_ts": item.get("created_ts"),
|
||||
"updated_ts": item.get("updated_ts"),
|
||||
"started_ts": item.get("started_ts"),
|
||||
"duration_ms": item.get("duration_ms"),
|
||||
}
|
||||
data = item.get("data")
|
||||
if isinstance(data, list):
|
||||
task["data"] = data
|
||||
usage = item.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
task["usage"] = usage
|
||||
timings = item.get("stage_durations_ms")
|
||||
if isinstance(timings, dict):
|
||||
task["stage_durations_ms"] = {
|
||||
str(stage): duration for stage, duration in timings.items()
|
||||
if isinstance(duration, int) and duration >= 0
|
||||
}
|
||||
error = _clean(item.get("error"))
|
||||
if error:
|
||||
task["error"] = error
|
||||
tasks[_task_key(owner, task_id)] = task
|
||||
return tasks
|
||||
|
||||
def _save_locked(self) -> None:
|
||||
items = sorted(self._tasks.values(), key=lambda item: str(item.get("updated_at") or ""), reverse=True)
|
||||
tmp_path = self.path.with_suffix(self.path.suffix + ".tmp")
|
||||
tmp_path.write_text(json.dumps({"tasks": items}, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
tmp_path.replace(self.path)
|
||||
|
||||
def _recover_unfinished_locked(self) -> bool:
|
||||
changed = False
|
||||
for task in self._tasks.values():
|
||||
if task.get("status") in UNFINISHED_STATUSES:
|
||||
task["status"] = TASK_STATUS_ERROR
|
||||
task["error"] = "服务已重启，未完成的图片任务已中断"
|
||||
task["updated_at"] = _now_iso()
|
||||
changed = True
|
||||
return changed
|
||||
|
||||
def _cleanup_locked(self) -> bool:
|
||||
try:
|
||||
retention_days = max(1, int(self.retention_days_getter()))
|
||||
except Exception:
|
||||
retention_days = 30
|
||||
cutoff = time.time() - retention_days * 86400
|
||||
removed_keys = [
|
||||
key
|
||||
for key, task in self._tasks.items()
|
||||
if task.get("status") in TERMINAL_STATUSES and _timestamp(task.get("updated_at")) < cutoff
|
||||
]
|
||||
for key in removed_keys:
|
||||
self._tasks.pop(key, None)
|
||||
return bool(removed_keys)
|
||||
|
||||
def resume_poll(
|
||||
self,
|
||||
identity: dict[str, object],
|
||||
task_id: str,
|
||||
extra_timeout_secs: float = 30.0,
|
||||
) -> dict[str, Any]:
|
||||
"""恢复对已超时任务的轮询，额外等待 extra_timeout_secs 秒。"""
|
||||
owner = _owner_id(identity)
|
||||
key = _task_key(owner, _clean(task_id))
|
||||
with self._lock:
|
||||
task = self._tasks.get(key)
|
||||
if task is None:
|
||||
raise ValueError("task not found")
|
||||
if task.get("status") != TASK_STATUS_ERROR:
|
||||
raise ValueError("task is not in error state")
|
||||
error_msg = _clean(task.get("error"))
|
||||
if "超时" not in error_msg:
|
||||
raise ValueError("task error is not a timeout error")
|
||||
conversation_id = _clean(task.get("conversation_id"))
|
||||
if not conversation_id:
|
||||
raise ValueError("task has no conversation_id")
|
||||
mode = task.get("mode", "generate")
|
||||
model = task.get("model", "gpt-image-2")
|
||||
# 将任务状态重置为 running
|
||||
self._update_task(key, status=TASK_STATUS_RUNNING, error="")
|
||||
|
||||
# 启动新线程继续轮询
|
||||
thread = threading.Thread(
|
||||
target=self._run_resume_poll,
|
||||
args=(key, conversation_id, extra_timeout_secs, dict(identity), mode, model),
|
||||
name=f"image-resume-{_clean(task_id)[:16]}",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
return _public_task(task)
|
||||
|
||||
def _run_resume_poll(
|
||||
self,
|
||||
key: str,
|
||||
conversation_id: str,
|
||||
extra_timeout_secs: float,
|
||||
identity: dict[str, object],
|
||||
mode: str,
|
||||
model: str,
|
||||
) -> None:
|
||||
"""后台线程：继续轮询已有 conversation_id 的图片结果。"""
|
||||
started = time.time()
|
||||
backend = None
|
||||
try:
|
||||
from services.openai_backend_api import OpenAIBackendAPI
|
||||
from services.protocol.conversation import format_image_result
|
||||
|
||||
backend = OpenAIBackendAPI(proxy_url=config.proxy_url or None)
|
||||
file_ids, sediment_ids = backend._poll_image_results(
|
||||
conversation_id,
|
||||
extra_timeout_secs,
|
||||
)
|
||||
if not file_ids and not sediment_ids:
|
||||
raise RuntimeError(
|
||||
f"继续等待 {extra_timeout_secs} 秒后仍未找到图片结果。"
|
||||
)
|
||||
|
||||
image_urls = backend.resolve_conversation_image_urls(
|
||||
conversation_id, file_ids, sediment_ids, poll=False,
|
||||
)
|
||||
if not image_urls:
|
||||
raise RuntimeError("图片 URL 解析失败")
|
||||
|
||||
image_items = [
|
||||
{"b64_json": __import__("base64").b64encode(image_data).decode("ascii")}
|
||||
for image_data in backend.download_image_bytes(image_urls)
|
||||
]
|
||||
# 获取 task 的原始 prompt（从 _public_task 的 mode 判断）
|
||||
with self._lock:
|
||||
task = self._tasks.get(key)
|
||||
quality = _clean(task.get("quality"), "auto") if task else "auto"
|
||||
size = _clean(task.get("size")) if task else None
|
||||
data = format_image_result(
|
||||
image_items,
|
||||
"", # prompt 已不重要，结果已经拿到了
|
||||
"b64_json",
|
||||
"",
|
||||
int(time.time()),
|
||||
)["data"]
|
||||
self._update_task(key, status=TASK_STATUS_SUCCESS, data=data, error="", duration_ms=int((time.time() - started) * 1000))
|
||||
self._log_call(
|
||||
identity,
|
||||
mode,
|
||||
model,
|
||||
started,
|
||||
"调用完成（续轮询）",
|
||||
status="success",
|
||||
urls=_collect_image_urls(data),
|
||||
)
|
||||
except Exception as exc:
|
||||
error_message = str(exc) or "resume poll failed"
|
||||
duration_ms = int((time.time() - started) * 1000)
|
||||
self._update_task(key, status=TASK_STATUS_ERROR, error=error_message, data=[], duration_ms=duration_ms)
|
||||
self._log_call(
|
||||
identity,
|
||||
mode,
|
||||
model,
|
||||
started,
|
||||
"调用失败（续轮询）",
|
||||
status="failed",
|
||||
error=error_message,
|
||||
)
|
||||
finally:
|
||||
if backend is not None:
|
||||
backend.close()
|
||||
|
||||
|
||||
image_task_service = ImageTaskService(DATA_DIR / "image_tasks.json")
|
||||
@@ -0,0 +1,337 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import itertools
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import HTTPException
|
||||
from fastapi.concurrency import run_in_threadpool
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
|
||||
from services.config import DATA_DIR
|
||||
from services.protocol.error_response import anthropic_error_response, openai_error_response
|
||||
from utils.helper import anthropic_sse_stream, sse_json_stream
|
||||
|
||||
LOG_TYPE_CALL = "call"
|
||||
LOG_TYPE_ACCOUNT = "account"
|
||||
INTERNAL_RESPONSE_KEYS = {"_account_email", "_conversation_id"}
|
||||
|
||||
|
||||
class LogService:
|
||||
def __init__(self, path: Path):
|
||||
self.path = path
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _legacy_id(raw_line: str, line_number: int) -> str:
|
||||
payload = f"{line_number}:{raw_line}".encode("utf-8", errors="ignore")
|
||||
return hashlib.sha1(payload).hexdigest()[:24]
|
||||
|
||||
def _parse_line(self, raw_line: str, line_number: int) -> dict[str, Any] | None:
|
||||
try:
|
||||
item = json.loads(raw_line)
|
||||
except Exception:
|
||||
return None
|
||||
if not isinstance(item, dict):
|
||||
return None
|
||||
parsed = dict(item)
|
||||
parsed["id"] = str(parsed.get("id") or self._legacy_id(raw_line, line_number))
|
||||
return parsed
|
||||
|
||||
@staticmethod
|
||||
def _serialize_item(item: dict[str, Any]) -> str:
|
||||
return json.dumps(item, ensure_ascii=False, separators=(",", ":"))
|
||||
|
||||
@staticmethod
|
||||
def _matches_filters(item: dict[str, Any], *, type: str = "", start_date: str = "", end_date: str = "") -> bool:
|
||||
t = str(item.get("time") or "")
|
||||
day = t[:10]
|
||||
if type and item.get("type") != type:
|
||||
return False
|
||||
if start_date and day < start_date:
|
||||
return False
|
||||
if end_date and day > end_date:
|
||||
return False
|
||||
return True
|
||||
|
||||
def add(self, type: str, summary: str = "", detail: dict[str, Any] | None = None, **data: Any) -> None:
|
||||
item = {
|
||||
"id": uuid4().hex,
|
||||
"time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"type": type,
|
||||
"summary": summary,
|
||||
"detail": detail or data,
|
||||
}
|
||||
with self.path.open("a", encoding="utf-8") as file:
|
||||
file.write(self._serialize_item(item) + "\n")
|
||||
|
||||
def list(self, type: str = "", start_date: str = "", end_date: str = "", limit: int = 200) -> list[dict[str, Any]]:
|
||||
if not self.path.exists():
|
||||
return []
|
||||
items: list[dict[str, Any]] = []
|
||||
lines = self.path.read_text(encoding="utf-8").splitlines()
|
||||
for line_number in range(len(lines) - 1, -1, -1):
|
||||
item = self._parse_line(lines[line_number], line_number)
|
||||
if item is None:
|
||||
continue
|
||||
if not self._matches_filters(item, type=type, start_date=start_date, end_date=end_date):
|
||||
continue
|
||||
items.append(item)
|
||||
if len(items) >= limit:
|
||||
break
|
||||
return items
|
||||
|
||||
def delete(self, ids: list[str]) -> dict[str, int]:
|
||||
target_ids = {str(item or "").strip() for item in ids if str(item or "").strip()}
|
||||
if not self.path.exists() or not target_ids:
|
||||
return {"removed": 0}
|
||||
lines = self.path.read_text(encoding="utf-8").splitlines()
|
||||
kept_lines: list[str] = []
|
||||
removed = 0
|
||||
for line_number, raw_line in enumerate(lines):
|
||||
item = self._parse_line(raw_line, line_number)
|
||||
if item is None:
|
||||
kept_lines.append(raw_line)
|
||||
continue
|
||||
if str(item.get("id") or "") in target_ids:
|
||||
removed += 1
|
||||
continue
|
||||
kept_lines.append(self._serialize_item(item))
|
||||
content = "\n".join(kept_lines)
|
||||
if content:
|
||||
content += "\n"
|
||||
self.path.write_text(content, encoding="utf-8")
|
||||
return {"removed": removed}
|
||||
|
||||
|
||||
log_service = LogService(DATA_DIR / "logs.jsonl")
|
||||
|
||||
|
||||
def _collect_urls(value: object) -> list[str]:
|
||||
urls: list[str] = []
|
||||
if isinstance(value, dict):
|
||||
for key, item in value.items():
|
||||
if key == "url" and isinstance(item, str):
|
||||
urls.append(item)
|
||||
elif key == "urls" and isinstance(item, list):
|
||||
urls.extend(str(url) for url in item if isinstance(url, str))
|
||||
else:
|
||||
urls.extend(_collect_urls(item))
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
urls.extend(_collect_urls(item))
|
||||
return urls
|
||||
|
||||
|
||||
def _collect_account_emails(value: object) -> list[str]:
|
||||
emails: list[str] = []
|
||||
if isinstance(value, dict):
|
||||
for key, item in value.items():
|
||||
if key in {"_account_email", "account_email"} and isinstance(item, str) and item.strip():
|
||||
emails.append(item.strip())
|
||||
else:
|
||||
emails.extend(_collect_account_emails(item))
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
emails.extend(_collect_account_emails(item))
|
||||
return emails
|
||||
|
||||
|
||||
def _collect_conversation_ids(value: object) -> list[str]:
|
||||
ids: list[str] = []
|
||||
if isinstance(value, dict):
|
||||
for key, item in value.items():
|
||||
if key == "_conversation_id" and isinstance(item, str) and item.strip():
|
||||
ids.append(item.strip())
|
||||
else:
|
||||
ids.extend(_collect_conversation_ids(item))
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
ids.extend(_collect_conversation_ids(item))
|
||||
return ids
|
||||
|
||||
|
||||
def _strip_internal_response_fields(value: object) -> object:
|
||||
if isinstance(value, dict):
|
||||
return {
|
||||
key: _strip_internal_response_fields(item)
|
||||
for key, item in value.items()
|
||||
if key not in INTERNAL_RESPONSE_KEYS
|
||||
}
|
||||
if isinstance(value, list):
|
||||
return [_strip_internal_response_fields(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def _request_excerpt(text: object, limit: int = 1000) -> str:
|
||||
value = str(text or "").strip()
|
||||
if not value:
|
||||
return ""
|
||||
normalized = " ".join(value.split())
|
||||
if len(normalized) <= limit:
|
||||
return normalized
|
||||
return normalized[: limit - 1].rstrip() + "…"
|
||||
|
||||
|
||||
def _image_error_response(exc: Exception) -> JSONResponse:
|
||||
from services.protocol.conversation import public_image_error_message
|
||||
|
||||
message = public_image_error_message(str(exc))
|
||||
if "no available image quota" in message.lower():
|
||||
return openai_error_response(
|
||||
{
|
||||
"error": {
|
||||
"message": "no available image quota",
|
||||
"type": "insufficient_quota",
|
||||
"param": None,
|
||||
"code": "insufficient_quota",
|
||||
}
|
||||
},
|
||||
429,
|
||||
)
|
||||
if hasattr(exc, "to_openai_error") and hasattr(exc, "status_code"):
|
||||
return JSONResponse(status_code=int(exc.status_code), content=exc.to_openai_error())
|
||||
return openai_error_response(message, 502)
|
||||
|
||||
|
||||
def _protocol_error_response(exc: Exception, status_code: int, sse: str) -> JSONResponse:
|
||||
message = str(exc)
|
||||
if sse == "anthropic":
|
||||
return anthropic_error_response(message, status_code)
|
||||
return openai_error_response(message, status_code)
|
||||
|
||||
|
||||
def _next_item(items):
|
||||
try:
|
||||
return True, next(items)
|
||||
except StopIteration:
|
||||
return False, None
|
||||
|
||||
|
||||
@dataclass
|
||||
class LoggedCall:
|
||||
identity: dict[str, object]
|
||||
endpoint: str
|
||||
model: str
|
||||
summary: str
|
||||
started: float = field(default_factory=time.time)
|
||||
request_text: str = ""
|
||||
request_shape: dict[str, int] | None = None
|
||||
|
||||
async def run(self, handler, *args, sse: str = "openai"):
|
||||
from services.protocol.conversation import ImageGenerationError
|
||||
|
||||
try:
|
||||
result = await run_in_threadpool(handler, *args)
|
||||
except ImageGenerationError as exc:
|
||||
self.log("调用失败", status="failed", error=str(exc), account_email=getattr(exc, "account_email", ""),
|
||||
conversation_id=getattr(exc, "conversation_id", ""))
|
||||
return _image_error_response(exc)
|
||||
except HTTPException as exc:
|
||||
self.log("调用失败", status="failed", error=str(exc.detail))
|
||||
raise
|
||||
except Exception as exc:
|
||||
self.log("调用失败", status="failed", error=str(exc), account_email=getattr(exc, "account_email", ""))
|
||||
if self.endpoint.startswith("/v1/images"):
|
||||
return _image_error_response(exc)
|
||||
return _protocol_error_response(exc, 502, sse)
|
||||
|
||||
if isinstance(result, dict):
|
||||
self.log("调用完成", result)
|
||||
response = dict(result)
|
||||
response.pop("_account_email", None)
|
||||
return response
|
||||
|
||||
sender = anthropic_sse_stream if sse == "anthropic" else sse_json_stream
|
||||
try:
|
||||
has_first, first = await run_in_threadpool(_next_item, result)
|
||||
except ImageGenerationError as exc:
|
||||
self.log("调用失败", status="failed", error=str(exc), account_email=getattr(exc, "account_email", ""),
|
||||
conversation_id=getattr(exc, "conversation_id", ""))
|
||||
return _image_error_response(exc)
|
||||
except HTTPException as exc:
|
||||
self.log("调用失败", status="failed", error=str(exc.detail))
|
||||
raise
|
||||
except Exception as exc:
|
||||
self.log("调用失败", status="failed", error=str(exc), account_email=getattr(exc, "account_email", ""))
|
||||
if self.endpoint.startswith("/v1/images"):
|
||||
return _image_error_response(exc)
|
||||
return _protocol_error_response(exc, 502, sse)
|
||||
if not has_first:
|
||||
self.log("流式调用结束")
|
||||
return StreamingResponse(sender(()), media_type="text/event-stream")
|
||||
return StreamingResponse(sender(self.stream(itertools.chain([first], result))), media_type="text/event-stream")
|
||||
|
||||
def stream(self, items):
|
||||
urls: list[str] = []
|
||||
account_emails: list[str] = []
|
||||
conversation_ids: list[str] = []
|
||||
failed = False
|
||||
try:
|
||||
for item in items:
|
||||
urls.extend(_collect_urls(item))
|
||||
account_emails.extend(_collect_account_emails(item))
|
||||
conversation_ids.extend(_collect_conversation_ids(item))
|
||||
yield _strip_internal_response_fields(item)
|
||||
except Exception as exc:
|
||||
failed = True
|
||||
self.log(
|
||||
"流式调用失败",
|
||||
status="failed",
|
||||
error=str(exc),
|
||||
urls=urls,
|
||||
account_email=(account_emails[0] if account_emails else getattr(exc, "account_email", "")),
|
||||
conversation_id=(conversation_ids[0] if conversation_ids else getattr(exc, "conversation_id", "")),
|
||||
)
|
||||
if self.endpoint.startswith("/v1/images") and not hasattr(exc, "to_openai_error"):
|
||||
from services.protocol.conversation import ImageGenerationError, public_image_error_message
|
||||
|
||||
raise ImageGenerationError(public_image_error_message(str(exc))) from exc
|
||||
raise
|
||||
finally:
|
||||
if not failed:
|
||||
self.log("流式调用结束", urls=urls, account_email=account_emails[0] if account_emails else "",
|
||||
conversation_id=conversation_ids[0] if conversation_ids else "")
|
||||
|
||||
def log(self, suffix: str, result: object = None, status: str = "success", error: str = "",
|
||||
urls: list[str] | None = None, account_email: str = "", conversation_id: str = "") -> None:
|
||||
detail = {
|
||||
"key_id": self.identity.get("id"),
|
||||
"key_name": self.identity.get("name"),
|
||||
"role": self.identity.get("role"),
|
||||
"endpoint": self.endpoint,
|
||||
"model": self.model,
|
||||
"started_at": datetime.fromtimestamp(self.started).strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"ended_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
||||
"duration_ms": int((time.time() - self.started) * 1000),
|
||||
"status": status,
|
||||
}
|
||||
request_excerpt = _request_excerpt(self.request_text)
|
||||
if request_excerpt:
|
||||
detail["request_text"] = request_excerpt
|
||||
if self.request_shape:
|
||||
detail["request_shape"] = self.request_shape
|
||||
if error:
|
||||
detail["error"] = error
|
||||
email = str(account_email or "").strip()
|
||||
if not email:
|
||||
emails = _collect_account_emails(result)
|
||||
email = emails[0] if emails else ""
|
||||
if email:
|
||||
detail["account_email"] = email
|
||||
conv_id = str(conversation_id or "").strip()
|
||||
if not conv_id:
|
||||
conv_ids = _collect_conversation_ids(result)
|
||||
conv_id = conv_ids[0] if conv_ids else ""
|
||||
if conv_id:
|
||||
detail["conversation_id"] = conv_id
|
||||
collected_urls = [*(urls or []), *_collect_urls(result)]
|
||||
if collected_urls and not self.endpoint.startswith("/v1/search"):
|
||||
detail["urls"] = list(dict.fromkeys(collected_urls))
|
||||
log_service.add(LOG_TYPE_CALL, f"{self.summary}{suffix}", detail)
|
||||
@@ -0,0 +1,182 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from threading import RLock
|
||||
from typing import Any
|
||||
|
||||
from services.account_service import AccountService, account_service
|
||||
from services.openai_backend_api import OpenAIBackendAPI
|
||||
from utils.log import logger
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelRoute:
|
||||
account_types: frozenset[str]
|
||||
allow_anonymous: bool = False
|
||||
|
||||
|
||||
class ModelUnavailableError(RuntimeError):
|
||||
pass
|
||||
|
||||
|
||||
class ModelCatalogService:
|
||||
"""Caches the model catalogs advertised to each active account type."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
accounts: AccountService,
|
||||
*,
|
||||
backend_factory: Callable[..., Any] = OpenAIBackendAPI,
|
||||
cache_ttl_seconds: float = 300,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> None:
|
||||
self._accounts = accounts
|
||||
self._backend_factory = backend_factory
|
||||
self._cache_ttl_seconds = max(1.0, float(cache_ttl_seconds))
|
||||
self._clock = clock
|
||||
self._lock = RLock()
|
||||
self._expires_at = 0.0
|
||||
self._account_signature: tuple[tuple[str, int], ...] = ()
|
||||
self._anonymous_models: dict[str, dict[str, Any]] = {}
|
||||
self._models_by_account_type: dict[str, dict[str, dict[str, Any]]] = {}
|
||||
|
||||
@staticmethod
|
||||
def _model_map(result: object) -> dict[str, dict[str, Any]]:
|
||||
if not isinstance(result, dict) or not isinstance(result.get("data"), list):
|
||||
raise TypeError("upstream model response has no data list")
|
||||
models: dict[str, dict[str, Any]] = {}
|
||||
for item in result["data"]:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
model_id = str(item.get("id") or "").strip()
|
||||
if model_id and model_id not in models:
|
||||
models[model_id] = dict(item)
|
||||
return models
|
||||
|
||||
def _active_accounts_by_type(self) -> dict[str, list[str]]:
|
||||
groups: dict[str, list[str]] = {}
|
||||
for account in self._accounts.list_accounts():
|
||||
if not isinstance(account, dict) or account.get("status") in {"禁用", "异常"}:
|
||||
continue
|
||||
access_token = str(account.get("access_token") or "").strip()
|
||||
account_type = self._accounts._normalize_account_type(account.get("type"))
|
||||
if access_token and account_type:
|
||||
groups.setdefault(account_type, []).append(access_token)
|
||||
return groups
|
||||
|
||||
@staticmethod
|
||||
def _signature(groups: dict[str, list[str]]) -> tuple[tuple[str, int], ...]:
|
||||
return tuple(
|
||||
(account_type, len(tokens))
|
||||
for account_type, tokens in sorted(groups.items())
|
||||
)
|
||||
|
||||
def _fetch_models(self, access_token: str = "") -> dict[str, dict[str, Any]]:
|
||||
backend = self._backend_factory(access_token=access_token)
|
||||
try:
|
||||
return self._model_map(backend.list_models())
|
||||
finally:
|
||||
backend.close()
|
||||
|
||||
def _fetch_account_type_models(
|
||||
self,
|
||||
account_type: str,
|
||||
access_tokens: list[str],
|
||||
) -> dict[str, dict[str, Any]] | None:
|
||||
attempted_tokens: set[str] = set()
|
||||
last_error: Exception | None = None
|
||||
for access_token in access_tokens:
|
||||
try:
|
||||
resolved_token = self._accounts.refresh_access_token(
|
||||
access_token,
|
||||
event="list_models",
|
||||
) or access_token
|
||||
if resolved_token in attempted_tokens:
|
||||
continue
|
||||
attempted_tokens.add(resolved_token)
|
||||
return self._fetch_models(resolved_token)
|
||||
except Exception as exc: # noqa: BLE001 - try the next account for any upstream failure
|
||||
last_error = exc
|
||||
if last_error is not None:
|
||||
logger.warning({
|
||||
"event": "model_catalog_account_type_failed",
|
||||
"account_type": account_type,
|
||||
"error_type": type(last_error).__name__,
|
||||
})
|
||||
return None
|
||||
|
||||
def _refresh(self, groups: dict[str, list[str]], signature: tuple[tuple[str, int], ...]) -> None:
|
||||
models_by_account_type: dict[str, dict[str, dict[str, Any]]] = {}
|
||||
with ThreadPoolExecutor(max_workers=min(4, len(groups) + 1)) as executor:
|
||||
anonymous_future = executor.submit(self._fetch_models)
|
||||
account_futures = {
|
||||
account_type: executor.submit(
|
||||
self._fetch_account_type_models,
|
||||
account_type,
|
||||
access_tokens,
|
||||
)
|
||||
for account_type, access_tokens in groups.items()
|
||||
}
|
||||
try:
|
||||
anonymous_models = anonymous_future.result()
|
||||
except Exception as exc: # noqa: BLE001 - retain cached models on upstream failure
|
||||
logger.warning({
|
||||
"event": "model_catalog_anonymous_failed",
|
||||
"error_type": type(exc).__name__,
|
||||
})
|
||||
anonymous_models = self._anonymous_models
|
||||
|
||||
for account_type, future in account_futures.items():
|
||||
models = future.result()
|
||||
if models is not None:
|
||||
models_by_account_type[account_type] = models
|
||||
elif account_type in self._models_by_account_type:
|
||||
models_by_account_type[account_type] = self._models_by_account_type[account_type]
|
||||
|
||||
self._anonymous_models = anonymous_models
|
||||
self._models_by_account_type = models_by_account_type
|
||||
self._account_signature = signature
|
||||
self._expires_at = self._clock() + self._cache_ttl_seconds
|
||||
|
||||
def _ensure_catalog(self) -> None:
|
||||
groups = self._active_accounts_by_type()
|
||||
signature = self._signature(groups)
|
||||
with self._lock:
|
||||
if signature == self._account_signature and self._clock() < self._expires_at:
|
||||
return
|
||||
self._refresh(groups, signature)
|
||||
|
||||
def list_models(self) -> dict[str, Any]:
|
||||
self._ensure_catalog()
|
||||
with self._lock:
|
||||
union: dict[str, dict[str, Any]] = {
|
||||
model_id: dict(item)
|
||||
for model_id, item in self._anonymous_models.items()
|
||||
}
|
||||
for account_type in sorted(self._models_by_account_type):
|
||||
for model_id, item in self._models_by_account_type[account_type].items():
|
||||
union.setdefault(model_id, dict(item))
|
||||
return {
|
||||
"object": "list",
|
||||
"data": [union[model_id] for model_id in sorted(union)],
|
||||
}
|
||||
|
||||
def route_for_model(self, model: str) -> ModelRoute:
|
||||
model = str(model or "").strip()
|
||||
self._ensure_catalog()
|
||||
with self._lock:
|
||||
account_types = frozenset(
|
||||
account_type
|
||||
for account_type, models in self._models_by_account_type.items()
|
||||
if model in models
|
||||
)
|
||||
return ModelRoute(
|
||||
account_types=account_types,
|
||||
allow_anonymous=model in self._anonymous_models,
|
||||
)
|
||||
|
||||
|
||||
model_catalog_service = ModelCatalogService(account_service)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""手动 OAuth 桥服务
|
||||
|
||||
让用户用自己浏览器走一遍 OpenAI 的标准 OAuth + PKCE 授权码流程：
|
||||
1. 后端生成 code_verifier / code_challenge / state，构造 authorize URL。
|
||||
2. 用户在浏览器登录，浏览器最终被 OpenAI 重定向到 platform.openai.com 的
|
||||
callback 地址；用户从地址栏或 devtools 抓出 code，回填到前端。
|
||||
3. 后端拿之前存好的 code_verifier + 回填的 code 调用 /api/accounts/oauth/token
|
||||
得到 {access_token, refresh_token, id_token}。
|
||||
|
||||
得到的 refresh_token 跟 account_service 自动刷新机制用的 client_id 是同一个
|
||||
（app_2SKx67EdpoN0G6j64rFvigXD），所以落盘后能直接进入 keepalive 周期。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any
|
||||
from urllib.parse import parse_qs, urlencode, urlparse
|
||||
|
||||
from curl_cffi import requests
|
||||
|
||||
from services.openai_oauth import (
|
||||
auth_base,
|
||||
common_headers,
|
||||
platform_auth0_client,
|
||||
platform_base,
|
||||
platform_oauth_audience,
|
||||
platform_oauth_client_id,
|
||||
platform_oauth_redirect_uri,
|
||||
sec_ch_ua,
|
||||
user_agent,
|
||||
)
|
||||
from services.proxy_service import proxy_settings
|
||||
|
||||
|
||||
class OAuthLoginError(Exception):
|
||||
"""OAuth 桥流程中的可预期错误，会被 API 层翻译成 400。"""
|
||||
|
||||
|
||||
class OAuthLoginService:
|
||||
"""维护 PKCE 临时会话，并完成 code → token 的兑换。"""
|
||||
|
||||
_SESSION_TTL_SECONDS = 10 * 60 # 用户点开浏览器 + 拿 code 给的时间上限
|
||||
_MAX_SESSIONS = 64 # 防止异常累积；超过容量时清理最老的
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._sessions: dict[str, dict[str, Any]] = {}
|
||||
|
||||
@staticmethod
|
||||
def _generate_pkce() -> tuple[str, str]:
|
||||
"""生成 PKCE code_verifier 与对应的 code_challenge（S256）。"""
|
||||
from utils.pkce import generate_pkce
|
||||
return generate_pkce()
|
||||
|
||||
def _purge_expired_locked(self) -> None:
|
||||
"""清理过期或溢出容量的会话，必须在持锁状态下调用。"""
|
||||
now = time.time()
|
||||
expired = [sid for sid, item in self._sessions.items() if now - item["created_at"] > self._SESSION_TTL_SECONDS]
|
||||
for sid in expired:
|
||||
self._sessions.pop(sid, None)
|
||||
if len(self._sessions) > self._MAX_SESSIONS:
|
||||
ordered = sorted(self._sessions.items(), key=lambda kv: kv[1]["created_at"])
|
||||
for sid, _ in ordered[: len(self._sessions) - self._MAX_SESSIONS]:
|
||||
self._sessions.pop(sid, None)
|
||||
|
||||
def start(self, email_hint: str = "") -> dict[str, str]:
|
||||
"""登记一个新的 PKCE 会话，返回 session_id 与可让用户打开的 authorize_url。
|
||||
|
||||
state 形如 "<session_id>.<nonce>"，让 callback URL 自带 session_id，
|
||||
finish 时即便前端 React 状态被覆盖也能从 URL 恢复正确的 verifier。
|
||||
"""
|
||||
verifier, challenge = self._generate_pkce()
|
||||
nonce = secrets.token_urlsafe(32)
|
||||
device_id = str(uuid.uuid4())
|
||||
session_id = uuid.uuid4().hex
|
||||
state = f"{session_id}.{secrets.token_urlsafe(16)}"
|
||||
|
||||
params = {
|
||||
"issuer": auth_base,
|
||||
"client_id": platform_oauth_client_id,
|
||||
"audience": platform_oauth_audience,
|
||||
"redirect_uri": platform_oauth_redirect_uri,
|
||||
"device_id": device_id,
|
||||
"screen_hint": "login_or_signup",
|
||||
"max_age": "0",
|
||||
"scope": "openid profile email offline_access",
|
||||
"response_type": "code",
|
||||
"response_mode": "query",
|
||||
"state": state,
|
||||
"nonce": nonce,
|
||||
"code_challenge": challenge,
|
||||
"code_challenge_method": "S256",
|
||||
"auth0Client": platform_auth0_client,
|
||||
}
|
||||
email_hint = str(email_hint or "").strip()
|
||||
if email_hint:
|
||||
params["login_hint"] = email_hint
|
||||
|
||||
authorize_url = f"{auth_base}/api/accounts/authorize?{urlencode(params)}"
|
||||
|
||||
with self._lock:
|
||||
self._purge_expired_locked()
|
||||
self._sessions[session_id] = {
|
||||
"code_verifier": verifier,
|
||||
"state": state,
|
||||
"created_at": time.time(),
|
||||
"redirect_uri": platform_oauth_redirect_uri,
|
||||
}
|
||||
|
||||
return {
|
||||
"session_id": session_id,
|
||||
"authorize_url": authorize_url,
|
||||
"expires_in": str(self._SESSION_TTL_SECONDS),
|
||||
"redirect_uri_prefix": platform_oauth_redirect_uri,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _extract_code_from_callback(value: str) -> tuple[str, str]:
|
||||
"""从 callback URL 或 raw code 中提取 (code, state)。
|
||||
|
||||
既允许用户粘贴整段 platform.openai.com/auth/callback?code=...&state=... 的 URL，
|
||||
也允许只粘 code 本身。
|
||||
"""
|
||||
raw = str(value or "").strip()
|
||||
if not raw:
|
||||
return "", ""
|
||||
if raw.startswith("http://") or raw.startswith("https://"):
|
||||
try:
|
||||
parsed = parse_qs(urlparse(raw).query)
|
||||
except Exception as exc:
|
||||
raise OAuthLoginError(f"无法解析 callback URL: {exc}") from exc
|
||||
code = str((parsed.get("code") or [""])[0]).strip()
|
||||
state = str((parsed.get("state") or [""])[0]).strip()
|
||||
if not code:
|
||||
err = str((parsed.get("error_description") or parsed.get("error") or [""])[0]).strip()
|
||||
raise OAuthLoginError(err or "callback URL 中没有 code 参数")
|
||||
return code, state
|
||||
# 用户可能直接粘了 code 字符串
|
||||
return raw, ""
|
||||
|
||||
def finish(self, session_id: str, callback: str) -> dict[str, str]:
|
||||
"""用 session_id 配对的 code_verifier 把 callback 里的 code 换成 token 三件套。
|
||||
|
||||
- 优先用 callback URL 自带 state 里的 session_id（更可靠），
|
||||
找不到才用前端传来的 session_id；
|
||||
- 失败时不立刻销毁 session（OAuth code 错配换 token 失败通常不会消耗 code），
|
||||
只有成功兑换才 pop，便于用户用同一 verifier 重试。
|
||||
"""
|
||||
body_sid = str(session_id or "").strip()
|
||||
code, state = self._extract_code_from_callback(callback)
|
||||
if not code:
|
||||
raise OAuthLoginError("缺少 code 或 callback URL")
|
||||
|
||||
# state 里嵌的 session_id 优先级最高
|
||||
state_sid = state.split(".", 1)[0] if state else ""
|
||||
candidate_sids = [sid for sid in (state_sid, body_sid) if sid]
|
||||
if not candidate_sids:
|
||||
raise OAuthLoginError("既未提供 session_id，callback URL 中也未携带 state")
|
||||
|
||||
with self._lock:
|
||||
self._purge_expired_locked()
|
||||
session = None
|
||||
picked_sid = ""
|
||||
for sid in candidate_sids:
|
||||
cur = self._sessions.get(sid)
|
||||
if cur is not None:
|
||||
session = cur
|
||||
picked_sid = sid
|
||||
break
|
||||
if session is None:
|
||||
raise OAuthLoginError(
|
||||
"OAuth 会话已过期或不存在，请回到导入对话框点\"重新生成\"再走一次"
|
||||
)
|
||||
|
||||
if state and session.get("state") and state != session["state"]:
|
||||
raise OAuthLoginError(
|
||||
"state 不匹配。常见原因：你点过两次\"打开授权页面\"，但浏览器里登录的还是前一次的窗口。请点\"重新生成\"重来。"
|
||||
)
|
||||
|
||||
tokens = self._exchange_code(
|
||||
code,
|
||||
session["code_verifier"],
|
||||
session.get("redirect_uri") or platform_oauth_redirect_uri,
|
||||
)
|
||||
# 仅在成功兑换之后才消耗 session
|
||||
with self._lock:
|
||||
self._sessions.pop(picked_sid, None)
|
||||
return tokens
|
||||
|
||||
@staticmethod
|
||||
def _exchange_code(code: str, code_verifier: str, redirect_uri: str) -> dict[str, str]:
|
||||
"""调用 /api/accounts/oauth/token 用 code+verifier 换 token 三件套。"""
|
||||
kwargs = proxy_settings.build_session_kwargs(impersonate="chrome", verify=False)
|
||||
session = requests.Session(**kwargs)
|
||||
try:
|
||||
response = session.post(
|
||||
f"{auth_base}/api/accounts/oauth/token",
|
||||
headers={
|
||||
**common_headers,
|
||||
"referer": f"{platform_base}/",
|
||||
"origin": platform_base,
|
||||
"auth0-client": platform_auth0_client,
|
||||
"sec-ch-ua": sec_ch_ua,
|
||||
"user-agent": user_agent,
|
||||
},
|
||||
json={
|
||||
"client_id": platform_oauth_client_id,
|
||||
"code_verifier": code_verifier,
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": redirect_uri,
|
||||
},
|
||||
timeout=60,
|
||||
)
|
||||
except Exception as exc:
|
||||
raise OAuthLoginError(f"换 token 网络异常: {exc}") from exc
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
try:
|
||||
data = response.json() if response.text else {}
|
||||
except Exception:
|
||||
data = {}
|
||||
|
||||
if response.status_code != 200 or not isinstance(data, dict) or not data.get("access_token"):
|
||||
detail = ""
|
||||
if isinstance(data, dict):
|
||||
detail = str(data.get("error_description") or data.get("error") or data.get("message") or "")
|
||||
if not detail:
|
||||
try:
|
||||
detail = str(response.text or "")[:300]
|
||||
except Exception:
|
||||
detail = ""
|
||||
# 打到 docker logs 方便排错——OAuth 换 token 的失败原因往往只有这里能看到
|
||||
print(
|
||||
f"[oauth-login] /api/accounts/oauth/token rejected: "
|
||||
f"status={response.status_code} detail={detail!r} "
|
||||
f"raw_body={(getattr(response, 'text', '') or '')[:500]!r}",
|
||||
flush=True,
|
||||
)
|
||||
raise OAuthLoginError(
|
||||
f"OpenAI 拒绝换 token (HTTP {response.status_code}){': ' + detail if detail else ''}"
|
||||
)
|
||||
|
||||
access_token = str(data.get("access_token") or "").strip()
|
||||
refresh_token = str(data.get("refresh_token") or "").strip()
|
||||
id_token = str(data.get("id_token") or "").strip()
|
||||
|
||||
if not access_token:
|
||||
raise OAuthLoginError("OpenAI 返回的 access_token 为空")
|
||||
if not refresh_token:
|
||||
# scope 含 offline_access 时正常会下发 refresh_token；这里给出明确提示
|
||||
raise OAuthLoginError(
|
||||
"OpenAI 没有返回 refresh_token（可能 scope 未包含 offline_access 或 code 已使用过）"
|
||||
)
|
||||
|
||||
return {
|
||||
"access_token": access_token,
|
||||
"refresh_token": refresh_token,
|
||||
"id_token": id_token,
|
||||
}
|
||||
|
||||
|
||||
oauth_login_service = OAuthLoginService()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,32 @@
|
||||
auth_base = "https://auth.openai.com"
|
||||
platform_base = "https://platform.openai.com"
|
||||
platform_oauth_client_id = "app_2SKx67EdpoN0G6j64rFvigXD"
|
||||
platform_oauth_redirect_uri = f"{platform_base}/auth/callback"
|
||||
platform_oauth_audience = "https://api.openai.com/v1"
|
||||
platform_auth0_client = "eyJuYW1lIjoiYXV0aDAtc3BhLWpzIiwidmVyc2lvbiI6IjEuMjEuMCJ9"
|
||||
user_agent = (
|
||||
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
|
||||
"AppleWebKit/537.36 (KHTML, like Gecko) "
|
||||
"Chrome/145.0.0.0 Safari/537.36"
|
||||
)
|
||||
sec_ch_ua = '"Google Chrome";v="145", "Not?A_Brand";v="8", "Chromium";v="145"'
|
||||
|
||||
common_headers = {
|
||||
"accept": "application/json",
|
||||
"accept-encoding": "gzip, deflate, br",
|
||||
"accept-language": "en-US,en;q=0.9",
|
||||
"cache-control": "no-cache",
|
||||
"connection": "keep-alive",
|
||||
"content-type": "application/json",
|
||||
"dnt": "1",
|
||||
"origin": auth_base,
|
||||
"priority": "u=1, i",
|
||||
"sec-gpc": "1",
|
||||
"sec-ch-ua": sec_ch_ua,
|
||||
"sec-ch-ua-mobile": "?0",
|
||||
"sec-ch-ua-platform": '"Windows"',
|
||||
"sec-fetch-dest": "empty",
|
||||
"sec-fetch-mode": "cors",
|
||||
"sec-fetch-site": "same-origin",
|
||||
"user-agent": user_agent,
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
"""Protocol converters for OpenAI-compatible endpoints."""
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterable, Iterator
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.openai_backend_api import OpenAIBackendAPI
|
||||
from services.protocol.conversation import count_message_tokens, count_text_tokens, normalize_messages
|
||||
from services.protocol.openai_v1_chat_complete import collect_chat_content, stream_text_chat_completion
|
||||
|
||||
XML_TOOL_RULE = """Tool output adapter: when calling tools, output ONLY this XML and no prose/markdown:
|
||||
<tool_calls><tool_call><tool_name>TOOL_NAME</tool_name><parameters><PARAM><![CDATA[value]]></PARAM></parameters></tool_call></tool_calls>"""
|
||||
|
||||
|
||||
@dataclass
|
||||
class MessageRequest:
|
||||
backend: OpenAIBackendAPI
|
||||
messages: list[dict[str, Any]]
|
||||
model: str
|
||||
tools: Any = None
|
||||
|
||||
|
||||
def _tool_meta(tool: dict[str, object]) -> tuple[str, str, object]:
|
||||
fn = tool.get("function") if isinstance(tool.get("function"), dict) else {}
|
||||
name = str(tool.get("name") or fn.get("name") or "").strip()
|
||||
desc = str(tool.get("description") or fn.get("description") or "").strip()
|
||||
schema = tool.get("input_schema") or tool.get("parameters") or fn.get("input_schema") or fn.get("parameters") or {}
|
||||
return name, desc, schema
|
||||
|
||||
|
||||
def build_tool_prompt(tools: object) -> str:
|
||||
if not isinstance(tools, list):
|
||||
return ""
|
||||
blocks = []
|
||||
for tool in tools:
|
||||
if not isinstance(tool, dict):
|
||||
continue
|
||||
name, desc, schema = _tool_meta(tool)
|
||||
if name:
|
||||
blocks.append(f"Tool: {name}\nDescription: {desc}\nParameters: {json.dumps(schema, ensure_ascii=False)}")
|
||||
if not blocks:
|
||||
return ""
|
||||
return "Available tools:\n" + "\n".join(blocks) + """
|
||||
|
||||
Tool use rules:
|
||||
- If the user asks to list/read/search files, inspect project state, run a command, or answer from local code, you MUST call a suitable tool first. Do not say you cannot access files.
|
||||
- To call tools, output ONLY XML and no prose/markdown:
|
||||
<tool_calls><tool_call><tool_name>TOOL_NAME</tool_name><parameters><PARAM><![CDATA[value]]></PARAM></parameters></tool_call></tool_calls>
|
||||
- Put parameters under <parameters> using the exact schema names.
|
||||
""".strip()
|
||||
|
||||
|
||||
def merge_system(system: object, extra: str) -> object:
|
||||
system = compact_system(system)
|
||||
if _has_claude_code_system(system):
|
||||
extra = XML_TOOL_RULE
|
||||
if not extra:
|
||||
return system
|
||||
if isinstance(system, str) and system.strip():
|
||||
return f"{system.strip()}\n\n{extra}"
|
||||
if isinstance(system, list):
|
||||
return [*system, {"type": "text", "text": extra}]
|
||||
return extra
|
||||
|
||||
|
||||
def _has_claude_code_system(system: object) -> bool:
|
||||
if isinstance(system, str):
|
||||
return "You are Claude Code" in system
|
||||
if isinstance(system, list):
|
||||
return any(isinstance(item, dict) and "You are Claude Code" in str(item.get("text") or "") for item in system)
|
||||
return False
|
||||
|
||||
|
||||
def compact_system(system: object) -> object:
|
||||
if isinstance(system, str):
|
||||
return _compact_system_text(system)
|
||||
if isinstance(system, list):
|
||||
result = []
|
||||
for item in system:
|
||||
if isinstance(item, dict) and str(item.get("type") or "") == "text":
|
||||
copied = dict(item)
|
||||
copied["text"] = _compact_system_text(str(item.get("text") or ""))
|
||||
result.append(copied)
|
||||
else:
|
||||
result.append(item)
|
||||
return result
|
||||
return system
|
||||
|
||||
|
||||
def _compact_system_text(text: str) -> str:
|
||||
return text or ""
|
||||
|
||||
|
||||
def _compact_message_text(text: str) -> str:
|
||||
return text or ""
|
||||
|
||||
|
||||
def preprocess_payload(payload: dict[str, object], text_mapper: Callable[[str], str] | None = None) -> dict[str, object]:
|
||||
payload["messages"] = preprocess_messages(payload.get("messages"), text_mapper)
|
||||
payload["system"] = merge_system(payload.get("system"), build_tool_prompt(payload.get("tools")))
|
||||
return payload
|
||||
|
||||
|
||||
def message_request(body: dict[str, Any]) -> MessageRequest:
|
||||
payload = preprocess_payload(dict(body))
|
||||
model = str(payload.get("model") or "auto").strip() or "auto"
|
||||
return MessageRequest(
|
||||
backend=OpenAIBackendAPI(access_token=account_service.get_text_access_token(model=model)),
|
||||
messages=normalize_messages(payload.get("messages"), payload.get("system")),
|
||||
model=model,
|
||||
tools=payload.get("tools"),
|
||||
)
|
||||
|
||||
|
||||
def preprocess_messages(messages: object, text_mapper: Callable[[str], str] | None = None) -> object:
|
||||
if not isinstance(messages, list):
|
||||
return messages
|
||||
mapper = text_mapper or (lambda text: text)
|
||||
result = []
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
item = dict(message)
|
||||
content = item.get("content")
|
||||
if isinstance(content, str):
|
||||
item["content"] = _compact_message_text(mapper(content))
|
||||
elif isinstance(content, list):
|
||||
item["content"] = [_preprocess_block(block, mapper) for block in content]
|
||||
result.append(item)
|
||||
return result
|
||||
|
||||
|
||||
def _preprocess_block(block: object, text_mapper: Callable[[str], str]) -> object:
|
||||
if not isinstance(block, dict):
|
||||
return block
|
||||
block_type = str(block.get("type") or "")
|
||||
if block_type == "text":
|
||||
item = dict(block)
|
||||
item["text"] = _compact_message_text(text_mapper(str(block.get("text") or "")))
|
||||
return item
|
||||
if block_type == "tool_use":
|
||||
return {"type": "text", "text": f"<tool_calls><tool_call><tool_name>{block.get('name') or ''}</tool_name><parameters>{json.dumps(block.get('input') or {}, ensure_ascii=False)}</parameters></tool_call></tool_calls>"}
|
||||
if block_type == "tool_result":
|
||||
return {"type": "text", "text": f"Tool result {block.get('tool_use_id') or ''}: {block.get('content') or ''}"}
|
||||
return block
|
||||
|
||||
|
||||
def message_response(model: str, text: str, input_tokens: int, output_tokens: int, tools: object = None) -> dict[str, object]:
|
||||
content, stop_reason = content_blocks(text, tools)
|
||||
return {
|
||||
"id": f"msg_{uuid.uuid4()}",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": model,
|
||||
"content": content,
|
||||
"stop_reason": stop_reason,
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
|
||||
}
|
||||
|
||||
|
||||
def content_blocks(text: str, tools: object = None) -> tuple[list[dict[str, object]], str]:
|
||||
calls = parse_tool_calls(text) if isinstance(tools, list) and tools else []
|
||||
text = strip_tool_markup(text)
|
||||
if calls:
|
||||
content = ([{"type": "text", "text": text}] if text else []) + [{"type": "tool_use", "id": f"toolu_{uuid.uuid4()}", "name": name, "input": args} for name, args in calls]
|
||||
return content, "tool_use"
|
||||
return [{"type": "text", "text": text}], "end_turn"
|
||||
|
||||
|
||||
def strip_tool_markup(text: str) -> str:
|
||||
return re.sub(r"(?is)<tool_calls\b[^>]*>.*?</tool_calls>|<tool_call\b[^>]*>.*?</tool_call>|<function_call\b[^>]*>.*?</function_call>|<invoke\b[^>]*>.*?</invoke>", "", text or "").strip()
|
||||
|
||||
|
||||
def streamable_text(text: str) -> str:
|
||||
text = text or ""
|
||||
match = re.search(r"(?is)<tool_calls\b|<tool_call\b|<function_call\b|<invoke\b", text)
|
||||
return text[:match.start()].rstrip() if match else text
|
||||
|
||||
|
||||
def parse_tool_calls(text: str) -> list[tuple[str, dict[str, object]]]:
|
||||
text = re.sub(r"(?is)```.*?```", "", text or "").strip()
|
||||
blocks = re.findall(r"(?is)<tool_call\b[^>]*>(.*?)</tool_call>|<function_call\b[^>]*>(.*?)</function_call>|<invoke\b[^>]*>(.*?)</invoke>", text)
|
||||
result = []
|
||||
for block in (next((part for part in match if part), "") for match in blocks):
|
||||
name = xml_value(block, "tool_name") or xml_value(block, "name") or xml_value(block, "function")
|
||||
params = xml_value(block, "parameters") or xml_value(block, "input") or xml_value(block, "arguments") or "{}"
|
||||
if name:
|
||||
result.append((name, parse_tool_params(params)))
|
||||
return result
|
||||
|
||||
|
||||
def xml_value(text: str, tag: str) -> str:
|
||||
match = re.search(rf"(?is)<{tag}\b[^>]*>(.*?)</{tag}>", text)
|
||||
if not match:
|
||||
return ""
|
||||
value = match.group(1).strip()
|
||||
cdata = re.fullmatch(r"(?is)<!\[CDATA\[(.*?)]]>", value)
|
||||
return html.unescape(cdata.group(1) if cdata else value).strip()
|
||||
|
||||
|
||||
def parse_tool_params(raw: str) -> dict[str, object]:
|
||||
raw = raw.strip()
|
||||
try:
|
||||
parsed = json.loads(raw)
|
||||
return parsed if isinstance(parsed, dict) else {}
|
||||
except Exception:
|
||||
return {m.group(1): parse_tool_value(m.group(2)) for m in re.finditer(r"(?is)<([\w.-]+)\b[^>]*>(.*?)</\1>", raw)}
|
||||
|
||||
|
||||
def parse_tool_value(raw: str) -> object:
|
||||
value = xml_value(f"<x>{raw}</x>", "x")
|
||||
try:
|
||||
return json.loads(value)
|
||||
except Exception:
|
||||
return value
|
||||
|
||||
|
||||
def stream_events(chunks: Iterable[dict[str, object]], model: str, input_tokens: int, output_tokens: Callable[[str], int], tools: object = None) -> Iterator[dict[str, object]]:
|
||||
message_id = f"msg_{uuid.uuid4()}"
|
||||
created = int(time.time())
|
||||
current_text = ""
|
||||
streamed_text = ""
|
||||
tool_mode = isinstance(tools, list) and bool(tools)
|
||||
tool_started = False
|
||||
text_open = False
|
||||
yield {"type": "message_start", "message": {"id": message_id, "type": "message", "role": "assistant", "model": model, "content": [], "stop_reason": None, "stop_sequence": None, "usage": {"input_tokens": input_tokens, "output_tokens": 0}}}
|
||||
if not tool_mode:
|
||||
text_open = True
|
||||
yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
|
||||
for chunk in chunks:
|
||||
choice = (chunk.get("choices") or [{}])[0]
|
||||
delta = choice.get("delta") or {}
|
||||
text_delta = delta.get("content", "") if isinstance(delta, dict) else ""
|
||||
if text_delta:
|
||||
current_text += text_delta
|
||||
if not tool_started:
|
||||
visible_text = current_text if not tool_mode else streamable_text(current_text)
|
||||
if visible_text.startswith(streamed_text):
|
||||
text_delta = visible_text[len(streamed_text):]
|
||||
if text_delta:
|
||||
if not text_open:
|
||||
text_open = True
|
||||
yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
|
||||
streamed_text = visible_text
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text_delta}}
|
||||
tool_started = tool_mode and visible_text != current_text
|
||||
if choice.get("finish_reason"):
|
||||
content, stop_reason = content_blocks(current_text, tools)
|
||||
if text_open:
|
||||
yield {"type": "content_block_stop", "index": 0}
|
||||
if stop_reason == "tool_use":
|
||||
start_index = 1 if text_open else 0
|
||||
if content and content[0]["type"] == "text":
|
||||
remaining = str(content[0].get("text") or "")[len(streamed_text):]
|
||||
if remaining:
|
||||
if not text_open:
|
||||
yield {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}
|
||||
yield {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": remaining}}
|
||||
if not text_open:
|
||||
yield {"type": "content_block_stop", "index": 0}
|
||||
start_index = 1
|
||||
content = content[1:]
|
||||
yield from _stream_buffered_blocks(content, start_index)
|
||||
yield {"type": "message_delta", "delta": {"stop_reason": stop_reason, "stop_sequence": None}, "usage": {"output_tokens": output_tokens(current_text)}}
|
||||
break
|
||||
yield {"type": "message_stop", "created": created}
|
||||
|
||||
|
||||
def _stream_buffered_blocks(content: list[dict[str, object]], start_index: int = 0) -> Iterator[dict[str, object]]:
|
||||
for offset, block in enumerate(content):
|
||||
index = start_index + offset
|
||||
if block["type"] == "tool_use":
|
||||
start = {"type": "tool_use", "id": block["id"], "name": block["name"], "input": {}}
|
||||
delta = {"type": "input_json_delta", "partial_json": json.dumps(block.get("input") or {}, ensure_ascii=False)}
|
||||
else:
|
||||
start = {"type": "text", "text": ""}
|
||||
delta = {"type": "text_delta", "text": block.get("text") or ""}
|
||||
yield {"type": "content_block_start", "index": index, "content_block": start}
|
||||
yield {"type": "content_block_delta", "index": index, "delta": delta}
|
||||
yield {"type": "content_block_stop", "index": index}
|
||||
|
||||
|
||||
def handle(body: dict[str, Any]) -> dict[str, Any] | Iterator[dict[str, Any]]:
|
||||
request = message_request(body)
|
||||
if body.get("stream"):
|
||||
return stream_events(
|
||||
stream_text_chat_completion(request.backend, request.messages, request.model),
|
||||
request.model,
|
||||
count_message_tokens(request.messages, request.model),
|
||||
lambda text: count_text_tokens(text, request.model),
|
||||
request.tools,
|
||||
)
|
||||
text = collect_chat_content(stream_text_chat_completion(request.backend, request.messages, request.model))
|
||||
return message_response(
|
||||
request.model,
|
||||
text,
|
||||
count_message_tokens(request.messages, request.model),
|
||||
count_text_tokens(text, request.model),
|
||||
request.tools,
|
||||
)
|
||||
@@ -0,0 +1,238 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import hashlib
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Iterable, Iterator
|
||||
|
||||
from services.config import config
|
||||
|
||||
CACHEABLE_TEXT_KEYS = {
|
||||
"frequency_penalty",
|
||||
"max_completion_tokens",
|
||||
"max_tokens",
|
||||
"metadata",
|
||||
"model",
|
||||
"presence_penalty",
|
||||
"reasoning_effort",
|
||||
"response_format",
|
||||
"seed",
|
||||
"stop",
|
||||
"temperature",
|
||||
"thinking_effort",
|
||||
"tool_choice",
|
||||
"tools",
|
||||
"top_p",
|
||||
"user",
|
||||
"reasoning",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CacheEntry:
|
||||
expires_at: float
|
||||
value: Any
|
||||
|
||||
|
||||
@dataclass
|
||||
class InflightCall:
|
||||
condition: threading.Condition = field(default_factory=lambda: threading.Condition(threading.RLock()))
|
||||
done: bool = False
|
||||
value: Any = None
|
||||
error: BaseException | None = None
|
||||
|
||||
|
||||
def _json_safe(value: Any) -> Any:
|
||||
if isinstance(value, bytes):
|
||||
return {"__bytes_sha256__": hashlib.sha256(value).hexdigest(), "length": len(value)}
|
||||
if isinstance(value, bytearray):
|
||||
data = bytes(value)
|
||||
return {"__bytes_sha256__": hashlib.sha256(data).hexdigest(), "length": len(data)}
|
||||
if isinstance(value, dict):
|
||||
return {str(key): _json_safe(item) for key, item in value.items()}
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [_json_safe(item) for item in value]
|
||||
return value
|
||||
|
||||
|
||||
def canonical_body(body: dict[str, Any], messages: list[dict[str, Any]], *, stream: bool) -> dict[str, Any]:
|
||||
payload = {key: body.get(key) for key in CACHEABLE_TEXT_KEYS if key in body}
|
||||
payload["messages"] = messages
|
||||
payload["stream"] = bool(stream)
|
||||
return payload
|
||||
|
||||
|
||||
def cache_key(body: dict[str, Any], messages: list[dict[str, Any]], *, stream: bool) -> str:
|
||||
encoded = json.dumps(
|
||||
_json_safe(canonical_body(body, messages, stream=stream)),
|
||||
ensure_ascii=False,
|
||||
sort_keys=True,
|
||||
separators=(",", ":"),
|
||||
).encode("utf-8")
|
||||
return hashlib.sha256(encoded).hexdigest()
|
||||
|
||||
|
||||
def _message_signature(message: dict[str, Any]) -> str:
|
||||
return json.dumps(_json_safe(message), ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
||||
|
||||
|
||||
def normalize_text_messages(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
settings = config.get_chat_completion_cache_settings()
|
||||
if not settings.get("normalize_messages"):
|
||||
return messages
|
||||
|
||||
normalized: list[dict[str, Any]] = []
|
||||
previous_signature = ""
|
||||
for message in messages:
|
||||
if settings.get("drop_assistant_history") and str(message.get("role") or "") == "assistant":
|
||||
continue
|
||||
signature = _message_signature(message)
|
||||
if settings.get("drop_adjacent_duplicates") and signature == previous_signature:
|
||||
continue
|
||||
normalized.append(message)
|
||||
previous_signature = signature
|
||||
return normalized
|
||||
|
||||
|
||||
class ChatCompletionCache:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.RLock()
|
||||
self._entries: dict[str, CacheEntry] = {}
|
||||
self._inflight: dict[str, InflightCall] = {}
|
||||
|
||||
def clear(self) -> None:
|
||||
with self._lock:
|
||||
self._entries.clear()
|
||||
self._inflight.clear()
|
||||
|
||||
def _settings(self) -> dict[str, object]:
|
||||
return config.get_chat_completion_cache_settings()
|
||||
|
||||
def _prune_locked(self, now: float, max_entries: int) -> None:
|
||||
expired = [key for key, item in self._entries.items() if item.expires_at <= now]
|
||||
for key in expired:
|
||||
self._entries.pop(key, None)
|
||||
while len(self._entries) > max_entries:
|
||||
oldest_key = min(self._entries, key=lambda key: self._entries[key].expires_at)
|
||||
self._entries.pop(oldest_key, None)
|
||||
|
||||
@staticmethod
|
||||
def _copy(value: Any) -> Any:
|
||||
return copy.deepcopy(value)
|
||||
|
||||
def get_or_compute_response(self, key: str, compute: Callable[[], dict[str, Any]]) -> dict[str, Any]:
|
||||
settings = self._settings()
|
||||
if not settings.get("enabled") or int(settings.get("ttl_seconds") or 0) <= 0:
|
||||
return compute()
|
||||
|
||||
now = time.time()
|
||||
max_entries = int(settings.get("max_entries") or 1)
|
||||
with self._lock:
|
||||
self._prune_locked(now, max_entries)
|
||||
entry = self._entries.get(key)
|
||||
if entry and entry.expires_at > now:
|
||||
return self._copy(entry.value)
|
||||
inflight = self._inflight.get(key) if settings.get("dedupe_inflight") else None
|
||||
if inflight is None:
|
||||
inflight = InflightCall()
|
||||
if settings.get("dedupe_inflight"):
|
||||
self._inflight[key] = inflight
|
||||
owner = True
|
||||
else:
|
||||
owner = False
|
||||
|
||||
if not owner:
|
||||
with inflight.condition:
|
||||
while not inflight.done:
|
||||
inflight.condition.wait()
|
||||
if inflight.error:
|
||||
raise inflight.error
|
||||
return self._copy(inflight.value)
|
||||
|
||||
try:
|
||||
value = compute()
|
||||
except BaseException as exc:
|
||||
with self._lock:
|
||||
self._inflight.pop(key, None)
|
||||
with inflight.condition:
|
||||
inflight.error = exc
|
||||
inflight.done = True
|
||||
inflight.condition.notify_all()
|
||||
raise
|
||||
|
||||
expires_at = time.time() + int(settings.get("ttl_seconds") or 0)
|
||||
with self._lock:
|
||||
self._entries[key] = CacheEntry(expires_at=expires_at, value=self._copy(value))
|
||||
self._prune_locked(time.time(), max_entries)
|
||||
self._inflight.pop(key, None)
|
||||
with inflight.condition:
|
||||
inflight.value = self._copy(value)
|
||||
inflight.done = True
|
||||
inflight.condition.notify_all()
|
||||
return value
|
||||
|
||||
def get_or_compute_stream(self, key: str, compute: Callable[[], Iterable[dict[str, Any]]]) -> Iterator[dict[str, Any]]:
|
||||
settings = self._settings()
|
||||
if (
|
||||
not settings.get("enabled")
|
||||
or not settings.get("stream_cache")
|
||||
or int(settings.get("ttl_seconds") or 0) <= 0
|
||||
):
|
||||
yield from compute()
|
||||
return
|
||||
|
||||
now = time.time()
|
||||
max_entries = int(settings.get("max_entries") or 1)
|
||||
with self._lock:
|
||||
self._prune_locked(now, max_entries)
|
||||
entry = self._entries.get(key)
|
||||
if entry and entry.expires_at > now:
|
||||
yield from self._copy(entry.value)
|
||||
return
|
||||
inflight = self._inflight.get(key) if settings.get("dedupe_inflight") else None
|
||||
if inflight is None:
|
||||
inflight = InflightCall()
|
||||
if settings.get("dedupe_inflight"):
|
||||
self._inflight[key] = inflight
|
||||
owner = True
|
||||
else:
|
||||
owner = False
|
||||
|
||||
if not owner:
|
||||
with inflight.condition:
|
||||
while not inflight.done:
|
||||
inflight.condition.wait()
|
||||
if inflight.error:
|
||||
raise inflight.error
|
||||
yield from self._copy(inflight.value)
|
||||
return
|
||||
|
||||
chunks: list[dict[str, Any]] = []
|
||||
try:
|
||||
for chunk in compute():
|
||||
chunks.append(self._copy(chunk))
|
||||
yield chunk
|
||||
except BaseException as exc:
|
||||
with self._lock:
|
||||
self._inflight.pop(key, None)
|
||||
with inflight.condition:
|
||||
inflight.error = exc
|
||||
inflight.done = True
|
||||
inflight.condition.notify_all()
|
||||
raise
|
||||
|
||||
expires_at = time.time() + int(settings.get("ttl_seconds") or 0)
|
||||
with self._lock:
|
||||
self._entries[key] = CacheEntry(expires_at=expires_at, value=self._copy(chunks))
|
||||
self._prune_locked(time.time(), max_entries)
|
||||
self._inflight.pop(key, None)
|
||||
with inflight.condition:
|
||||
inflight.value = self._copy(chunks)
|
||||
inflight.done = True
|
||||
inflight.condition.notify_all()
|
||||
|
||||
|
||||
chat_completion_cache = ChatCompletionCache()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,124 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
|
||||
def _message_from_value(value: object) -> str:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if not isinstance(value, dict):
|
||||
return ""
|
||||
message = value.get("message")
|
||||
if isinstance(message, str) and message:
|
||||
return message
|
||||
return _message_from_value(value.get("error"))
|
||||
|
||||
|
||||
def error_message_from_detail(detail: object) -> str:
|
||||
if isinstance(detail, list):
|
||||
messages = []
|
||||
for item in detail:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
location = ".".join(str(part) for part in item.get("loc", []) if part != "body")
|
||||
message = str(item.get("msg") or "").strip()
|
||||
if location and message:
|
||||
messages.append(f"{location}: {message}")
|
||||
elif message:
|
||||
messages.append(message)
|
||||
return "; ".join(messages)
|
||||
if isinstance(detail, dict):
|
||||
message = _message_from_value(detail.get("error")) or _message_from_value(detail)
|
||||
if message:
|
||||
return message
|
||||
return str(detail or "").strip()
|
||||
|
||||
|
||||
def _default_error_type(status_code: int) -> str:
|
||||
if status_code == 401:
|
||||
return "authentication_error"
|
||||
if status_code == 403:
|
||||
return "permission_error"
|
||||
if status_code == 429:
|
||||
return "rate_limit_error"
|
||||
if 400 <= status_code < 500:
|
||||
return "invalid_request_error"
|
||||
return "server_error"
|
||||
|
||||
|
||||
def _default_error_code(status_code: int) -> str:
|
||||
if status_code == 401:
|
||||
return "invalid_api_key"
|
||||
if status_code == 403:
|
||||
return "permission_denied"
|
||||
if status_code == 429:
|
||||
return "rate_limit_exceeded"
|
||||
if 400 <= status_code < 500:
|
||||
return "bad_request"
|
||||
return "upstream_error"
|
||||
|
||||
|
||||
def openai_error_payload(
|
||||
detail: object,
|
||||
status_code: int,
|
||||
*,
|
||||
error_type: str | None = None,
|
||||
code: object | None = None,
|
||||
param: object | None = None,
|
||||
) -> dict[str, Any]:
|
||||
error_detail = detail.get("error") if isinstance(detail, dict) else None
|
||||
if isinstance(error_detail, dict):
|
||||
return {
|
||||
"error": {
|
||||
"message": error_message_from_detail(error_detail) or "request failed",
|
||||
"type": str(error_detail.get("type") or error_type or _default_error_type(status_code)),
|
||||
"param": error_detail.get("param", param),
|
||||
"code": error_detail.get("code", code if code is not None else _default_error_code(status_code)),
|
||||
}
|
||||
}
|
||||
return {
|
||||
"error": {
|
||||
"message": error_message_from_detail(detail) or "request failed",
|
||||
"type": error_type or _default_error_type(status_code),
|
||||
"param": param,
|
||||
"code": code if code is not None else _default_error_code(status_code),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def openai_error_response(
|
||||
detail: object,
|
||||
status_code: int,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
error_type: str | None = None,
|
||||
code: object | None = None,
|
||||
param: object | None = None,
|
||||
) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content=openai_error_payload(detail, status_code, error_type=error_type, code=code, param=param),
|
||||
headers=headers,
|
||||
)
|
||||
|
||||
|
||||
def anthropic_error_response(
|
||||
detail: object,
|
||||
status_code: int,
|
||||
*,
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> JSONResponse:
|
||||
error_type = "api_error" if status_code >= 500 else _default_error_type(status_code)
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
content={
|
||||
"type": "error",
|
||||
"error": {
|
||||
"type": error_type,
|
||||
"message": error_message_from_detail(detail) or "request failed",
|
||||
},
|
||||
},
|
||||
headers=headers,
|
||||
)
|
||||
@@ -0,0 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.openai_backend_api import OpenAIBackendAPI, SEARCH_MODEL
|
||||
|
||||
MODEL = SEARCH_MODEL
|
||||
|
||||
|
||||
def handle(body: dict[str, object]) -> dict[str, object]:
|
||||
token = account_service.get_text_access_token()
|
||||
account = account_service.get_account(token) or {}
|
||||
backend = OpenAIBackendAPI(token)
|
||||
try:
|
||||
result = backend.search(str(body["prompt"]))
|
||||
finally:
|
||||
backend.close()
|
||||
account_service.mark_text_used(token)
|
||||
result["_account_email"] = str(account.get("email") or "")
|
||||
return result
|
||||
@@ -0,0 +1,316 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Iterable, Iterator
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from services.protocol.chat_completion_cache import cache_key, chat_completion_cache, normalize_text_messages
|
||||
from services.protocol.conversation import (
|
||||
ConversationRequest,
|
||||
ImageOutput,
|
||||
collect_image_outputs,
|
||||
collect_text,
|
||||
count_message_image_tokens,
|
||||
count_message_text_tokens,
|
||||
count_text_tokens,
|
||||
encode_images,
|
||||
normalize_messages,
|
||||
stream_image_outputs_with_pool,
|
||||
stream_text_deltas,
|
||||
text_backend,
|
||||
)
|
||||
from services.protocol.web_search_tool import (
|
||||
WEB_SEARCH_TOOL_TYPES,
|
||||
has_unsupported_tools,
|
||||
is_web_search_chat_request,
|
||||
run_web_search,
|
||||
search_query_from_messages,
|
||||
text_with_url_citations,
|
||||
)
|
||||
from utils.helper import build_chat_image_markdown_content, extract_chat_image, extract_chat_prompt, is_image_chat_request, parse_image_count
|
||||
from utils.image_tokens import (
|
||||
chat_usage_from_image_usage,
|
||||
count_image_inputs_tokens,
|
||||
count_image_output_items_tokens,
|
||||
image_usage,
|
||||
)
|
||||
|
||||
TOOL_UNAVAILABLE_SYSTEM_MESSAGE = (
|
||||
"This compatibility backend cannot execute local tools, shell commands, non-search tools, "
|
||||
"or file operations. Do not claim to have run tools or inspected external resources. "
|
||||
"If a user asks you to use a tool, say that tool execution is unavailable through this backend."
|
||||
)
|
||||
|
||||
|
||||
def normalize_thinking_effort(value: object) -> str:
|
||||
normalized = str(value or "").strip().lower()
|
||||
if normalized in {"", "none"}:
|
||||
return ""
|
||||
if normalized in {"low", "medium", "high"}:
|
||||
return normalized
|
||||
if normalized in {"xhigh", "extended"}:
|
||||
return "extended"
|
||||
return ""
|
||||
|
||||
|
||||
def thinking_effort_from_body(body: dict[str, Any]) -> str:
|
||||
if "thinking_effort" in body:
|
||||
return normalize_thinking_effort(body.get("thinking_effort"))
|
||||
if "reasoning_effort" in body:
|
||||
return normalize_thinking_effort(body.get("reasoning_effort"))
|
||||
reasoning = body.get("reasoning")
|
||||
if isinstance(reasoning, dict):
|
||||
return normalize_thinking_effort(reasoning.get("effort"))
|
||||
return ""
|
||||
|
||||
|
||||
def completion_chunk(model: str, delta: dict[str, Any], finish_reason: str | None = None, completion_id: str = "", created: int | None = None) -> dict[str, Any]:
|
||||
return {
|
||||
"id": completion_id or f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": created or int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
|
||||
}
|
||||
|
||||
|
||||
def completion_response(
|
||||
model: str,
|
||||
content: str,
|
||||
created: int | None = None,
|
||||
messages: list[dict[str, Any]] | None = None,
|
||||
annotations: list[dict[str, Any]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
prompt_text_tokens = count_message_text_tokens(messages, model) if messages else 0
|
||||
prompt_image_tokens = count_message_image_tokens(messages, model) if messages else 0
|
||||
prompt_tokens = prompt_text_tokens + prompt_image_tokens
|
||||
completion_tokens = count_text_tokens(content, model) if messages else 0
|
||||
message = {"role": "assistant", "content": content}
|
||||
if annotations:
|
||||
message["annotations"] = annotations
|
||||
return {
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": created or int(time.time()),
|
||||
"model": model,
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": message,
|
||||
"finish_reason": "stop",
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": prompt_tokens,
|
||||
"completion_tokens": completion_tokens,
|
||||
"total_tokens": prompt_tokens + completion_tokens,
|
||||
"prompt_tokens_details": {
|
||||
"text_tokens": prompt_text_tokens,
|
||||
"image_tokens": prompt_image_tokens,
|
||||
"cached_tokens": 0,
|
||||
},
|
||||
"completion_tokens_details": {
|
||||
"text_tokens": completion_tokens,
|
||||
"image_tokens": 0,
|
||||
"reasoning_tokens": 0,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def stream_text_chat_completion(
|
||||
backend,
|
||||
messages: list[dict[str, Any]],
|
||||
model: str,
|
||||
thinking_effort: str = "",
|
||||
) -> Iterator[dict[str, Any]]:
|
||||
completion_id = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
sent_role = False
|
||||
request = ConversationRequest(model=model, messages=messages, thinking_effort=thinking_effort)
|
||||
for delta_text in stream_text_deltas(backend, request):
|
||||
if not sent_role:
|
||||
sent_role = True
|
||||
yield completion_chunk(model, {"role": "assistant", "content": delta_text}, None, completion_id, created)
|
||||
else:
|
||||
yield completion_chunk(model, {"content": delta_text}, None, completion_id, created)
|
||||
if not sent_role:
|
||||
yield completion_chunk(model, {"role": "assistant", "content": ""}, None, completion_id, created)
|
||||
yield completion_chunk(model, {}, "stop", completion_id, created)
|
||||
|
||||
|
||||
def collect_chat_content(chunks: Iterable[dict[str, Any]]) -> str:
|
||||
parts: list[str] = []
|
||||
for chunk in chunks:
|
||||
choices = chunk.get("choices")
|
||||
first = choices[0] if isinstance(choices, list) and choices and isinstance(choices[0], dict) else {}
|
||||
delta = first.get("delta") if isinstance(first.get("delta"), dict) else {}
|
||||
content = str(delta.get("content") or "")
|
||||
if content:
|
||||
parts.append(content)
|
||||
return "".join(parts)
|
||||
|
||||
|
||||
def chat_messages_from_body(body: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
messages = body.get("messages")
|
||||
if isinstance(messages, list) and messages:
|
||||
return [message for message in messages if isinstance(message, dict)]
|
||||
prompt = str(body.get("prompt") or "").strip()
|
||||
if prompt:
|
||||
return [{"role": "user", "content": prompt}]
|
||||
raise HTTPException(status_code=400, detail={"error": "messages or prompt is required"})
|
||||
|
||||
|
||||
def chat_image_args(body: dict[str, Any]) -> tuple[str, str, int, list[tuple[bytes, str, str]]]:
|
||||
model = str(body.get("model") or "gpt-image-2").strip() or "gpt-image-2"
|
||||
prompt = extract_chat_prompt(body)
|
||||
if not prompt:
|
||||
raise HTTPException(status_code=400, detail={"error": "prompt is required"})
|
||||
images = [
|
||||
(data, f"image_{idx}.png", mime)
|
||||
for idx, (data, mime) in enumerate(extract_chat_image(body), start=1)
|
||||
]
|
||||
return model, prompt, parse_image_count(body.get("n")), images
|
||||
|
||||
|
||||
def text_chat_parts(body: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]:
|
||||
model = str(body.get("model") or "auto").strip() or "auto"
|
||||
messages = normalize_text_messages(normalize_messages(chat_messages_from_body(body)))
|
||||
if has_unsupported_tools(body, WEB_SEARCH_TOOL_TYPES):
|
||||
messages.insert(0, {"role": "system", "content": TOOL_UNAVAILABLE_SYSTEM_MESSAGE})
|
||||
return model, messages
|
||||
|
||||
|
||||
def chat_completion_annotations(annotations: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
output = []
|
||||
for item in annotations:
|
||||
if item.get("type") != "url_citation":
|
||||
continue
|
||||
output.append({
|
||||
"type": "url_citation",
|
||||
"url_citation": {
|
||||
"start_index": item.get("start_index", 0),
|
||||
"end_index": item.get("end_index", 0),
|
||||
"url": item.get("url", ""),
|
||||
"title": item.get("title", ""),
|
||||
},
|
||||
})
|
||||
return output
|
||||
|
||||
|
||||
def web_search_chat_response(messages: list[dict[str, Any]], model: str) -> dict[str, Any]:
|
||||
query = search_query_from_messages(messages)
|
||||
if not query:
|
||||
raise HTTPException(status_code=400, detail={"error": "messages or prompt is required for web search"})
|
||||
text, annotations = text_with_url_citations(run_web_search(query))
|
||||
return completion_response(
|
||||
model,
|
||||
text,
|
||||
messages=messages,
|
||||
annotations=chat_completion_annotations(annotations),
|
||||
)
|
||||
|
||||
|
||||
def stream_web_search_chat_completion(messages: list[dict[str, Any]], model: str) -> Iterator[dict[str, Any]]:
|
||||
query = search_query_from_messages(messages)
|
||||
if not query:
|
||||
raise HTTPException(status_code=400, detail={"error": "messages or prompt is required for web search"})
|
||||
text, _annotations = text_with_url_citations(run_web_search(query))
|
||||
completion_id = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
yield completion_chunk(model, {"role": "assistant", "content": text}, None, completion_id, created)
|
||||
yield completion_chunk(model, {}, "stop", completion_id, created)
|
||||
|
||||
|
||||
def image_result_content(result: dict[str, Any]) -> str:
|
||||
data = result.get("data")
|
||||
if isinstance(data, list) and data:
|
||||
return build_chat_image_markdown_content(result)
|
||||
return str(result.get("message") or "Image generation completed.")
|
||||
|
||||
|
||||
def image_chat_response(body: dict[str, Any]) -> dict[str, Any]:
|
||||
model, prompt, n, images = chat_image_args(body)
|
||||
result = collect_image_outputs(stream_image_outputs_with_pool(ConversationRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
n=n,
|
||||
response_format="b64_json",
|
||||
images=encode_images(images) or None,
|
||||
)))
|
||||
response = completion_response(model, image_result_content(result), int(result.get("created") or 0) or None)
|
||||
usage = image_usage(
|
||||
input_text_tokens=count_text_tokens(prompt, model),
|
||||
input_image_tokens=count_image_inputs_tokens(images, model),
|
||||
output_tokens=count_image_output_items_tokens(result.get("data")),
|
||||
)
|
||||
response["usage"] = chat_usage_from_image_usage(usage)
|
||||
return response
|
||||
|
||||
|
||||
def image_chat_events(body: dict[str, Any]) -> Iterator[dict[str, Any]]:
|
||||
model, prompt, n, images = chat_image_args(body)
|
||||
image_outputs = stream_image_outputs_with_pool(ConversationRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
n=n,
|
||||
response_format="b64_json",
|
||||
images=encode_images(images) or None,
|
||||
))
|
||||
yield from stream_image_chat_completion(image_outputs, model)
|
||||
|
||||
|
||||
def stream_image_chat_completion(image_outputs: Iterable[ImageOutput], model: str) -> Iterator[dict[str, Any]]:
|
||||
completion_id = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
sent_role = False
|
||||
sent_text = ""
|
||||
for output in image_outputs:
|
||||
content = ""
|
||||
if output.kind == "progress":
|
||||
content = output.text
|
||||
sent_text += content
|
||||
elif output.kind == "result":
|
||||
content = build_chat_image_markdown_content({"data": output.data})
|
||||
elif output.kind == "message":
|
||||
content = output.text[len(sent_text):] if output.text.startswith(sent_text) else output.text
|
||||
if not content:
|
||||
continue
|
||||
if not sent_role:
|
||||
sent_role = True
|
||||
yield completion_chunk(model, {"role": "assistant", "content": content}, None, completion_id, created)
|
||||
else:
|
||||
yield completion_chunk(model, {"content": content}, None, completion_id, created)
|
||||
if not sent_role:
|
||||
yield completion_chunk(model, {"role": "assistant", "content": ""}, None, completion_id, created)
|
||||
yield completion_chunk(model, {}, "stop", completion_id, created)
|
||||
|
||||
|
||||
def handle(body: dict[str, Any]) -> dict[str, Any] | Iterator[dict[str, Any]]:
|
||||
if body.get("stream"):
|
||||
if is_image_chat_request(body):
|
||||
return image_chat_events(body)
|
||||
model, messages = text_chat_parts(body)
|
||||
if is_web_search_chat_request(body) and not has_unsupported_tools(body, WEB_SEARCH_TOOL_TYPES):
|
||||
return stream_web_search_chat_completion(messages, model)
|
||||
thinking_effort = thinking_effort_from_body(body)
|
||||
key = cache_key(body, messages, stream=True)
|
||||
return chat_completion_cache.get_or_compute_stream(
|
||||
key,
|
||||
lambda: stream_text_chat_completion(text_backend(model), messages, model, thinking_effort),
|
||||
)
|
||||
if is_image_chat_request(body):
|
||||
return image_chat_response(body)
|
||||
model, messages = text_chat_parts(body)
|
||||
if is_web_search_chat_request(body) and not has_unsupported_tools(body, WEB_SEARCH_TOOL_TYPES):
|
||||
return web_search_chat_response(messages, model)
|
||||
thinking_effort = thinking_effort_from_body(body)
|
||||
key = cache_key(body, messages, stream=False)
|
||||
return chat_completion_cache.get_or_compute_response(
|
||||
key,
|
||||
lambda: completion_response(
|
||||
model,
|
||||
collect_text(text_backend(model), ConversationRequest(model=model, messages=messages, thinking_effort=thinking_effort)),
|
||||
messages=messages,
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,86 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from io import BytesIO
|
||||
from typing import Any, Iterator
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from services.protocol.conversation import (
|
||||
ConversationRequest,
|
||||
ImageGenerationError,
|
||||
collect_image_outputs,
|
||||
count_text_tokens,
|
||||
encode_images,
|
||||
stream_image_chunks,
|
||||
stream_image_outputs_with_pool,
|
||||
)
|
||||
from utils.image_tokens import count_image_inputs_tokens, count_image_output_items_tokens, image_usage
|
||||
|
||||
|
||||
def _composite_mask(
|
||||
images: list[tuple[bytes, str, str]],
|
||||
masks: list[tuple[bytes, str, str]],
|
||||
) -> list[tuple[bytes, str, str]]:
|
||||
"""将 mask 的 alpha 通道合成到图片中，标识需要编辑的区域。
|
||||
|
||||
mask 的透明区域（低 alpha）= 需要编辑的区域，
|
||||
mask 的不透明区域（高 alpha）= 保留的区域。
|
||||
如果无 mask 则返回原图。
|
||||
"""
|
||||
if not masks:
|
||||
return images
|
||||
result: list[tuple[bytes, str, str]] = []
|
||||
for i, (data, filename, mime_type) in enumerate(images):
|
||||
mask_data = masks[i][0] if i < len(masks) else masks[-1][0]
|
||||
img = Image.open(BytesIO(data)).convert("RGBA")
|
||||
mask_img = Image.open(BytesIO(mask_data))
|
||||
if mask_img.mode == "RGBA":
|
||||
alpha = mask_img.split()[3]
|
||||
elif mask_img.mode == "L":
|
||||
alpha = mask_img
|
||||
else:
|
||||
alpha = mask_img.convert("L")
|
||||
alpha = alpha.resize(img.size, Image.LANCZOS)
|
||||
img.putalpha(alpha)
|
||||
buf = BytesIO()
|
||||
img.save(buf, format="PNG")
|
||||
result.append((buf.getvalue(), filename, "image/png"))
|
||||
return result
|
||||
|
||||
|
||||
def handle(body: dict[str, Any]) -> dict[str, Any] | Iterator[dict[str, Any]]:
|
||||
prompt = str(body.get("prompt") or "")
|
||||
images = body.get("images") or []
|
||||
masks = body.get("mask") or []
|
||||
images = _composite_mask(images, masks)
|
||||
model = str(body.get("model") or "gpt-image-2")
|
||||
n = int(body.get("n") or 1)
|
||||
size = body.get("size")
|
||||
quality = str(body.get("quality") or "auto")
|
||||
response_format = str(body.get("response_format") or "b64_json")
|
||||
base_url = str(body.get("base_url") or "") or None
|
||||
progress_callback = body.get("progress_callback")
|
||||
encoded_images = encode_images(images)
|
||||
if not encoded_images:
|
||||
raise ImageGenerationError("image is required")
|
||||
outputs = stream_image_outputs_with_pool(ConversationRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
n=n,
|
||||
size=size,
|
||||
quality=quality,
|
||||
response_format=response_format,
|
||||
base_url=base_url,
|
||||
images=encoded_images,
|
||||
message_as_error=True,
|
||||
progress_callback=progress_callback,
|
||||
))
|
||||
if body.get("stream"):
|
||||
return stream_image_chunks(outputs)
|
||||
result = collect_image_outputs(outputs)
|
||||
result["usage"] = image_usage(
|
||||
input_text_tokens=count_text_tokens(prompt, model),
|
||||
input_image_tokens=count_image_inputs_tokens(images, model),
|
||||
output_tokens=count_image_output_items_tokens(result.get("data"), size, quality),
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,42 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Iterator
|
||||
|
||||
from services.protocol.conversation import (
|
||||
ConversationRequest,
|
||||
collect_image_outputs,
|
||||
count_text_tokens,
|
||||
stream_image_chunks,
|
||||
stream_image_outputs_with_pool,
|
||||
)
|
||||
from utils.image_tokens import count_image_output_items_tokens, image_usage
|
||||
|
||||
|
||||
def handle(body: dict[str, Any]) -> dict[str, Any] | Iterator[dict[str, Any]]:
|
||||
prompt = str(body.get("prompt") or "")
|
||||
model = str(body.get("model") or "gpt-image-2")
|
||||
n = int(body.get("n") or 1)
|
||||
size = body.get("size")
|
||||
quality = str(body.get("quality") or "auto")
|
||||
response_format = str(body.get("response_format") or "b64_json")
|
||||
base_url = str(body.get("base_url") or "") or None
|
||||
progress_callback = body.get("progress_callback")
|
||||
outputs = stream_image_outputs_with_pool(ConversationRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
n=n,
|
||||
size=size,
|
||||
quality=quality,
|
||||
response_format=response_format,
|
||||
base_url=base_url,
|
||||
message_as_error=True,
|
||||
progress_callback=progress_callback,
|
||||
))
|
||||
if body.get("stream"):
|
||||
return stream_image_chunks(outputs)
|
||||
result = collect_image_outputs(outputs)
|
||||
result["usage"] = image_usage(
|
||||
input_text_tokens=count_text_tokens(prompt, model),
|
||||
output_tokens=count_image_output_items_tokens(result.get("data"), size, quality),
|
||||
)
|
||||
return result
|
||||
@@ -0,0 +1,53 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.model_service import model_catalog_service
|
||||
from utils.helper import CODEX_IMAGE_MODEL
|
||||
|
||||
|
||||
def list_models() -> dict[str, Any]:
|
||||
result = model_catalog_service.list_models()
|
||||
data = result.get("data")
|
||||
if not isinstance(data, list):
|
||||
return result
|
||||
seen = {str(item.get("id") or "").strip() for item in data if isinstance(item, dict)}
|
||||
dynamic_models: set[str] = set()
|
||||
accounts = account_service.list_accounts()
|
||||
web_image_accounts = [
|
||||
account
|
||||
for account in accounts
|
||||
if isinstance(account, dict)
|
||||
]
|
||||
codex_types = {
|
||||
normalized
|
||||
for account in accounts
|
||||
if isinstance(account, dict)
|
||||
and account_service._normalize_source_type(account.get("source_type")) == "codex"
|
||||
and (normalized := account_service._normalize_account_type(account.get("type")))
|
||||
}
|
||||
|
||||
if web_image_accounts:
|
||||
dynamic_models.add("gpt-image-2")
|
||||
if codex_types & {"Plus", "Team", "Pro"}:
|
||||
dynamic_models.add(CODEX_IMAGE_MODEL)
|
||||
if "Plus" in codex_types:
|
||||
dynamic_models.add(f"plus-{CODEX_IMAGE_MODEL}")
|
||||
if "Team" in codex_types:
|
||||
dynamic_models.add(f"team-{CODEX_IMAGE_MODEL}")
|
||||
if "Pro" in codex_types:
|
||||
dynamic_models.add(f"pro-{CODEX_IMAGE_MODEL}")
|
||||
|
||||
for model in sorted(dynamic_models):
|
||||
if model not in seen:
|
||||
data.append({
|
||||
"id": model,
|
||||
"object": "model",
|
||||
"created": 0,
|
||||
"owned_by": "chatgpt2api",
|
||||
"permission": [],
|
||||
"root": model,
|
||||
"parent": None,
|
||||
})
|
||||
return result
|
||||
@@ -0,0 +1,453 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Iterable, Iterator
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
from services.protocol.chat_completion_cache import cache_key, chat_completion_cache, normalize_text_messages
|
||||
from services.protocol.conversation import (
|
||||
ConversationRequest,
|
||||
ImageOutput,
|
||||
count_message_image_tokens,
|
||||
count_message_text_tokens,
|
||||
count_text_tokens,
|
||||
encode_images,
|
||||
normalize_messages,
|
||||
stream_image_outputs_with_pool,
|
||||
stream_text_deltas,
|
||||
text_backend,
|
||||
)
|
||||
from services.protocol.web_search_tool import (
|
||||
WEB_SEARCH_TOOL_TYPES,
|
||||
has_unsupported_tools,
|
||||
has_web_search_tool,
|
||||
normalized_sources,
|
||||
run_web_search,
|
||||
search_query_from_messages,
|
||||
text_with_url_citations,
|
||||
)
|
||||
from utils.helper import extract_image_from_message_content, extract_response_prompt, has_response_image_generation_tool
|
||||
from utils.image_tokens import (
|
||||
count_image_content_tokens,
|
||||
count_image_output_items_tokens,
|
||||
image_usage,
|
||||
token_usage,
|
||||
)
|
||||
|
||||
TOOL_UNAVAILABLE_SYSTEM_MESSAGE = (
|
||||
"This compatibility backend cannot execute local tools, shell commands, non-search tools, "
|
||||
"or file operations. Do not claim to have run tools or inspected external resources. "
|
||||
"If a user asks you to use a tool, say that tool execution is unavailable through this backend."
|
||||
)
|
||||
|
||||
RESPONSE_CONTENT_PART_TYPES = {"text", "input_text", "output_text", "image_url", "input_image", "image"}
|
||||
|
||||
|
||||
def normalize_thinking_effort(value: object) -> str:
|
||||
normalized = str(value or "").strip().lower()
|
||||
if normalized in {"", "none"}:
|
||||
return ""
|
||||
if normalized in {"low", "medium", "high"}:
|
||||
return normalized
|
||||
if normalized in {"xhigh", "extended"}:
|
||||
return "extended"
|
||||
return ""
|
||||
|
||||
|
||||
def thinking_effort_from_body(body: dict[str, Any]) -> str:
|
||||
reasoning = body.get("reasoning")
|
||||
if isinstance(reasoning, dict):
|
||||
return normalize_thinking_effort(reasoning.get("effort"))
|
||||
if "thinking_effort" in body:
|
||||
return normalize_thinking_effort(body.get("thinking_effort"))
|
||||
if "reasoning_effort" in body:
|
||||
return normalize_thinking_effort(body.get("reasoning_effort"))
|
||||
return ""
|
||||
|
||||
|
||||
def is_text_response_request(body: dict[str, Any]) -> bool:
|
||||
return not has_response_image_generation_tool(body)
|
||||
|
||||
|
||||
def has_unsupported_response_tools(body: dict[str, Any]) -> bool:
|
||||
return has_unsupported_tools(body, {"image_generation", *WEB_SEARCH_TOOL_TYPES})
|
||||
|
||||
|
||||
def response_image_tool(body: dict[str, Any]) -> dict[str, object]:
|
||||
for tool in body.get("tools") or []:
|
||||
if isinstance(tool, dict) and tool.get("type") == "image_generation":
|
||||
return tool
|
||||
return {}
|
||||
|
||||
|
||||
def extract_response_image(input_value: object) -> tuple[bytes, str] | None:
|
||||
if isinstance(input_value, dict):
|
||||
if str(input_value.get("type") or "").strip() == "input_image":
|
||||
images = extract_image_from_message_content([input_value])
|
||||
return images[0] if images else None
|
||||
images = extract_image_from_message_content(input_value.get("content"))
|
||||
return images[0] if images else None
|
||||
if not isinstance(input_value, list):
|
||||
return None
|
||||
for item in reversed(input_value):
|
||||
if isinstance(item, dict):
|
||||
if str(item.get("type") or "").strip() == "input_image":
|
||||
images = extract_image_from_message_content([item])
|
||||
if images:
|
||||
return images[0]
|
||||
images = extract_image_from_message_content(item.get("content"))
|
||||
if images:
|
||||
return images[0]
|
||||
return None
|
||||
|
||||
|
||||
def _input_image_parts(input_value: object) -> list[dict[str, Any]]:
|
||||
parts: list[dict[str, Any]] = []
|
||||
if isinstance(input_value, dict):
|
||||
content = input_value.get("content")
|
||||
if isinstance(content, list):
|
||||
parts.extend(item for item in content if isinstance(item, dict))
|
||||
return parts
|
||||
if not isinstance(input_value, list):
|
||||
return parts
|
||||
if all(isinstance(item, dict) and item.get("type") for item in input_value):
|
||||
return [item for item in input_value if isinstance(item, dict)]
|
||||
for item in input_value:
|
||||
if isinstance(item, dict):
|
||||
content = item.get("content")
|
||||
if isinstance(content, list):
|
||||
parts.extend(part for part in content if isinstance(part, dict))
|
||||
return parts
|
||||
|
||||
|
||||
def _is_response_content_part(value: object) -> bool:
|
||||
if not isinstance(value, dict):
|
||||
return False
|
||||
part_type = str(value.get("type") or "").strip()
|
||||
return part_type in RESPONSE_CONTENT_PART_TYPES or ("image_url" in value and part_type != "message")
|
||||
|
||||
|
||||
def _message_content_from_response_item(item: dict[str, Any]) -> object:
|
||||
content = item.get("content")
|
||||
if isinstance(content, list):
|
||||
return [dict(part) if isinstance(part, dict) else part for part in content]
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
return extract_response_prompt([item]) or content or ""
|
||||
|
||||
|
||||
def _append_response_message(messages: list[dict[str, Any]], role: object, content: object) -> None:
|
||||
if isinstance(content, str):
|
||||
if content.strip():
|
||||
messages.append({"role": str(role or "user"), "content": content.strip()})
|
||||
return
|
||||
if isinstance(content, list) and content:
|
||||
messages.append({"role": str(role or "user"), "content": content})
|
||||
|
||||
|
||||
def messages_from_input(input_value: object, instructions: object = None) -> list[dict[str, Any]]:
|
||||
messages: list[dict[str, Any]] = []
|
||||
system_text = str(instructions or "").strip()
|
||||
if system_text:
|
||||
messages.append({"role": "system", "content": system_text})
|
||||
if isinstance(input_value, str):
|
||||
if input_value.strip():
|
||||
messages.append({"role": "user", "content": input_value.strip()})
|
||||
return messages
|
||||
if isinstance(input_value, dict):
|
||||
if _is_response_content_part(input_value):
|
||||
_append_response_message(messages, "user", [dict(input_value)])
|
||||
return messages
|
||||
_append_response_message(
|
||||
messages,
|
||||
input_value.get("role") or "user",
|
||||
_message_content_from_response_item(input_value),
|
||||
)
|
||||
return messages
|
||||
if isinstance(input_value, list):
|
||||
if all(_is_response_content_part(item) for item in input_value):
|
||||
_append_response_message(messages, "user", [dict(item) for item in input_value if isinstance(item, dict)])
|
||||
return messages
|
||||
pending_parts: list[dict[str, Any]] = []
|
||||
for item in input_value:
|
||||
if _is_response_content_part(item):
|
||||
pending_parts.append(dict(item))
|
||||
continue
|
||||
if pending_parts:
|
||||
_append_response_message(messages, "user", pending_parts)
|
||||
pending_parts = []
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
_append_response_message(
|
||||
messages,
|
||||
item.get("role") or "user",
|
||||
_message_content_from_response_item(item),
|
||||
)
|
||||
if pending_parts:
|
||||
_append_response_message(messages, "user", pending_parts)
|
||||
return messages
|
||||
|
||||
|
||||
def text_output_item(
|
||||
text: str,
|
||||
item_id: str | None = None,
|
||||
status: str = "completed",
|
||||
annotations: list[dict[str, Any]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"id": item_id or f"msg_{uuid.uuid4().hex}",
|
||||
"type": "message",
|
||||
"status": status,
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": text, "annotations": annotations or []}],
|
||||
}
|
||||
|
||||
|
||||
def web_search_call_item(
|
||||
query: str,
|
||||
item_id: str | None = None,
|
||||
status: str = "completed",
|
||||
sources: list[dict[str, str]] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
action: dict[str, Any] = {
|
||||
"type": "search",
|
||||
"query": query,
|
||||
"queries": [query],
|
||||
}
|
||||
if sources:
|
||||
action["sources"] = [
|
||||
{"type": "url", "url": source["url"]}
|
||||
for source in sources
|
||||
if source.get("url")
|
||||
]
|
||||
return {
|
||||
"id": item_id or f"ws_{uuid.uuid4().hex}",
|
||||
"type": "web_search_call",
|
||||
"status": status,
|
||||
"action": action,
|
||||
}
|
||||
|
||||
|
||||
def image_output_items(prompt: str, data: list[dict[str, Any]], item_id: str | None = None) -> list[dict[str, Any]]:
|
||||
output = []
|
||||
for item in data:
|
||||
b64_json = str(item.get("b64_json") or "").strip()
|
||||
if b64_json:
|
||||
output.append({
|
||||
"id": item_id or f"ig_{len(output) + 1}",
|
||||
"type": "image_generation_call",
|
||||
"status": "completed",
|
||||
"result": b64_json,
|
||||
"revised_prompt": str(item.get("revised_prompt") or prompt).strip() or prompt,
|
||||
})
|
||||
return output
|
||||
|
||||
|
||||
def response_created(response_id: str, model: str, created: int) -> dict[str, Any]:
|
||||
return {
|
||||
"type": "response.created",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": created,
|
||||
"status": "in_progress",
|
||||
"error": None,
|
||||
"incomplete_details": None,
|
||||
"model": model,
|
||||
"output": [],
|
||||
"parallel_tool_calls": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def response_completed(
|
||||
response_id: str,
|
||||
model: str,
|
||||
created: int,
|
||||
output: list[dict[str, Any]],
|
||||
usage: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
response = {
|
||||
"type": "response.completed",
|
||||
"response": {
|
||||
"id": response_id,
|
||||
"object": "response",
|
||||
"created_at": created,
|
||||
"status": "completed",
|
||||
"error": None,
|
||||
"incomplete_details": None,
|
||||
"model": model,
|
||||
"output": output,
|
||||
"parallel_tool_calls": False,
|
||||
},
|
||||
}
|
||||
if usage:
|
||||
response["response"]["usage"] = usage
|
||||
return response
|
||||
|
||||
|
||||
def text_response_parts(body: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]:
|
||||
model = str(body.get("model") or "auto").strip() or "auto"
|
||||
messages = normalize_text_messages(normalize_messages(messages_from_input(body.get("input"), body.get("instructions"))))
|
||||
if has_unsupported_response_tools(body):
|
||||
messages.insert(0, {"role": "system", "content": TOOL_UNAVAILABLE_SYSTEM_MESSAGE})
|
||||
return model, messages
|
||||
|
||||
|
||||
def stream_text_response(backend, body: dict[str, Any], messages: list[dict[str, Any]] | None = None) -> Iterator[dict[str, Any]]:
|
||||
model = str(body.get("model") or "auto").strip() or "auto"
|
||||
messages = messages if messages is not None else messages_from_input(body.get("input"), body.get("instructions"))
|
||||
thinking_effort = thinking_effort_from_body(body)
|
||||
response_id = f"resp_{uuid.uuid4().hex}"
|
||||
item_id = f"msg_{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
full_text = ""
|
||||
yield response_created(response_id, model, created)
|
||||
yield {"type": "response.output_item.added", "output_index": 0, "item": text_output_item("", item_id, "in_progress")}
|
||||
request = ConversationRequest(model=model, messages=messages, thinking_effort=thinking_effort)
|
||||
for delta in stream_text_deltas(backend, request):
|
||||
full_text += delta
|
||||
yield {"type": "response.output_text.delta", "item_id": item_id, "output_index": 0, "content_index": 0, "delta": delta}
|
||||
yield {"type": "response.output_text.done", "item_id": item_id, "output_index": 0, "content_index": 0, "text": full_text}
|
||||
item = text_output_item(full_text, item_id, "completed")
|
||||
yield {"type": "response.output_item.done", "output_index": 0, "item": item}
|
||||
usage = token_usage(
|
||||
input_text_tokens=count_message_text_tokens(messages, model),
|
||||
input_image_tokens=count_message_image_tokens(messages, model),
|
||||
output_text_tokens=count_text_tokens(full_text, model),
|
||||
)
|
||||
yield response_completed(response_id, model, created, [item], usage)
|
||||
|
||||
|
||||
def stream_web_search_response(body: dict[str, Any], messages: list[dict[str, Any]] | None = None) -> Iterator[dict[str, Any]]:
|
||||
model = str(body.get("model") or "auto").strip() or "auto"
|
||||
messages = messages if messages is not None else messages_from_input(body.get("input"), body.get("instructions"))
|
||||
query = search_query_from_messages(messages) or extract_response_prompt(body.get("input"))
|
||||
if not query:
|
||||
raise HTTPException(status_code=400, detail={"error": "input text is required for web_search"})
|
||||
|
||||
response_id = f"resp_{uuid.uuid4().hex}"
|
||||
search_id = f"ws_{uuid.uuid4().hex}"
|
||||
item_id = f"msg_{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
yield response_created(response_id, model, created)
|
||||
|
||||
searching_item = web_search_call_item(query, search_id, "in_progress")
|
||||
yield {"type": "response.output_item.added", "output_index": 0, "item": searching_item}
|
||||
yield {"type": "response.web_search_call.in_progress", "output_index": 0, "item_id": search_id}
|
||||
yield {"type": "response.web_search_call.searching", "output_index": 0, "item_id": search_id}
|
||||
result = run_web_search(query)
|
||||
search_item = web_search_call_item(query, search_id, "completed", normalized_sources(result))
|
||||
yield {"type": "response.web_search_call.completed", "output_index": 0, "item_id": search_id}
|
||||
yield {"type": "response.output_item.done", "output_index": 0, "item": search_item}
|
||||
|
||||
text, annotations = text_with_url_citations(result)
|
||||
message_item = text_output_item("", item_id, "in_progress", annotations)
|
||||
yield {"type": "response.output_item.added", "output_index": 1, "item": message_item}
|
||||
if text:
|
||||
yield {"type": "response.output_text.delta", "item_id": item_id, "output_index": 1, "content_index": 0, "delta": text}
|
||||
yield {"type": "response.output_text.done", "item_id": item_id, "output_index": 1, "content_index": 0, "text": text}
|
||||
message_item = text_output_item(text, item_id, "completed", annotations)
|
||||
yield {"type": "response.output_item.done", "output_index": 1, "item": message_item}
|
||||
usage = token_usage(
|
||||
input_text_tokens=count_message_text_tokens(messages, model),
|
||||
input_image_tokens=count_message_image_tokens(messages, model),
|
||||
output_text_tokens=count_text_tokens(text, model),
|
||||
)
|
||||
yield response_completed(response_id, model, created, [search_item, message_item], usage)
|
||||
|
||||
|
||||
def stream_image_response(
|
||||
image_outputs: Iterable[ImageOutput],
|
||||
prompt: str,
|
||||
model: str,
|
||||
input_image_tokens: int = 0,
|
||||
size: object = None,
|
||||
quality: str = "auto",
|
||||
) -> Iterator[dict[str, Any]]:
|
||||
response_id = f"resp_{uuid.uuid4().hex}"
|
||||
created = int(time.time())
|
||||
yield response_created(response_id, model, created)
|
||||
for output in image_outputs:
|
||||
if output.kind == "message":
|
||||
text = output.text
|
||||
item = text_output_item(text)
|
||||
usage = token_usage(
|
||||
input_text_tokens=count_text_tokens(prompt, model),
|
||||
input_image_tokens=input_image_tokens,
|
||||
output_text_tokens=count_text_tokens(text, model),
|
||||
)
|
||||
yield {"type": "response.output_text.delta", "item_id": item["id"], "output_index": 0, "content_index": 0, "delta": text}
|
||||
yield {"type": "response.output_text.done", "item_id": item["id"], "output_index": 0, "content_index": 0, "text": text}
|
||||
yield {"type": "response.output_item.done", "output_index": 0, "item": item}
|
||||
yield response_completed(response_id, model, created, [item], usage)
|
||||
return
|
||||
if output.kind != "result":
|
||||
continue
|
||||
items = image_output_items(prompt, output.data)
|
||||
if items:
|
||||
usage = image_usage(
|
||||
input_text_tokens=count_text_tokens(prompt, model),
|
||||
input_image_tokens=input_image_tokens,
|
||||
output_tokens=count_image_output_items_tokens(output.data, size, quality),
|
||||
)
|
||||
for output_index, item in enumerate(items):
|
||||
yield {"type": "response.output_item.done", "output_index": output_index, "item": item}
|
||||
yield response_completed(response_id, model, created, items, usage)
|
||||
return
|
||||
raise RuntimeError("image generation failed")
|
||||
|
||||
|
||||
def collect_response(events: Iterable[dict[str, Any]]) -> dict[str, Any]:
|
||||
completed = {}
|
||||
for event in events:
|
||||
if event.get("type") == "response.completed":
|
||||
completed = event.get("response") if isinstance(event.get("response"), dict) else {}
|
||||
if not completed:
|
||||
raise RuntimeError("response generation failed")
|
||||
return completed
|
||||
|
||||
|
||||
def response_events(body: dict[str, Any]) -> Iterator[dict[str, Any]]:
|
||||
if is_text_response_request(body):
|
||||
model, messages = text_response_parts(body)
|
||||
if has_web_search_tool(body) and not has_unsupported_response_tools(body):
|
||||
yield from stream_web_search_response(body, messages)
|
||||
return
|
||||
key = cache_key(body, messages, stream=bool(body.get("stream")))
|
||||
yield from chat_completion_cache.get_or_compute_stream(
|
||||
key,
|
||||
lambda: stream_text_response(text_backend(model), body, messages),
|
||||
)
|
||||
return
|
||||
|
||||
prompt = extract_response_prompt(body.get("input"))
|
||||
if not prompt:
|
||||
raise HTTPException(status_code=400, detail={"error": "input text is required"})
|
||||
model = str(body.get("model") or "gpt-image-2").strip() or "gpt-image-2"
|
||||
image_info = extract_response_image(body.get("input"))
|
||||
if image_info:
|
||||
image_data, mime_type = image_info
|
||||
images = encode_images([(image_data, "image.png", mime_type)])
|
||||
else:
|
||||
images = None
|
||||
input_image_tokens = count_image_content_tokens(_input_image_parts(body.get("input")), model)
|
||||
tool = response_image_tool(body)
|
||||
image_outputs = stream_image_outputs_with_pool(ConversationRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
size=tool.get("size"),
|
||||
quality=str(tool.get("quality") or "auto"),
|
||||
response_format="b64_json",
|
||||
images=images,
|
||||
))
|
||||
yield from stream_image_response(image_outputs, prompt, model, input_image_tokens, tool.get("size"), str(tool.get("quality") or "auto"))
|
||||
|
||||
|
||||
def handle(body: dict[str, Any]) -> dict[str, Any] | Iterator[dict[str, Any]]:
|
||||
events = response_events(body)
|
||||
if body.get("stream"):
|
||||
return events
|
||||
return collect_response(events)
|
||||
@@ -0,0 +1,160 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.openai_backend_api import OpenAIBackendAPI
|
||||
|
||||
WEB_SEARCH_TOOL_TYPES = {"web_search", "web_search_preview", "web_search_preview_2025_03_11"}
|
||||
SEARCH_CHAT_MODEL_PREFIXES = (
|
||||
"gpt-4o-search-preview",
|
||||
"gpt-4o-mini-search-preview",
|
||||
"gpt-5-search-api",
|
||||
)
|
||||
|
||||
|
||||
def _tool_type(tool: object) -> str:
|
||||
return str(tool.get("type") or "").strip() if isinstance(tool, dict) else ""
|
||||
|
||||
|
||||
def has_web_search_tool(body: dict[str, Any]) -> bool:
|
||||
tools = body.get("tools")
|
||||
if isinstance(tools, list):
|
||||
return any(_tool_type(tool) in WEB_SEARCH_TOOL_TYPES for tool in tools)
|
||||
tool_choice = body.get("tool_choice")
|
||||
return _tool_type(tool_choice) in WEB_SEARCH_TOOL_TYPES
|
||||
|
||||
|
||||
def is_web_search_chat_request(body: dict[str, Any]) -> bool:
|
||||
model = str(body.get("model") or "").strip()
|
||||
return (
|
||||
has_web_search_tool(body)
|
||||
or isinstance(body.get("web_search_options"), dict)
|
||||
or any(
|
||||
model == prefix or model.startswith(f"{prefix}-")
|
||||
for prefix in SEARCH_CHAT_MODEL_PREFIXES
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def has_unsupported_tools(body: dict[str, Any], allowed_types: set[str]) -> bool:
|
||||
tools = body.get("tools")
|
||||
if not isinstance(tools, list):
|
||||
return False
|
||||
return any(_tool_type(tool) not in allowed_types for tool in tools if isinstance(tool, dict))
|
||||
|
||||
|
||||
def message_text(value: object) -> str:
|
||||
if isinstance(value, str):
|
||||
return value.strip()
|
||||
if isinstance(value, list):
|
||||
parts: list[str] = []
|
||||
for item in value:
|
||||
if isinstance(item, str):
|
||||
text = item.strip()
|
||||
elif isinstance(item, dict):
|
||||
text = str(item.get("text") or item.get("input_text") or "").strip()
|
||||
else:
|
||||
text = ""
|
||||
if text:
|
||||
parts.append(text)
|
||||
return "\n".join(parts).strip()
|
||||
return ""
|
||||
|
||||
|
||||
def search_query_from_messages(messages: list[dict[str, Any]]) -> str:
|
||||
for message in reversed(messages):
|
||||
if str(message.get("role") or "").strip().lower() != "user":
|
||||
continue
|
||||
text = message_text(message.get("content"))
|
||||
if text:
|
||||
return text
|
||||
return ""
|
||||
|
||||
|
||||
def _readable_annotation_part(parts: list[str]) -> str:
|
||||
for part in parts:
|
||||
value = part.strip()
|
||||
lower = value.lower()
|
||||
if value and not (
|
||||
lower.startswith(("turn", "source", "sources"))
|
||||
or re.fullmatch(r"\d+", value)
|
||||
):
|
||||
return value
|
||||
return ""
|
||||
|
||||
|
||||
def clean_search_text(text: str) -> str:
|
||||
def replace_annotation(match: re.Match[str]) -> str:
|
||||
parts = [part.strip() for part in match.group(1).split("\ue202")]
|
||||
kind = (parts[0] if parts else "").lower()
|
||||
data = parts[1:]
|
||||
if kind == "url":
|
||||
label = data[0] if data else ""
|
||||
url = data[1] if len(data) > 1 else ""
|
||||
if label and url.startswith(("http://", "https://")):
|
||||
return f"{label} ({url})"
|
||||
return label or url
|
||||
if kind == "cite":
|
||||
return _readable_annotation_part(data)
|
||||
return _readable_annotation_part(data)
|
||||
|
||||
text = re.sub(r"\ue200([^\ue201]*)\ue201", replace_annotation, text)
|
||||
text = re.sub(r"\ue200[^\ue201]*$", "", text)
|
||||
return re.sub(r"\s+([.,;:!?])", r"\1", text).strip()
|
||||
|
||||
|
||||
def normalized_sources(result: dict[str, Any]) -> list[dict[str, str]]:
|
||||
sources = result.get("sources")
|
||||
if not isinstance(sources, list):
|
||||
return []
|
||||
output: list[dict[str, str]] = []
|
||||
seen: set[str] = set()
|
||||
for item in sources:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
url = str(item.get("url") or "").strip()
|
||||
title = str(item.get("title") or "").strip()
|
||||
snippet = str(item.get("snippet") or "").strip()
|
||||
if not url or url in seen:
|
||||
continue
|
||||
seen.add(url)
|
||||
output.append({"title": title, "url": url, "snippet": snippet})
|
||||
return output
|
||||
|
||||
|
||||
def text_with_url_citations(result: dict[str, Any]) -> tuple[str, list[dict[str, Any]]]:
|
||||
text = clean_search_text(str(result.get("answer") or ""))
|
||||
annotations: list[dict[str, Any]] = []
|
||||
sources = normalized_sources(result)
|
||||
if sources:
|
||||
text = text.rstrip()
|
||||
if text:
|
||||
text += "\n\n"
|
||||
text += "Sources:\n"
|
||||
for index, source in enumerate(sources, start=1):
|
||||
title = source["title"] or source["url"]
|
||||
line_prefix = f"{index}. {title}"
|
||||
text += line_prefix
|
||||
if source["url"]:
|
||||
if source["title"]:
|
||||
text += " - "
|
||||
start = len(text)
|
||||
text += source["url"]
|
||||
annotations.append({
|
||||
"type": "url_citation",
|
||||
"start_index": start,
|
||||
"end_index": len(text),
|
||||
"url": source["url"],
|
||||
"title": source["title"] or source["url"],
|
||||
})
|
||||
text += "\n"
|
||||
return text.strip(), annotations
|
||||
|
||||
|
||||
def run_web_search(query: str) -> dict[str, Any]:
|
||||
token = account_service.get_text_access_token()
|
||||
result = OpenAIBackendAPI(token).search(query)
|
||||
account_service.mark_text_used(token)
|
||||
return result
|
||||
@@ -0,0 +1,648 @@
|
||||
"""Global outbound proxy and Cloudflare clearance helpers."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field, replace
|
||||
import json
|
||||
import re
|
||||
import threading
|
||||
import time
|
||||
from typing import Callable, Mapping
|
||||
from urllib import request as urllib_request
|
||||
from urllib.parse import quote, urlparse
|
||||
|
||||
from curl_cffi.requests import Session
|
||||
|
||||
from services.config import config
|
||||
|
||||
|
||||
FlareSolverrRequestMethod = Callable[[str, bytes, dict[str, str], float], bytes]
|
||||
|
||||
|
||||
def normalize_proxy_url(url: str) -> str:
|
||||
"""Normalize proxy URLs for curl_cffi.
|
||||
|
||||
SOCKS proxies should use remote-DNS resolution by default, so generic
|
||||
``socks://`` and ``socks5://`` inputs are upgraded to ``socks5h://``.
|
||||
HTTP/HTTPS/socks5h inputs are otherwise left untouched except trimming.
|
||||
"""
|
||||
candidate = str(url or "").strip()
|
||||
if candidate and "://" not in candidate:
|
||||
candidate = _colon_proxy_to_url(candidate)
|
||||
lowered = candidate.lower()
|
||||
if lowered.startswith("socks://"):
|
||||
return "socks5h://" + candidate[len("socks://") :]
|
||||
if lowered.startswith("socks5://"):
|
||||
return "socks5h://" + candidate[len("socks5://") :]
|
||||
return candidate
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ProxyRuntimeProfile:
|
||||
proxy_url: str = ""
|
||||
proxy_source: str = "direct"
|
||||
resource: bool = False
|
||||
runtime_enabled: bool = False
|
||||
egress_mode: str = "direct"
|
||||
skip_ssl_verify: bool = False
|
||||
reset_session_status_codes: tuple[int, ...] = field(default_factory=lambda: (403,))
|
||||
clearance: dict[str, object] = field(default_factory=dict, repr=False)
|
||||
|
||||
@property
|
||||
def clearance_enabled(self) -> bool:
|
||||
return (
|
||||
self.runtime_enabled
|
||||
and bool(self.clearance.get("enabled"))
|
||||
and self.clearance_mode in {"manual", "flaresolverr"}
|
||||
)
|
||||
|
||||
@property
|
||||
def clearance_mode(self) -> str:
|
||||
return str(self.clearance.get("mode") or "none").strip().lower()
|
||||
|
||||
@property
|
||||
def refresh_interval(self) -> int:
|
||||
try:
|
||||
return max(0, int(self.clearance.get("refresh_interval") or 0))
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
@property
|
||||
def timeout_sec(self) -> int:
|
||||
try:
|
||||
return max(1, int(self.clearance.get("timeout_sec") or 60))
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
return 60
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClearanceBundle:
|
||||
target_host: str
|
||||
proxy_url: str = ""
|
||||
cookies: dict[str, str] = field(default_factory=dict, repr=False)
|
||||
user_agent: str = ""
|
||||
created_at: float = field(default_factory=time.time)
|
||||
expires_at: float | None = None
|
||||
|
||||
def is_valid_for(self, target_host: str, proxy_url: str, *, now: float | None = None) -> bool:
|
||||
host = _normalize_host(target_host)
|
||||
if self.target_host and host and _normalize_host(self.target_host) != host:
|
||||
return False
|
||||
if normalize_proxy_url(self.proxy_url) != normalize_proxy_url(proxy_url):
|
||||
return False
|
||||
if self.expires_at is not None and (time.time() if now is None else now) >= self.expires_at:
|
||||
return False
|
||||
return bool(self.cookies or self.user_agent)
|
||||
|
||||
def cookie_header(self) -> str:
|
||||
return _cookies_to_header(self.cookies)
|
||||
|
||||
|
||||
class FlareSolverrClearanceProvider:
|
||||
def __init__(self, flaresolverr_url: str, request_method: FlareSolverrRequestMethod | None = None) -> None:
|
||||
self.flaresolverr_url = str(flaresolverr_url or "").strip().rstrip("/")
|
||||
self._request_method = request_method or self._urllib_post
|
||||
|
||||
def get_clearance(self, target_url: str, proxy_url: str = "", timeout_sec: int = 60) -> ClearanceBundle | None:
|
||||
if not self.flaresolverr_url:
|
||||
return None
|
||||
|
||||
timeout = _coerce_timeout(timeout_sec)
|
||||
payload: dict[str, object] = {
|
||||
"cmd": "request.get",
|
||||
"url": str(target_url or ""),
|
||||
"maxTimeout": int(timeout * 1000),
|
||||
}
|
||||
proxy_url = normalize_proxy_url(proxy_url)
|
||||
if proxy_url:
|
||||
payload["proxy"] = {"url": proxy_url}
|
||||
|
||||
endpoint = f"{self.flaresolverr_url}/v1"
|
||||
try:
|
||||
body = json.dumps(payload).encode("utf-8")
|
||||
raw_response = self._request_method(
|
||||
endpoint,
|
||||
body,
|
||||
{"Content-Type": "application/json"},
|
||||
timeout,
|
||||
)
|
||||
data = json.loads(raw_response.decode("utf-8") if isinstance(raw_response, bytes) else raw_response)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
if not isinstance(data, dict) or str(data.get("status") or "").lower() != "ok":
|
||||
return None
|
||||
solution = data.get("solution")
|
||||
if not isinstance(solution, dict):
|
||||
return None
|
||||
|
||||
target_host = _host_from_url(target_url)
|
||||
cookies = _filter_flaresolverr_cookies(solution.get("cookies"), target_host)
|
||||
user_agent = str(solution.get("userAgent") or "").strip()
|
||||
if not cookies and not user_agent:
|
||||
return None
|
||||
return ClearanceBundle(
|
||||
target_host=target_host,
|
||||
proxy_url=proxy_url,
|
||||
cookies=cookies,
|
||||
user_agent=user_agent,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _urllib_post(endpoint: str, body: bytes, headers: dict[str, str], timeout: float) -> bytes:
|
||||
req = urllib_request.Request(endpoint, data=body, headers=headers, method="POST")
|
||||
with urllib_request.urlopen(req, timeout=timeout) as response:
|
||||
return response.read()
|
||||
|
||||
|
||||
class ProxySettingsStore:
|
||||
def __init__(
|
||||
self,
|
||||
config_store=None,
|
||||
clearance_provider_factory: Callable[[str], FlareSolverrClearanceProvider] | None = None,
|
||||
) -> None:
|
||||
self._config = config_store or config
|
||||
self._clearance_provider_factory = clearance_provider_factory or FlareSolverrClearanceProvider
|
||||
self._clearance_cache: dict[tuple[str, str], ClearanceBundle] = {}
|
||||
self._provider_cache: dict[str, FlareSolverrClearanceProvider] = {}
|
||||
self._flight_locks: dict[tuple[str, str], threading.Lock] = {}
|
||||
self._lock = threading.RLock()
|
||||
|
||||
def get_profile(
|
||||
self,
|
||||
account: dict | None = None,
|
||||
proxy: str = "",
|
||||
resource: bool = False,
|
||||
upstream: bool = False,
|
||||
) -> ProxyRuntimeProfile:
|
||||
runtime = self._get_runtime_settings()
|
||||
clearance = dict(runtime.get("clearance") if isinstance(runtime.get("clearance"), dict) else {})
|
||||
runtime_enabled = bool(runtime.get("enabled"))
|
||||
egress_mode = str(runtime.get("egress_mode") or "direct").strip().lower()
|
||||
|
||||
account_proxy = _clean((account or {}).get("proxy") if isinstance(account, dict) else "")
|
||||
explicit_proxy = _clean(proxy)
|
||||
legacy_proxy = _clean(self._config.get_proxy_settings())
|
||||
|
||||
runtime_proxy = ""
|
||||
runtime_proxy_source = "runtime"
|
||||
if upstream and runtime_enabled and egress_mode == "single_proxy":
|
||||
resource_proxy = _clean(runtime.get("resource_proxy_url")) if resource else ""
|
||||
runtime_proxy = resource_proxy or _clean(runtime.get("proxy_url"))
|
||||
runtime_proxy_source = "runtime_resource" if resource_proxy else "runtime"
|
||||
|
||||
selected_proxy = ""
|
||||
source = "direct"
|
||||
if account_proxy:
|
||||
selected_proxy = account_proxy
|
||||
source = "account"
|
||||
elif runtime_proxy:
|
||||
selected_proxy = runtime_proxy
|
||||
source = runtime_proxy_source
|
||||
elif explicit_proxy:
|
||||
selected_proxy = explicit_proxy
|
||||
source = "explicit"
|
||||
elif legacy_proxy:
|
||||
selected_proxy = legacy_proxy
|
||||
source = "global"
|
||||
|
||||
return ProxyRuntimeProfile(
|
||||
proxy_url=normalize_proxy_url(selected_proxy),
|
||||
proxy_source=source,
|
||||
resource=bool(resource),
|
||||
runtime_enabled=runtime_enabled,
|
||||
egress_mode=egress_mode,
|
||||
skip_ssl_verify=bool(runtime.get("skip_ssl_verify")) if runtime_enabled else False,
|
||||
reset_session_status_codes=_status_codes_tuple(runtime.get("reset_session_status_codes")),
|
||||
clearance=clearance,
|
||||
)
|
||||
|
||||
def build_session_kwargs(
|
||||
self,
|
||||
account: dict | None = None,
|
||||
proxy: str = "",
|
||||
resource: bool = False,
|
||||
upstream: bool = False,
|
||||
**session_kwargs,
|
||||
) -> dict[str, object]:
|
||||
profile = self.get_profile(account=account, proxy=proxy, resource=resource, upstream=upstream)
|
||||
if profile.proxy_url:
|
||||
session_kwargs["proxy"] = profile.proxy_url
|
||||
if profile.runtime_enabled and profile.skip_ssl_verify:
|
||||
session_kwargs["verify"] = False
|
||||
return session_kwargs
|
||||
|
||||
def build_headers(
|
||||
self,
|
||||
headers: Mapping[str, object] | None = None,
|
||||
target_url: str = "https://chatgpt.com",
|
||||
account: dict | None = None,
|
||||
proxy: str = "",
|
||||
resource: bool = False,
|
||||
upstream: bool = True,
|
||||
) -> dict[str, object]:
|
||||
merged_headers: dict[str, object] = dict(headers or {})
|
||||
profile = self.get_profile(account=account, proxy=proxy, resource=resource, upstream=upstream)
|
||||
if not profile.clearance_enabled:
|
||||
return merged_headers
|
||||
|
||||
target_host = _host_from_url(target_url)
|
||||
bundle = self._bundle_for_headers(profile, target_host)
|
||||
if bundle is None or not bundle.is_valid_for(target_host, profile.proxy_url):
|
||||
return merged_headers
|
||||
|
||||
if bundle.user_agent and _find_header_key(merged_headers, "user-agent") is None:
|
||||
merged_headers["User-Agent"] = bundle.user_agent
|
||||
|
||||
if bundle.cookies:
|
||||
cookie_key = _find_header_key(merged_headers, "cookie") or "Cookie"
|
||||
existing_cookie = str(merged_headers.get(cookie_key) or "")
|
||||
cookie_header = _merge_cookie_header(existing_cookie, bundle.cookies)
|
||||
if cookie_header:
|
||||
merged_headers[cookie_key] = cookie_header
|
||||
return merged_headers
|
||||
|
||||
def refresh_clearance(
|
||||
self,
|
||||
target_url: str = "https://chatgpt.com",
|
||||
account: dict | None = None,
|
||||
proxy: str = "",
|
||||
resource: bool = False,
|
||||
force: bool = False,
|
||||
upstream: bool = True,
|
||||
) -> ClearanceBundle | None:
|
||||
profile = self.get_profile(account=account, proxy=proxy, resource=resource, upstream=upstream)
|
||||
if not profile.clearance_enabled:
|
||||
return None
|
||||
|
||||
target_host = _host_from_url(target_url)
|
||||
key = self._cache_key(profile.proxy_url, target_host)
|
||||
if profile.clearance_mode == "manual":
|
||||
bundle = self._build_manual_bundle(profile, target_host)
|
||||
if bundle is not None:
|
||||
self._set_cached_bundle(key, bundle)
|
||||
return bundle
|
||||
if profile.clearance_mode != "flaresolverr":
|
||||
return None
|
||||
|
||||
cached_before = self._get_cached_bundle(key)
|
||||
if cached_before is not None and not force and cached_before.is_valid_for(target_host, profile.proxy_url):
|
||||
return cached_before
|
||||
|
||||
lock = self._get_flight_lock(key)
|
||||
if not lock.acquire(blocking=False):
|
||||
with lock:
|
||||
pass
|
||||
return self._get_cached_bundle(key) or cached_before
|
||||
|
||||
try:
|
||||
cached_now = self._get_cached_bundle(key)
|
||||
if cached_now is not None and not force and cached_now.is_valid_for(target_host, profile.proxy_url):
|
||||
return cached_now
|
||||
|
||||
flaresolverr_url = str(profile.clearance.get("flaresolverr_url") or "").strip()
|
||||
provider = self._get_provider(flaresolverr_url)
|
||||
new_bundle = provider.get_clearance(target_url, proxy_url=profile.proxy_url, timeout_sec=profile.timeout_sec)
|
||||
if new_bundle is not None:
|
||||
expires_at = time.time() + profile.refresh_interval if profile.refresh_interval else None
|
||||
if (
|
||||
not new_bundle.target_host
|
||||
or normalize_proxy_url(new_bundle.proxy_url) != normalize_proxy_url(profile.proxy_url)
|
||||
or new_bundle.expires_at != expires_at
|
||||
):
|
||||
new_bundle = replace(
|
||||
new_bundle,
|
||||
target_host=new_bundle.target_host or target_host,
|
||||
proxy_url=profile.proxy_url,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
self._set_cached_bundle(key, new_bundle)
|
||||
return new_bundle
|
||||
return cached_now or cached_before
|
||||
finally:
|
||||
lock.release()
|
||||
|
||||
def invalidate_clearance(
|
||||
self,
|
||||
target_url: str = "https://chatgpt.com",
|
||||
account: dict | None = None,
|
||||
proxy: str = "",
|
||||
resource: bool = False,
|
||||
upstream: bool = True,
|
||||
) -> None:
|
||||
profile = self.get_profile(account=account, proxy=proxy, resource=resource, upstream=upstream)
|
||||
target_host = _host_from_url(target_url)
|
||||
key = self._cache_key(profile.proxy_url, target_host)
|
||||
with self._lock:
|
||||
self._clearance_cache.pop(key, None)
|
||||
|
||||
def get_runtime_status(self) -> dict[str, object]:
|
||||
profile = self.get_profile(upstream=True)
|
||||
with self._lock:
|
||||
cached_hosts = [host for _proxy, host in self._clearance_cache]
|
||||
cached_count = len(self._clearance_cache)
|
||||
return {
|
||||
"enabled": profile.runtime_enabled,
|
||||
"egress_mode": profile.egress_mode,
|
||||
"proxy_source": profile.proxy_source,
|
||||
"has_proxy": bool(profile.proxy_url),
|
||||
"clearance_enabled": profile.clearance_enabled,
|
||||
"clearance_mode": profile.clearance_mode,
|
||||
"has_clearance_bundle": cached_count > 0,
|
||||
"cached_clearance_hosts": sorted(set(cached_hosts)),
|
||||
}
|
||||
|
||||
def _get_runtime_settings(self) -> dict[str, object]:
|
||||
try:
|
||||
runtime = self._config.get_proxy_runtime_settings()
|
||||
except AttributeError:
|
||||
runtime = {}
|
||||
return runtime if isinstance(runtime, dict) else {}
|
||||
|
||||
def _bundle_for_headers(self, profile: ProxyRuntimeProfile, target_host: str) -> ClearanceBundle | None:
|
||||
key = self._cache_key(profile.proxy_url, target_host)
|
||||
if profile.clearance_mode == "manual":
|
||||
bundle = self._build_manual_bundle(profile, target_host)
|
||||
if bundle is not None:
|
||||
self._set_cached_bundle(key, bundle)
|
||||
return bundle
|
||||
if profile.clearance_mode == "flaresolverr":
|
||||
return self._get_cached_bundle(key)
|
||||
return None
|
||||
|
||||
def _build_manual_bundle(self, profile: ProxyRuntimeProfile, target_host: str) -> ClearanceBundle | None:
|
||||
cookies = _parse_cookie_header(str(profile.clearance.get("cf_cookies") or ""))
|
||||
cf_clearance = str(profile.clearance.get("cf_clearance") or "").strip()
|
||||
if cf_clearance and "cf_clearance" not in cookies:
|
||||
cookies["cf_clearance"] = cf_clearance
|
||||
user_agent = str(profile.clearance.get("user_agent") or "").strip()
|
||||
if not cookies and not user_agent:
|
||||
return None
|
||||
|
||||
now = time.time()
|
||||
expires_at = now + profile.refresh_interval if profile.refresh_interval else None
|
||||
return ClearanceBundle(
|
||||
target_host=target_host,
|
||||
proxy_url=profile.proxy_url,
|
||||
cookies=cookies,
|
||||
user_agent=user_agent,
|
||||
created_at=now,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
|
||||
def _get_provider(self, flaresolverr_url: str) -> FlareSolverrClearanceProvider:
|
||||
url = str(flaresolverr_url or "").strip().rstrip("/")
|
||||
with self._lock:
|
||||
provider = self._provider_cache.get(url)
|
||||
if provider is None:
|
||||
provider = self._clearance_provider_factory(url)
|
||||
self._provider_cache[url] = provider
|
||||
return provider
|
||||
|
||||
def _get_flight_lock(self, key: tuple[str, str]) -> threading.Lock:
|
||||
with self._lock:
|
||||
lock = self._flight_locks.get(key)
|
||||
if lock is None:
|
||||
lock = threading.Lock()
|
||||
self._flight_locks[key] = lock
|
||||
return lock
|
||||
|
||||
def _get_cached_bundle(self, key: tuple[str, str]) -> ClearanceBundle | None:
|
||||
with self._lock:
|
||||
return self._clearance_cache.get(key)
|
||||
|
||||
def _set_cached_bundle(self, key: tuple[str, str], bundle: ClearanceBundle) -> None:
|
||||
with self._lock:
|
||||
self._clearance_cache[key] = bundle
|
||||
|
||||
@staticmethod
|
||||
def _cache_key(proxy_url: str, target_host: str) -> tuple[str, str]:
|
||||
return (normalize_proxy_url(proxy_url), _normalize_host(target_host))
|
||||
|
||||
|
||||
def _clean(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def _colon_proxy_to_url(url: str) -> str:
|
||||
parts = url.split(":", 3)
|
||||
if len(parts) == 4 and parts[1].isdigit():
|
||||
host, port, username, password = parts
|
||||
return f"http://{quote(username, safe='')}:{quote(password, safe='')}@{host}:{port}"
|
||||
if len(parts) == 2 and parts[1].isdigit():
|
||||
return f"http://{url}"
|
||||
return url
|
||||
|
||||
|
||||
def _normalize_host(host: str) -> str:
|
||||
return str(host or "").strip().strip(".").lower()
|
||||
|
||||
|
||||
def _host_from_url(url: str) -> str:
|
||||
candidate = str(url or "").strip()
|
||||
parsed = urlparse(candidate)
|
||||
if not parsed.hostname and candidate and "://" not in candidate:
|
||||
parsed = urlparse(f"https://{candidate}")
|
||||
return _normalize_host(parsed.hostname or "")
|
||||
|
||||
|
||||
def _status_codes_tuple(value: object) -> tuple[int, ...]:
|
||||
source = value if isinstance(value, list) else [403]
|
||||
codes: list[int] = []
|
||||
for item in source:
|
||||
if isinstance(item, bool):
|
||||
continue
|
||||
try:
|
||||
code = int(item)
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
continue
|
||||
if 100 <= code <= 599 and code not in codes:
|
||||
codes.append(code)
|
||||
return tuple(codes or [403])
|
||||
|
||||
|
||||
def _coerce_timeout(value: object) -> float:
|
||||
try:
|
||||
timeout = float(value)
|
||||
except (OverflowError, TypeError, ValueError):
|
||||
timeout = 60.0
|
||||
return max(1.0, timeout)
|
||||
|
||||
|
||||
def _is_valid_proxy_url(url: str) -> bool:
|
||||
parsed = urlparse(normalize_proxy_url(url))
|
||||
return parsed.scheme in {"http", "https", "socks5", "socks5h"} and bool(parsed.netloc)
|
||||
|
||||
|
||||
def _domain_matches(host: str, domain: str) -> bool:
|
||||
normalized_host = _normalize_host(host)
|
||||
normalized_domain = _normalize_host(domain.lstrip("."))
|
||||
if not normalized_domain:
|
||||
return True
|
||||
return normalized_host == normalized_domain or normalized_host.endswith(f".{normalized_domain}")
|
||||
|
||||
|
||||
def _filter_flaresolverr_cookies(raw_cookies: object, target_host: str) -> dict[str, str]:
|
||||
if not isinstance(raw_cookies, list):
|
||||
return {}
|
||||
|
||||
filtered_cookies: dict[str, str] = {}
|
||||
for item in raw_cookies:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
name = str(item.get("name") or "").strip()
|
||||
if not name:
|
||||
continue
|
||||
value = str(item.get("value") or "")
|
||||
domain = str(item.get("domain") or "").strip()
|
||||
if not domain or _domain_matches(target_host, domain):
|
||||
filtered_cookies[name] = value
|
||||
return filtered_cookies
|
||||
|
||||
|
||||
def _parse_cookie_header(header: str) -> dict[str, str]:
|
||||
cookies: dict[str, str] = {}
|
||||
for part in str(header or "").split(";"):
|
||||
name, sep, value = part.strip().partition("=")
|
||||
if sep and name:
|
||||
cookies[name.strip()] = value.strip()
|
||||
return cookies
|
||||
|
||||
|
||||
def _cookies_to_header(cookies: Mapping[str, str]) -> str:
|
||||
return "; ".join(f"{name}={value}" for name, value in cookies.items() if name)
|
||||
|
||||
|
||||
def _merge_cookie_header(existing_header: str, cookies: Mapping[str, str]) -> str:
|
||||
existing = str(existing_header or "").strip()
|
||||
existing_names = set(_parse_cookie_header(existing).keys())
|
||||
additions = [f"{name}={value}" for name, value in cookies.items() if name and name not in existing_names]
|
||||
if existing and additions:
|
||||
return existing.rstrip("; ") + "; " + "; ".join(additions)
|
||||
if existing:
|
||||
return existing
|
||||
return "; ".join(additions)
|
||||
|
||||
|
||||
def _find_header_key(headers: Mapping[str, object], name: str) -> str | None:
|
||||
target = name.lower()
|
||||
for key in headers:
|
||||
if str(key).lower() == target:
|
||||
return str(key)
|
||||
return None
|
||||
|
||||
|
||||
def _redact_url_credentials(text: str) -> str:
|
||||
return re.sub(
|
||||
r"((?:https?|socks5h?|socks)://)([^\s/@:]+):([^\s/@]+)@",
|
||||
r"\1[REDACTED]@",
|
||||
str(text or ""),
|
||||
flags=re.IGNORECASE,
|
||||
)
|
||||
|
||||
|
||||
def test_proxy(url: str = "", *, timeout: float = 15.0) -> dict:
|
||||
candidate = normalize_proxy_url(_clean(url))
|
||||
proxy_source = "input"
|
||||
if not candidate:
|
||||
profile = proxy_settings.get_profile(upstream=True)
|
||||
candidate = profile.proxy_url
|
||||
proxy_source = profile.proxy_source
|
||||
result_base = {"proxy_source": proxy_source, "has_proxy": bool(candidate)}
|
||||
if not candidate:
|
||||
return {
|
||||
"ok": False,
|
||||
"status": 0,
|
||||
"latency_ms": 0,
|
||||
"error": "no active proxy configured",
|
||||
**result_base,
|
||||
}
|
||||
if not _is_valid_proxy_url(candidate):
|
||||
return {
|
||||
"ok": False,
|
||||
"status": 0,
|
||||
"latency_ms": 0,
|
||||
"error": "invalid proxy url",
|
||||
**result_base,
|
||||
}
|
||||
session = Session(impersonate="edge101", verify=True, proxy=candidate)
|
||||
started = time.perf_counter()
|
||||
try:
|
||||
response = session.get(
|
||||
"https://chatgpt.com/api/auth/csrf",
|
||||
headers={"user-agent": "Mozilla/5.0 (chatgpt2api proxy test)"},
|
||||
timeout=timeout,
|
||||
)
|
||||
latency_ms = int((time.perf_counter() - started) * 1000)
|
||||
return {
|
||||
"ok": response.status_code < 500,
|
||||
"status": int(response.status_code),
|
||||
"latency_ms": latency_ms,
|
||||
"error": None if response.status_code < 500 else f"HTTP {response.status_code}",
|
||||
**result_base,
|
||||
}
|
||||
except Exception as exc:
|
||||
latency_ms = int((time.perf_counter() - started) * 1000)
|
||||
return {
|
||||
"ok": False,
|
||||
"status": 0,
|
||||
"latency_ms": latency_ms,
|
||||
"error": _redact_url_credentials(str(exc) or exc.__class__.__name__),
|
||||
**result_base,
|
||||
}
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
|
||||
def test_clearance(target_url: str = "https://chatgpt.com") -> dict:
|
||||
target_url = str(target_url or "https://chatgpt.com").strip() or "https://chatgpt.com"
|
||||
started = time.perf_counter()
|
||||
status = proxy_settings.get_runtime_status()
|
||||
if not status.get("clearance_enabled"):
|
||||
return {
|
||||
"ok": False,
|
||||
"status": "disabled",
|
||||
"latency_ms": 0,
|
||||
"has_cookies": False,
|
||||
"user_agent": "",
|
||||
"error": "clearance is disabled",
|
||||
"runtime": status,
|
||||
}
|
||||
try:
|
||||
bundle = proxy_settings.refresh_clearance(target_url=target_url, force=True, upstream=True)
|
||||
except Exception as exc:
|
||||
latency_ms = int((time.perf_counter() - started) * 1000)
|
||||
return {
|
||||
"ok": False,
|
||||
"status": "error",
|
||||
"latency_ms": latency_ms,
|
||||
"has_cookies": False,
|
||||
"user_agent": "",
|
||||
"error": _redact_url_credentials(str(exc) or exc.__class__.__name__),
|
||||
"runtime": proxy_settings.get_runtime_status(),
|
||||
}
|
||||
|
||||
latency_ms = int((time.perf_counter() - started) * 1000)
|
||||
runtime = proxy_settings.get_runtime_status()
|
||||
if bundle is None:
|
||||
return {
|
||||
"ok": False,
|
||||
"status": "failed",
|
||||
"latency_ms": latency_ms,
|
||||
"has_cookies": False,
|
||||
"user_agent": "",
|
||||
"error": "clearance refresh returned no bundle",
|
||||
"runtime": runtime,
|
||||
}
|
||||
return {
|
||||
"ok": True,
|
||||
"status": "ok",
|
||||
"latency_ms": latency_ms,
|
||||
"has_cookies": bool(bundle.cookies),
|
||||
"user_agent": bundle.user_agent or "",
|
||||
"error": None,
|
||||
"runtime": runtime,
|
||||
}
|
||||
|
||||
|
||||
proxy_settings = ProxySettingsStore()
|
||||
@@ -0,0 +1,5 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from services.storage.factory import create_storage_backend
|
||||
|
||||
__all__ = ["create_storage_backend"]
|
||||
@@ -0,0 +1,38 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any
|
||||
|
||||
|
||||
class StorageBackend(ABC):
|
||||
"""抽象存储后端基类"""
|
||||
|
||||
@abstractmethod
|
||||
def load_accounts(self) -> list[dict[str, Any]]:
|
||||
"""加载所有账号数据"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
|
||||
"""保存所有账号数据"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def load_auth_keys(self) -> list[dict[str, Any]]:
|
||||
"""加载所有鉴权密钥数据"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
|
||||
"""保存所有鉴权密钥数据"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def health_check(self) -> dict[str, Any]:
|
||||
"""健康检查，返回存储后端状态"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_backend_info(self) -> dict[str, Any]:
|
||||
"""获取存储后端信息"""
|
||||
pass
|
||||
@@ -0,0 +1,194 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import Column, String, Text, create_engine, Integer, text
|
||||
from sqlalchemy.ext.declarative import declarative_base
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
|
||||
from services.storage.base import StorageBackend
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class AccountModel(Base):
|
||||
"""账号数据模型"""
|
||||
__tablename__ = "accounts"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
access_token = Column(String(2048), unique=True, nullable=False, index=True)
|
||||
data = Column(Text, nullable=False) # JSON 格式存储完整账号数据
|
||||
|
||||
|
||||
class AuthKeyModel(Base):
|
||||
"""鉴权密钥数据模型"""
|
||||
__tablename__ = "auth_keys"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
key_id = Column(String(255), unique=True, nullable=False, index=True)
|
||||
data = Column(Text, nullable=False)
|
||||
|
||||
|
||||
class DatabaseStorageBackend(StorageBackend):
|
||||
"""数据库存储后端（支持 SQLite、PostgreSQL、MySQL 等）"""
|
||||
|
||||
def __init__(self, database_url: str):
|
||||
self.database_url = database_url
|
||||
self.engine = create_engine(
|
||||
database_url,
|
||||
pool_pre_ping=True, # 自动检测连接是否有效
|
||||
pool_recycle=3600, # 1小时回收连接
|
||||
)
|
||||
Base.metadata.create_all(self.engine)
|
||||
self.Session = sessionmaker(bind=self.engine)
|
||||
|
||||
def load_accounts(self) -> list[dict[str, Any]]:
|
||||
"""从数据库加载账号数据"""
|
||||
session = self.Session()
|
||||
try:
|
||||
accounts = []
|
||||
for row in session.query(AccountModel).all():
|
||||
try:
|
||||
account_data = json.loads(row.data)
|
||||
if isinstance(account_data, dict):
|
||||
accounts.append(account_data)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return accounts
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
|
||||
"""保存账号数据到数据库"""
|
||||
self._save_rows(AccountModel, accounts, "access_token")
|
||||
|
||||
def load_auth_keys(self) -> list[dict[str, Any]]:
|
||||
"""从数据库加载鉴权密钥数据"""
|
||||
return self._load_rows(AuthKeyModel)
|
||||
|
||||
def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
|
||||
"""保存鉴权密钥数据到数据库"""
|
||||
self._save_rows(AuthKeyModel, auth_keys, "id", "key_id")
|
||||
|
||||
def _load_rows(self, model: type[AccountModel] | type[AuthKeyModel]) -> list[dict[str, Any]]:
|
||||
session = self.Session()
|
||||
try:
|
||||
items = []
|
||||
for row in session.query(model).all():
|
||||
try:
|
||||
item_data = json.loads(row.data)
|
||||
if isinstance(item_data, dict):
|
||||
items.append(item_data)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return items
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def _save_rows(
|
||||
self,
|
||||
model: type[AccountModel] | type[AuthKeyModel],
|
||||
items: list[dict[str, Any]],
|
||||
source_key: str,
|
||||
target_key: str | None = None,
|
||||
) -> None:
|
||||
session = self.Session()
|
||||
try:
|
||||
key_column = target_key or source_key
|
||||
existing_rows = {
|
||||
str(getattr(row, key_column)): row
|
||||
for row in session.query(model).all()
|
||||
}
|
||||
incoming_keys: set[str] = set()
|
||||
|
||||
for item in items:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
key_value = str(item.get(source_key) or "").strip()
|
||||
if not key_value:
|
||||
continue
|
||||
if key_value in incoming_keys:
|
||||
raise ValueError(f"Duplicate {source_key} in storage snapshot")
|
||||
|
||||
incoming_keys.add(key_value)
|
||||
serialized_data = json.dumps(item, ensure_ascii=False)
|
||||
existing_row = existing_rows.get(key_value)
|
||||
if existing_row is None:
|
||||
session.add(
|
||||
model(
|
||||
**{key_column: key_value},
|
||||
data=serialized_data,
|
||||
)
|
||||
)
|
||||
elif existing_row.data != serialized_data:
|
||||
existing_row.data = serialized_data
|
||||
|
||||
for key_value, row in existing_rows.items():
|
||||
if key_value not in incoming_keys:
|
||||
session.delete(row)
|
||||
|
||||
session.commit()
|
||||
except Exception:
|
||||
session.rollback()
|
||||
raise
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
def health_check(self) -> dict[str, Any]:
|
||||
"""健康检查"""
|
||||
try:
|
||||
session = self.Session()
|
||||
try:
|
||||
# 尝试执行简单查询
|
||||
session.execute(text("SELECT 1"))
|
||||
count = session.query(AccountModel).count()
|
||||
auth_key_count = session.query(AuthKeyModel).count()
|
||||
return {
|
||||
"status": "healthy",
|
||||
"backend": "database",
|
||||
"database_url": self._mask_password(self.database_url),
|
||||
"account_count": count,
|
||||
"auth_key_count": auth_key_count,
|
||||
}
|
||||
finally:
|
||||
session.close()
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"backend": "database",
|
||||
"error": str(e),
|
||||
}
|
||||
|
||||
def get_backend_info(self) -> dict[str, Any]:
|
||||
"""获取存储后端信息"""
|
||||
db_type = "unknown"
|
||||
if "sqlite" in self.database_url:
|
||||
db_type = "sqlite"
|
||||
elif "postgresql" in self.database_url or "postgres" in self.database_url:
|
||||
db_type = "postgresql"
|
||||
elif "mysql" in self.database_url:
|
||||
db_type = "mysql"
|
||||
|
||||
return {
|
||||
"type": "database",
|
||||
"db_type": db_type,
|
||||
"description": f"数据库存储 ({db_type})",
|
||||
"database_url": self._mask_password(self.database_url),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _mask_password(url: str) -> str:
|
||||
"""隐藏数据库连接字符串中的密码"""
|
||||
if "://" not in url:
|
||||
return url
|
||||
try:
|
||||
protocol, rest = url.split("://", 1)
|
||||
if "@" in rest:
|
||||
credentials, host = rest.split("@", 1)
|
||||
if ":" in credentials:
|
||||
username, _ = credentials.split(":", 1)
|
||||
return f"{protocol}://{username}:****@{host}"
|
||||
return url
|
||||
except Exception:
|
||||
return url
|
||||
@@ -0,0 +1,104 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
from services.storage.base import StorageBackend
|
||||
from services.storage.database_storage import DatabaseStorageBackend
|
||||
from services.storage.git_storage import GitStorageBackend
|
||||
from services.storage.json_storage import JSONStorageBackend
|
||||
|
||||
|
||||
def create_storage_backend(data_dir: Path) -> StorageBackend:
|
||||
"""
|
||||
根据环境变量创建存储后端
|
||||
|
||||
环境变量：
|
||||
- STORAGE_BACKEND: json|sqlite|postgres|git (默认 json)
|
||||
- DATABASE_URL: 数据库连接字符串 (用于 sqlite/postgres)
|
||||
- GIT_REPO_URL: Git 仓库地址 (用于 git)
|
||||
- GIT_TOKEN: Git 访问令牌 (用于 git)
|
||||
- GIT_BRANCH: Git 分支 (默认 main)
|
||||
- GIT_FILE_PATH: Git 仓库中的文件路径 (默认 accounts.json)
|
||||
"""
|
||||
backend_type = os.getenv("STORAGE_BACKEND", "json").lower().strip()
|
||||
|
||||
print(f"[storage] Initializing storage backend: {backend_type}")
|
||||
|
||||
if backend_type == "json":
|
||||
# 本地 JSON 文件存储
|
||||
file_path = data_dir / "accounts.json"
|
||||
auth_keys_path = data_dir / "auth_keys.json"
|
||||
print(f"[storage] Using JSON storage: {file_path}")
|
||||
return JSONStorageBackend(file_path, auth_keys_path)
|
||||
|
||||
elif backend_type in ("sqlite", "postgres", "postgresql", "mysql", "database"):
|
||||
# 数据库存储
|
||||
database_url = os.getenv("DATABASE_URL", "").strip()
|
||||
|
||||
if not database_url:
|
||||
# 如果没有指定 DATABASE_URL，使用本地 SQLite
|
||||
database_url = f"sqlite:///{data_dir / 'accounts.db'}"
|
||||
print(f"[storage] No DATABASE_URL provided, using local SQLite: {database_url}")
|
||||
else:
|
||||
print(f"[storage] Using database storage: {_mask_password(database_url)}")
|
||||
|
||||
return DatabaseStorageBackend(database_url)
|
||||
|
||||
elif backend_type == "git":
|
||||
# Git 仓库存储
|
||||
repo_url = os.getenv("GIT_REPO_URL", "").strip()
|
||||
token = os.getenv("GIT_TOKEN", "").strip()
|
||||
branch = os.getenv("GIT_BRANCH", "main").strip()
|
||||
file_path = os.getenv("GIT_FILE_PATH", "accounts.json").strip()
|
||||
auth_keys_file_path = os.getenv("GIT_AUTH_KEYS_FILE_PATH", "auth_keys.json").strip()
|
||||
|
||||
if not repo_url:
|
||||
raise ValueError(
|
||||
"GIT_REPO_URL is required when using git storage backend. "
|
||||
"Please set GIT_REPO_URL environment variable."
|
||||
)
|
||||
|
||||
print(f"[storage] Using Git storage: {_mask_token(repo_url)}, branch: {branch}, file: {file_path}")
|
||||
|
||||
cache_dir = data_dir / "git_cache"
|
||||
return GitStorageBackend(
|
||||
repo_url=repo_url,
|
||||
token=token,
|
||||
branch=branch,
|
||||
file_path=file_path,
|
||||
auth_keys_file_path=auth_keys_file_path,
|
||||
local_cache_dir=cache_dir,
|
||||
)
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown storage backend: {backend_type}. "
|
||||
f"Supported backends: json, sqlite, postgres, git"
|
||||
)
|
||||
|
||||
|
||||
def _mask_password(url: str) -> str:
|
||||
"""隐藏数据库连接字符串中的密码"""
|
||||
if "://" not in url:
|
||||
return url
|
||||
try:
|
||||
protocol, rest = url.split("://", 1)
|
||||
if "@" in rest:
|
||||
credentials, host = rest.split("@", 1)
|
||||
if ":" in credentials:
|
||||
username, _ = credentials.split(":", 1)
|
||||
return f"{protocol}://{username}:****@{host}"
|
||||
return url
|
||||
except Exception:
|
||||
return url
|
||||
|
||||
|
||||
def _mask_token(url: str) -> str:
|
||||
"""隐藏 URL 中的 token"""
|
||||
if "@" in url and "://" in url:
|
||||
protocol, rest = url.split("://", 1)
|
||||
if "@" in rest:
|
||||
_, host = rest.split("@", 1)
|
||||
return f"{protocol}://****@{host}"
|
||||
return url
|
||||
@@ -0,0 +1,183 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from git import Repo
|
||||
from git.exc import GitCommandError
|
||||
|
||||
from services.storage.base import StorageBackend
|
||||
|
||||
|
||||
class GitStorageBackend(StorageBackend):
|
||||
"""Git 私有仓库存储后端"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
repo_url: str,
|
||||
token: str,
|
||||
branch: str = "main",
|
||||
file_path: str = "accounts.json",
|
||||
auth_keys_file_path: str = "auth_keys.json",
|
||||
local_cache_dir: Path | None = None,
|
||||
):
|
||||
self.repo_url = repo_url
|
||||
self.token = token
|
||||
self.branch = branch
|
||||
self.file_path = file_path
|
||||
self.auth_keys_file_path = auth_keys_file_path
|
||||
|
||||
# 本地缓存目录
|
||||
if local_cache_dir is None:
|
||||
local_cache_dir = Path(tempfile.gettempdir()) / "chatgpt2api_git_cache"
|
||||
self.local_cache_dir = local_cache_dir
|
||||
self.local_cache_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# 构建带认证的 Git URL
|
||||
self.auth_repo_url = self._build_auth_url(repo_url, token)
|
||||
|
||||
@staticmethod
|
||||
def _build_auth_url(repo_url: str, token: str) -> str:
|
||||
"""构建带认证的 Git URL"""
|
||||
if not token:
|
||||
return repo_url
|
||||
|
||||
# 支持 HTTPS 格式：https://github.com/user/repo.git
|
||||
if repo_url.startswith("https://"):
|
||||
# 插入 token
|
||||
return repo_url.replace("https://", f"https://{token}@")
|
||||
|
||||
# 支持 git@ 格式：git@github.com:user/repo.git
|
||||
# 转换为 HTTPS 格式
|
||||
if repo_url.startswith("git@"):
|
||||
repo_url = repo_url.replace("git@", "https://")
|
||||
repo_url = repo_url.replace(".com:", ".com/")
|
||||
return repo_url.replace("https://", f"https://{token}@")
|
||||
|
||||
return repo_url
|
||||
|
||||
def _clone_or_pull(self) -> Repo:
|
||||
"""克隆或拉取仓库"""
|
||||
repo_path = self.local_cache_dir / "repo"
|
||||
|
||||
if repo_path.exists() and (repo_path / ".git").exists():
|
||||
# 仓库已存在，拉取最新代码
|
||||
try:
|
||||
repo = Repo(repo_path)
|
||||
origin = repo.remote("origin")
|
||||
origin.pull(self.branch)
|
||||
return repo
|
||||
except GitCommandError:
|
||||
# 拉取失败，删除重新克隆
|
||||
shutil.rmtree(repo_path)
|
||||
|
||||
# 克隆仓库
|
||||
repo = Repo.clone_from(
|
||||
self.auth_repo_url,
|
||||
repo_path,
|
||||
branch=self.branch,
|
||||
)
|
||||
return repo
|
||||
|
||||
def load_accounts(self) -> list[dict[str, Any]]:
|
||||
"""从 Git 仓库加载账号数据"""
|
||||
try:
|
||||
return self._load_json_file(self.file_path)
|
||||
except Exception as e:
|
||||
print(f"[git-storage] load failed: {e}")
|
||||
raise
|
||||
|
||||
def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
|
||||
"""保存账号数据到 Git 仓库"""
|
||||
try:
|
||||
self._save_json_file(self.file_path, accounts, "Update accounts data")
|
||||
except Exception as e:
|
||||
print(f"[git-storage] save failed: {e}")
|
||||
raise e
|
||||
|
||||
def load_auth_keys(self) -> list[dict[str, Any]]:
|
||||
"""从 Git 仓库加载鉴权密钥数据"""
|
||||
try:
|
||||
data = self._load_json_value(self.auth_keys_file_path)
|
||||
if isinstance(data, dict):
|
||||
data = data.get("items")
|
||||
return data if isinstance(data, list) else []
|
||||
except Exception as e:
|
||||
print(f"[git-storage] load failed: {e}")
|
||||
raise
|
||||
|
||||
def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
|
||||
"""保存鉴权密钥数据到 Git 仓库"""
|
||||
try:
|
||||
self._save_json_file(self.auth_keys_file_path, {"items": auth_keys}, "Update auth keys data")
|
||||
except Exception as e:
|
||||
print(f"[git-storage] save failed: {e}")
|
||||
raise e
|
||||
|
||||
def _load_json_file(self, file_path: str) -> list[dict[str, Any]]:
|
||||
data = self._load_json_value(file_path)
|
||||
return data if isinstance(data, list) else []
|
||||
|
||||
def _load_json_value(self, file_path: str) -> Any:
|
||||
repo = self._clone_or_pull()
|
||||
file_full_path = Path(repo.working_dir) / file_path
|
||||
if not file_full_path.exists():
|
||||
return None
|
||||
return json.loads(file_full_path.read_text(encoding="utf-8"))
|
||||
|
||||
def _save_json_file(self, file_path: str, items: Any, message: str) -> None:
|
||||
repo = self._clone_or_pull()
|
||||
file_full_path = Path(repo.working_dir) / file_path
|
||||
file_full_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_full_path.write_text(
|
||||
json.dumps(items, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
repo.index.add([file_path])
|
||||
if repo.is_dirty():
|
||||
repo.index.commit(message)
|
||||
repo.remote("origin").push(self.branch)
|
||||
|
||||
def health_check(self) -> dict[str, Any]:
|
||||
"""健康检查"""
|
||||
try:
|
||||
repo = self._clone_or_pull()
|
||||
return {
|
||||
"status": "healthy",
|
||||
"backend": "git",
|
||||
"repo_url": self._mask_token(self.repo_url),
|
||||
"branch": self.branch,
|
||||
"file_path": self.file_path,
|
||||
"auth_keys_file_path": self.auth_keys_file_path,
|
||||
"last_commit": repo.head.commit.hexsha[:8],
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"backend": "git",
|
||||
"error": str(e),
|
||||
}
|
||||
|
||||
def get_backend_info(self) -> dict[str, Any]:
|
||||
"""获取存储后端信息"""
|
||||
return {
|
||||
"type": "git",
|
||||
"description": "Git 私有仓库存储",
|
||||
"repo_url": self._mask_token(self.repo_url),
|
||||
"branch": self.branch,
|
||||
"file_path": self.file_path,
|
||||
"auth_keys_file_path": self.auth_keys_file_path,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _mask_token(url: str) -> str:
|
||||
"""隐藏 URL 中的 token"""
|
||||
if "@" in url and "://" in url:
|
||||
protocol, rest = url.split("://", 1)
|
||||
if "@" in rest:
|
||||
_, host = rest.split("@", 1)
|
||||
return f"{protocol}://****@{host}"
|
||||
return url
|
||||
@@ -0,0 +1,95 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from services.storage.base import StorageBackend
|
||||
|
||||
|
||||
class JSONStorageBackend(StorageBackend):
|
||||
"""本地 JSON 文件存储后端"""
|
||||
|
||||
def __init__(self, file_path: Path, auth_keys_path: Path | None = None):
|
||||
self.file_path = file_path
|
||||
self.auth_keys_path = auth_keys_path or file_path.with_name("auth_keys.json")
|
||||
self.file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.auth_keys_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
@staticmethod
|
||||
def _load_json_list(file_path: Path) -> list[dict[str, Any]]:
|
||||
if not file_path.exists():
|
||||
return []
|
||||
try:
|
||||
data = json.loads(file_path.read_text(encoding="utf-8"))
|
||||
return data if isinstance(data, list) else []
|
||||
except (json.JSONDecodeError, Exception):
|
||||
return []
|
||||
|
||||
@staticmethod
|
||||
def _save_json_list(file_path: Path, items: list[dict[str, Any]]) -> None:
|
||||
file_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
file_path.write_text(
|
||||
json.dumps(items, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def load_accounts(self) -> list[dict[str, Any]]:
|
||||
"""从 JSON 文件加载账号数据"""
|
||||
return self._load_json_list(self.file_path)
|
||||
|
||||
def save_accounts(self, accounts: list[dict[str, Any]]) -> None:
|
||||
"""保存账号数据到 JSON 文件"""
|
||||
self._save_json_list(self.file_path, accounts)
|
||||
|
||||
def load_auth_keys(self) -> list[dict[str, Any]]:
|
||||
"""从 JSON 文件加载鉴权密钥数据"""
|
||||
if not self.auth_keys_path.exists():
|
||||
return []
|
||||
try:
|
||||
data = json.loads(self.auth_keys_path.read_text(encoding="utf-8"))
|
||||
except (json.JSONDecodeError, Exception):
|
||||
return []
|
||||
if isinstance(data, dict):
|
||||
data = data.get("items")
|
||||
return data if isinstance(data, list) else []
|
||||
|
||||
def save_auth_keys(self, auth_keys: list[dict[str, Any]]) -> None:
|
||||
"""保存鉴权密钥数据到 JSON 文件"""
|
||||
self.auth_keys_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.auth_keys_path.write_text(
|
||||
json.dumps({"items": auth_keys}, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def health_check(self) -> dict[str, Any]:
|
||||
"""健康检查"""
|
||||
try:
|
||||
# 检查文件是否可读写
|
||||
if self.file_path.exists():
|
||||
self.file_path.read_text(encoding="utf-8")
|
||||
return {
|
||||
"status": "healthy",
|
||||
"backend": "json",
|
||||
"file_exists": self.file_path.exists(),
|
||||
"file_path": str(self.file_path),
|
||||
"auth_keys_file_exists": self.auth_keys_path.exists(),
|
||||
"auth_keys_file_path": str(self.auth_keys_path),
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"status": "unhealthy",
|
||||
"backend": "json",
|
||||
"error": str(e),
|
||||
}
|
||||
|
||||
def get_backend_info(self) -> dict[str, Any]:
|
||||
"""获取存储后端信息"""
|
||||
return {
|
||||
"type": "json",
|
||||
"description": "本地 JSON 文件存储",
|
||||
"file_path": str(self.file_path),
|
||||
"file_exists": self.file_path.exists(),
|
||||
"auth_keys_file_path": str(self.auth_keys_path),
|
||||
"auth_keys_file_exists": self.auth_keys_path.exists(),
|
||||
}
|
||||
@@ -0,0 +1,531 @@
|
||||
"""Sub2API integration for browsing and importing ChatGPT OAuth accounts from a sub2api admin."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import threading
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from threading import Lock
|
||||
|
||||
from curl_cffi.requests import Session
|
||||
|
||||
from services.account_service import account_service
|
||||
from services.config import DATA_DIR
|
||||
|
||||
|
||||
SUB2API_CONFIG_FILE = DATA_DIR / "sub2api_config.json"
|
||||
|
||||
# Cached JWT per server to avoid re-login on every list/import call.
|
||||
# Token lifetime on sub2api defaults to 24h; we refresh 5 min before expiry.
|
||||
_TOKEN_REFRESH_SKEW = 5 * 60
|
||||
|
||||
|
||||
def _new_id() -> str:
|
||||
return uuid.uuid4().hex[:12]
|
||||
|
||||
|
||||
def _now_iso() -> str:
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
def _clean(value: object) -> str:
|
||||
return str(value or "").strip()
|
||||
|
||||
|
||||
def _normalize_import_job(raw: object, *, fail_unfinished: bool) -> dict | None:
|
||||
if not isinstance(raw, dict):
|
||||
return None
|
||||
status = _clean(raw.get("status")) or "failed"
|
||||
if fail_unfinished and status in {"pending", "running"}:
|
||||
status = "failed"
|
||||
return {
|
||||
"job_id": _clean(raw.get("job_id")) or uuid.uuid4().hex,
|
||||
"status": status,
|
||||
"created_at": _clean(raw.get("created_at")) or _now_iso(),
|
||||
"updated_at": _clean(raw.get("updated_at")) or _clean(raw.get("created_at")) or _now_iso(),
|
||||
"total": int(raw.get("total") or 0),
|
||||
"completed": int(raw.get("completed") or 0),
|
||||
"added": int(raw.get("added") or 0),
|
||||
"skipped": int(raw.get("skipped") or 0),
|
||||
"refreshed": int(raw.get("refreshed") or 0),
|
||||
"failed": int(raw.get("failed") or 0),
|
||||
"errors": raw.get("errors") if isinstance(raw.get("errors"), list) else [],
|
||||
}
|
||||
|
||||
|
||||
def _normalize_server(raw: dict) -> dict:
|
||||
return {
|
||||
"id": _clean(raw.get("id")) or _new_id(),
|
||||
"name": _clean(raw.get("name")),
|
||||
"base_url": _clean(raw.get("base_url")),
|
||||
"email": _clean(raw.get("email")),
|
||||
"password": _clean(raw.get("password")),
|
||||
"api_key": _clean(raw.get("api_key")),
|
||||
"group_id": _clean(raw.get("group_id")),
|
||||
"import_job": _normalize_import_job(raw.get("import_job"), fail_unfinished=True),
|
||||
}
|
||||
|
||||
|
||||
class Sub2APIConfig:
|
||||
def __init__(self, store_file: Path):
|
||||
self._store_file = store_file
|
||||
self._lock = Lock()
|
||||
self._servers: list[dict] = self._load()
|
||||
|
||||
def _load(self) -> list[dict]:
|
||||
if not self._store_file.exists():
|
||||
return []
|
||||
try:
|
||||
raw = json.loads(self._store_file.read_text(encoding="utf-8"))
|
||||
if isinstance(raw, list):
|
||||
return [_normalize_server(item) for item in raw if isinstance(item, dict)]
|
||||
except Exception:
|
||||
pass
|
||||
return []
|
||||
|
||||
def _save(self) -> None:
|
||||
self._store_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._store_file.write_text(
|
||||
json.dumps(self._servers, ensure_ascii=False, indent=2) + "\n",
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
def list_servers(self) -> list[dict]:
|
||||
with self._lock:
|
||||
return [dict(server) for server in self._servers]
|
||||
|
||||
def get_server(self, server_id: str) -> dict | None:
|
||||
with self._lock:
|
||||
for server in self._servers:
|
||||
if server["id"] == server_id:
|
||||
return dict(server)
|
||||
return None
|
||||
|
||||
def add_server(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
base_url: str,
|
||||
email: str,
|
||||
password: str,
|
||||
api_key: str,
|
||||
group_id: str = "",
|
||||
) -> dict:
|
||||
server = _normalize_server({
|
||||
"id": _new_id(),
|
||||
"name": name,
|
||||
"base_url": base_url,
|
||||
"email": email,
|
||||
"password": password,
|
||||
"api_key": api_key,
|
||||
"group_id": group_id,
|
||||
})
|
||||
with self._lock:
|
||||
self._servers.append(server)
|
||||
self._save()
|
||||
_token_cache.pop(server["id"], None)
|
||||
return dict(server)
|
||||
|
||||
def update_server(self, server_id: str, updates: dict) -> dict | None:
|
||||
with self._lock:
|
||||
for index, server in enumerate(self._servers):
|
||||
if server["id"] != server_id:
|
||||
continue
|
||||
merged = {**server, **{k: v for k, v in updates.items() if v is not None}, "id": server_id}
|
||||
self._servers[index] = _normalize_server(merged)
|
||||
self._save()
|
||||
result = dict(self._servers[index])
|
||||
break
|
||||
else:
|
||||
return None
|
||||
_token_cache.pop(server_id, None)
|
||||
return result
|
||||
|
||||
def delete_server(self, server_id: str) -> bool:
|
||||
with self._lock:
|
||||
before = len(self._servers)
|
||||
self._servers = [server for server in self._servers if server["id"] != server_id]
|
||||
removed = len(self._servers) < before
|
||||
if removed:
|
||||
self._save()
|
||||
if removed:
|
||||
_token_cache.pop(server_id, None)
|
||||
return removed
|
||||
|
||||
def set_import_job(self, server_id: str, import_job: dict | None) -> dict | None:
|
||||
with self._lock:
|
||||
for index, server in enumerate(self._servers):
|
||||
if server["id"] != server_id:
|
||||
continue
|
||||
next_server = dict(server)
|
||||
next_server["import_job"] = _normalize_import_job(import_job, fail_unfinished=False)
|
||||
self._servers[index] = next_server
|
||||
self._save()
|
||||
return dict(next_server)
|
||||
return None
|
||||
|
||||
def get_import_job(self, server_id: str) -> dict | None:
|
||||
with self._lock:
|
||||
for server in self._servers:
|
||||
if server["id"] == server_id:
|
||||
job = server.get("import_job")
|
||||
return dict(job) if isinstance(job, dict) else None
|
||||
return None
|
||||
|
||||
|
||||
# Per-server cached access token: {server_id: (jwt, expires_at_epoch)}
|
||||
_token_cache: dict[str, tuple[str, float]] = {}
|
||||
_token_cache_lock = Lock()
|
||||
|
||||
|
||||
def _login(base_url: str, email: str, password: str) -> tuple[str, float]:
|
||||
url = f"{base_url.rstrip('/')}/api/v1/auth/login"
|
||||
session = Session(verify=True)
|
||||
try:
|
||||
response = session.post(
|
||||
url,
|
||||
json={"email": email, "password": password},
|
||||
headers={"Accept": "application/json", "Content-Type": "application/json"},
|
||||
timeout=30,
|
||||
)
|
||||
if not response.ok:
|
||||
raise RuntimeError(f"sub2api login failed: HTTP {response.status_code} {response.text[:200]}")
|
||||
payload = response.json()
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
body = _unwrap_envelope(payload)
|
||||
if not isinstance(body, dict):
|
||||
raise RuntimeError("sub2api login payload is invalid")
|
||||
|
||||
token = _clean(body.get("access_token"))
|
||||
if not token:
|
||||
raise RuntimeError("sub2api login did not return access_token")
|
||||
|
||||
expires_in = int(body.get("expires_in") or 3600)
|
||||
expires_at = time.time() + max(60, expires_in) - _TOKEN_REFRESH_SKEW
|
||||
return token, expires_at
|
||||
|
||||
|
||||
def _auth_headers(server: dict) -> dict[str, str]:
|
||||
api_key = _clean(server.get("api_key"))
|
||||
if api_key:
|
||||
return {"x-api-key": api_key, "Accept": "application/json"}
|
||||
|
||||
email = _clean(server.get("email"))
|
||||
password = _clean(server.get("password"))
|
||||
if not email or not password:
|
||||
raise RuntimeError("sub2api server requires email+password or api_key")
|
||||
|
||||
server_id = _clean(server.get("id"))
|
||||
base_url = _clean(server.get("base_url"))
|
||||
|
||||
with _token_cache_lock:
|
||||
cached = _token_cache.get(server_id)
|
||||
if cached and cached[1] > time.time():
|
||||
return {"Authorization": f"Bearer {cached[0]}", "Accept": "application/json"}
|
||||
|
||||
token, expires_at = _login(base_url, email, password)
|
||||
with _token_cache_lock:
|
||||
_token_cache[server_id] = (token, expires_at)
|
||||
return {"Authorization": f"Bearer {token}", "Accept": "application/json"}
|
||||
|
||||
|
||||
def _extract_access_token(credentials: object) -> str:
|
||||
if not isinstance(credentials, dict):
|
||||
return ""
|
||||
for key in ("access_token", "accessToken", "token"):
|
||||
value = _clean(credentials.get(key))
|
||||
if value:
|
||||
return value
|
||||
return ""
|
||||
|
||||
|
||||
def _unwrap_envelope(payload: object) -> object:
|
||||
"""Peel sub2api's `{code, message, data}` envelope, returning the inner `data` field
|
||||
when present. Also handles unwrapped responses from older/alt versions."""
|
||||
if isinstance(payload, dict) and "data" in payload and "code" in payload:
|
||||
return payload.get("data")
|
||||
return payload
|
||||
|
||||
|
||||
def _extract_paged_items(payload: object) -> tuple[list, int]:
|
||||
"""Return (items, total) from a paginated sub2api response.
|
||||
|
||||
Handles both the wrapped shape `{code,data:{items,total,...}}` and a few looser
|
||||
variants (`{data:[...]}`, `[...]`, `{items:[...],total:N}`)."""
|
||||
inner = _unwrap_envelope(payload)
|
||||
if isinstance(inner, list):
|
||||
return inner, len(inner)
|
||||
if isinstance(inner, dict):
|
||||
for key in ("items", "data", "list"):
|
||||
value = inner.get(key)
|
||||
if isinstance(value, list):
|
||||
return value, int(inner.get("total") or len(value))
|
||||
return [], 0
|
||||
|
||||
|
||||
def list_remote_accounts(server: dict) -> list[dict]:
|
||||
"""Return a flat list of OpenAI OAuth accounts from a sub2api server."""
|
||||
base_url = _clean(server.get("base_url"))
|
||||
if not base_url:
|
||||
return []
|
||||
|
||||
headers = _auth_headers(server)
|
||||
group_id = _clean(server.get("group_id"))
|
||||
|
||||
session = Session(verify=True)
|
||||
items: list[dict] = []
|
||||
try:
|
||||
page = 1
|
||||
while True:
|
||||
params: dict[str, object] = {
|
||||
"platform": "openai",
|
||||
"type": "oauth",
|
||||
"page": page,
|
||||
"page_size": 200,
|
||||
}
|
||||
if group_id:
|
||||
params["group"] = group_id
|
||||
response = session.get(
|
||||
f"{base_url.rstrip('/')}/api/v1/admin/accounts",
|
||||
headers=headers,
|
||||
params=params,
|
||||
timeout=30,
|
||||
)
|
||||
if not response.ok:
|
||||
raise RuntimeError(f"sub2api list failed: HTTP {response.status_code} {response.text[:200]}")
|
||||
payload = response.json()
|
||||
|
||||
data, total = _extract_paged_items(payload)
|
||||
if not data:
|
||||
break
|
||||
|
||||
for account in data:
|
||||
if not isinstance(account, dict):
|
||||
continue
|
||||
credentials = account.get("credentials") if isinstance(account.get("credentials"), dict) else {}
|
||||
account_id = account.get("id")
|
||||
if account_id is None:
|
||||
continue
|
||||
items.append({
|
||||
"id": str(account_id),
|
||||
"name": _clean(account.get("name")),
|
||||
"email": _clean(credentials.get("email")) or _clean(account.get("name")),
|
||||
"plan_type": _clean(credentials.get("plan_type")),
|
||||
"status": _clean(account.get("status")),
|
||||
"expires_at": _clean(credentials.get("expires_at")),
|
||||
"has_refresh_token": bool(_clean(credentials.get("refresh_token"))),
|
||||
})
|
||||
|
||||
if page * 200 >= total or len(data) < 200:
|
||||
break
|
||||
page += 1
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
return items
|
||||
|
||||
|
||||
def list_remote_groups(server: dict) -> list[dict]:
|
||||
"""Return OpenAI account groups from a sub2api server."""
|
||||
base_url = _clean(server.get("base_url"))
|
||||
if not base_url:
|
||||
return []
|
||||
|
||||
headers = _auth_headers(server)
|
||||
|
||||
session = Session(verify=True)
|
||||
items: list[dict] = []
|
||||
try:
|
||||
page = 1
|
||||
while True:
|
||||
response = session.get(
|
||||
f"{base_url.rstrip('/')}/api/v1/admin/groups",
|
||||
headers=headers,
|
||||
params={
|
||||
"page": page,
|
||||
"page_size": 200,
|
||||
},
|
||||
timeout=30,
|
||||
)
|
||||
if not response.ok:
|
||||
raise RuntimeError(f"sub2api groups failed: HTTP {response.status_code} {response.text[:200]}")
|
||||
payload = response.json()
|
||||
|
||||
data, total = _extract_paged_items(payload)
|
||||
if not data:
|
||||
break
|
||||
|
||||
for group in data:
|
||||
if not isinstance(group, dict):
|
||||
continue
|
||||
group_id = group.get("id")
|
||||
if group_id is None:
|
||||
continue
|
||||
items.append({
|
||||
"id": str(group_id),
|
||||
"name": _clean(group.get("name")),
|
||||
"description": _clean(group.get("description")),
|
||||
"platform": _clean(group.get("platform")),
|
||||
"status": _clean(group.get("status")),
|
||||
"account_count": int(group.get("account_count") or 0),
|
||||
"active_account_count": int(group.get("active_account_count") or 0),
|
||||
})
|
||||
|
||||
if page * 200 >= total or len(data) < 200:
|
||||
break
|
||||
page += 1
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
return items
|
||||
|
||||
|
||||
def _fetch_access_tokens_for_accounts(server: dict, account_ids: list[str]) -> tuple[list[str], list[dict]]:
|
||||
"""Return exported access tokens and per-account errors from sub2api."""
|
||||
base_url = _clean(server.get("base_url"))
|
||||
headers = _auth_headers(server)
|
||||
ids = [_clean(item) for item in account_ids if _clean(item)]
|
||||
if not ids:
|
||||
return [], []
|
||||
|
||||
session = Session(verify=True)
|
||||
try:
|
||||
response = session.get(
|
||||
f"{base_url.rstrip('/')}/api/v1/admin/accounts/data",
|
||||
headers=headers,
|
||||
params={"ids": ",".join(ids), "timezone": "Asia/Shanghai"},
|
||||
timeout=30,
|
||||
)
|
||||
if not response.ok:
|
||||
raise RuntimeError(f"HTTP {response.status_code}")
|
||||
payload = response.json()
|
||||
finally:
|
||||
session.close()
|
||||
|
||||
data = _unwrap_envelope(payload)
|
||||
accounts = data.get("accounts") if isinstance(data, dict) else None
|
||||
if not isinstance(accounts, list):
|
||||
raise RuntimeError("invalid export payload")
|
||||
|
||||
tokens: list[str] = []
|
||||
errors: list[dict] = []
|
||||
for account in accounts:
|
||||
if not isinstance(account, dict):
|
||||
continue
|
||||
credentials = account.get("credentials") if isinstance(account.get("credentials"), dict) else {}
|
||||
token = _extract_access_token(credentials)
|
||||
account_id = _clean(account.get("id")) or _clean(credentials.get("chatgpt_account_id")) or _clean(account.get("name"))
|
||||
if token:
|
||||
tokens.append(token)
|
||||
else:
|
||||
errors.append({"name": account_id or _clean(account.get("name")) or "unknown", "error": "missing access_token"})
|
||||
|
||||
if len(accounts) < len(ids):
|
||||
errors.append({"name": ",".join(ids), "error": f"exported {len(accounts)}/{len(ids)} accounts"})
|
||||
|
||||
return tokens, errors
|
||||
|
||||
|
||||
class Sub2APIImportService:
|
||||
def __init__(self, sub2api_config: Sub2APIConfig):
|
||||
self._config = sub2api_config
|
||||
|
||||
def start_import(self, server: dict, account_ids: list[str]) -> dict:
|
||||
ids = [_clean(item) for item in account_ids if _clean(item)]
|
||||
if not ids:
|
||||
raise ValueError("account ids is required")
|
||||
|
||||
server_id = _clean(server.get("id"))
|
||||
job = {
|
||||
"job_id": uuid.uuid4().hex,
|
||||
"status": "pending",
|
||||
"created_at": _now_iso(),
|
||||
"updated_at": _now_iso(),
|
||||
"total": len(ids),
|
||||
"completed": 0,
|
||||
"added": 0,
|
||||
"skipped": 0,
|
||||
"refreshed": 0,
|
||||
"failed": 0,
|
||||
"errors": [],
|
||||
}
|
||||
saved = self._config.set_import_job(server_id, job)
|
||||
if saved is None:
|
||||
raise ValueError("server not found")
|
||||
|
||||
thread = threading.Thread(
|
||||
target=self._run_import,
|
||||
args=(server_id, server, ids),
|
||||
name=f"sub2api-import-{server_id}",
|
||||
daemon=True,
|
||||
)
|
||||
thread.start()
|
||||
return dict(saved.get("import_job") or job)
|
||||
|
||||
def _update_job(self, server_id: str, **updates) -> None:
|
||||
current = self._config.get_import_job(server_id)
|
||||
if current is None:
|
||||
return
|
||||
next_job = {**current, **updates, "updated_at": _now_iso()}
|
||||
self._config.set_import_job(server_id, next_job)
|
||||
|
||||
def _append_error(self, server_id: str, account_id: str, message: str) -> None:
|
||||
current = self._config.get_import_job(server_id)
|
||||
if current is None:
|
||||
return
|
||||
errors = list(current.get("errors") or [])
|
||||
errors.append({"name": account_id, "error": message})
|
||||
self._update_job(server_id, errors=errors, failed=len(errors))
|
||||
|
||||
def _run_import(self, server_id: str, server: dict, account_ids: list[str]) -> None:
|
||||
self._update_job(server_id, status="running")
|
||||
|
||||
try:
|
||||
tokens, errors = _fetch_access_tokens_for_accounts(server, account_ids)
|
||||
except Exception as exc:
|
||||
message = str(exc) or "unknown error"
|
||||
for account_id in account_ids:
|
||||
self._append_error(server_id, account_id, message)
|
||||
tokens = []
|
||||
else:
|
||||
for error in errors:
|
||||
self._append_error(server_id, _clean(error.get("name")), _clean(error.get("error")) or "unknown error")
|
||||
|
||||
current = self._config.get_import_job(server_id) or {}
|
||||
self._update_job(
|
||||
server_id,
|
||||
completed=len(account_ids),
|
||||
failed=len(current.get("errors") or []),
|
||||
)
|
||||
|
||||
if not tokens:
|
||||
current = self._config.get_import_job(server_id) or {}
|
||||
self._update_job(
|
||||
server_id,
|
||||
status="failed",
|
||||
completed=int(current.get("total") or 0),
|
||||
failed=len(current.get("errors") or []),
|
||||
)
|
||||
return
|
||||
|
||||
add_result = account_service.add_accounts(tokens, source_type="codex")
|
||||
refresh_result = account_service.refresh_accounts(tokens)
|
||||
current = self._config.get_import_job(server_id) or {}
|
||||
self._update_job(
|
||||
server_id,
|
||||
status="completed",
|
||||
completed=len(account_ids),
|
||||
added=int(add_result.get("added") or 0),
|
||||
skipped=int(add_result.get("skipped") or 0),
|
||||
refreshed=int(refresh_result.get("refreshed") or 0),
|
||||
failed=len(current.get("errors") or []),
|
||||
)
|
||||
|
||||
|
||||
sub2api_config = Sub2APIConfig(SUB2API_CONFIG_FILE)
|
||||
sub2api_import_service = Sub2APIImportService(sub2api_config)
|
||||
@@ -0,0 +1,68 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
from utils.helper import UpstreamHTTPError
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class UploadLimit:
|
||||
blocked_until: datetime
|
||||
reason: str
|
||||
|
||||
|
||||
class UploadUnavailableError(RuntimeError):
|
||||
def __init__(self, blocked_until: datetime | None = None, *, attempts_exhausted: bool = False):
|
||||
message = (
|
||||
"多个账号上传受限，本次自动切换已达到上限，请稍后重试。"
|
||||
if attempts_exhausted else "当前没有可用于参考图上传的账号，请等待上传额度恢复或添加可用账号。"
|
||||
)
|
||||
if blocked_until:
|
||||
message += f" 已知上传限制最早预计恢复时间：{blocked_until.astimezone().strftime('%Y-%m-%d %H:%M:%S %Z')}。"
|
||||
super().__init__(message)
|
||||
|
||||
|
||||
def parse_upload_limit(error: Exception, now: datetime | None = None) -> UploadLimit | None:
|
||||
"""Recognize file quota failures only at the upstream upload allocation endpoint."""
|
||||
if not isinstance(error, UpstreamHTTPError) or error.status_code != 429 or error.context != "/backend-api/files":
|
||||
return None
|
||||
body = error.body
|
||||
if isinstance(body, str):
|
||||
try:
|
||||
body = json.loads(body)
|
||||
except (ValueError, TypeError):
|
||||
return None
|
||||
if not isinstance(body, dict):
|
||||
return None
|
||||
detail = body.get("detail") or body.get("error") or body
|
||||
if not isinstance(detail, dict):
|
||||
return None
|
||||
reason = str(detail.get("message") or "文件上传额度已用尽")
|
||||
codes = {str(detail.get(key) or "").lower() for key in ("code", "error_code", "type")}
|
||||
if not codes.intersection({"throttled", "file_upload_limit_exceeded", "upload_limit_exceeded"}) and not (
|
||||
"上传上限" in reason or "upload limit" in reason.lower()
|
||||
):
|
||||
return None
|
||||
now = now or datetime.now(timezone.utc)
|
||||
for key in ("retry_at", "reset_at"):
|
||||
value = detail.get(key)
|
||||
if value:
|
||||
try:
|
||||
reset = datetime.fromisoformat(str(value).replace("Z", "+00:00"))
|
||||
if reset.tzinfo and reset > now:
|
||||
return UploadLimit(reset, reason)
|
||||
except (ValueError, TypeError):
|
||||
pass
|
||||
if error.retry_after is not None and error.retry_after > 0:
|
||||
seconds = min(error.retry_after, 365 * 86400)
|
||||
else:
|
||||
units = {"小时": 3600, "分钟": 60, "秒": 1, "天": 86400,
|
||||
"day": 86400, "hour": 3600, "minute": 60, "second": 1}
|
||||
matches = re.findall(r"(\d+(?:\.\d+)?)\s*(小时|分钟|秒|天|days?|hours?|minutes?|seconds?)", reason.lower())
|
||||
seconds = sum(float(number) * units[unit.rstrip("s")] for number, unit in matches)
|
||||
# Unknown reset time: temporary cooling, never a fabricated daily reset.
|
||||
seconds = min(seconds, 365 * 86400) if seconds > 0 else 15 * 60
|
||||
return UploadLimit(now + timedelta(seconds=max(1, seconds)), reason)
|
||||
Reference in New Issue
Block a user