bwstudio
db8cea02f8
- AiSetting 模型新增 6 个提示词字段 + 3 个阶段温度字段 - main.py 迁移逻辑改为检测 ai_settings 表的新字段并填充默认值 - generate.py 从数据库读取生成提示词，去掉硬编码 - optimizer 从 API 读取各阶段提示词和温度，删除 prompts.py - crawler 从 API 读取提取/改写提示词和温度，删除 prompts.py - settings/active 端点去掉 token 认证（供爬虫/优化器使用） - 后台设置页新增提示词编辑区和温度调节控件 - 新增 _ensure_default_settings 自动创建默认配置
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
"""LLM 处理：调用 NVIDIA NIM（OpenAI 兼容 API）进行笑话提取、改写。"""
|
||
|
||
import json
|
||
import httpx
|
||
from openai import OpenAI
|
||
|
||
|
||
class AiService:
|
||
def __init__(self, api_base: str, username: str, password: str):
|
||
self.api_base = api_base.rstrip("/")
|
||
self.username = username
|
||
self.password = password
|
||
self.token = None
|
||
self.client = None
|
||
self.model = ""
|
||
self.ai_config = None
|
||
|
||
def _login(self) -> str:
|
||
resp = httpx.post(
|
||
f"{self.api_base}/api/auth/login",
|
||
json={"username": self.username, "password": self.password},
|
||
timeout=30,
|
||
)
|
||
resp.raise_for_status()
|
||
return resp.json()["access_token"]
|
||
|
||
def _get(self, path: str) -> dict | list:
|
||
resp = httpx.get(
|
||
f"{self.api_base}{path}",
|
||
headers={"Authorization": f"Bearer {self.token}"},
|
||
timeout=30,
|
||
)
|
||
resp.raise_for_status()
|
||
return resp.json()
|
||
|
||
def setup(self):
|
||
"""登录并从 API 获取 AI 配置（含提示词和温度）"""
|
||
self.token = self._login()
|
||
self.ai_config = self._get("/api/admin/settings/active")
|
||
self.client = OpenAI(
|
||
base_url=self.ai_config["api_base"],
|
||
api_key=self.ai_config["api_key"],
|
||
)
|
||
self.model = self.ai_config["model_name"]
|
||
|
||
def _get_prompt(self, key: str, default: str) -> str:
|
||
if self.ai_config:
|
||
val = self.ai_config.get(key)
|
||
if val and val.strip():
|
||
return val
|
||
return default
|
||
|
||
def extract_jokes(self, page_content: str, known_types: list[str], known_crowds: list[str]) -> list[dict]:
|
||
"""从页面内容中提取笑话，返回结构化数据。"""
|
||
system_prompt = self._get_prompt("crawler_extract_prompt",
|
||
"""你是一个笑话提取专家。从给定的网页文本中识别并提取所有笑话、幽默段子或有趣内容。
|
||
要求：
|
||
1. 只返回真正的笑话内容，不要提取普通文章或新闻
|
||
2. 每条笑话需要包含：title（简短标题）、content（完整笑话内容）、type（类型）、crowd（人群）
|
||
3. 如果网页中没有笑话，返回空数组 []
|
||
4. 永远返回合法的 JSON 格式，根节点为数组或包含 jokes 键的对象""")
|
||
|
||
user_prompt = f"""网页内容：
|
||
---
|
||
{page_content[:8000]}
|
||
---
|
||
|
||
已知笑话类型：{', '.join(known_types)}
|
||
已知人群分类：{', '.join(known_crowds)}
|
||
|
||
请提取所有笑话，以 JSON 格式返回，示例：
|
||
[
|
||
{{"title": "程序员的幽默", "content": "程序员去相亲...", "types": ["谐音梗", "段子"], "crowds": ["职场", "大学生"]}},
|
||
{{"title": "...", "content": "...", "types": ["..."], "crowds": ["..."]}}
|
||
]
|
||
|
||
注意：types 和 crowds 是数组，可以填多个。
|
||
只返回 JSON，不要其他文字。"""
|
||
|
||
response = self.client.chat.completions.create(
|
||
model=self.model,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_prompt},
|
||
],
|
||
temperature=self.ai_config.get("temperature", 0.7) if self.ai_config else 0.7,
|
||
max_tokens=self.ai_config.get("max_tokens", 2048) if self.ai_config else 2048,
|
||
)
|
||
|
||
raw = response.choices[0].message.content
|
||
return self._parse_json(raw)
|
||
|
||
def rewrite_joke(self, content: str) -> str:
|
||
"""润色单条笑话内容。"""
|
||
system_prompt = self._get_prompt("crawler_rewrite_prompt",
|
||
"""你是一个幽默作家，负责润色和改写笑话。
|
||
要求：
|
||
1. 保持笑话的核心笑点不变
|
||
2. 语言更通顺、更幽默
|
||
3. 字数控制在原内容的 80%-120% 之间
|
||
4. 不要添加任何解释说明""")
|
||
|
||
response = self.client.chat.completions.create(
|
||
model=self.model,
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": f"请润色以下笑话：\n\n{content}"},
|
||
],
|
||
temperature=0.8,
|
||
max_tokens=500,
|
||
)
|
||
return response.choices[0].message.content.strip()
|
||
|
||
def _parse_json(self, raw: str) -> list[dict]:
|
||
"""安全解析 LLM 返回的 JSON。"""
|
||
try:
|
||
data = json.loads(raw)
|
||
if isinstance(data, dict):
|
||
jokes = data.get("jokes") or data.get("items") or [data]
|
||
else:
|
||
jokes = data
|
||
return [j for j in jokes if isinstance(j, dict) and j.get("content")]
|
||
except json.JSONDecodeError:
|
||
if "```json" in raw:
|
||
raw = raw.split("```json")[1].split("```")[0]
|
||
elif "```" in raw:
|
||
raw = raw.split("```")[1].split("```")[0]
|
||
try:
|
||
return json.loads(raw.strip())
|
||
except Exception:
|
||
return [] |