init repo

app/agents/data_agent.py

@@ -0,0 +1,78 @@
+"""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}