110 lines
3.4 KiB
Python
110 lines
3.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Test script to verify Gemini integration with TradingAgents
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
# Add the project root to the Python path
|
|
project_root = Path(__file__).parent.parent
|
|
sys.path.insert(0, str(project_root))
|
|
|
|
from tradingagents.graph.trading_graph import TradingAgentsGraph
|
|
from tradingagents.default_config import DEFAULT_CONFIG
|
|
|
|
def test_gemini_integration():
|
|
"""Test the Gemini integration"""
|
|
|
|
# Check if Google API key is available
|
|
if not os.getenv("GOOGLE_API_KEY"):
|
|
print("❌ GOOGLE_API_KEY environment variable not set")
|
|
print("Please set it with: export GOOGLE_API_KEY=your_key_here")
|
|
return False
|
|
|
|
# Create a config for Gemini
|
|
config = DEFAULT_CONFIG.copy()
|
|
config["llm_provider"] = "google"
|
|
config["backend_url"] = "https://generativelanguage.googleapis.com/v1"
|
|
config["deep_think_llm"] = "gemini-2.0-flash"
|
|
config["quick_think_llm"] = "gemini-2.0-flash"
|
|
config["max_debate_rounds"] = 1
|
|
config["online_tools"] = True
|
|
|
|
try:
|
|
print("🚀 Testing Gemini integration...")
|
|
|
|
# Initialize the graph with Gemini
|
|
graph = TradingAgentsGraph(
|
|
selected_analysts=["market", "news"], # Use fewer analysts for testing
|
|
debug=True,
|
|
config=config
|
|
)
|
|
|
|
print("✅ TradingAgentsGraph initialized successfully with Gemini")
|
|
|
|
# Test a simple propagation
|
|
print("🔄 Testing propagation with SPY...")
|
|
_, decision = graph.propagate("SPY", "2024-12-01")
|
|
|
|
print(f"✅ Propagation completed successfully!")
|
|
print(f"📊 Final decision: {decision}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error during Gemini integration test: {str(e)}")
|
|
return False
|
|
|
|
def test_memory_system():
|
|
"""Test the memory system with Gemini embeddings"""
|
|
|
|
if not os.getenv("GOOGLE_API_KEY"):
|
|
print("❌ GOOGLE_API_KEY environment variable not set")
|
|
return False
|
|
|
|
try:
|
|
from tradingagents.agents.utils.memory import FinancialSituationMemory
|
|
|
|
config = {
|
|
"llm_provider": "google",
|
|
"backend_url": "https://generativelanguage.googleapis.com/v1"
|
|
}
|
|
|
|
print("🧠 Testing memory system with Gemini embeddings...")
|
|
|
|
memory = FinancialSituationMemory("test_memory", config)
|
|
|
|
# Test adding situations
|
|
test_data = [
|
|
("High inflation with rising interest rates", "Consider defensive sectors"),
|
|
("Tech sector volatility increasing", "Reduce exposure to growth stocks")
|
|
]
|
|
|
|
memory.add_situations(test_data)
|
|
print("✅ Memory system working with Gemini embeddings")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
print(f"❌ Error in memory system test: {str(e)}")
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
print("🧪 Testing TradingAgents Gemini Integration")
|
|
print("=" * 50)
|
|
|
|
# Test memory system first
|
|
memory_success = test_memory_system()
|
|
|
|
if memory_success:
|
|
# Test full integration
|
|
integration_success = test_gemini_integration()
|
|
|
|
if integration_success:
|
|
print("\n🎉 All Gemini integration tests passed!")
|
|
else:
|
|
print("\n❌ Gemini integration test failed")
|
|
else:
|
|
print("\n❌ Memory system test failed") |