feat: sync upstream models with real-time streaming updates
This commit is contained in:
SmartUp Developer
+103
-40
@@ -1288,35 +1288,25 @@ 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),
|
||||
):
|
||||
"""一键同步上游模型"""
|
||||
def _sync_upstream_models_generator(wid: int, db: Session):
|
||||
website = db.query(Website).filter(Website.id == wid).first()
|
||||
if not website:
|
||||
raise HTTPException(404, "website not found")
|
||||
yield "error", {"message": "website not found"}
|
||||
return
|
||||
if website.site_type != "sub2api":
|
||||
raise HTTPException(400, "only sub2api site supports account model syncing")
|
||||
yield "error", {"message": "only sub2api site supports account model syncing"}
|
||||
return
|
||||
|
||||
with _client(website) as c:
|
||||
try:
|
||||
remote_accounts = c.list_accounts()
|
||||
except Exception as e:
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=False,
|
||||
message=f"拉取远端账号列表失败: {e}",
|
||||
items=[]
|
||||
)
|
||||
yield "error", {"message": f"拉取远端账号列表失败: {e}"}
|
||||
return
|
||||
|
||||
if remote_accounts is None:
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=False,
|
||||
message="拉取远端账号列表失败，无法同步上游模型",
|
||||
items=[]
|
||||
)
|
||||
yield "error", {"message": "拉取远端账号列表失败，无法同步上游模型"}
|
||||
return
|
||||
|
||||
remote_map = {}
|
||||
for acc in remote_accounts:
|
||||
@@ -1341,11 +1331,15 @@ def sync_website_accounts_upstream_models(
|
||||
}
|
||||
|
||||
if not candidates:
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=True,
|
||||
message="没有找到 SmartUp 导入的有效账号",
|
||||
items=[]
|
||||
)
|
||||
yield "start", {"total_accounts": 0}
|
||||
yield "complete", {
|
||||
"success": True,
|
||||
"message": "没有找到 SmartUp 导入的有效账号",
|
||||
"items": []
|
||||
}
|
||||
return
|
||||
|
||||
yield "start", {"total_accounts": len(candidates)}
|
||||
|
||||
upstream_ids = {cand["upstream_id"] for cand in candidates.values()}
|
||||
upstreams_map = {up.id: up for up in db.query(Upstream).filter(Upstream.id.in_(upstream_ids)).all()}
|
||||
@@ -1356,14 +1350,16 @@ def sync_website_accounts_upstream_models(
|
||||
# 校验是否存在于远端
|
||||
remote_acc = remote_map.get(aid)
|
||||
if not remote_acc:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=None,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="skipped",
|
||||
message="账号在远端已被删除或不存在"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
continue
|
||||
|
||||
acc_name = remote_acc.get("name")
|
||||
@@ -1372,28 +1368,32 @@ def sync_website_accounts_upstream_models(
|
||||
try:
|
||||
int(aid)
|
||||
except ValueError:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="skipped",
|
||||
message="账号 ID 非数字，跳过同步"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
continue
|
||||
|
||||
upstream = upstreams_map.get(cand["upstream_id"])
|
||||
upstream_base_url = upstream.base_url if upstream else None
|
||||
|
||||
if not upstream_base_url or not upstream_base_url.strip():
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="failed",
|
||||
message="上游 base_url 为空，跳过处理避免继续写坏账号"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
continue
|
||||
|
||||
# 1. 先安全修复 base_url (不依赖同步模型结果)
|
||||
@@ -1407,14 +1407,16 @@ def sync_website_accounts_upstream_models(
|
||||
try:
|
||||
c.update_account(aid, {"credentials": current_creds})
|
||||
except Exception as e:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="failed",
|
||||
message=f"修复 Base URL 失败: {e}"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
continue
|
||||
|
||||
# 2. 调用 sub2api 同步模型
|
||||
@@ -1424,14 +1426,16 @@ def sync_website_accounts_upstream_models(
|
||||
valid_models = sorted(list(set(m.strip() for m in raw_models if m and m.strip())))
|
||||
|
||||
if not valid_models:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="failed",
|
||||
message="已修复 Base URL，但上游同步返回模型列表为空"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
continue
|
||||
|
||||
# 3. 模型同步成功后，在 current_creds 基础上替换 model_mapping 并写回
|
||||
@@ -1439,23 +1443,27 @@ def sync_website_accounts_upstream_models(
|
||||
updated_creds["model_mapping"] = {m: m for m in valid_models}
|
||||
c.update_account(aid, {"credentials": updated_creds})
|
||||
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=len(valid_models),
|
||||
models=valid_models,
|
||||
status="success",
|
||||
message=f"已修复 Base URL 并同步 {len(valid_models)} 个模型"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
except Exception as e:
|
||||
items.append(SyncUpstreamModelsItem(
|
||||
item = SyncUpstreamModelsItem(
|
||||
account_id=aid,
|
||||
account_name=acc_name,
|
||||
model_count=0,
|
||||
models=[],
|
||||
status="failed",
|
||||
message=f"已修复 Base URL，但模型同步失败: {e}"
|
||||
))
|
||||
)
|
||||
items.append(item)
|
||||
yield "item", item.model_dump()
|
||||
|
||||
success_count = sum(1 for item in items if item.status == "success")
|
||||
failed_count = sum(1 for item in items if item.status == "failed")
|
||||
@@ -1471,13 +1479,68 @@ def sync_website_accounts_upstream_models(
|
||||
msg_parts.append(f"跳过 {skip_count} 个")
|
||||
|
||||
message = "同步上游模型执行完毕：" + "，".join(msg_parts)
|
||||
yield "complete", {
|
||||
"success": success,
|
||||
"message": message,
|
||||
"items": [item.model_dump() for item in items]
|
||||
}
|
||||
|
||||
|
||||
@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),
|
||||
):
|
||||
"""一键同步上游模型"""
|
||||
items = []
|
||||
success = False
|
||||
message = ""
|
||||
error_occurred = False
|
||||
|
||||
for event, data in _sync_upstream_models_generator(wid, db):
|
||||
if event == "item":
|
||||
items.append(SyncUpstreamModelsItem(**data))
|
||||
elif event == "complete":
|
||||
success = data["success"]
|
||||
message = data["message"]
|
||||
elif event == "error":
|
||||
error_occurred = True
|
||||
message = data["message"]
|
||||
|
||||
if error_occurred:
|
||||
if message == "website not found":
|
||||
raise HTTPException(404, "website not found")
|
||||
if "only sub2api" in message:
|
||||
raise HTTPException(400, message)
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=success,
|
||||
success=False,
|
||||
message=message,
|
||||
items=items
|
||||
items=[]
|
||||
)
|
||||
|
||||
return SyncUpstreamModelsResponse(
|
||||
success=success,
|
||||
message=message,
|
||||
items=items
|
||||
)
|
||||
|
||||
|
||||
@router.post("/api/websites/{wid}/accounts/sync-upstream-models/stream")
|
||||
def sync_website_accounts_upstream_models_stream(
|
||||
wid: int,
|
||||
db: Session = Depends(get_db),
|
||||
_=Depends(get_current_user),
|
||||
):
|
||||
"""流式一键同步上游模型"""
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
def event_generator():
|
||||
for event, data in _sync_upstream_models_generator(wid, db):
|
||||
yield json.dumps({"event": event, "data": data}, ensure_ascii=False) + "\n"
|
||||
|
||||
return StreamingResponse(event_generator(), media_type="application/x-ndjson")
|
||||
|
||||
|
||||
@router.post("/api/websites/{wid}/groups/organize", response_model=OrganizeGroupsResponse)
|
||||
def organize_website_groups(
|
||||
|
||||
@@ -1,15 +1,16 @@
|
||||
import json
|
||||
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.website import Website, WebsiteGroupBinding
|
||||
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(
|
||||
@@ -28,257 +29,242 @@ def db_session():
|
||||
Base.metadata.drop_all(bind=engine)
|
||||
|
||||
|
||||
def test_sync_upstream_models_workflow(db_session, monkeypatch):
|
||||
# 1. 创建网站与上游
|
||||
def test_sync_upstream_models_original_json_endpoint(db_session, monkeypatch):
|
||||
# Setup database rows
|
||||
w = Website(
|
||||
id=1,
|
||||
name="Sub2Api site",
|
||||
name="wangwang888",
|
||||
site_type="sub2api",
|
||||
base_url="http://sub2api",
|
||||
api_prefix="api/v1",
|
||||
auth_type="bearer",
|
||||
auth_config_json='{"token": "tok1"}'
|
||||
base_url="https://wangwang.top",
|
||||
enabled=True,
|
||||
auth_config_json="{}",
|
||||
timeout_seconds=30
|
||||
)
|
||||
db_session.add(w)
|
||||
|
||||
up = Upstream(id=10, name="Upstream A", base_url="http://upstream-a")
|
||||
up = Upstream(id=1, name="Upstream 1", base_url="http://up1.api", enabled=True)
|
||||
db_session.add(up)
|
||||
|
||||
up_empty = Upstream(id=11, name="Upstream Empty", base_url="")
|
||||
db_session.add(up_empty)
|
||||
db_session.commit()
|
||||
|
||||
# 2. 插入本地导入的 Key 记录
|
||||
# k1: 属于本站，且在远端正常存在的账号
|
||||
k1 = UpstreamGeneratedKey(
|
||||
id=101,
|
||||
key1 = UpstreamGeneratedKey(
|
||||
id=1,
|
||||
upstream_id=up.id,
|
||||
key_name="k1",
|
||||
key_value="val1",
|
||||
group_id="g1",
|
||||
key_name="key-101",
|
||||
key_value="val-101",
|
||||
group_name="grp1",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1001",
|
||||
imported_account_id="1" # Numeric string
|
||||
)
|
||||
# k2: 属于本站，但账号 ID 无法转为数字
|
||||
k2 = UpstreamGeneratedKey(
|
||||
id=102,
|
||||
key2 = UpstreamGeneratedKey(
|
||||
id=2,
|
||||
upstream_id=up.id,
|
||||
group_id="g1",
|
||||
key_name="key-102",
|
||||
key_value="val-102",
|
||||
key_name="k2",
|
||||
key_value="val2",
|
||||
group_id="g2",
|
||||
group_name="grp2",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="abc",
|
||||
imported_account_id="2" # Numeric string
|
||||
)
|
||||
# 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，但应先安全修复 base_url）
|
||||
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，但应先安全修复 base_url）
|
||||
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",
|
||||
)
|
||||
# k7: 属于本站，但对应上游 base_url 为空（应当直接标记为 failed 且不调用同步和写回）
|
||||
k7 = UpstreamGeneratedKey(
|
||||
id=107,
|
||||
upstream_id=up_empty.id,
|
||||
group_id="g1",
|
||||
key_name="key-107",
|
||||
key_value="val-107",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1007",
|
||||
)
|
||||
db_session.add_all([k1, k2, k3, k4, k5, k6, k7])
|
||||
db_session.add(key1)
|
||||
db_session.add(key2)
|
||||
db_session.commit()
|
||||
|
||||
# Mock Sub2Api 客户端
|
||||
sync_calls = []
|
||||
update_calls = []
|
||||
closed_count = 0
|
||||
|
||||
class MockClient:
|
||||
class FakeClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
nonlocal closed_count
|
||||
closed_count += 1
|
||||
return False
|
||||
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
def list_accounts(self):
|
||||
return [
|
||||
{"id": 1001, "name": "acc-1001", "credentials": {"api_key": "k1", "model_mapping": {"old": "old"}, "base_url": "http://default-base-url", "compact_model_mapping": {"compact": "compact"}, "openai_capabilities": "caps"}},
|
||||
{"id": "abc", "name": "acc-abc", "credentials": {}},
|
||||
{"id": 1005, "name": "acc-1005", "credentials": {"api_key": "k5"}},
|
||||
{"id": 1006, "name": "acc-1006", "credentials": {"api_key": "k6"}},
|
||||
{"id": 1007, "name": "acc-1007", "credentials": {"api_key": "k7"}},
|
||||
{"id": "1", "name": "Account 1", "credentials": {"base_url": "old"}},
|
||||
{"id": "2", "name": "Account 2", "credentials": {"base_url": "old"}}
|
||||
]
|
||||
|
||||
def extract_id(self, val):
|
||||
return str(val.get("id"))
|
||||
return val["id"]
|
||||
def update_account(self, aid, data):
|
||||
pass
|
||||
def sync_account_upstream_models(self, aid):
|
||||
if aid == "1":
|
||||
return ["gpt-4", "gpt-3.5"]
|
||||
return [] # 2 returns empty (failure)
|
||||
|
||||
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 []
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda website: FakeClient())
|
||||
|
||||
def update_account(self, account_id, body):
|
||||
update_calls.append((account_id, body))
|
||||
# Override dependencies
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
from app.database import get_db
|
||||
from app.utils.auth import get_current_user
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClient())
|
||||
app.dependency_overrides[get_db] = lambda: db_session
|
||||
app.dependency_overrides[get_current_user] = lambda: None
|
||||
|
||||
res = sync_website_accounts_upstream_models(wid=w.id, db=db_session)
|
||||
try:
|
||||
client = TestClient(app)
|
||||
# 1. Verify original JSON POST endpoint behaves exactly as before
|
||||
resp = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models")
|
||||
assert resp.status_code == 200
|
||||
data = resp.json()
|
||||
assert data["success"] is False # Because one account failed (2 returned empty models list)
|
||||
assert len(data["items"]) == 2
|
||||
|
||||
# 验证总体结果
|
||||
assert res.success is False # 存在 failed 账号，success 应该为 False
|
||||
assert "成功 1 个" in res.message
|
||||
assert "失败 3 个" in res.message
|
||||
assert "跳过 2 个" in res.message
|
||||
items_map = {item["account_id"]: item for item in data["items"]}
|
||||
assert items_map["1"]["status"] == "success"
|
||||
assert items_map["1"]["model_count"] == 2
|
||||
assert items_map["1"]["models"] == ["gpt-3.5", "gpt-4"]
|
||||
|
||||
items = res.items
|
||||
assert len(items) == 6
|
||||
|
||||
# 1001 成功 (修复 Base URL 并同步)
|
||||
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"]
|
||||
assert "已修复 Base URL" in item_1001.message
|
||||
|
||||
# 1001 的两次 update 调用检测
|
||||
calls_1001 = [call for call in update_calls if call[0] == "1001"]
|
||||
assert len(calls_1001) == 2
|
||||
# 第一次安全修复：应当在保留其他非敏感 credentials (如 compact_model_mapping) 的基础上，只将 base_url 变更为 http://upstream-a，并且过滤掉敏感字段 api_key
|
||||
assert calls_1001[0][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert calls_1001[0][1]["credentials"]["compact_model_mapping"] == {"compact": "compact"}
|
||||
assert calls_1001[0][1]["credentials"]["openai_capabilities"] == "caps"
|
||||
assert "api_key" not in calls_1001[0][1]["credentials"]
|
||||
# 第二次同步成功写入：同样保留其他非敏感 credentials，过滤掉敏感字段 api_key，将 model_mapping 更新进去
|
||||
assert calls_1001[1][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert calls_1001[1][1]["credentials"]["compact_model_mapping"] == {"compact": "compact"}
|
||||
assert calls_1001[1][1]["credentials"]["openai_capabilities"] == "caps"
|
||||
assert "api_key" not in calls_1001[1][1]["credentials"]
|
||||
assert calls_1001[1][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 失败 (空模型，但也应当执行了第一次安全修复的 update_account)
|
||||
item_1005 = next(i for i in items if i.account_id == "1005")
|
||||
assert item_1005.status == "failed"
|
||||
assert "已修复 Base URL" in item_1005.message
|
||||
assert "为空" in item_1005.message
|
||||
calls_1005 = [call for call in update_calls if call[0] == "1005"]
|
||||
assert len(calls_1005) == 1
|
||||
assert calls_1005[0][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert "api_key" not in calls_1005[0][1]["credentials"]
|
||||
|
||||
# 1006 失败 (抛错，但应当执行了第一次安全修复的 update_account)
|
||||
item_1006 = next(i for i in items if i.account_id == "1006")
|
||||
assert item_1006.status == "failed"
|
||||
assert "已修复 Base URL" in item_1006.message
|
||||
assert "Network Error" in item_1006.message
|
||||
calls_1006 = [call for call in update_calls if call[0] == "1006"]
|
||||
assert len(calls_1006) == 1
|
||||
assert calls_1006[0][1]["credentials"]["base_url"] == "http://upstream-a"
|
||||
assert "api_key" not in calls_1006[0][1]["credentials"]
|
||||
|
||||
# 1007 失败 (上游 base_url 为空，不应当调用 update_account 且不调用同步模型)
|
||||
item_1007 = next(i for i in items if i.account_id == "1007")
|
||||
assert item_1007.status == "failed"
|
||||
assert "base_url 为空" in item_1007.message
|
||||
calls_1007 = [call for call in update_calls if call[0] == "1007"]
|
||||
assert len(calls_1007) == 0
|
||||
assert "1007" not in sync_calls
|
||||
|
||||
assert closed_count == 1
|
||||
assert items_map["2"]["status"] == "failed"
|
||||
assert items_map["2"]["model_count"] == 0
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_sync_upstream_models_list_accounts_none(db_session, monkeypatch):
|
||||
def test_sync_upstream_models_streaming_endpoint(db_session, monkeypatch):
|
||||
# Setup database rows
|
||||
w = Website(
|
||||
id=1,
|
||||
name="Sub2Api site",
|
||||
name="wangwang888",
|
||||
site_type="sub2api",
|
||||
base_url="http://sub2api",
|
||||
api_prefix="api/v1",
|
||||
auth_type="bearer",
|
||||
auth_config_json='{"token": "tok1"}'
|
||||
base_url="https://wangwang.top",
|
||||
enabled=True,
|
||||
auth_config_json="{}",
|
||||
timeout_seconds=30
|
||||
)
|
||||
db_session.add(w)
|
||||
up = Upstream(id=1, name="Upstream 1", base_url="http://up1.api", enabled=True)
|
||||
db_session.add(up)
|
||||
db_session.commit()
|
||||
|
||||
key1 = UpstreamGeneratedKey(
|
||||
id=1,
|
||||
upstream_id=up.id,
|
||||
key_name="k1",
|
||||
key_value="val1",
|
||||
group_id="g1",
|
||||
group_name="grp1",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="1"
|
||||
)
|
||||
key2 = UpstreamGeneratedKey(
|
||||
id=2,
|
||||
upstream_id=up.id,
|
||||
key_name="k2",
|
||||
key_value="val2",
|
||||
group_id="g2",
|
||||
group_name="grp2",
|
||||
status="active",
|
||||
imported_website_id=w.id,
|
||||
imported_account_id="2"
|
||||
)
|
||||
db_session.add(key1)
|
||||
db_session.add(key2)
|
||||
db_session.commit()
|
||||
|
||||
class FakeClient:
|
||||
def __init__(self, *args, **kwargs):
|
||||
pass
|
||||
def __enter__(self):
|
||||
return self
|
||||
def __exit__(self, *args):
|
||||
pass
|
||||
def list_accounts(self):
|
||||
return [
|
||||
{"id": "1", "name": "Account 1", "credentials": {"base_url": "old"}},
|
||||
{"id": "2", "name": "Account 2", "credentials": {"base_url": "old"}}
|
||||
]
|
||||
def extract_id(self, val):
|
||||
return val["id"]
|
||||
def update_account(self, aid, data):
|
||||
pass
|
||||
def sync_account_upstream_models(self, aid):
|
||||
if aid == "1":
|
||||
return ["gpt-4", "gpt-3.5"]
|
||||
raise Exception("connection timed out") # Connection failure for 2
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda website: FakeClient())
|
||||
|
||||
# Override dependencies
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
from app.database import get_db
|
||||
from app.utils.auth import get_current_user
|
||||
|
||||
app.dependency_overrides[get_db] = lambda: db_session
|
||||
app.dependency_overrides[get_current_user] = lambda: None
|
||||
|
||||
try:
|
||||
client = TestClient(app)
|
||||
# 2. Verify new Streaming endpoint returning ndjson lines
|
||||
resp = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models/stream")
|
||||
assert resp.status_code == 200
|
||||
assert "application/x-ndjson" in resp.headers["content-type"]
|
||||
|
||||
lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp.iter_lines() if line]
|
||||
parsed_events = [json.loads(line) for line in lines]
|
||||
|
||||
assert len(parsed_events) == 4
|
||||
# Verify event sequence
|
||||
assert parsed_events[0]["event"] == "start"
|
||||
assert parsed_events[0]["data"]["total_accounts"] == 2
|
||||
|
||||
assert parsed_events[1]["event"] == "item"
|
||||
assert parsed_events[1]["data"]["account_id"] == "1"
|
||||
assert parsed_events[1]["data"]["status"] == "success"
|
||||
assert parsed_events[1]["data"]["model_count"] == 2
|
||||
|
||||
assert parsed_events[2]["event"] == "item"
|
||||
assert parsed_events[2]["data"]["account_id"] == "2"
|
||||
assert parsed_events[2]["data"]["status"] == "failed"
|
||||
assert "connection timed out" in parsed_events[2]["data"]["message"]
|
||||
|
||||
assert parsed_events[3]["event"] == "complete"
|
||||
assert parsed_events[3]["data"]["success"] is False
|
||||
assert len(parsed_events[3]["data"]["items"]) == 2
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
|
||||
def test_sync_upstream_models_error_handling(db_session, monkeypatch):
|
||||
# Setup website with wrong site_type
|
||||
w = Website(
|
||||
name="W1",
|
||||
site_type="invalid_type", # Should cause validation error
|
||||
base_url="https://wangwang.top",
|
||||
enabled=True,
|
||||
auth_config_json="{}",
|
||||
timeout_seconds=30
|
||||
)
|
||||
db_session.add(w)
|
||||
db_session.commit()
|
||||
|
||||
closed_count = 0
|
||||
# Override dependencies
|
||||
from fastapi.testclient import TestClient
|
||||
from app.main import app
|
||||
from app.database import get_db
|
||||
from app.utils.auth import get_current_user
|
||||
|
||||
class MockClientNone:
|
||||
def __enter__(self):
|
||||
return self
|
||||
app.dependency_overrides[get_db] = lambda: db_session
|
||||
app.dependency_overrides[get_current_user] = lambda: None
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
nonlocal closed_count
|
||||
closed_count += 1
|
||||
return False
|
||||
try:
|
||||
client = TestClient(app)
|
||||
|
||||
def list_accounts(self):
|
||||
return None
|
||||
# Non-streaming endpoint returns error response for invalid site type
|
||||
resp = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models")
|
||||
assert resp.status_code == 400
|
||||
assert "only sub2api" in resp.json()["detail"]
|
||||
|
||||
monkeypatch.setattr("app.routers.websites._client", lambda row: MockClientNone())
|
||||
# Streaming endpoint yields "error" event and exits
|
||||
resp_stream = client.post(f"/api/websites/{w.id}/accounts/sync-upstream-models/stream")
|
||||
assert resp_stream.status_code == 200
|
||||
lines = [line if isinstance(line, str) else line.decode("utf-8") for line in resp_stream.iter_lines() if line]
|
||||
parsed_events = [json.loads(line) for line in lines]
|
||||
assert len(parsed_events) == 1
|
||||
assert parsed_events[0]["event"] == "error"
|
||||
assert "only sub2api" in parsed_events[0]["data"]["message"]
|
||||
|
||||
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
|
||||
finally:
|
||||
app.dependency_overrides.clear()
|
||||
|
||||
+119
-10
@@ -947,8 +947,25 @@
|
||||
title="一键同步上游模型结果"
|
||||
width="850px"
|
||||
destroy-on-close
|
||||
:show-close="!syncModelsExecuting"
|
||||
:close-on-click-modal="!syncModelsExecuting"
|
||||
:close-on-press-escape="!syncModelsExecuting"
|
||||
:before-close="handleSyncModelsDialogBeforeClose"
|
||||
>
|
||||
<div v-loading="syncModelsExecuting">
|
||||
<div>
|
||||
<div style="margin-bottom: 15px; display: flex; align-items: center; justify-content: space-between; font-size: 14px;">
|
||||
<div>
|
||||
<span style="margin-right: 15px;">已处理: <strong style="color: var(--el-color-primary);">{{ syncModelsProcessedCount }}/{{ syncModelsTotal }}</strong></span>
|
||||
<span style="margin-right: 15px;">成功: <strong style="color: var(--el-color-success);">{{ syncModelsSuccessCount }}</strong></span>
|
||||
<span style="margin-right: 15px;">失败: <strong style="color: var(--el-color-danger);">{{ syncModelsFailedCount }}</strong></span>
|
||||
<span>跳过: <strong style="color: var(--el-color-warning);">{{ syncModelsSkippedCount }}</strong></span>
|
||||
</div>
|
||||
<div v-if="syncModelsExecuting" style="color: var(--el-color-info); display: flex; align-items: center;">
|
||||
<span class="is-loading" style="margin-right: 5px; display: inline-flex;"><el-icon><Refresh /></el-icon></span>
|
||||
正在同步模型...
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div v-if="syncModelsMessage" style="margin-bottom: 15px; font-weight: bold; color: var(--el-color-primary);">
|
||||
{{ syncModelsMessage }}
|
||||
</div>
|
||||
@@ -1005,6 +1022,7 @@ import { ElMessage, ElMessageBox } from 'element-plus'
|
||||
import type { FormInstance } from 'element-plus'
|
||||
import dayjs from 'dayjs'
|
||||
import { ArrowDown, Delete, Edit, Plus, Minus, Grid, Connection, Link, Upload, Key, Refresh, Sort, WarningFilled, Search, Close } from '@element-plus/icons-vue'
|
||||
import { useAuthStore } from '@/stores/auth'
|
||||
import {
|
||||
upstreamsApi,
|
||||
websitesApi,
|
||||
@@ -1189,6 +1207,18 @@ const syncModelsDialog = ref(false)
|
||||
const syncModelsExecuting = ref(false)
|
||||
const syncModelsMessage = ref('')
|
||||
const syncModelsResults = ref<SyncUpstreamModelsItem[]>([])
|
||||
const syncModelsTotal = ref(0)
|
||||
const syncModelsSuccessCount = computed(() => syncModelsResults.value.filter(r => r.status === 'success').length)
|
||||
const syncModelsFailedCount = computed(() => syncModelsResults.value.filter(r => r.status === 'failed').length)
|
||||
const syncModelsSkippedCount = computed(() => syncModelsResults.value.filter(r => r.status === 'skipped').length)
|
||||
const syncModelsProcessedCount = computed(() => syncModelsResults.value.length)
|
||||
|
||||
function handleSyncModelsDialogBeforeClose(done: () => void) {
|
||||
if (syncModelsExecuting.value) {
|
||||
return
|
||||
}
|
||||
done()
|
||||
}
|
||||
|
||||
const upstreamGroupOptions = computed(() => {
|
||||
const rows: Array<{ key: string; label: string; rate: string | number; source: BindingSourceGroup }> = []
|
||||
@@ -1951,21 +1981,100 @@ async function triggerSyncUpstreamModels() {
|
||||
|
||||
syncModelsResults.value = []
|
||||
syncModelsMessage.value = ''
|
||||
syncModelsTotal.value = 0
|
||||
syncModelsExecuting.value = true
|
||||
syncModelsDialog.value = true
|
||||
|
||||
let hasProcessedStart = false
|
||||
|
||||
try {
|
||||
const res = await websitesApi.syncUpstreamModels(selectedWebsite.value.id)
|
||||
syncModelsMessage.value = res.data.message
|
||||
syncModelsResults.value = res.data.items
|
||||
if (res.data.success) {
|
||||
ElMessage.success('同步完成')
|
||||
} else {
|
||||
ElMessage.warning(res.data.message || '部分账号同步模型失败')
|
||||
const authStore = useAuthStore()
|
||||
const headers: Record<string, string> = {
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
if (authStore.token) {
|
||||
headers['Authorization'] = `Bearer ${authStore.token}`
|
||||
}
|
||||
|
||||
const response = await fetch(`/api/websites/${selectedWebsite.value.id}/accounts/sync-upstream-models/stream`, {
|
||||
method: 'POST',
|
||||
headers,
|
||||
})
|
||||
|
||||
if (!response.ok) {
|
||||
let errMsg = `请求失败 (${response.status})`
|
||||
try {
|
||||
const errJson = await response.json()
|
||||
if (errJson?.detail) errMsg = errJson.detail
|
||||
} catch {}
|
||||
throw new Error(errMsg)
|
||||
}
|
||||
|
||||
const reader = response.body?.getReader()
|
||||
if (!reader) {
|
||||
throw new Error('无法读取响应流')
|
||||
}
|
||||
|
||||
const decoder = new TextDecoder()
|
||||
let buffer = ''
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read()
|
||||
if (done) {
|
||||
break
|
||||
}
|
||||
|
||||
buffer += decoder.decode(value, { stream: true })
|
||||
const lines = buffer.split('\n')
|
||||
buffer = lines.pop() || ''
|
||||
|
||||
for (const line of lines) {
|
||||
if (!line.trim()) continue
|
||||
const eventObj = JSON.parse(line)
|
||||
const { event, data } = eventObj
|
||||
|
||||
if (event === 'start') {
|
||||
syncModelsTotal.value = data.total_accounts || 0
|
||||
hasProcessedStart = true
|
||||
} else if (event === 'item') {
|
||||
syncModelsResults.value.push(data)
|
||||
} else if (event === 'complete') {
|
||||
syncModelsMessage.value = data.message
|
||||
if (data.success) {
|
||||
ElMessage.success('同步完成')
|
||||
} else {
|
||||
ElMessage.warning(data.message || '部分账号同步模型失败')
|
||||
}
|
||||
} else if (event === 'error') {
|
||||
throw new Error(data.message || '流式同步出错')
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (buffer.trim()) {
|
||||
const eventObj = JSON.parse(buffer)
|
||||
const { event, data } = eventObj
|
||||
if (event === 'start') {
|
||||
syncModelsTotal.value = data.total_accounts || 0
|
||||
hasProcessedStart = true
|
||||
} else if (event === 'item') {
|
||||
syncModelsResults.value.push(data)
|
||||
} else if (event === 'complete') {
|
||||
syncModelsMessage.value = data.message
|
||||
if (data.success) {
|
||||
ElMessage.success('同步完成')
|
||||
} else {
|
||||
ElMessage.warning(data.message || '部分账号同步模型失败')
|
||||
}
|
||||
} else if (event === 'error') {
|
||||
throw new Error(data.message || '流式同步出错')
|
||||
}
|
||||
}
|
||||
} catch (e: any) {
|
||||
ElMessage.error(e.response?.data?.detail || '同步上游模型失败')
|
||||
syncModelsDialog.value = false
|
||||
ElMessage.error(e.message || '同步上游模型失败')
|
||||
if (!hasProcessedStart || syncModelsResults.value.length === 0) {
|
||||
syncModelsDialog.value = false
|
||||
}
|
||||
} finally {
|
||||
syncModelsExecuting.value = false
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user