workflow_state.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391
  1. """
  2. 整合的工作流状态定义
  3. ===================
  4. 此文件定义了整合了多个Agent的工作流状态,兼容现有的Big Agent状态管理和新增的报告生成Agent状态。
  5. 状态层次:
  6. 1. 输入层:用户查询和数据
  7. 2. 意图层:意图识别结果
  8. 3. 规划层:规划决策和大纲生成
  9. 4. 计算层:指标计算结果
  10. 5. 结果层:最终报告生成
  11. 6. 对话层:消息历史和错误处理
  12. 兼容性:
  13. - 兼容现有的Big Agent WorkflowState
  14. - 整合来自other_agents的AgentState
  15. - 支持扩展新的Agent状态需求
  16. 作者: Big Agent Team
  17. 版本: 1.0.0
  18. 创建时间: 2024-12-20
  19. """
  20. from typing import TypedDict, List, Dict, Any, Optional
  21. from datetime import datetime
  22. from langchain_core.messages import BaseMessage
  23. from pydantic import BaseModel, Field
  24. # ============= 数据模型 =============
  25. class MetricRequirement(BaseModel):
  26. """指标需求定义"""
  27. metric_id: str = Field(description="指标唯一标识,如 'total_income_jan'")
  28. metric_name: str = Field(description="指标中文名称")
  29. calculation_logic: str = Field(description="计算逻辑描述")
  30. required_fields: List[str] = Field(description="所需字段")
  31. dependencies: List[str] = Field(default_factory=list, description="依赖的其他指标ID")
  32. class ReportSection(BaseModel):
  33. """报告大纲章节"""
  34. section_id: str = Field(description="章节ID")
  35. title: str = Field(description="章节标题")
  36. description: str = Field(description="章节内容要求")
  37. metrics_needed: List[str] = Field(description="所需指标ID列表")
  38. class ReportOutline(BaseModel):
  39. """完整报告大纲"""
  40. report_title: str = Field(description="报告标题")
  41. sections: List[ReportSection] = Field(description="章节列表")
  42. global_metrics: List[MetricRequirement] = Field(description="全局指标列表")
  43. # ============= 序列化工具函数 =============
  44. def convert_numpy_types(obj: Any) -> Any:
  45. """
  46. 递归转换所有numpy类型为Python原生类型
  47. 确保所有数据可序列化
  48. """
  49. if isinstance(obj, dict):
  50. return {str(k): convert_numpy_types(v) for k, v in obj.items()}
  51. elif isinstance(obj, list):
  52. return [convert_numpy_types(item) for item in obj]
  53. elif isinstance(obj, tuple):
  54. return tuple(convert_numpy_types(item) for item in obj)
  55. elif isinstance(obj, set):
  56. return {convert_numpy_types(item) for item in obj}
  57. elif hasattr(obj, 'item') and hasattr(obj, 'dtype'): # numpy scalar
  58. return convert_numpy_types(obj.item())
  59. else:
  60. return obj
  61. # ============= 整合的工作流状态定义 =============
  62. class IntegratedWorkflowState(TypedDict):
  63. """整合的工作流状态定义 - 兼容多个Agent系统"""
  64. # === 基础输入层 (兼容Big Agent) ===
  65. user_input: str
  66. question: str # 别名,兼容报告生成Agent
  67. industry: str # 行业
  68. # === 数据层 ===
  69. data_set: List[Dict[str, Any]] # 报告生成Agent的数据格式
  70. transactions_df: Optional[Any] # 可选的数据框格式
  71. file_name: str # 数据文件名称
  72. data_set_classified: List[Dict[str, Any]] # 分类打标后的数据集
  73. original_file_path: str # 上传文件绝对路径
  74. is_standardized: int # 数据是否已经标准化 0: 否 1: 是
  75. standardized_file_path: str # 数据标准化的文件路径
  76. # === 意图识别层 (Big Agent原有) ===
  77. intent_result: Optional[Dict[str, Any]]
  78. # === 规划和大纲层 (新增) ===
  79. planning_step: int
  80. plan_history: List[str]
  81. outline_draft: Optional[ReportOutline]
  82. outline_version: int
  83. outline_ready: bool
  84. # === 指标计算层 ===
  85. metrics_requirements: List[MetricRequirement] # 报告生成Agent格式
  86. computed_metrics: Dict[str, Any] # 计算结果
  87. metrics_cache: Dict[str, Any] # 缓存
  88. pending_metric_ids: List[str] # 待计算指标ID
  89. failed_metric_attempts: Dict[str, int] # 失败统计
  90. calculation_results: Optional[Dict[str, Any]] # Big Agent格式的计算结果
  91. # === 结果层 ===
  92. report_draft: Dict[str, Any] # 报告草稿
  93. knowledge_result: Optional[Dict[str, Any]] # Big Agent知识沉淀结果
  94. is_complete: bool
  95. completeness_score: float
  96. answer: Optional[str] # 最终答案
  97. # === 对话和消息层 ===
  98. messages: List[Dict[str, Any]] # Big Agent消息格式
  99. current_node: str
  100. session_id: str
  101. next_route: str
  102. # === 错误处理层 ===
  103. errors: List[str]
  104. last_decision: str
  105. # === 时间跟踪层 ===
  106. start_time: str
  107. end_time: Optional[str]
  108. api_result: Dict[str, Any] # 存储所有API调用结果
  109. # ============= 状态创建和初始化函数 =============
  110. def create_initial_integrated_state(question: str, industry: str, original_file_path: str, session_id: str = None) -> IntegratedWorkflowState:
  111. """
  112. 创建初始的整合状态
  113. Args:
  114. question: 用户查询
  115. industry: 行业
  116. original_file_path: 原始数据文件
  117. session_id: 会话ID
  118. Returns:
  119. 初始化后的状态
  120. """
  121. current_time = datetime.now().isoformat()
  122. session = session_id or f"session_{int(datetime.now().timestamp())}"
  123. return {
  124. # 基础输入
  125. "user_input": question,
  126. "question": question,
  127. "industry": industry,
  128. "original_file_path": original_file_path,
  129. # 数据层
  130. "is_standardized": 0, # 未标准化
  131. "data_set": [],
  132. "data_set_classified": [], # 分类打标后的数据集
  133. "transactions_df": None,
  134. "file_name": "", # 文件名称
  135. # 意图识别层
  136. "intent_result": None,
  137. # 规划和大纲层
  138. "planning_step": 0,
  139. "plan_history": [],
  140. "outline_draft": None,
  141. "outline_version": 0,
  142. "outline_ready": False,
  143. # 指标计算层
  144. "metrics_requirements": [],
  145. "computed_metrics": {},
  146. "metrics_cache": {},
  147. "pending_metric_ids": [],
  148. "failed_metric_attempts": {},
  149. "calculation_results": None,
  150. # 结果层
  151. "report_draft": {},
  152. "knowledge_result": None,
  153. "is_complete": False,
  154. "completeness_score": 0.0,
  155. "answer": None,
  156. # 对话和消息层
  157. "messages": [{
  158. "role": "user",
  159. "content": question,
  160. "timestamp": current_time
  161. }],
  162. "current_node": "start",
  163. "session_id": session,
  164. "next_route": "planning_node",
  165. # 错误处理层
  166. "errors": [],
  167. "last_decision": "init",
  168. # 时间跟踪层
  169. "start_time": current_time,
  170. "end_time": None,
  171. "api_result": {}, # 存储所有API调用结果
  172. # 计算模式配置层
  173. "use_rules_engine_only": False,
  174. "use_traditional_engine_only": False
  175. }
  176. def is_state_ready_for_calculation(state: IntegratedWorkflowState) -> bool:
  177. """
  178. 检查状态是否准备好进行指标计算
  179. Args:
  180. state: 当前状态
  181. Returns:
  182. 是否准备好
  183. """
  184. return (
  185. state.get("outline_draft") is not None and
  186. len(state.get("metrics_requirements", [])) > 0 and
  187. len(state.get("pending_metric_ids", [])) > 0
  188. )
  189. def get_calculation_progress(state: IntegratedWorkflowState) -> Dict[str, Any]:
  190. """
  191. 获取指标计算进度信息
  192. Args:
  193. state: 当前状态
  194. Returns:
  195. 进度信息
  196. """
  197. required = len(state.get("metrics_requirements", []))
  198. computed = len(state.get("computed_metrics", {}))
  199. pending = len(state.get("pending_metric_ids", []))
  200. return {
  201. "required_count": required,
  202. "computed_count": computed,
  203. "pending_count": pending,
  204. "coverage_rate": computed / required if required > 0 else 0,
  205. "is_complete": computed >= required * 0.8 # 80%覆盖率视为完成
  206. }
  207. def update_state_with_outline_generation(state: IntegratedWorkflowState, outline: ReportOutline) -> IntegratedWorkflowState:
  208. """
  209. 使用大纲生成结果更新状态
  210. Args:
  211. state: 当前状态
  212. outline: 生成的大纲
  213. Returns:
  214. 更新后的状态
  215. """
  216. new_state = state.copy()
  217. new_state["outline_draft"] = outline
  218. new_state["outline_version"] += 1
  219. new_state["outline_ready"] = True
  220. new_state["metrics_requirements"] = outline.global_metrics
  221. new_state["pending_metric_ids"] = [m.metric_id for m in outline.global_metrics]
  222. # 添加消息
  223. new_state["messages"].append({
  224. "role": "assistant",
  225. "content": f"✅ 大纲生成完成 v{new_state['outline_version']}:{outline.report_title}",
  226. "timestamp": datetime.now().isoformat()
  227. })
  228. return new_state
  229. def update_state_with_planning_decision(state: IntegratedWorkflowState, decision: Dict[str, Any]) -> IntegratedWorkflowState:
  230. """
  231. 使用规划决策结果更新状态
  232. Args:
  233. state: 当前状态
  234. decision: 规划决策
  235. Returns:
  236. 更新后的状态
  237. """
  238. new_state = state.copy()
  239. new_state["planning_step"] += 1
  240. new_state["last_decision"] = decision.get("decision", "unknown")
  241. new_state["next_route"] = decision.get("next_route", "planning_node")
  242. # 如果有待计算指标,更新待计算列表
  243. if decision.get("metrics_to_compute"):
  244. new_state["pending_metric_ids"] = decision["metrics_to_compute"]
  245. # 添加规划历史
  246. new_state["plan_history"].append(
  247. f"Step {new_state['planning_step']}: {decision.get('decision', 'unknown')}"
  248. )
  249. return new_state
  250. def update_state_with_report(state: IntegratedWorkflowState, final_report: Dict[str, Any]) -> IntegratedWorkflowState:
  251. """
  252. 使用最终报告完成状态
  253. Args:
  254. state: 当前状态
  255. final_report: 最终报告
  256. Returns:
  257. 完成的状态
  258. """
  259. new_state = state.copy()
  260. new_state["report_draft"] = final_report
  261. new_state["is_complete"] = True
  262. new_state["answer"] = final_report
  263. new_state["end_time"] = datetime.now().isoformat()
  264. # 计算完整性分数
  265. progress = get_calculation_progress(new_state)
  266. new_state["completeness_score"] = progress["coverage_rate"]
  267. return new_state
  268. def update_state_with_data_classified(state: IntegratedWorkflowState, data_set_classified: List[Dict]) -> IntegratedWorkflowState:
  269. """
  270. 使用分类打标结果更新状态
  271. Args:
  272. state: 当前状态
  273. data_set_classified: 分类打标的数据
  274. Returns:
  275. 更新后的状态
  276. """
  277. new_state = state.copy()
  278. new_state["data_set_classified"] = data_set_classified
  279. # 添加消息
  280. new_state["messages"].append({
  281. "role": "assistant",
  282. "content": f"✅ 数据分类打标已完成",
  283. "timestamp": datetime.now().isoformat()
  284. })
  285. return new_state
  286. def update_state_with_data_standardize(state: IntegratedWorkflowState, is_succ: int, standardized_file_path: str) -> IntegratedWorkflowState:
  287. """
  288. 根据数据标准化结果更新状态
  289. Args:
  290. state: 当前状态
  291. is_succ: 是否标准化数据成功 0:否 1:是
  292. standardized_file_path: 标准化后的文件路径
  293. Returns:
  294. 更新后的状态
  295. """
  296. import os
  297. new_state = state.copy()
  298. new_state["is_standardized"] = is_succ
  299. new_state["standardized_file_path"] = standardized_file_path
  300. new_state["file_name"] = os.path.basename(standardized_file_path)
  301. msg = "成功" if is_succ else "失败"
  302. # 添加消息
  303. new_state["messages"].append({
  304. "role": "assistant",
  305. "content": f"✅ 数据标准化完成,处理结果:{msg}",
  306. "timestamp": datetime.now().isoformat()
  307. })
  308. return new_state