chat_model.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980
  1. # -*- coding: utf-8 -*-
  2. """
  3. LLM 模型实例缓存 —— 避免每次请求重复 init_chat_model
  4. 直接读取 .env 配置,确保模型字符串带有 provider 前缀(如 openai:),
  5. 避免 init_chat_model 根据模型名猜测错误的 provider(如将 deepseek-v4-pro
  6. 解析为 ChatDeepSeek 而非 ChatOpenAI)。
  7. """
  8. # Fast 模式模型实例缓存(模块级单例,避免每次请求重复 init_chat_model)
  9. # _fast_model: fast_mode=True 使用 SUMMARY_MODEL(低延迟 flash 模型)
  10. # _chat_model: fast_mode=False 使用 DEEPAGENT_MODEL(pro 模型)
  11. _fast_model = None
  12. _chat_model = None
  13. def get_chat_model(fast_mode: bool = True):
  14. """获取 LangChain ChatModel 实例,用于 Fast 模式直接调用 LLM。
  15. fast_mode=True 时使用 SUMMARY_MODEL(默认 deepseek-v4-flash),
  16. 延迟更低,适合直接模板填充任务;fast_mode=False 时使用 DEEPAGENT_MODEL。
  17. 模型实例按 mode 分别缓存,后续请求直接复用。
  18. """
  19. global _fast_model, _chat_model
  20. if fast_mode and _fast_model is not None:
  21. return _fast_model
  22. if not fast_mode and _chat_model is not None:
  23. return _chat_model
  24. from dotenv import dotenv_values
  25. from pathlib import Path as _Path
  26. from langchain.chat_models import init_chat_model
  27. _env_file = _Path(__file__).parent.parent / ".env"
  28. _cfg = dotenv_values(str(_env_file))
  29. # Fast 模式优先使用 SUMMARY_MODEL(低延迟),否则 fallback 到 DEEPAGENT_MODEL
  30. if fast_mode:
  31. model_str = _cfg.get("SUMMARY_MODEL", "") or _cfg.get("DEEPAGENT_MODEL", "openai:gpt-4o")
  32. else:
  33. model_str = _cfg.get("DEEPAGENT_MODEL", "openai:gpt-4o")
  34. base_url = _cfg.get("OPENAI_BASE_URL", "")
  35. api_key = _cfg.get("OPENAI_API_KEY", "")
  36. timeout_str = _cfg.get("OPENAI_TIMEOUT", "180")
  37. # 清理可能带入的引号
  38. api_key = api_key.strip().strip('"').strip("'") if api_key else ""
  39. base_url = base_url.strip().strip('"').strip("'") if base_url else ""
  40. # 解析超时时间
  41. try:
  42. timeout = float(timeout_str.strip())
  43. except (ValueError, TypeError):
  44. timeout = 180.0
  45. # 确保模型字符串带有 provider 前缀,否则 init_chat_model 可能猜错 provider
  46. if ":" not in model_str:
  47. model_str = f"openai:{model_str}"
  48. if base_url and api_key:
  49. model_instance = init_chat_model(
  50. model_str,
  51. openai_api_key=api_key,
  52. openai_api_base=base_url,
  53. temperature=0,
  54. timeout=timeout,
  55. max_retries=2,
  56. )
  57. else:
  58. model_instance = init_chat_model(model_str, temperature=0, timeout=timeout, max_retries=2)
  59. # 按模式分别缓存
  60. if fast_mode:
  61. _fast_model = model_instance
  62. else:
  63. _chat_model = model_instance
  64. return model_instance