273 lines
10 KiB
Python
273 lines
10 KiB
Python
"""供应商与模型的服务层：CRUD、密钥加解密、模型发现与落库。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
from typing import Any
|
|
|
|
from .. import db, security
|
|
from ..providers import ADAPTER_LABELS, ProviderConfig, ProviderError, get_adapter
|
|
|
|
|
|
def now() -> str:
|
|
return datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def _row_to_dict(row) -> dict[str, Any]:
|
|
models = _models_of(row["id"])
|
|
return {
|
|
"id": row["id"],
|
|
"name": row["name"],
|
|
"adapter": row["adapter"],
|
|
"adapter_label": ADAPTER_LABELS.get(row["adapter"], row["adapter"]),
|
|
"base_url": row["base_url"],
|
|
"api_key_masked": security.mask(security.decrypt(row["api_key"] or "")),
|
|
"default_model": row["default_model"],
|
|
"enabled": bool(row["enabled"]),
|
|
"note": row["note"],
|
|
"model_count": len(models),
|
|
"enabled_model_count": sum(1 for m in models if m["enabled"]),
|
|
"models_discovered": _latest_discovered(row["id"]),
|
|
"created_at": row["created_at"],
|
|
"updated_at": row["updated_at"],
|
|
}
|
|
|
|
|
|
def _models_of(provider_id: str) -> list[dict[str, Any]]:
|
|
with db.session() as conn:
|
|
rows = conn.execute(
|
|
"SELECT * FROM models WHERE provider_id = ? ORDER BY model_id", (provider_id,)
|
|
).fetchall()
|
|
return [_model_to_dict(r) for r in rows]
|
|
|
|
|
|
def _model_to_dict(row) -> dict[str, Any]:
|
|
return {
|
|
"id": row["id"],
|
|
"provider_id": row["provider_id"],
|
|
"model_id": row["model_id"],
|
|
"display_name": row["display_name"],
|
|
"enabled": bool(row["enabled"]),
|
|
"kind": row["kind"],
|
|
"size_hint": row["size_hint"],
|
|
"supports_image_input": bool(row["supports_image_input"]),
|
|
"note": row["note"],
|
|
"discovered_at": row["discovered_at"],
|
|
}
|
|
|
|
|
|
def _latest_discovered(provider_id: str) -> str:
|
|
with db.session() as conn:
|
|
row = conn.execute(
|
|
"SELECT MAX(discovered_at) AS d FROM models WHERE provider_id = ?", (provider_id,)
|
|
).fetchone()
|
|
return (row["d"] if row and row["d"] else "—") or "—"
|
|
|
|
|
|
# --------------------------------------------------------------- 供应商 CRUD
|
|
def list_providers() -> list[dict[str, Any]]:
|
|
with db.session() as conn:
|
|
rows = conn.execute("SELECT * FROM providers ORDER BY created_at, id").fetchall()
|
|
return [_row_to_dict(r) for r in rows]
|
|
|
|
|
|
def get_provider_row(provider_id: str):
|
|
with db.session() as conn:
|
|
return conn.execute("SELECT * FROM providers WHERE id = ?", (provider_id,)).fetchone()
|
|
|
|
|
|
def get_provider(provider_id: str) -> dict[str, Any] | None:
|
|
row = get_provider_row(provider_id)
|
|
return _row_to_dict(row) if row else None
|
|
|
|
|
|
def save_provider(data: dict[str, Any]) -> dict[str, Any]:
|
|
"""新增或更新供应商。api_key 留空表示不修改。"""
|
|
pid = data["id"].strip()
|
|
ts = now()
|
|
existing = get_provider_row(pid)
|
|
if existing:
|
|
api_key_enc = existing["api_key"]
|
|
if data.get("api_key"):
|
|
api_key_enc = security.encrypt(data["api_key"])
|
|
with db.session() as conn:
|
|
conn.execute(
|
|
"UPDATE providers SET name=?, adapter=?, base_url=?, api_key=?, default_model=?,"
|
|
" enabled=?, note=?, updated_at=? WHERE id=?",
|
|
(
|
|
data["name"],
|
|
data["adapter"],
|
|
data["base_url"],
|
|
api_key_enc,
|
|
data.get("default_model", ""),
|
|
1 if data.get("enabled", True) else 0,
|
|
data.get("note", ""),
|
|
ts,
|
|
pid,
|
|
),
|
|
)
|
|
else:
|
|
with db.session() as conn:
|
|
conn.execute(
|
|
"INSERT INTO providers (id,name,adapter,base_url,api_key,default_model,enabled,"
|
|
"note,created_at,updated_at) VALUES (?,?,?,?,?,?,?,?,?,?)",
|
|
(
|
|
pid,
|
|
data["name"],
|
|
data["adapter"],
|
|
data["base_url"],
|
|
security.encrypt(data.get("api_key", "")),
|
|
data.get("default_model", ""),
|
|
1 if data.get("enabled", True) else 0,
|
|
data.get("note", ""),
|
|
ts,
|
|
ts,
|
|
),
|
|
)
|
|
return get_provider(pid)
|
|
|
|
|
|
def delete_provider(provider_id: str) -> None:
|
|
with db.session() as conn:
|
|
conn.execute("DELETE FROM providers WHERE id = ?", (provider_id,))
|
|
|
|
|
|
def provider_config(provider_id: str, *, require_enabled: bool = False) -> ProviderConfig:
|
|
row = get_provider_row(provider_id)
|
|
if row is None:
|
|
raise ProviderError(f"供应商不存在: {provider_id}", status=404)
|
|
if require_enabled and not row["enabled"]:
|
|
raise ProviderError("供应商已停用", status=400)
|
|
return ProviderConfig(
|
|
id=row["id"],
|
|
name=row["name"],
|
|
adapter=row["adapter"],
|
|
base_url=row["base_url"],
|
|
api_key=security.decrypt(row["api_key"] or ""),
|
|
default_model=row["default_model"],
|
|
)
|
|
|
|
|
|
# ----------------------------------------------------------------- 模型
|
|
def list_models(provider_id: str | None = None) -> list[dict[str, Any]]:
|
|
with db.session() as conn:
|
|
if provider_id:
|
|
rows = conn.execute(
|
|
"SELECT * FROM models WHERE provider_id = ? ORDER BY model_id", (provider_id,)
|
|
).fetchall()
|
|
else:
|
|
rows = conn.execute("SELECT * FROM models ORDER BY provider_id, model_id").fetchall()
|
|
return [_model_to_dict(r) for r in rows]
|
|
|
|
|
|
def set_model_enabled(model_id: str, enabled: bool) -> dict[str, Any]:
|
|
with db.session() as conn:
|
|
conn.execute(
|
|
"UPDATE models SET enabled=?, updated_at=? WHERE id=?", (1 if enabled else 0, now(), model_id)
|
|
)
|
|
return get_model(model_id)
|
|
|
|
|
|
def set_model_note(model_id: str, note: str) -> dict[str, Any]:
|
|
with db.session() as conn:
|
|
conn.execute(
|
|
"UPDATE models SET note=?, updated_at=? WHERE id=?", (note, now(), model_id)
|
|
)
|
|
return get_model(model_id)
|
|
|
|
|
|
def set_models_enabled(model_ids: list[str], enabled: bool) -> list[dict[str, Any]]:
|
|
ts = now()
|
|
with db.session() as conn:
|
|
for mid in model_ids:
|
|
conn.execute(
|
|
"UPDATE models SET enabled=?, updated_at=? WHERE id=?", (1 if enabled else 0, ts, mid)
|
|
)
|
|
return list_models()
|
|
|
|
|
|
def get_model(model_id: str) -> dict[str, Any]:
|
|
with db.session() as conn:
|
|
row = conn.execute("SELECT * FROM models WHERE id = ?", (model_id,)).fetchone()
|
|
if row is None:
|
|
raise ProviderError(f"模型不存在: {model_id}", status=404)
|
|
return _model_to_dict(row)
|
|
|
|
|
|
def upsert_models(
|
|
provider_id: str,
|
|
model_ids: list[str],
|
|
enabled: bool = True,
|
|
note: str = "自动发现",
|
|
exclusive: bool = False,
|
|
preserve_enabled: bool = False,
|
|
) -> list[dict[str, Any]]:
|
|
"""批量写入模型；已存在的仅更新启用态与发现时间。
|
|
|
|
exclusive=True 时，把该供应商下未被选中的模型置为停用。
|
|
"""
|
|
cfg = provider_config(provider_id)
|
|
capabilities = get_adapter(cfg).capabilities
|
|
ts = now()
|
|
today = ts[:10]
|
|
with db.session() as conn:
|
|
if exclusive:
|
|
keep = {f"{provider_id}:{m}" for m in model_ids}
|
|
for row in conn.execute(
|
|
"SELECT id FROM models WHERE provider_id = ?", (provider_id,)
|
|
).fetchall():
|
|
if row["id"] not in keep:
|
|
conn.execute(
|
|
"UPDATE models SET enabled=0, updated_at=? WHERE id=?", (ts, row["id"])
|
|
)
|
|
for mid in model_ids:
|
|
model_id = f"{provider_id}:{mid}"
|
|
exists = conn.execute("SELECT id FROM models WHERE id = ?", (model_id,)).fetchone()
|
|
if exists:
|
|
conn.execute(
|
|
"UPDATE models SET enabled=CASE WHEN ? THEN enabled ELSE ? END,"
|
|
" discovered_at=?, updated_at=?, kind=?, size_hint=?, supports_image_input=? WHERE id=?",
|
|
(preserve_enabled, int(enabled), today, ts, capabilities.kind,
|
|
capabilities.size_hint, int(capabilities.supports_image_input), model_id),
|
|
)
|
|
else:
|
|
conn.execute(
|
|
"INSERT INTO models (id,provider_id,model_id,display_name,enabled,kind,size_hint,"
|
|
"supports_image_input,note,discovered_at,created_at,updated_at)"
|
|
" VALUES (?,?,?,?,?,?,?,?,?,?,?,?)",
|
|
(model_id, provider_id, mid, mid, int(enabled), capabilities.kind,
|
|
capabilities.size_hint, int(capabilities.supports_image_input), note, today, ts, ts),
|
|
)
|
|
return list_models(provider_id)
|
|
|
|
|
|
# ------------------------------------------------------------- 模型发现
|
|
def discover(provider_id: str, *, require_enabled: bool = False) -> dict[str, Any]:
|
|
"""按已保存的供应商拉取上游模型，返回 {provider_id,total,matched}。"""
|
|
cfg = provider_config(provider_id, require_enabled=require_enabled)
|
|
adapter = get_adapter(cfg)
|
|
all_ids = adapter.list_models()
|
|
matched = adapter.filter_image_models(all_ids)
|
|
upsert_models(provider_id, matched, enabled=True, preserve_enabled=True)
|
|
if require_enabled:
|
|
enabled = {m["model_id"] for m in list_models(provider_id) if m["enabled"]}
|
|
matched = [mid for mid in matched if mid in enabled]
|
|
return {"provider_id": provider_id, "total": len(all_ids), "matched": matched}
|
|
|
|
|
|
def discover_with_credentials(
|
|
adapter_name: str, base_url: str, api_key: str
|
|
) -> dict[str, Any]:
|
|
"""尚未保存供应商时的临时发现。"""
|
|
cfg = ProviderConfig(
|
|
id="__tmp__",
|
|
name="__tmp__",
|
|
adapter=adapter_name,
|
|
base_url=base_url,
|
|
api_key=api_key,
|
|
)
|
|
adapter = get_adapter(cfg)
|
|
all_ids = adapter.list_models()
|
|
matched = adapter.filter_image_models(all_ids)
|
|
return {"total": len(all_ids), "matched": matched}
|