405 lines
14 KiB
Python
405 lines
14 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import os
|
||
import json
|
||
import logging
|
||
from typing import Dict, Any, List, Optional, Union
|
||
from datetime import datetime
|
||
|
||
from sqlalchemy import create_engine, Column, Integer, String, Text, DateTime, Index, text
|
||
from sqlalchemy.ext.declarative import declarative_base
|
||
from sqlalchemy.orm import sessionmaker
|
||
from sqlalchemy.dialects.mysql import JSON
|
||
from sqlalchemy.pool import QueuePool
|
||
|
||
from cryptoai.utils.config_loader import ConfigLoader
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger('db_manager')
|
||
|
||
# 创建模型基类
|
||
Base = declarative_base()
|
||
|
||
# 定义分析结果模型
|
||
class AnalysisResult(Base):
|
||
"""分析结果表模型"""
|
||
__tablename__ = 'analysis_results'
|
||
|
||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
agent = Column(String(50), nullable=False, comment='智能体类型(crypto, gold)')
|
||
symbol = Column(String(50), nullable=False, comment='交易对符号')
|
||
time_interval = Column(String(20), nullable=False, comment='时间间隔')
|
||
completion_result = Column(JSON, nullable=False, comment='分析结果JSON')
|
||
created_at = Column(DateTime, nullable=False, default=datetime.now, comment='创建时间')
|
||
updated_at = Column(DateTime, nullable=False, default=datetime.now, onupdate=datetime.now, comment='更新时间')
|
||
|
||
# 索引
|
||
__table_args__ = (
|
||
Index('idx_agent', 'agent'),
|
||
Index('idx_symbol', 'symbol'),
|
||
Index('idx_time_interval', 'time_interval'),
|
||
Index('idx_created_at', 'created_at'),
|
||
)
|
||
|
||
# 定义AI Agent信息流模型
|
||
class AgentFeed(Base):
|
||
"""AI Agent信息流表模型"""
|
||
__tablename__ = 'agent_feeds'
|
||
|
||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||
agent_name = Column(String(50), nullable=False, comment='AI Agent名称')
|
||
avatar_url = Column(String(255), nullable=True, comment='头像URL')
|
||
content = Column(Text, nullable=False, comment='内容')
|
||
create_time = Column(DateTime, nullable=False, default=datetime.now, comment='创建时间')
|
||
|
||
# 索引和表属性
|
||
__table_args__ = (
|
||
Index('idx_agent_name', 'agent_name'),
|
||
Index('idx_create_time', 'create_time'),
|
||
{'mysql_charset': 'utf8mb4', 'mysql_collate': 'utf8mb4_unicode_ci'}
|
||
)
|
||
|
||
class DBManager:
|
||
"""数据库管理工具，用于连接MySQL数据库并保存智能体分析结果"""
|
||
|
||
def __init__(self, host: str, port: int, user: str, password: str, db_name: str):
|
||
"""
|
||
初始化数据库管理器
|
||
|
||
Args:
|
||
host: 数据库主机地址
|
||
port: 数据库端口
|
||
user: 用户名
|
||
password: 密码
|
||
db_name: 数据库名
|
||
"""
|
||
self.host = host
|
||
self.port = port
|
||
self.user = user
|
||
self.password = password
|
||
self.db_name = db_name
|
||
self.engine = None
|
||
self.Session = None
|
||
|
||
# 初始化数据库连接
|
||
self._init_db()
|
||
|
||
def _init_db(self) -> None:
|
||
"""初始化数据库连接和表"""
|
||
try:
|
||
# 创建数据库连接
|
||
connection_string = f"mysql+pymysql://{self.user}:{self.password}@{self.host}:{self.port}/{self.db_name}?charset=utf8mb4"
|
||
|
||
# 创建引擎，设置连接池
|
||
self.engine = create_engine(
|
||
connection_string,
|
||
echo=False, # 设置为True可以输出SQL语句（调试用）
|
||
pool_size=5, # 连接池大小
|
||
max_overflow=10, # 最大溢出连接数
|
||
pool_timeout=30, # 连接超时时间
|
||
pool_recycle=1800, # 连接回收时间（秒）
|
||
pool_pre_ping=True, # 在使用连接前先ping一下，确保连接有效
|
||
connect_args={'charset': 'utf8mb4'}
|
||
)
|
||
|
||
# 创建会话工厂
|
||
self.Session = sessionmaker(bind=self.engine)
|
||
|
||
# 创建表（如果不存在）
|
||
Base.metadata.create_all(self.engine)
|
||
|
||
logger.info(f"成功连接到数据库 {self.db_name}")
|
||
|
||
except Exception as e:
|
||
logger.error(f"数据库初始化失败: {e}")
|
||
self.engine = None
|
||
|
||
def save_analysis_result(self, agent: str, symbol: str, time_interval: str,
|
||
analysis_result: Dict[str, Any]) -> bool:
|
||
"""
|
||
保存分析结果到数据库
|
||
|
||
Args:
|
||
agent: 智能体类型，例如 'crypto' 或 'gold'
|
||
symbol: 交易对符号，例如 'BTCUSDT'
|
||
time_interval: 时间间隔，例如 '1h', '4h', '1d'
|
||
analysis_result: 分析结果数据
|
||
|
||
Returns:
|
||
保存是否成功
|
||
"""
|
||
if not self.engine:
|
||
try:
|
||
self._init_db()
|
||
except Exception as e:
|
||
logger.error(f"重新连接数据库失败: {e}")
|
||
return False
|
||
|
||
try:
|
||
# 创建会话
|
||
session = self.Session()
|
||
|
||
try:
|
||
# 创建新记录
|
||
new_result = AnalysisResult(
|
||
agent=agent,
|
||
symbol=symbol,
|
||
time_interval=time_interval,
|
||
completion_result=analysis_result,
|
||
created_at=datetime.now(),
|
||
updated_at=datetime.now()
|
||
)
|
||
|
||
# 添加并提交
|
||
session.add(new_result)
|
||
session.commit()
|
||
|
||
logger.info(f"成功保存 {agent} 分析结果，交易对: {symbol}, 时间间隔: {time_interval}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
session.rollback()
|
||
logger.error(f"保存分析结果失败: {e}")
|
||
return False
|
||
|
||
finally:
|
||
session.close()
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建数据库会话失败: {e}")
|
||
# 如果是连接错误，尝试重新初始化
|
||
try:
|
||
self._init_db()
|
||
except:
|
||
pass
|
||
return False
|
||
|
||
def save_agent_feed(self, agent_name: str, content: str, avatar_url: Optional[str] = None) -> bool:
|
||
"""
|
||
保存AI Agent信息流到数据库
|
||
|
||
Args:
|
||
agent_name: AI Agent名称
|
||
content: 内容
|
||
avatar_url: 头像URL，可选
|
||
|
||
Returns:
|
||
保存是否成功
|
||
"""
|
||
if not self.engine:
|
||
try:
|
||
self._init_db()
|
||
except Exception as e:
|
||
logger.error(f"重新连接数据库失败: {e}")
|
||
return False
|
||
|
||
try:
|
||
# 创建会话
|
||
session = self.Session()
|
||
|
||
try:
|
||
# 创建新记录
|
||
new_feed = AgentFeed(
|
||
agent_name=agent_name,
|
||
content=content,
|
||
avatar_url=avatar_url
|
||
)
|
||
|
||
# 添加并提交
|
||
session.add(new_feed)
|
||
session.commit()
|
||
|
||
logger.info(f"成功保存 {agent_name} 的信息流")
|
||
return True
|
||
|
||
except Exception as e:
|
||
session.rollback()
|
||
logger.error(f"保存信息流失败: {e}")
|
||
return False
|
||
|
||
finally:
|
||
session.close()
|
||
|
||
except Exception as e:
|
||
logger.error(f"创建数据库会话失败: {e}")
|
||
# 如果是连接错误，尝试重新初始化
|
||
try:
|
||
self._init_db()
|
||
except:
|
||
pass
|
||
return False
|
||
|
||
def get_agent_feeds(self, agent_name: Optional[str] = None, limit: int = 20, skip: int = 0) -> List[Dict[str, Any]]:
|
||
"""
|
||
获取AI Agent信息流
|
||
|
||
Args:
|
||
agent_name: 可选，指定获取特定Agent的信息流
|
||
limit: 返回的最大记录数，默认20条
|
||
|
||
Returns:
|
||
信息流列表，如果查询失败则返回空列表
|
||
"""
|
||
if not self.engine:
|
||
try:
|
||
self._init_db()
|
||
except Exception as e:
|
||
logger.error(f"重新连接数据库失败: {e}")
|
||
return []
|
||
|
||
try:
|
||
# 创建会话
|
||
session = self.Session()
|
||
|
||
try:
|
||
# 构建查询
|
||
query = session.query(AgentFeed)
|
||
|
||
# 如果指定了agent_name，则筛选
|
||
if agent_name:
|
||
query = query.filter(AgentFeed.agent_name == agent_name)
|
||
|
||
# 按创建时间降序排序并限制数量
|
||
results = query.order_by(AgentFeed.create_time.desc()).offset(skip).limit(limit).all()
|
||
|
||
# 转换为字典列表
|
||
feeds = []
|
||
for result in results:
|
||
feeds.append({
|
||
'id': result.id,
|
||
'agent_name': result.agent_name,
|
||
'avatar_url': result.avatar_url,
|
||
'content': result.content,
|
||
'create_time': result.create_time
|
||
})
|
||
|
||
return feeds
|
||
|
||
finally:
|
||
session.close()
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取信息流失败: {e}")
|
||
return []
|
||
|
||
def get_latest_result(self, agent: str, symbol: str, time_interval: str) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
获取最新的分析结果
|
||
|
||
Args:
|
||
agent: 智能体类型，例如 'crypto' 或 'gold'
|
||
symbol: 交易对符号，例如 'BTCUSDT'
|
||
time_interval: 时间间隔，例如 '1h', '4h', '1d'
|
||
|
||
Returns:
|
||
最新分析结果，如果查询失败则返回None
|
||
"""
|
||
if not self.engine:
|
||
try:
|
||
self._init_db()
|
||
except Exception as e:
|
||
logger.error(f"重新连接数据库失败: {e}")
|
||
return None
|
||
|
||
try:
|
||
# 创建会话
|
||
session = self.Session()
|
||
|
||
try:
|
||
# 查询最新的结果
|
||
result = session.query(AnalysisResult).filter(
|
||
AnalysisResult.agent == agent,
|
||
AnalysisResult.symbol == symbol,
|
||
AnalysisResult.time_interval == time_interval
|
||
).order_by(AnalysisResult.created_at.desc()).first()
|
||
|
||
if result:
|
||
# 转换为字典
|
||
return {
|
||
'id': result.id,
|
||
'agent': result.agent,
|
||
'symbol': result.symbol,
|
||
'time_interval': result.time_interval,
|
||
'completion_result': result.completion_result,
|
||
'created_at': result.created_at
|
||
}
|
||
else:
|
||
return None
|
||
|
||
finally:
|
||
session.close()
|
||
|
||
except Exception as e:
|
||
logger.error(f"获取最新分析结果失败: {e}")
|
||
return None
|
||
|
||
def close(self) -> None:
|
||
"""关闭数据库连接"""
|
||
if self.engine:
|
||
self.engine.dispose()
|
||
self.engine = None
|
||
logger.info("数据库连接已关闭")
|
||
|
||
|
||
# 单例模式
|
||
_db_instance = None
|
||
|
||
def get_db_manager(host: Optional[str] = None,
|
||
port: Optional[int] = None,
|
||
user: Optional[str] = None,
|
||
password: Optional[str] = None,
|
||
db_name: Optional[str] = None) -> DBManager:
|
||
"""
|
||
获取数据库管理器实例（单例模式）
|
||
|
||
Args:
|
||
host: 数据库主机地址
|
||
port: 数据库端口
|
||
user: 用户名
|
||
password: 密码
|
||
db_name: 数据库名
|
||
|
||
Returns:
|
||
数据库管理器实例
|
||
"""
|
||
global _db_instance
|
||
|
||
# 如果已经初始化过，直接返回
|
||
if _db_instance is not None:
|
||
return _db_instance
|
||
|
||
# 如果未指定参数，从配置加载器获取数据库配置
|
||
if host is None or port is None or user is None or password is None or db_name is None:
|
||
config_loader = ConfigLoader()
|
||
db_config = config_loader.get_database_config()
|
||
|
||
# 使用配置中的值或默认值
|
||
db_host = host or db_config.get('host')
|
||
db_port = port or db_config.get('port')
|
||
db_user = user or db_config.get('user')
|
||
db_password = password or db_config.get('password')
|
||
db_name = db_name or db_config.get('db_name')
|
||
|
||
logger.info(f"从配置加载数据库连接信息: {db_host}:{db_port}/{db_name}")
|
||
else:
|
||
db_host = host
|
||
db_port = port
|
||
db_user = user
|
||
db_password = password
|
||
db_name = db_name
|
||
|
||
# 创建实例
|
||
_db_instance = DBManager(
|
||
host=db_host,
|
||
port=db_port,
|
||
user=db_user,
|
||
password=db_password,
|
||
db_name=db_name
|
||
)
|
||
|
||
return _db_instance |