feat: add admin usage statistics and LLM configuration management
This commit is contained in:
@@ -1,7 +1,7 @@
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict
|
||||
from typing import Any, Callable, Dict, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
@@ -14,7 +14,7 @@ class LLMService:
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
|
||||
async def stream_chat(self, system_prompt: str, user_message: str, on_token=None, max_tokens: int = 8192) -> str:
|
||||
async def stream_chat(self, system_prompt: str, user_message: str, on_token=None, max_tokens: int = 8192, usage_callback: Optional[Callable[[int, int], None]] = None) -> str:
|
||||
url = f"{self.base_url}/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
@@ -29,8 +29,10 @@ class LLMService:
|
||||
"temperature": 0.3,
|
||||
"max_tokens": max_tokens,
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True},
|
||||
}
|
||||
full_text = ""
|
||||
_usage = None
|
||||
timeout = httpx.Timeout(connect=10.0, read=90.0, write=10.0, pool=5.0)
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
async with client.stream("POST", url, json=payload, headers=headers) as resp:
|
||||
@@ -46,6 +48,9 @@ class LLMService:
|
||||
break
|
||||
try:
|
||||
chunk = json.loads(data)
|
||||
if chunk.get("usage"):
|
||||
_usage = chunk["usage"]
|
||||
continue
|
||||
delta = chunk["choices"][0]["delta"].get("content", "")
|
||||
if delta:
|
||||
full_text += delta
|
||||
@@ -53,10 +58,12 @@ class LLMService:
|
||||
on_token(delta)
|
||||
except (json.JSONDecodeError, KeyError, IndexError):
|
||||
continue
|
||||
if _usage and usage_callback:
|
||||
usage_callback(_usage.get("prompt_tokens", 0), _usage.get("completion_tokens", 0))
|
||||
return full_text
|
||||
|
||||
async def stream_chat_json(self, system_prompt: str, user_message: str, on_token=None, max_tokens: int = 8192):
|
||||
raw = await self.stream_chat(system_prompt, user_message, on_token, max_tokens=max_tokens)
|
||||
async def stream_chat_json(self, system_prompt: str, user_message: str, on_token=None, max_tokens: int = 8192, usage_callback: Optional[Callable[[int, int], None]] = None):
|
||||
raw = await self.stream_chat(system_prompt, user_message, on_token, max_tokens=max_tokens, usage_callback=usage_callback)
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
raise ValueError("LLM returned empty response")
|
||||
@@ -74,7 +81,7 @@ class LLMService:
|
||||
logger.error(f"JSON parse failed. Raw (first 500): {raw[:500]}")
|
||||
raise
|
||||
|
||||
async def chat(self, system_prompt: str, user_message: str) -> str:
|
||||
async def chat(self, system_prompt: str, user_message: str, usage_callback: Optional[Callable[[int, int], None]] = None) -> str:
|
||||
url = f"{self.base_url}/chat/completions"
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
@@ -97,10 +104,13 @@ class LLMService:
|
||||
logger.error(f"LLM API error {resp.status_code}: {resp.text}")
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
usage = data.get("usage", {})
|
||||
if usage and usage_callback:
|
||||
usage_callback(usage.get("prompt_tokens", 0), usage.get("completion_tokens", 0))
|
||||
return data["choices"][0]["message"]["content"]
|
||||
|
||||
async def chat_json(self, system_prompt: str, user_message: str) -> Any:
|
||||
raw = await self.chat(system_prompt, user_message)
|
||||
async def chat_json(self, system_prompt: str, user_message: str, usage_callback: Optional[Callable[[int, int], None]] = None) -> Any:
|
||||
raw = await self.chat(system_prompt, user_message, usage_callback=usage_callback)
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
raise ValueError("LLM returned empty response")
|
||||
@@ -118,7 +128,7 @@ class LLMService:
|
||||
logger.error(f"JSON parse failed. Raw response (first 500 chars): {raw[:500]}")
|
||||
raise
|
||||
|
||||
async def extract_characters(self, text_samples: list[str], on_token=None, on_sample=None, turbo: bool = False) -> list[Dict]:
|
||||
async def extract_characters(self, text_samples: list[str], on_token=None, on_sample=None, turbo: bool = False, usage_callback: Optional[Callable[[int, int], None]] = None) -> list[Dict]:
|
||||
system_prompt = (
|
||||
"你是一个专业的小说分析助手兼声音导演。请分析给定的小说文本,提取所有出现的角色(包括旁白narrator)。\n"
|
||||
"gender字段必须明确标注性别,只能取以下三个值之一:\"男\"、\"女\"、\"未知\"。\n"
|
||||
@@ -145,7 +155,7 @@ class LLMService:
|
||||
|
||||
async def _extract_one(i: int, sample: str) -> list[Dict]:
|
||||
user_message = f"请分析以下小说文本并提取角色:\n\n{sample}"
|
||||
result = await self.stream_chat_json(system_prompt, user_message, None)
|
||||
result = await self.stream_chat_json(system_prompt, user_message, None, usage_callback=usage_callback)
|
||||
if on_sample:
|
||||
on_sample(i, len(text_samples))
|
||||
return result.get("characters", [])
|
||||
@@ -160,14 +170,14 @@ class LLMService:
|
||||
logger.warning(f"Character extraction failed for sample {i+1}: {r}")
|
||||
else:
|
||||
raw_all.extend(r)
|
||||
return await self.merge_characters(raw_all)
|
||||
return await self.merge_characters(raw_all, usage_callback=usage_callback)
|
||||
|
||||
raw_all: list[Dict] = []
|
||||
for i, sample in enumerate(text_samples):
|
||||
logger.info(f"Extracting characters from sample {i+1}/{len(text_samples)}")
|
||||
user_message = f"请分析以下小说文本并提取角色:\n\n{sample}"
|
||||
try:
|
||||
result = await self.stream_chat_json(system_prompt, user_message, on_token)
|
||||
result = await self.stream_chat_json(system_prompt, user_message, on_token, usage_callback=usage_callback)
|
||||
raw_all.extend(result.get("characters", []))
|
||||
except Exception as e:
|
||||
logger.warning(f"Character extraction failed for sample {i+1}: {e}")
|
||||
@@ -175,9 +185,9 @@ class LLMService:
|
||||
on_sample(i, len(text_samples))
|
||||
if len(text_samples) == 1:
|
||||
return raw_all
|
||||
return await self.merge_characters(raw_all)
|
||||
return await self.merge_characters(raw_all, usage_callback=usage_callback)
|
||||
|
||||
async def merge_characters(self, raw_characters: list[Dict]) -> list[Dict]:
|
||||
async def merge_characters(self, raw_characters: list[Dict], usage_callback: Optional[Callable[[int, int], None]] = None) -> list[Dict]:
|
||||
system_prompt = (
|
||||
"你是一个专业的小说角色整合助手。你收到的是从同一本书不同段落中提取的角色列表,其中可能存在重复。\n"
|
||||
"请完成以下任务:\n"
|
||||
@@ -191,7 +201,7 @@ class LLMService:
|
||||
)
|
||||
user_message = f"请整合以下角色列表:\n\n{json.dumps(raw_characters, ensure_ascii=False, indent=2)}"
|
||||
try:
|
||||
result = await self.chat_json(system_prompt, user_message)
|
||||
result = await self.chat_json(system_prompt, user_message, usage_callback=usage_callback)
|
||||
return result.get("characters", [])
|
||||
except Exception as e:
|
||||
logger.warning(f"Character merge failed, falling back to name-dedup: {e}")
|
||||
@@ -202,7 +212,7 @@ class LLMService:
|
||||
seen[name] = c
|
||||
return list(seen.values())
|
||||
|
||||
async def parse_chapter_segments(self, chapter_text: str, character_names: list[str], on_token=None) -> list[Dict]:
|
||||
async def parse_chapter_segments(self, chapter_text: str, character_names: list[str], on_token=None, usage_callback: Optional[Callable[[int, int], None]] = None) -> list[Dict]:
|
||||
names_str = "、".join(character_names)
|
||||
system_prompt = (
|
||||
"你是一个专业的有声书制作助手。请将给定的章节文本解析为对话片段列表。"
|
||||
@@ -217,7 +227,7 @@ class LLMService:
|
||||
'{"character": "角色名", "text": "对话内容", "emo_text": "开心", "emo_alpha": 0.6}, ...]'
|
||||
)
|
||||
user_message = f"请解析以下章节文本:\n\n{chapter_text}"
|
||||
result = await self.stream_chat_json(system_prompt, user_message, on_token, max_tokens=16384)
|
||||
result = await self.stream_chat_json(system_prompt, user_message, on_token, max_tokens=16384, usage_callback=usage_callback)
|
||||
if isinstance(result, list):
|
||||
return result
|
||||
return []
|
||||
|
||||
Reference in New Issue
Block a user