259 lines
10 KiB
Python
259 lines
10 KiB
Python
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}")
|
|
|