"""供应商与模型的服务层：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}