from __future__ import annotations from .shared import * class LLMEngineAdminMixin: async def list_engines(self, db: AsyncSession, tenant_id: int = 1) -> dict[str, Any]: _, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id) engines = bundle.get("engines", []) return { "active_engine_id": bundle.get("active_engine_id"), "engines": [self._sanitize_engine_for_response(item) for item in engines], "settings": self._normalize_engine_settings(bundle.get("settings")), } async def update_engine_settings( self, db: AsyncSession, *, tenant_id: int, enable_local_kb_qa: Optional[bool] = None, ) -> dict[str, Any]: integration, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id) settings = self._normalize_engine_settings(bundle.get("settings")) if enable_local_kb_qa is not None: settings["enable_local_kb_qa"] = bool(enable_local_kb_qa) bundle["settings"] = settings integration.config = json.loads(json.dumps(bundle)) await db.commit() return settings async def create_engine( self, db: AsyncSession, *, tenant_id: int, name: str, provider: str, api_key: Optional[str], model: Optional[str] = None, base_url: Optional[str] = None, set_active: bool = True, ) -> dict[str, Any]: provider = _normalize_provider(provider) api_key = str(api_key or "").strip() if provider != "ollama" and not api_key: raise ValueError("api_key 不能为空") integration, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id) now = now_cn_aware().isoformat() engine_id = str(uuid4()) defaults = PROVIDER_DEFAULTS[provider] engine = { "id": engine_id, "name": (name or defaults["name"]).strip()[:80] or defaults["name"], "provider": provider, "api_key": api_key, "model": (model or defaults["model"]).strip() or defaults["model"], "base_url": (base_url or defaults["base_url"]).strip().rstrip("/") or defaults["base_url"], "created_at": now, "updated_at": now, } engines = bundle.get("engines", []) engines.append(engine) bundle["engines"] = engines if set_active or not bundle.get("active_engine_id"): bundle["active_engine_id"] = engine_id integration.config = json.loads(json.dumps(bundle)) integration.is_active = True await db.commit() return self._sanitize_engine_for_response(engine) async def update_engine( self, db: AsyncSession, *, tenant_id: int, engine_id: str, name: Optional[str] = None, provider: Optional[str] = None, api_key: Optional[str] = None, model: Optional[str] = None, base_url: Optional[str] = None, ) -> dict[str, Any]: integration, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id) engines = bundle.get("engines", []) target: Optional[dict[str, Any]] = None for item in engines: if str(item.get("id")) == str(engine_id): target = item break if not target: raise ValueError("engine 不存在") if provider is not None: target["provider"] = _normalize_provider(provider) active_provider = _normalize_provider(target.get("provider")) defaults = PROVIDER_DEFAULTS[active_provider] if name is not None: target["name"] = (name or "").strip()[:80] or target.get("name") or defaults["name"] if api_key is not None: api_key_val = str(api_key).strip() if api_key_val: target["api_key"] = api_key_val if model is not None: target["model"] = (model or "").strip() or defaults["model"] if base_url is not None: target["base_url"] = (base_url or "").strip().rstrip("/") or defaults["base_url"] target["updated_at"] = now_cn_aware().isoformat() integration.config = json.loads(json.dumps(bundle)) await db.commit() return self._sanitize_engine_for_response(target) async def delete_engine( self, db: AsyncSession, *, tenant_id: int, engine_id: str, ) -> dict[str, Any]: integration, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id) engines = bundle.get("engines", []) kept = [item for item in engines if str(item.get("id")) != str(engine_id)] if len(kept) == len(engines): raise ValueError("engine 不存在") bundle["engines"] = kept active_engine_id = str(bundle.get("active_engine_id") or "") if active_engine_id == str(engine_id): bundle["active_engine_id"] = kept[0]["id"] if kept else None integration.config = json.loads(json.dumps(bundle)) await db.commit() return { "deleted_engine_id": str(engine_id), "active_engine_id": bundle.get("active_engine_id"), "remaining_count": len(kept), } async def activate_engine( self, db: AsyncSession, *, tenant_id: int, engine_id: str, ) -> dict[str, Any]: integration, bundle = await self._get_or_init_bundle(db, tenant_id=tenant_id) engines = bundle.get("engines", []) exists = any(str(item.get("id")) == str(engine_id) for item in engines) if not exists: raise ValueError("engine 不存在") bundle["active_engine_id"] = str(engine_id) integration.config = json.loads(json.dumps(bundle)) await db.commit() return {"active_engine_id": str(engine_id)} async def test_engine_connection( self, db: AsyncSession, *, tenant_id: int, engine_id: Optional[str] = None, prompt: Optional[str] = None, max_tokens: int = 64, ) -> dict[str, Any]: test_prompt = (prompt or "你是连通性测试助手。请只返回: OK").strip() if not test_prompt: raise ValueError("测试 prompt 不能为空") engine = await self._resolve_engine_for_test( db, tenant_id=tenant_id, engine_id=engine_id, ) response_text = "" status = "success" started = time.perf_counter() try: transient_retry_budget = 2 transient_retry_attempt = 0 while True: try: response_text = await self._chat_completion( provider=str(engine.get("provider") or ""), prompt=test_prompt, system_prompt="", api_key=str(engine.get("api_key") or ""), base_url=str(engine.get("base_url") or ""), model=str(engine.get("model") or ""), temperature=0.0, max_tokens=max_tokens, ) break except Exception as transient_exc: status_code = self._http_status_from_exception(transient_exc) can_retry_transient = ( transient_retry_budget > 0 and ( self._is_retryable_http_status(status_code) or isinstance(transient_exc, httpx.RequestError) ) ) if not can_retry_transient: raise retry_after_header = None if isinstance(transient_exc, httpx.HTTPStatusError) and transient_exc.response is not None: retry_after_header = transient_exc.response.headers.get("Retry-After") retry_delay = self._http_retry_delay_seconds( status_code=status_code, attempt=transient_retry_attempt, retry_after_header=retry_after_header, ) _runtime_log( "LLM", ( "连通测试重试：" f"engine={engine.get('provider')}/{engine.get('model')}, " f"status={status_code or 'request_error'}, " f"retry_in={retry_delay:.2f}s, left={transient_retry_budget - 1}" ), ) transient_retry_budget -= 1 transient_retry_attempt += 1 await asyncio.sleep(retry_delay) if not response_text.strip(): raise ValueError("测试响应为空") latency_ms = int((time.perf_counter() - started) * 1000) return { "ok": True, "engine_id": str(engine.get("id") or ""), "name": str(engine.get("name") or ""), "provider": str(engine.get("provider") or ""), "model": str(engine.get("model") or ""), "base_url": str(engine.get("base_url") or ""), "api_key_masked": _mask_secret(str(engine.get("api_key") or "")), "latency_ms": latency_ms, "response_preview": response_text[:300], } except Exception as exc: status = "failed" response_text = f"[error] {exc}" status_code = self._http_status_from_exception(exc) if int(status_code or 0) == 429: raise ValueError("引擎连通测试失败: 上游限流(429)，请稍后重试或更换 API Key。") from exc raise ValueError(f"引擎连通测试失败: {exc}") from exc finally: try: await self._log_ai_call( tenant_id=tenant_id, function_name="test_engine_connection", prompt=test_prompt, response=response_text, status=status, provider=str(engine.get("provider") or ""), model=str(engine.get("model") or ""), ) except Exception as log_exc: _runtime_log("LLM", f"测试调用日志写入失败：{log_exc}")