79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
"""Data Agent — processes data, generates chart specs and table data."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from typing import Any
|
||
|
||
from .base import BaseAgent
|
||
from app.config import settings
|
||
|
||
|
||
class DataAgent(BaseAgent):
|
||
name = "data"
|
||
description = "处理数据、生成图表规格和表格数据"
|
||
system_prompt = """\
|
||
你是一位数据分析专家。你的任务是根据报告草稿中标注的图表和表格需求，
|
||
生成具体的数据和图表规格。
|
||
|
||
输出要求（JSON 格式）：
|
||
{
|
||
"charts": [
|
||
{
|
||
"id": "chart_1",
|
||
"title": "图表标题",
|
||
"type": "bar|line|pie|area|scatter",
|
||
"description": "图表说明",
|
||
"data": {
|
||
"labels": ["标签1", "标签2"],
|
||
"datasets": [
|
||
{"label": "数据集名", "data": [100, 200]}
|
||
]
|
||
}
|
||
}
|
||
],
|
||
"tables": [
|
||
{
|
||
"id": "table_1",
|
||
"title": "表格标题",
|
||
"headers": ["列1", "列2", "列3"],
|
||
"rows": [["数据1", "数据2", "数据3"]]
|
||
}
|
||
]
|
||
}"""
|
||
|
||
def __init__(self):
|
||
super().__init__(model=settings.model_for_domain("fast"))
|
||
|
||
async def run(self, context: dict[str, Any]) -> dict[str, Any]:
|
||
draft = context["draft"]
|
||
extra_data = context.get("extra_data", "")
|
||
|
||
# Collect chart/table needs from draft
|
||
chart_needs = []
|
||
table_needs = []
|
||
for ch in draft.get("chapters", []):
|
||
chart_needs.extend(ch.get("charts", []))
|
||
table_needs.extend(ch.get("tables", []))
|
||
|
||
if not chart_needs and not table_needs:
|
||
return {"data_assets": {"charts": [], "tables": []}}
|
||
|
||
prompt = f"""\
|
||
## 报告标题
|
||
{draft.get("title", "")}
|
||
|
||
## 需要生成的图表
|
||
{json.dumps(chart_needs, ensure_ascii=False)}
|
||
|
||
## 需要生成的表格
|
||
{json.dumps(table_needs, ensure_ascii=False)}
|
||
|
||
## 补充数据源
|
||
{extra_data if extra_data else "（无额外数据，请根据行业常识生成合理的示例数据）"}
|
||
|
||
请为以上需求生成具体的图表规格和表格数据。输出 JSON。"""
|
||
|
||
result = await self.call_llm_json(prompt)
|
||
return {"data_assets": result}
|