src/webui/services/generation.py
Rama main · réplica de lectura del repositorio original.
"""AI 生成服务：连接测试 / 世界 / 角色 / 文本生成。"""
from __future__ import annotations
import logging
import json
import re
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable
from src.engine.language import DEFAULT_LANGUAGE, normalize_language
from src.generation import creator
from src.llm.openai_transport import OpenAIChatTransport
from src.rules.rule_system import RuleSystem
from src.webui.config_update import connection_test_timeout
logger = logging.getLogger("trpg")
@dataclass(frozen=True)
class ConnectionDependencies:
llm_client: Any | None
config_state: Callable[[], dict[str, Any]]
@dataclass(frozen=True)
class GenerationDependencies:
llm_client: Any | None
llm_configuration_error: Callable[[str], dict[str, Any] | None]
worlds_dir: Path
lorebook_store: Any
rules_dir: Path
registry: Any
parse_game_key: Callable[[str], tuple[str, ...]]
get_instance: Callable[[tuple[str, ...]], Any | None]
load_rule_by_id: Callable[[str, str], RuleSystem | None]
load_rule_for_game: Callable[[Any], RuleSystem | None]
character_gen_max_tokens: int
text_gen_max_tokens: int
def _model_list_url(base_url: str, api_format: str = "openai") -> str:
url = base_url.rstrip("/")
if url.endswith("/chat/completions"):
return f"{url[:-len('/chat/completions')]}/models"
if url.endswith("/messages"):
url = url[:-len("/messages")]
if url.endswith("/models"):
return url
if (api_format or "openai").strip().lower() == "anthropic" and not url.endswith("/v1"):
return f"{url}/v1/models"
return f"{url}/models"
def _extract_model_ids(payload: Any) -> list[str]:
if isinstance(payload, dict):
raw_models = payload.get("data")
if not isinstance(raw_models, list):
raw_models = payload.get("models")
else:
raw_models = payload
if not isinstance(raw_models, list):
return []
models: list[str] = []
seen: set[str] = set()
for item in raw_models:
if isinstance(item, dict):
model = str(item.get("id") or item.get("name") or item.get("model") or "").strip()
else:
model = str(item or "").strip()
if not model or model in seen:
continue
seen.add(model)
models.append(model[:160])
if len(models) >= 300:
break
return sorted(models, key=str.casefold)
async def list_models(
dependencies: ConnectionDependencies,
base_url: str,
api_key: str,
proxy_url: str = "",
api_format: str = "openai",
) -> dict[str, Any]:
"""读取 OpenAI / Anthropic 兼容接口的模型目录。"""
import aiohttp
url = _model_list_url(base_url, api_format)
anthropic = (api_format or "openai").strip().lower() == "anthropic"
headers = (
{"x-api-key": api_key, "anthropic-version": "2023-06-01"}
if anthropic else {"Authorization": f"Bearer {api_key}"}
)
try:
if not dependencies.llm_client:
return {"ok": False, "error": "LLM 客户端未初始化", "models": []}
session = await dependencies.llm_client._get_session()
active_proxy = proxy_url or dependencies.llm_client.proxy_url
request_kwargs = {"proxy": active_proxy} if active_proxy else {}
async with session.get(
url,
headers=headers,
timeout=aiohttp.ClientTimeout(
total=connection_test_timeout(dependencies.config_state())
),
**request_kwargs,
) as response:
if response.status != 200:
error_text = await response.text()
return {"ok": False, "error": f"HTTP {response.status}: {error_text[:200]}", "models": []}
try:
payload = await response.json()
except (aiohttp.ContentTypeError, json.JSONDecodeError):
return {"ok": False, "error": "模型列表响应不是有效 JSON", "models": []}
models = _extract_model_ids(payload)
if not models:
return {"ok": False, "error": "接口未返回可识别的模型列表", "models": []}
return {"ok": True, "models": models, "count": len(models)}
except Exception as exc:
return {"ok": False, "error": str(exc), "models": []}
async def test_connection(
dependencies: ConnectionDependencies,
base_url: str,
api_key: str,
model: str,
proxy_url: str = "",
api_format: str = "openai",
) -> dict[str, Any]:
"""测试 LLM API 连接是否正常。"""
import aiohttp
if (api_format or "openai").strip().lower() == "anthropic":
return await _test_anthropic_connection(
dependencies, base_url, api_key, model, proxy_url,
)
return await _test_openai_connection(
dependencies, base_url, api_key, model, proxy_url,
)
async def _test_openai_connection(
dependencies: ConnectionDependencies,
base_url: str,
api_key: str,
model: str,
proxy_url: str = "",
) -> dict[str, Any]:
import aiohttp
url = base_url.rstrip("/")
if not url.endswith("/chat/completions"):
url += "/chat/completions"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {api_key}",
}
body = {
"model": model,
"messages": [{"role": "user", "content": "回复：OK"}],
"max_tokens": 10,
}
start = time.time()
try:
if not dependencies.llm_client:
return {"ok": False, "error": "LLM 客户端未初始化", "elapsed": 0}
session = await dependencies.llm_client._get_session()
active_proxy = proxy_url or dependencies.llm_client.proxy_url
request_kwargs = {"proxy": active_proxy} if active_proxy else {}
async with OpenAIChatTransport().post(
session, url, json=body, headers=headers,
timeout=aiohttp.ClientTimeout(
total=connection_test_timeout(dependencies.config_state())
),
**request_kwargs,
) as resp:
return await _parse_connection_test_response(resp, start)
except Exception as e:
return {"ok": False, "error": str(e), "elapsed": round(time.time() - start, 2)}
async def _test_anthropic_connection(
dependencies: ConnectionDependencies,
base_url: str,
api_key: str,
model: str,
proxy_url: str = "",
) -> dict[str, Any]:
import aiohttp
from src.llm.client import _anthropic_messages_url, _anthropic_text_content
url = _anthropic_messages_url(base_url)
headers = {
"Content-Type": "application/json",
"x-api-key": api_key,
"anthropic-version": "2023-06-01",
}
body = {
"model": model,
"messages": [{"role": "user", "content": "回复：OK"}],
"max_tokens": 10,
}
start = time.time()
try:
if not dependencies.llm_client:
return {"ok": False, "error": "LLM 客户端未初始化", "elapsed": 0}
session = await dependencies.llm_client._get_session()
active_proxy = proxy_url or dependencies.llm_client.proxy_url
request_kwargs = {"proxy": active_proxy} if active_proxy else {}
async with session.post(
url, json=body, headers=headers,
timeout=aiohttp.ClientTimeout(
total=connection_test_timeout(dependencies.config_state())
),
**request_kwargs,
) as resp:
elapsed = round(time.time() - start, 2)
if resp.status == 200:
data = await resp.json()
usage = data.get("usage", {})
tokens = int(usage.get("input_tokens", 0) or 0) + int(usage.get("output_tokens", 0) or 0)
return {
"ok": True,
"elapsed": elapsed,
"response": _anthropic_text_content(data).strip(),
"tokens": tokens,
}
error_text = await resp.text()
return {
"ok": False,
"error": f"HTTP {resp.status}: {error_text[:200]}",
"elapsed": elapsed,
}
except Exception as e:
return {"ok": False, "error": str(e), "elapsed": round(time.time() - start, 2)}
async def _parse_connection_test_response(resp, start: float) -> dict[str, Any]:
elapsed = round(time.time() - start, 2)
if resp.status == 200:
data = await resp.json()
content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
return {
"ok": True,
"elapsed": elapsed,
"response": content.strip(),
"tokens": data.get("usage", {}).get("total_tokens", 0),
}
error_text = await resp.text()
return {
"ok": False,
"error": f"HTTP {resp.status}: {error_text[:200]}",
"elapsed": elapsed,
}
async def generate_world(
dependencies: GenerationDependencies,
prompt: str,
rule_id: str = "",
language: str = DEFAULT_LANGUAGE,
) -> dict[str, Any]:
"""使用 AI 根据用户描述生成完整世界模板。"""
if config_error := dependencies.llm_configuration_error(language):
return config_error
try:
return await _gen_world(
dependencies, prompt, rule_id, normalize_language(language),
)
except Exception as e:
logger.exception("AI 生成世界失败")
return {"ok": False, "error": str(e)}
async def _gen_world(
dependencies: GenerationDependencies,
prompt: str,
rule_id: str,
language: str,
) -> dict[str, Any]:
return await creator.generate_world(
dependencies.llm_client, prompt, rule_id or "freeform_fantasy",
worlds_dir=dependencies.worlds_dir,
lorebook_store=dependencies.lorebook_store,
max_tokens=dependencies.character_gen_max_tokens,
language=language,
)
def _slug(value: str) -> str:
raw = re.sub(r"[^A-Za-z0-9_-]+", "_", value.strip())[:32].strip("_")
return raw or "custom"
async def generate_rule(
dependencies: GenerationDependencies,
prompt: str,
source_rule_id: str = "",
language: str = DEFAULT_LANGUAGE,
) -> dict[str, Any]:
"""按母版规则生成并保存一套 AI 自定义规则。"""
if config_error := dependencies.llm_configuration_error(language):
return config_error
prompt = (prompt or "").strip()
if not prompt:
return {"ok": False, "error": "请输入规则题材描述"}
source_rule_id = (source_rule_id or "freeform_fantasy").strip()
source_path = RuleSystem.path_for(
dependencies.rules_dir, source_rule_id, language,
)
if not source_path.exists():
return {"ok": False, "error": f"母版规则不存在: {source_rule_id}"}
try:
language = normalize_language(language)
source_rule = json.loads(source_path.read_text(encoding="utf-8"))
rule_id = f"ai_rule_{_slug(prompt)}_{int(time.time())}"
data = await creator.generate_rule(
dependencies.llm_client,
prompt,
source_rule=source_rule,
source_rule_id=source_rule_id,
rule_id=rule_id,
max_tokens=max(dependencies.character_gen_max_tokens, 4096),
language=language,
)
if not data:
return {"ok": False, "error": "AI 返回规则 JSON 解析失败，请重试"}
data["rule_id"] = rule_id
data["custom"] = True
data["source_rule_id"] = source_rule_id
dependencies.rules_dir.mkdir(parents=True, exist_ok=True)
target = RuleSystem.path_for(dependencies.rules_dir, rule_id)
tmp = target.with_suffix(".json.tmp")
tmp.write_text(json.dumps(data, ensure_ascii=False, indent=2), encoding="utf-8")
# 读取一次确保 JSON 与规则系统兼容；不兼容则不落正式文件。
RuleSystem.load(tmp)
tmp.replace(target)
return {
"ok": True,
"rule_id": rule_id,
"rule_name": data.get("rule_name", rule_id),
"description": data.get("description", ""),
"source_rule_id": source_rule_id,
"rule": data,
}
except Exception as e:
logger.exception("AI 生成规则失败")
return {"ok": False, "error": str(e)}
async def generate_character(
dependencies: GenerationDependencies,
prompt: str,
game_key: str = "",
rule_id: str = "",
language: str = DEFAULT_LANGUAGE,
) -> dict[str, Any]:
"""使用 AI 根据用户描述生成角色卡。"""
if config_error := dependencies.llm_configuration_error(language):
return config_error
try:
resolved_language = normalize_language(language)
rule = dependencies.load_rule_by_id(rule_id, resolved_language)
if game_key:
inst = dependencies.get_instance(
dependencies.parse_game_key(game_key)
)
if inst:
resolved_language = normalize_language(language or getattr(inst, "language", DEFAULT_LANGUAGE))
rule = rule or dependencies.load_rule_for_game(inst)
data = await creator.generate_character(
dependencies.llm_client,
prompt,
game_key,
dependencies.registry,
rule=rule,
max_tokens=dependencies.character_gen_max_tokens,
language=resolved_language,
)
if not data:
return {"ok": False, "error": "AI 返回内容解析失败，请重试"}
return {"ok": True, "character": data}
except Exception as e:
logger.exception("AI 生成角色失败")
return {"ok": False, "error": str(e)}
async def generate_text(
dependencies: GenerationDependencies,
prompt: str,
system_hint: str = "",
language: str = DEFAULT_LANGUAGE,
) -> dict[str, Any]:
"""轻量文字生成：直接发 prompt 给 LLM，返回原始文本，不解析 JSON。"""
if config_error := dependencies.llm_configuration_error(language):
return config_error
llm_client = dependencies.llm_client
if llm_client is None:
return {"ok": False, "error": "LLM 客户端未初始化"}
language = normalize_language(language)
system = system_hint or (
"You are a TRPG character and setting assistant. Answer briefly and practically in natural English. Do not output JSON."
if language == "en"
else "你是一个TRPG角色设定助手。根据用户描述，生成简短实用的回答。不要输出JSON，直接输出文字。"
)
try:
response = await llm_client.call(
system_prompt=system,
user_message=prompt,
temperature=0.7,
max_tokens=dependencies.text_gen_max_tokens,
)
text = response.narration or response.content
return {"ok": True, "text": text.strip()}
except Exception as e:
logger.exception("文字生成失败")
return {"ok": False, "error": str(e)}