575 lines
22 KiB
Python
575 lines
22 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
配置管理器
|
||
管理API密钥、模型配置、费率设置等
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
from datetime import datetime
|
||
from typing import Dict, List, Optional, Any
|
||
from dataclasses import dataclass, asdict
|
||
from pathlib import Path
|
||
from dotenv import load_dotenv
|
||
|
||
try:
|
||
from .mongodb_storage import MongoDBStorage
|
||
MONGODB_AVAILABLE = True
|
||
except ImportError:
|
||
MONGODB_AVAILABLE = False
|
||
MongoDBStorage = None
|
||
|
||
|
||
@dataclass
|
||
class ModelConfig:
|
||
"""模型配置"""
|
||
provider: str # 供应商：dashscope, openai, google, etc.
|
||
model_name: str # 模型名称
|
||
api_key: str # API密钥
|
||
base_url: Optional[str] = None # 自定义API地址
|
||
max_tokens: int = 4000 # 最大token数
|
||
temperature: float = 0.7 # 温度参数
|
||
enabled: bool = True # 是否启用
|
||
|
||
|
||
@dataclass
|
||
class PricingConfig:
|
||
"""定价配置"""
|
||
provider: str # 供应商
|
||
model_name: str # 模型名称
|
||
input_price_per_1k: float # 输入token价格（每1000个token）
|
||
output_price_per_1k: float # 输出token价格（每1000个token）
|
||
currency: str = "CNY" # 货币单位
|
||
|
||
|
||
@dataclass
|
||
class UsageRecord:
|
||
"""使用记录"""
|
||
timestamp: str # 时间戳
|
||
provider: str # 供应商
|
||
model_name: str # 模型名称
|
||
input_tokens: int # 输入token数
|
||
output_tokens: int # 输出token数
|
||
cost: float # 成本
|
||
session_id: str # 会话ID
|
||
analysis_type: str # 分析类型
|
||
|
||
|
||
class ConfigManager:
|
||
"""配置管理器"""
|
||
|
||
def __init__(self, config_dir: str = "config"):
|
||
self.config_dir = Path(config_dir)
|
||
self.config_dir.mkdir(exist_ok=True)
|
||
|
||
self.models_file = self.config_dir / "models.json"
|
||
self.pricing_file = self.config_dir / "pricing.json"
|
||
self.usage_file = self.config_dir / "usage.json"
|
||
self.settings_file = self.config_dir / "settings.json"
|
||
|
||
# 加载.env文件（保持向后兼容）
|
||
self._load_env_file()
|
||
|
||
# 初始化MongoDB存储（如果可用）
|
||
self.mongodb_storage = None
|
||
self._init_mongodb_storage()
|
||
|
||
self._init_default_configs()
|
||
|
||
def _load_env_file(self):
|
||
"""加载.env文件（保持向后兼容）"""
|
||
# 尝试从项目根目录加载.env文件
|
||
project_root = Path(__file__).parent.parent.parent
|
||
env_file = project_root / ".env"
|
||
|
||
if env_file.exists():
|
||
load_dotenv(env_file, override=True)
|
||
|
||
def _get_env_api_key(self, provider: str) -> str:
|
||
"""从环境变量获取API密钥"""
|
||
env_key_map = {
|
||
"dashscope": "DASHSCOPE_API_KEY",
|
||
"openai": "OPENAI_API_KEY",
|
||
"google": "GOOGLE_API_KEY",
|
||
"anthropic": "ANTHROPIC_API_KEY",
|
||
"deepseek": "DEEPSEEK_API_KEY"
|
||
}
|
||
|
||
env_key = env_key_map.get(provider.lower())
|
||
if env_key:
|
||
return os.getenv(env_key, "")
|
||
return ""
|
||
|
||
def _init_mongodb_storage(self):
|
||
"""初始化MongoDB存储"""
|
||
if not MONGODB_AVAILABLE:
|
||
return
|
||
|
||
# 检查是否启用MongoDB存储
|
||
use_mongodb = os.getenv("USE_MONGODB_STORAGE", "false").lower() == "true"
|
||
if not use_mongodb:
|
||
return
|
||
|
||
try:
|
||
connection_string = os.getenv("MONGODB_CONNECTION_STRING")
|
||
database_name = os.getenv("MONGODB_DATABASE_NAME", "tradingagents")
|
||
|
||
self.mongodb_storage = MongoDBStorage(
|
||
connection_string=connection_string,
|
||
database_name=database_name
|
||
)
|
||
|
||
if self.mongodb_storage.is_connected():
|
||
print("✅ MongoDB存储已启用")
|
||
else:
|
||
self.mongodb_storage = None
|
||
print("⚠️ MongoDB连接失败，将使用JSON文件存储")
|
||
|
||
except Exception as e:
|
||
print(f"❌ MongoDB初始化失败: {e}")
|
||
self.mongodb_storage = None
|
||
|
||
def _init_default_configs(self):
|
||
"""初始化默认配置"""
|
||
# 默认模型配置
|
||
if not self.models_file.exists():
|
||
default_models = [
|
||
ModelConfig(
|
||
provider="dashscope",
|
||
model_name="qwen-turbo",
|
||
api_key="",
|
||
max_tokens=4000,
|
||
temperature=0.7
|
||
),
|
||
ModelConfig(
|
||
provider="dashscope",
|
||
model_name="qwen-plus-latest",
|
||
api_key="",
|
||
max_tokens=8000,
|
||
temperature=0.7
|
||
),
|
||
ModelConfig(
|
||
provider="openai",
|
||
model_name="gpt-3.5-turbo",
|
||
api_key="",
|
||
max_tokens=4000,
|
||
temperature=0.7,
|
||
enabled=False
|
||
),
|
||
ModelConfig(
|
||
provider="openai",
|
||
model_name="gpt-4",
|
||
api_key="",
|
||
max_tokens=8000,
|
||
temperature=0.7,
|
||
enabled=False
|
||
),
|
||
ModelConfig(
|
||
provider="google",
|
||
model_name="gemini-pro",
|
||
api_key="",
|
||
max_tokens=4000,
|
||
temperature=0.7,
|
||
enabled=False
|
||
)
|
||
]
|
||
self.save_models(default_models)
|
||
|
||
# 默认定价配置
|
||
if not self.pricing_file.exists():
|
||
default_pricing = [
|
||
# 阿里百炼定价 (人民币)
|
||
PricingConfig("dashscope", "qwen-turbo", 0.002, 0.006, "CNY"),
|
||
PricingConfig("dashscope", "qwen-plus-latest", 0.004, 0.012, "CNY"),
|
||
PricingConfig("dashscope", "qwen-max", 0.02, 0.06, "CNY"),
|
||
|
||
# OpenAI定价 (美元)
|
||
PricingConfig("openai", "gpt-3.5-turbo", 0.0015, 0.002, "USD"),
|
||
PricingConfig("openai", "gpt-4", 0.03, 0.06, "USD"),
|
||
PricingConfig("openai", "gpt-4-turbo", 0.01, 0.03, "USD"),
|
||
|
||
# Google定价 (美元)
|
||
PricingConfig("google", "gemini-pro", 0.00025, 0.0005, "USD"),
|
||
PricingConfig("google", "gemini-pro-vision", 0.00025, 0.0005, "USD"),
|
||
]
|
||
self.save_pricing(default_pricing)
|
||
|
||
# 默认设置
|
||
if not self.settings_file.exists():
|
||
# 导入默认数据目录配置
|
||
import os
|
||
default_data_dir = os.path.join(os.path.expanduser("~"), "Documents", "TradingAgents", "data")
|
||
|
||
default_settings = {
|
||
"default_provider": "dashscope",
|
||
"default_model": "qwen-turbo",
|
||
"enable_cost_tracking": True,
|
||
"cost_alert_threshold": 100.0, # 成本警告阈值
|
||
"currency_preference": "CNY",
|
||
"auto_save_usage": True,
|
||
"max_usage_records": 10000,
|
||
"data_dir": default_data_dir, # 数据目录配置
|
||
"cache_dir": os.path.join(default_data_dir, "cache"), # 缓存目录
|
||
"results_dir": os.path.join(os.path.expanduser("~"), "Documents", "TradingAgents", "results"), # 结果目录
|
||
"auto_create_dirs": True # 自动创建目录
|
||
}
|
||
self.save_settings(default_settings)
|
||
|
||
def load_models(self) -> List[ModelConfig]:
|
||
"""加载模型配置，优先使用.env中的API密钥"""
|
||
try:
|
||
with open(self.models_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
models = [ModelConfig(**item) for item in data]
|
||
|
||
# 合并.env中的API密钥（优先级更高）
|
||
for model in models:
|
||
env_api_key = self._get_env_api_key(model.provider)
|
||
if env_api_key:
|
||
model.api_key = env_api_key
|
||
# 如果.env中有API密钥，自动启用该模型
|
||
if not model.enabled:
|
||
model.enabled = True
|
||
|
||
return models
|
||
except Exception as e:
|
||
print(f"加载模型配置失败: {e}")
|
||
return []
|
||
|
||
def save_models(self, models: List[ModelConfig]):
|
||
"""保存模型配置"""
|
||
try:
|
||
data = [asdict(model) for model in models]
|
||
with open(self.models_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
except Exception as e:
|
||
print(f"保存模型配置失败: {e}")
|
||
|
||
def load_pricing(self) -> List[PricingConfig]:
|
||
"""加载定价配置"""
|
||
try:
|
||
with open(self.pricing_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
return [PricingConfig(**item) for item in data]
|
||
except Exception as e:
|
||
print(f"加载定价配置失败: {e}")
|
||
return []
|
||
|
||
def save_pricing(self, pricing: List[PricingConfig]):
|
||
"""保存定价配置"""
|
||
try:
|
||
data = [asdict(price) for price in pricing]
|
||
with open(self.pricing_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
except Exception as e:
|
||
print(f"保存定价配置失败: {e}")
|
||
|
||
def load_usage_records(self) -> List[UsageRecord]:
|
||
"""加载使用记录"""
|
||
try:
|
||
if not self.usage_file.exists():
|
||
return []
|
||
with open(self.usage_file, 'r', encoding='utf-8') as f:
|
||
data = json.load(f)
|
||
return [UsageRecord(**item) for item in data]
|
||
except Exception as e:
|
||
print(f"加载使用记录失败: {e}")
|
||
return []
|
||
|
||
def save_usage_records(self, records: List[UsageRecord]):
|
||
"""保存使用记录"""
|
||
try:
|
||
data = [asdict(record) for record in records]
|
||
with open(self.usage_file, 'w', encoding='utf-8') as f:
|
||
json.dump(data, f, ensure_ascii=False, indent=2)
|
||
except Exception as e:
|
||
print(f"保存使用记录失败: {e}")
|
||
|
||
def add_usage_record(self, provider: str, model_name: str, input_tokens: int,
|
||
output_tokens: int, session_id: str, analysis_type: str = "stock_analysis"):
|
||
"""添加使用记录"""
|
||
# 计算成本
|
||
cost = self.calculate_cost(provider, model_name, input_tokens, output_tokens)
|
||
|
||
record = UsageRecord(
|
||
timestamp=datetime.now().isoformat(),
|
||
provider=provider,
|
||
model_name=model_name,
|
||
input_tokens=input_tokens,
|
||
output_tokens=output_tokens,
|
||
cost=cost,
|
||
session_id=session_id,
|
||
analysis_type=analysis_type
|
||
)
|
||
|
||
# 优先使用MongoDB存储
|
||
if self.mongodb_storage and self.mongodb_storage.is_connected():
|
||
success = self.mongodb_storage.save_usage_record(record)
|
||
if success:
|
||
return record
|
||
else:
|
||
print("⚠️ MongoDB保存失败，回退到JSON文件存储")
|
||
|
||
# 回退到JSON文件存储
|
||
records = self.load_usage_records()
|
||
records.append(record)
|
||
|
||
# 限制记录数量
|
||
settings = self.load_settings()
|
||
max_records = settings.get("max_usage_records", 10000)
|
||
if len(records) > max_records:
|
||
records = records[-max_records:]
|
||
|
||
self.save_usage_records(records)
|
||
return record
|
||
|
||
def calculate_cost(self, provider: str, model_name: str, input_tokens: int, output_tokens: int) -> float:
|
||
"""计算使用成本"""
|
||
pricing_configs = self.load_pricing()
|
||
|
||
for pricing in pricing_configs:
|
||
if pricing.provider == provider and pricing.model_name == model_name:
|
||
input_cost = (input_tokens / 1000) * pricing.input_price_per_1k
|
||
output_cost = (output_tokens / 1000) * pricing.output_price_per_1k
|
||
return round(input_cost + output_cost, 6)
|
||
|
||
return 0.0
|
||
|
||
def load_settings(self) -> Dict[str, Any]:
|
||
"""加载设置，合并.env中的配置"""
|
||
try:
|
||
with open(self.settings_file, 'r', encoding='utf-8') as f:
|
||
settings = json.load(f)
|
||
except Exception as e:
|
||
print(f"加载设置失败: {e}")
|
||
settings = {}
|
||
|
||
# 合并.env中的其他配置
|
||
env_settings = {
|
||
"finnhub_api_key": os.getenv("FINNHUB_API_KEY", ""),
|
||
"reddit_client_id": os.getenv("REDDIT_CLIENT_ID", ""),
|
||
"reddit_client_secret": os.getenv("REDDIT_CLIENT_SECRET", ""),
|
||
"reddit_user_agent": os.getenv("REDDIT_USER_AGENT", ""),
|
||
"results_dir": os.getenv("TRADINGAGENTS_RESULTS_DIR", ""),
|
||
"log_level": os.getenv("TRADINGAGENTS_LOG_LEVEL", "INFO"),
|
||
"data_dir": os.getenv("TRADINGAGENTS_DATA_DIR", ""), # 数据目录环境变量
|
||
"cache_dir": os.getenv("TRADINGAGENTS_CACHE_DIR", ""), # 缓存目录环境变量
|
||
}
|
||
|
||
# 只有当环境变量存在且不为空时才覆盖
|
||
for key, value in env_settings.items():
|
||
if value:
|
||
settings[key] = value
|
||
|
||
return settings
|
||
|
||
def get_env_config_status(self) -> Dict[str, Any]:
|
||
"""获取.env配置状态"""
|
||
return {
|
||
"env_file_exists": (Path(__file__).parent.parent.parent / ".env").exists(),
|
||
"api_keys": {
|
||
"dashscope": bool(os.getenv("DASHSCOPE_API_KEY")),
|
||
"openai": bool(os.getenv("OPENAI_API_KEY")),
|
||
"google": bool(os.getenv("GOOGLE_API_KEY")),
|
||
"anthropic": bool(os.getenv("ANTHROPIC_API_KEY")),
|
||
"finnhub": bool(os.getenv("FINNHUB_API_KEY")),
|
||
},
|
||
"other_configs": {
|
||
"reddit_configured": bool(os.getenv("REDDIT_CLIENT_ID") and os.getenv("REDDIT_CLIENT_SECRET")),
|
||
"results_dir": os.getenv("TRADINGAGENTS_RESULTS_DIR", "./results"),
|
||
"log_level": os.getenv("TRADINGAGENTS_LOG_LEVEL", "INFO"),
|
||
}
|
||
}
|
||
|
||
def save_settings(self, settings: Dict[str, Any]):
|
||
"""保存设置"""
|
||
try:
|
||
with open(self.settings_file, 'w', encoding='utf-8') as f:
|
||
json.dump(settings, f, ensure_ascii=False, indent=2)
|
||
except Exception as e:
|
||
print(f"保存设置失败: {e}")
|
||
|
||
def get_enabled_models(self) -> List[ModelConfig]:
|
||
"""获取启用的模型"""
|
||
models = self.load_models()
|
||
return [model for model in models if model.enabled and model.api_key]
|
||
|
||
def get_model_by_name(self, provider: str, model_name: str) -> Optional[ModelConfig]:
|
||
"""根据名称获取模型配置"""
|
||
models = self.load_models()
|
||
for model in models:
|
||
if model.provider == provider and model.model_name == model_name:
|
||
return model
|
||
return None
|
||
|
||
def get_usage_statistics(self, days: int = 30) -> Dict[str, Any]:
|
||
"""获取使用统计"""
|
||
# 优先使用MongoDB获取统计
|
||
if self.mongodb_storage and self.mongodb_storage.is_connected():
|
||
try:
|
||
# 从MongoDB获取基础统计
|
||
stats = self.mongodb_storage.get_usage_statistics(days)
|
||
# 获取供应商统计
|
||
provider_stats = self.mongodb_storage.get_provider_statistics(days)
|
||
|
||
if stats:
|
||
stats["provider_stats"] = provider_stats
|
||
stats["records_count"] = stats.get("total_requests", 0)
|
||
return stats
|
||
except Exception as e:
|
||
print(f"⚠️ MongoDB统计获取失败，回退到JSON文件: {e}")
|
||
|
||
# 回退到JSON文件统计
|
||
records = self.load_usage_records()
|
||
|
||
# 过滤最近N天的记录
|
||
from datetime import datetime, timedelta
|
||
cutoff_date = datetime.now() - timedelta(days=days)
|
||
|
||
recent_records = []
|
||
for record in records:
|
||
try:
|
||
record_date = datetime.fromisoformat(record.timestamp)
|
||
if record_date >= cutoff_date:
|
||
recent_records.append(record)
|
||
except:
|
||
continue
|
||
|
||
# 统计数据
|
||
total_cost = sum(record.cost for record in recent_records)
|
||
total_input_tokens = sum(record.input_tokens for record in recent_records)
|
||
total_output_tokens = sum(record.output_tokens for record in recent_records)
|
||
|
||
# 按供应商统计
|
||
provider_stats = {}
|
||
for record in recent_records:
|
||
if record.provider not in provider_stats:
|
||
provider_stats[record.provider] = {
|
||
"cost": 0,
|
||
"input_tokens": 0,
|
||
"output_tokens": 0,
|
||
"requests": 0
|
||
}
|
||
provider_stats[record.provider]["cost"] += record.cost
|
||
provider_stats[record.provider]["input_tokens"] += record.input_tokens
|
||
provider_stats[record.provider]["output_tokens"] += record.output_tokens
|
||
provider_stats[record.provider]["requests"] += 1
|
||
|
||
return {
|
||
"period_days": days,
|
||
"total_cost": round(total_cost, 4),
|
||
"total_input_tokens": total_input_tokens,
|
||
"total_output_tokens": total_output_tokens,
|
||
"total_requests": len(recent_records),
|
||
"provider_stats": provider_stats,
|
||
"records_count": len(recent_records)
|
||
}
|
||
|
||
def get_data_dir(self) -> str:
|
||
"""获取数据目录路径"""
|
||
settings = self.load_settings()
|
||
data_dir = settings.get("data_dir")
|
||
if not data_dir:
|
||
# 如果没有配置，使用默认路径
|
||
data_dir = os.path.join(os.path.expanduser("~"), "Documents", "TradingAgents", "data")
|
||
return data_dir
|
||
|
||
def set_data_dir(self, data_dir: str):
|
||
"""设置数据目录路径"""
|
||
settings = self.load_settings()
|
||
settings["data_dir"] = data_dir
|
||
# 同时更新缓存目录
|
||
settings["cache_dir"] = os.path.join(data_dir, "cache")
|
||
self.save_settings(settings)
|
||
|
||
# 如果启用自动创建目录，则创建目录
|
||
if settings.get("auto_create_dirs", True):
|
||
self.ensure_directories_exist()
|
||
|
||
def ensure_directories_exist(self):
|
||
"""确保必要的目录存在"""
|
||
settings = self.load_settings()
|
||
|
||
directories = [
|
||
settings.get("data_dir"),
|
||
settings.get("cache_dir"),
|
||
settings.get("results_dir"),
|
||
os.path.join(settings.get("data_dir", ""), "finnhub_data"),
|
||
os.path.join(settings.get("data_dir", ""), "finnhub_data", "news_data"),
|
||
os.path.join(settings.get("data_dir", ""), "finnhub_data", "insider_sentiment"),
|
||
os.path.join(settings.get("data_dir", ""), "finnhub_data", "insider_transactions")
|
||
]
|
||
|
||
for directory in directories:
|
||
if directory and not os.path.exists(directory):
|
||
try:
|
||
os.makedirs(directory, exist_ok=True)
|
||
print(f"✅ 创建目录: {directory}")
|
||
except Exception as e:
|
||
print(f"❌ 创建目录失败 {directory}: {e}")
|
||
|
||
|
||
class TokenTracker:
|
||
"""Token使用跟踪器"""
|
||
|
||
def __init__(self, config_manager: ConfigManager):
|
||
self.config_manager = config_manager
|
||
|
||
def track_usage(self, provider: str, model_name: str, input_tokens: int,
|
||
output_tokens: int, session_id: str = None, analysis_type: str = "stock_analysis"):
|
||
"""跟踪Token使用"""
|
||
if session_id is None:
|
||
session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
|
||
|
||
# 检查是否启用成本跟踪
|
||
settings = self.config_manager.load_settings()
|
||
if not settings.get("enable_cost_tracking", True):
|
||
return None
|
||
|
||
# 添加使用记录
|
||
record = self.config_manager.add_usage_record(
|
||
provider=provider,
|
||
model_name=model_name,
|
||
input_tokens=input_tokens,
|
||
output_tokens=output_tokens,
|
||
session_id=session_id,
|
||
analysis_type=analysis_type
|
||
)
|
||
|
||
# 检查成本警告
|
||
self._check_cost_alert(record.cost)
|
||
|
||
return record
|
||
|
||
def _check_cost_alert(self, current_cost: float):
|
||
"""检查成本警告"""
|
||
settings = self.config_manager.load_settings()
|
||
threshold = settings.get("cost_alert_threshold", 100.0)
|
||
|
||
# 获取今日总成本
|
||
today_stats = self.config_manager.get_usage_statistics(1)
|
||
total_today = today_stats["total_cost"]
|
||
|
||
if total_today >= threshold:
|
||
print(f"⚠️ 成本警告: 今日成本已达到 ¥{total_today:.4f}，超过阈值 ¥{threshold}")
|
||
|
||
def get_session_cost(self, session_id: str) -> float:
|
||
"""获取会话成本"""
|
||
records = self.config_manager.load_usage_records()
|
||
session_cost = sum(record.cost for record in records if record.session_id == session_id)
|
||
return session_cost
|
||
|
||
def estimate_cost(self, provider: str, model_name: str, estimated_input_tokens: int,
|
||
estimated_output_tokens: int) -> float:
|
||
"""估算成本"""
|
||
return self.config_manager.calculate_cost(
|
||
provider, model_name, estimated_input_tokens, estimated_output_tokens
|
||
)
|
||
|
||
|
||
|
||
|
||
# 全局配置管理器实例
|
||
config_manager = ConfigManager()
|
||
token_tracker = TokenTracker(config_manager)
|