init: RAG 知识库服务初始提交
- 后端 API（Flask + Gunicorn） - RAG 引擎（混合检索 + 云端 Reranker + 引用溯源） - 文档解析（MinerU + 多格式支持） - Docker 生产部署配置 - 排除前端项目、敏感配置、模型文件
This commit is contained in:
lacerate551
290
core/llm_utils.py
Normal file
290
core/llm_utils.py
Normal file
@@ -0,0 +1,290 @@
|
||||
"""
|
||||
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)
|
||||
Reference in New Issue
Block a user