247 lines
8.5 KiB
Python
247 lines
8.5 KiB
Python
"""
|
||
大模型API集成 — 支持 SiliconFlow 和 DeepSeek 官方 + 视觉模型
|
||
"""
|
||
import logging
|
||
import os
|
||
import base64
|
||
from typing import List, Dict, Any, Optional, AsyncGenerator, Tuple
|
||
from langchain_openai import ChatOpenAI
|
||
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
|
||
from langchain_core.callbacks.base import BaseCallbackHandler
|
||
from langchain_core.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
|
||
import openai
|
||
|
||
from ..core.config import get_settings
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
settings = get_settings()
|
||
|
||
# 图片描述提示词
|
||
IMAGE_DESCRIPTION_PROMPT = """你是一个国土空间规划专家。请详细描述这张PDF文档中的图片内容。
|
||
|
||
图片周围文字上下文（来自PDF页面）：{context_text}
|
||
|
||
要求：
|
||
1. 说明图片类型（地图/规划图/图表/流程图/示意图/照片等）
|
||
2. 描述图片中的关键信息、数据和空间关系
|
||
3. 提取图中所有文字标注
|
||
4. 描述控制在200-300字"""
|
||
|
||
# DeepSeek 官方模型 ID 前缀（用于自动路由）
|
||
DEEPSEEK_OFFICIAL_MODELS = {
|
||
"deepseek-chat",
|
||
"deepseek-reasoner",
|
||
}
|
||
|
||
# 模型 ID → (api_key, base_url) 的路由映射
|
||
def _resolve_provider(model: str) -> tuple[str, str]:
|
||
"""根据模型 ID 选择 API provider，返回 (api_key, base_url)"""
|
||
# DeepSeek 官方模型（不带 siliconflow 前缀的纯 deepseek-xxx）
|
||
if model in DEEPSEEK_OFFICIAL_MODELS or model.startswith("deepseek/"):
|
||
if not settings.deepseek_api_key:
|
||
raise ValueError(
|
||
f"模型 {model} 需要 DeepSeek 官方 API Key，"
|
||
"请在 .env 中配置 DEEPSEEK_API_KEY"
|
||
)
|
||
actual_model = model.replace("deepseek/", "")
|
||
return settings.deepseek_api_key, settings.deepseek_base_url, actual_model
|
||
|
||
# 默认走 SiliconFlow
|
||
return settings.siliconflow_api_key, settings.siliconflow_base_url, model
|
||
|
||
|
||
class SiliconFlowLLM:
|
||
"""大模型客户端 — 自动路由 SiliconFlow / DeepSeek 官方"""
|
||
|
||
def __init__(self, model: Optional[str] = None):
|
||
raw_model = model or settings.siliconflow_model
|
||
api_key, base_url, resolved_model = _resolve_provider(raw_model)
|
||
|
||
self.model_name = resolved_model
|
||
logger.info(f"模型: {resolved_model}, API: {base_url}")
|
||
|
||
self.llm = ChatOpenAI(
|
||
model=resolved_model,
|
||
api_key=api_key,
|
||
base_url=base_url,
|
||
temperature=0.7,
|
||
max_tokens=2000,
|
||
streaming=True
|
||
)
|
||
|
||
self._api_key = api_key
|
||
self._base_url = base_url
|
||
|
||
def chat(self, messages: List[BaseMessage], **kwargs) -> str:
|
||
"""同步聊天"""
|
||
try:
|
||
response = self.llm.invoke(messages, **kwargs)
|
||
return response.content
|
||
except Exception as e:
|
||
raise Exception(f"LLM调用失败: {str(e)}")
|
||
|
||
async def achat(self, messages: List[BaseMessage], **kwargs) -> str:
|
||
"""异步聊天"""
|
||
try:
|
||
response = await self.llm.ainvoke(messages, **kwargs)
|
||
return response.content
|
||
except Exception as e:
|
||
raise Exception(f"LLM异步调用失败: {str(e)}")
|
||
|
||
async def stream_chat(self, messages: List[BaseMessage], **kwargs) -> AsyncGenerator[str, None]:
|
||
"""流式聊天"""
|
||
try:
|
||
async for chunk in self.llm.astream(messages, **kwargs):
|
||
if hasattr(chunk, 'content') and chunk.content:
|
||
yield chunk.content
|
||
except Exception as e:
|
||
raise Exception(f"LLM流式调用失败: {str(e)}")
|
||
|
||
async def stream_with_reasoning(
|
||
self, messages: List[BaseMessage], **kwargs
|
||
) -> AsyncGenerator[Tuple[str, str], None]:
|
||
"""流式调用（直接使用 OpenAI SDK，捕获推理内容）
|
||
|
||
Yields:
|
||
(type, content) — type 为 "reasoning" 或 "content"
|
||
"""
|
||
client = openai.AsyncOpenAI(
|
||
api_key=self._api_key,
|
||
base_url=self._base_url
|
||
)
|
||
|
||
openai_messages = []
|
||
for msg in messages:
|
||
if isinstance(msg, SystemMessage):
|
||
openai_messages.append({"role": "system", "content": msg.content})
|
||
elif isinstance(msg, HumanMessage):
|
||
openai_messages.append({"role": "user", "content": msg.content})
|
||
elif isinstance(msg, AIMessage):
|
||
openai_messages.append({"role": "assistant", "content": msg.content})
|
||
else:
|
||
openai_messages.append({"role": "user", "content": str(msg.content)})
|
||
|
||
stream = await client.chat.completions.create(
|
||
model=self.model_name,
|
||
messages=openai_messages,
|
||
stream=True,
|
||
**kwargs
|
||
)
|
||
|
||
async for chunk in stream:
|
||
if not chunk.choices:
|
||
continue
|
||
delta = chunk.choices[0].delta
|
||
|
||
rc = getattr(delta, 'reasoning_content', None)
|
||
if rc:
|
||
yield ("reasoning", rc)
|
||
|
||
if delta.content:
|
||
yield ("content", delta.content)
|
||
|
||
def create_messages(
|
||
self,
|
||
user_message: str,
|
||
system_prompt: Optional[str] = None,
|
||
chat_history: Optional[List[Dict[str, str]]] = None
|
||
) -> List[BaseMessage]:
|
||
"""创建消息列表"""
|
||
messages = []
|
||
|
||
if system_prompt:
|
||
messages.append(SystemMessage(content=system_prompt))
|
||
|
||
if chat_history:
|
||
for msg in chat_history:
|
||
if msg["role"] == "user":
|
||
messages.append(HumanMessage(content=msg["content"]))
|
||
elif msg["role"] == "assistant":
|
||
messages.append(AIMessage(content=msg["content"]))
|
||
|
||
messages.append(HumanMessage(content=user_message))
|
||
|
||
return messages
|
||
|
||
|
||
async def describe_image(self, image_path: str, context_text: str = "") -> str:
|
||
"""使用VLM模型描述图片内容
|
||
|
||
Args:
|
||
image_path: 图片文件路径
|
||
context_text: 图片周围的PDF文本上下文
|
||
|
||
Returns:
|
||
图片的文字描述
|
||
"""
|
||
import asyncio
|
||
import time
|
||
|
||
# 读取图片并编码为base64，不支持的格式先转为PNG
|
||
ext = os.path.splitext(image_path)[1].lower()
|
||
mime_map = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".gif": "image/gif", ".webp": "image/webp"}
|
||
|
||
if ext not in mime_map:
|
||
from PIL import Image
|
||
import io
|
||
img = Image.open(image_path)
|
||
if img.mode in ("CMYK", "P"):
|
||
img = img.convert("RGB")
|
||
buf = io.BytesIO()
|
||
img.save(buf, format="PNG")
|
||
image_data = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||
mime_type = "image/png"
|
||
else:
|
||
with open(image_path, "rb") as f:
|
||
image_data = base64.b64encode(f.read()).decode("utf-8")
|
||
mime_type = mime_map[ext]
|
||
|
||
prompt = IMAGE_DESCRIPTION_PROMPT.format(context_text=context_text[:600])
|
||
|
||
vision_model = "Qwen/Qwen3-VL-8B-Instruct"
|
||
api_key, base_url, _ = _resolve_provider(vision_model)
|
||
|
||
client = openai.AsyncOpenAI(api_key=api_key, base_url=base_url)
|
||
|
||
max_retries = 3
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = await client.chat.completions.create(
|
||
model=vision_model,
|
||
messages=[{
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": prompt},
|
||
{"type": "image_url", "image_url": {"url": f"data:{mime_type};base64,{image_data}"}},
|
||
],
|
||
}],
|
||
max_tokens=600,
|
||
temperature=0.3,
|
||
timeout=90.0,
|
||
)
|
||
return response.choices[0].message.content or ""
|
||
|
||
except Exception as e:
|
||
logger.warning(f"描述失败 attempt={attempt+1}: {e}")
|
||
if attempt < max_retries - 1:
|
||
await asyncio.sleep(2 ** attempt)
|
||
|
||
return ""
|
||
|
||
|
||
# 全局LLM实例（使用默认模型）
|
||
llm_client = SiliconFlowLLM()
|
||
|
||
|
||
def get_llm_client(model: Optional[str] = None) -> SiliconFlowLLM:
|
||
"""获取LLM客户端实例
|
||
|
||
Args:
|
||
model: 可选的模型名称，如果提供则创建新的实例，否则返回默认实例
|
||
"""
|
||
if model is None:
|
||
return llm_client
|
||
else:
|
||
return SiliconFlowLLM(model=model)
|