""" 大模型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)