100d1a06eb
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
291 lines
7.7 KiB
Python
291 lines
7.7 KiB
Python
"""
|
||
LLM 调用工具函数
|
||
|
||
统一封装 LLM 调用模式，减少代码重复。
|
||
"""
|
||
|
||
import json
|
||
import re
|
||
import logging
|
||
from typing import List, Optional, Union, Iterator, Callable
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def call_llm(
|
||
client,
|
||
prompt: str,
|
||
model: str,
|
||
temperature: float = 0.3,
|
||
max_tokens: int = 1000,
|
||
messages: List[dict] = None,
|
||
stream: bool = False,
|
||
**kwargs
|
||
) -> Union[str, Iterator, None]:
|
||
"""
|
||
统一的 LLM 调用封装
|
||
|
||
Args:
|
||
client: OpenAI 客户端实例
|
||
prompt: 用户提示（当 messages 为 None 时使用）
|
||
model: 模型名称
|
||
temperature: 温度参数 (0-1)
|
||
max_tokens: 最大 token 数
|
||
messages: 完整消息列表（优先于 prompt）
|
||
stream: 是否启用流式输出
|
||
**kwargs: 其他参数（如 response_format, tools 等）
|
||
|
||
Returns:
|
||
流式模式返回 stream 迭代器，否则返回响应内容字符串
|
||
调用失败返回 None
|
||
|
||
Example:
|
||
# 简单调用
|
||
response = call_llm(client, "你好", model, temperature=0.3)
|
||
|
||
# 使用消息列表
|
||
response = call_llm(client, "", model, messages=[
|
||
{"role": "system", "content": "你是助手"},
|
||
{"role": "user", "content": "你好"}
|
||
])
|
||
|
||
# 流式调用
|
||
for chunk in call_llm(client, prompt, model, stream=True):
|
||
if chunk.choices and chunk.choices[0].delta.content:
|
||
yield chunk.choices[0].delta.content
|
||
"""
|
||
if messages is None:
|
||
messages = [{"role": "user", "content": prompt}]
|
||
|
||
try:
|
||
response = client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
stream=stream,
|
||
**kwargs
|
||
)
|
||
|
||
if stream:
|
||
return response
|
||
|
||
content = response.choices[0].message.content
|
||
|
||
# 推理模型兼容：content 为空时尝试从 reasoning_content 提取
|
||
if not content or not content.strip():
|
||
reasoning = getattr(response.choices[0].message, 'reasoning_content', None)
|
||
if reasoning and reasoning.strip():
|
||
# 从思维链中提取 JSON 块作为内容
|
||
json_match = re.search(r'\{[\s\S]*\}', reasoning)
|
||
if json_match:
|
||
logger.info("LLM: content为空，从reasoning_content提取JSON")
|
||
return json_match.group().strip()
|
||
logger.warning("LLM 返回空 content（可能需要增大 max_tokens）")
|
||
return None
|
||
|
||
return content.strip()
|
||
except Exception as e:
|
||
logger.warning(f"LLM 调用失败: {e}")
|
||
return None
|
||
|
||
|
||
def call_llm_stream(
|
||
client,
|
||
prompt: str,
|
||
model: str,
|
||
temperature: float = 0.3,
|
||
max_tokens: int = 1000,
|
||
messages: List[dict] = None,
|
||
error_prefix: str = "[错误]",
|
||
**kwargs
|
||
) -> Iterator[str]:
|
||
"""
|
||
流式 LLM 调用（生成器封装）
|
||
|
||
自动处理流式响应，逐块 yield 文本内容。
|
||
|
||
Args:
|
||
client: OpenAI 客户端实例
|
||
prompt: 用户提示
|
||
model: 模型名称
|
||
temperature: 温度参数
|
||
max_tokens: 最大 token 数
|
||
messages: 完整消息列表
|
||
error_prefix: 错误时的前缀
|
||
**kwargs: 其他参数
|
||
|
||
Yields:
|
||
响应文本片段
|
||
|
||
Example:
|
||
for text in call_llm_stream(client, "讲个故事", model):
|
||
print(text, end="", flush=True)
|
||
"""
|
||
if messages is None:
|
||
messages = [{"role": "user", "content": prompt}]
|
||
|
||
try:
|
||
stream = client.chat.completions.create(
|
||
model=model,
|
||
messages=messages,
|
||
temperature=temperature,
|
||
max_tokens=max_tokens,
|
||
stream=True,
|
||
**kwargs
|
||
)
|
||
|
||
for chunk in stream:
|
||
if chunk.choices and chunk.choices[0].delta.content:
|
||
yield chunk.choices[0].delta.content
|
||
|
||
except Exception as e:
|
||
logger.error(f"LLM 流式调用失败: {e}")
|
||
yield f"{error_prefix} 调用大模型失败: {str(e)}"
|
||
|
||
|
||
def call_llm_with_retry(
|
||
client,
|
||
prompt: str,
|
||
model: str,
|
||
max_retries: int = 3,
|
||
retry_delay: float = 1.0,
|
||
**kwargs
|
||
) -> Optional[str]:
|
||
"""
|
||
带重试的 LLM 调用
|
||
|
||
Args:
|
||
client: OpenAI 客户端实例
|
||
prompt: 用户提示
|
||
model: 模型名称
|
||
max_retries: 最大重试次数
|
||
retry_delay: 重试间隔（秒）
|
||
**kwargs: 传递给 call_llm 的其他参数
|
||
|
||
Returns:
|
||
响应内容，全部失败返回 None
|
||
"""
|
||
import time
|
||
|
||
for attempt in range(max_retries):
|
||
result = call_llm(client, prompt, model, **kwargs)
|
||
if result is not None:
|
||
return result
|
||
|
||
if attempt < max_retries - 1:
|
||
logger.debug(f"LLM 调用重试 {attempt + 1}/{max_retries}")
|
||
time.sleep(retry_delay)
|
||
|
||
logger.warning(f"LLM 调用失败，已重试 {max_retries} 次")
|
||
return None
|
||
|
||
|
||
def parse_json_from_response(content: str) -> Optional[dict]:
|
||
"""
|
||
从 LLM 响应中解析 JSON
|
||
|
||
自动处理 ```json 或 ``` 代码块格式
|
||
|
||
Args:
|
||
content: LLM 返回的原始内容
|
||
|
||
Returns:
|
||
解析后的字典，失败返回 None
|
||
"""
|
||
if not content:
|
||
return None
|
||
|
||
# 提取 JSON 代码块（支持 ```json 或 ```）
|
||
json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
|
||
json_str = json_match.group(1) if json_match else content
|
||
|
||
try:
|
||
return json.loads(json_str.strip())
|
||
except (json.JSONDecodeError, TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
def parse_json_list_from_response(content: str) -> Optional[List[dict]]:
|
||
"""
|
||
从 LLM 响应中解析 JSON 数组
|
||
|
||
Args:
|
||
content: LLM 返回的原始内容
|
||
|
||
Returns:
|
||
解析后的列表，失败返回 None
|
||
"""
|
||
if not content:
|
||
return None
|
||
|
||
# 提取 JSON 代码块
|
||
json_match = re.search(r'```(?:json)?\s*([\s\S]*?)\s*```', content)
|
||
json_str = json_match.group(1) if json_match else content
|
||
|
||
try:
|
||
result = json.loads(json_str.strip())
|
||
if isinstance(result, list):
|
||
return result
|
||
# 如果是 {"items": [...]} 格式，提取 items
|
||
if isinstance(result, dict):
|
||
for key in ['items', 'data', 'results', 'questions']:
|
||
if key in result and isinstance(result[key], list):
|
||
return result[key]
|
||
return None
|
||
except (json.JSONDecodeError, TypeError, ValueError):
|
||
return None
|
||
|
||
|
||
# ==================== 便捷函数 ====================
|
||
|
||
def quick_ask(
|
||
client,
|
||
prompt: str,
|
||
model: str,
|
||
temperature: float = 0.3
|
||
) -> str:
|
||
"""
|
||
快速提问（简化版）
|
||
|
||
Args:
|
||
client: OpenAI 客户端
|
||
prompt: 问题
|
||
model: 模型
|
||
temperature: 温度
|
||
|
||
Returns:
|
||
回答内容，失败返回空字符串
|
||
"""
|
||
result = call_llm(client, prompt, model, temperature=temperature)
|
||
return result or ""
|
||
|
||
|
||
def quick_yes_no(
|
||
client,
|
||
prompt: str,
|
||
model: str,
|
||
keywords: List[str] = None
|
||
) -> bool:
|
||
"""
|
||
快速是/否判断
|
||
|
||
Args:
|
||
client: OpenAI 客户端
|
||
prompt: 问题（需包含判断标准）
|
||
model: 模型
|
||
keywords: 判断为 True 的关键词（默认 ["是", "需要", "yes", "true"]）
|
||
|
||
Returns:
|
||
布尔判断结果
|
||
"""
|
||
if keywords is None:
|
||
keywords = ["是", "需要", "yes", "true"]
|
||
|
||
result = call_llm(client, prompt, model, temperature=0, max_tokens=10)
|
||
if result is None:
|
||
return False
|
||
|
||
result_lower = result.lower()
|
||
return any(kw.lower() in result_lower for kw in keywords)
|