""" LLM 服务 - 会话客户端 提供一个有状态的、面向会话的 LLM 客户端，用于进行多轮对话和复杂交互。 """ import copy from dataclasses import dataclass from typing import Any from nonebot_plugin_alconna.uniseg import UniMessage from zhenxun.services.log import logger from .config import CommonOverrides, LLMGenerationConfig from .config.providers import get_ai_config from .manager import get_global_default_model_name, get_model_instance from .tools import tool_registry from .types import ( EmbeddingTaskType, LLMContentPart, LLMErrorCode, LLMException, LLMMessage, LLMResponse, LLMTool, ModelName, ) from .utils import unimsg_to_llm_parts @dataclass class AIConfig: """AI配置类 - 简化版本""" model: ModelName = None default_embedding_model: ModelName = None temperature: float | None = None max_tokens: int | None = None enable_cache: bool = False enable_code: bool = False enable_search: bool = False timeout: int | None = None enable_gemini_json_mode: bool = False enable_gemini_thinking: bool = False enable_gemini_safe_mode: bool = False enable_gemini_multimodal: bool = False enable_gemini_grounding: bool = False default_preserve_media_in_history: bool = False def __post_init__(self): """初始化后从配置中读取默认值""" ai_config = get_ai_config() if self.model is None: self.model = ai_config.get("default_model_name") if self.timeout is None: self.timeout = ai_config.get("timeout", 180) class AI: """统一的AI服务类 - 平衡设计版本 提供三层API： 1. 简单方法：ai.chat(), ai.code(), ai.search() 2. 标准方法：ai.analyze() 支持复杂参数 3. 高级方法：通过get_model_instance()直接访问 """ def __init__( self, config: AIConfig | None = None, history: list[LLMMessage] | None = None ): """ 初始化AI服务 参数: config: AI 配置. history: 可选的初始对话历史. """ self.config = config or AIConfig() self.history = history or [] def clear_history(self): """清空当前会话的历史记录""" self.history = [] logger.info("AI session history cleared.") def _sanitize_message_for_history(self, message: LLMMessage) -> LLMMessage: """ 净化用于存入历史记录的消息。 将非文本的多模态内容部分替换为文本占位符，以避免重复处理。 """ if not isinstance(message.content, list): return message sanitized_message = copy.deepcopy(message) content_list = sanitized_message.content if not isinstance(content_list, list): return sanitized_message new_content_parts: list[LLMContentPart] = [] has_multimodal_content = False for part in content_list: if isinstance(part, LLMContentPart) and part.type == "text": new_content_parts.append(part) else: has_multimodal_content = True if has_multimodal_content: placeholder = "[用户发送了媒体文件，内容已在首次分析时处理]" text_part_found = False for part in new_content_parts: if part.type == "text": part.text = f"{placeholder} {part.text or ''}".strip() text_part_found = True break if not text_part_found: new_content_parts.insert(0, LLMContentPart.text_part(placeholder)) sanitized_message.content = new_content_parts return sanitized_message async def chat( self, message: str | LLMMessage | list[LLMContentPart], *, model: ModelName = None, preserve_media_in_history: bool | None = None, tools: list[LLMTool] | None = None, tool_choice: str | dict[str, Any] | None = None, **kwargs: Any, ) -> LLMResponse: """ 进行一次聊天对话，支持工具调用。 此方法会自动使用和更新会话内的历史记录。 参数: message: 用户输入的消息。 model: 本次对话要使用的模型。 preserve_media_in_history: 是否在历史记录中保留原始多模态信息。 - True: 保留，用于深度多轮媒体分析。 - False: 不保留，替换为占位符，提高效率。 - None (默认): 使用AI实例配置的默认值。 tools: 本次对话可用的工具列表。 tool_choice: 强制模型使用的工具。 **kwargs: 传递给模型的其他生成参数。 返回: LLMResponse: 模型的完整响应，可能包含文本或工具调用请求。 """ current_message: LLMMessage if isinstance(message, str): current_message = LLMMessage.user(message) elif isinstance(message, list) and all( isinstance(part, LLMContentPart) for part in message ): current_message = LLMMessage.user(message) elif isinstance(message, LLMMessage): current_message = message else: raise LLMException( f"AI.chat 不支持的消息类型: {type(message)}. " "请使用 str, LLMMessage, 或 list[LLMContentPart]. " "对于更复杂的多模态输入或文件路径，请使用 AI.analyze().", code=LLMErrorCode.API_REQUEST_FAILED, ) final_messages = [*self.history, current_message] response = await self._execute_generation( messages=final_messages, model_name=model, error_message="聊天失败", config_overrides=kwargs, llm_tools=tools, tool_choice=tool_choice, ) should_preserve = ( preserve_media_in_history if preserve_media_in_history is not None else self.config.default_preserve_media_in_history ) if should_preserve: logger.debug("深度分析模式：在历史记录中保留原始多模态消息。") self.history.append(current_message) else: logger.debug("高效模式：净化历史记录中的多模态消息。") sanitized_user_message = self._sanitize_message_for_history(current_message) self.history.append(sanitized_user_message) self.history.append( LLMMessage( role="assistant", content=response.text, tool_calls=response.tool_calls ) ) return response async def code( self, prompt: str, *, model: ModelName = None, timeout: int | None = None, **kwargs: Any, ) -> dict[str, Any]: """ 代码执行 参数: prompt: 代码执行的提示词。 model: 要使用的模型名称。 timeout: 代码执行超时时间（秒）。 **kwargs: 传递给模型的其他参数。 返回: dict[str, Any]: 包含执行结果的字典，包含text、code_executions和success字段。 """ resolved_model = model or self.config.model or "Gemini/gemini-2.0-flash" config = CommonOverrides.gemini_code_execution() if timeout: config.custom_params = config.custom_params or {} config.custom_params["code_execution_timeout"] = timeout messages = [LLMMessage.user(prompt)] response = await self._execute_generation( messages=messages, model_name=resolved_model, error_message="代码执行失败", config_overrides=kwargs, base_config=config, ) return { "text": response.text, "code_executions": response.code_executions or [], "success": True, } async def search( self, query: str | UniMessage, *, model: ModelName = None, instruction: str = "", **kwargs: Any, ) -> dict[str, Any]: """ 信息搜索 - 支持多模态输入 参数: query: 搜索查询内容，支持文本或多模态消息。 model: 要使用的模型名称。 instruction: 搜索指令。 **kwargs: 传递给模型的其他参数。 返回: dict[str, Any]: 包含搜索结果的字典，包含text、sources、queries和success字段 """ from nonebot_plugin_alconna.uniseg import UniMessage resolved_model = model or self.config.model or "Gemini/gemini-2.0-flash" config = CommonOverrides.gemini_grounding() if isinstance(query, str): messages = [LLMMessage.user(query)] elif isinstance(query, UniMessage): content_parts = await unimsg_to_llm_parts(query) final_messages: list[LLMMessage] = [] if instruction: final_messages.append(LLMMessage.system(instruction)) if not content_parts: if instruction: final_messages.append(LLMMessage.user(instruction)) else: raise LLMException( "搜索内容为空或无法处理。", code=LLMErrorCode.API_REQUEST_FAILED ) else: final_messages.append(LLMMessage.user(content_parts)) messages = final_messages else: raise LLMException( f"不支持的搜索输入类型: {type(query)}. 请使用 str 或 UniMessage.", code=LLMErrorCode.API_REQUEST_FAILED, ) response = await self._execute_generation( messages=messages, model_name=resolved_model, error_message="信息搜索失败", config_overrides=kwargs, base_config=config, ) result = { "text": response.text, "sources": [], "queries": [], "success": True, } if response.grounding_metadata: result["sources"] = response.grounding_metadata.grounding_attributions or [] result["queries"] = response.grounding_metadata.web_search_queries or [] return result async def analyze( self, message: UniMessage | None, *, instruction: str = "", model: ModelName = None, use_tools: list[str] | None = None, tool_config: dict[str, Any] | None = None, activated_tools: list[LLMTool] | None = None, history: list[LLMMessage] | None = None, **kwargs: Any, ) -> LLMResponse: """ 内容分析 - 接收 UniMessage 物件进行多模态分析和工具呼叫。 参数: message: 要分析的消息内容（支持多模态）。 instruction: 分析指令。 model: 要使用的模型名称。 use_tools: 要使用的工具名称列表。 tool_config: 工具配置。 activated_tools: 已激活的工具列表。 history: 对话历史记录。 **kwargs: 传递给模型的其他参数。 返回: LLMResponse: 模型的完整响应结果。 """ from nonebot_plugin_alconna.uniseg import UniMessage content_parts = await unimsg_to_llm_parts(message or UniMessage()) final_messages: list[LLMMessage] = [] if history: final_messages.extend(history) if instruction: if not any(msg.role == "system" for msg in final_messages): final_messages.insert(0, LLMMessage.system(instruction)) if not content_parts: if instruction and not history: final_messages.append(LLMMessage.user(instruction)) elif not history: raise LLMException( "分析内容为空或无法处理。", code=LLMErrorCode.API_REQUEST_FAILED ) else: final_messages.append(LLMMessage.user(content_parts)) llm_tools: list[LLMTool] | None = activated_tools if not llm_tools and use_tools: try: llm_tools = tool_registry.get_tools(use_tools) logger.debug(f"已从注册表加载工具定义: {use_tools}") except ValueError as e: raise LLMException( f"加载工具定义失败: {e}", code=LLMErrorCode.CONFIGURATION_ERROR, cause=e, ) tool_choice = None if tool_config: mode = tool_config.get("mode", "auto") if mode in ["auto", "any", "none"]: tool_choice = mode response = await self._execute_generation( messages=final_messages, model_name=model, error_message="内容分析失败", config_overrides=kwargs, llm_tools=llm_tools, tool_choice=tool_choice, ) return response async def _execute_generation( self, messages: list[LLMMessage], model_name: ModelName, error_message: str, config_overrides: dict[str, Any], llm_tools: list[LLMTool] | None = None, tool_choice: str | dict[str, Any] | None = None, base_config: LLMGenerationConfig | None = None, ) -> LLMResponse: """通用的生成执行方法，封装模型获取和单次API调用""" try: resolved_model_name = self._resolve_model_name( model_name or self.config.model ) final_config_dict = self._merge_config( config_overrides, base_config=base_config ) async with await get_model_instance( resolved_model_name, override_config=final_config_dict ) as model_instance: return await model_instance.generate_response( messages, tools=llm_tools, tool_choice=tool_choice, ) except LLMException: raise except Exception as e: logger.error(f"{error_message}: {e}", e=e) raise LLMException(f"{error_message}: {e}", cause=e) def _resolve_model_name(self, model_name: ModelName) -> str: """解析模型名称""" if model_name: return model_name default_model = get_global_default_model_name() if default_model: return default_model raise LLMException( "未指定模型名称且未设置全局默认模型", code=LLMErrorCode.MODEL_NOT_FOUND, ) def _merge_config( self, user_config: dict[str, Any], base_config: LLMGenerationConfig | None = None, ) -> dict[str, Any]: """合并配置""" final_config = {} if base_config: final_config.update(base_config.to_dict()) if self.config.temperature is not None: final_config["temperature"] = self.config.temperature if self.config.max_tokens is not None: final_config["max_tokens"] = self.config.max_tokens if self.config.enable_cache: final_config["enable_caching"] = True if self.config.enable_code: final_config["enable_code_execution"] = True if self.config.enable_search: final_config["enable_grounding"] = True if self.config.enable_gemini_json_mode: final_config["response_mime_type"] = "application/json" if self.config.enable_gemini_thinking: final_config["thinking_budget"] = 0.8 if self.config.enable_gemini_safe_mode: final_config["safety_settings"] = ( CommonOverrides.gemini_safe().safety_settings ) if self.config.enable_gemini_multimodal: final_config.update(CommonOverrides.gemini_multimodal().to_dict()) if self.config.enable_gemini_grounding: final_config["enable_grounding"] = True final_config.update(user_config) return final_config async def embed( self, texts: list[str] | str, *, model: ModelName = None, task_type: EmbeddingTaskType | str = EmbeddingTaskType.RETRIEVAL_DOCUMENT, **kwargs: Any, ) -> list[list[float]]: """ 生成文本嵌入向量 参数: texts: 要生成嵌入向量的文本或文本列表。 model: 要使用的嵌入模型名称。 task_type: 嵌入任务类型。 **kwargs: 传递给模型的其他参数。 返回: list[list[float]]: 文本的嵌入向量列表。 """ if isinstance(texts, str): texts = [texts] if not texts: return [] try: resolved_model_str = ( model or self.config.default_embedding_model or self.config.model ) if not resolved_model_str: raise LLMException( "使用 embed 功能时必须指定嵌入模型名称，" "或在 AIConfig 中配置 default_embedding_model。", code=LLMErrorCode.MODEL_NOT_FOUND, ) resolved_model_str = self._resolve_model_name(resolved_model_str) async with await get_model_instance( resolved_model_str, override_config=None, ) as embedding_model_instance: return await embedding_model_instance.generate_embeddings( texts, task_type=task_type, **kwargs ) except LLMException: raise except Exception as e: logger.error(f"文本嵌入失败: {e}", e=e) raise LLMException( f"文本嵌入失败: {e}", code=LLMErrorCode.EMBEDDING_FAILED, cause=e )