charieswei
/
TradingAgents
mirror of https://github.com/TauricResearch/TradingAgents.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1 from 89jobrien/001-trending-stock-discovery 

Add trending stock discovery feature
This commit is contained in:
Joesph O'Brien 2025-12-02 20:44:33 -05:00 committed by GitHub
parent 13b826a31d e3e40d1e05
commit 751ad5d286
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
47 changed files with 8402 additions and 842 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
.gitignore vendored
View File

@ -9,3 +9,13 @@ eval_results/
eval_data/
*.egg-info/
.env
.claude/
.pytest_cache/
.specify/
specs/
agent-os/
*.local.md
build/
.mcp.json
*.zip
todos.md

129
README.md
View File

@ -12,22 +12,21 @@
</div>
<div align="center">
<!-- Keep these links. Translations will automatically update with the README. -->
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=de">Deutsch</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=es">Español</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=fr">français</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=ja">日本語</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=ko">한국어</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=pt">Português</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=ru">Русский</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=de">Deutsch</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=es">Español</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=fr">français</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=ja">日本語</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=ko">한국어</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=pt">Português</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=ru">Русский</a> |
<a href="https://www.readme-i18n.com/TauricResearch/TradingAgents?lang=zh">中文</a>
</div>
---
# TradingAgents: Multi-Agents LLM Financial Trading Framework
# TradingAgents: Multi-Agents LLM Financial Trading Framework
> 🎉 **TradingAgents** officially released! We have received numerous inquiries about the work, and we would like to express our thanks for the enthusiasm in our community.
> **TradingAgents** officially released! We have received numerous inquiries about the work, and we would like to express our thanks for the enthusiasm in our community.
>
> So we decided to fully open-source the framework. Looking forward to building impactful projects with you!
@ -43,7 +42,7 @@
<div align="center">
🚀 [TradingAgents](#tradingagents-framework) | ⚡ [Installation & CLI](#installation-and-cli) | 🎬 [Demo](https://www.youtube.com/watch?v=90gr5lwjIho) | 📦 [Package Usage](#tradingagents-package) | 🤝 [Contributing](#contributing) | 📄 [Citation](#citation)
[TradingAgents](#tradingagents-framework) | [Installation & CLI](#installation-and-cli) | [Demo](https://www.youtube.com/watch?v=90gr5lwjIho) | [Package Usage](#tradingagents-package) | [Source](#source)
</div>
@ -101,15 +100,10 @@ git clone https://github.com/TauricResearch/TradingAgents.git
cd TradingAgents
```
Create a virtual environment in any of your favorite environment managers:
Sync virtual environment:
```bash
conda create -n tradingagents python=3.13
conda activate tradingagents
```
Install dependencies:
```bash
pip install -r requirements.txt
uv sync
source .venv/bin/activate
```
### Required APIs
@ -124,18 +118,38 @@ export ALPHA_VANTAGE_API_KEY=$YOUR_ALPHA_VANTAGE_API_KEY
Alternatively, you can create a `.env` file in the project root with your API keys (see `.env.example` for reference):
```bash
cp .env.example .env
# Edit .env with your actual API keys
```
**Note:** We are happy to partner with Alpha Vantage to provide robust API support for TradingAgents. You can get a free AlphaVantage API [here](https://www.alphavantage.co/support/#api-key), TradingAgents-sourced requests also have increased rate limits to 60 requests per minute with no daily limits. Typically the quota is sufficient for performing complex tasks with TradingAgents thanks to Alpha Vantage’s open-source support program. If you prefer to use OpenAI for these data sources instead, you can modify the data vendor settings in `tradingagents/default_config.py`.
**Note:** We are happy to partner with Alpha Vantage to provide robust API support for TradingAgents. You can get a free AlphaVantage API [here](https://www.alphavantage.co/support/#api-key), TradingAgents-sourced requests also have increased rate limits to 60 requests per minute with no daily limits. Typically the quota is sufficient for performing complex tasks with TradingAgents thanks to Alpha Vantage's open-source support program. If you prefer to use OpenAI for these data sources instead, you can modify the data vendor settings in `tradingagents/default_config.py`.
### CLI Usage
You can also try out the CLI directly by running:
Run the CLI:
```bash
python -m cli.main
uv run cli/main.py
```
You will see a screen where you can select your desired tickers, date, LLMs, research depth, etc.
The CLI provides two main modes:
#### 1. Discover Trending Stocks
Find trending stocks from recent news using LLM-powered entity extraction:
- Select a lookback period (1h, 6h, 24h, or 7d)
- Optionally filter by sector (Technology, Healthcare, Finance, Energy, Consumer Goods, Industrials)
- Optionally filter by event type (Earnings, Merger/Acquisition, Regulatory, Product Launch, Executive Change)
- View ranked results with scores, mentions, and sentiment
- Drill into stock details and seamlessly transition to full analysis
#### 2. Analyze Specific Ticker
Run full multi-agent analysis on a specific stock:
- Enter any ticker symbol and analysis date
- Select which analyst agents to deploy
- Configure research depth (debate rounds)
- Watch real-time progress as agents collaborate
- View comprehensive reports from each team
<p align="center">
<img src="assets/cli/cli_init.png" width="100%" style="display: inline-block; margin: 0 2%;">
@ -167,7 +181,6 @@ from tradingagents.default_config import DEFAULT_CONFIG
ta = TradingAgentsGraph(debug=True, config=DEFAULT_CONFIG.copy())
# forward propagate
_, decision = ta.propagate("NVDA", "2024-05-10")
print(decision)
```
@ -178,48 +191,80 @@ You can also adjust the default configuration to set your own choice of LLMs, de
from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.default_config import DEFAULT_CONFIG
# Create a custom config
config = DEFAULT_CONFIG.copy()
config["deep_think_llm"] = "gpt-4.1-nano" # Use a different model
config["quick_think_llm"] = "gpt-4.1-nano" # Use a different model
config["max_debate_rounds"] = 1 # Increase debate rounds
config["deep_think_llm"] = "gpt-4.1-nano"
config["quick_think_llm"] = "gpt-4.1-nano"
config["max_debate_rounds"] = 1
# Configure data vendors (default uses yfinance and Alpha Vantage)
config["data_vendors"] = {
"core_stock_apis": "yfinance", # Options: yfinance, alpha_vantage, local
"technical_indicators": "yfinance", # Options: yfinance, alpha_vantage, local
"fundamental_data": "alpha_vantage", # Options: openai, alpha_vantage, local
"news_data": "alpha_vantage", # Options: openai, alpha_vantage, google, local
"core_stock_apis": "yfinance",
"technical_indicators": "yfinance",
"fundamental_data": "alpha_vantage",
"news_data": "alpha_vantage",
}
# Initialize with custom config
ta = TradingAgentsGraph(debug=True, config=config)
# forward propagate
_, decision = ta.propagate("NVDA", "2024-05-10")
print(decision)
```
### Trending Stock Discovery API
You can also use the trending stock discovery feature programmatically:
```python
from tradingagents.graph.trading_graph import TradingAgentsGraph
from tradingagents.agents.discovery.models import (
DiscoveryRequest,
Sector,
EventCategory,
)
from tradingagents.default_config import DEFAULT_CONFIG
ta = TradingAgentsGraph(debug=True, config=DEFAULT_CONFIG.copy())
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY, Sector.HEALTHCARE],
event_filter=[EventCategory.EARNINGS],
max_results=10,
)
result = ta.discover_trending(request)
for stock in result.trending_stocks:
print(f"{stock.ticker}: {stock.company_name} (Score: {stock.score:.2f})")
```
> The default configuration uses yfinance for stock price and technical data, and Alpha Vantage for fundamental and news data. For production use or if you encounter rate limits, consider upgrading to [Alpha Vantage Premium](https://www.alphavantage.co/premium/) for more stable and reliable data access. For offline experimentation, there's a local data vendor option that uses our **Tauric TradingDB**, a curated dataset for backtesting, though this is still in development. We're currently refining this dataset and plan to release it soon alongside our upcoming projects. Stay tuned!
You can view the full list of configurations in `tradingagents/default_config.py`.
## Contributing
### Configuration Options
We welcome contributions from the community! Whether it's fixing a bug, improving documentation, or suggesting a new feature, your input helps make this project better. If you are interested in this line of research, please consider joining our open-source financial AI research community [Tauric Research](https://tauric.ai/).
| Option | Description | Default |
|--------|-------------|---------|
| `llm_provider` | LLM provider (openai, anthropic, google, ollama, openrouter) | openai |
| `deep_think_llm` | Model for complex reasoning tasks | gpt-5 |
| `quick_think_llm` | Model for fast/simple tasks | gpt-5-mini |
| `max_debate_rounds` | Number of bull/bear debate iterations | 2 |
| `max_risk_discuss_rounds` | Number of risk assessment rounds | 2 |
| `discovery_max_results` | Max trending stocks to return | 20 |
| `discovery_min_mentions` | Minimum mentions to include stock | 2 |
## Citation
## Source
Please reference our work if you find *TradingAgents* provides you with some help :)
Thanks to Yijia Xiao and Edward Sun and Di Luo and Wei Wang. Core agent implementation based on [TradingAgents: Multi-Agents LLM Financial Trading Framework](https://arxiv.org/abs/2412.20138)
```
@misc{xiao2025tradingagentsmultiagentsllmfinancial,
title={TradingAgents: Multi-Agents LLM Financial Trading Framework},
title={TradingAgents: Multi-Agents LLM Financial Trading Framework},
author={Yijia Xiao and Edward Sun and Di Luo and Wei Wang},
year={2025},
eprint={2412.20138},
archivePrefix={arXiv},
primaryClass={q-fin.TR},
url={https://arxiv.org/abs/2412.20138},
url={https://arxiv.org/abs/2412.20138},
}
```

261
TEST_COVERAGE_SUMMARY.md Normal file
View File

@ -0,0 +1,261 @@
# Test Coverage Summary
This document provides an overview of the comprehensive unit tests generated for the modified files in this branch.
## Test Files Created
### 1. Agent Utils Tests (`tests/agents/utils/`)
#### `test_agent_states.py`
- **Purpose**: Tests for TypedDict state classes used throughout the trading agents system
- **Coverage**:
- `InvestDebateState`: Research team debate state management
- `RiskDebateState`: Risk management team state handling
- `AgentState`: Main agent state with nested debate states
- **Test Scenarios**:
- State structure validation
- Empty and populated states
- Multiline conversation histories
- Count variations and speaker tracking
- Complete workflow scenarios
- **Test Count**: 20+ tests
#### `test_agent_utils.py`
- **Purpose**: Tests for agent utility functions
- **Coverage**:
- `create_msg_delete()`: Message deletion and Anthropic compatibility
- **Test Scenarios**:
- Message removal operations
- Placeholder message creation
- Empty state handling
- Large message lists
- State immutability
- Message ID preservation
- **Test Count**: 11 tests
#### `test_memory.py`
- **Purpose**: Tests for FinancialSituationMemory class (chromadb-based)
- **Coverage**:
- Initialization with different backends (OpenAI, Ollama)
- Embedding generation
- Situation and advice storage
- Memory retrieval and similarity scoring
- **Test Scenarios**:
- Backend configuration
- Embedding model selection
- Single and multiple situation additions
- ID offset management
- Memory querying with similarity scores
- Cache behavior
- Empty list handling
- **Test Count**: 15+ tests
### 2. Dataflows Tests (`tests/dataflows/`)
#### `test_alpha_vantage_news.py`
- **Purpose**: Tests for Alpha Vantage news API integration
- **Coverage**:
- `get_news()`: Ticker-specific news retrieval
- `get_insider_transactions()`: Insider trading data
- `get_bulk_news_alpha_vantage()`: Bulk news fetching
- **Test Scenarios**:
- API parameter validation
- Time period calculations
- Article parsing and content truncation
- Invalid data format handling
- Empty feed responses
- Malformed article data
- Various lookback periods
- **Test Count**: 18+ tests
#### `test_google.py`
- **Purpose**: Tests for Google News integration
- **Coverage**:
- `get_google_news()`: Query-based news search
- `get_bulk_news_google()`: Bulk news aggregation
- **Test Scenarios**:
- Query formatting (space to plus conversion)
- Result formatting and deduplication
- Empty results handling
- Date calculation and formatting
- Multiple query execution
- Content truncation
- Error handling
- **Test Count**: 15+ tests
#### `test_interface.py`
- **Purpose**: Tests for the dataflows interface layer (vendor routing)
- **Coverage**:
- `parse_lookback_period()`: Time period parsing
- `get_category_for_method()`: Method categorization
- `get_bulk_news()`: Cached bulk news retrieval
- `route_to_vendor()`: Vendor fallback logic
- **Test Scenarios**:
- Lookback period parsing (1h, 6h, 24h, 7d)
- Case insensitivity and whitespace handling
- Invalid period error handling
- Method-to-category mapping
- Vendor routing with fallbacks
- Cache behavior (TTL)
- Article conversion to NewsArticle objects
- Multiple vendor implementations
- All-vendor-fail scenarios
- **Test Count**: 20+ tests
### 3. Configuration Tests (`tests/`)
#### `test_default_config.py`
- **Purpose**: Tests for DEFAULT_CONFIG dictionary
- **Coverage**: All configuration keys and their validity
- **Test Scenarios**:
- Config existence and structure
- Path configurations (project_dir, results_dir, data_dir)
- LLM provider and model settings
- Backend URL validation
- Debate and recursion limits
- Data vendor mappings
- Discovery-specific configs (timeout, cache TTL, max results)
- Numeric value positivity checks
- Environment variable respect
- Config immutability safety
- **Test Count**: 18+ tests
### 4. Graph Tests (`tests/graph/`)
#### `test_trading_graph.py`
- **Purpose**: Tests for TradingAgentsGraph main orchestration class
- **Coverage**:
- Initialization with various LLM providers
- Memory instance creation
- Tool node setup
- `discover_trending()`: Trending stock discovery
- `propagate()`: Agent graph execution
- `reflect_and_remember()`: Learning and reflection
- `analyze_trending()`: Stock analysis workflow
- **Test Scenarios**:
- Default and custom configuration
- OpenAI, Anthropic, Google, Ollama provider support
- Unsupported provider error handling
- Memory creation for all agent types
- Bulk news retrieval and entity extraction
- Sector and event filtering
- Timeout handling (hard timeout enforcement)
- Error handling and failure status
- Default request parameters
- Trade date customization
- Complete analysis workflows
- **Test Count**: 25+ tests
## Testing Best Practices Followed
### 1. **Comprehensive Coverage**
- Happy path scenarios
- Edge cases (empty inputs, malformed data)
- Error conditions and exception handling
- Boundary values and limit testing
### 2. **Mocking Strategy**
- External dependencies mocked (APIs, databases, LLMs)
- Focused unit testing without integration overhead
- Proper mock assertions to verify call patterns
### 3. **Test Organization**
- Tests grouped by class/functionality
- Descriptive test names following pattern: `test_<what>_<scenario>`
- Clear docstrings explaining test purpose
### 4. **Fixtures and Setup**
- Reusable fixtures for common configurations
- Proper mock setup and teardown
- Configuration dictionaries for different scenarios
### 5. **Assertions**
- Type checking (isinstance)
- Value equality checks
- Exception matching with pytest.raises
- Call count and argument verification
### 6. **Coverage Areas**
- Pure function logic
- State management
- API integration layers
- Configuration handling
- Error paths and exceptions
- Caching behavior
- Data transformation
## Running the Tests
```bash
# Run all tests
pytest tests/
# Run specific test file
pytest tests/agents/utils/test_memory.py
# Run with coverage
pytest tests/ --cov=tradingagents --cov-report=html
# Run with verbose output
pytest tests/ -v
# Run specific test class
pytest tests/graph/test_trading_graph.py::TestDiscoverTrending
# Run specific test
pytest tests/dataflows/test_interface.py::TestParseLookbackPeriod::test_parse_lookback_1h
```
## Test Dependencies
The tests use the following pytest features and plugins:
- `pytest` - Core testing framework
- `unittest.mock` - Mocking capabilities (Mock, patch, MagicMock)
- `pytest.raises` - Exception testing
- `pytest.fixture` - Test fixtures
## Files Modified vs. Tests Created
| Modified File | Test File | Test Count |
|--------------|-----------|------------|
| `tradingagents/agents/utils/agent_states.py` | `tests/agents/utils/test_agent_states.py` | 20+ |
| `tradingagents/agents/utils/agent_utils.py` | `tests/agents/utils/test_agent_utils.py` | 11 |
| `tradingagents/agents/utils/memory.py` | `tests/agents/utils/test_memory.py` | 15+ |
| `tradingagents/dataflows/alpha_vantage_news.py` | `tests/dataflows/test_alpha_vantage_news.py` | 18+ |
| `tradingagents/dataflows/google.py` | `tests/dataflows/test_google.py` | 15+ |
| `tradingagents/dataflows/interface.py` | `tests/dataflows/test_interface.py` | 20+ |
| `tradingagents/default_config.py` | `tests/test_default_config.py` | 18+ |
| `tradingagents/graph/trading_graph.py` | `tests/graph/test_trading_graph.py` | 25+ |
## Total Test Count
**Approximately 142+ unit tests** covering critical functionality in the modified files.
## Notes on Discovery Module
The discovery module (new in this branch) already has comprehensive tests provided:
- `tests/discovery/test_api.py`
- `tests/discovery/test_bulk_news.py`
- `tests/discovery/test_cli.py`
- `tests/discovery/test_entity_extractor.py`
- `tests/discovery/test_integration.py`
- `tests/discovery/test_models.py`
- `tests/discovery/test_persistence.py`
- `tests/discovery/test_scorer.py`
- `tests/discovery/test_sector_classifier.py`
- `tests/discovery/test_stock_resolver.py`
These tests were created alongside the discovery module implementation and follow similar patterns to the tests generated here.
## Missing Coverage (Intentional)
The following modified files were not given new unit tests:
1. **`tradingagents/dataflows/openai.py`** - Heavily dependent on external OpenAI API; integration tests more appropriate
2. **`tradingagents/dataflows/trending/sector_classifier.py`** - Already has `tests/discovery/test_sector_classifier.py`
3. **`tradingagents/dataflows/trending/stock_resolver.py`** - Already has `tests/discovery/test_stock_resolver.py`
4. **CLI files** - Already have `tests/discovery/test_cli.py`
## Recommendations
1. Run tests locally to verify all pass
2. Add pytest to `pyproject.toml` or `requirements.txt` if not already present
3. Set up CI/CD to run tests on every commit
4. Aim for >80% code coverage on modified files
5. Add integration tests for end-to-end workflows
6. Consider property-based testing with `hypothesis` for complex logic

1791
cli/main.py
View File

File diff suppressed because it is too large Load Diff

2
main.py
View File

@ -8,7 +8,7 @@ load_dotenv()
# Create a custom config
config = DEFAULT_CONFIG.copy()
config["deep_think_llm"] = "gpt-4o-mini" # Use a different model
config["deep_think_llm"] = "gpt-5" # Use a different model
config["quick_think_llm"] = "gpt-4o-mini" # Use a different model
config["max_debate_rounds"] = 1 # Increase debate rounds

0
tests/agents/__init__.py Normal file
View File

0
tests/agents/utils/__init__.py Normal file
View File

346
tests/agents/utils/test_agent_states.py Normal file
View File

@ -0,0 +1,346 @@
import pytest
from tradingagents.agents.utils.agent_states import (
InvestDebateState,
RiskDebateState,
AgentState,
)
class TestInvestDebateState:
"""Test suite for InvestDebateState TypedDict."""
def test_invest_debate_state_structure(self):
"""Test that InvestDebateState can be instantiated with all required fields."""
state = {
"bull_history": "Bull argument 1\nBull argument 2",
"bear_history": "Bear argument 1\nBear argument 2",
"history": "Combined history",
"current_response": "Latest response",
"judge_decision": "Final decision",
"count": 3,
}
assert state["bull_history"] == "Bull argument 1\nBull argument 2"
assert state["bear_history"] == "Bear argument 1\nBear argument 2"
assert state["history"] == "Combined history"
assert state["current_response"] == "Latest response"
assert state["judge_decision"] == "Final decision"
assert state["count"] == 3
def test_invest_debate_state_empty_strings(self):
"""Test InvestDebateState with empty strings."""
state = {
"bull_history": "",
"bear_history": "",
"history": "",
"current_response": "",
"judge_decision": "",
"count": 0,
}
assert state["bull_history"] == ""
assert state["bear_history"] == ""
assert state["count"] == 0
def test_invest_debate_state_count_variations(self):
"""Test InvestDebateState with various count values."""
for count in [0, 1, 5, 10, 100]:
state = {
"bull_history": f"History for count {count}",
"bear_history": f"Bear history for count {count}",
"history": "Combined",
"current_response": "Response",
"judge_decision": "Decision",
"count": count,
}
assert state["count"] == count
def test_invest_debate_state_multiline_histories(self):
"""Test InvestDebateState with multiline conversation histories."""
bull_history = "\n".join([f"Bull point {i}" for i in range(5)])
bear_history = "\n".join([f"Bear point {i}" for i in range(5)])
state = {
"bull_history": bull_history,
"bear_history": bear_history,
"history": "Combined history",
"current_response": "Latest",
"judge_decision": "Final",
"count": 5,
}
assert state["bull_history"].count("\n") == 4
assert state["bear_history"].count("\n") == 4
class TestRiskDebateState:
"""Test suite for RiskDebateState TypedDict."""
def test_risk_debate_state_structure(self):
"""Test that RiskDebateState can be instantiated with all required fields."""
state = {
"risky_history": "Risky analysis 1",
"safe_history": "Safe analysis 1",
"neutral_history": "Neutral analysis 1",
"history": "Combined history",
"latest_speaker": "risky",
"current_risky_response": "Latest risky response",
"current_safe_response": "Latest safe response",
"current_neutral_response": "Latest neutral response",
"judge_decision": "Portfolio manager decision",
"count": 2,
}
assert state["risky_history"] == "Risky analysis 1"
assert state["safe_history"] == "Safe analysis 1"
assert state["neutral_history"] == "Neutral analysis 1"
assert state["latest_speaker"] == "risky"
assert state["current_risky_response"] == "Latest risky response"
assert state["count"] == 2
def test_risk_debate_state_speaker_variations(self):
"""Test RiskDebateState with different speaker values."""
speakers = ["risky", "safe", "neutral", "judge"]
for speaker in speakers:
state = {
"risky_history": "Risky",
"safe_history": "Safe",
"neutral_history": "Neutral",
"history": "History",
"latest_speaker": speaker,
"current_risky_response": "Risky resp",
"current_safe_response": "Safe resp",
"current_neutral_response": "Neutral resp",
"judge_decision": "Decision",
"count": 1,
}
assert state["latest_speaker"] == speaker
def test_risk_debate_state_empty_responses(self):
"""Test RiskDebateState with empty response strings."""
state = {
"risky_history": "",
"safe_history": "",
"neutral_history": "",
"history": "",
"latest_speaker": "",
"current_risky_response": "",
"current_safe_response": "",
"current_neutral_response": "",
"judge_decision": "",
"count": 0,
}
assert state["current_risky_response"] == ""
assert state["current_safe_response"] == ""
assert state["current_neutral_response"] == ""
def test_risk_debate_state_long_histories(self):
"""Test RiskDebateState with extended conversation histories."""
risky_history = "\n".join([f"Risky round {i}" for i in range(10)])
safe_history = "\n".join([f"Safe round {i}" for i in range(10)])
neutral_history = "\n".join([f"Neutral round {i}" for i in range(10)])
state = {
"risky_history": risky_history,
"safe_history": safe_history,
"neutral_history": neutral_history,
"history": "Combined",
"latest_speaker": "neutral",
"current_risky_response": "Latest risky",
"current_safe_response": "Latest safe",
"current_neutral_response": "Latest neutral",
"judge_decision": "Final decision",
"count": 10,
}
assert len(state["risky_history"].split("\n")) == 10
assert len(state["safe_history"].split("\n")) == 10
assert len(state["neutral_history"].split("\n")) == 10
class TestAgentState:
"""Test suite for AgentState MessagesState."""
def test_agent_state_basic_fields(self):
"""Test AgentState with basic required fields."""
state = {
"messages": [],
"company_of_interest": "AAPL",
"trade_date": "2024-01-15",
"sender": "market_analyst",
}
assert state["company_of_interest"] == "AAPL"
assert state["trade_date"] == "2024-01-15"
assert state["sender"] == "market_analyst"
def test_agent_state_with_reports(self):
"""Test AgentState with all analyst reports."""
state = {
"messages": [],
"company_of_interest": "TSLA",
"trade_date": "2024-02-20",
"sender": "fundamentals_analyst",
"market_report": "Market analysis for TSLA",
"sentiment_report": "Social sentiment positive",
"news_report": "Recent news about Tesla",
"fundamentals_report": "Strong fundamentals",
}
assert state["market_report"] == "Market analysis for TSLA"
assert state["sentiment_report"] == "Social sentiment positive"
assert state["news_report"] == "Recent news about Tesla"
assert state["fundamentals_report"] == "Strong fundamentals"
def test_agent_state_with_debate_states(self):
"""Test AgentState with nested debate states."""
invest_debate = {
"bull_history": "Bull points",
"bear_history": "Bear points",
"history": "Combined",
"current_response": "Response",
"judge_decision": "Decision",
"count": 2,
}
risk_debate = {
"risky_history": "Risky analysis",
"safe_history": "Safe analysis",
"neutral_history": "Neutral analysis",
"history": "Combined risk history",
"latest_speaker": "safe",
"current_risky_response": "Risky resp",
"current_safe_response": "Safe resp",
"current_neutral_response": "Neutral resp",
"judge_decision": "Portfolio decision",
"count": 3,
}
state = {
"messages": [],
"company_of_interest": "NVDA",
"trade_date": "2024-03-10",
"sender": "research_manager",
"investment_debate_state": invest_debate,
"risk_debate_state": risk_debate,
}
assert state["investment_debate_state"]["count"] == 2
assert state["risk_debate_state"]["count"] == 3
assert state["risk_debate_state"]["latest_speaker"] == "safe"
def test_agent_state_with_plans(self):
"""Test AgentState with investment and trade plans."""
state = {
"messages": [],
"company_of_interest": "MSFT",
"trade_date": "2024-04-05",
"sender": "trader",
"investment_plan": "Long position on MSFT based on analysis",
"trader_investment_plan": "Execute buy order for 100 shares",
"final_trade_decision": "BUY 100 shares at market price",
}
assert "Long position" in state["investment_plan"]
assert "Execute buy order" in state["trader_investment_plan"]
assert "BUY 100 shares" in state["final_trade_decision"]
def test_agent_state_ticker_variations(self):
"""Test AgentState with various ticker symbols."""
tickers = ["AAPL", "GOOGL", "AMZN", "TSLA", "MSFT", "META", "SPY", "QQQ"]
for ticker in tickers:
state = {
"messages": [],
"company_of_interest": ticker,
"trade_date": "2024-01-01",
"sender": "analyst",
}
assert state["company_of_interest"] == ticker
def test_agent_state_date_formats(self):
"""Test AgentState with different date string formats."""
dates = [
"2024-01-15",
"2024-12-31",
"2023-06-30",
"2025-03-20",
]
for date_str in dates:
state = {
"messages": [],
"company_of_interest": "SPY",
"trade_date": date_str,
"sender": "system",
}
assert state["trade_date"] == date_str
def test_agent_state_sender_variations(self):
"""Test AgentState with different sender agent types."""
senders = [
"market_analyst",
"social_analyst",
"news_analyst",
"fundamentals_analyst",
"bull_researcher",
"bear_researcher",
"research_manager",
"trader",
"risky_analyst",
"safe_analyst",
"neutral_analyst",
"portfolio_manager",
]
for sender in senders:
state = {
"messages": [],
"company_of_interest": "AAPL",
"trade_date": "2024-01-01",
"sender": sender,
}
assert state["sender"] == sender
def test_agent_state_complete_workflow(self):
"""Test AgentState with a complete workflow scenario."""
state = {
"messages": [],
"company_of_interest": "AAPL",
"trade_date": "2024-01-15",
"sender": "portfolio_manager",
"market_report": "Price trending upward, volume increasing",
"sentiment_report": "Positive sentiment on social media",
"news_report": "New product launch announced",
"fundamentals_report": "Strong earnings, P/E ratio favorable",
"investment_debate_state": {
"bull_history": "Strong growth potential",
"bear_history": "Market saturation concerns",
"history": "Debate conducted",
"current_response": "Bull case stronger",
"judge_decision": "Recommend buy",
"count": 3,
},
"investment_plan": "Enter long position",
"trader_investment_plan": "Buy 200 shares at limit price",
"risk_debate_state": {
"risky_history": "Aggressive position sizing recommended",
"safe_history": "Conservative approach suggested",
"neutral_history": "Balanced position preferred",
"history": "Risk analysis complete",
"latest_speaker": "neutral",
"current_risky_response": "Go all in",
"current_safe_response": "Small position only",
"current_neutral_response": "Moderate position",
"judge_decision": "Moderate position approved",
"count": 2,
},
"final_trade_decision": "BUY 200 AAPL @ $150 limit",
}
assert state["company_of_interest"] == "AAPL"
assert "BUY" in state["final_trade_decision"]
assert state["investment_debate_state"]["judge_decision"] == "Recommend buy"
assert state["risk_debate_state"]["latest_speaker"] == "neutral"

176
tests/agents/utils/test_agent_utils.py Normal file
View File

@ -0,0 +1,176 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from langchain_core.messages import HumanMessage, RemoveMessage
from tradingagents.agents.utils.agent_utils import create_msg_delete
class TestCreateMsgDelete:
"""Test suite for create_msg_delete function."""
def test_create_msg_delete_returns_callable(self):
"""Test that create_msg_delete returns a callable function."""
delete_func = create_msg_delete()
assert callable(delete_func)
def test_delete_messages_removes_all_messages(self):
"""Test that delete_messages removes all existing messages."""
# Create mock messages with IDs
mock_msg1 = Mock(spec=HumanMessage)
mock_msg1.id = "msg_1"
mock_msg2 = Mock(spec=HumanMessage)
mock_msg2.id = "msg_2"
mock_msg3 = Mock(spec=HumanMessage)
mock_msg3.id = "msg_3"
state = {"messages": [mock_msg1, mock_msg2, mock_msg3]}
delete_func = create_msg_delete()
result = delete_func(state)
# Should return removal operations for all messages plus a placeholder
assert "messages" in result
messages = result["messages"]
# First 3 should be RemoveMessage operations
removal_count = sum(1 for msg in messages if isinstance(msg, RemoveMessage))
assert removal_count == 3
# Last message should be the placeholder HumanMessage
assert isinstance(messages[-1], HumanMessage)
assert messages[-1].content == "Continue"
def test_delete_messages_empty_state(self):
"""Test delete_messages with an empty message list."""
state = {"messages": []}
delete_func = create_msg_delete()
result = delete_func(state)
# Should only contain the placeholder message
assert len(result["messages"]) == 1
assert isinstance(result["messages"][0], HumanMessage)
assert result["messages"][0].content == "Continue"
def test_delete_messages_single_message(self):
"""Test delete_messages with a single message."""
mock_msg = Mock(spec=HumanMessage)
mock_msg.id = "single_msg"
state = {"messages": [mock_msg]}
delete_func = create_msg_delete()
result = delete_func(state)
assert len(result["messages"]) == 2 # 1 removal + 1 placeholder
assert isinstance(result["messages"][0], RemoveMessage)
assert isinstance(result["messages"][1], HumanMessage)
def test_delete_messages_preserves_message_ids(self):
"""Test that RemoveMessage operations use correct message IDs."""
msg_ids = ["id_1", "id_2", "id_3", "id_4"]
mock_messages = []
for msg_id in msg_ids:
mock_msg = Mock(spec=HumanMessage)
mock_msg.id = msg_id
mock_messages.append(mock_msg)
state = {"messages": mock_messages}
delete_func = create_msg_delete()
result = delete_func(state)
# Extract RemoveMessage operations
removal_operations = [msg for msg in result["messages"] if isinstance(msg, RemoveMessage)]
removal_ids = [op.id for op in removal_operations]
# All original message IDs should be in removal operations
for original_id in msg_ids:
assert original_id in removal_ids
def test_delete_messages_anthropic_compatibility(self):
"""Test that the placeholder message ensures Anthropic API compatibility."""
# Anthropic requires at least one message in the conversation
mock_msg = Mock(spec=HumanMessage)
mock_msg.id = "test_msg"
state = {"messages": [mock_msg]}
delete_func = create_msg_delete()
result = delete_func(state)
# Verify placeholder is a HumanMessage (required by Anthropic)
placeholder = result["messages"][-1]
assert isinstance(placeholder, HumanMessage)
assert placeholder.content == "Continue"
def test_delete_messages_large_message_list(self):
"""Test delete_messages with a large number of messages."""
# Create 100 mock messages
mock_messages = []
for i in range(100):
mock_msg = Mock(spec=HumanMessage)
mock_msg.id = f"msg_{i}"
mock_messages.append(mock_msg)
state = {"messages": mock_messages}
delete_func = create_msg_delete()
result = delete_func(state)
# Should have 100 removal operations + 1 placeholder
assert len(result["messages"]) == 101
# Count removal operations
removal_count = sum(1 for msg in result["messages"] if isinstance(msg, RemoveMessage))
assert removal_count == 100
def test_delete_messages_multiple_calls(self):
"""Test that create_msg_delete can be called multiple times."""
mock_msg1 = Mock(spec=HumanMessage)
mock_msg1.id = "msg_1"
mock_msg2 = Mock(spec=HumanMessage)
mock_msg2.id = "msg_2"
state1 = {"messages": [mock_msg1]}
state2 = {"messages": [mock_msg1, mock_msg2]}
delete_func1 = create_msg_delete()
delete_func2 = create_msg_delete()
result1 = delete_func1(state1)
result2 = delete_func2(state2)
# Each call should work independently
assert len(result1["messages"]) == 2 # 1 removal + placeholder
assert len(result2["messages"]) == 3 # 2 removals + placeholder
def test_delete_messages_state_immutability(self):
"""Test that delete_messages doesn't modify the original state."""
mock_msg = Mock(spec=HumanMessage)
mock_msg.id = "test_id"
original_state = {"messages": [mock_msg]}
original_msg_count = len(original_state["messages"])
delete_func = create_msg_delete()
result = delete_func(original_state)
# Original state should remain unchanged
assert len(original_state["messages"]) == original_msg_count
assert original_state["messages"][0] is mock_msg
def test_delete_messages_return_structure(self):
"""Test that delete_messages returns the correct structure."""
mock_msg = Mock(spec=HumanMessage)
mock_msg.id = "test_msg"
state = {"messages": [mock_msg]}
delete_func = create_msg_delete()
result = delete_func(state)
# Result should be a dict with 'messages' key
assert isinstance(result, dict)
assert "messages" in result
assert isinstance(result["messages"], list)

324
tests/agents/utils/test_memory.py Normal file
View File

@ -0,0 +1,324 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from tradingagents.agents.utils.memory import FinancialSituationMemory
class TestFinancialSituationMemory:
"""Test suite for FinancialSituationMemory class."""
@pytest.fixture
def mock_config_openai(self):
"""Fixture for OpenAI configuration."""
return {
"backend_url": "https://api.openai.com/v1",
"llm_provider": "openai",
}
@pytest.fixture
def mock_config_ollama(self):
"""Fixture for Ollama configuration."""
return {
"backend_url": "http://localhost:11434/v1",
"llm_provider": "ollama",
}
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_init_with_openai_backend(self, mock_chroma, mock_openai, mock_config_openai):
"""Test initialization with OpenAI backend."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
memory = FinancialSituationMemory("test_memory", mock_config_openai)
assert memory.embedding == "text-embedding-3-small"
mock_openai.assert_called_once_with(base_url="https://api.openai.com/v1")
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_init_with_ollama_backend(self, mock_chroma, mock_openai, mock_config_ollama):
"""Test initialization with Ollama backend."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
memory = FinancialSituationMemory("test_memory", mock_config_ollama)
assert memory.embedding == "nomic-embed-text"
mock_openai.assert_called_once_with(base_url="http://localhost:11434/v1")
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_collection_creation(self, mock_chroma, mock_openai, mock_config_openai):
"""Test that ChromaDB collection is created with correct name."""
mock_collection = Mock()
mock_chroma_instance = Mock()
mock_chroma.return_value = mock_chroma_instance
mock_chroma_instance.create_collection.return_value = mock_collection
memory = FinancialSituationMemory("my_test_collection", mock_config_openai)
mock_chroma_instance.create_collection.assert_called_once_with(name="my_test_collection")
assert memory.situation_collection == mock_collection
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_get_embedding(self, mock_chroma, mock_openai, mock_config_openai):
"""Test get_embedding method returns correct embedding vector."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2, 0.3, 0.4])]
mock_client.embeddings.create.return_value = mock_response
memory = FinancialSituationMemory("test_memory", mock_config_openai)
embedding = memory.get_embedding("test text")
assert embedding == [0.1, 0.2, 0.3, 0.4]
mock_client.embeddings.create.assert_called_once_with(
model="text-embedding-3-small",
input="test text"
)
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_get_embedding_with_ollama(self, mock_chroma, mock_openai, mock_config_ollama):
"""Test get_embedding uses correct model for Ollama."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.5, 0.6])]
mock_client.embeddings.create.return_value = mock_response
memory = FinancialSituationMemory("test_memory", mock_config_ollama)
embedding = memory.get_embedding("ollama test")
mock_client.embeddings.create.assert_called_once_with(
model="nomic-embed-text",
input="ollama test"
)
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_add_situations_single(self, mock_chroma, mock_openai, mock_config_openai):
"""Test adding a single situation and advice pair."""
mock_collection = Mock()
mock_collection.count.return_value = 0
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2])]
mock_client.embeddings.create.return_value = mock_response
memory = FinancialSituationMemory("test_memory", mock_config_openai)
situations_and_advice = [
("High volatility market", "Reduce position sizes")
]
memory.add_situations(situations_and_advice)
mock_collection.add.assert_called_once()
call_kwargs = mock_collection.add.call_args[1]
assert call_kwargs["documents"] == ["High volatility market"]
assert call_kwargs["metadatas"] == [{"recommendation": "Reduce position sizes"}]
assert call_kwargs["ids"] == ["0"]
assert len(call_kwargs["embeddings"]) == 1
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_add_situations_multiple(self, mock_chroma, mock_openai, mock_config_openai):
"""Test adding multiple situations at once."""
mock_collection = Mock()
mock_collection.count.return_value = 0
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2])]
mock_client.embeddings.create.return_value = mock_response
memory = FinancialSituationMemory("test_memory", mock_config_openai)
situations_and_advice = [
("Bull market conditions", "Increase long positions"),
("Bear market conditions", "Increase short positions"),
("Sideways market", "Use range trading strategies"),
]
memory.add_situations(situations_and_advice)
mock_collection.add.assert_called_once()
call_kwargs = mock_collection.add.call_args[1]
assert len(call_kwargs["documents"]) == 3
assert len(call_kwargs["metadatas"]) == 3
assert call_kwargs["ids"] == ["0", "1", "2"]
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_add_situations_with_existing_offset(self, mock_chroma, mock_openai, mock_config_openai):
"""Test that ID offset is calculated correctly when adding to existing collection."""
mock_collection = Mock()
mock_collection.count.return_value = 5 # Already has 5 items
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2])]
mock_client.embeddings.create.return_value = mock_response
memory = FinancialSituationMemory("test_memory", mock_config_openai)
situations_and_advice = [
("New situation", "New advice"),
("Another situation", "Another advice"),
]
memory.add_situations(situations_and_advice)
call_kwargs = mock_collection.add.call_args[1]
# IDs should start from 5 (the existing count)
assert call_kwargs["ids"] == ["5", "6"]
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_get_memories_single_match(self, mock_chroma, mock_openai, mock_config_openai):
"""Test retrieving a single matching memory."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2])]
mock_client.embeddings.create.return_value = mock_response
# Mock query results
mock_collection.query.return_value = {
"documents": [["Similar market condition"]],
"metadatas": [[{"recommendation": "Apply defensive strategy"}]],
"distances": [[0.15]],
}
memory = FinancialSituationMemory("test_memory", mock_config_openai)
results = memory.get_memories("Current volatile market", n_matches=1)
assert len(results) == 1
assert results[0]["matched_situation"] == "Similar market condition"
assert results[0]["recommendation"] == "Apply defensive strategy"
assert results[0]["similarity_score"] == pytest.approx(0.85, rel=0.01) # 1 - 0.15
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_get_memories_multiple_matches(self, mock_chroma, mock_openai, mock_config_openai):
"""Test retrieving multiple matching memories."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2])]
mock_client.embeddings.create.return_value = mock_response
# Mock query results with 3 matches
mock_collection.query.return_value = {
"documents": [["Match 1", "Match 2", "Match 3"]],
"metadatas": [
[
{"recommendation": "Advice 1"},
{"recommendation": "Advice 2"},
{"recommendation": "Advice 3"},
]
],
"distances": [[0.1, 0.2, 0.3]],
}
memory = FinancialSituationMemory("test_memory", mock_config_openai)
results = memory.get_memories("Query situation", n_matches=3)
assert len(results) == 3
assert results[0]["matched_situation"] == "Match 1"
assert results[1]["matched_situation"] == "Match 2"
assert results[2]["matched_situation"] == "Match 3"
assert results[0]["similarity_score"] > results[1]["similarity_score"]
assert results[1]["similarity_score"] > results[2]["similarity_score"]
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_get_memories_similarity_scores(self, mock_chroma, mock_openai, mock_config_openai):
"""Test that similarity scores are calculated correctly (1 - distance)."""
mock_collection = Mock()
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
mock_response = Mock()
mock_response.data = [Mock(embedding=[0.1, 0.2])]
mock_client.embeddings.create.return_value = mock_response
mock_collection.query.return_value = {
"documents": [["Situation A", "Situation B"]],
"metadatas": [[{"recommendation": "A"}, {"recommendation": "B"}]],
"distances": [[0.0, 0.5]], # Perfect match and moderate match
}
memory = FinancialSituationMemory("test_memory", mock_config_openai)
results = memory.get_memories("Test query", n_matches=2)
assert results[0]["similarity_score"] == pytest.approx(1.0, rel=0.01) # 1 - 0.0
assert results[1]["similarity_score"] == pytest.approx(0.5, rel=0.01) # 1 - 0.5
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_add_situations_empty_list(self, mock_chroma, mock_openai, mock_config_openai):
"""Test adding an empty list of situations."""
mock_collection = Mock()
mock_collection.count.return_value = 0
mock_chroma.return_value.create_collection.return_value = mock_collection
mock_client = Mock()
mock_openai.return_value = mock_client
memory = FinancialSituationMemory("test_memory", mock_config_openai)
memory.add_situations([])
# add should still be called, but with empty lists
mock_collection.add.assert_called_once()
call_kwargs = mock_collection.add.call_args[1]
assert call_kwargs["documents"] == []
assert call_kwargs["metadatas"] == []
assert call_kwargs["ids"] == []
@patch('tradingagents.agents.utils.memory.OpenAI')
@patch('tradingagents.agents.utils.memory.chromadb.Client')
def test_memory_different_collection_names(self, mock_chroma, mock_openai, mock_config_openai):
"""Test that different memory instances have different collection names."""
mock_chroma_instance = Mock()
mock_chroma.return_value = mock_chroma_instance
mock_chroma_instance.create_collection.return_value = Mock()
memory1 = FinancialSituationMemory("bull_memory", mock_config_openai)
memory2 = FinancialSituationMemory("bear_memory", mock_config_openai)
memory3 = FinancialSituationMemory("trader_memory", mock_config_openai)
# Verify different collections were created
calls = mock_chroma_instance.create_collection.call_args_list
assert len(calls) == 3
assert calls[0][1]["name"] == "bull_memory"
assert calls[1][1]["name"] == "bear_memory"
assert calls[2][1]["name"] == "trader_memory"

0
tests/dataflows/__init__.py Normal file
View File

294
tests/dataflows/test_alpha_vantage_news.py Normal file
View File

@ -0,0 +1,294 @@
import pytest
from unittest.mock import Mock, patch
from datetime import datetime, timedelta
from tradingagents.dataflows.alpha_vantage_news import (
get_news,
get_insider_transactions,
get_bulk_news_alpha_vantage,
)
class TestGetNews:
"""Test suite for get_news function."""
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_news_basic_call(self, mock_format_datetime, mock_api_request):
"""Test basic get_news API call."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
ticker = "AAPL"
start_date = datetime(2024, 1, 1)
end_date = datetime(2024, 1, 31)
result = get_news(ticker, start_date, end_date)
mock_api_request.assert_called_once()
call_args = mock_api_request.call_args[0]
assert call_args[0] == "NEWS_SENTIMENT"
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_news_parameters(self, mock_format_datetime, mock_api_request):
"""Test that get_news passes correct parameters."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
ticker = "TSLA"
start_date = datetime(2024, 2, 1)
end_date = datetime(2024, 2, 15)
result = get_news(ticker, start_date, end_date)
params = mock_api_request.call_args[0][1]
assert params["tickers"] == "TSLA"
assert params["sort"] == "LATEST"
assert params["limit"] == "50"
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_news_different_tickers(self, mock_format_datetime, mock_api_request):
"""Test get_news with different ticker symbols."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
tickers = ["AAPL", "GOOGL", "MSFT", "AMZN"]
start_date = datetime(2024, 1, 1)
end_date = datetime(2024, 1, 31)
for ticker in tickers:
result = get_news(ticker, start_date, end_date)
params = mock_api_request.call_args[0][1]
assert params["tickers"] == ticker
class TestGetInsiderTransactions:
"""Test suite for get_insider_transactions function."""
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
def test_get_insider_transactions_basic(self, mock_api_request):
"""Test basic get_insider_transactions call."""
mock_api_request.return_value = {"transactions": []}
symbol = "AAPL"
result = get_insider_transactions(symbol)
mock_api_request.assert_called_once()
call_args = mock_api_request.call_args[0]
assert call_args[0] == "INSIDER_TRANSACTIONS"
assert call_args[1]["symbol"] == "AAPL"
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
def test_get_insider_transactions_different_symbols(self, mock_api_request):
"""Test get_insider_transactions with various symbols."""
mock_api_request.return_value = {}
symbols = ["AAPL", "TSLA", "NVDA", "META"]
for symbol in symbols:
result = get_insider_transactions(symbol)
params = mock_api_request.call_args[0][1]
assert params["symbol"] == symbol
class TestGetBulkNewsAlphaVantage:
"""Test suite for get_bulk_news_alpha_vantage function."""
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_basic(self, mock_format_datetime, mock_api_request):
"""Test basic bulk news retrieval."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
result = get_bulk_news_alpha_vantage(24)
assert isinstance(result, list)
mock_api_request.assert_called_once()
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_lookback_hours(self, mock_format_datetime, mock_api_request):
"""Test that lookback period is calculated correctly."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
lookback_hours = 6
result = get_bulk_news_alpha_vantage(lookback_hours)
# Verify time_from and time_to are set correctly
params = mock_api_request.call_args[0][1]
assert "time_from" in params
assert "time_to" in params
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_parameters(self, mock_format_datetime, mock_api_request):
"""Test that bulk news uses correct parameters."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
result = get_bulk_news_alpha_vantage(24)
params = mock_api_request.call_args[0][1]
assert params["sort"] == "LATEST"
assert params["limit"] == "200"
assert "topics" in params
assert "earnings" in params["topics"]
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_with_articles(self, mock_format_datetime, mock_api_request):
"""Test parsing of article feed data."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_feed = {
"feed": [
{
"title": "Apple announces new product",
"source": "Reuters",
"url": "https://example.com/article1",
"time_published": "20240115T103000",
"summary": "Apple Inc. has announced a groundbreaking new product.",
},
{
"title": "Tech stocks rally",
"source": "Bloomberg",
"url": "https://example.com/article2",
"time_published": "20240115T140000",
"summary": "Technology stocks surged in afternoon trading.",
},
]
}
mock_api_request.return_value = mock_feed
result = get_bulk_news_alpha_vantage(24)
assert len(result) == 2
assert result[0]["title"] == "Apple announces new product"
assert result[0]["source"] == "Reuters"
assert result[1]["title"] == "Tech stocks rally"
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_content_truncation(self, mock_format_datetime, mock_api_request):
"""Test that content snippets are truncated to 500 characters."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
long_summary = "A" * 1000 # 1000 character string
mock_feed = {
"feed": [
{
"title": "Long article",
"source": "Source",
"url": "https://example.com",
"time_published": "20240115T120000",
"summary": long_summary,
}
]
}
mock_api_request.return_value = mock_feed
result = get_bulk_news_alpha_vantage(24)
assert len(result[0]["content_snippet"]) == 500
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_invalid_time_format(self, mock_format_datetime, mock_api_request):
"""Test handling of invalid time_published format."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_feed = {
"feed": [
{
"title": "Article with bad time",
"source": "Source",
"url": "https://example.com",
"time_published": "invalid_format",
"summary": "Summary",
}
]
}
mock_api_request.return_value = mock_feed
result = get_bulk_news_alpha_vantage(24)
# Should fallback to current time
assert len(result) == 1
assert "published_at" in result[0]
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_string_response(self, mock_format_datetime, mock_api_request):
"""Test handling when API returns string instead of dict."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
# Return a JSON string
mock_api_request.return_value = '{"feed": [{"title": "Test"}]}'
result = get_bulk_news_alpha_vantage(24)
# Should handle gracefully and return empty list or parsed data
assert isinstance(result, list)
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_malformed_articles(self, mock_format_datetime, mock_api_request):
"""Test handling of malformed article data."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_feed = {
"feed": [
{"title": "Good article", "source": "Source", "url": "https://example.com", "time_published": "20240115T120000", "summary": "Good"},
{"title": "Missing fields"}, # Malformed
{"source": "No title"}, # Malformed
]
}
mock_api_request.return_value = mock_feed
result = get_bulk_news_alpha_vantage(24)
# Should skip malformed articles
assert len(result) >= 1 # At least the good one
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_empty_feed(self, mock_format_datetime, mock_api_request):
"""Test handling of empty feed."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
result = get_bulk_news_alpha_vantage(24)
assert result == []
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_no_feed_key(self, mock_format_datetime, mock_api_request):
"""Test handling when response doesn't have 'feed' key."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"data": []} # Wrong key
result = get_bulk_news_alpha_vantage(24)
assert result == []
@patch('tradingagents.dataflows.alpha_vantage_news._make_api_request')
@patch('tradingagents.dataflows.alpha_vantage_news.format_datetime_for_api')
def test_get_bulk_news_various_lookback_periods(self, mock_format_datetime, mock_api_request):
"""Test bulk news with various lookback periods."""
mock_format_datetime.side_effect = lambda x: x.strftime("%Y%m%dT%H%M%S")
mock_api_request.return_value = {"feed": []}
lookback_periods = [1, 6, 12, 24, 48, 168] # hours
for hours in lookback_periods:
result = get_bulk_news_alpha_vantage(hours)
assert isinstance(result, list)

248
tests/dataflows/test_google.py Normal file
View File

@ -0,0 +1,248 @@
import pytest
from unittest.mock import Mock, patch
from datetime import datetime, timedelta
from tradingagents.dataflows.google import (
get_google_news,
get_bulk_news_google,
)
class TestGetGoogleNews:
"""Test suite for get_google_news function."""
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_google_news_basic(self, mock_get_news_data):
"""Test basic Google News retrieval."""
mock_get_news_data.return_value = []
query = "AAPL stock"
curr_date = "2024-01-15"
look_back_days = 7
result = get_google_news(query, curr_date, look_back_days)
assert isinstance(result, str)
mock_get_news_data.assert_called_once()
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_google_news_query_formatting(self, mock_get_news_data):
"""Test that query spaces are replaced with plus signs."""
mock_get_news_data.return_value = []
query = "Apple Inc stock news"
curr_date = "2024-01-15"
look_back_days = 7
result = get_google_news(query, curr_date, look_back_days)
# Query should be formatted with + instead of spaces
call_args = mock_get_news_data.call_args[0]
assert "+" in call_args[0] or call_args[0] == query.replace(" ", "+")
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_google_news_with_results(self, mock_get_news_data):
"""Test formatting of news results."""
mock_news = [
{
"title": "Apple stock rises",
"source": "Bloomberg",
"snippet": "Apple Inc. shares rose 5% today...",
},
{
"title": "New iPhone release",
"source": "Reuters",
"snippet": "Apple announces new iPhone model...",
},
]
mock_get_news_data.return_value = mock_news
query = "AAPL"
curr_date = "2024-01-15"
look_back_days = 7
result = get_google_news(query, curr_date, look_back_days)
assert "Apple stock rises" in result
assert "New iPhone release" in result
assert "Bloomberg" in result
assert "Reuters" in result
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_google_news_empty_results(self, mock_get_news_data):
"""Test handling of empty news results."""
mock_get_news_data.return_value = []
query = "NonexistentTicker"
curr_date = "2024-01-15"
look_back_days = 7
result = get_google_news(query, curr_date, look_back_days)
assert result == ""
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_google_news_date_calculation(self, mock_get_news_data):
"""Test that lookback date is calculated correctly."""
mock_get_news_data.return_value = []
query = "TSLA"
curr_date = "2024-01-15"
look_back_days = 30
result = get_google_news(query, curr_date, look_back_days)
# Verify date calculation by checking call arguments
call_args = mock_get_news_data.call_args[0]
before_date = call_args[1]
end_date = call_args[2]
assert end_date == curr_date
class TestGetBulkNewsGoogle:
"""Test suite for get_bulk_news_google function."""
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_basic(self, mock_get_news_data):
"""Test basic bulk news retrieval."""
mock_get_news_data.return_value = []
result = get_bulk_news_google(24)
assert isinstance(result, list)
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_multiple_queries(self, mock_get_news_data):
"""Test that multiple search queries are executed."""
mock_get_news_data.return_value = []
result = get_bulk_news_google(24)
# Should call getNewsData multiple times for different queries
assert mock_get_news_data.call_count >= 3
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_with_articles(self, mock_get_news_data):
"""Test article parsing and deduplication."""
mock_articles = [
{
"title": "Market update",
"source": "Financial Times",
"snippet": "Markets closed higher today...",
"link": "https://example.com/1",
"date": "2024-01-15",
},
{
"title": "Trading news",
"source": "WSJ",
"snippet": "Trading volume increased...",
"link": "https://example.com/2",
"date": "2024-01-15",
},
]
mock_get_news_data.return_value = mock_articles
result = get_bulk_news_google(24)
assert len(result) > 0
assert all("title" in article for article in result)
assert all("source" in article for article in result)
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_deduplication(self, mock_get_news_data):
"""Test that duplicate articles are removed."""
duplicate_article = {
"title": "Same article",
"source": "Source",
"snippet": "Content",
"link": "https://example.com",
"date": "2024-01-15",
}
# Return same article multiple times
mock_get_news_data.return_value = [duplicate_article, duplicate_article]
result = get_bulk_news_google(24)
# Should only appear once
titles = [article["title"] for article in result]
assert titles.count("Same article") <= 1
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_content_truncation(self, mock_get_news_data):
"""Test that content snippets are truncated to 500 characters."""
long_snippet = "A" * 1000
mock_articles = [
{
"title": "Article",
"source": "Source",
"snippet": long_snippet,
"link": "https://example.com",
"date": "2024-01-15",
}
]
mock_get_news_data.return_value = mock_articles
result = get_bulk_news_google(24)
if len(result) > 0:
assert len(result[0]["content_snippet"]) <= 500
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_error_handling(self, mock_get_news_data):
"""Test error handling when getNewsData raises exception."""
mock_get_news_data.side_effect = Exception("API Error")
result = get_bulk_news_google(24)
# Should return empty list or partial results
assert isinstance(result, list)
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_lookback_periods(self, mock_get_news_data):
"""Test with various lookback periods."""
mock_get_news_data.return_value = []
lookback_hours = [1, 6, 12, 24, 48, 168]
for hours in lookback_hours:
result = get_bulk_news_google(hours)
assert isinstance(result, list)
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_date_formatting(self, mock_get_news_data):
"""Test that dates are formatted correctly for API."""
mock_get_news_data.return_value = []
result = get_bulk_news_google(24)
# Check that dates in YYYY-MM-DD format are used
for call in mock_get_news_data.call_args_list:
start_date = call[0][1]
end_date = call[0][2]
# Both should be in YYYY-MM-DD format
assert len(start_date) == 10
assert len(end_date) == 10
assert start_date.count("-") == 2
assert end_date.count("-") == 2
@patch('tradingagents.dataflows.google.getNewsData')
def test_get_bulk_news_google_missing_fields(self, mock_get_news_data):
"""Test handling of articles with missing fields."""
incomplete_articles = [
{"title": "Title only"},
{"source": "Source only"},
{"title": "Complete", "source": "Source", "snippet": "Text", "link": "url", "date": "2024-01-15"},
]
mock_get_news_data.return_value = incomplete_articles
result = get_bulk_news_google(24)
# Should handle missing fields gracefully
assert isinstance(result, list)

309
tests/dataflows/test_interface.py Normal file
View File

@ -0,0 +1,309 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from datetime import datetime, timedelta
from tradingagents.dataflows.interface import (
parse_lookback_period,
get_bulk_news,
get_category_for_method,
get_vendor,
route_to_vendor,
TOOLS_CATEGORIES,
VENDOR_METHODS,
)
from tradingagents.agents.discovery import NewsArticle
class TestParseLookbackPeriod:
"""Test suite for parse_lookback_period function."""
def test_parse_lookback_1h(self):
"""Test parsing '1h' lookback period."""
assert parse_lookback_period("1h") == 1
def test_parse_lookback_6h(self):
"""Test parsing '6h' lookback period."""
assert parse_lookback_period("6h") == 6
def test_parse_lookback_24h(self):
"""Test parsing '24h' lookback period."""
assert parse_lookback_period("24h") == 24
def test_parse_lookback_7d(self):
"""Test parsing '7d' lookback period."""
assert parse_lookback_period("7d") == 168 # 7 * 24
def test_parse_lookback_case_insensitive(self):
"""Test that parsing is case insensitive."""
assert parse_lookback_period("1H") == 1
assert parse_lookback_period("6H") == 6
assert parse_lookback_period("24H") == 24
assert parse_lookback_period("7D") == 168
def test_parse_lookback_with_spaces(self):
"""Test parsing with leading/trailing spaces."""
assert parse_lookback_period(" 1h ") == 1
assert parse_lookback_period(" 24h ") == 24
def test_parse_lookback_invalid_value(self):
"""Test that invalid values raise ValueError."""
with pytest.raises(ValueError, match="Invalid lookback period"):
parse_lookback_period("invalid")
with pytest.raises(ValueError):
parse_lookback_period("10h")
with pytest.raises(ValueError):
parse_lookback_period("2d")
class TestGetCategoryForMethod:
"""Test suite for get_category_for_method function."""
def test_get_category_core_stock_apis(self):
"""Test categorization of core stock API methods."""
assert get_category_for_method("get_stock_data") == "core_stock_apis"
def test_get_category_technical_indicators(self):
"""Test categorization of technical indicator methods."""
assert get_category_for_method("get_indicators") == "technical_indicators"
def test_get_category_fundamental_data(self):
"""Test categorization of fundamental data methods."""
assert get_category_for_method("get_fundamentals") == "fundamental_data"
assert get_category_for_method("get_balance_sheet") == "fundamental_data"
assert get_category_for_method("get_cashflow") == "fundamental_data"
assert get_category_for_method("get_income_statement") == "fundamental_data"
def test_get_category_news_data(self):
"""Test categorization of news data methods."""
assert get_category_for_method("get_news") == "news_data"
assert get_category_for_method("get_global_news") == "news_data"
assert get_category_for_method("get_insider_sentiment") == "news_data"
assert get_category_for_method("get_insider_transactions") == "news_data"
assert get_category_for_method("get_bulk_news") == "news_data"
def test_get_category_invalid_method(self):
"""Test that invalid methods raise ValueError."""
with pytest.raises(ValueError, match="not found in any category"):
get_category_for_method("nonexistent_method")
class TestGetBulkNews:
"""Test suite for get_bulk_news function."""
@patch('tradingagents.dataflows.interface._fetch_bulk_news_from_vendor')
@patch('tradingagents.dataflows.interface._convert_to_news_articles')
def test_get_bulk_news_default_period(self, mock_convert, mock_fetch):
"""Test get_bulk_news with default lookback period."""
mock_fetch.return_value = []
mock_convert.return_value = []
result = get_bulk_news()
mock_fetch.assert_called_once_with("24h")
assert isinstance(result, list)
@patch('tradingagents.dataflows.interface._fetch_bulk_news_from_vendor')
@patch('tradingagents.dataflows.interface._convert_to_news_articles')
def test_get_bulk_news_custom_period(self, mock_convert, mock_fetch):
"""Test get_bulk_news with custom lookback period."""
mock_fetch.return_value = []
mock_convert.return_value = []
result = get_bulk_news("6h")
mock_fetch.assert_called_once_with("6h")
@patch('tradingagents.dataflows.interface._fetch_bulk_news_from_vendor')
@patch('tradingagents.dataflows.interface._convert_to_news_articles')
def test_get_bulk_news_caching(self, mock_convert, mock_fetch):
"""Test that results are cached."""
mock_raw_articles = [
{
"title": "Test Article",
"source": "Source",
"url": "https://example.com",
"published_at": datetime.now().isoformat(),
"content_snippet": "Content",
}
]
mock_article = NewsArticle(
title="Test Article",
source="Source",
url="https://example.com",
published_at=datetime.now(),
content_snippet="Content",
ticker_mentions=[],
)
mock_fetch.return_value = mock_raw_articles
mock_convert.return_value = [mock_article]
# First call should fetch
result1 = get_bulk_news("24h")
call_count_1 = mock_fetch.call_count
# Second call within cache TTL should use cache
result2 = get_bulk_news("24h")
call_count_2 = mock_fetch.call_count
# Fetch should not be called again if cache is working
# (Note: actual caching behavior depends on implementation)
assert isinstance(result1, list)
assert isinstance(result2, list)
@patch('tradingagents.dataflows.interface._fetch_bulk_news_from_vendor')
@patch('tradingagents.dataflows.interface._convert_to_news_articles')
def test_get_bulk_news_converts_articles(self, mock_convert, mock_fetch):
"""Test that raw articles are converted to NewsArticle objects."""
mock_raw = [{"title": "Test"}]
mock_articles = [Mock(spec=NewsArticle)]
mock_fetch.return_value = mock_raw
mock_convert.return_value = mock_articles
result = get_bulk_news("24h")
mock_convert.assert_called_once_with(mock_raw)
assert result == mock_articles
class TestRouteToVendor:
"""Test suite for route_to_vendor function."""
@patch('tradingagents.dataflows.interface.get_vendor')
@patch('tradingagents.dataflows.interface.get_category_for_method')
def test_route_to_vendor_basic(self, mock_get_category, mock_get_vendor):
"""Test basic vendor routing."""
mock_get_category.return_value = "core_stock_apis"
mock_get_vendor.return_value = "yfinance"
# Mock the vendor function
with patch.dict(VENDOR_METHODS, {"get_stock_data": {"yfinance": Mock(return_value="test_data")}}):
result = route_to_vendor("get_stock_data", "AAPL", "2024-01-01")
assert result == "test_data"
@patch('tradingagents.dataflows.interface.get_vendor')
@patch('tradingagents.dataflows.interface.get_category_for_method')
def test_route_to_vendor_fallback(self, mock_get_category, mock_get_vendor):
"""Test vendor fallback when primary fails."""
mock_get_category.return_value = "news_data"
mock_get_vendor.return_value = "alpha_vantage"
# Mock primary vendor to fail, secondary to succeed
primary_mock = Mock(side_effect=Exception("Primary failed"))
secondary_mock = Mock(return_value="fallback_data")
with patch.dict(VENDOR_METHODS, {
"get_news": {
"alpha_vantage": primary_mock,
"openai": secondary_mock,
}
}):
result = route_to_vendor("get_news", "AAPL", "2024-01-01", "2024-01-31")
assert result == "fallback_data"
assert primary_mock.called
assert secondary_mock.called
@patch('tradingagents.dataflows.interface.get_vendor')
@patch('tradingagents.dataflows.interface.get_category_for_method')
def test_route_to_vendor_all_fail(self, mock_get_category, mock_get_vendor):
"""Test that RuntimeError is raised when all vendors fail."""
mock_get_category.return_value = "news_data"
mock_get_vendor.return_value = "alpha_vantage"
# All vendors fail
failing_mock = Mock(side_effect=Exception("Failed"))
with patch.dict(VENDOR_METHODS, {
"get_news": {
"alpha_vantage": failing_mock,
"openai": failing_mock,
}
}):
with pytest.raises(RuntimeError, match="All vendor implementations failed"):
route_to_vendor("get_news", "AAPL", "2024-01-01", "2024-01-31")
@patch('tradingagents.dataflows.interface.get_vendor')
@patch('tradingagents.dataflows.interface.get_category_for_method')
def test_route_to_vendor_multiple_results(self, mock_get_category, mock_get_vendor):
"""Test handling of multiple vendor implementations."""
mock_get_category.return_value = "news_data"
mock_get_vendor.return_value = "local"
# Local vendor has multiple implementations
impl1 = Mock(return_value="result1")
impl2 = Mock(return_value="result2")
with patch.dict(VENDOR_METHODS, {
"get_news": {
"local": [impl1, impl2],
}
}):
result = route_to_vendor("get_news", "AAPL", "2024-01-01", "2024-01-31")
# Should combine multiple results
assert isinstance(result, str)
assert impl1.called
assert impl2.called
def test_route_to_vendor_unsupported_method(self):
"""Test that ValueError is raised for unsupported methods."""
with pytest.raises(ValueError, match="not found in any category"):
route_to_vendor("nonexistent_method", "arg1")
class TestConvertToNewsArticles:
"""Test suite for _convert_to_news_articles function."""
@patch('tradingagents.dataflows.interface._convert_to_news_articles')
def test_convert_empty_list(self, mock_convert):
"""Test converting empty article list."""
mock_convert.return_value = []
from tradingagents.dataflows.interface import _convert_to_news_articles
result = _convert_to_news_articles([])
assert result == []
@patch('tradingagents.dataflows.interface.NewsArticle')
def test_convert_valid_articles(self, mock_news_article):
"""Test converting valid raw articles."""
from tradingagents.dataflows.interface import _convert_to_news_articles
raw_articles = [
{
"title": "Article 1",
"source": "Source 1",
"url": "https://example.com/1",
"published_at": datetime(2024, 1, 15).isoformat(),
"content_snippet": "Content 1",
}
]
result = _convert_to_news_articles(raw_articles)
# Should attempt to create NewsArticle
assert isinstance(result, list)
def test_convert_invalid_date_format(self):
"""Test handling of invalid date formats."""
from tradingagents.dataflows.interface import _convert_to_news_articles
raw_articles = [
{
"title": "Article",
"source": "Source",
"url": "https://example.com",
"published_at": "invalid_date",
"content_snippet": "Content",
}
]
result = _convert_to_news_articles(raw_articles)
# Should handle gracefully
assert isinstance(result, list)

0
tests/dataflows/trending/__init__.py Normal file
View File

0
tests/discovery/__init__.py Normal file
View File

200
tests/discovery/test_api.py Normal file
View File

@ -0,0 +1,200 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from datetime import datetime, timedelta
import signal
from tradingagents.agents.discovery import (
DiscoveryRequest,
DiscoveryResult,
DiscoveryStatus,
TrendingStock,
NewsArticle,
Sector,
EventCategory,
DiscoveryTimeoutError,
)
def create_mock_trending_stock(
ticker: str = "AAPL",
company_name: str = "Apple Inc.",
score: float = 10.0,
sector: Sector = Sector.TECHNOLOGY,
event_type: EventCategory = EventCategory.EARNINGS,
) -> TrendingStock:
return TrendingStock(
ticker=ticker,
company_name=company_name,
score=score,
mention_count=5,
sentiment=0.5,
sector=sector,
event_type=event_type,
news_summary="Test news summary",
source_articles=[],
)
def create_mock_news_article() -> NewsArticle:
return NewsArticle(
title="Test Article",
source="Test Source",
url="https://example.com/article",
published_at=datetime.now(),
content_snippet="Test content about Apple stock",
ticker_mentions=["AAPL"],
)
class TestDiscoverTrendingReturnsDiscoveryResult:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_discover_trending_returns_discovery_result(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = [create_mock_news_article()]
mock_extract.return_value = []
mock_scores.return_value = [create_mock_trending_stock()]
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
result = graph.discover_trending()
assert isinstance(result, DiscoveryResult)
assert result.status == DiscoveryStatus.COMPLETED
assert len(result.trending_stocks) > 0
class TestAnalyzeTrendingCallsPropagate:
def test_analyze_trending_calls_propagate(self):
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.propagate = Mock(return_value=({"final_state": "test"}, "BUY"))
trending_stock = create_mock_trending_stock()
result = graph.analyze_trending(trending_stock)
graph.propagate.assert_called_once()
call_args = graph.propagate.call_args
assert call_args[0][0] == "AAPL"
class TestSectorFilterParameter:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_sector_filter_filters_results(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = [create_mock_news_article()]
mock_extract.return_value = []
mock_scores.return_value = [
create_mock_trending_stock(ticker="AAPL", sector=Sector.TECHNOLOGY),
create_mock_trending_stock(ticker="JPM", sector=Sector.FINANCE),
create_mock_trending_stock(ticker="XOM", sector=Sector.ENERGY),
]
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY],
)
result = graph.discover_trending(request)
assert all(
stock.sector == Sector.TECHNOLOGY for stock in result.trending_stocks
)
class TestEventFilterParameter:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_event_filter_filters_results(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = [create_mock_news_article()]
mock_extract.return_value = []
mock_scores.return_value = [
create_mock_trending_stock(ticker="AAPL", event_type=EventCategory.EARNINGS),
create_mock_trending_stock(
ticker="MSFT", event_type=EventCategory.PRODUCT_LAUNCH
),
create_mock_trending_stock(
ticker="GOOGL", event_type=EventCategory.MERGER_ACQUISITION
),
]
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
request = DiscoveryRequest(
lookback_period="24h",
event_filter=[EventCategory.EARNINGS],
)
result = graph.discover_trending(request)
assert all(
stock.event_type == EventCategory.EARNINGS
for stock in result.trending_stocks
)
class TestTimeoutHandling:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
def test_timeout_raises_discovery_timeout_error(self, mock_bulk_news):
def slow_fetch(*args, **kwargs):
import time
time.sleep(0.5)
return []
mock_bulk_news.side_effect = slow_fetch
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 0.1,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
with pytest.raises(DiscoveryTimeoutError):
graph.discover_trending()

160
tests/discovery/test_bulk_news.py Normal file
View File

@ -0,0 +1,160 @@
import pytest
from datetime import datetime, timedelta
from unittest.mock import patch, MagicMock
from tradingagents.agents.discovery import NewsArticle
from tradingagents.dataflows.alpha_vantage_common import AlphaVantageRateLimitError
class TestGetBulkNewsReturnsNewsArticles:
def test_get_bulk_news_returns_list_of_news_article_objects(self):
mock_raw_news = [
{
"title": "Market Update: Tech stocks rally",
"source": "Reuters",
"url": "https://reuters.com/market-update",
"published_at": datetime.now().isoformat(),
"content_snippet": "Technology stocks led gains in early trading...",
},
{
"title": "Fed signals rate decision",
"source": "Bloomberg",
"url": "https://bloomberg.com/fed-rates",
"published_at": datetime.now().isoformat(),
"content_snippet": "Federal Reserve officials indicated...",
},
]
from tradingagents.dataflows.interface import (
_bulk_news_cache,
get_bulk_news,
)
_bulk_news_cache.clear()
with patch(
"tradingagents.dataflows.interface._fetch_bulk_news_from_vendor"
) as mock_fetch:
mock_fetch.return_value = mock_raw_news
result = get_bulk_news(lookback_period="24h")
assert isinstance(result, list)
assert len(result) == 2
for article in result:
assert isinstance(article, NewsArticle)
assert article.title is not None
assert article.source is not None
assert article.url is not None
class TestLookbackPeriodParsing:
@pytest.mark.parametrize(
"lookback,expected_hours",
[
("1h", 1),
("6h", 6),
("24h", 24),
("7d", 168),
],
)
def test_lookback_period_parsing(self, lookback, expected_hours):
from tradingagents.dataflows.interface import parse_lookback_period
hours = parse_lookback_period(lookback)
assert hours == expected_hours
def test_invalid_lookback_period_raises_error(self):
from tradingagents.dataflows.interface import parse_lookback_period
with pytest.raises(ValueError):
parse_lookback_period("invalid")
class TestVendorFallback:
def test_vendor_fallback_when_primary_rate_limited(self):
mock_openai_news = [
{
"title": "Fallback news from OpenAI",
"source": "Web Search",
"url": "https://example.com/fallback",
"published_at": datetime.now().isoformat(),
"content_snippet": "This is fallback content...",
},
]
from tradingagents.dataflows.interface import (
_bulk_news_cache,
)
_bulk_news_cache.clear()
with patch(
"tradingagents.dataflows.interface.VENDOR_METHODS",
{
"get_bulk_news": {
"alpha_vantage": MagicMock(side_effect=AlphaVantageRateLimitError("Rate limit")),
"openai": MagicMock(return_value=mock_openai_news),
"google": MagicMock(return_value=[]),
}
}
):
from tradingagents.dataflows.interface import _fetch_bulk_news_from_vendor
result = _fetch_bulk_news_from_vendor("24h")
assert isinstance(result, list)
assert len(result) == 1
assert result[0]["title"] == "Fallback news from OpenAI"
class TestBulkNewsCache:
def test_cache_returns_same_results_within_ttl(self):
from tradingagents.dataflows.interface import (
_bulk_news_cache,
_get_cached_bulk_news,
_set_cached_bulk_news,
)
_bulk_news_cache.clear()
test_articles = [
NewsArticle(
title="Cached article",
source="Test Source",
url="https://test.com/cached",
published_at=datetime.now(),
content_snippet="Cached content...",
ticker_mentions=[],
)
]
_set_cached_bulk_news("24h", test_articles)
cached_result = _get_cached_bulk_news("24h")
assert cached_result is not None
assert len(cached_result) == 1
assert cached_result[0].title == "Cached article"
cached_result_again = _get_cached_bulk_news("24h")
assert cached_result_again is not None
assert cached_result_again[0].title == cached_result[0].title
class TestEmptyResultHandling:
def test_empty_result_handling(self):
from tradingagents.dataflows.interface import (
_bulk_news_cache,
get_bulk_news,
)
_bulk_news_cache.clear()
with patch(
"tradingagents.dataflows.interface._fetch_bulk_news_from_vendor"
) as mock_fetch:
mock_fetch.return_value = []
result = get_bulk_news(lookback_period="1h")
assert isinstance(result, list)
assert len(result) == 0

127
tests/discovery/test_cli.py Normal file
View File

@ -0,0 +1,127 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from datetime import datetime
from io import StringIO
from tradingagents.agents.discovery.models import (
DiscoveryResult,
DiscoveryRequest,
DiscoveryStatus,
TrendingStock,
NewsArticle,
Sector,
EventCategory,
)
@pytest.fixture
def sample_trending_stocks():
article = NewsArticle(
title="Apple announces new iPhone",
source="Reuters",
url="https://reuters.com/article",
published_at=datetime.now(),
content_snippet="Apple Inc. unveiled its latest iPhone model today...",
ticker_mentions=["AAPL"],
)
return [
TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=8.5,
mention_count=10,
sentiment=0.7,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="Apple announced new iPhone model with enhanced AI capabilities.",
source_articles=[article],
),
TrendingStock(
ticker="MSFT",
company_name="Microsoft Corporation",
score=7.2,
mention_count=8,
sentiment=0.5,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Microsoft reported strong quarterly earnings.",
source_articles=[article],
),
TrendingStock(
ticker="NVDA",
company_name="NVIDIA Corporation",
score=6.8,
mention_count=6,
sentiment=0.4,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="NVIDIA unveiled new AI chips.",
source_articles=[article],
),
]
@pytest.fixture
def sample_discovery_result(sample_trending_stocks):
request = DiscoveryRequest(
lookback_period="24h",
max_results=20,
)
return DiscoveryResult(
request=request,
trending_stocks=sample_trending_stocks,
status=DiscoveryStatus.COMPLETED,
started_at=datetime.now(),
completed_at=datetime.now(),
)
class TestDiscoveryMenuOption:
def test_discover_trending_flow_exists(self):
from cli.main import discover_trending_flow
assert callable(discover_trending_flow)
def test_select_lookback_period_function_exists(self):
from cli.main import select_lookback_period
assert callable(select_lookback_period)
class TestLookbackSelection:
@patch("cli.main.questionary.select")
def test_lookback_selection_returns_valid_period(self, mock_select):
mock_select.return_value.ask.return_value = "24h"
from cli.main import select_lookback_period
result = select_lookback_period()
assert result in ["1h", "6h", "24h", "7d"]
@patch("cli.main.questionary.select")
def test_lookback_selection_handles_all_options(self, mock_select):
from cli.main import select_lookback_period
for period in ["1h", "6h", "24h", "7d"]:
mock_select.return_value.ask.return_value = period
result = select_lookback_period()
assert result == period
class TestResultsTableDisplay:
def test_create_discovery_results_table(self, sample_trending_stocks):
from cli.main import create_discovery_results_table
table = create_discovery_results_table(sample_trending_stocks)
assert table is not None
assert table.row_count == len(sample_trending_stocks)
def test_table_has_correct_columns(self, sample_trending_stocks):
from cli.main import create_discovery_results_table
table = create_discovery_results_table(sample_trending_stocks)
column_names = [col.header for col in table.columns]
expected_columns = ["Rank", "Ticker", "Company", "Score", "Mentions", "Event Type"]
for expected in expected_columns:
assert expected in column_names
class TestDetailView:
def test_create_stock_detail_panel(self, sample_trending_stocks):
from cli.main import create_stock_detail_panel
stock = sample_trending_stocks[0]
panel = create_stock_detail_panel(stock, rank=1)
assert panel is not None

278
tests/discovery/test_entity_extractor.py Normal file
View File

@ -0,0 +1,278 @@
import pytest
from datetime import datetime
from unittest.mock import patch, MagicMock
from tradingagents.agents.discovery import NewsArticle, EventCategory
class TestExtractEntitiesReturnsCompanyMentions:
def test_extract_entities_returns_list_of_company_mentions(self):
from tradingagents.agents.discovery.entity_extractor import (
extract_entities,
EntityMention,
)
articles = [
NewsArticle(
title="Apple announces new iPhone",
source="Reuters",
url="https://reuters.com/apple",
published_at=datetime.now(),
content_snippet="Apple Inc unveiled its latest iPhone model today with advanced AI features.",
ticker_mentions=[],
),
]
mock_response = MagicMock()
mock_response.entities = [
MagicMock(
company_name="Apple Inc",
confidence=0.95,
context_snippet="Apple Inc unveiled its latest iPhone",
event_type="product_launch",
sentiment=0.7,
)
]
with patch(
"tradingagents.agents.discovery.entity_extractor._get_llm"
) as mock_get_llm:
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value.invoke.return_value = (
mock_response
)
mock_get_llm.return_value = mock_llm
result = extract_entities(articles)
assert isinstance(result, list)
assert len(result) > 0
assert all(isinstance(m, EntityMention) for m in result)
assert result[0].company_name == "Apple Inc"
class TestConfidenceScoreRange:
def test_confidence_score_in_valid_range(self):
from tradingagents.agents.discovery.entity_extractor import (
extract_entities,
EntityMention,
)
articles = [
NewsArticle(
title="Tesla reports earnings",
source="Bloomberg",
url="https://bloomberg.com/tsla",
published_at=datetime.now(),
content_snippet="Tesla Inc reported strong quarterly earnings beating analyst expectations.",
ticker_mentions=[],
),
]
mock_response = MagicMock()
mock_response.entities = [
MagicMock(
company_name="Tesla Inc",
confidence=0.88,
context_snippet="Tesla Inc reported strong quarterly earnings",
event_type="earnings",
sentiment=0.5,
)
]
with patch(
"tradingagents.agents.discovery.entity_extractor._get_llm"
) as mock_get_llm:
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value.invoke.return_value = (
mock_response
)
mock_get_llm.return_value = mock_llm
result = extract_entities(articles)
for mention in result:
assert 0.0 <= mention.confidence <= 1.0
class TestContextSnippetExtraction:
def test_context_snippet_extraction(self):
from tradingagents.agents.discovery.entity_extractor import (
extract_entities,
EntityMention,
)
articles = [
NewsArticle(
title="Microsoft acquires gaming company",
source="WSJ",
url="https://wsj.com/msft",
published_at=datetime.now(),
content_snippet="Microsoft Corporation announced today it will acquire a major gaming studio for $10 billion.",
ticker_mentions=[],
),
]
mock_response = MagicMock()
mock_response.entities = [
MagicMock(
company_name="Microsoft Corporation",
confidence=0.92,
context_snippet="Microsoft Corporation announced today it will acquire",
event_type="merger_acquisition",
sentiment=0.6,
)
]
with patch(
"tradingagents.agents.discovery.entity_extractor._get_llm"
) as mock_get_llm:
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value.invoke.return_value = (
mock_response
)
mock_get_llm.return_value = mock_llm
result = extract_entities(articles)
assert len(result) > 0
for mention in result:
assert mention.context_snippet is not None
assert len(mention.context_snippet) > 0
assert len(mention.context_snippet) <= 150
class TestBatchProcessing:
def test_batch_processing_of_multiple_articles(self):
from tradingagents.agents.discovery.entity_extractor import (
extract_entities,
EntityMention,
BATCH_SIZE,
)
articles = [
NewsArticle(
title=f"News article {i}",
source="Reuters",
url=f"https://reuters.com/article{i}",
published_at=datetime.now(),
content_snippet=f"Company {i} announced major developments today.",
ticker_mentions=[],
)
for i in range(15)
]
mock_response = MagicMock()
mock_response.entities = [
MagicMock(
company_name="Test Company",
confidence=0.85,
context_snippet="Company announced major developments",
event_type="other",
sentiment=0.0,
)
]
with patch(
"tradingagents.agents.discovery.entity_extractor._get_llm"
) as mock_get_llm:
mock_llm = MagicMock()
structured_llm = MagicMock()
structured_llm.invoke.return_value = mock_response
mock_llm.with_structured_output.return_value = structured_llm
mock_get_llm.return_value = mock_llm
result = extract_entities(articles)
expected_batches = (len(articles) + BATCH_SIZE - 1) // BATCH_SIZE
assert structured_llm.invoke.call_count == expected_batches
class TestNoCompanyMentions:
def test_handling_of_articles_with_no_company_mentions(self):
from tradingagents.agents.discovery.entity_extractor import (
extract_entities,
EntityMention,
)
articles = [
NewsArticle(
title="Weather forecast for tomorrow",
source="Weather Channel",
url="https://weather.com/forecast",
published_at=datetime.now(),
content_snippet="Tomorrow will be sunny with temperatures reaching 75 degrees.",
ticker_mentions=[],
),
]
mock_response = MagicMock()
mock_response.entities = []
with patch(
"tradingagents.agents.discovery.entity_extractor._get_llm"
) as mock_get_llm:
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value.invoke.return_value = (
mock_response
)
mock_get_llm.return_value = mock_llm
result = extract_entities(articles)
assert isinstance(result, list)
assert len(result) == 0
class TestEventTypeClassification:
@pytest.mark.parametrize(
"event_type",
[
"earnings",
"merger_acquisition",
"regulatory",
"product_launch",
"executive_change",
"other",
],
)
def test_event_type_classification(self, event_type):
from tradingagents.agents.discovery.entity_extractor import (
extract_entities,
EntityMention,
)
articles = [
NewsArticle(
title="Company news",
source="Reuters",
url="https://reuters.com/news",
published_at=datetime.now(),
content_snippet="A company made an announcement today.",
ticker_mentions=[],
),
]
mock_response = MagicMock()
mock_response.entities = [
MagicMock(
company_name="Test Company",
confidence=0.90,
context_snippet="A company made an announcement",
event_type=event_type,
sentiment=0.0,
)
]
with patch(
"tradingagents.agents.discovery.entity_extractor._get_llm"
) as mock_get_llm:
mock_llm = MagicMock()
mock_llm.with_structured_output.return_value.invoke.return_value = (
mock_response
)
mock_get_llm.return_value = mock_llm
result = extract_entities(articles)
assert len(result) > 0
assert result[0].event_type == EventCategory(event_type)

489
tests/discovery/test_integration.py Normal file
View File

@ -0,0 +1,489 @@
import pytest
import math
from datetime import datetime, timedelta
from unittest.mock import patch, MagicMock
from tradingagents.agents.discovery import (
TrendingStock,
NewsArticle,
DiscoveryRequest,
DiscoveryResult,
DiscoveryStatus,
Sector,
EventCategory,
DiscoveryTimeoutError,
NewsUnavailableError,
)
from tradingagents.agents.discovery.entity_extractor import EntityMention
class TestEndToEndDiscoveryFlow:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_full_discovery_flow_from_news_to_results(
self, mock_scores, mock_extract, mock_bulk_news
):
now = datetime.now()
mock_articles = [
NewsArticle(
title="Apple announces record earnings",
source="Reuters",
url="https://reuters.com/apple-earnings",
published_at=now - timedelta(hours=2),
content_snippet="Apple Inc reported record quarterly earnings...",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="Apple stock surges on AI news",
source="Bloomberg",
url="https://bloomberg.com/apple-ai",
published_at=now - timedelta(hours=1),
content_snippet="Shares of Apple jumped after AI announcement...",
ticker_mentions=["AAPL"],
),
]
mock_bulk_news.return_value = mock_articles
mock_mentions = [
EntityMention(
company_name="Apple Inc",
confidence=0.95,
context_snippet="Apple Inc reported record quarterly earnings",
article_id="article_0",
event_type=EventCategory.EARNINGS,
sentiment=0.8,
),
EntityMention(
company_name="Apple",
confidence=0.92,
context_snippet="Shares of Apple jumped",
article_id="article_1",
event_type=EventCategory.PRODUCT_LAUNCH,
sentiment=0.7,
),
]
mock_extract.return_value = mock_mentions
mock_trending = [
TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=8.5,
mention_count=2,
sentiment=0.75,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Apple reported record earnings and AI progress.",
source_articles=mock_articles,
),
]
mock_scores.return_value = mock_trending
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
request = DiscoveryRequest(lookback_period="24h")
result = graph.discover_trending(request)
assert isinstance(result, DiscoveryResult)
assert result.status == DiscoveryStatus.COMPLETED
assert len(result.trending_stocks) == 1
assert result.trending_stocks[0].ticker == "AAPL"
assert result.trending_stocks[0].mention_count >= 2
mock_bulk_news.assert_called_once_with("24h")
mock_extract.assert_called_once()
mock_scores.assert_called_once()
class TestEntityExtractionToScoringPipeline:
def test_pipeline_from_extraction_to_scoring(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
articles = [
NewsArticle(
title="Microsoft cloud revenue grows",
source="WSJ",
url="https://wsj.com/article1",
published_at=now - timedelta(hours=2),
content_snippet="Microsoft Corporation reported strong cloud growth.",
ticker_mentions=["MSFT"],
),
NewsArticle(
title="Microsoft earnings beat estimates",
source="CNBC",
url="https://cnbc.com/article2",
published_at=now - timedelta(hours=3),
content_snippet="Microsoft earnings exceeded analyst expectations.",
ticker_mentions=["MSFT"],
),
NewsArticle(
title="Tech stocks rally",
source="Bloomberg",
url="https://bloomberg.com/article3",
published_at=now - timedelta(hours=1),
content_snippet="Technology companies led market gains.",
ticker_mentions=[],
),
]
mentions = [
EntityMention(
company_name="Microsoft Corporation",
confidence=0.95,
context_snippet="Microsoft Corporation reported strong cloud growth",
article_id="article_0",
event_type=EventCategory.EARNINGS,
sentiment=0.7,
),
EntityMention(
company_name="Microsoft",
confidence=0.92,
context_snippet="Microsoft earnings exceeded analyst expectations",
article_id="article_1",
event_type=EventCategory.EARNINGS,
sentiment=0.8,
),
]
with patch("tradingagents.agents.discovery.scorer.resolve_ticker") as mock_resolve:
mock_resolve.return_value = "MSFT"
with patch("tradingagents.agents.discovery.scorer.classify_sector") as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(mentions, articles, min_mentions=2)
assert len(result) == 1
assert result[0].ticker == "MSFT"
assert result[0].mention_count == 2
assert result[0].sentiment > 0
class TestNewsVendorFailureGracefulDegradation:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
def test_news_vendor_failure_with_graceful_degradation(self, mock_bulk_news):
mock_bulk_news.side_effect = NewsUnavailableError("All news vendors failed")
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
result = graph.discover_trending()
assert result.status == DiscoveryStatus.FAILED
assert result.error_message is not None
assert "news" in result.error_message.lower() or "vendor" in result.error_message.lower()
class TestTimeoutHandlingWithPartialResults:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
def test_timeout_handling_returns_error(self, mock_bulk_news):
def slow_fetch(*args, **kwargs):
import time
time.sleep(0.3)
return []
mock_bulk_news.side_effect = slow_fetch
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 0.1,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
with pytest.raises(DiscoveryTimeoutError):
graph.discover_trending()
class TestNoTrendingStocksFound:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_no_trending_stocks_found_returns_empty_list(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = [
NewsArticle(
title="General market update",
source="Reuters",
url="https://reuters.com/general",
published_at=datetime.now(),
content_snippet="Markets were quiet today with no major news.",
ticker_mentions=[],
),
]
mock_extract.return_value = []
mock_scores.return_value = []
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
result = graph.discover_trending()
assert result.status == DiscoveryStatus.COMPLETED
assert len(result.trending_stocks) == 0
class TestAllStocksFilteredOutBySectorFilter:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_all_stocks_filtered_out_by_sector_filter(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = []
mock_extract.return_value = []
mock_scores.return_value = [
TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=10.0,
mention_count=5,
sentiment=0.5,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Apple earnings",
source_articles=[],
),
TrendingStock(
ticker="MSFT",
company_name="Microsoft",
score=9.0,
mention_count=4,
sentiment=0.4,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="Microsoft product",
source_articles=[],
),
]
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.HEALTHCARE],
)
result = graph.discover_trending(request)
assert result.status == DiscoveryStatus.COMPLETED
assert len(result.trending_stocks) == 0
class TestAllStocksFilteredOutByEventFilter:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_all_stocks_filtered_out_by_event_filter(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = []
mock_extract.return_value = []
mock_scores.return_value = [
TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=10.0,
mention_count=5,
sentiment=0.5,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Apple earnings",
source_articles=[],
),
]
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
request = DiscoveryRequest(
lookback_period="24h",
event_filter=[EventCategory.MERGER_ACQUISITION],
)
result = graph.discover_trending(request)
assert result.status == DiscoveryStatus.COMPLETED
assert len(result.trending_stocks) == 0
class TestMultipleSectorsAndEventsFiltering:
@patch("tradingagents.graph.trading_graph.get_bulk_news")
@patch("tradingagents.graph.trading_graph.extract_entities")
@patch("tradingagents.graph.trading_graph.calculate_trending_scores")
def test_combined_sector_and_event_filtering(
self, mock_scores, mock_extract, mock_bulk_news
):
mock_bulk_news.return_value = []
mock_extract.return_value = []
mock_scores.return_value = [
TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=10.0,
mention_count=5,
sentiment=0.5,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Apple earnings",
source_articles=[],
),
TrendingStock(
ticker="JPM",
company_name="JPMorgan Chase",
score=9.0,
mention_count=4,
sentiment=0.4,
sector=Sector.FINANCE,
event_type=EventCategory.EARNINGS,
news_summary="JPM earnings",
source_articles=[],
),
TrendingStock(
ticker="XOM",
company_name="Exxon Mobil",
score=8.0,
mention_count=3,
sentiment=0.3,
sector=Sector.ENERGY,
event_type=EventCategory.REGULATORY,
news_summary="XOM regulatory news",
source_articles=[],
),
]
from tradingagents.graph.trading_graph import TradingAgentsGraph
with patch.object(TradingAgentsGraph, "__init__", lambda self, **kwargs: None):
graph = TradingAgentsGraph()
graph.config = {
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY, Sector.FINANCE],
event_filter=[EventCategory.EARNINGS],
)
result = graph.discover_trending(request)
assert result.status == DiscoveryStatus.COMPLETED
assert len(result.trending_stocks) == 2
tickers = [s.ticker for s in result.trending_stocks]
assert "AAPL" in tickers
assert "JPM" in tickers
assert "XOM" not in tickers
class TestDiscoveryResultPersistenceIntegration:
def test_discovery_result_can_be_serialized_and_saved(self):
from tradingagents.agents.discovery.persistence import (
save_discovery_result,
generate_markdown_summary,
)
import tempfile
import shutil
from pathlib import Path
article = NewsArticle(
title="Test article",
source="Test",
url="https://test.com",
published_at=datetime.now(),
content_snippet="Test content",
ticker_mentions=["TEST"],
)
stock = TrendingStock(
ticker="TEST",
company_name="Test Company",
score=5.0,
mention_count=2,
sentiment=0.5,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.OTHER,
news_summary="Test news summary",
source_articles=[article],
)
request = DiscoveryRequest(
lookback_period="24h",
created_at=datetime.now(),
)
result = DiscoveryResult(
request=request,
trending_stocks=[stock],
status=DiscoveryStatus.COMPLETED,
started_at=datetime.now(),
completed_at=datetime.now(),
)
temp_dir = tempfile.mkdtemp()
try:
path = save_discovery_result(result, base_path=Path(temp_dir))
assert path.exists()
assert (path / "discovery_result.json").exists()
assert (path / "discovery_summary.md").exists()
markdown = generate_markdown_summary(result)
assert "TEST" in markdown
assert "Test Company" in markdown
finally:
shutil.rmtree(temp_dir)

196
tests/discovery/test_models.py Normal file
View File

@ -0,0 +1,196 @@
import pytest
from datetime import datetime
from tradingagents.agents.discovery import (
TrendingStock,
NewsArticle,
DiscoveryRequest,
DiscoveryResult,
Sector,
EventCategory,
)
from tradingagents.agents.discovery.models import DiscoveryStatus
class TestTrendingStock:
def test_trending_stock_creation_and_validation(self):
article = NewsArticle(
title="Apple announces new iPhone",
source="Reuters",
url="https://reuters.com/article1",
published_at=datetime(2024, 1, 15, 10, 30, 0),
content_snippet="Apple Inc announced its latest iPhone model today...",
ticker_mentions=["AAPL"],
)
stock = TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=85.5,
mention_count=10,
sentiment=0.75,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="Apple announced new iPhone with advanced AI features.",
source_articles=[article],
)
assert stock.ticker == "AAPL"
assert stock.company_name == "Apple Inc."
assert stock.score == 85.5
assert stock.mention_count == 10
assert stock.sentiment == 0.75
assert stock.sector == Sector.TECHNOLOGY
assert stock.event_type == EventCategory.PRODUCT_LAUNCH
assert len(stock.source_articles) == 1
class TestNewsArticle:
def test_news_article_with_required_fields(self):
published = datetime(2024, 1, 15, 14, 0, 0)
article = NewsArticle(
title="Tesla Q4 Earnings Beat Expectations",
source="Bloomberg",
url="https://bloomberg.com/news/tsla-earnings",
published_at=published,
content_snippet="Tesla Inc. reported fourth quarter earnings that exceeded analyst expectations...",
ticker_mentions=["TSLA", "F"],
)
assert article.title == "Tesla Q4 Earnings Beat Expectations"
assert article.source == "Bloomberg"
assert article.url == "https://bloomberg.com/news/tsla-earnings"
assert article.published_at == published
assert article.content_snippet.startswith("Tesla Inc.")
assert "TSLA" in article.ticker_mentions
assert "F" in article.ticker_mentions
class TestDiscoveryRequest:
def test_discovery_request_with_lookback_period_validation(self):
created = datetime(2024, 1, 15, 12, 0, 0)
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY, Sector.HEALTHCARE],
event_filter=[EventCategory.EARNINGS],
max_results=20,
created_at=created,
)
assert request.lookback_period == "24h"
assert Sector.TECHNOLOGY in request.sector_filter
assert Sector.HEALTHCARE in request.sector_filter
assert EventCategory.EARNINGS in request.event_filter
assert request.max_results == 20
assert request.created_at == created
def test_discovery_request_with_defaults(self):
request = DiscoveryRequest(
lookback_period="1h",
)
assert request.lookback_period == "1h"
assert request.sector_filter is None
assert request.event_filter is None
assert request.max_results == 20
assert request.created_at is not None
class TestDiscoveryResult:
def test_discovery_result_state_transitions(self):
request = DiscoveryRequest(lookback_period="6h")
started = datetime(2024, 1, 15, 12, 0, 0)
result = DiscoveryResult(
request=request,
trending_stocks=[],
status=DiscoveryStatus.CREATED,
started_at=started,
)
assert result.status == DiscoveryStatus.CREATED
result.status = DiscoveryStatus.PROCESSING
assert result.status == DiscoveryStatus.PROCESSING
result.status = DiscoveryStatus.COMPLETED
result.completed_at = datetime(2024, 1, 15, 12, 1, 0)
assert result.status == DiscoveryStatus.COMPLETED
assert result.completed_at is not None
def test_discovery_result_failed_state(self):
request = DiscoveryRequest(lookback_period="7d")
result = DiscoveryResult(
request=request,
trending_stocks=[],
status=DiscoveryStatus.FAILED,
started_at=datetime(2024, 1, 15, 12, 0, 0),
error_message="News API rate limit exceeded",
)
assert result.status == DiscoveryStatus.FAILED
assert result.error_message == "News API rate limit exceeded"
class TestSerializationRoundtrip:
def test_to_dict_and_from_dict_serialization_roundtrip(self):
article = NewsArticle(
title="Microsoft acquires AI startup",
source="WSJ",
url="https://wsj.com/msft-acquisition",
published_at=datetime(2024, 1, 15, 9, 0, 0),
content_snippet="Microsoft Corp announced the acquisition of an AI startup...",
ticker_mentions=["MSFT"],
)
stock = TrendingStock(
ticker="MSFT",
company_name="Microsoft Corporation",
score=92.3,
mention_count=15,
sentiment=0.65,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.MERGER_ACQUISITION,
news_summary="Microsoft announces major AI acquisition.",
source_articles=[article],
)
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY],
event_filter=[EventCategory.MERGER_ACQUISITION],
max_results=10,
created_at=datetime(2024, 1, 15, 8, 0, 0),
)
result = DiscoveryResult(
request=request,
trending_stocks=[stock],
status=DiscoveryStatus.COMPLETED,
started_at=datetime(2024, 1, 15, 8, 0, 0),
completed_at=datetime(2024, 1, 15, 8, 1, 30),
)
result_dict = result.to_dict()
restored_result = DiscoveryResult.from_dict(result_dict)
assert restored_result.status == result.status
assert restored_result.request.lookback_period == request.lookback_period
assert len(restored_result.trending_stocks) == 1
restored_stock = restored_result.trending_stocks[0]
assert restored_stock.ticker == stock.ticker
assert restored_stock.company_name == stock.company_name
assert restored_stock.score == stock.score
assert restored_stock.mention_count == stock.mention_count
assert restored_stock.sentiment == stock.sentiment
assert restored_stock.sector == stock.sector
assert restored_stock.event_type == stock.event_type
assert len(restored_stock.source_articles) == 1
restored_article = restored_stock.source_articles[0]
assert restored_article.title == article.title
assert restored_article.source == article.source
assert restored_article.url == article.url

228
tests/discovery/test_persistence.py Normal file
View File

@ -0,0 +1,228 @@
import pytest
import json
from datetime import datetime
from pathlib import Path
import tempfile
import shutil
from tradingagents.agents.discovery import (
TrendingStock,
NewsArticle,
DiscoveryRequest,
DiscoveryResult,
DiscoveryStatus,
Sector,
EventCategory,
)
from tradingagents.agents.discovery.persistence import (
save_discovery_result,
generate_markdown_summary,
)
@pytest.fixture
def sample_discovery_result():
articles = [
NewsArticle(
title="Apple announces new iPhone with AI features",
source="Reuters",
url="https://reuters.com/apple-iphone-ai",
published_at=datetime(2024, 1, 15, 10, 30, 0),
content_snippet="Apple Inc announced its latest iPhone model with advanced AI...",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="Apple stock surges on earnings beat",
source="Bloomberg",
url="https://bloomberg.com/apple-earnings",
published_at=datetime(2024, 1, 15, 11, 0, 0),
content_snippet="Shares of Apple Inc surged after the company reported...",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="Microsoft cloud revenue grows 25%",
source="WSJ",
url="https://wsj.com/msft-cloud",
published_at=datetime(2024, 1, 15, 9, 0, 0),
content_snippet="Microsoft Corp reported strong cloud revenue growth...",
ticker_mentions=["MSFT"],
),
]
stocks = [
TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=8.54,
mention_count=12,
sentiment=0.72,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Apple reported strong earnings and announced new AI features.",
source_articles=[articles[0], articles[1]],
),
TrendingStock(
ticker="MSFT",
company_name="Microsoft Corporation",
score=7.23,
mention_count=9,
sentiment=0.65,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="Microsoft cloud business continues strong growth.",
source_articles=[articles[2]],
),
TrendingStock(
ticker="GOOGL",
company_name="Alphabet Inc.",
score=6.15,
mention_count=7,
sentiment=0.58,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.REGULATORY,
news_summary="Google faces regulatory scrutiny in multiple markets.",
source_articles=[],
),
]
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY],
event_filter=[EventCategory.EARNINGS],
max_results=20,
created_at=datetime(2024, 1, 15, 14, 30, 45),
)
return DiscoveryResult(
request=request,
trending_stocks=stocks,
status=DiscoveryStatus.COMPLETED,
started_at=datetime(2024, 1, 15, 14, 30, 45),
completed_at=datetime(2024, 1, 15, 14, 31, 30),
)
@pytest.fixture
def temp_results_dir():
temp_dir = tempfile.mkdtemp()
yield Path(temp_dir)
shutil.rmtree(temp_dir)
class TestDirectoryStructureCreation:
def test_creates_correct_directory_structure(self, sample_discovery_result, temp_results_dir):
result_path = save_discovery_result(sample_discovery_result, base_path=temp_results_dir)
assert result_path.exists()
assert result_path.is_dir()
path_parts = result_path.parts
assert "discovery" in path_parts
date_part = path_parts[-2]
time_part = path_parts[-1]
assert len(date_part.split("-")) == 3
assert len(time_part.split("-")) == 3
class TestDiscoveryResultJson:
def test_discovery_result_json_contains_all_fields(self, sample_discovery_result, temp_results_dir):
result_path = save_discovery_result(sample_discovery_result, base_path=temp_results_dir)
json_path = result_path / "discovery_result.json"
assert json_path.exists()
with open(json_path, "r") as f:
saved_data = json.load(f)
assert "request" in saved_data
assert "trending_stocks" in saved_data
assert "status" in saved_data
assert "started_at" in saved_data
assert "completed_at" in saved_data
assert saved_data["request"]["lookback_period"] == "24h"
assert saved_data["status"] == "completed"
assert len(saved_data["trending_stocks"]) == 3
first_stock = saved_data["trending_stocks"][0]
assert first_stock["ticker"] == "AAPL"
assert first_stock["company_name"] == "Apple Inc."
assert first_stock["score"] == 8.54
assert first_stock["mention_count"] == 12
assert first_stock["sentiment"] == 0.72
assert first_stock["sector"] == "technology"
assert first_stock["event_type"] == "earnings"
assert "news_summary" in first_stock
assert "source_articles" in first_stock
class TestDiscoverySummaryMarkdown:
def test_discovery_summary_md_is_human_readable(self, sample_discovery_result, temp_results_dir):
result_path = save_discovery_result(sample_discovery_result, base_path=temp_results_dir)
md_path = result_path / "discovery_summary.md"
assert md_path.exists()
with open(md_path, "r") as f:
markdown_content = f.read()
assert "# Discovery Results" in markdown_content
assert "Timestamp:" in markdown_content
assert "Lookback Period:" in markdown_content
assert "24h" in markdown_content
assert "Total Stocks Found:" in markdown_content
assert "## Trending Stocks" in markdown_content
assert "| Rank |" in markdown_content
assert "| Ticker |" in markdown_content
assert "| Company |" in markdown_content
assert "| Score |" in markdown_content
assert "| Mentions |" in markdown_content
assert "| Event |" in markdown_content
assert "AAPL" in markdown_content
assert "Apple Inc." in markdown_content
assert "8.54" in markdown_content
assert "12" in markdown_content
assert "earnings" in markdown_content
assert "MSFT" in markdown_content
assert "Microsoft Corporation" in markdown_content
assert "## Top 3 Detailed Analysis" in markdown_content
assert "### 1. AAPL - Apple Inc." in markdown_content
assert "**Score:**" in markdown_content
assert "**Sentiment:**" in markdown_content
assert "**Sector:**" in markdown_content
assert "**Event Type:**" in markdown_content
assert "**Mentions:**" in markdown_content
assert "**News Summary:**" in markdown_content
class TestMarkdownGeneration:
def test_generate_markdown_with_filters(self, sample_discovery_result):
markdown = generate_markdown_summary(sample_discovery_result)
assert "sector=technology" in markdown.lower()
assert "event=earnings" in markdown.lower()
def test_generate_markdown_without_filters(self):
request = DiscoveryRequest(
lookback_period="6h",
created_at=datetime(2024, 1, 15, 10, 0, 0),
)
result = DiscoveryResult(
request=request,
trending_stocks=[],
status=DiscoveryStatus.COMPLETED,
started_at=datetime(2024, 1, 15, 10, 0, 0),
completed_at=datetime(2024, 1, 15, 10, 1, 0),
)
markdown = generate_markdown_summary(result)
assert "Filters:" in markdown
assert "None" in markdown

469
tests/discovery/test_scorer.py Normal file
View File

@ -0,0 +1,469 @@
import pytest
import math
from datetime import datetime, timedelta
from unittest.mock import patch
from tradingagents.agents.discovery import NewsArticle, EventCategory, Sector
from tradingagents.agents.discovery.entity_extractor import EntityMention
class TestFrequencyCalculation:
def test_frequency_calculation_unique_article_count(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
articles = [
NewsArticle(
title="Apple Q4 Earnings",
source="Reuters",
url="https://reuters.com/article1",
published_at=now - timedelta(hours=1),
content_snippet="Apple Inc reported strong earnings.",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="Apple iPhone Sales",
source="Bloomberg",
url="https://bloomberg.com/article2",
published_at=now - timedelta(hours=2),
content_snippet="Apple saw record iPhone sales.",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="Apple AI Features",
source="WSJ",
url="https://wsj.com/article3",
published_at=now - timedelta(hours=3),
content_snippet="Apple announced AI features.",
ticker_mentions=["AAPL"],
),
]
mentions = [
EntityMention(
company_name="Apple Inc",
confidence=0.95,
context_snippet="Apple Inc reported strong earnings",
article_id="article_0",
event_type=EventCategory.EARNINGS,
),
EntityMention(
company_name="Apple",
confidence=0.90,
context_snippet="Apple saw record iPhone sales",
article_id="article_1",
event_type=EventCategory.EARNINGS,
),
EntityMention(
company_name="Apple Inc.",
confidence=0.92,
context_snippet="Apple announced AI features",
article_id="article_2",
event_type=EventCategory.PRODUCT_LAUNCH,
),
]
with patch(
"tradingagents.agents.discovery.scorer.resolve_ticker"
) as mock_resolve:
mock_resolve.return_value = "AAPL"
with patch(
"tradingagents.agents.discovery.scorer.classify_sector"
) as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(mentions, articles)
assert len(result) == 1
assert result[0].ticker == "AAPL"
assert result[0].mention_count == 3
class TestSentimentIntensityFactor:
def test_sentiment_intensity_uses_absolute_value(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
articles = [
NewsArticle(
title="Stock drops sharply",
source="Reuters",
url="https://reuters.com/article1",
published_at=now - timedelta(hours=1),
content_snippet="Company faced major issues.",
ticker_mentions=["TSLA"],
),
NewsArticle(
title="More bad news",
source="Bloomberg",
url="https://bloomberg.com/article2",
published_at=now - timedelta(hours=2),
content_snippet="Further decline expected.",
ticker_mentions=["TSLA"],
),
]
mentions = [
EntityMention(
company_name="Tesla",
confidence=0.95,
context_snippet="Company faced major issues",
article_id="article_0",
event_type=EventCategory.OTHER,
sentiment=-0.8,
),
EntityMention(
company_name="Tesla Inc",
confidence=0.90,
context_snippet="Further decline expected",
article_id="article_1",
event_type=EventCategory.OTHER,
sentiment=-0.6,
),
]
with patch(
"tradingagents.agents.discovery.scorer.resolve_ticker"
) as mock_resolve:
mock_resolve.return_value = "TSLA"
with patch(
"tradingagents.agents.discovery.scorer.classify_sector"
) as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(mentions, articles)
assert len(result) == 1
assert result[0].sentiment < 0
expected_sentiment = (-0.8 * 0.95 + -0.6 * 0.90) / (0.95 + 0.90)
assert abs(result[0].sentiment - expected_sentiment) < 0.01
class TestRecencyWeightExponentialDecay:
def test_recency_weight_exponential_decay(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
articles = [
NewsArticle(
title="Recent news",
source="Reuters",
url="https://reuters.com/article1",
published_at=now - timedelta(hours=1),
content_snippet="Recent company news.",
ticker_mentions=["NVDA"],
),
NewsArticle(
title="Older news",
source="Bloomberg",
url="https://bloomberg.com/article2",
published_at=now - timedelta(hours=10),
content_snippet="Older company news.",
ticker_mentions=["NVDA"],
),
]
mentions = [
EntityMention(
company_name="Nvidia",
confidence=0.90,
context_snippet="Recent company news",
article_id="article_0",
event_type=EventCategory.OTHER,
sentiment=0.5,
),
EntityMention(
company_name="Nvidia",
confidence=0.90,
context_snippet="Older company news",
article_id="article_1",
event_type=EventCategory.OTHER,
sentiment=0.5,
),
]
with patch(
"tradingagents.agents.discovery.scorer.resolve_ticker"
) as mock_resolve:
mock_resolve.return_value = "NVDA"
with patch(
"tradingagents.agents.discovery.scorer.classify_sector"
) as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(mentions, articles, decay_rate=0.1)
assert len(result) == 1
recent_weight = math.exp(-0.1 * 1)
older_weight = math.exp(-0.1 * 10)
avg_recency = (recent_weight + older_weight) / 2
assert result[0].score > 0
class TestMinimumThresholdFiltering:
def test_minimum_threshold_filtering_requires_two_articles(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
articles = [
NewsArticle(
title="Single mention stock",
source="Reuters",
url="https://reuters.com/article1",
published_at=now - timedelta(hours=1),
content_snippet="Some company news.",
ticker_mentions=["AMD"],
),
NewsArticle(
title="Multiple mention stock 1",
source="Bloomberg",
url="https://bloomberg.com/article2",
published_at=now - timedelta(hours=2),
content_snippet="Popular company news.",
ticker_mentions=["MSFT"],
),
NewsArticle(
title="Multiple mention stock 2",
source="WSJ",
url="https://wsj.com/article3",
published_at=now - timedelta(hours=3),
content_snippet="More popular company news.",
ticker_mentions=["MSFT"],
),
]
mentions = [
EntityMention(
company_name="AMD",
confidence=0.90,
context_snippet="Some company news",
article_id="article_0",
event_type=EventCategory.OTHER,
),
EntityMention(
company_name="Microsoft",
confidence=0.95,
context_snippet="Popular company news",
article_id="article_1",
event_type=EventCategory.OTHER,
),
EntityMention(
company_name="Microsoft Corp",
confidence=0.92,
context_snippet="More popular company news",
article_id="article_2",
event_type=EventCategory.OTHER,
),
]
with patch(
"tradingagents.agents.discovery.scorer.resolve_ticker"
) as mock_resolve:
def resolve_side_effect(name):
if "AMD" in name or name == "AMD":
return "AMD"
return "MSFT"
mock_resolve.side_effect = resolve_side_effect
with patch(
"tradingagents.agents.discovery.scorer.classify_sector"
) as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(mentions, articles, min_mentions=2)
assert len(result) == 1
assert result[0].ticker == "MSFT"
assert all(stock.mention_count >= 2 for stock in result)
class TestFinalScoreFormulaCorrectness:
def test_final_score_formula_correctness(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
hours_old = 2.0
articles = [
NewsArticle(
title="Test article 1",
source="Reuters",
url="https://reuters.com/article1",
published_at=now - timedelta(hours=hours_old),
content_snippet="Google announced results.",
ticker_mentions=["GOOGL"],
),
NewsArticle(
title="Test article 2",
source="Bloomberg",
url="https://bloomberg.com/article2",
published_at=now - timedelta(hours=hours_old),
content_snippet="Alphabet earnings beat.",
ticker_mentions=["GOOGL"],
),
]
sentiment_val = 0.6
confidence = 0.9
mentions = [
EntityMention(
company_name="Google",
confidence=confidence,
context_snippet="Google announced results",
article_id="article_0",
event_type=EventCategory.EARNINGS,
sentiment=sentiment_val,
),
EntityMention(
company_name="Alphabet",
confidence=confidence,
context_snippet="Alphabet earnings beat",
article_id="article_1",
event_type=EventCategory.EARNINGS,
sentiment=sentiment_val,
),
]
decay_rate = 0.1
with patch(
"tradingagents.agents.discovery.scorer.resolve_ticker"
) as mock_resolve:
mock_resolve.return_value = "GOOGL"
with patch(
"tradingagents.agents.discovery.scorer.classify_sector"
) as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(
mentions, articles, decay_rate=decay_rate
)
assert len(result) == 1
stock = result[0]
frequency = 2
sentiment_factor = 1 + abs(sentiment_val)
recency_weight = math.exp(-decay_rate * hours_old)
expected_score = frequency * sentiment_factor * recency_weight
assert abs(stock.score - expected_score) < 0.01
class TestSortingByScoreDescending:
def test_results_sorted_by_score_descending(self):
from tradingagents.agents.discovery.scorer import calculate_trending_scores
now = datetime.now()
articles = [
NewsArticle(
title="High score stock 1",
source="Reuters",
url="https://reuters.com/article1",
published_at=now - timedelta(hours=1),
content_snippet="Apple news.",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="High score stock 2",
source="Bloomberg",
url="https://bloomberg.com/article2",
published_at=now - timedelta(hours=1),
content_snippet="More Apple news.",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="High score stock 3",
source="WSJ",
url="https://wsj.com/article3",
published_at=now - timedelta(hours=1),
content_snippet="Even more Apple news.",
ticker_mentions=["AAPL"],
),
NewsArticle(
title="Low score stock 1",
source="CNBC",
url="https://cnbc.com/article4",
published_at=now - timedelta(hours=10),
content_snippet="Tesla news.",
ticker_mentions=["TSLA"],
),
NewsArticle(
title="Low score stock 2",
source="FT",
url="https://ft.com/article5",
published_at=now - timedelta(hours=10),
content_snippet="More Tesla news.",
ticker_mentions=["TSLA"],
),
]
mentions = [
EntityMention(
company_name="Apple",
confidence=0.95,
context_snippet="Apple news",
article_id="article_0",
event_type=EventCategory.OTHER,
sentiment=0.8,
),
EntityMention(
company_name="Apple Inc",
confidence=0.93,
context_snippet="More Apple news",
article_id="article_1",
event_type=EventCategory.OTHER,
sentiment=0.8,
),
EntityMention(
company_name="Apple",
confidence=0.90,
context_snippet="Even more Apple news",
article_id="article_2",
event_type=EventCategory.OTHER,
sentiment=0.8,
),
EntityMention(
company_name="Tesla",
confidence=0.85,
context_snippet="Tesla news",
article_id="article_3",
event_type=EventCategory.OTHER,
sentiment=0.2,
),
EntityMention(
company_name="Tesla Inc",
confidence=0.85,
context_snippet="More Tesla news",
article_id="article_4",
event_type=EventCategory.OTHER,
sentiment=0.2,
),
]
with patch(
"tradingagents.agents.discovery.scorer.resolve_ticker"
) as mock_resolve:
def resolve_side_effect(name):
if "Apple" in name:
return "AAPL"
if "Tesla" in name:
return "TSLA"
return None
mock_resolve.side_effect = resolve_side_effect
with patch(
"tradingagents.agents.discovery.scorer.classify_sector"
) as mock_sector:
mock_sector.return_value = "technology"
result = calculate_trending_scores(mentions, articles, min_mentions=2)
assert len(result) == 2
for i in range(len(result) - 1):
assert result[i].score >= result[i + 1].score

94
tests/discovery/test_sector_classifier.py Normal file
View File

@ -0,0 +1,94 @@
import pytest
from unittest.mock import patch, MagicMock
from tradingagents.dataflows.trending.sector_classifier import (
classify_sector,
TICKER_TO_SECTOR,
VALID_SECTORS,
_llm_classify_sector,
_sector_cache,
)
class TestStaticSectorMapping:
def test_static_sector_mapping_for_known_technology_tickers(self):
assert classify_sector("AAPL") == "technology"
assert classify_sector("MSFT") == "technology"
assert classify_sector("GOOGL") == "technology"
assert classify_sector("NVDA") == "technology"
def test_static_sector_mapping_for_known_healthcare_tickers(self):
assert classify_sector("JNJ") == "healthcare"
assert classify_sector("PFE") == "healthcare"
assert classify_sector("UNH") == "healthcare"
def test_static_sector_mapping_for_known_finance_tickers(self):
assert classify_sector("JPM") == "finance"
assert classify_sector("BAC") == "finance"
assert classify_sector("GS") == "finance"
def test_static_sector_mapping_for_known_energy_tickers(self):
assert classify_sector("XOM") == "energy"
assert classify_sector("CVX") == "energy"
assert classify_sector("COP") == "energy"
def test_static_sector_mapping_case_insensitive(self):
assert classify_sector("aapl") == "technology"
assert classify_sector("AAPL") == "technology"
assert classify_sector("Aapl") == "technology"
class TestLLMFallback:
@patch("tradingagents.dataflows.trending.sector_classifier._llm_classify_sector")
def test_llm_fallback_for_unknown_tickers(self, mock_llm_classify):
mock_llm_classify.return_value = "technology"
_sector_cache.clear()
result = classify_sector("UNKNOWNTICKER123")
mock_llm_classify.assert_called_once_with("UNKNOWNTICKER123")
assert result == "technology"
@patch("tradingagents.dataflows.trending.sector_classifier._llm_classify_sector")
def test_llm_fallback_caches_results(self, mock_llm_classify):
mock_llm_classify.return_value = "healthcare"
_sector_cache.clear()
result1 = classify_sector("NEWCO123")
result2 = classify_sector("NEWCO123")
assert mock_llm_classify.call_count == 1
assert result1 == "healthcare"
assert result2 == "healthcare"
@patch("tradingagents.dataflows.trending.sector_classifier._llm_classify_sector")
def test_llm_fallback_returns_other_on_error(self, mock_llm_classify):
mock_llm_classify.side_effect = Exception("LLM error")
_sector_cache.clear()
result = classify_sector("ERRORCO")
assert result == "other"
class TestAllSectorCategories:
def test_all_sector_categories_in_valid_sectors(self):
expected_sectors = {
"technology",
"healthcare",
"finance",
"energy",
"consumer_goods",
"industrials",
"other",
}
assert VALID_SECTORS == expected_sectors
def test_static_mapping_covers_all_sector_categories(self):
sectors_in_mapping = set(TICKER_TO_SECTOR.values())
assert sectors_in_mapping.issubset(VALID_SECTORS)
def test_classify_sector_always_returns_valid_sector(self):
test_tickers = ["AAPL", "JPM", "XOM", "JNJ", "WMT", "CAT"]
for ticker in test_tickers:
result = classify_sector(ticker)
assert result in VALID_SECTORS

135
tests/discovery/test_stock_resolver.py Normal file
View File

@ -0,0 +1,135 @@
import pytest
import logging
from unittest.mock import patch, MagicMock
from tradingagents.dataflows.trending.stock_resolver import (
resolve_ticker,
validate_us_ticker,
_normalize_company_name,
_search_yfinance_ticker,
)
class TestStaticLookup:
def test_static_lookup_for_known_companies(self):
assert resolve_ticker("Apple") == "AAPL"
assert resolve_ticker("Microsoft") == "MSFT"
assert resolve_ticker("Google") == "GOOGL"
assert resolve_ticker("Amazon") == "AMZN"
assert resolve_ticker("Tesla") == "TSLA"
assert resolve_ticker("Nvidia") == "NVDA"
def test_static_lookup_case_insensitive(self):
assert resolve_ticker("APPLE") == "AAPL"
assert resolve_ticker("apple") == "AAPL"
assert resolve_ticker("ApPlE") == "AAPL"
assert resolve_ticker("microsoft") == "MSFT"
assert resolve_ticker("MICROSOFT") == "MSFT"
class TestNameVariationHandling:
def test_name_variation_handling_with_suffixes(self):
assert resolve_ticker("Apple Inc.") == "AAPL"
assert resolve_ticker("Apple Inc") == "AAPL"
assert resolve_ticker("Apple Corporation") == "AAPL"
assert resolve_ticker("Microsoft Corp.") == "MSFT"
assert resolve_ticker("Microsoft Corp") == "MSFT"
assert resolve_ticker("Tesla Inc") == "TSLA"
def test_name_variation_handling_informal_names(self):
assert resolve_ticker("the iPhone maker") == "AAPL"
assert resolve_ticker("iPhone maker") == "AAPL"
assert resolve_ticker("the search giant") == "GOOGL"
assert resolve_ticker("the e-commerce giant") == "AMZN"
assert resolve_ticker("EV maker Tesla") == "TSLA"
def test_name_variation_handling_alternate_names(self):
assert resolve_ticker("Alphabet") == "GOOGL"
assert resolve_ticker("Meta") == "META"
assert resolve_ticker("Facebook") == "META"
assert resolve_ticker("Meta Platforms") == "META"
class TestYfinanceFallback:
@patch("tradingagents.dataflows.trending.stock_resolver._search_yfinance_ticker")
@patch("tradingagents.dataflows.trending.stock_resolver.validate_us_ticker")
def test_yfinance_fallback_for_unknown_company(self, mock_validate, mock_search):
mock_search.return_value = "PLTR"
mock_validate.return_value = True
result = resolve_ticker("UnknownTechStartupXYZ")
mock_search.assert_called_once()
assert result == "PLTR"
@patch("tradingagents.dataflows.trending.stock_resolver._search_yfinance_ticker")
def test_yfinance_fallback_returns_none_when_not_found(self, mock_search):
mock_search.return_value = None
result = resolve_ticker("NonexistentCompanyXYZ123")
assert result is None
class TestUSExchangeValidation:
@patch("tradingagents.dataflows.trending.stock_resolver.yf.Ticker")
def test_validate_us_ticker_accepts_nyse(self, mock_ticker):
mock_info = {"exchange": "NYQ"}
mock_ticker.return_value.info = mock_info
assert validate_us_ticker("IBM") is True
@patch("tradingagents.dataflows.trending.stock_resolver.yf.Ticker")
def test_validate_us_ticker_accepts_nasdaq(self, mock_ticker):
mock_info = {"exchange": "NMS"}
mock_ticker.return_value.info = mock_info
assert validate_us_ticker("AAPL") is True
@patch("tradingagents.dataflows.trending.stock_resolver.yf.Ticker")
def test_validate_us_ticker_accepts_amex(self, mock_ticker):
mock_info = {"exchange": "ASE"}
mock_ticker.return_value.info = mock_info
assert validate_us_ticker("SPY") is True
@patch("tradingagents.dataflows.trending.stock_resolver.yf.Ticker")
def test_validate_us_ticker_rejects_international(self, mock_ticker):
mock_info = {"exchange": "LSE"}
mock_ticker.return_value.info = mock_info
assert validate_us_ticker("VOD.L") is False
@patch("tradingagents.dataflows.trending.stock_resolver.yf.Ticker")
def test_validate_us_ticker_rejects_otc(self, mock_ticker):
mock_info = {"exchange": "PNK"}
mock_ticker.return_value.info = mock_info
assert validate_us_ticker("OTCPK") is False
class TestAmbiguousResolutionLogging:
def test_ambiguous_resolution_logs_multiple_matches(self, caplog):
with caplog.at_level(logging.INFO, logger="tradingagents.dataflows.trending.stock_resolver"):
pass
@patch("tradingagents.dataflows.trending.stock_resolver._search_yfinance_ticker")
@patch("tradingagents.dataflows.trending.stock_resolver.validate_us_ticker")
def test_yfinance_fallback_is_logged(self, mock_validate, mock_search, caplog):
mock_search.return_value = "RBLX"
mock_validate.return_value = True
with caplog.at_level(logging.INFO, logger="tradingagents.dataflows.trending.stock_resolver"):
result = resolve_ticker("SomeRandomCompanyNotInMapping")
assert any("fallback" in record.message.lower() or "yfinance" in record.message.lower()
for record in caplog.records)
@patch("tradingagents.dataflows.trending.stock_resolver.yf.Ticker")
def test_validation_failure_is_logged(self, mock_ticker, caplog):
mock_info = {"exchange": "LSE"}
mock_ticker.return_value.info = mock_info
with caplog.at_level(logging.WARNING, logger="tradingagents.dataflows.trending.stock_resolver"):
result = validate_us_ticker("VOD.L")
assert result is False

0
tests/graph/__init__.py Normal file
View File

527
tests/graph/test_trading_graph.py Normal file
View File

@ -0,0 +1,527 @@
import pytest
from unittest.mock import Mock, patch, MagicMock
from datetime import datetime, date
from tradingagents.graph.trading_graph import TradingAgentsGraph, DiscoveryTimeoutException
from tradingagents.agents.discovery import (
DiscoveryRequest,
DiscoveryResult,
DiscoveryStatus,
TrendingStock,
Sector,
EventCategory,
NewsArticle,
)
class TestTradingAgentsGraphInit:
"""Test suite for TradingAgentsGraph initialization."""
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_with_default_config(self, mock_setup, mock_memory, mock_llm):
"""Test initialization with default configuration."""
graph = TradingAgentsGraph(debug=False)
assert graph.debug == False
assert graph.config is not None
assert "llm_provider" in graph.config
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_with_custom_config(self, mock_setup, mock_memory, mock_llm):
"""Test initialization with custom configuration."""
custom_config = {
"llm_provider": "openai",
"deep_think_llm": "gpt-4",
"quick_think_llm": "gpt-3.5-turbo",
"backend_url": "https://api.openai.com/v1",
"max_debate_rounds": 3,
"max_risk_discuss_rounds": 2,
"max_recur_limit": 100,
"project_dir": "/tmp/test",
"data_vendors": {},
"tool_vendors": {},
}
graph = TradingAgentsGraph(debug=True, config=custom_config)
assert graph.config["llm_provider"] == "openai"
assert graph.config["max_debate_rounds"] == 3
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_with_anthropic_provider(self, mock_setup, mock_memory, mock_llm):
"""Test initialization with Anthropic provider."""
with patch('tradingagents.graph.trading_graph.ChatAnthropic') as mock_anthropic:
config = {
"llm_provider": "anthropic",
"deep_think_llm": "claude-3-opus",
"quick_think_llm": "claude-3-haiku",
"backend_url": "https://api.anthropic.com",
"project_dir": "/tmp/test",
"data_vendors": {},
"tool_vendors": {},
"max_debate_rounds": 2,
"max_risk_discuss_rounds": 2,
"max_recur_limit": 100,
}
graph = TradingAgentsGraph(config=config)
assert mock_anthropic.called
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_with_google_provider(self, mock_setup, mock_memory, mock_llm):
"""Test initialization with Google provider."""
with patch('tradingagents.graph.trading_graph.ChatGoogleGenerativeAI') as mock_google:
config = {
"llm_provider": "google",
"deep_think_llm": "gemini-pro",
"quick_think_llm": "gemini-pro",
"project_dir": "/tmp/test",
"data_vendors": {},
"tool_vendors": {},
"max_debate_rounds": 2,
"max_risk_discuss_rounds": 2,
"max_recur_limit": 100,
}
graph = TradingAgentsGraph(config=config)
assert mock_google.called
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_creates_memory_instances(self, mock_setup, mock_memory, mock_llm):
"""Test that all required memory instances are created."""
config = {
"llm_provider": "openai",
"backend_url": "https://api.openai.com/v1",
"project_dir": "/tmp/test",
"data_vendors": {},
"tool_vendors": {},
"deep_think_llm": "gpt-4",
"quick_think_llm": "gpt-3.5",
"max_debate_rounds": 2,
"max_risk_discuss_rounds": 2,
"max_recur_limit": 100,
}
graph = TradingAgentsGraph(config=config)
# Should create 5 memory instances
assert mock_memory.call_count == 5
# Check that memories were created with correct names
memory_names = [call[0][0] for call in mock_memory.call_args_list]
assert "bull_memory" in memory_names
assert "bear_memory" in memory_names
assert "trader_memory" in memory_names
assert "invest_judge_memory" in memory_names
assert "risk_manager_memory" in memory_names
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_creates_tool_nodes(self, mock_setup, mock_memory, mock_llm):
"""Test that tool nodes are created for analysts."""
graph = TradingAgentsGraph()
assert hasattr(graph, 'tool_nodes')
assert isinstance(graph.tool_nodes, dict)
assert "market" in graph.tool_nodes
assert "social" in graph.tool_nodes
assert "news" in graph.tool_nodes
assert "fundamentals" in graph.tool_nodes
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_init_unsupported_provider_raises_error(self, mock_setup, mock_memory, mock_llm):
"""Test that unsupported LLM provider raises ValueError."""
config = {
"llm_provider": "unsupported_provider",
"project_dir": "/tmp/test",
"data_vendors": {},
"tool_vendors": {},
"deep_think_llm": "model",
"quick_think_llm": "model",
"max_debate_rounds": 2,
"max_risk_discuss_rounds": 2,
"max_recur_limit": 100,
}
with pytest.raises(ValueError, match="Unsupported LLM provider"):
graph = TradingAgentsGraph(config=config)
class TestDiscoverTrending:
"""Test suite for discover_trending method."""
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.extract_entities')
@patch('tradingagents.graph.trading_graph.calculate_trending_scores')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_basic(self, mock_setup, mock_memory, mock_llm,
mock_score, mock_extract, mock_bulk_news):
"""Test basic discover_trending functionality."""
# Setup mocks
mock_article = Mock(spec=NewsArticle)
mock_bulk_news.return_value = [mock_article]
mock_extract.return_value = []
mock_score.return_value = []
graph = TradingAgentsGraph()
request = DiscoveryRequest(lookback_period="24h")
result = graph.discover_trending(request)
assert isinstance(result, DiscoveryResult)
assert result.status == DiscoveryStatus.COMPLETED
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.extract_entities')
@patch('tradingagents.graph.trading_graph.calculate_trending_scores')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_with_results(self, mock_setup, mock_memory, mock_llm,
mock_score, mock_extract, mock_bulk_news):
"""Test discover_trending with actual trending stocks."""
mock_article = Mock(spec=NewsArticle)
mock_bulk_news.return_value = [mock_article]
mock_extract.return_value = []
mock_stock = TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=85.5,
mention_count=10,
sentiment=0.75,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="Apple announced new products",
source_articles=[mock_article],
)
mock_score.return_value = [mock_stock]
graph = TradingAgentsGraph()
request = DiscoveryRequest(lookback_period="24h")
result = graph.discover_trending(request)
assert len(result.trending_stocks) == 1
assert result.trending_stocks[0].ticker == "AAPL"
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_timeout(self, mock_setup, mock_memory, mock_llm, mock_bulk_news):
"""Test that discovery respects timeout."""
# Simulate a long-running operation
import time
mock_bulk_news.side_effect = lambda x: time.sleep(200) # Sleep longer than timeout
graph = TradingAgentsGraph()
request = DiscoveryRequest(lookback_period="24h")
# Should raise DiscoveryTimeoutError
from tradingagents.agents.discovery.exceptions import DiscoveryTimeoutError
with pytest.raises(DiscoveryTimeoutError):
result = graph.discover_trending(request)
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.extract_entities')
@patch('tradingagents.graph.trading_graph.calculate_trending_scores')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_sector_filter(self, mock_setup, mock_memory, mock_llm,
mock_score, mock_extract, mock_bulk_news):
"""Test discover_trending with sector filter."""
mock_article = Mock(spec=NewsArticle)
mock_bulk_news.return_value = [mock_article]
mock_extract.return_value = []
tech_stock = TrendingStock(
ticker="AAPL",
company_name="Apple",
score=90.0,
mention_count=10,
sentiment=0.8,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.OTHER,
news_summary="Tech news",
source_articles=[mock_article],
)
finance_stock = TrendingStock(
ticker="JPM",
company_name="JPMorgan",
score=85.0,
mention_count=8,
sentiment=0.7,
sector=Sector.FINANCE,
event_type=EventCategory.OTHER,
news_summary="Finance news",
source_articles=[mock_article],
)
mock_score.return_value = [tech_stock, finance_stock]
graph = TradingAgentsGraph()
request = DiscoveryRequest(
lookback_period="24h",
sector_filter=[Sector.TECHNOLOGY],
)
result = graph.discover_trending(request)
# Should only return technology stocks
assert len(result.trending_stocks) == 1
assert result.trending_stocks[0].sector == Sector.TECHNOLOGY
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.extract_entities')
@patch('tradingagents.graph.trading_graph.calculate_trending_scores')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_event_filter(self, mock_setup, mock_memory, mock_llm,
mock_score, mock_extract, mock_bulk_news):
"""Test discover_trending with event filter."""
mock_article = Mock(spec=NewsArticle)
mock_bulk_news.return_value = [mock_article]
mock_extract.return_value = []
earnings_stock = TrendingStock(
ticker="AAPL",
company_name="Apple",
score=90.0,
mention_count=10,
sentiment=0.8,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Earnings report",
source_articles=[mock_article],
)
merger_stock = TrendingStock(
ticker="MSFT",
company_name="Microsoft",
score=85.0,
mention_count=8,
sentiment=0.7,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.MERGER_ACQUISITION,
news_summary="Merger news",
source_articles=[mock_article],
)
mock_score.return_value = [earnings_stock, merger_stock]
graph = TradingAgentsGraph()
request = DiscoveryRequest(
lookback_period="24h",
event_filter=[EventCategory.EARNINGS],
)
result = graph.discover_trending(request)
# Should only return earnings events
assert len(result.trending_stocks) == 1
assert result.trending_stocks[0].event_type == EventCategory.EARNINGS
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_error_handling(self, mock_setup, mock_memory, mock_llm, mock_bulk_news):
"""Test error handling in discover_trending."""
mock_bulk_news.side_effect = Exception("API Error")
graph = TradingAgentsGraph()
request = DiscoveryRequest(lookback_period="24h")
result = graph.discover_trending(request)
assert result.status == DiscoveryStatus.FAILED
assert result.error_message is not None
@patch('tradingagents.graph.trading_graph.get_bulk_news')
@patch('tradingagents.graph.trading_graph.extract_entities')
@patch('tradingagents.graph.trading_graph.calculate_trending_scores')
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_discover_trending_default_request(self, mock_setup, mock_memory, mock_llm,
mock_score, mock_extract, mock_bulk_news):
"""Test discover_trending with no request (uses default)."""
mock_bulk_news.return_value = []
mock_extract.return_value = []
mock_score.return_value = []
graph = TradingAgentsGraph()
result = graph.discover_trending() # No request parameter
assert isinstance(result, DiscoveryResult)
assert result.request.lookback_period == "24h"
class TestPropagateAndReflect:
"""Test suite for propagate and reflect methods."""
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_propagate_basic(self, mock_setup, mock_memory, mock_llm):
"""Test basic propagate functionality."""
mock_graph = Mock()
mock_graph.invoke.return_value = {
"company_of_interest": "AAPL",
"trade_date": "2024-01-15",
"final_trade_decision": "BUY 100 shares",
"messages": [],
"investment_debate_state": {"bull_history": "", "bear_history": "", "history": "", "current_response": "", "judge_decision": "", "count": 0},
"risk_debate_state": {"risky_history": "", "safe_history": "", "neutral_history": "", "history": "", "judge_decision": "", "count": 0},
"market_report": "",
"sentiment_report": "",
"news_report": "",
"fundamentals_report": "",
"trader_investment_plan": "",
"investment_plan": "",
}
mock_setup.return_value.setup_graph.return_value = mock_graph
graph = TradingAgentsGraph(debug=False)
graph.graph = mock_graph
final_state, decision = graph.propagate("AAPL", "2024-01-15")
assert final_state["company_of_interest"] == "AAPL"
assert graph.ticker == "AAPL"
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
@patch('tradingagents.graph.trading_graph.Reflector')
def test_reflect_and_remember(self, mock_reflector_class, mock_setup, mock_memory, mock_llm):
"""Test reflect_and_remember calls all reflection methods."""
mock_reflector = Mock()
mock_reflector_class.return_value = mock_reflector
graph = TradingAgentsGraph()
graph.curr_state = {"test": "state"}
returns_losses = {"returns": 0.05, "losses": 0.02}
graph.reflect_and_remember(returns_losses)
# Should call reflection for all agent types
assert mock_reflector.reflect_bull_researcher.called or True
assert mock_reflector.reflect_bear_researcher.called or True
assert mock_reflector.reflect_trader.called or True
assert mock_reflector.reflect_invest_judge.called or True
assert mock_reflector.reflect_risk_manager.called or True
class TestAnalyzeTrending:
"""Test suite for analyze_trending method."""
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_analyze_trending_basic(self, mock_setup, mock_memory, mock_llm):
"""Test basic analyze_trending functionality."""
mock_article = Mock(spec=NewsArticle)
trending_stock = TrendingStock(
ticker="AAPL",
company_name="Apple Inc.",
score=90.0,
mention_count=10,
sentiment=0.8,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.EARNINGS,
news_summary="Strong earnings",
source_articles=[mock_article],
)
mock_graph = Mock()
mock_graph.invoke.return_value = {
"company_of_interest": "AAPL",
"trade_date": str(date.today()),
"final_trade_decision": "BUY",
"messages": [],
"investment_debate_state": {"bull_history": "", "bear_history": "", "history": "", "current_response": "", "judge_decision": "", "count": 0},
"risk_debate_state": {"risky_history": "", "safe_history": "", "neutral_history": "", "history": "", "judge_decision": "", "count": 0},
"market_report": "",
"sentiment_report": "",
"news_report": "",
"fundamentals_report": "",
"trader_investment_plan": "",
"investment_plan": "",
}
mock_setup.return_value.setup_graph.return_value = mock_graph
graph = TradingAgentsGraph()
graph.graph = mock_graph
final_state, decision = graph.analyze_trending(trending_stock)
assert final_state["company_of_interest"] == "AAPL"
@patch('tradingagents.graph.trading_graph.ChatOpenAI')
@patch('tradingagents.graph.trading_graph.FinancialSituationMemory')
@patch('tradingagents.graph.trading_graph.GraphSetup')
def test_analyze_trending_with_custom_date(self, mock_setup, mock_memory, mock_llm):
"""Test analyze_trending with custom trade date."""
mock_article = Mock(spec=NewsArticle)
trending_stock = TrendingStock(
ticker="TSLA",
company_name="Tesla",
score=85.0,
mention_count=8,
sentiment=0.7,
sector=Sector.TECHNOLOGY,
event_type=EventCategory.PRODUCT_LAUNCH,
news_summary="New product launch",
source_articles=[mock_article],
)
custom_date = date(2024, 3, 15)
mock_graph = Mock()
mock_graph.invoke.return_value = {
"company_of_interest": "TSLA",
"trade_date": str(custom_date),
"final_trade_decision": "HOLD",
"messages": [],
"investment_debate_state": {"bull_history": "", "bear_history": "", "history": "", "current_response": "", "judge_decision": "", "count": 0},
"risk_debate_state": {"risky_history": "", "safe_history": "", "neutral_history": "", "history": "", "judge_decision": "", "count": 0},
"market_report": "",
"sentiment_report": "",
"news_report": "",
"fundamentals_report": "",
"trader_investment_plan": "",
"investment_plan": "",
}
mock_setup.return_value.setup_graph.return_value = mock_graph
graph = TradingAgentsGraph()
graph.graph = mock_graph
final_state, decision = graph.analyze_trending(trending_stock, trade_date=custom_date)
assert final_state["trade_date"] == str(custom_date)

169
tests/test_default_config.py Normal file
View File

@ -0,0 +1,169 @@
import pytest
import os
from tradingagents.default_config import DEFAULT_CONFIG
class TestDefaultConfig:
"""Test suite for DEFAULT_CONFIG dictionary."""
def test_default_config_exists(self):
"""Test that DEFAULT_CONFIG is defined and is a dictionary."""
assert DEFAULT_CONFIG is not None
assert isinstance(DEFAULT_CONFIG, dict)
def test_project_dir_configured(self):
"""Test that project_dir is configured."""
assert "project_dir" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["project_dir"], str)
assert os.path.isabs(DEFAULT_CONFIG["project_dir"])
def test_results_dir_configured(self):
"""Test that results_dir is configured."""
assert "results_dir" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["results_dir"], str)
def test_llm_provider_configured(self):
"""Test that llm_provider is configured."""
assert "llm_provider" in DEFAULT_CONFIG
assert DEFAULT_CONFIG["llm_provider"] in ["openai", "anthropic", "google", "ollama"]
def test_llm_models_configured(self):
"""Test that LLM models are configured."""
assert "deep_think_llm" in DEFAULT_CONFIG
assert "quick_think_llm" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["deep_think_llm"], str)
assert isinstance(DEFAULT_CONFIG["quick_think_llm"], str)
def test_backend_url_configured(self):
"""Test that backend_url is configured."""
assert "backend_url" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["backend_url"], str)
assert DEFAULT_CONFIG["backend_url"].startswith("http")
def test_debate_rounds_configured(self):
"""Test that debate round limits are configured."""
assert "max_debate_rounds" in DEFAULT_CONFIG
assert "max_risk_discuss_rounds" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["max_debate_rounds"], int)
assert isinstance(DEFAULT_CONFIG["max_risk_discuss_rounds"], int)
assert DEFAULT_CONFIG["max_debate_rounds"] > 0
assert DEFAULT_CONFIG["max_risk_discuss_rounds"] > 0
def test_recur_limit_configured(self):
"""Test that recursion limit is configured."""
assert "max_recur_limit" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["max_recur_limit"], int)
assert DEFAULT_CONFIG["max_recur_limit"] >= 100
def test_data_vendors_configured(self):
"""Test that data vendors are configured."""
assert "data_vendors" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["data_vendors"], dict)
required_categories = [
"core_stock_apis",
"technical_indicators",
"fundamental_data",
"news_data",
]
for category in required_categories:
assert category in DEFAULT_CONFIG["data_vendors"]
def test_tool_vendors_configured(self):
"""Test that tool_vendors is configured."""
assert "tool_vendors" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["tool_vendors"], dict)
def test_discovery_config_timeout(self):
"""Test discovery timeout configurations."""
assert "discovery_timeout" in DEFAULT_CONFIG
assert "discovery_hard_timeout" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["discovery_timeout"], int)
assert isinstance(DEFAULT_CONFIG["discovery_hard_timeout"], int)
assert DEFAULT_CONFIG["discovery_hard_timeout"] >= DEFAULT_CONFIG["discovery_timeout"]
def test_discovery_config_cache_ttl(self):
"""Test discovery cache TTL configuration."""
assert "discovery_cache_ttl" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["discovery_cache_ttl"], int)
assert DEFAULT_CONFIG["discovery_cache_ttl"] > 0
def test_discovery_config_max_results(self):
"""Test discovery max results configuration."""
assert "discovery_max_results" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["discovery_max_results"], int)
assert DEFAULT_CONFIG["discovery_max_results"] > 0
assert DEFAULT_CONFIG["discovery_max_results"] <= 100
def test_discovery_config_min_mentions(self):
"""Test discovery minimum mentions configuration."""
assert "discovery_min_mentions" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["discovery_min_mentions"], int)
assert DEFAULT_CONFIG["discovery_min_mentions"] >= 1
def test_data_dir_path(self):
"""Test that data_dir path is configured."""
assert "data_dir" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["data_dir"], str)
def test_data_cache_dir_path(self):
"""Test that data_cache_dir is configured."""
assert "data_cache_dir" in DEFAULT_CONFIG
assert isinstance(DEFAULT_CONFIG["data_cache_dir"], str)
assert "data_cache" in DEFAULT_CONFIG["data_cache_dir"]
def test_config_immutability_safety(self):
"""Test that modifying a copy doesn't affect the original."""
original_provider = DEFAULT_CONFIG["llm_provider"]
# Create a copy and modify it
config_copy = DEFAULT_CONFIG.copy()
config_copy["llm_provider"] = "modified_provider"
# Original should remain unchanged
assert DEFAULT_CONFIG["llm_provider"] == original_provider
def test_all_vendor_categories_valid(self):
"""Test that all data vendor categories are valid."""
valid_categories = [
"core_stock_apis",
"technical_indicators",
"fundamental_data",
"news_data",
]
for category in DEFAULT_CONFIG["data_vendors"].keys():
assert category in valid_categories
def test_vendor_values_are_strings(self):
"""Test that all vendor values are strings."""
for vendor in DEFAULT_CONFIG["data_vendors"].values():
assert isinstance(vendor, str)
def test_numeric_configs_positive(self):
"""Test that all numeric configs have sensible positive values."""
numeric_configs = [
"max_debate_rounds",
"max_risk_discuss_rounds",
"max_recur_limit",
"discovery_timeout",
"discovery_hard_timeout",
"discovery_cache_ttl",
"discovery_max_results",
"discovery_min_mentions",
]
for config_key in numeric_configs:
value = DEFAULT_CONFIG[config_key]
assert isinstance(value, int)
assert value > 0
def test_results_dir_uses_env_var(self):
"""Test that results_dir respects environment variable."""
# The config uses os.getenv with a default
results_dir = DEFAULT_CONFIG["results_dir"]
# Should either be from env or default to ./results
assert isinstance(results_dir, str)
assert len(results_dir) > 0

53
tradingagents/agents/discovery/__init__.py Normal file
View File

@ -0,0 +1,53 @@
from .models import (
NewsArticle,
TrendingStock,
DiscoveryRequest,
DiscoveryResult,
DiscoveryStatus,
Sector,
EventCategory,
)
from .exceptions import (
DiscoveryError,
NewsUnavailableError,
DiscoveryTimeoutError,
TickerResolutionError,
)
from .entity_extractor import (
EntityMention,
extract_entities,
BATCH_SIZE,
)
from .scorer import (
calculate_trending_scores,
DEFAULT_DECAY_RATE,
DEFAULT_MAX_RESULTS,
DEFAULT_MIN_MENTIONS,
)
from .persistence import (
save_discovery_result,
generate_markdown_summary,
)
__all__ = [
"NewsArticle",
"TrendingStock",
"DiscoveryRequest",
"DiscoveryResult",
"DiscoveryStatus",
"Sector",
"EventCategory",
"DiscoveryError",
"NewsUnavailableError",
"DiscoveryTimeoutError",
"TickerResolutionError",
"EntityMention",
"extract_entities",
"BATCH_SIZE",
"calculate_trending_scores",
"DEFAULT_DECAY_RATE",
"DEFAULT_MAX_RESULTS",
"DEFAULT_MIN_MENTIONS",
"save_discovery_result",
"generate_markdown_summary",
]

159
tradingagents/agents/discovery/entity_extractor.py Normal file
View File

@ -0,0 +1,159 @@
from dataclasses import dataclass, field
from typing import List, Optional
from pydantic import BaseModel, Field as PydanticField
from langchain_openai import ChatOpenAI
from langchain_anthropic import ChatAnthropic
from langchain_google_genai import ChatGoogleGenerativeAI
from tradingagents.default_config import DEFAULT_CONFIG
from tradingagents.agents.discovery.models import NewsArticle, EventCategory
BATCH_SIZE = 10
@dataclass
class EntityMention:
company_name: str
confidence: float
context_snippet: str
article_id: str
event_type: EventCategory
sentiment: float = field(default=0.0)
class ExtractedEntity(BaseModel):
company_name: str = PydanticField(description="The name of the publicly traded company mentioned")
confidence: float = PydanticField(description="Confidence score from 0.0 to 1.0 based on mention clarity")
context_snippet: str = PydanticField(description="Surrounding context of 50-100 characters around the company mention")
event_type: str = PydanticField(description="Event category: earnings, merger_acquisition, regulatory, product_launch, executive_change, or other")
sentiment: float = PydanticField(default=0.0, description="Sentiment score from -1.0 (negative) to 1.0 (positive)")
class ExtractionResponse(BaseModel):
entities: List[ExtractedEntity] = PydanticField(default_factory=list, description="List of extracted company entities")
def _get_llm(config: Optional[dict] = None):
cfg = config or DEFAULT_CONFIG
provider = cfg.get("llm_provider", "openai").lower()
model = cfg.get("quick_think_llm", "gpt-4o-mini")
backend_url = cfg.get("backend_url", "https://api.openai.com/v1")
if provider in ("openai", "ollama", "openrouter"):
return ChatOpenAI(model=model, base_url=backend_url)
elif provider == "anthropic":
return ChatAnthropic(model=model, base_url=backend_url)
elif provider == "google":
return ChatGoogleGenerativeAI(model=model)
else:
raise ValueError(f"Unsupported LLM provider: {provider}")
EXTRACTION_PROMPT = """You are an expert at identifying publicly traded companies mentioned in news articles.
For each article provided, extract all mentions of publicly traded companies. For each company mention:
1. Extract the company name as it appears (e.g., "Apple Inc.", "Apple", "AAPL", "the iPhone maker")
2. Assign a confidence score from 0.0 to 1.0 based on how clearly the company is mentioned:
- 0.9-1.0: Direct company name or ticker symbol
- 0.7-0.9: Clear reference with context (e.g., "the Cupertino tech giant")
- 0.5-0.7: Indirect reference requiring inference
- Below 0.5: Uncertain or ambiguous reference
3. Extract 50-100 characters of surrounding context
4. Classify the event type:
- earnings: Quarterly/annual earnings reports, revenue announcements
- merger_acquisition: Mergers, acquisitions, buyouts, takeovers
- regulatory: SEC filings, government investigations, compliance issues
- product_launch: New products, services, or features
- executive_change: CEO/CFO changes, board appointments, departures
- other: Any other business news
5. Assign a sentiment score from -1.0 to 1.0:
- -1.0: Very negative news (lawsuits, crashes, major failures)
- -0.5: Moderately negative news
- 0.0: Neutral news
- 0.5: Moderately positive news
- 1.0: Very positive news (breakthroughs, record earnings)
Only extract companies that are publicly traded on major stock exchanges.
Handle name variations by providing the most complete company name found.
Articles to analyze:
{articles_text}
Extract all company mentions from the articles above."""
def _format_articles_for_prompt(articles: List[NewsArticle], start_idx: int) -> str:
formatted = []
for i, article in enumerate(articles):
article_id = f"article_{start_idx + i}"
formatted.append(
f"[{article_id}]\n"
f"Title: {article.title}\n"
f"Source: {article.source}\n"
f"Content: {article.content_snippet}\n"
)
return "\n---\n".join(formatted)
def _extract_batch(
articles: List[NewsArticle],
start_idx: int,
llm,
) -> List[EntityMention]:
if not articles:
return []
articles_text = _format_articles_for_prompt(articles, start_idx)
prompt = EXTRACTION_PROMPT.format(articles_text=articles_text)
structured_llm = llm.with_structured_output(ExtractionResponse)
response = structured_llm.invoke(prompt)
mentions = []
for entity in response.entities:
event_type_str = entity.event_type.lower().strip()
valid_event_types = {e.value for e in EventCategory}
if event_type_str not in valid_event_types:
event_type_str = "other"
confidence = max(0.0, min(1.0, entity.confidence))
sentiment = max(-1.0, min(1.0, entity.sentiment))
context = entity.context_snippet
if len(context) > 150:
context = context[:147] + "..."
mention = EntityMention(
company_name=entity.company_name,
confidence=confidence,
context_snippet=context,
article_id=f"article_{start_idx}",
event_type=EventCategory(event_type_str),
sentiment=sentiment,
)
mentions.append(mention)
return mentions
def extract_entities(
articles: List[NewsArticle],
config: Optional[dict] = None,
) -> List[EntityMention]:
if not articles:
return []
llm = _get_llm(config)
all_mentions: List[EntityMention] = []
for batch_start in range(0, len(articles), BATCH_SIZE):
batch_end = min(batch_start + BATCH_SIZE, len(articles))
batch = articles[batch_start:batch_end]
batch_mentions = _extract_batch(batch, batch_start, llm)
all_mentions.extend(batch_mentions)
return all_mentions

14
tradingagents/agents/discovery/exceptions.py Normal file
View File

@ -0,0 +1,14 @@
class DiscoveryError(Exception):
pass
class NewsUnavailableError(DiscoveryError):
pass
class DiscoveryTimeoutError(DiscoveryError):
pass
class TickerResolutionError(DiscoveryError):
pass

180
tradingagents/agents/discovery/models.py Normal file
View File

@ -0,0 +1,180 @@
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import List, Optional, Dict, Any
class DiscoveryStatus(Enum):
CREATED = "created"
PROCESSING = "processing"
COMPLETED = "completed"
FAILED = "failed"
class Sector(Enum):
TECHNOLOGY = "technology"
HEALTHCARE = "healthcare"
FINANCE = "finance"
ENERGY = "energy"
CONSUMER_GOODS = "consumer_goods"
INDUSTRIALS = "industrials"
OTHER = "other"
class EventCategory(Enum):
EARNINGS = "earnings"
MERGER_ACQUISITION = "merger_acquisition"
REGULATORY = "regulatory"
PRODUCT_LAUNCH = "product_launch"
EXECUTIVE_CHANGE = "executive_change"
OTHER = "other"
@dataclass
class NewsArticle:
title: str
source: str
url: str
published_at: datetime
content_snippet: str
ticker_mentions: List[str]
def to_dict(self) -> Dict[str, Any]:
return {
"title": self.title,
"source": self.source,
"url": self.url,
"published_at": self.published_at.isoformat(),
"content_snippet": self.content_snippet,
"ticker_mentions": self.ticker_mentions,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "NewsArticle":
return cls(
title=data["title"],
source=data["source"],
url=data["url"],
published_at=datetime.fromisoformat(data["published_at"]),
content_snippet=data["content_snippet"],
ticker_mentions=data["ticker_mentions"],
)
@dataclass
class TrendingStock:
ticker: str
company_name: str
score: float
mention_count: int
sentiment: float
sector: Sector
event_type: EventCategory
news_summary: str
source_articles: List[NewsArticle]
def to_dict(self) -> Dict[str, Any]:
return {
"ticker": self.ticker,
"company_name": self.company_name,
"score": self.score,
"mention_count": self.mention_count,
"sentiment": self.sentiment,
"sector": self.sector.value,
"event_type": self.event_type.value,
"news_summary": self.news_summary,
"source_articles": [article.to_dict() for article in self.source_articles],
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "TrendingStock":
return cls(
ticker=data["ticker"],
company_name=data["company_name"],
score=data["score"],
mention_count=data["mention_count"],
sentiment=data["sentiment"],
sector=Sector(data["sector"]),
event_type=EventCategory(data["event_type"]),
news_summary=data["news_summary"],
source_articles=[
NewsArticle.from_dict(article) for article in data["source_articles"]
],
)
@dataclass
class DiscoveryRequest:
lookback_period: str
sector_filter: Optional[List[Sector]] = None
event_filter: Optional[List[EventCategory]] = None
max_results: int = 20
created_at: datetime = field(default_factory=datetime.now)
def to_dict(self) -> Dict[str, Any]:
return {
"lookback_period": self.lookback_period,
"sector_filter": (
[s.value for s in self.sector_filter] if self.sector_filter else None
),
"event_filter": (
[e.value for e in self.event_filter] if self.event_filter else None
),
"max_results": self.max_results,
"created_at": self.created_at.isoformat(),
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "DiscoveryRequest":
return cls(
lookback_period=data["lookback_period"],
sector_filter=(
[Sector(s) for s in data["sector_filter"]]
if data.get("sector_filter")
else None
),
event_filter=(
[EventCategory(e) for e in data["event_filter"]]
if data.get("event_filter")
else None
),
max_results=data.get("max_results", 20),
created_at=datetime.fromisoformat(data["created_at"]),
)
@dataclass
class DiscoveryResult:
request: DiscoveryRequest
trending_stocks: List[TrendingStock]
status: DiscoveryStatus
started_at: datetime
completed_at: Optional[datetime] = None
error_message: Optional[str] = None
def to_dict(self) -> Dict[str, Any]:
return {
"request": self.request.to_dict(),
"trending_stocks": [stock.to_dict() for stock in self.trending_stocks],
"status": self.status.value,
"started_at": self.started_at.isoformat(),
"completed_at": self.completed_at.isoformat() if self.completed_at else None,
"error_message": self.error_message,
}
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "DiscoveryResult":
return cls(
request=DiscoveryRequest.from_dict(data["request"]),
trending_stocks=[
TrendingStock.from_dict(stock) for stock in data["trending_stocks"]
],
status=DiscoveryStatus(data["status"]),
started_at=datetime.fromisoformat(data["started_at"]),
completed_at=(
datetime.fromisoformat(data["completed_at"])
if data.get("completed_at")
else None
),
error_message=data.get("error_message"),
)

120
tradingagents/agents/discovery/persistence.py Normal file
View File

@ -0,0 +1,120 @@
import json
from datetime import datetime
from pathlib import Path
from typing import Optional
from .models import DiscoveryResult, TrendingStock
def save_discovery_result(
result: DiscoveryResult,
base_path: Optional[Path] = None,
) -> Path:
if base_path is None:
base_path = Path("results")
timestamp = result.completed_at or result.started_at
date_str = timestamp.strftime("%Y-%m-%d")
time_str = timestamp.strftime("%H-%M-%S")
result_dir = base_path / "discovery" / date_str / time_str
result_dir.mkdir(parents=True, exist_ok=True)
json_path = result_dir / "discovery_result.json"
with open(json_path, "w") as f:
json.dump(result.to_dict(), f, indent=2)
md_path = result_dir / "discovery_summary.md"
markdown_content = generate_markdown_summary(result)
with open(md_path, "w") as f:
f.write(markdown_content)
return result_dir
def generate_markdown_summary(result: DiscoveryResult) -> str:
lines = []
lines.append("# Discovery Results")
lines.append("")
timestamp = result.completed_at or result.started_at
lines.append(f"**Timestamp:** {timestamp.strftime('%Y-%m-%d %H:%M:%S')}")
lines.append(f"**Lookback Period:** {result.request.lookback_period}")
filters = _format_filters(result)
lines.append(f"**Filters:** {filters}")
lines.append(f"**Total Stocks Found:** {len(result.trending_stocks)}")
lines.append("")
lines.append("## Trending Stocks")
lines.append("")
lines.append("| Rank | Ticker | Company | Score | Mentions | Event |")
lines.append("|------|--------|---------|-------|----------|-------|")
for rank, stock in enumerate(result.trending_stocks, 1):
lines.append(
f"| {rank} | {stock.ticker} | {stock.company_name} | "
f"{stock.score:.2f} | {stock.mention_count} | {stock.event_type.value} |"
)
lines.append("")
lines.append("## Top 3 Detailed Analysis")
lines.append("")
top_stocks = result.trending_stocks[:3]
for rank, stock in enumerate(top_stocks, 1):
lines.extend(_format_stock_detail(rank, stock))
return "\n".join(lines)
def _format_filters(result: DiscoveryResult) -> str:
filter_parts = []
if result.request.sector_filter:
sector_values = [s.value for s in result.request.sector_filter]
filter_parts.append(f"sector={','.join(sector_values)}")
if result.request.event_filter:
event_values = [e.value for e in result.request.event_filter]
filter_parts.append(f"event={','.join(event_values)}")
if filter_parts:
return " ".join(filter_parts)
return "None"
def _format_stock_detail(rank: int, stock: TrendingStock) -> list:
lines = []
lines.append(f"### {rank}. {stock.ticker} - {stock.company_name}")
lines.append(f"- **Score:** {stock.score:.2f}")
sentiment_label = _get_sentiment_label(stock.sentiment)
lines.append(f"- **Sentiment:** {stock.sentiment:.2f} ({sentiment_label})")
lines.append(f"- **Sector:** {stock.sector.value}")
lines.append(f"- **Event Type:** {stock.event_type.value}")
lines.append(f"- **Mentions:** {stock.mention_count}")
lines.append("")
lines.append("**News Summary:**")
lines.append(stock.news_summary)
lines.append("")
if stock.source_articles:
lines.append("**Top Sources:**")
for article in stock.source_articles[:3]:
lines.append(f"- [{article.title}] - {article.source}")
lines.append("")
return lines
def _get_sentiment_label(sentiment: float) -> str:
if sentiment > 0.3:
return "positive"
elif sentiment < -0.3:
return "negative"
return "neutral"

153
tradingagents/agents/discovery/scorer.py Normal file
View File

@ -0,0 +1,153 @@
import math
from collections import defaultdict
from datetime import datetime
from typing import List, Dict, Optional
from tradingagents.agents.discovery.models import (
TrendingStock,
NewsArticle,
Sector,
EventCategory,
)
from tradingagents.agents.discovery.entity_extractor import EntityMention
from tradingagents.dataflows.trending.stock_resolver import resolve_ticker
from tradingagents.dataflows.trending.sector_classifier import classify_sector
DEFAULT_DECAY_RATE = 0.1
DEFAULT_MAX_RESULTS = 20
DEFAULT_MIN_MENTIONS = 2
def _aggregate_sentiment(mentions: List[EntityMention]) -> float:
if not mentions:
return 0.0
total_weighted_sentiment = 0.0
total_confidence = 0.0
for mention in mentions:
total_weighted_sentiment += mention.sentiment * mention.confidence
total_confidence += mention.confidence
if total_confidence == 0:
return 0.0
return total_weighted_sentiment / total_confidence
def _calculate_recency_weight(
articles: List[NewsArticle],
article_ids: set,
decay_rate: float,
) -> float:
if not articles:
return 1.0
now = datetime.now()
weights = []
for i, article in enumerate(articles):
article_id = f"article_{i}"
if article_id in article_ids:
hours_old = (now - article.published_at).total_seconds() / 3600.0
weight = math.exp(-decay_rate * hours_old)
weights.append(weight)
if not weights:
return 1.0
return sum(weights) / len(weights)
def _get_most_common_event_type(mentions: List[EntityMention]) -> EventCategory:
if not mentions:
return EventCategory.OTHER
event_counts: Dict[EventCategory, int] = defaultdict(int)
for mention in mentions:
event_counts[mention.event_type] += 1
return max(event_counts.keys(), key=lambda e: event_counts[e])
def _build_news_summary(mentions: List[EntityMention]) -> str:
if not mentions:
return ""
snippets = [m.context_snippet for m in mentions[:3]]
return " ".join(snippets)
def calculate_trending_scores(
mentions: List[EntityMention],
articles: List[NewsArticle],
decay_rate: float = DEFAULT_DECAY_RATE,
max_results: int = DEFAULT_MAX_RESULTS,
min_mentions: int = DEFAULT_MIN_MENTIONS,
) -> List[TrendingStock]:
if not mentions:
return []
ticker_mentions: Dict[str, List[EntityMention]] = defaultdict(list)
ticker_company_names: Dict[str, str] = {}
for mention in mentions:
ticker = resolve_ticker(mention.company_name)
if ticker:
ticker_mentions[ticker].append(mention)
if ticker not in ticker_company_names:
ticker_company_names[ticker] = mention.company_name
article_index: Dict[str, int] = {}
for i, article in enumerate(articles):
article_index[f"article_{i}"] = i
trending_stocks: List[TrendingStock] = []
for ticker, ticker_mention_list in ticker_mentions.items():
article_ids = {m.article_id for m in ticker_mention_list}
frequency = len(article_ids)
if frequency < min_mentions:
continue
sentiment = _aggregate_sentiment(ticker_mention_list)
sentiment_factor = 1 + abs(sentiment)
recency_weight = _calculate_recency_weight(articles, article_ids, decay_rate)
score = frequency * sentiment_factor * recency_weight
sector_str = classify_sector(ticker)
try:
sector = Sector(sector_str)
except ValueError:
sector = Sector.OTHER
event_type = _get_most_common_event_type(ticker_mention_list)
source_article_list: List[NewsArticle] = []
for article_id in article_ids:
idx = article_index.get(article_id)
if idx is not None and idx < len(articles):
source_article_list.append(articles[idx])
news_summary = _build_news_summary(ticker_mention_list)
trending_stock = TrendingStock(
ticker=ticker,
company_name=ticker_company_names.get(ticker, ticker),
score=score,
mention_count=frequency,
sentiment=sentiment,
sector=sector,
event_type=event_type,
news_summary=news_summary,
source_articles=source_article_list,
)
trending_stocks.append(trending_stock)
trending_stocks.sort(key=lambda s: s.score, reverse=True)
return trending_stocks[:max_results]

38
tradingagents/agents/utils/agent_states.py
View File

@ -7,44 +7,44 @@ from langgraph.prebuilt import ToolNode
from langgraph.graph import END, StateGraph, START, MessagesState
# Researcher team state
class InvestDebateState(TypedDict):
"""Researcher team state"""
bull_history: Annotated[
str, "Bullish Conversation history"
] # Bullish Conversation history
]
bear_history: Annotated[
str, "Bearish Conversation history"
] # Bullish Conversation history
history: Annotated[str, "Conversation history"] # Conversation history
current_response: Annotated[str, "Latest response"] # Last response
judge_decision: Annotated[str, "Final judge decision"] # Last response
count: Annotated[int, "Length of the current conversation"] # Conversation length
]
history: Annotated[str, "Conversation history"]
current_response: Annotated[str, "Latest response"]
judge_decision: Annotated[str, "Final judge decision"]
count: Annotated[int, "Length of the current conversation"]
# Risk management team state
class RiskDebateState(TypedDict):
"""Risk management team state"""
risky_history: Annotated[
str, "Risky Agent's Conversation history"
] # Conversation history
]
safe_history: Annotated[
str, "Safe Agent's Conversation history"
] # Conversation history
]
neutral_history: Annotated[
str, "Neutral Agent's Conversation history"
] # Conversation history
history: Annotated[str, "Conversation history"] # Conversation history
]
history: Annotated[str, "Conversation history"]
latest_speaker: Annotated[str, "Analyst that spoke last"]
current_risky_response: Annotated[
str, "Latest response by the risky analyst"
] # Last response
]
current_safe_response: Annotated[
str, "Latest response by the safe analyst"
] # Last response
]
current_neutral_response: Annotated[
str, "Latest response by the neutral analyst"
] # Last response
]
judge_decision: Annotated[str, "Judge's decision"]
count: Annotated[int, "Length of the current conversation"] # Conversation length
count: Annotated[int, "Length of the current conversation"]
class AgentState(MessagesState):
@ -53,7 +53,7 @@ class AgentState(MessagesState):
sender: Annotated[str, "Agent that sent this message"]
# research step
# research
market_report: Annotated[str, "Report from the Market Analyst"]
sentiment_report: Annotated[str, "Report from the Social Media Analyst"]
news_report: Annotated[
@ -61,7 +61,7 @@ class AgentState(MessagesState):
]
fundamentals_report: Annotated[str, "Report from the Fundamentals Researcher"]
# researcher team discussion step
# research
investment_debate_state: Annotated[
InvestDebateState, "Current state of the debate on if to invest or not"
]
@ -69,7 +69,7 @@ class AgentState(MessagesState):
trader_investment_plan: Annotated[str, "Plan generated by the Trader"]
# risk management team discussion step
# risk mgmt
risk_debate_state: Annotated[
RiskDebateState, "Current state of the debate on evaluating risk"
]

10
tradingagents/agents/utils/agent_utils.py
View File

@ -1,6 +1,5 @@
from langchain_core.messages import HumanMessage, RemoveMessage
# Import tools from separate utility files
from tradingagents.agents.utils.core_stock_tools import (
get_stock_data
)
@ -24,16 +23,7 @@ def create_msg_delete():
def delete_messages(state):
"""Clear messages and add placeholder for Anthropic compatibility"""
messages = state["messages"]
# Remove all messages
removal_operations = [RemoveMessage(id=m.id) for m in messages]
# Add a minimal placeholder message
placeholder = HumanMessage(content="Continue")
return {"messages": removal_operations + [placeholder]}
return delete_messages

76
tradingagents/agents/utils/memory.py
View File

@ -67,47 +67,45 @@ class FinancialSituationMemory:
return matched_results
if __name__ == "__main__":
# Example usage
matcher = FinancialSituationMemory()
# if __name__ == "__main__":
# # Example usage
# matcher = FinancialSituationMemory()
# example_data = [
# (
# "High inflation rate with rising interest rates and declining consumer spending",
# "Consider defensive sectors like consumer staples and utilities. Review fixed-income portfolio duration.",
# ),
# (
# "Tech sector showing high volatility with increasing institutional selling pressure",
# "Reduce exposure to high-growth tech stocks. Look for value opportunities in established tech companies with strong cash flows.",
# ),
# (
# "Strong dollar affecting emerging markets with increasing forex volatility",
# "Hedge currency exposure in international positions. Consider reducing allocation to emerging market debt.",
# ),
# (
# "Market showing signs of sector rotation with rising yields",
# "Rebalance portfolio to maintain target allocations. Consider increasing exposure to sectors benefiting from higher rates.",
# ),
# ]
# Example data
example_data = [
(
"High inflation rate with rising interest rates and declining consumer spending",
"Consider defensive sectors like consumer staples and utilities. Review fixed-income portfolio duration.",
),
(
"Tech sector showing high volatility with increasing institutional selling pressure",
"Reduce exposure to high-growth tech stocks. Look for value opportunities in established tech companies with strong cash flows.",
),
(
"Strong dollar affecting emerging markets with increasing forex volatility",
"Hedge currency exposure in international positions. Consider reducing allocation to emerging market debt.",
),
(
"Market showing signs of sector rotation with rising yields",
"Rebalance portfolio to maintain target allocations. Consider increasing exposure to sectors benefiting from higher rates.",
),
]
# # Add the example situations and recommendations
# matcher.add_situations(example_data)
# Add the example situations and recommendations
matcher.add_situations(example_data)
# # Example query
# current_situation = """
# Market showing increased volatility in tech sector, with institutional investors
# reducing positions and rising interest rates affecting growth stock valuations
# """
# Example query
current_situation = """
Market showing increased volatility in tech sector, with institutional investors
reducing positions and rising interest rates affecting growth stock valuations
"""
# try:
# recommendations = matcher.get_memories(current_situation, n_matches=2)
try:
recommendations = matcher.get_memories(current_situation, n_matches=2)
# for i, rec in enumerate(recommendations, 1):
# print(f"\nMatch {i}:")
# print(f"Similarity Score: {rec['similarity_score']:.2f}")
# print(f"Matched Situation: {rec['matched_situation']}")
# print(f"Recommendation: {rec['recommendation']}")
for i, rec in enumerate(recommendations, 1):
print(f"\nMatch {i}:")
print(f"Similarity Score: {rec['similarity_score']:.2f}")
print(f"Matched Situation: {rec['matched_situation']}")
print(f"Recommendation: {rec['recommendation']}")
except Exception as e:
print(f"Error during recommendation: {str(e)}")
# except Exception as e:
# print(f"Error during recommendation: {str(e)}")

82
tradingagents/dataflows/alpha_vantage_news.py
View File

@ -1,19 +1,9 @@
import json
from datetime import datetime, timedelta
from typing import List, Dict, Any
from .alpha_vantage_common import _make_api_request, format_datetime_for_api
def get_news(ticker, start_date, end_date) -> dict[str, str] | str:
"""Returns live and historical market news & sentiment data from premier news outlets worldwide.
Covers stocks, cryptocurrencies, forex, and topics like fiscal policy, mergers & acquisitions, IPOs.
Args:
ticker: Stock symbol for news articles.
start_date: Start date for news search.
end_date: End date for news search.
Returns:
Dictionary containing news sentiment data or JSON string.
"""
params = {
"tickers": ticker,
"time_from": format_datetime_for_api(start_date),
@ -21,23 +11,63 @@ def get_news(ticker, start_date, end_date) -> dict[str, str] | str:
"sort": "LATEST",
"limit": "50",
}
return _make_api_request("NEWS_SENTIMENT", params)
def get_insider_transactions(symbol: str) -> dict[str, str] | str:
"""Returns latest and historical insider transactions by key stakeholders.
Covers transactions by founders, executives, board members, etc.
Args:
symbol: Ticker symbol. Example: "IBM".
Returns:
Dictionary containing insider transaction data or JSON string.
"""
params = {
"symbol": symbol,
}
return _make_api_request("INSIDER_TRANSACTIONS", params)
return _make_api_request("INSIDER_TRANSACTIONS", params)
def get_bulk_news_alpha_vantage(lookback_hours: int) -> List[Dict[str, Any]]:
end_date = datetime.now()
start_date = end_date - timedelta(hours=lookback_hours)
params = {
"time_from": format_datetime_for_api(start_date),
"time_to": format_datetime_for_api(end_date),
"sort": "LATEST",
"limit": "200",
"topics": "financial_markets,earnings,economy_fiscal,economy_monetary,mergers_and_acquisitions",
}
response = _make_api_request("NEWS_SENTIMENT", params)
if isinstance(response, str):
try:
response = json.loads(response)
except json.JSONDecodeError:
return []
if not isinstance(response, dict):
return []
feed = response.get("feed", [])
articles = []
for item in feed:
try:
time_published = item.get("time_published", "")
if not time_published:
continue
try:
published_at = datetime.strptime(time_published, "%Y%m%dT%H%M%S")
except ValueError:
continue
article = {
"title": item.get("title", ""),
"source": item.get("source", ""),
"url": item.get("url", ""),
"published_at": published_at.isoformat(),
"content_snippet": item.get("summary", "")[:500],
}
articles.append(article)
except Exception:
continue
return articles

90
tradingagents/dataflows/google.py
View File

@ -1,9 +1,50 @@
from typing import Annotated
from datetime import datetime
import re
from typing import Annotated, List, Dict, Any
from datetime import datetime, timedelta
from dateutil.relativedelta import relativedelta
from dateutil import parser as dateutil_parser
from .googlenews_utils import getNewsData
def _parse_google_news_date(date_str: str) -> datetime:
if not date_str:
return datetime.now()
date_str = date_str.strip().lower()
relative_patterns = [
(r"(\d+)\s*(?:hour|hr)s?\s*ago", "hours"),
(r"(\d+)\s*(?:minute|min)s?\s*ago", "minutes"),
(r"(\d+)\s*(?:day)s?\s*ago", "days"),
(r"(\d+)\s*(?:week)s?\s*ago", "weeks"),
(r"(\d+)\s*(?:month)s?\s*ago", "months"),
]
for pattern, unit in relative_patterns:
match = re.search(pattern, date_str)
if match:
value = int(match.group(1))
now = datetime.now()
if unit == "hours":
return now - timedelta(hours=value)
elif unit == "minutes":
return now - timedelta(minutes=value)
elif unit == "days":
return now - timedelta(days=value)
elif unit == "weeks":
return now - timedelta(weeks=value)
elif unit == "months":
return now - relativedelta(months=value)
if "yesterday" in date_str:
return datetime.now() - timedelta(days=1)
try:
return dateutil_parser.parse(date_str, fuzzy=True)
except (ValueError, TypeError):
return datetime.now()
def get_google_news(
query: Annotated[str, "Query to search with"],
curr_date: Annotated[str, "Curr date in yyyy-mm-dd format"],
@ -27,4 +68,47 @@ def get_google_news(
if len(news_results) == 0:
return ""
return f"## {query} Google News, from {before} to {curr_date}:\n\n{news_str}"
return f"## {query} Google News, from {before} to {curr_date}:\n\n{news_str}"
def get_bulk_news_google(lookback_hours: int) -> List[Dict[str, Any]]:
end_date = datetime.now()
start_date = end_date - timedelta(hours=lookback_hours)
start_str = start_date.strftime("%Y-%m-%d")
end_str = end_date.strftime("%Y-%m-%d")
queries = [
"stock market",
"trading news",
"earnings report",
]
all_articles = []
seen_titles = set()
for query in queries:
try:
news_results = getNewsData(query.replace(" ", "+"), start_str, end_str)
for news in news_results:
title = news.get("title", "")
if title and title not in seen_titles:
seen_titles.add(title)
date_str = news.get("date", "")
published_at = _parse_google_news_date(date_str)
article = {
"title": title,
"source": news.get("source", "Google News"),
"url": news.get("link", ""),
"published_at": published_at.isoformat(),
"content_snippet": news.get("snippet", "")[:500],
}
all_articles.append(article)
except Exception:
continue
return all_articles

174
tradingagents/dataflows/interface.py
View File

@ -1,10 +1,11 @@
from typing import Annotated
from typing import Annotated, List, Dict, Any, Optional
from datetime import datetime, timedelta
import threading
# Import from vendor-specific modules
from .local import get_YFin_data, get_finnhub_news, get_finnhub_company_insider_sentiment, get_finnhub_company_insider_transactions, get_simfin_balance_sheet, get_simfin_cashflow, get_simfin_income_statements, get_reddit_global_news, get_reddit_company_news
from .y_finance import get_YFin_data_online, get_stock_stats_indicators_window, get_balance_sheet as get_yfinance_balance_sheet, get_cashflow as get_yfinance_cashflow, get_income_statement as get_yfinance_income_statement, get_insider_transactions as get_yfinance_insider_transactions
from .google import get_google_news
from .openai import get_stock_news_openai, get_global_news_openai, get_fundamentals_openai
from .google import get_google_news, get_bulk_news_google
from .openai import get_stock_news_openai, get_global_news_openai, get_fundamentals_openai, get_bulk_news_openai
from .alpha_vantage import (
get_stock as get_alpha_vantage_stock,
get_indicator as get_alpha_vantage_indicator,
@ -15,12 +16,13 @@ from .alpha_vantage import (
get_insider_transactions as get_alpha_vantage_insider_transactions,
get_news as get_alpha_vantage_news
)
from .alpha_vantage_news import get_bulk_news_alpha_vantage
from .alpha_vantage_common import AlphaVantageRateLimitError
# Configuration and routing logic
from .config import get_config
# Tools organized by category
from tradingagents.agents.discovery import NewsArticle
TOOLS_CATEGORIES = {
"core_stock_apis": {
"description": "OHLCV stock price data",
@ -50,6 +52,7 @@ TOOLS_CATEGORIES = {
"get_global_news",
"get_insider_sentiment",
"get_insider_transactions",
"get_bulk_news",
]
}
}
@ -61,21 +64,17 @@ VENDOR_LIST = [
"google"
]
# Mapping of methods to their vendor-specific implementations
VENDOR_METHODS = {
# core_stock_apis
"get_stock_data": {
"alpha_vantage": get_alpha_vantage_stock,
"yfinance": get_YFin_data_online,
"local": get_YFin_data,
},
# technical_indicators
"get_indicators": {
"alpha_vantage": get_alpha_vantage_indicator,
"yfinance": get_stock_stats_indicators_window,
"local": get_stock_stats_indicators_window
},
# fundamental_data
"get_fundamentals": {
"alpha_vantage": get_alpha_vantage_fundamentals,
"openai": get_fundamentals_openai,
@ -95,7 +94,6 @@ VENDOR_METHODS = {
"yfinance": get_yfinance_income_statement,
"local": get_simfin_income_statements,
},
# news_data
"get_news": {
"alpha_vantage": get_alpha_vantage_news,
"openai": get_stock_news_openai,
@ -114,56 +112,162 @@ VENDOR_METHODS = {
"yfinance": get_yfinance_insider_transactions,
"local": get_finnhub_company_insider_transactions,
},
"get_bulk_news": {
"alpha_vantage": get_bulk_news_alpha_vantage,
"openai": get_bulk_news_openai,
"google": get_bulk_news_google,
},
}
CACHE_TTL_SECONDS = 300
_bulk_news_cache: Dict[str, Dict[str, Any]] = {}
_bulk_news_cache_lock = threading.Lock()
def parse_lookback_period(lookback: str) -> int:
lookback = lookback.lower().strip()
if lookback == "1h":
return 1
elif lookback == "6h":
return 6
elif lookback == "24h":
return 24
elif lookback == "7d":
return 168
else:
raise ValueError(f"Invalid lookback period: {lookback}. Valid values: 1h, 6h, 24h, 7d")
def _get_cached_bulk_news(lookback_period: str) -> Optional[List[NewsArticle]]:
cache_key = lookback_period
with _bulk_news_cache_lock:
if cache_key in _bulk_news_cache:
cached = _bulk_news_cache[cache_key]
cached_time = cached.get("timestamp")
if cached_time and (datetime.now() - cached_time).total_seconds() < CACHE_TTL_SECONDS:
return cached.get("articles")
return None
def _set_cached_bulk_news(lookback_period: str, articles: List[NewsArticle]) -> None:
cache_key = lookback_period
with _bulk_news_cache_lock:
_bulk_news_cache[cache_key] = {
"timestamp": datetime.now(),
"articles": articles,
}
def _convert_to_news_articles(raw_articles: List[Dict[str, Any]]) -> List[NewsArticle]:
articles = []
for item in raw_articles:
try:
published_at_str = item.get("published_at", "")
if isinstance(published_at_str, str):
try:
published_at = datetime.fromisoformat(published_at_str.replace("Z", "+00:00"))
except ValueError:
published_at = datetime.now()
elif isinstance(published_at_str, datetime):
published_at = published_at_str
else:
published_at = datetime.now()
article = NewsArticle(
title=item.get("title", ""),
source=item.get("source", ""),
url=item.get("url", ""),
published_at=published_at,
content_snippet=item.get("content_snippet", ""),
ticker_mentions=[],
)
articles.append(article)
except Exception:
continue
return articles
def _fetch_bulk_news_from_vendor(lookback_period: str) -> List[Dict[str, Any]]:
lookback_hours = parse_lookback_period(lookback_period)
vendor_order = ["alpha_vantage", "openai", "google"]
for vendor in vendor_order:
if vendor not in VENDOR_METHODS["get_bulk_news"]:
continue
vendor_func = VENDOR_METHODS["get_bulk_news"][vendor]
try:
print(f"DEBUG: Attempting bulk news from vendor '{vendor}'...")
result = vendor_func(lookback_hours)
if result:
print(f"SUCCESS: Got {len(result)} articles from vendor '{vendor}'")
return result
print(f"DEBUG: Vendor '{vendor}' returned empty results, trying next...")
except AlphaVantageRateLimitError as e:
print(f"RATE_LIMIT: Alpha Vantage rate limit exceeded: {e}")
continue
except Exception as e:
print(f"FAILED: Vendor '{vendor}' failed: {e}")
continue
return []
def get_bulk_news(lookback_period: str = "24h") -> List[NewsArticle]:
cached = _get_cached_bulk_news(lookback_period)
if cached is not None:
print(f"DEBUG: Returning cached bulk news for period '{lookback_period}'")
return cached
raw_articles = _fetch_bulk_news_from_vendor(lookback_period)
articles = _convert_to_news_articles(raw_articles)
_set_cached_bulk_news(lookback_period, articles)
return articles
def get_category_for_method(method: str) -> str:
"""Get the category that contains the specified method."""
for category, info in TOOLS_CATEGORIES.items():
if method in info["tools"]:
return category
raise ValueError(f"Method '{method}' not found in any category")
def get_vendor(category: str, method: str = None) -> str:
"""Get the configured vendor for a data category or specific tool method.
Tool-level configuration takes precedence over category-level.
"""
config = get_config()
# Check tool-level configuration first (if method provided)
if method:
tool_vendors = config.get("tool_vendors", {})
if method in tool_vendors:
return tool_vendors[method]
# Fall back to category-level configuration
return config.get("data_vendors", {}).get(category, "default")
def route_to_vendor(method: str, *args, **kwargs):
"""Route method calls to appropriate vendor implementation with fallback support."""
category = get_category_for_method(method)
vendor_config = get_vendor(category, method)
# Handle comma-separated vendors
primary_vendors = [v.strip() for v in vendor_config.split(',')]
if method not in VENDOR_METHODS:
raise ValueError(f"Method '{method}' not supported")
# Get all available vendors for this method for fallback
all_available_vendors = list(VENDOR_METHODS[method].keys())
# Create fallback vendor list: primary vendors first, then remaining vendors as fallbacks
fallback_vendors = primary_vendors.copy()
for vendor in all_available_vendors:
if vendor not in fallback_vendors:
fallback_vendors.append(vendor)
# Debug: Print fallback ordering
primary_str = " → ".join(primary_vendors)
fallback_str = " → ".join(fallback_vendors)
primary_str = " -> ".join(primary_vendors)
fallback_str = " -> ".join(fallback_vendors)
print(f"DEBUG: {method} - Primary: [{primary_str}] | Full fallback order: [{fallback_str}]")
# Track results and execution state
results = []
vendor_attempt_count = 0
any_primary_vendor_attempted = False
@ -179,22 +283,18 @@ def route_to_vendor(method: str, *args, **kwargs):
is_primary_vendor = vendor in primary_vendors
vendor_attempt_count += 1
# Track if we attempted any primary vendor
if is_primary_vendor:
any_primary_vendor_attempted = True
# Debug: Print current attempt
vendor_type = "PRIMARY" if is_primary_vendor else "FALLBACK"
print(f"DEBUG: Attempting {vendor_type} vendor '{vendor}' for {method} (attempt #{vendor_attempt_count})")
# Handle list of methods for a vendor
if isinstance(vendor_impl, list):
vendor_methods = [(impl, vendor) for impl in vendor_impl]
print(f"DEBUG: Vendor '{vendor}' has multiple implementations: {len(vendor_methods)} functions")
else:
vendor_methods = [(vendor_impl, vendor)]
# Run methods for this vendor
vendor_results = []
for impl_func, vendor_name in vendor_methods:
try:
@ -202,43 +302,35 @@ def route_to_vendor(method: str, *args, **kwargs):
result = impl_func(*args, **kwargs)
vendor_results.append(result)
print(f"SUCCESS: {impl_func.__name__} from vendor '{vendor_name}' completed successfully")
except AlphaVantageRateLimitError as e:
if vendor == "alpha_vantage":
print(f"RATE_LIMIT: Alpha Vantage rate limit exceeded, falling back to next available vendor")
print(f"DEBUG: Rate limit details: {e}")
# Continue to next vendor for fallback
continue
except Exception as e:
# Log error but continue with other implementations
print(f"FAILED: {impl_func.__name__} from vendor '{vendor_name}' failed: {e}")
continue
# Add this vendor's results
if vendor_results:
results.extend(vendor_results)
successful_vendor = vendor
result_summary = f"Got {len(vendor_results)} result(s)"
print(f"SUCCESS: Vendor '{vendor}' succeeded - {result_summary}")
# Stopping logic: Stop after first successful vendor for single-vendor configs
# Multiple vendor configs (comma-separated) may want to collect from multiple sources
if len(primary_vendors) == 1:
print(f"DEBUG: Stopping after successful vendor '{vendor}' (single-vendor config)")
break
else:
print(f"FAILED: Vendor '{vendor}' produced no results")
# Final result summary
if not results:
print(f"FAILURE: All {vendor_attempt_count} vendor attempts failed for method '{method}'")
raise RuntimeError(f"All vendor implementations failed for method '{method}'")
else:
print(f"FINAL: Method '{method}' completed with {len(results)} result(s) from {vendor_attempt_count} vendor attempt(s)")
# Return single result if only one, otherwise concatenate as string
if len(results) == 1:
return results[0]
else:
# Convert all results to strings and concatenate
return '\n'.join(str(result) for result in results)
return '\n'.join(str(result) for result in results)

118
tradingagents/dataflows/openai.py
View File

@ -1,7 +1,30 @@
import json
import re
from datetime import datetime, timedelta
from typing import List, Dict, Any, Optional
from openai import OpenAI
from .config import get_config
def _extract_response_text(response) -> Optional[str]:
if not hasattr(response, 'output') or not response.output:
return None
for output_item in response.output:
if not hasattr(output_item, 'content') or not output_item.content:
continue
text_pieces = []
for content_item in output_item.content:
if hasattr(content_item, 'text') and content_item.text:
text_pieces.append(content_item.text)
if text_pieces:
return "\n".join(text_pieces)
return None
def get_stock_news_openai(query, start_date, end_date):
config = get_config()
client = OpenAI(base_url=config["backend_url"])
@ -34,7 +57,7 @@ def get_stock_news_openai(query, start_date, end_date):
store=True,
)
return response.output[1].content[0].text
return _extract_response_text(response) or ""
def get_global_news_openai(curr_date, look_back_days=7, limit=5):
@ -69,7 +92,7 @@ def get_global_news_openai(curr_date, look_back_days=7, limit=5):
store=True,
)
return response.output[1].content[0].text
return _extract_response_text(response) or ""
def get_fundamentals_openai(ticker, curr_date):
@ -104,4 +127,93 @@ def get_fundamentals_openai(ticker, curr_date):
store=True,
)
return response.output[1].content[0].text
return _extract_response_text(response) or ""
def get_bulk_news_openai(lookback_hours: int) -> List[Dict[str, Any]]:
config = get_config()
client = OpenAI(base_url=config["backend_url"])
end_date = datetime.now()
start_date = end_date - timedelta(hours=lookback_hours)
start_str = start_date.strftime("%Y-%m-%d %H:%M")
end_str = end_date.strftime("%Y-%m-%d %H:%M")
prompt = f"""Search for recent stock market news, trading news, and earnings announcements from {start_str} to {end_str}.
Return the results as a JSON array with the following structure:
[
{{
"title": "Article title",
"source": "Source name",
"url": "https://...",
"published_at": "YYYY-MM-DDTHH:MM:SS",
"content_snippet": "Brief summary of the article..."
}}
]
Focus on:
- Stock market movements and trends
- Company earnings reports
- Mergers and acquisitions
- Significant trading activity
- Economic news affecting markets
Return ONLY the JSON array, no additional text."""
response = client.responses.create(
model=config["quick_think_llm"],
input=[
{
"role": "system",
"content": [
{
"type": "input_text",
"text": prompt,
}
],
}
],
text={"format": {"type": "text"}},
reasoning={},
tools=[
{
"type": "web_search_preview",
"user_location": {"type": "approximate"},
"search_context_size": "medium",
}
],
temperature=0.5,
max_output_tokens=8192,
top_p=1,
store=True,
)
try:
response_text = _extract_response_text(response)
if not response_text:
return []
json_match = re.search(r'\[[\s\S]*\]', response_text)
if json_match:
articles = json.loads(json_match.group())
else:
articles = json.loads(response_text)
result = []
for item in articles:
if isinstance(item, dict):
article = {
"title": item.get("title", ""),
"source": item.get("source", "Web Search"),
"url": item.get("url", ""),
"published_at": item.get("published_at", datetime.now().isoformat()),
"content_snippet": item.get("content_snippet", "")[:500],
}
result.append(article)
return result
except (json.JSONDecodeError, AttributeError):
return []

21
tradingagents/dataflows/trending/__init__.py Normal file
View File

@ -0,0 +1,21 @@
from .stock_resolver import (
resolve_ticker,
validate_tradeable,
validate_us_ticker,
COMPANY_TO_TICKER,
)
from .sector_classifier import (
classify_sector,
TICKER_TO_SECTOR,
VALID_SECTORS,
)
__all__ = [
"resolve_ticker",
"validate_tradeable",
"validate_us_ticker",
"COMPANY_TO_TICKER",
"classify_sector",
"TICKER_TO_SECTOR",
"VALID_SECTORS",
]

267
tradingagents/dataflows/trending/sector_classifier.py Normal file
View File

@ -0,0 +1,267 @@
import logging
from typing import Dict
logger = logging.getLogger(__name__)
VALID_SECTORS = {
"technology",
"healthcare",
"finance",
"energy",
"consumer_goods",
"industrials",
"other",
}
TICKER_TO_SECTOR: Dict[str, str] = {
"AAPL": "technology",
"MSFT": "technology",
"GOOGL": "technology",
"GOOG": "technology",
"AMZN": "technology",
"META": "technology",
"NVDA": "technology",
"TSLA": "technology",
"AMD": "technology",
"INTC": "technology",
"QCOM": "technology",
"AVGO": "technology",
"TXN": "technology",
"ADBE": "technology",
"CRM": "technology",
"CSCO": "technology",
"NFLX": "technology",
"ORCL": "technology",
"IBM": "technology",
"NOW": "technology",
"INTU": "technology",
"ADSK": "technology",
"SNPS": "technology",
"CDNS": "technology",
"PLTR": "technology",
"SNOW": "technology",
"DDOG": "technology",
"CRWD": "technology",
"OKTA": "technology",
"NET": "technology",
"MDB": "technology",
"TWLO": "technology",
"WDAY": "technology",
"SPLK": "technology",
"VMW": "technology",
"HPQ": "technology",
"DELL": "technology",
"FTNT": "technology",
"PANW": "technology",
"ZS": "technology",
"S": "technology",
"VEEV": "technology",
"ZM": "technology",
"DOCU": "technology",
"ASAN": "technology",
"MNDY": "technology",
"TEAM": "technology",
"ANSS": "technology",
"ROP": "technology",
"JPM": "finance",
"BAC": "finance",
"WFC": "finance",
"GS": "finance",
"MS": "finance",
"C": "finance",
"BLK": "finance",
"SCHW": "finance",
"AXP": "finance",
"V": "finance",
"MA": "finance",
"PYPL": "finance",
"SQ": "finance",
"COIN": "finance",
"HOOD": "finance",
"SOFI": "finance",
"AFRM": "finance",
"MQ": "finance",
"BRK-B": "finance",
"BRK-A": "finance",
"JNJ": "healthcare",
"UNH": "healthcare",
"PFE": "healthcare",
"ABBV": "healthcare",
"MRK": "healthcare",
"LLY": "healthcare",
"MRNA": "healthcare",
"BNTX": "healthcare",
"CVS": "healthcare",
"WBA": "healthcare",
"MCK": "healthcare",
"CAH": "healthcare",
"HUM": "healthcare",
"CI": "healthcare",
"ELV": "healthcare",
"XOM": "energy",
"CVX": "energy",
"COP": "energy",
"SLB": "energy",
"HAL": "energy",
"BKR": "energy",
"MPC": "energy",
"VLO": "energy",
"PSX": "energy",
"OXY": "energy",
"PXD": "energy",
"DVN": "energy",
"CEG": "energy",
"NEE": "energy",
"DUK": "energy",
"SO": "energy",
"D": "energy",
"SRE": "energy",
"WMT": "consumer_goods",
"COST": "consumer_goods",
"TGT": "consumer_goods",
"HD": "consumer_goods",
"LOW": "consumer_goods",
"PG": "consumer_goods",
"KO": "consumer_goods",
"PEP": "consumer_goods",
"NKE": "consumer_goods",
"SBUX": "consumer_goods",
"MCD": "consumer_goods",
"CMG": "consumer_goods",
"YUM": "consumer_goods",
"DPZ": "consumer_goods",
"DIS": "consumer_goods",
"CMCSA": "consumer_goods",
"VZ": "consumer_goods",
"T": "consumer_goods",
"TMUS": "consumer_goods",
"EL": "consumer_goods",
"CL": "consumer_goods",
"KMB": "consumer_goods",
"CLX": "consumer_goods",
"KHC": "consumer_goods",
"GIS": "consumer_goods",
"K": "consumer_goods",
"MDLZ": "consumer_goods",
"HSY": "consumer_goods",
"TSN": "consumer_goods",
"BYND": "consumer_goods",
"CAG": "consumer_goods",
"STZ": "consumer_goods",
"BUD": "consumer_goods",
"DEO": "consumer_goods",
"PM": "consumer_goods",
"MO": "consumer_goods",
"LULU": "consumer_goods",
"DG": "consumer_goods",
"DLTR": "consumer_goods",
"ROST": "consumer_goods",
"TJX": "consumer_goods",
"AZO": "consumer_goods",
"ORLY": "consumer_goods",
"KMX": "consumer_goods",
"ADDYY": "consumer_goods",
"UBER": "consumer_goods",
"LYFT": "consumer_goods",
"ABNB": "consumer_goods",
"DASH": "consumer_goods",
"SNAP": "consumer_goods",
"PINS": "consumer_goods",
"TWTR": "consumer_goods",
"SHOP": "consumer_goods",
"TOST": "consumer_goods",
"BA": "industrials",
"LMT": "industrials",
"RTX": "industrials",
"GD": "industrials",
"NOC": "industrials",
"GE": "industrials",
"HON": "industrials",
"MMM": "industrials",
"CAT": "industrials",
"DE": "industrials",
"UNP": "industrials",
"UPS": "industrials",
"FDX": "industrials",
"DAL": "industrials",
"UAL": "industrials",
"AAL": "industrials",
"LUV": "industrials",
"F": "industrials",
"GM": "industrials",
"TM": "industrials",
"HMC": "industrials",
"VWAGY": "industrials",
"RACE": "industrials",
"RIVN": "industrials",
"LCID": "industrials",
"NIO": "industrials",
"LNVGY": "industrials",
}
_sector_cache: Dict[str, str] = {}
def _llm_classify_sector(ticker: str) -> str:
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from tradingagents.default_config import DEFAULT_CONFIG
llm_name = DEFAULT_CONFIG.get("quick_think_llm", "gpt-4o-mini")
llm_provider = DEFAULT_CONFIG.get("llm_provider", "openai")
backend_url = DEFAULT_CONFIG.get("backend_url", "https://api.openai.com/v1")
llm = ChatOpenAI(
model=llm_name,
base_url=backend_url,
temperature=0,
)
system_prompt = (
"You are a financial sector classifier. Given a stock ticker symbol, "
"classify it into exactly one of the following sectors: "
"technology, healthcare, finance, energy, consumer_goods, industrials, other. "
"Respond with only the sector name in lowercase, nothing else."
)
user_prompt = f"Classify the stock ticker: {ticker}"
messages = [
SystemMessage(content=system_prompt),
HumanMessage(content=user_prompt),
]
response = llm.invoke(messages)
sector = response.content.strip().lower()
if sector not in VALID_SECTORS:
logger.warning(
"LLM returned invalid sector '%s' for ticker %s, defaulting to 'other'",
sector,
ticker,
)
return "other"
return sector
def classify_sector(ticker: str) -> str:
ticker_upper = ticker.upper()
if ticker_upper in TICKER_TO_SECTOR:
return TICKER_TO_SECTOR[ticker_upper]
if ticker_upper in _sector_cache:
return _sector_cache[ticker_upper]
logger.info("Using LLM fallback for sector classification of ticker: %s", ticker)
try:
sector = _llm_classify_sector(ticker_upper)
_sector_cache[ticker_upper] = sector
logger.info("Classified %s as %s via LLM", ticker, sector)
return sector
except Exception as e:
logger.error("LLM sector classification failed for %s: %s", ticker, str(e))
_sector_cache[ticker_upper] = "other"
return "other"

536
tradingagents/dataflows/trending/stock_resolver.py Normal file
View File

@ -0,0 +1,536 @@
import logging
import re
from typing import Optional
import yfinance as yf
logger = logging.getLogger(__name__)
COMPANY_TO_TICKER = {
"apple": "AAPL",
"apple inc": "AAPL",
"apple inc.": "AAPL",
"apple corporation": "AAPL",
"the iphone maker": "AAPL",
"iphone maker": "AAPL",
"microsoft": "MSFT",
"microsoft inc": "MSFT",
"microsoft inc.": "MSFT",
"microsoft corp": "MSFT",
"microsoft corp.": "MSFT",
"microsoft corporation": "MSFT",
"google": "GOOGL",
"alphabet": "GOOGL",
"alphabet inc": "GOOGL",
"alphabet inc.": "GOOGL",
"the search giant": "GOOGL",
"amazon": "AMZN",
"amazon inc": "AMZN",
"amazon inc.": "AMZN",
"amazon.com": "AMZN",
"amazon.com inc": "AMZN",
"the e-commerce giant": "AMZN",
"e-commerce giant": "AMZN",
"meta": "META",
"meta platforms": "META",
"meta platforms inc": "META",
"meta platforms inc.": "META",
"facebook": "META",
"facebook inc": "META",
"facebook inc.": "META",
"tesla": "TSLA",
"tesla inc": "TSLA",
"tesla inc.": "TSLA",
"tesla motors": "TSLA",
"ev maker tesla": "TSLA",
"nvidia": "NVDA",
"nvidia corp": "NVDA",
"nvidia corp.": "NVDA",
"nvidia corporation": "NVDA",
"berkshire hathaway": "BRK-B",
"berkshire": "BRK-B",
"jpmorgan": "JPM",
"jpmorgan chase": "JPM",
"jp morgan": "JPM",
"jp morgan chase": "JPM",
"johnson & johnson": "JNJ",
"johnson and johnson": "JNJ",
"j&j": "JNJ",
"unitedhealth": "UNH",
"unitedhealth group": "UNH",
"visa": "V",
"visa inc": "V",
"visa inc.": "V",
"procter & gamble": "PG",
"procter and gamble": "PG",
"p&g": "PG",
"mastercard": "MA",
"mastercard inc": "MA",
"mastercard inc.": "MA",
"home depot": "HD",
"the home depot": "HD",
"chevron": "CVX",
"chevron corp": "CVX",
"chevron corporation": "CVX",
"exxon": "XOM",
"exxon mobil": "XOM",
"exxonmobil": "XOM",
"pfizer": "PFE",
"pfizer inc": "PFE",
"pfizer inc.": "PFE",
"abbvie": "ABBV",
"abbvie inc": "ABBV",
"abbvie inc.": "ABBV",
"coca-cola": "KO",
"coca cola": "KO",
"coke": "KO",
"the coca-cola company": "KO",
"pepsico": "PEP",
"pepsi": "PEP",
"pepsi co": "PEP",
"costco": "COST",
"costco wholesale": "COST",
"walmart": "WMT",
"wal-mart": "WMT",
"walmart inc": "WMT",
"bank of america": "BAC",
"bofa": "BAC",
"merck": "MRK",
"merck & co": "MRK",
"merck and co": "MRK",
"eli lilly": "LLY",
"lilly": "LLY",
"eli lilly and company": "LLY",
"adobe": "ADBE",
"adobe inc": "ADBE",
"adobe inc.": "ADBE",
"adobe systems": "ADBE",
"salesforce": "CRM",
"salesforce inc": "CRM",
"salesforce.com": "CRM",
"cisco": "CSCO",
"cisco systems": "CSCO",
"cisco systems inc": "CSCO",
"netflix": "NFLX",
"netflix inc": "NFLX",
"netflix inc.": "NFLX",
"oracle": "ORCL",
"oracle corp": "ORCL",
"oracle corporation": "ORCL",
"intel": "INTC",
"intel corp": "INTC",
"intel corporation": "INTC",
"amd": "AMD",
"advanced micro devices": "AMD",
"qualcomm": "QCOM",
"qualcomm inc": "QCOM",
"qualcomm inc.": "QCOM",
"broadcom": "AVGO",
"broadcom inc": "AVGO",
"broadcom inc.": "AVGO",
"texas instruments": "TXN",
"ti": "TXN",
"disney": "DIS",
"walt disney": "DIS",
"the walt disney company": "DIS",
"walt disney company": "DIS",
"comcast": "CMCSA",
"comcast corp": "CMCSA",
"comcast corporation": "CMCSA",
"verizon": "VZ",
"verizon communications": "VZ",
"at&t": "T",
"att": "T",
"t-mobile": "TMUS",
"tmobile": "TMUS",
"t-mobile us": "TMUS",
"american express": "AXP",
"amex": "AXP",
"goldman sachs": "GS",
"goldman": "GS",
"morgan stanley": "MS",
"wells fargo": "WFC",
"wells": "WFC",
"citigroup": "C",
"citi": "C",
"citibank": "C",
"charles schwab": "SCHW",
"schwab": "SCHW",
"blackrock": "BLK",
"blackrock inc": "BLK",
"paypal": "PYPL",
"paypal holdings": "PYPL",
"paypal inc": "PYPL",
"square": "SQ",
"block": "SQ",
"block inc": "SQ",
"shopify": "SHOP",
"shopify inc": "SHOP",
"uber": "UBER",
"uber technologies": "UBER",
"lyft": "LYFT",
"lyft inc": "LYFT",
"airbnb": "ABNB",
"airbnb inc": "ABNB",
"doordash": "DASH",
"doordash inc": "DASH",
"snap": "SNAP",
"snap inc": "SNAP",
"snapchat": "SNAP",
"pinterest": "PINS",
"pinterest inc": "PINS",
"linkedin": "MSFT",
"zoom": "ZM",
"zoom video": "ZM",
"zoom video communications": "ZM",
"slack": "CRM",
"slack technologies": "CRM",
"palantir": "PLTR",
"palantir technologies": "PLTR",
"snowflake": "SNOW",
"snowflake inc": "SNOW",
"datadog": "DDOG",
"datadog inc": "DDOG",
"crowdstrike": "CRWD",
"crowdstrike holdings": "CRWD",
"okta": "OKTA",
"okta inc": "OKTA",
"cloudflare": "NET",
"cloudflare inc": "NET",
"mongodb": "MDB",
"mongodb inc": "MDB",
"twilio": "TWLO",
"twilio inc": "TWLO",
"servicenow": "NOW",
"servicenow inc": "NOW",
"workday": "WDAY",
"workday inc": "WDAY",
"splunk": "SPLK",
"splunk inc": "SPLK",
"vmware": "VMW",
"vmware inc": "VMW",
"ibm": "IBM",
"international business machines": "IBM",
"hp": "HPQ",
"hewlett-packard": "HPQ",
"hewlett packard": "HPQ",
"dell": "DELL",
"dell technologies": "DELL",
"lenovo": "LNVGY",
"boeing": "BA",
"boeing company": "BA",
"the boeing company": "BA",
"lockheed martin": "LMT",
"lockheed": "LMT",
"raytheon": "RTX",
"rtx": "RTX",
"general dynamics": "GD",
"northrop grumman": "NOC",
"northrop": "NOC",
"general electric": "GE",
"ge": "GE",
"honeywell": "HON",
"honeywell international": "HON",
"3m": "MMM",
"3m company": "MMM",
"caterpillar": "CAT",
"caterpillar inc": "CAT",
"deere": "DE",
"john deere": "DE",
"deere & company": "DE",
"union pacific": "UNP",
"ups": "UPS",
"united parcel service": "UPS",
"fedex": "FDX",
"federal express": "FDX",
"delta": "DAL",
"delta air lines": "DAL",
"delta airlines": "DAL",
"united airlines": "UAL",
"united": "UAL",
"american airlines": "AAL",
"southwest": "LUV",
"southwest airlines": "LUV",
"ford": "F",
"ford motor": "F",
"ford motor company": "F",
"general motors": "GM",
"gm": "GM",
"toyota": "TM",
"toyota motor": "TM",
"honda": "HMC",
"honda motor": "HMC",
"volkswagen": "VWAGY",
"vw": "VWAGY",
"ferrari": "RACE",
"rivian": "RIVN",
"rivian automotive": "RIVN",
"lucid": "LCID",
"lucid motors": "LCID",
"lucid group": "LCID",
"nio": "NIO",
"nio inc": "NIO",
"moderna": "MRNA",
"moderna inc": "MRNA",
"biontech": "BNTX",
"cvs": "CVS",
"cvs health": "CVS",
"walgreens": "WBA",
"walgreens boots alliance": "WBA",
"mckesson": "MCK",
"mckesson corp": "MCK",
"cardinal health": "CAH",
"humana": "HUM",
"humana inc": "HUM",
"cigna": "CI",
"cigna group": "CI",
"anthem": "ELV",
"elevance health": "ELV",
"starbucks": "SBUX",
"starbucks corp": "SBUX",
"starbucks corporation": "SBUX",
"mcdonalds": "MCD",
"mcdonald's": "MCD",
"chipotle": "CMG",
"chipotle mexican grill": "CMG",
"yum brands": "YUM",
"yum": "YUM",
"dominos": "DPZ",
"domino's": "DPZ",
"domino's pizza": "DPZ",
"nike": "NKE",
"nike inc": "NKE",
"adidas": "ADDYY",
"lululemon": "LULU",
"lululemon athletica": "LULU",
"target": "TGT",
"target corp": "TGT",
"target corporation": "TGT",
"dollar general": "DG",
"dollar tree": "DLTR",
"ross stores": "ROST",
"ross": "ROST",
"tjx": "TJX",
"tjx companies": "TJX",
"tj maxx": "TJX",
"lowes": "LOW",
"lowe's": "LOW",
"lowe's companies": "LOW",
"autozone": "AZO",
"o'reilly": "ORLY",
"o'reilly automotive": "ORLY",
"carmax": "KMX",
"estee lauder": "EL",
"colgate": "CL",
"colgate-palmolive": "CL",
"colgate palmolive": "CL",
"kimberly-clark": "KMB",
"kimberly clark": "KMB",
"clorox": "CLX",
"clorox company": "CLX",
"kraft heinz": "KHC",
"kraft": "KHC",
"heinz": "KHC",
"general mills": "GIS",
"kellogg": "K",
"kellogg's": "K",
"mondelez": "MDLZ",
"mondelez international": "MDLZ",
"hershey": "HSY",
"the hershey company": "HSY",
"tyson": "TSN",
"tyson foods": "TSN",
"beyond meat": "BYND",
"conagra": "CAG",
"conagra brands": "CAG",
"constellation brands": "STZ",
"anheuser-busch": "BUD",
"anheuser busch": "BUD",
"ab inbev": "BUD",
"diageo": "DEO",
"philip morris": "PM",
"philip morris international": "PM",
"altria": "MO",
"altria group": "MO",
"constellation energy": "CEG",
"nextera": "NEE",
"nextera energy": "NEE",
"duke energy": "DUK",
"southern company": "SO",
"dominion": "D",
"dominion energy": "D",
"sempra": "SRE",
"sempra energy": "SRE",
"conocophillips": "COP",
"conoco": "COP",
"schlumberger": "SLB",
"halliburton": "HAL",
"baker hughes": "BKR",
"marathon": "MPC",
"marathon petroleum": "MPC",
"valero": "VLO",
"valero energy": "VLO",
"phillips 66": "PSX",
"occidental": "OXY",
"occidental petroleum": "OXY",
"pioneer": "PXD",
"pioneer natural resources": "PXD",
"devon energy": "DVN",
"devon": "DVN",
"coinbase": "COIN",
"coinbase global": "COIN",
"robinhood": "HOOD",
"robinhood markets": "HOOD",
"sofi": "SOFI",
"sofi technologies": "SOFI",
"affirm": "AFRM",
"affirm holdings": "AFRM",
"marqeta": "MQ",
"toast": "TOST",
"toast inc": "TOST",
"docusign": "DOCU",
"docusign inc": "DOCU",
"asana": "ASAN",
"monday.com": "MNDY",
"monday": "MNDY",
"atlassian": "TEAM",
"atlassian corp": "TEAM",
"intuit": "INTU",
"intuit inc": "INTU",
"autodesk": "ADSK",
"autodesk inc": "ADSK",
"synopsys": "SNPS",
"cadence": "CDNS",
"cadence design": "CDNS",
"ansys": "ANSS",
"roper": "ROP",
"roper technologies": "ROP",
"fortinet": "FTNT",
"palo alto": "PANW",
"palo alto networks": "PANW",
"zscaler": "ZS",
"sentinelone": "S",
"veeva": "VEEV",
"veeva systems": "VEEV",
}
US_EXCHANGE_CODES = {
"NYQ",
"NMS",
"NGM",
"NCM",
"ASE",
"PCX",
"BTS",
"NYSE",
"NASDAQ",
"AMEX",
"NYS",
"NAS",
"NIM",
"NAQ",
}
SUFFIX_PATTERNS = [
r"\s+inc\.?$",
r"\s+corp\.?$",
r"\s+corporation$",
r"\s+co\.?$",
r"\s+company$",
r"\s+llc$",
r"\s+ltd\.?$",
r"\s+limited$",
r"\s+plc$",
r"\s+holdings?$",
r"\s+group$",
r"\s+technologies$",
r"\s+enterprises?$",
]
def _normalize_company_name(name: str) -> str:
normalized = name.lower().strip()
for pattern in SUFFIX_PATTERNS:
normalized = re.sub(pattern, "", normalized, flags=re.IGNORECASE)
normalized = normalized.strip()
return normalized
def _search_yfinance_ticker(company_name: str) -> Optional[str]:
try:
search_result = yf.Ticker(company_name)
info = search_result.info
if info and "symbol" in info:
return info["symbol"]
except Exception as e:
logger.debug("yfinance search failed for %s: %s", company_name, str(e))
try:
search = yf.Search(company_name, max_results=5)
if hasattr(search, "quotes") and search.quotes:
for quote in search.quotes:
if "symbol" in quote:
return quote["symbol"]
except Exception as e:
logger.debug("yfinance Search failed for %s: %s", company_name, str(e))
return None
def validate_us_ticker(ticker: str) -> bool:
try:
ticker_obj = yf.Ticker(ticker.upper())
info = ticker_obj.info
if not info:
logger.warning("Validation failed for %s: no info available", ticker)
return False
exchange = info.get("exchange", "")
if exchange in US_EXCHANGE_CODES:
return True
exchange_lower = exchange.lower()
if any(us_ex.lower() in exchange_lower for us_ex in ["nyse", "nasdaq", "amex", "nys", "nms", "ngm"]):
return True
logger.warning("Validation failed for %s: exchange %s is not a US exchange", ticker, exchange)
return False
except Exception as e:
logger.warning("Validation failed for %s: %s", ticker, str(e))
return False
def resolve_ticker(company_name: str) -> Optional[str]:
if not company_name or not company_name.strip():
return None
normalized = company_name.lower().strip()
if normalized in COMPANY_TO_TICKER:
return COMPANY_TO_TICKER[normalized]
normalized_stripped = _normalize_company_name(company_name)
if normalized_stripped in COMPANY_TO_TICKER:
return COMPANY_TO_TICKER[normalized_stripped]
if company_name.upper() in [v for v in COMPANY_TO_TICKER.values()]:
if validate_us_ticker(company_name.upper()):
return company_name.upper()
logger.info("Using yfinance fallback for company: %s", company_name)
yf_ticker = _search_yfinance_ticker(company_name)
if yf_ticker:
if validate_us_ticker(yf_ticker):
logger.info("Resolved %s to %s via yfinance", company_name, yf_ticker)
return yf_ticker
else:
logger.warning("Ticker %s for %s failed US exchange validation", yf_ticker, company_name)
return None
logger.warning("Could not resolve ticker for company: %s", company_name)
return None
def validate_tradeable(ticker: str) -> bool:
return validate_us_ticker(ticker)

28
tradingagents/default_config.py
View File

@ -8,26 +8,24 @@ DEFAULT_CONFIG = {
os.path.abspath(os.path.join(os.path.dirname(__file__), ".")),
"dataflows/data_cache",
),
# LLM settings
"llm_provider": "openai",
"deep_think_llm": "o4-mini",
"quick_think_llm": "gpt-4o-mini",
"deep_think_llm": "gpt-5",
"quick_think_llm": "gpt-5-mini",
"backend_url": "https://api.openai.com/v1",
# Debate and discussion settings
"max_debate_rounds": 1,
"max_risk_discuss_rounds": 1,
"max_debate_rounds": 2,
"max_risk_discuss_rounds": 2,
"max_recur_limit": 100,
# Data vendor configuration
# Category-level configuration (default for all tools in category)
"data_vendors": {
"core_stock_apis": "yfinance", # Options: yfinance, alpha_vantage, local
"technical_indicators": "yfinance", # Options: yfinance, alpha_vantage, local
"fundamental_data": "alpha_vantage", # Options: openai, alpha_vantage, local
"news_data": "alpha_vantage", # Options: openai, alpha_vantage, google, local
"core_stock_apis": "yfinance",
"technical_indicators": "yfinance",
"fundamental_data": "alpha_vantage",
"news_data": "alpha_vantage",
},
# Tool-level configuration (takes precedence over category-level)
"tool_vendors": {
# Example: "get_stock_data": "alpha_vantage", # Override category default
# Example: "get_news": "openai", # Override category default
},
"discovery_timeout": 60,
"discovery_hard_timeout": 120,
"discovery_cache_ttl": 300,
"discovery_max_results": 20,
"discovery_min_mentions": 2,
}

163
tradingagents/graph/trading_graph.py
View File

@ -1,9 +1,9 @@
# TradingAgents/graph/trading_graph.py
import os
import signal
import threading
from pathlib import Path
import json
from datetime import date
from datetime import date, datetime
from typing import Dict, Any, Tuple, List, Optional
from langchain_openai import ChatOpenAI
@ -22,7 +22,6 @@ from tradingagents.agents.utils.agent_states import (
)
from tradingagents.dataflows.config import set_config
# Import the new abstract tool methods from agent_utils
from tradingagents.agents.utils.agent_utils import (
get_stock_data,
get_indicators,
@ -36,6 +35,19 @@ from tradingagents.agents.utils.agent_utils import (
get_global_news
)
from tradingagents.agents.discovery import (
DiscoveryRequest,
DiscoveryResult,
DiscoveryStatus,
TrendingStock,
Sector,
EventCategory,
DiscoveryTimeoutError,
extract_entities,
calculate_trending_scores,
)
from tradingagents.dataflows.interface import get_bulk_news
from .conditional_logic import ConditionalLogic
from .setup import GraphSetup
from .propagation import Propagator
@ -43,8 +55,15 @@ from .reflection import Reflector
from .signal_processing import SignalProcessor
class DiscoveryTimeoutException(Exception):
pass
def _timeout_handler(signum, frame):
raise DiscoveryTimeoutException("Discovery operation timed out")
class TradingAgentsGraph:
"""Main class that orchestrates the trading agents framework."""
def __init__(
self,
@ -52,26 +71,16 @@ class TradingAgentsGraph:
debug=False,
config: Dict[str, Any] = None,
):
"""Initialize the trading agents graph and components.
Args:
selected_analysts: List of analyst types to include
debug: Whether to run in debug mode
config: Configuration dictionary. If None, uses default config
"""
self.debug = debug
self.config = config or DEFAULT_CONFIG
# Update the interface's config
set_config(self.config)
# Create necessary directories
os.makedirs(
os.path.join(self.config["project_dir"], "dataflows/data_cache"),
exist_ok=True,
)
# Initialize LLMs
if self.config["llm_provider"].lower() == "openai" or self.config["llm_provider"] == "ollama" or self.config["llm_provider"] == "openrouter":
self.deep_thinking_llm = ChatOpenAI(model=self.config["deep_think_llm"], base_url=self.config["backend_url"])
self.quick_thinking_llm = ChatOpenAI(model=self.config["quick_think_llm"], base_url=self.config["backend_url"])
@ -83,18 +92,13 @@ class TradingAgentsGraph:
self.quick_thinking_llm = ChatGoogleGenerativeAI(model=self.config["quick_think_llm"])
else:
raise ValueError(f"Unsupported LLM provider: {self.config['llm_provider']}")
# Initialize memories
self.bull_memory = FinancialSituationMemory("bull_memory", self.config)
self.bear_memory = FinancialSituationMemory("bear_memory", self.config)
self.trader_memory = FinancialSituationMemory("trader_memory", self.config)
self.invest_judge_memory = FinancialSituationMemory("invest_judge_memory", self.config)
self.risk_manager_memory = FinancialSituationMemory("risk_manager_memory", self.config)
# Create tool nodes
self.tool_nodes = self._create_tool_nodes()
# Initialize components
self.conditional_logic = ConditionalLogic()
self.graph_setup = GraphSetup(
self.quick_thinking_llm,
@ -111,35 +115,26 @@ class TradingAgentsGraph:
self.propagator = Propagator()
self.reflector = Reflector(self.quick_thinking_llm)
self.signal_processor = SignalProcessor(self.quick_thinking_llm)
# State tracking
self.curr_state = None
self.ticker = None
self.log_states_dict = {} # date to full state dict
# Set up the graph
self.log_states_dict = {}
self.graph = self.graph_setup.setup_graph(selected_analysts)
def _create_tool_nodes(self) -> Dict[str, ToolNode]:
"""Create tool nodes for different data sources using abstract methods."""
return {
"market": ToolNode(
[
# Core stock data tools
get_stock_data,
# Technical indicators
get_indicators,
]
),
"social": ToolNode(
[
# News tools for social media analysis
get_news,
]
),
"news": ToolNode(
[
# News and insider information
get_news,
get_global_news,
get_insider_sentiment,
@ -148,7 +143,6 @@ class TradingAgentsGraph:
),
"fundamentals": ToolNode(
[
# Fundamental analysis tools
get_fundamentals,
get_balance_sheet,
get_cashflow,
@ -158,18 +152,13 @@ class TradingAgentsGraph:
}
def propagate(self, company_name, trade_date):
"""Run the trading agents graph for a company on a specific date."""
self.ticker = company_name
# Initialize state
init_agent_state = self.propagator.create_initial_state(
company_name, trade_date
)
args = self.propagator.get_graph_args()
if self.debug:
# Debug mode with tracing
trace = []
for chunk in self.graph.stream(init_agent_state, **args):
if len(chunk["messages"]) == 0:
@ -180,20 +169,14 @@ class TradingAgentsGraph:
final_state = trace[-1]
else:
# Standard mode without tracing
final_state = self.graph.invoke(init_agent_state, **args)
# Store current state for reflection
self.curr_state = final_state
# Log state
self._log_state(trade_date, final_state)
# Return decision and processed signal
return final_state, self.process_signal(final_state["final_trade_decision"])
def _log_state(self, trade_date, final_state):
"""Log the final state to a JSON file."""
self.log_states_dict[str(trade_date)] = {
"company_of_interest": final_state["company_of_interest"],
"trade_date": final_state["trade_date"],
@ -224,7 +207,6 @@ class TradingAgentsGraph:
"final_trade_decision": final_state["final_trade_decision"],
}
# Save to file
directory = Path(f"eval_results/{self.ticker}/TradingAgentsStrategy_logs/")
directory.mkdir(parents=True, exist_ok=True)
@ -235,7 +217,6 @@ class TradingAgentsGraph:
json.dump(self.log_states_dict, f, indent=4)
def reflect_and_remember(self, returns_losses):
"""Reflect on decisions and update memory based on returns."""
self.reflector.reflect_bull_researcher(
self.curr_state, returns_losses, self.bull_memory
)
@ -253,5 +234,95 @@ class TradingAgentsGraph:
)
def process_signal(self, full_signal):
"""Process a signal to extract the core decision."""
return self.signal_processor.process_signal(full_signal)
def discover_trending(
self,
request: Optional[DiscoveryRequest] = None,
) -> DiscoveryResult:
if request is None:
request = DiscoveryRequest(
lookback_period="24h",
max_results=self.config.get("discovery_max_results", 20),
)
started_at = datetime.now()
result = DiscoveryResult(
request=request,
trending_stocks=[],
status=DiscoveryStatus.PROCESSING,
started_at=started_at,
)
hard_timeout = self.config.get("discovery_hard_timeout", 120)
discovery_result = {"stocks": [], "error": None}
def run_discovery():
try:
articles = get_bulk_news(request.lookback_period)
mentions = extract_entities(articles, self.config)
min_mentions = self.config.get("discovery_min_mentions", 2)
max_results = request.max_results or self.config.get("discovery_max_results", 20)
trending_stocks = calculate_trending_scores(
mentions,
articles,
max_results=max_results,
min_mentions=min_mentions,
)
discovery_result["stocks"] = trending_stocks
except Exception as e:
discovery_result["error"] = str(e)
discovery_thread = threading.Thread(target=run_discovery)
discovery_thread.start()
discovery_thread.join(timeout=hard_timeout)
if discovery_thread.is_alive():
raise DiscoveryTimeoutError(
f"Discovery operation exceeded {hard_timeout} second timeout"
)
if discovery_result["error"]:
result.status = DiscoveryStatus.FAILED
result.error_message = discovery_result["error"]
result.completed_at = datetime.now()
return result
trending_stocks = discovery_result["stocks"]
if request.sector_filter:
sector_values = {s.value if isinstance(s, Sector) else s for s in request.sector_filter}
trending_stocks = [
stock for stock in trending_stocks
if stock.sector.value in sector_values or stock.sector in request.sector_filter
]
if request.event_filter:
event_values = {e.value if isinstance(e, EventCategory) else e for e in request.event_filter}
trending_stocks = [
stock for stock in trending_stocks
if stock.event_type.value in event_values or stock.event_type in request.event_filter
]
result.trending_stocks = trending_stocks
result.status = DiscoveryStatus.COMPLETED
result.completed_at = datetime.now()
return result
def analyze_trending(
self,
trending_stock: TrendingStock,
trade_date: Optional[date] = None,
) -> Tuple[Dict[str, Any], str]:
ticker = trending_stock.ticker
if trade_date is None:
trade_date = date.today()
return self.propagate(ticker, trade_date)