feat: 实现一键同步网站账号上游模型功能
This commit is contained in:
liumangmang
@@ -44,6 +44,8 @@ from app.schemas.website import (
|
||||
SetConcurrencyRequest,
|
||||
SetConcurrencyItem,
|
||||
SetConcurrencyResponse,
|
||||
SyncUpstreamModelsItem,
|
||||
SyncUpstreamModelsResponse,
|
||||
)
|
||||
|
||||
from app.services.website_client import Sub2ApiWebsiteClient, _extract_id
|
||||
@@ -1280,6 +1282,165 @@ def set_website_accounts_concurrency(
|
||||
)
|
||||
|
||||
|
||||
@router.post("/api/websites/{wid}/accounts/sync-upstream-models", response_model=SyncUpstreamModelsResponse)
|
||||
def sync_website_accounts_upstream_models(
|
||||
wid: int,
|
||||
db: Session = Depends(get_db),
|
||||
_=Depends(get_current_user),
|
||||
):
|
||||
"""一键同步上游模型"""
|
||||
website = db.query(Website).filter(Website.id == wid).first()
|
||||
if not website:
|
||||
raise HTTPException(404, "website not found")
|
||||
if website.site_type != "sub2api":
|
||||
raise HTTPException(400, "only sub2api site supports account model syncing")
|
||||
|
||||
with _client(website) as c:
|
||||
try:
|
||||
remote_accounts = c.list_accounts()
|
||||
except Exception as e:
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=False,
|
||||
message=f"拉取远端账号列表失败: {e}",
|
||||
items=[]
|
||||
)
|
||||
|
||||
if remote_accounts is None:
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=False,
|
||||
message="拉取远端账号列表失败，无法同步上游模型",
|
||||
items=[]
|
||||
)
|
||||
|
||||
remote_map = {}
|
||||
for acc in remote_accounts:
|
||||
aid = c.extract_id(acc)
|
||||
if aid:
|
||||
remote_map[aid] = acc
|
||||
|
||||
# 1. 查找候选账号：本地 UpstreamGeneratedKey.imported_website_id == wid 且 imported_account_id 非空
|
||||
keys = db.query(UpstreamGeneratedKey).filter(
|
||||
UpstreamGeneratedKey.imported_website_id == wid,
|
||||
UpstreamGeneratedKey.imported_account_id.isnot(None)
|
||||
).all()
|
||||
|
||||
candidates = {}
|
||||
for key in keys:
|
||||
aid = key.imported_account_id
|
||||
if aid not in candidates:
|
||||
candidates[aid] = {
|
||||
"account_id": aid,
|
||||
"db_key": key
|
||||
}
|
||||
|
||||
if not candidates:
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=True,
|
||||
message="没有找到 SmartUp 导入的有效账号",
|
||||
items=[]
|
||||
)
|
||||
|
||||
items = []
|
||||
|
||||
for aid, cand in candidates.items():
|
||||
# 校验是否存在于远端
|
||||
remote_acc = remote_map.get(aid)
|
||||
if not remote_acc:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=None,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="skipped",
|
||||
message="账号在远端已被删除或不存在"
|
||||
))
|
||||
continue
|
||||
|
||||
acc_name = remote_acc.get("name")
|
||||
|
||||
# 校验账号 ID 是否为数字
|
||||
try:
|
||||
int(aid)
|
||||
except ValueError:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="skipped",
|
||||
message="账号 ID 非数字，跳过同步"
|
||||
))
|
||||
continue
|
||||
|
||||
# 开始同步该账号的模型
|
||||
try:
|
||||
raw_models = c.sync_account_upstream_models(aid)
|
||||
# 过滤空、去重、排序
|
||||
valid_models = sorted(list(set(m.strip() for m in raw_models if m and m.strip())))
|
||||
|
||||
if not valid_models:
|
||||
# 如果返回空模型，保守处理为失败，不清空已有的模型白名单配置
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="failed",
|
||||
message="上游同步返回模型列表为空"
|
||||
))
|
||||
continue
|
||||
|
||||
# 构造 model_mapping
|
||||
model_mapping = {m: m for m in valid_models}
|
||||
|
||||
# 优先读取账号现有的 credentials 并 merge，以防覆盖掉 api_key/base_url 等其他配置
|
||||
current_creds = remote_acc.get("credentials") or {}
|
||||
updated_creds = dict(current_creds)
|
||||
updated_creds["model_mapping"] = model_mapping
|
||||
|
||||
# 调用更新接口
|
||||
c.update_account(aid, {"credentials": updated_creds})
|
||||
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=len(valid_models),
|
||||
models=valid_models,
|
||||
status="success",
|
||||
message=f"成功同步 {len(valid_models)} 个模型"
|
||||
))
|
||||
except Exception as e:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="failed",
|
||||
message=f"同步/保存上游模型失败: {e}"
|
||||
))
|
||||
|
||||
success_count = sum(1 for item in items if item.status == "success")
|
||||
failed_count = sum(1 for item in items if item.status == "failed")
|
||||
skip_count = sum(1 for item in items if item.status == "skipped")
|
||||
|
||||
success = (failed_count == 0)
|
||||
msg_parts = []
|
||||
if success_count:
|
||||
msg_parts.append(f"成功 {success_count} 个")
|
||||
if failed_count:
|
||||
msg_parts.append(f"失败 {failed_count} 个")
|
||||
if skip_count:
|
||||
msg_parts.append(f"跳过 {skip_count} 个")
|
||||
|
||||
message = "同步上游模型执行完毕：" + "，".join(msg_parts)
|
||||
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=success,
|
||||
message=message,
|
||||
items=items
|
||||
)
|
||||
|
||||
|
||||
@router.post("/api/websites/{wid}/groups/organize", response_model=OrganizeGroupsResponse)
|
||||
def organize_website_groups(
|
||||
wid: int,
|
||||
|
||||
@@ -287,3 +287,18 @@ class SetConcurrencyResponse(BaseModel):
|
||||
success: bool
|
||||
message: str
|
||||
items: list[SetConcurrencyItem]
|
||||
|
||||
|
||||
class SyncUpstreamModelsItem(BaseModel):
|
||||
account_id: str
|
||||
account_name: Optional[str] = None
|
||||
model_count: int
|
||||
models: list[str] = Field(default_factory=list)
|
||||
status: str # "success" | "skipped" | "failed"
|
||||
message: str
|
||||
|
||||
|
||||
class SyncUpstreamModelsResponse(BaseModel):
|
||||
success: bool
|
||||
message: str
|
||||
items: list[SyncUpstreamModelsItem]
|
||||
|
||||
@@ -312,6 +312,21 @@ class Sub2ApiWebsiteClient:
|
||||
data = _unwrap_data(resp)
|
||||
return data if isinstance(data, dict) else {"value": data}
|
||||
|
||||
def sync_account_upstream_models(self, account_id: str, endpoint: str = "/accounts") -> list[str]:
|
||||
"""拉取上游真实支持模型并返回列表。"""
|
||||
quoted_id = quote(account_id, safe="")
|
||||
path = f"{endpoint}/{quoted_id}/models/sync-upstream"
|
||||
resp = self._request("POST", path)
|
||||
data = _unwrap_data(resp)
|
||||
if isinstance(data, list):
|
||||
return [str(m) for m in data if m]
|
||||
if isinstance(data, dict):
|
||||
for key in ("models", "items", "data", "list"):
|
||||
val = data.get(key)
|
||||
if isinstance(val, list):
|
||||
return [str(m) for m in val if m]
|
||||
raise WebsiteError(f"同步上游模型接口返回的数据格式不正确: {resp}")
|
||||
|
||||
@staticmethod
|
||||
def _unwrap_list(value: dict) -> list | None:
|
||||
"""递归展开嵌套的列表包装：data.items、data.data、items、accounts 等。"""
|
||||
|
||||
@@ -0,0 +1,241 @@
|
||||
import pytest
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from sqlalchemy.pool import StaticPool
|
||||
|
||||
from app.database import Base
|
||||
from app.models.upstream import Upstream
|
||||
from app.models.upstream_key import UpstreamGeneratedKey
|
||||
from app.models.website import Website
|
||||
from app.routers.websites import sync_website_accounts_upstream_models
|
||||
from app.schemas.website import SyncUpstreamModelsResponse
|
||||
|
||||
@pytest.fixture()
|
||||
def db_session():
|
||||
engine = create_engine(
|
||||
"sqlite://",
|
||||
connect_args={"check_same_thread": False},
|
||||
poolclass=StaticPool,
|
||||
)
|
||||
from app.models import admin_user, upstream, snapshot, webhook_config, notification_log, custom_page, website, revoked_token, upstream_key
|
||||
Base.metadata.create_all(bind=engine)
|
||||
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
db = TestingSessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
|
||||
def test_sync_upstream_models_workflow(db_session, monkeypatch):
|
||||
# 1. 创建网站与上游
|
||||
w = Website(
|
||||
id=1,
|
||||
name="Sub2Api site",
|
||||
site_type="sub2api",
|
||||
base_url="http://sub2api",
|
||||
api_prefix="api/v1",
|
||||
auth_type="bearer",
|
||||
auth_config_json='{"token": "tok1"}'
|
||||
)
|
||||
db_session.add(w)
|
||||
|
||||
up = Upstream(id=10, name="Upstream A", base_url="http://upstream-a")
|
||||
db_session.add(up)
|
||||
db_session.commit()
|
||||
|
||||
# 2. 插入本地导入的 Key 记录
|
||||
# k1: 属于本站，且在远端正常存在的账号
|
||||
k1 = UpstreamGeneratedKey(
|
||||
id=101,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-101",
|
||||
key_value="val-101",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1001",
|
||||
)
|
||||
# k2: 属于本站，但账号 ID 无法转为数字
|
||||
k2 = UpstreamGeneratedKey(
|
||||
id=102,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-102",
|
||||
key_value="val-102",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="abc",
|
||||
)
|
||||
# k3: 属于其他站的账号
|
||||
k3 = UpstreamGeneratedKey(
|
||||
id=103,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-103",
|
||||
key_value="val-103",
|
||||
status="active",
|
||||
imported_website_id=999,
|
||||
imported_account_id="1003",
|
||||
)
|
||||
# k4: 属于本站，但远端已被删除的账号
|
||||
k4 = UpstreamGeneratedKey(
|
||||
id=104,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-104",
|
||||
key_value="val-104",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1004",
|
||||
)
|
||||
# k5: 属于本站，但同步接口返回空模型（应当标记为 failed，且不覆盖）
|
||||
k5 = UpstreamGeneratedKey(
|
||||
id=105,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-105",
|
||||
key_value="val-105",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1005",
|
||||
)
|
||||
# k6: 属于本站，但同步接口抛错（应当标记为 failed，且不覆盖）
|
||||
k6 = UpstreamGeneratedKey(
|
||||
id=106,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-106",
|
||||
key_value="val-106",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1006",
|
||||
)
|
||||
db_session.add_all([k1, k2, k3, k4, k5, k6])
|
||||
db_session.commit()
|
||||
|
||||
# Mock Sub2Api 客户端
|
||||
sync_calls = []
|
||||
update_calls = []
|
||||
closed_count = 0
|
||||
|
||||
class MockClient:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
nonlocal closed_count
|
||||
closed_count += 1
|
||||
return False
|
||||
|
||||
def list_accounts(self):
|
||||
return [
|
||||
{"id": 1001, "name": "acc-1001", "credentials": {"api_key": "k1", "model_mapping": {"old": "old"}}},
|
||||
{"id": "abc", "name": "acc-abc", "credentials": {}},
|
||||
{"id": 1005, "name": "acc-1005", "credentials": {"api_key": "k5"}},
|
||||
{"id": 1006, "name": "acc-1006", "credentials": {"api_key": "k6"}},
|
||||
]
|
||||
|
||||
def extract_id(self, val):
|
||||
return str(val.get("id"))
|
||||
|
||||
def sync_account_upstream_models(self, account_id):
|
||||
sync_calls.append(account_id)
|
||||
if account_id == "1001":
|
||||
return ["gpt-4", "gpt-3.5-turbo", "", "gpt-4 "]
|
||||
elif account_id == "1005":
|
||||
return []
|
||||
elif account_id == "1006":
|
||||
raise Exception("Network Error")
|
||||
return []
|
||||
|
||||
def update_account(self, account_id, body):
|
||||
update_calls.append((account_id, body))
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClient())
|
||||
|
||||
res = sync_website_accounts_upstream_models(wid=w.id, db=db_session)
|
||||
|
||||
# 验证总体结果
|
||||
assert res.success is False # 存在 failed 账号，success 应该为 False
|
||||
assert "成功 1 个" in res.message
|
||||
assert "失败 2 个" in res.message
|
||||
assert "跳过 2 个" in res.message
|
||||
|
||||
items = res.items
|
||||
assert len(items) == 5
|
||||
|
||||
# 1001 成功
|
||||
item_1001 = next(i for i in items if i.account_id == "1001")
|
||||
assert item_1001.status == "success"
|
||||
assert item_1001.model_count == 2
|
||||
assert item_1001.models == ["gpt-3.5-turbo", "gpt-4"]
|
||||
|
||||
# 1001 的 update 应当 merge credentials
|
||||
assert len(update_calls) == 1
|
||||
assert update_calls[0][0] == "1001"
|
||||
assert update_calls[0][1]["credentials"]["api_key"] == "k1"
|
||||
assert update_calls[0][1]["credentials"]["model_mapping"] == {
|
||||
"gpt-3.5-turbo": "gpt-3.5-turbo",
|
||||
"gpt-4": "gpt-4"
|
||||
}
|
||||
|
||||
# abc 跳过
|
||||
item_abc = next(i for i in items if i.account_id == "abc")
|
||||
assert item_abc.status == "skipped"
|
||||
assert "非数字" in item_abc.message
|
||||
|
||||
# 1004 跳过
|
||||
item_1004 = next(i for i in items if i.account_id == "1004")
|
||||
assert item_1004.status == "skipped"
|
||||
assert "不存在" in item_1004.message
|
||||
|
||||
# 1005 失败 (空模型)
|
||||
item_1005 = next(i for i in items if i.account_id == "1005")
|
||||
assert item_1005.status == "failed"
|
||||
assert "为空" in item_1005.message
|
||||
|
||||
# 1006 失败 (抛错)
|
||||
item_1006 = next(i for i in items if i.account_id == "1006")
|
||||
assert item_1006.status == "failed"
|
||||
assert "Network Error" in item_1006.message
|
||||
|
||||
assert closed_count == 1
|
||||
|
||||
|
||||
def test_sync_upstream_models_list_accounts_none(db_session, monkeypatch):
|
||||
w = Website(
|
||||
id=1,
|
||||
name="Sub2Api site",
|
||||
site_type="sub2api",
|
||||
base_url="http://sub2api",
|
||||
api_prefix="api/v1",
|
||||
auth_type="bearer",
|
||||
auth_config_json='{"token": "tok1"}'
|
||||
)
|
||||
db_session.add(w)
|
||||
db_session.commit()
|
||||
|
||||
closed_count = 0
|
||||
|
||||
class MockClientNone:
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
nonlocal closed_count
|
||||
closed_count += 1
|
||||
return False
|
||||
|
||||
def list_accounts(self):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClientNone())
|
||||
|
||||
res = sync_website_accounts_upstream_models(wid=w.id, db=db_session)
|
||||
|
||||
assert res.success is False
|
||||
assert "拉取远端账号列表失败" in res.message
|
||||
assert len(res.items) == 0
|
||||
assert closed_count == 1
|
||||
Reference in New Issue
Block a user