119 lines
3.5 KiB
Python
119 lines
3.5 KiB
Python
"""
|
||
数据可视化技能
|
||
生成图表配置数据
|
||
"""
|
||
from typing import Dict, Any, List
|
||
from app.skills.base import BaseSkill, SkillParameter
|
||
from app.services.tushare_service import tushare_service
|
||
from app.utils.logger import logger
|
||
|
||
|
||
class VisualizationSkill(BaseSkill):
|
||
"""数据可视化技能"""
|
||
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.name = "visualization"
|
||
self.description = "生成K线图和技术指标图表配置"
|
||
self.parameters = [
|
||
SkillParameter(
|
||
name="stock_code",
|
||
type="string",
|
||
description="股票代码",
|
||
required=True
|
||
),
|
||
SkillParameter(
|
||
name="chart_type",
|
||
type="string",
|
||
description="图表类型：candlestick（K线图）",
|
||
required=False,
|
||
default="candlestick"
|
||
),
|
||
SkillParameter(
|
||
name="period",
|
||
type="integer",
|
||
description="数据周期（天数）",
|
||
required=False,
|
||
default=60
|
||
)
|
||
]
|
||
|
||
async def execute(self, **kwargs) -> Dict[str, Any]:
|
||
"""
|
||
生成图表配置
|
||
|
||
Args:
|
||
stock_code: 股票代码
|
||
chart_type: 图表类型
|
||
period: 数据周期
|
||
|
||
Returns:
|
||
图表配置数据
|
||
"""
|
||
stock_code = kwargs.get("stock_code")
|
||
chart_type = kwargs.get("chart_type", "candlestick")
|
||
period = kwargs.get("period", 60)
|
||
|
||
logger.info(f"生成图表配置: {stock_code}, 类型: {chart_type}")
|
||
|
||
# 获取K线数据
|
||
kline_data = tushare_service.get_kline_data(stock_code)
|
||
|
||
if not kline_data:
|
||
return {
|
||
"error": f"未找到K线数据: {stock_code}"
|
||
}
|
||
|
||
# 限制数据量
|
||
if len(kline_data) > period:
|
||
kline_data = kline_data[-period:]
|
||
|
||
if chart_type == "candlestick":
|
||
return self._generate_candlestick_config(kline_data)
|
||
else:
|
||
return {
|
||
"error": f"不支持的图表类型: {chart_type}"
|
||
}
|
||
|
||
def _generate_candlestick_config(self, kline_data: List[dict]) -> Dict[str, Any]:
|
||
"""
|
||
生成K线图配置（Lightweight Charts格式）
|
||
|
||
Args:
|
||
kline_data: K线数据列表
|
||
|
||
Returns:
|
||
图表配置
|
||
"""
|
||
# 转换为Lightweight Charts格式
|
||
candlestick_data = []
|
||
volume_data = []
|
||
|
||
for item in kline_data:
|
||
# 转换日期格式 YYYYMMDD -> YYYY-MM-DD
|
||
date_str = item['trade_date']
|
||
formatted_date = f"{date_str[:4]}-{date_str[4:6]}-{date_str[6:8]}"
|
||
|
||
# K线数据
|
||
candlestick_data.append({
|
||
"time": formatted_date,
|
||
"open": item['open'],
|
||
"high": item['high'],
|
||
"low": item['low'],
|
||
"close": item['close']
|
||
})
|
||
|
||
# 成交量数据
|
||
volume_data.append({
|
||
"time": formatted_date,
|
||
"value": item['vol'],
|
||
"color": "rgba(0, 150, 136, 0.8)" if item['close'] >= item['open'] else "rgba(255, 82, 82, 0.8)"
|
||
})
|
||
|
||
return {
|
||
"chart_type": "candlestick",
|
||
"candlestick_data": candlestick_data,
|
||
"volume_data": volume_data,
|
||
"stock_code": kline_data[0]['ts_code'] if kline_data else None
|
||
}
|