replenishment.py
11.5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
"""
补货建议 Agent
重构版本：使用 part_ratio + SQL Agent + LangGraph
"""
import logging
import time
import uuid
from typing import Optional, List
from datetime import date, datetime
from decimal import Decimal
from langgraph.graph import StateGraph, END
from .state import AgentState
from .nodes import (
fetch_part_ratio_node,
sql_agent_node,
allocate_budget_node,
should_retry_sql,
)
from .analysis_report_node import generate_analysis_report_node
from ..models import ReplenishmentTask, TaskStatus, TaskExecutionLog, LogStatus, ReplenishmentPartSummary
from ..services import ResultWriter
logger = logging.getLogger(__name__)
class ReplenishmentAgent:
"""补货建议 Agent"""
def __init__(self):
self._graph = None
self._result_writer = ResultWriter()
@property
def graph(self) -> StateGraph:
"""获取工作流图"""
if self._graph is None:
self._graph = self._build_graph()
return self._graph
def _build_graph(self) -> StateGraph:
"""
构建 LangGraph 工作流
工作流结构：
fetch_part_ratio → sql_agent → allocate_budget → generate_analysis_report → END
"""
workflow = StateGraph(AgentState)
# 添加核心节点
workflow.add_node("fetch_part_ratio", fetch_part_ratio_node)
workflow.add_node("sql_agent", sql_agent_node)
workflow.add_node("allocate_budget", allocate_budget_node)
workflow.add_node("generate_analysis_report", generate_analysis_report_node)
# 设置入口
workflow.set_entry_point("fetch_part_ratio")
# 添加边
workflow.add_edge("fetch_part_ratio", "sql_agent")
# SQL Agent 条件边（支持重试）
workflow.add_conditional_edges(
"sql_agent",
should_retry_sql,
{
"retry": "sql_agent",
"continue": "allocate_budget",
}
)
# allocate_budget → generate_analysis_report → END
workflow.add_edge("allocate_budget", "generate_analysis_report")
workflow.add_edge("generate_analysis_report", END)
return workflow.compile()
def run(
self,
group_id: int,
dealer_grouping_id: int,
dealer_grouping_name: str,
brand_grouping_id: Optional[int] = None,
brand_grouping_name: str = "",
statistics_date: Optional[str] = None,
) -> AgentState:
"""
执行补货建议生成
Args:
group_id: 集团ID
dealer_grouping_id: 商家组合ID
dealer_grouping_name: 商家组合名称
brand_grouping_id: 品牌组合ID
brand_grouping_name: 品牌组合名称
statistics_date: 统计日期
"""
task_no = f"AI-{uuid.uuid4().hex[:12].upper()}"
if statistics_date is None:
statistics_date = date.today().strftime("%Y-%m-%d")
logger.info(
f"开始执行补货建议: task_no={task_no}, "
f"dealer_grouping={dealer_grouping_name}"
)
# 初始化状态
initial_state: AgentState = {
"task_no": task_no,
"group_id": group_id,
"brand_grouping_id": brand_grouping_id,
"brand_grouping_name": brand_grouping_name,
"dealer_grouping_id": dealer_grouping_id,
"dealer_grouping_name": dealer_grouping_name,
"statistics_date": statistics_date,
"part_ratios": [],
"sql_queries": [],
"sql_results": [],
"sql_retry_count": 0,
"sql_execution_logs": [],
"base_ratio": Decimal("1.1"),
"allocated_details": [],
"details": [],
"llm_suggestions": [],
"part_results": [],
"report": None,
"llm_provider": "",
"llm_model": "",
"llm_prompt_tokens": 0,
"llm_completion_tokens": 0,
"status": "running",
"error_message": "",
"start_time": time.time(),
"end_time": 0,
"current_node": "",
"next_node": "fetch_part_ratio",
}
# 创建任务记录
task = ReplenishmentTask(
task_no=task_no,
group_id=group_id,
dealer_grouping_id=dealer_grouping_id,
dealer_grouping_name=dealer_grouping_name,
brand_grouping_id=brand_grouping_id,
statistics_date=statistics_date,
status=TaskStatus.RUNNING,
)
self._result_writer.save_task(task)
try:
# 执行工作流
final_state = self.graph.invoke(initial_state)
# 更新任务状态
execution_time = int((final_state.get("end_time", time.time()) - final_state["start_time"]) * 1000)
actual_amount = sum(d.suggest_amount for d in final_state.get("details", []))
task.status = TaskStatus.SUCCESS
task.actual_amount = actual_amount
task.part_count = len(final_state.get("details", []))
task.shop_count = len(set(d.shop_id for d in final_state.get("details", [])))
task.base_ratio = final_state.get("base_ratio", Decimal("0"))
task.llm_provider = final_state.get("llm_provider", "")
task.llm_model = final_state.get("llm_model", "")
task.llm_prompt_tokens = final_state.get("llm_prompt_tokens", 0)
task.llm_completion_tokens = final_state.get("llm_completion_tokens", 0)
task.llm_total_tokens = task.llm_prompt_tokens + task.llm_completion_tokens
task.llm_analysis_summary = final_state.get("llm_analysis_summary", "")
task.execution_time_ms = execution_time
self._result_writer.update_task(task)
# 保存执行日志
if final_state.get("sql_execution_logs"):
self._save_execution_logs(
task_no=task_no,
group_id=group_id,
brand_grouping_id=brand_grouping_id,
brand_grouping_name=brand_grouping_name,
dealer_grouping_id=dealer_grouping_id,
dealer_grouping_name=dealer_grouping_name,
logs=final_state["sql_execution_logs"],
)
# 配件汇总已在 allocate_budget_node 中保存，此处跳过避免重复
# if final_state.get("part_results"):
# self._save_part_summaries(
# task_no=task_no,
# group_id=group_id,
# dealer_grouping_id=dealer_grouping_id,
# statistics_date=statistics_date,
# part_results=final_state["part_results"],
# )
logger.info(
f"补货建议执行完成: task_no={task_no}, "
f"parts={task.part_count}, amount={actual_amount}, "
f"time={execution_time}ms"
)
return final_state
except Exception as e:
logger.error(f"补货建议执行失败: task_no={task_no}, error={e}")
task.status = TaskStatus.FAILED
task.error_message = str(e)
task.execution_time_ms = int((time.time() - initial_state["start_time"]) * 1000)
self._result_writer.update_task(task)
raise
finally:
self._result_writer.close()
def _save_execution_logs(
self,
task_no: str,
group_id: int,
brand_grouping_id: Optional[int],
brand_grouping_name: str,
dealer_grouping_id: int,
dealer_grouping_name: str,
logs: List[dict],
):
"""保存执行日志"""
for log_data in logs:
log = TaskExecutionLog(
task_no=task_no,
group_id=group_id,
brand_grouping_id=brand_grouping_id,
brand_grouping_name=brand_grouping_name,
dealer_grouping_id=dealer_grouping_id,
dealer_grouping_name=dealer_grouping_name,
step_name=log_data.get("step_name", ""),
step_order=log_data.get("step_order", 0),
status=log_data.get("status", LogStatus.SUCCESS),
input_data=log_data.get("input_data", ""),
output_data=log_data.get("output_data", ""),
error_message=log_data.get("error_message", ""),
retry_count=log_data.get("retry_count", 0),
sql_query=log_data.get("sql_query", ""),
llm_prompt=log_data.get("llm_prompt", ""),
llm_response=log_data.get("llm_response", ""),
llm_tokens=log_data.get("llm_tokens", 0),
execution_time_ms=log_data.get("execution_time_ms", 0),
)
self._result_writer.save_execution_log(log)
def _save_part_summaries(
self,
task_no: str,
group_id: int,
dealer_grouping_id: int,
statistics_date: str,
part_results: list,
):
"""保存配件汇总"""
from .sql_agent import PartAnalysisResult
summaries = []
for pr in part_results:
if not isinstance(pr, PartAnalysisResult):
continue
summary = ReplenishmentPartSummary(
task_no=task_no,
group_id=group_id,
dealer_grouping_id=dealer_grouping_id,
part_code=pr.part_code,
part_name=pr.part_name,
unit=pr.unit,
cost_price=pr.cost_price,
total_storage_cnt=pr.total_storage_cnt,
total_avg_sales_cnt=pr.total_avg_sales_cnt,
group_current_ratio=pr.group_current_ratio,
total_suggest_cnt=pr.total_suggest_cnt,
total_suggest_amount=pr.total_suggest_amount,
shop_count=pr.shop_count,
need_replenishment_shop_count=pr.need_replenishment_shop_count,
part_decision_reason=pr.part_decision_reason,
priority=pr.priority,
llm_confidence=pr.confidence,
statistics_date=statistics_date,
)
summaries.append(summary)
if summaries:
self._result_writer.save_part_summaries(summaries)
logger.info(f"保存配件汇总: count={len(summaries)}")
def run_for_all_groupings(self, group_id: int):
"""
为所有商家组合执行补货建议
"""
from ..services import DataService
data_service = DataService()
try:
groupings = data_service.get_dealer_groupings(group_id)
logger.info(f"获取商家组合: group_id={group_id}, count={len(groupings)}")
for idx, grouping in enumerate(groupings):
logger.info(f"[{idx+1}/{len(groupings)}] 开始处理商家组合: {grouping['name']} (id={grouping['id']})")
try:
self.run(
group_id=group_id,
dealer_grouping_id=grouping["id"],
dealer_grouping_name=grouping["name"],
)
logger.info(f"[{grouping['name']}] 执行完成")
except Exception as e:
logger.error(f"商家组合执行失败: {grouping['name']}, error={e}", exc_info=True)
continue
finally:
data_service.close()