262 lines
8.9 KiB
Python
262 lines
8.9 KiB
Python
"""
|
||
AI-Driven Data Analysis Framework
|
||
==================================
|
||
Generic pipeline: works with any CSV file.
|
||
AI decides everything at runtime based on data characteristics.
|
||
|
||
Flow: Load CSV → AI understands metadata → AI plans analysis → AI executes tasks → AI generates report
|
||
"""
|
||
import sys
|
||
import os
|
||
sys.path.insert(0, os.path.dirname(__file__))
|
||
|
||
from src.env_loader import load_env_with_fallback
|
||
load_env_with_fallback(['.env'])
|
||
|
||
import logging
|
||
import json
|
||
from datetime import datetime
|
||
from typing import Optional, List
|
||
from openai import OpenAI
|
||
|
||
from src.engines.ai_data_understanding import ai_understand_data_with_dal
|
||
from src.engines.requirement_understanding import understand_requirement
|
||
from src.engines.analysis_planning import plan_analysis
|
||
from src.engines.task_execution import execute_task
|
||
from src.engines.report_generation import generate_report
|
||
from src.tools.tool_manager import ToolManager
|
||
from src.tools.base import _global_registry
|
||
from src.models import DataProfile, AnalysisResult
|
||
from src.config import get_config
|
||
|
||
# Register all tools
|
||
from src.tools import register_tool
|
||
from src.tools.query_tools import (
|
||
GetColumnDistributionTool, GetValueCountsTool,
|
||
GetTimeSeriesTool, GetCorrelationTool
|
||
)
|
||
from src.tools.stats_tools import (
|
||
CalculateStatisticsTool, PerformGroupbyTool,
|
||
DetectOutliersTool, CalculateTrendTool
|
||
)
|
||
from src.tools.viz_tools import (
|
||
CreateBarChartTool, CreateLineChartTool,
|
||
CreatePieChartTool, CreateHeatmapTool
|
||
)
|
||
|
||
for tool_cls in [
|
||
GetColumnDistributionTool, GetValueCountsTool, GetTimeSeriesTool, GetCorrelationTool,
|
||
CalculateStatisticsTool, PerformGroupbyTool, DetectOutliersTool, CalculateTrendTool,
|
||
CreateBarChartTool, CreateLineChartTool, CreatePieChartTool, CreateHeatmapTool
|
||
]:
|
||
register_tool(tool_cls())
|
||
|
||
logging.basicConfig(level=logging.WARNING)
|
||
|
||
|
||
def run_analysis(
|
||
data_file: str,
|
||
user_requirement: Optional[str] = None,
|
||
template_file: Optional[str] = None,
|
||
output_dir: str = "analysis_output"
|
||
):
|
||
"""
|
||
Run the full AI-driven analysis pipeline.
|
||
|
||
Args:
|
||
data_file: Path to any CSV file
|
||
user_requirement: Natural language requirement (optional)
|
||
template_file: Report template path (optional)
|
||
output_dir: Output directory
|
||
"""
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
config = get_config()
|
||
|
||
print("\n" + "=" * 70)
|
||
print("AI-Driven Data Analysis")
|
||
print("=" * 70)
|
||
print(f"Start: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print(f"Data: {data_file}")
|
||
if template_file:
|
||
print(f"Template: {template_file}")
|
||
print("=" * 70)
|
||
|
||
# ── Stage 1: AI Data Understanding ──
|
||
print("\n[1/5] AI Understanding Data...")
|
||
print(" (AI sees only metadata — never raw rows)")
|
||
profile, dal = ai_understand_data_with_dal(data_file)
|
||
print(f" Type: {profile.inferred_type}")
|
||
print(f" Quality: {profile.quality_score}/100")
|
||
print(f" Columns: {profile.column_count}, Rows: {profile.row_count}")
|
||
print(f" Summary: {profile.summary[:120]}...")
|
||
|
||
# ── Stage 2: Requirement Understanding ──
|
||
print("\n[2/5] Understanding Requirements...")
|
||
requirement = understand_requirement(
|
||
user_input=user_requirement or "对数据进行全面分析",
|
||
data_profile=profile,
|
||
template_path=template_file
|
||
)
|
||
print(f" Objectives: {len(requirement.objectives)}")
|
||
for obj in requirement.objectives:
|
||
print(f" - {obj.name} (priority: {obj.priority})")
|
||
|
||
# ── Stage 3: AI Analysis Planning ──
|
||
print("\n[3/5] AI Planning Analysis...")
|
||
tool_manager = ToolManager(_global_registry)
|
||
tools = tool_manager.select_tools(profile)
|
||
print(f" Available tools: {len(tools)}")
|
||
|
||
analysis_plan = plan_analysis(profile, requirement, available_tools=tools)
|
||
print(f" Tasks planned: {len(analysis_plan.tasks)}")
|
||
for task in sorted(analysis_plan.tasks, key=lambda t: t.priority, reverse=True):
|
||
print(f" [{task.priority}] {task.name}")
|
||
if task.required_tools:
|
||
print(f" tools: {', '.join(task.required_tools)}")
|
||
|
||
# ── Stage 4: AI Task Execution ──
|
||
print("\n[4/5] AI Executing Tasks...")
|
||
# Reuse DAL from Stage 1 — no need to load data again
|
||
results: List[AnalysisResult] = []
|
||
|
||
sorted_tasks = sorted(analysis_plan.tasks, key=lambda t: t.priority, reverse=True)
|
||
for i, task in enumerate(sorted_tasks, 1):
|
||
print(f"\n Task {i}/{len(sorted_tasks)}: {task.name}")
|
||
result = execute_task(task, tools, dal)
|
||
results.append(result)
|
||
|
||
if result.success:
|
||
print(f" ✓ Done ({result.execution_time:.1f}s), insights: {len(result.insights)}")
|
||
for insight in result.insights[:2]:
|
||
print(f" - {insight[:80]}")
|
||
else:
|
||
print(f" ✗ Failed: {result.error}")
|
||
|
||
successful = sum(1 for r in results if r.success)
|
||
print(f"\n Results: {successful}/{len(results)} tasks succeeded")
|
||
|
||
# ── Stage 5: Report Generation ──
|
||
print("\n[5/5] Generating Report...")
|
||
report_path = os.path.join(output_dir, "analysis_report.md")
|
||
|
||
if template_file and os.path.exists(template_file):
|
||
report = _generate_template_report(profile, results, template_file, config)
|
||
else:
|
||
report = generate_report(results, requirement, profile)
|
||
|
||
# Save report
|
||
with open(report_path, 'w', encoding='utf-8') as f:
|
||
f.write(report)
|
||
|
||
print(f" Report saved: {report_path}")
|
||
print(f" Report length: {len(report)} chars")
|
||
|
||
# ── Done ──
|
||
print("\n" + "=" * 70)
|
||
print("Analysis Complete!")
|
||
print(f"End: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
|
||
print(f"Output: {report_path}")
|
||
print("=" * 70)
|
||
|
||
return True
|
||
|
||
|
||
def _generate_template_report(
|
||
profile: DataProfile,
|
||
results: List[AnalysisResult],
|
||
template_path: str,
|
||
config
|
||
) -> str:
|
||
"""Use AI to fill a template with data from task execution results."""
|
||
client = OpenAI(api_key=config.llm.api_key, base_url=config.llm.base_url)
|
||
|
||
with open(template_path, 'r', encoding='utf-8') as f:
|
||
template = f.read()
|
||
|
||
# Collect all data from task results
|
||
all_data = {}
|
||
all_insights = []
|
||
for r in results:
|
||
if r.success:
|
||
all_data[r.task_name] = {
|
||
'data': json.dumps(r.data, ensure_ascii=False, default=str)[:1000],
|
||
'insights': r.insights
|
||
}
|
||
all_insights.extend(r.insights)
|
||
|
||
data_json = json.dumps(all_data, ensure_ascii=False, indent=1)
|
||
if len(data_json) > 12000:
|
||
data_json = data_json[:12000] + "\n... (truncated)"
|
||
|
||
prompt = f"""你是一位专业的数据分析师。请根据以下分析结果，按照模板格式生成完整的报告。
|
||
|
||
## 数据概况
|
||
- 类型: {profile.inferred_type}
|
||
- 行数: {profile.row_count}, 列数: {profile.column_count}
|
||
- 质量: {profile.quality_score}/100
|
||
- 摘要: {profile.summary[:500]}
|
||
|
||
## 关键字段
|
||
{json.dumps(profile.key_fields, ensure_ascii=False, indent=2)}
|
||
|
||
## 分析结果
|
||
{data_json}
|
||
|
||
## 关键洞察
|
||
{chr(10).join(f"- {i}" for i in all_insights[:20])}
|
||
|
||
## 报告模板
|
||
```markdown
|
||
{template}
|
||
```
|
||
|
||
## 要求
|
||
1. 用实际数据填充模板中所有占位符
|
||
2. 根据数据中的字段，智能映射到模板分类
|
||
3. 所有数字必须来自分析结果，不要编造
|
||
4. 如果某个模板分类在数据中没有对应，标注"本期无数据"
|
||
5. 保持Markdown格式
|
||
"""
|
||
|
||
print(" AI filling template with analysis results...")
|
||
response = client.chat.completions.create(
|
||
model=config.llm.model,
|
||
messages=[
|
||
{"role": "system", "content": "你是数据分析专家。根据分析结果填充报告模板，所有数字必须来自真实数据。输出纯Markdown。"},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
temperature=0.3,
|
||
max_tokens=4000
|
||
)
|
||
|
||
report = response.choices[0].message.content
|
||
|
||
header = f"""<!--
|
||
Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
|
||
Data: {profile.file_path} ({profile.row_count} rows x {profile.column_count} cols)
|
||
Quality: {profile.quality_score}/100
|
||
Template: {template_path}
|
||
AI never accessed raw data rows — only aggregated tool results
|
||
-->
|
||
|
||
"""
|
||
return header + report
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
parser = argparse.ArgumentParser(description="AI-Driven Data Analysis")
|
||
parser.add_argument("--data", default="cleaned_data.csv", help="CSV file path")
|
||
parser.add_argument("--requirement", default=None, help="Analysis requirement (natural language)")
|
||
parser.add_argument("--template", default=None, help="Report template path")
|
||
parser.add_argument("--output", default="analysis_output", help="Output directory")
|
||
args = parser.parse_args()
|
||
|
||
success = run_analysis(
|
||
data_file=args.data,
|
||
user_requirement=args.requirement,
|
||
template_file=args.template,
|
||
output_dir=args.output
|
||
)
|
||
sys.exit(0 if success else 1)
|