agent_builder.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. from typing import Annotated, TypedDict, Any
  2. from langchain.agents import create_agent, AgentState
  3. from langchain.chat_models import init_chat_model
  4. from langchain.tools import tool, InjectedState
  5. from langgraph.checkpoint.memory import InMemorySaver
  6. from config.settings import LLM_API_KEY, LLM_BASE_URL, LLM_MODEL, VENT_REVIEW_PROMPT, VENT_CALCULATION_PROMPT, VENT_COMPARISON_PROMPT, INTENT_RECOGNITION_PROMPT
  7. checkpointer = InMemorySaver()
  8. # 导入所有配风审查工具
  9. from tools.data_check_tools import (
  10. form_review,
  11. check_current_month_plan,
  12. check_data_by_gas_report,
  13. check_data_by_face_design,
  14. check_data_by_vent_report
  15. )
  16. from tools.clac_check_tools import (
  17. verify_coal_face_ventilation,
  18. verify_tunneling_face_ventilation,
  19. verify_chamber_ventilation,
  20. verify_other_points_ventilation
  21. )
  22. from tools.remote_clac_tools import (
  23. fetch_ventilation_data,
  24. recompute_ventilation,
  25. simulate_ventilation,
  26. fetch_field_translation
  27. )
  28. # 导入需风量对照与检查工具
  29. from tools.vent_comparison_tools import (
  30. get_plan_required_wind,
  31. get_model_calculated_wind,
  32. get_sensor_wind,
  33. merge_wind_by_location,
  34. check_wind_compliance
  35. )
  36. # ============================================================
  37. # SECTION 1: 自定义 Agent 状态
  38. # ============================================================
  39. class VentAgentState(AgentState):
  40. """配风审查 / 需风量对照 Agent 自定义状态"""
  41. vent_plan_data: dict[str, Any] | None # 挂载配风计划完整JSON数据
  42. plan_wind_result: dict[str, Any] | None # 工具1结果:配风计划需风量列表
  43. model_wind_result: dict[str, Any] | None # 工具2结果:模型解算风量列表
  44. sensor_wind_result: dict[str, Any] | None # 工具3结果:监测风量列表
  45. merged_wind_result: dict[str, Any] | None # 工具4结果:合并后的风量对照数据
  46. # ============================================================
  47. # SECTION 2: 模型初始化
  48. # ============================================================
  49. model = init_chat_model(
  50. model=LLM_MODEL,
  51. api_key=LLM_API_KEY,
  52. base_url=LLM_BASE_URL,
  53. temperature=0,
  54. streaming=True,
  55. model_provider="openai",
  56. extra_body={"enable_thinking": False},
  57. timeout=300
  58. )
  59. # ============================================================
  60. # SECTION 3: 工具列表
  61. # ============================================================
  62. review_tools = [
  63. form_review,
  64. check_current_month_plan,
  65. check_data_by_gas_report,
  66. check_data_by_face_design,
  67. check_data_by_vent_report,
  68. verify_coal_face_ventilation,
  69. verify_tunneling_face_ventilation,
  70. verify_chamber_ventilation,
  71. verify_other_points_ventilation,
  72. ]
  73. calculation_tools = [
  74. fetch_ventilation_data,
  75. recompute_ventilation,
  76. simulate_ventilation,
  77. fetch_field_translation,
  78. ]
  79. comparison_tools = [
  80. get_plan_required_wind, # 需风量提取
  81. get_model_calculated_wind, # 三维模型解算风量提取
  82. get_sensor_wind, # 巷道监测风量提取
  83. merge_wind_by_location, # 地点风量智能合并
  84. check_wind_compliance, # 风量检查
  85. ]
  86. # ============================================================
  87. # SECTION 4: 创建 Agent 实例
  88. # ============================================================
  89. review_agent = create_agent(
  90. model=model,
  91. tools=review_tools,
  92. system_prompt=VENT_REVIEW_PROMPT,
  93. debug=False,
  94. checkpointer=checkpointer,
  95. state_schema=VentAgentState
  96. )
  97. calculation_agent = create_agent(
  98. model=model,
  99. tools=calculation_tools,
  100. system_prompt=VENT_CALCULATION_PROMPT,
  101. checkpointer=checkpointer,
  102. debug=False
  103. )
  104. comparison_agent = create_agent(
  105. model=model,
  106. tools=comparison_tools,
  107. system_prompt=VENT_COMPARISON_PROMPT,
  108. debug=False,
  109. checkpointer=checkpointer,
  110. state_schema=VentAgentState
  111. )
  112. intent_agent = create_agent(
  113. model=model,
  114. tools=[],
  115. system_prompt=INTENT_RECOGNITION_PROMPT,
  116. checkpointer=checkpointer,
  117. debug=False
  118. )