from loguru import logger from typing import List, Dict, Optional, Union, AsyncGenerator, Generator, Any # 核心：导入 LangChain 的基础语言模型抽象类 from langchain_core.language_models import BaseLanguageModel from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatResult from langchain_core.callbacks import CallbackManagerForLLMRun from dataclasses import dataclass, field from typing import Optional, Dict, Any, List from datetime import datetime @dataclass class LLMConfig_DataClass: """ 统一的LLM配置基类，覆盖在线/本地/嵌入式模型所有配置，映射数据库完整字段 通过 provider + is_embedding 区分模型类型： - 在线模型：provider in ['openai', 'zhipu', 'baidu'] + is_embedding=False - 本地模型：provider in ['llama', 'qwen', 'yi'] + is_embedding=False - 嵌入式模型：provider in ['bge', 'text2vec'] + is_embedding=True """ # ====================== 数据库核心公共字段（必选/可选） ====================== # 基础标识字段 name: str # 模型自定义名称（如 "gpt-5"） model_name: str # 模型官方标识名（如 "gpt-5"、"BAAI/bge-small-zh-v1.5"） provider: str # 提供商（openai/llama/bge/zhipu 等） id: Optional[int] = None # 数据库主键ID description: Optional[str] = None # 模型描述 is_active: bool = True # 是否启用 is_default: bool = False # 是否默认模型 is_embedding: bool = False # 是否为嵌入式模型（核心区分标识） # ====================== 通用生成参数（所有推理模型共用） ====================== temperature: float = 0.7 # 生成温度（默认值对齐数据库示例） max_tokens: int = 3000 # 最大生成长度（默认值对齐数据库示例） top_p: float = 0.6 # 采样Top-P frequency_penalty: float = 0.0 # 频率惩罚 presence_penalty: float = 0.0 # 存在惩罚 # ====================== 在线模型专属参数（非必填，仅在线模型生效） ====================== api_key: Optional[str] = None # API密钥（在线模型必填） base_url: Optional[str] = None # API代理地址（如 https://api.openai-proxy.org/v1） # timeout: int = 30 # 请求超时时间（秒） max_retries: int = 3 # 最大重试次数 api_version: Optional[str] = None # API版本（如 OpenAI 的 2024-02-15-preview） # ====================== 本地模型专属参数（非必填，仅本地模型生效） ====================== model_path: Optional[str] = None # 本地模型文件路径（本地模型必填） device: str = "cpu" # 运行设备（cpu/cuda/mps） n_ctx: int = 2048 # 上下文窗口大小 n_threads: int = 8 # 推理线程数 quantization: str = "q4_0" # 量化级别（q4_0/q8_0/f16） load_in_8bit: bool = False # 是否8bit加载 load_in_4bit: bool = False # 是否4bit加载 prompt_template: Optional[str] = None # 自定义Prompt模板 # ====================== 嵌入式模型专属参数（非必填，仅嵌入式模型生效） ====================== normalize_embeddings: bool = True # 是否归一化向量 batch_size: int = 32 # 批量编码大小 encode_kwargs: Dict[str, Any] = field(default_factory=dict) # 编码扩展参数 dimension: Optional[int] = None # 向量维度（如 768） # ====================== 元数据字段（数据库自动维护） ====================== extra_config: Dict[str, Any] = field(default_factory=dict) # 额外扩展配置 usage_count: int = 0 # 使用次数 last_used_at: Optional[datetime] = None # 最后使用时间 created_at: Optional[datetime] = None # 创建时间 updated_at: Optional[datetime] = None # 更新时间 created_by: Optional[int] = None # 创建人ID updated_by: Optional[int] = None # 更新人ID api_key_masked: Optional[str] = "" # 掩码后的API密钥（数据库存储） # ====================== 核心工具方法 ====================== def __post_init__(self): """后置初始化：自动校验和修正配置""" # 1. 嵌入式模型强制清空推理参数（避免误用） if self.is_embedding: self.max_tokens = 0 self.temperature = 0.0 self.top_p = 0.0 # 2. 校验必填参数（按模型类型） self._validate_required_fields() def _validate_required_fields(self): """按模型类型校验必填参数""" # 在线模型校验 if not self.is_embedding and self.provider in ['openai', 'zhipu', 'baidu', 'anthropic']: if not self.api_key: raise ValueError(f"[{self.name}] 在线模型({self.provider})必须配置 api_key") # 本地模型校验 if not self.is_embedding and self.provider in ['llama', 'qwen', 'yi', 'glm', 'mistral']: if not self.model_path: raise ValueError(f"[{self.name}] 本地模型({self.provider})必须配置 model_path") def to_dict(self) -> Dict[str, Any]: """转换为字典（用于存入/更新数据库）""" return { key: value for key, value in self.__dict__.items() if not key.startswith('_') # 排除私有属性 } @classmethod def from_db_dict(cls, db_dict: Dict[str, Any]) -> "LLMConfig_DataClass": """从数据库字典初始化配置（核心方法）""" # 1. 时间字段转换：字符串 → datetime time_fields = ['last_used_at', 'created_at', 'updated_at'] for field_name in time_fields: val = db_dict.get(field_name) if val and isinstance(val, str): try: db_dict[field_name] = datetime.fromisoformat(val.replace('Z', '+00:00')) except (ValueError, TypeError): db_dict[field_name] = None # 2. 过滤数据库中无关字段（如 api_key_masked） valid_fields = cls.__dataclass_fields__.keys() filtered_dict = {k: v for k, v in db_dict.items() if k in valid_fields} # 3. 初始化并返回配置实例 return cls(**filtered_dict) def get_model_type(self) -> str: """快速判断模型类型（返回：online/local/embedding）""" if self.is_embedding: return "embedding" if self.provider in ['openai', 'zhipu', 'baidu', 'anthropic']: return "online" if self.provider in ['llama', 'qwen', 'yi', 'glm', 'mistral']: return "local" return "unknown" class BaseLLM(BaseChatModel): """ 继承 LangChain 的 BaseChatModel（BaseLanguageModel 的子类） 使其能直接用于 create_agent """ # 配置参数（通过 __init__ 初始化） config: Any = None model: Any = None def __init__(self, config): super().__init__() # 必须调用父类构造函数 self.config = config self.model = None self._validate_config() logger.info(f"初始化 {self.__class__.__name__}，模型: {config.model_name}") # ---------------------- 必须实现的核心抽象方法（LangChain 协议） ---------------------- def _generate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, **kwargs: Any, ) -> ChatResult: """ 核心同步生成方法（LangChain 要求必须实现） messages: 消息列表（如 [HumanMessage(content="你好")]） 返回 ChatResult 类型（LangChain 标准输出） """ logger.error(f"{self.__class__.__name__} 未实现 同步 _generate 方法") async def _agenerate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None, run_manager: Optional[CallbackManagerForLLMRun] = None, ** kwargs: Any, ) -> ChatResult: """异步生成方法（LangChain 异步协议）""" logger.error(f"{self.__class__.__name__} 未实现 异步 _agenerate 方法") @property def _llm_type(self) -> str: """返回模型类型标识（如 "openai"、"llama"、"bge"）""" return self.__class__.__name__ def load_model(self) -> None: """加载模型（自定义逻辑）""" logger.error(f"{self.__class__.__name__} 未实现 load_model 方法") def close(self) -> None: """释放资源（自定义逻辑）""" if self.model: logger.info(f"释放 {self.__class__.__name__} 模型资源") self.model = None def __enter__(self): self.load_model() return self def __exit__(self, exc_type, exc_val, exc_tb): self.close()