127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Script to build historical memories for TradingAgents
|
|
|
|
This script:
|
|
1. Fetches historical stock data for specified tickers
|
|
2. Analyzes outcomes to create agent-specific memories
|
|
3. Saves memories to disk for later use
|
|
|
|
Usage:
|
|
python scripts/build_historical_memories.py
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
import pickle
|
|
from datetime import datetime, timedelta
|
|
|
|
from tradingagents.agents.utils.historical_memory_builder import HistoricalMemoryBuilder
|
|
from tradingagents.default_config import DEFAULT_CONFIG
|
|
from tradingagents.utils.logger import get_logger
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
def main():
|
|
logger.info(
|
|
"""
|
|
╔══════════════════════════════════════════════════════════════╗
|
|
║ TradingAgents - Historical Memory Builder ║
|
|
╚══════════════════════════════════════════════════════════════╝
|
|
"""
|
|
)
|
|
|
|
# Configuration
|
|
tickers = [
|
|
"AAPL",
|
|
"GOOGL",
|
|
"MSFT",
|
|
"NVDA",
|
|
"TSLA", # Tech
|
|
"JPM",
|
|
"BAC",
|
|
"GS", # Finance
|
|
"XOM",
|
|
"CVX", # Energy
|
|
"JNJ",
|
|
"PFE", # Healthcare
|
|
"WMT",
|
|
"AMZN", # Retail
|
|
]
|
|
|
|
# Date range - last 2 years
|
|
end_date = datetime.now()
|
|
start_date = end_date - timedelta(days=730) # 2 years
|
|
|
|
logger.info(f"Tickers: {', '.join(tickers)}")
|
|
logger.info(f"Period: {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')}")
|
|
logger.info("Lookforward: 7 days (1 week returns)")
|
|
logger.info("Sample interval: 30 days (monthly)\n")
|
|
|
|
proceed = input("Proceed with memory building? (y/n): ")
|
|
if proceed.lower() != "y":
|
|
logger.info("Aborted.")
|
|
return
|
|
|
|
# Build memories
|
|
builder = HistoricalMemoryBuilder(DEFAULT_CONFIG)
|
|
|
|
memories = builder.populate_agent_memories(
|
|
tickers=tickers,
|
|
start_date=start_date.strftime("%Y-%m-%d"),
|
|
end_date=end_date.strftime("%Y-%m-%d"),
|
|
lookforward_days=7,
|
|
interval_days=30,
|
|
)
|
|
|
|
# Save to disk
|
|
memory_dir = os.path.join(DEFAULT_CONFIG["data_dir"], "memories")
|
|
os.makedirs(memory_dir, exist_ok=True)
|
|
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
|
|
for agent_type, memory in memories.items():
|
|
filename = os.path.join(memory_dir, f"{agent_type}_memory_{timestamp}.pkl")
|
|
|
|
# Save the ChromaDB collection data
|
|
# Note: ChromaDB doesn't serialize well, so we extract the data
|
|
collection = memory.situation_collection
|
|
|
|
# Get all items from collection
|
|
results = collection.get(include=["documents", "metadatas", "embeddings"])
|
|
|
|
with open(filename, "wb") as f:
|
|
pickle.dump(
|
|
{
|
|
"documents": results["documents"],
|
|
"metadatas": results["metadatas"],
|
|
"embeddings": results["embeddings"],
|
|
"ids": results["ids"],
|
|
"created_at": timestamp,
|
|
"tickers": tickers,
|
|
"config": {
|
|
"start_date": start_date.strftime("%Y-%m-%d"),
|
|
"end_date": end_date.strftime("%Y-%m-%d"),
|
|
"lookforward_days": 7,
|
|
"interval_days": 30,
|
|
},
|
|
},
|
|
f,
|
|
)
|
|
|
|
logger.info(f"✅ Saved {agent_type} memory to {filename}")
|
|
|
|
logger.info("\n🎉 Memory building complete!")
|
|
logger.info(f" Memories saved to: {memory_dir}")
|
|
logger.info("\n📝 To use these memories, update DEFAULT_CONFIG with:")
|
|
logger.info(f' "memory_dir": "{memory_dir}"')
|
|
logger.info(' "load_historical_memories": True')
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|