utils.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. '''
  4. @Project : lq-agent-api
  5. @File : utils.py
  6. @IDE : VsCode
  7. @Author :
  8. @Date : 2025-12-04 10:13:12
  9. =================================
  10. 📋 方法总览 (Method Overview)
  11. 🏗️ 核心模型管理:
  12. ├── ModelHandler() # 模型处理器类
  13. ├── get_models() # 获取模型的全局函数
  14. └── model_handler # 全局模型处理器实例
  15. 🔍 模型获取方法:
  16. ├── _get_doubao_model() # 获取豆包模型
  17. ├── _get_qwen_model() # 获取通义千问模型
  18. ├── _get_deepseek_model() # 获取DeepSeek模型
  19. ├── _get_gemini_model() # 获取Gemini模型
  20. ├── _get_lq_qwen3_8b_model() # 获取本地Qwen3-8B模型
  21. ├── _get_lq_qwen3_4b_model() # 获取本地Qwen3-4B模型
  22. └── _get_qwen_local_14b_model() # 获取本地Qwen3-14B模型
  23. '''
  24. from langchain_openai import ChatOpenAI
  25. from langchain_core.messages import HumanMessage
  26. from foundation.base.config import config_handler
  27. from foundation.logger.loggering import server_logger as logger
  28. class ModelHandler:
  29. """
  30. AI模型处理器类,用于管理多种AI模型的创建和配置
  31. 支持的模型类型:
  32. - doubao: 豆包模型
  33. - qwen: 通义千问模型
  34. - deepseek: DeepSeek模型
  35. - gemini: Gemini模型
  36. - lq_qwen3_8b: 本地Qwen3-8B模型
  37. - lq_qwen3_4b: 本地Qwen3-4B模型
  38. - qwen_local_14b: 本地Qwen3-14B模型
  39. """
  40. def __init__(self):
  41. """
  42. 初始化模型处理器
  43. 加载配置处理器,用于后续读取各种模型的配置信息
  44. """
  45. self.config = config_handler
  46. def get_models(self):
  47. """
  48. 获取AI模型实例
  49. Returns:
  50. ChatOpenAI: 配置好的AI模型实例
  51. Note:
  52. 根据配置文件中的MODEL_TYPE参数选择对应模型
  53. 支持的模型类型:doubao, qwen, deepseek, lq_qwen3_8b, lq_qwen3_4b, qwen_local_14b
  54. 默认返回豆包模型
  55. """
  56. model_type = self.config.get("model", "MODEL_TYPE")
  57. logger.info(f"正在初始化AI模型,模型类型: {model_type}")
  58. if model_type == "doubao":
  59. model = self._get_doubao_model()
  60. elif model_type == "gemini":
  61. model = self._get_gemini_model()
  62. elif model_type == "qwen":
  63. model = self._get_qwen_model()
  64. elif model_type == "deepseek":
  65. model = self._get_deepseek_model()
  66. elif model_type == "lq_qwen3_8b":
  67. model = self._get_lq_qwen3_8b_model()
  68. elif model_type == "lq_qwen3_4b":
  69. model = self._get_lq_qwen3_4b_model()
  70. elif model_type == "qwen_local_14b":
  71. model = self._get_qwen_local_14b_model()
  72. else:
  73. # 默认返回gemini
  74. logger.warning(f"未知的模型类型 '{model_type}',使用默认gemini模型")
  75. model = self._get_gemini_model()
  76. logger.info(f"AI模型初始化完成: {model_type}")
  77. return model
  78. def _get_doubao_model(self):
  79. """
  80. 获取豆包模型
  81. Returns:
  82. ChatOpenAI: 配置好的豆包模型实例
  83. """
  84. doubao_url = self.config.get("doubao", "DOUBAO_SERVER_URL")
  85. doubao_model_id = self.config.get("doubao", "DOUBAO_MODEL_ID")
  86. doubao_api_key = self.config.get("doubao", "DOUBAO_API_KEY")
  87. llm = ChatOpenAI(
  88. base_url=doubao_url,
  89. model=doubao_model_id,
  90. api_key=doubao_api_key,
  91. temperature=0.7,
  92. extra_body={
  93. "enable_thinking": False,
  94. })
  95. return llm
  96. def _get_qwen_model(self):
  97. """
  98. 获取通义千问模型
  99. Returns:
  100. ChatOpenAI: 配置好的通义千问模型实例
  101. """
  102. qwen_url = self.config.get("qwen", "QWEN_SERVER_URL")
  103. qwen_model_id = self.config.get("qwen", "QWEN_MODEL_ID")
  104. qwen_api_key = self.config.get("qwen", "QWEN_API_KEY")
  105. llm = ChatOpenAI(
  106. base_url=qwen_url,
  107. model=qwen_model_id,
  108. api_key=qwen_api_key,
  109. temperature=0.7,
  110. extra_body={
  111. "enable_thinking": False,
  112. })
  113. return llm
  114. def _get_deepseek_model(self):
  115. """
  116. 获取DeepSeek模型
  117. Returns:
  118. ChatOpenAI: 配置好的DeepSeek模型实例
  119. """
  120. deepseek_url = self.config.get("deepseek", "DEEPSEEK_SERVER_URL")
  121. deepseek_model_id = self.config.get("deepseek", "DEEPSEEK_MODEL_ID")
  122. deepseek_api_key = self.config.get("deepseek", "DEEPSEEK_API_KEY")
  123. llm = ChatOpenAI(
  124. base_url=deepseek_url,
  125. model=deepseek_model_id,
  126. api_key=deepseek_api_key,
  127. temperature=0.7,
  128. extra_body={
  129. "enable_thinking": False,
  130. })
  131. return llm
  132. def _get_gemini_model(self):
  133. """
  134. 获取Gemini模型
  135. Returns:
  136. ChatOpenAI: 配置好的Gemini模型实例
  137. """
  138. gemini_url = self.config.get("gemini", "GEMINI_SERVER_URL")
  139. gemini_model_id = self.config.get("gemini", "GEMINI_MODEL_ID")
  140. gemini_api_key = self.config.get("gemini", "GEMINI_API_KEY")
  141. llm = ChatOpenAI(
  142. base_url=gemini_url,
  143. model=gemini_model_id,
  144. api_key=gemini_api_key,
  145. temperature=0.7,
  146. # extra_body={
  147. # "enable_thinking": False,
  148. # }
  149. )
  150. return llm
  151. def _get_lq_qwen3_8b_model(self):
  152. """
  153. 获取本地Qwen3-8B-Instruct模型
  154. Returns:
  155. ChatOpenAI: 配置好的本地Qwen3-8B模型实例
  156. """
  157. llm = ChatOpenAI(
  158. base_url="http://192.168.91.253:9002/v1",
  159. model="Qwen3-8B",
  160. api_key="dummy", # 本地模型使用虚拟API key
  161. temperature=0.7,
  162. )
  163. return llm
  164. def _get_lq_qwen3_4b_model(self):
  165. """
  166. 获取本地Qwen3-4B-Instruct模型
  167. Returns:
  168. ChatOpenAI: 配置好的本地Qwen3-4B模型实例
  169. """
  170. llm = ChatOpenAI(
  171. base_url="http://192.168.91.253:9001/v1",
  172. model="Qwen3-4B",
  173. api_key="dummy", # 本地模型使用虚拟API key
  174. temperature=0.7,
  175. )
  176. return llm
  177. # 创建全局实例
  178. model_handler = ModelHandler()
  179. def get_models():
  180. """
  181. 获取模型的全局函数
  182. Returns:
  183. tuple: (llm, chat, embed) - LLM模型、聊天模型和嵌入模型实例
  184. 注意:当前llm和chat使用相同模型实例,embed暂时返回None
  185. Note:
  186. 这是一个便捷函数,直接使用全局model_handler实例获取模型
  187. """
  188. llm = model_handler.get_models()
  189. # 暂时返回相同的模型作为chat和embed
  190. return llm, llm, None