60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
|
#novel_generator/common.py
|
||
|
|
# -*- coding: utf-8 -*-
|
||
|
|
"""
|
||
|
|
通用重试、清洗、日志工具
|
||
|
|
"""
|
||
|
|
import logging
|
||
|
|
import re
|
||
|
|
import time
|
||
|
|
import traceback
|
||
|
|
|
||
|
|
def call_with_retry(func, max_retries=3, sleep_time=2, fallback_return=None, **kwargs):
|
||
|
|
"""
|
||
|
|
通用的重试机制封装。
|
||
|
|
:param func: 要执行的函数
|
||
|
|
:param max_retries: 最大重试次数
|
||
|
|
:param sleep_time: 重试前的等待秒数
|
||
|
|
:param fallback_return: 如果多次重试仍失败时的返回值
|
||
|
|
:param kwargs: 传给func的命名参数
|
||
|
|
:return: func的结果，若失败则返回 fallback_return
|
||
|
|
"""
|
||
|
|
for attempt in range(1, max_retries + 1):
|
||
|
|
try:
|
||
|
|
return func(**kwargs)
|
||
|
|
except Exception as e:
|
||
|
|
logging.warning(f"[call_with_retry] Attempt {attempt} failed with error: {e}")
|
||
|
|
traceback.print_exc()
|
||
|
|
if attempt < max_retries:
|
||
|
|
time.sleep(sleep_time)
|
||
|
|
else:
|
||
|
|
logging.error("Max retries reached, returning fallback_return.")
|
||
|
|
return fallback_return
|
||
|
|
|
||
|
|
def remove_think_tags(text: str) -> str:
|
||
|
|
"""移除 <think>...</think> 包裹的内容"""
|
||
|
|
return re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
|
||
|
|
|
||
|
|
def debug_log(prompt: str, response_content: str):
|
||
|
|
logging.info(
|
||
|
|
f"\n[######################################### Prompt #########################################]\n{prompt}\n"
|
||
|
|
)
|
||
|
|
logging.info(
|
||
|
|
f"\n[######################################### Response #########################################]\n{response_content}\n"
|
||
|
|
)
|
||
|
|
|
||
|
|
def invoke_with_cleaning(llm_adapter, prompt: str) -> str:
|
||
|
|
"""
|
||
|
|
调用 LLM，增加重试和清洗逻辑
|
||
|
|
如果多次失败，则返回空字符串以继续流程，而不是中断。
|
||
|
|
"""
|
||
|
|
def _invoke(prompt):
|
||
|
|
return llm_adapter.invoke(prompt)
|
||
|
|
|
||
|
|
response = call_with_retry(func=_invoke, max_retries=3, fallback_return="", prompt=prompt)
|
||
|
|
if not response:
|
||
|
|
logging.warning("No response from model after retry. Return empty.")
|
||
|
|
return ""
|
||
|
|
cleaned_text = remove_think_tags(response)
|
||
|
|
debug_log(prompt, cleaned_text)
|
||
|
|
return cleaned_text.strip()
|