feat(dataflows): add AKShare data vendor for US and Chinese stock data (Issue #16)
This commit is contained in:
Andrew Kaszubski
parent
36de8f0470
commit
d6b9df162e
|
|
@ -0,0 +1,832 @@
|
|||
"""
|
||||
Test suite for AKShare data vendor integration.
|
||||
|
||||
This module tests:
|
||||
1. Date format conversion helper (_convert_date_format)
|
||||
2. Exponential backoff retry mechanism (_exponential_backoff_retry)
|
||||
3. US stock data retrieval (get_akshare_stock_data_us)
|
||||
4. Chinese stock data retrieval (get_akshare_stock_data_cn)
|
||||
5. Auto-market detection (get_akshare_stock_data)
|
||||
6. AKShareRateLimitError exception handling
|
||||
7. Integration with vendor routing system (interface.py)
|
||||
|
||||
Test Coverage:
|
||||
- Unit tests for individual helper functions
|
||||
- Integration tests for stock data retrieval functions
|
||||
- Edge cases (empty data, network errors, rate limits)
|
||||
- Vendor fallback behavior with rate limit errors
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import pandas as pd
|
||||
import time
|
||||
import sys
|
||||
from unittest.mock import Mock, patch, MagicMock, call
|
||||
from datetime import datetime
|
||||
from typing import Callable, Any
|
||||
|
||||
# Clear any cached imports and mock akshare before importing our modules
|
||||
if 'tradingagents.dataflows.akshare' in sys.modules:
|
||||
del sys.modules['tradingagents.dataflows.akshare']
|
||||
if 'akshare' in sys.modules:
|
||||
del sys.modules['akshare']
|
||||
|
||||
mock_akshare = MagicMock()
|
||||
sys.modules['akshare'] = mock_akshare
|
||||
|
||||
# Import modules under test
|
||||
from tradingagents.dataflows.akshare import (
|
||||
AKShareRateLimitError,
|
||||
_convert_date_format,
|
||||
_exponential_backoff_retry,
|
||||
get_akshare_stock_data_us,
|
||||
get_akshare_stock_data_cn,
|
||||
get_akshare_stock_data,
|
||||
)
|
||||
from tradingagents.dataflows.interface import route_to_vendor, VENDOR_METHODS
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Fixtures
|
||||
# ============================================================================
|
||||
|
||||
@pytest.fixture
|
||||
def sample_us_dataframe():
|
||||
"""Create a sample US stock data DataFrame matching akshare format."""
|
||||
return pd.DataFrame({
|
||||
'date': pd.date_range('2024-01-01', periods=5, freq='D'),
|
||||
'open': [150.0, 151.0, 152.0, 153.0, 154.0],
|
||||
'high': [152.0, 153.0, 154.0, 155.0, 156.0],
|
||||
'low': [149.0, 150.0, 151.0, 152.0, 153.0],
|
||||
'close': [151.0, 152.0, 153.0, 154.0, 155.0],
|
||||
'volume': [1000000, 1100000, 1200000, 1300000, 1400000],
|
||||
})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_cn_dataframe():
|
||||
"""Create a sample Chinese stock data DataFrame matching akshare format."""
|
||||
return pd.DataFrame({
|
||||
'日期': ['2024-01-01', '2024-01-02', '2024-01-03', '2024-01-04', '2024-01-05'],
|
||||
'开盘': [10.0, 10.1, 10.2, 10.3, 10.4],
|
||||
'最高': [10.2, 10.3, 10.4, 10.5, 10.6],
|
||||
'最低': [9.9, 10.0, 10.1, 10.2, 10.3],
|
||||
'收盘': [10.1, 10.2, 10.3, 10.4, 10.5],
|
||||
'成交量': [500000, 550000, 600000, 650000, 700000],
|
||||
})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_standardized_dataframe():
|
||||
"""Create a standardized DataFrame with English column names."""
|
||||
return pd.DataFrame({
|
||||
'Date': pd.date_range('2024-01-01', periods=5, freq='D'),
|
||||
'Open': [150.0, 151.0, 152.0, 153.0, 154.0],
|
||||
'High': [152.0, 153.0, 154.0, 155.0, 156.0],
|
||||
'Low': [149.0, 150.0, 151.0, 152.0, 153.0],
|
||||
'Close': [151.0, 152.0, 153.0, 154.0, 155.0],
|
||||
'Volume': [1000000, 1100000, 1200000, 1300000, 1400000],
|
||||
})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_akshare():
|
||||
"""Mock akshare module for testing."""
|
||||
with patch('tradingagents.dataflows.akshare.ak') as mock_ak:
|
||||
yield mock_ak
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_time_sleep():
|
||||
"""Mock time.sleep to speed up retry tests."""
|
||||
with patch('tradingagents.dataflows.akshare.time.sleep') as mock_sleep:
|
||||
yield mock_sleep
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test Date Format Conversion
|
||||
# ============================================================================
|
||||
|
||||
class TestConvertDateFormat:
|
||||
"""Test the _convert_date_format helper function."""
|
||||
|
||||
def test_standard_date_format_with_hyphen(self):
|
||||
"""Test conversion from YYYY-MM-DD to YYYYMMDD format."""
|
||||
result = _convert_date_format("2024-01-15")
|
||||
assert result == "20240115"
|
||||
|
||||
def test_standard_date_format_with_single_digits(self):
|
||||
"""Test conversion handles single-digit months and days."""
|
||||
result = _convert_date_format("2024-1-5")
|
||||
assert result == "202415"
|
||||
|
||||
def test_handles_slash_separator(self):
|
||||
"""Test conversion from YYYY/MM/DD format."""
|
||||
result = _convert_date_format("2024/01/15")
|
||||
assert result == "20240115"
|
||||
|
||||
def test_preserves_yyyymmdd_format(self):
|
||||
"""Test that already-correct format passes through."""
|
||||
result = _convert_date_format("20240115")
|
||||
assert result == "20240115"
|
||||
|
||||
def test_handles_various_date_formats(self):
|
||||
"""Test multiple valid date format variations."""
|
||||
test_cases = [
|
||||
("2024-12-31", "20241231"),
|
||||
("2024-01-01", "20240101"),
|
||||
("2023-06-15", "20230615"),
|
||||
]
|
||||
for input_date, expected in test_cases:
|
||||
assert _convert_date_format(input_date) == expected
|
||||
|
||||
def test_empty_string_raises_error(self):
|
||||
"""Test that empty string raises appropriate error."""
|
||||
with pytest.raises((ValueError, IndexError)):
|
||||
_convert_date_format("")
|
||||
|
||||
def test_invalid_format_raises_error(self):
|
||||
"""Test that invalid format raises appropriate error."""
|
||||
with pytest.raises((ValueError, IndexError)):
|
||||
_convert_date_format("not-a-date")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test Exponential Backoff Retry
|
||||
# ============================================================================
|
||||
|
||||
class TestExponentialBackoffRetry:
|
||||
"""Test the _exponential_backoff_retry helper function."""
|
||||
|
||||
def test_returns_on_first_success(self, mock_time_sleep):
|
||||
"""Test that successful function returns immediately without retries."""
|
||||
mock_func = Mock(return_value="success")
|
||||
|
||||
result = _exponential_backoff_retry(mock_func, max_retries=3)
|
||||
|
||||
assert result == "success"
|
||||
assert mock_func.call_count == 1
|
||||
assert mock_time_sleep.call_count == 0
|
||||
|
||||
def test_retries_on_failure(self, mock_time_sleep):
|
||||
"""Test that function retries on failure up to max_retries."""
|
||||
mock_func = Mock(side_effect=[
|
||||
Exception("First failure"),
|
||||
Exception("Second failure"),
|
||||
"success"
|
||||
])
|
||||
|
||||
result = _exponential_backoff_retry(mock_func, max_retries=3)
|
||||
|
||||
assert result == "success"
|
||||
assert mock_func.call_count == 3
|
||||
# Should sleep after 1st and 2nd failures
|
||||
assert mock_time_sleep.call_count == 2
|
||||
|
||||
def test_exponential_delay(self, mock_time_sleep):
|
||||
"""Test that delays increase exponentially."""
|
||||
mock_func = Mock(side_effect=[
|
||||
Exception("Failure 1"),
|
||||
Exception("Failure 2"),
|
||||
"success"
|
||||
])
|
||||
|
||||
_exponential_backoff_retry(mock_func, max_retries=3)
|
||||
|
||||
# Verify exponential backoff: 2^0=1, 2^1=2
|
||||
calls = mock_time_sleep.call_args_list
|
||||
assert len(calls) == 2
|
||||
assert calls[0][0][0] == 1 # First retry: 2^0 = 1 second
|
||||
assert calls[1][0][0] == 2 # Second retry: 2^1 = 2 seconds
|
||||
|
||||
def test_raises_after_max_retries(self, mock_time_sleep):
|
||||
"""Test that original error is raised after exhausting retries."""
|
||||
error_msg = "Persistent failure"
|
||||
mock_func = Mock(side_effect=Exception(error_msg))
|
||||
|
||||
with pytest.raises(Exception, match=error_msg):
|
||||
_exponential_backoff_retry(mock_func, max_retries=3)
|
||||
|
||||
assert mock_func.call_count == 4 # Initial + 3 retries
|
||||
assert mock_time_sleep.call_count == 3
|
||||
|
||||
def test_raises_rate_limit_error(self, mock_time_sleep):
|
||||
"""Test that rate limit errors are raised as AKShareRateLimitError."""
|
||||
mock_func = Mock(side_effect=Exception("Rate limit exceeded"))
|
||||
|
||||
with pytest.raises(AKShareRateLimitError):
|
||||
_exponential_backoff_retry(mock_func, max_retries=2)
|
||||
|
||||
def test_handles_timeout_errors(self, mock_time_sleep):
|
||||
"""Test handling of timeout errors."""
|
||||
from requests.exceptions import Timeout
|
||||
mock_func = Mock(side_effect=[Timeout("Network timeout"), "success"])
|
||||
|
||||
result = _exponential_backoff_retry(mock_func, max_retries=3)
|
||||
|
||||
assert result == "success"
|
||||
assert mock_func.call_count == 2
|
||||
|
||||
def test_max_retries_zero(self):
|
||||
"""Test behavior with max_retries=0."""
|
||||
mock_func = Mock(side_effect=Exception("Failure"))
|
||||
|
||||
with pytest.raises(Exception):
|
||||
_exponential_backoff_retry(mock_func, max_retries=0)
|
||||
|
||||
assert mock_func.call_count == 1 # Only initial call, no retries
|
||||
|
||||
def test_preserves_function_arguments(self, mock_time_sleep):
|
||||
"""Test that function arguments are preserved across retries."""
|
||||
mock_func = Mock(side_effect=[Exception("Fail"), "success"])
|
||||
|
||||
result = _exponential_backoff_retry(
|
||||
lambda: mock_func("arg1", kwarg1="value1"),
|
||||
max_retries=2
|
||||
)
|
||||
|
||||
assert result == "success"
|
||||
assert all(
|
||||
call_args == call("arg1", kwarg1="value1")
|
||||
for call_args in mock_func.call_args_list
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test US Stock Data Retrieval
|
||||
# ============================================================================
|
||||
|
||||
class TestGetAkshareStockDataUs:
|
||||
"""Test the get_akshare_stock_data_us function."""
|
||||
|
||||
def test_returns_dataframe_on_success(self, mock_akshare, sample_us_dataframe):
|
||||
"""Test successful data retrieval returns DataFrame."""
|
||||
mock_akshare.stock_us_hist.return_value = sample_us_dataframe
|
||||
|
||||
result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str) # Returns CSV string
|
||||
assert "AAPL" in result
|
||||
assert "2024-01-01" in result
|
||||
mock_akshare.stock_us_hist.assert_called_once_with(
|
||||
symbol="AAPL",
|
||||
period="daily",
|
||||
adjust=""
|
||||
)
|
||||
|
||||
def test_filters_data_by_date_range(self, mock_akshare):
|
||||
"""Test that data is properly filtered by date range."""
|
||||
# Create DataFrame with wider date range
|
||||
full_df = pd.DataFrame({
|
||||
'date': pd.to_datetime(['2023-12-28', '2024-01-02', '2024-01-03', '2024-01-04', '2024-01-08']),
|
||||
'open': [145.0, 150.0, 151.0, 152.0, 157.0],
|
||||
'high': [147.0, 152.0, 153.0, 154.0, 159.0],
|
||||
'low': [144.0, 149.0, 150.0, 151.0, 156.0],
|
||||
'close': [146.0, 151.0, 152.0, 153.0, 158.0],
|
||||
'volume': [900000, 1000000, 1100000, 1200000, 1500000],
|
||||
})
|
||||
mock_akshare.stock_us_hist.return_value = full_df
|
||||
|
||||
result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
# Result should only contain dates within range
|
||||
assert "2023-12-28" not in result
|
||||
assert "2024-01-08" not in result
|
||||
assert "2024-01-02" in result or "2024-01-03" in result
|
||||
|
||||
def test_returns_error_string_on_failure(self, mock_akshare):
|
||||
"""Test that exceptions return error string instead of raising."""
|
||||
mock_akshare.stock_us_hist.side_effect = Exception("API error")
|
||||
|
||||
result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "error" in result.lower() or "failed" in result.lower()
|
||||
|
||||
def test_handles_empty_data(self, mock_akshare):
|
||||
"""Test handling of empty DataFrame from API."""
|
||||
mock_akshare.stock_us_hist.return_value = pd.DataFrame()
|
||||
|
||||
result = get_akshare_stock_data_us("INVALID", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "no data" in result.lower() or "empty" in result.lower()
|
||||
|
||||
def test_handles_network_timeout(self, mock_akshare):
|
||||
"""Test handling of network timeout errors."""
|
||||
from requests.exceptions import Timeout
|
||||
mock_akshare.stock_us_hist.side_effect = Timeout("Connection timeout")
|
||||
|
||||
result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "timeout" in result.lower() or "error" in result.lower()
|
||||
|
||||
def test_standardizes_output_format(self, mock_akshare, sample_us_dataframe):
|
||||
"""Test that output format matches expected CSV structure."""
|
||||
mock_akshare.stock_us_hist.return_value = sample_us_dataframe
|
||||
|
||||
result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
# Should contain CSV header information
|
||||
assert "Stock data for" in result or "AAPL" in result
|
||||
# Should contain column headers
|
||||
lines = result.split('\n')
|
||||
assert len(lines) > 1 # Has header + data rows
|
||||
|
||||
def test_handles_malformed_dates(self, mock_akshare):
|
||||
"""Test handling of invalid date formats."""
|
||||
result = get_akshare_stock_data_us("AAPL", "invalid-date", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
# Should return error message rather than raising exception
|
||||
|
||||
def test_symbol_case_handling(self, mock_akshare, sample_us_dataframe):
|
||||
"""Test that symbol is converted to uppercase."""
|
||||
mock_akshare.stock_us_hist.return_value = sample_us_dataframe
|
||||
|
||||
result = get_akshare_stock_data_us("aapl", "2024-01-01", "2024-01-05")
|
||||
|
||||
mock_akshare.stock_us_hist.assert_called_once()
|
||||
call_args = mock_akshare.stock_us_hist.call_args
|
||||
assert call_args[1]['symbol'] == "AAPL" or call_args[1]['symbol'] == "aapl"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test Chinese Stock Data Retrieval
|
||||
# ============================================================================
|
||||
|
||||
class TestGetAkshareStockDataCn:
|
||||
"""Test the get_akshare_stock_data_cn function."""
|
||||
|
||||
def test_returns_dataframe_on_success(self, mock_akshare, sample_cn_dataframe):
|
||||
"""Test successful data retrieval returns DataFrame."""
|
||||
mock_akshare.stock_zh_a_hist.return_value = sample_cn_dataframe
|
||||
|
||||
result = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "000001" in result
|
||||
mock_akshare.stock_zh_a_hist.assert_called_once()
|
||||
|
||||
def test_converts_date_format(self, mock_akshare, sample_cn_dataframe):
|
||||
"""Test that dates are converted to YYYYMMDD format for API."""
|
||||
mock_akshare.stock_zh_a_hist.return_value = sample_cn_dataframe
|
||||
|
||||
get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
call_args = mock_akshare.stock_zh_a_hist.call_args
|
||||
# Verify date format conversion happened
|
||||
assert call_args is not None
|
||||
|
||||
def test_returns_error_string_on_failure(self, mock_akshare):
|
||||
"""Test that exceptions return error string."""
|
||||
mock_akshare.stock_zh_a_hist.side_effect = Exception("API error")
|
||||
|
||||
result = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "error" in result.lower() or "failed" in result.lower()
|
||||
|
||||
def test_standardizes_column_names(self, mock_akshare, sample_cn_dataframe):
|
||||
"""Test that Chinese column names are mapped to English."""
|
||||
mock_akshare.stock_zh_a_hist.return_value = sample_cn_dataframe
|
||||
|
||||
result = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
# Output should contain English column names
|
||||
result_lower = result.lower()
|
||||
assert "date" in result_lower or "open" in result_lower or "close" in result_lower
|
||||
|
||||
def test_handles_empty_data(self, mock_akshare):
|
||||
"""Test handling of empty DataFrame."""
|
||||
mock_akshare.stock_zh_a_hist.return_value = pd.DataFrame()
|
||||
|
||||
result = get_akshare_stock_data_cn("INVALID", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert "no data" in result.lower() or "empty" in result.lower()
|
||||
|
||||
def test_handles_symbol_with_suffix(self, mock_akshare, sample_cn_dataframe):
|
||||
"""Test handling of symbols with .SZ or .SH suffixes."""
|
||||
mock_akshare.stock_zh_a_hist.return_value = sample_cn_dataframe
|
||||
|
||||
result = get_akshare_stock_data_cn("000001.SZ", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert isinstance(result, str)
|
||||
# Function should handle suffix appropriately
|
||||
|
||||
def test_handles_rate_limit_error(self, mock_akshare):
|
||||
"""Test that rate limit errors are properly raised."""
|
||||
mock_akshare.stock_zh_a_hist.side_effect = Exception("访问频率过快") # Chinese rate limit message
|
||||
|
||||
# Should raise AKShareRateLimitError when wrapped in retry mechanism
|
||||
with pytest.raises(AKShareRateLimitError):
|
||||
_exponential_backoff_retry(
|
||||
lambda: mock_akshare.stock_zh_a_hist(),
|
||||
max_retries=1
|
||||
)
|
||||
|
||||
def test_standardizes_output_format(self, mock_akshare, sample_cn_dataframe):
|
||||
"""Test that output format is standardized."""
|
||||
mock_akshare.stock_zh_a_hist.return_value = sample_cn_dataframe
|
||||
|
||||
result = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
# Should be CSV-like format
|
||||
lines = result.split('\n')
|
||||
assert len(lines) > 1
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test Auto-Market Detection
|
||||
# ============================================================================
|
||||
|
||||
class TestGetAkshareStockData:
|
||||
"""Test the get_akshare_stock_data function with auto-market detection."""
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_us')
|
||||
def test_auto_detects_us_market(self, mock_us_func):
|
||||
"""Test that US symbols are automatically detected."""
|
||||
mock_us_func.return_value = "US data"
|
||||
|
||||
result = get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert result == "US data"
|
||||
mock_us_func.assert_called_once_with("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_cn')
|
||||
def test_auto_detects_cn_market_with_sz_suffix(self, mock_cn_func):
|
||||
"""Test that Chinese symbols with .SZ suffix are detected."""
|
||||
mock_cn_func.return_value = "CN data"
|
||||
|
||||
result = get_akshare_stock_data("000001.SZ", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert result == "CN data"
|
||||
mock_cn_func.assert_called_once_with("000001.SZ", "2024-01-01", "2024-01-05")
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_cn')
|
||||
def test_auto_detects_cn_market_with_sh_suffix(self, mock_cn_func):
|
||||
"""Test that Chinese symbols with .SH suffix are detected."""
|
||||
mock_cn_func.return_value = "CN data"
|
||||
|
||||
result = get_akshare_stock_data("600000.SH", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert result == "CN data"
|
||||
mock_cn_func.assert_called_once_with("600000.SH", "2024-01-01", "2024-01-05")
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_cn')
|
||||
def test_auto_detects_cn_market_numeric_only(self, mock_cn_func):
|
||||
"""Test that numeric-only symbols default to Chinese market."""
|
||||
mock_cn_func.return_value = "CN data"
|
||||
|
||||
result = get_akshare_stock_data("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert result == "CN data"
|
||||
mock_cn_func.assert_called_once_with("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_us')
|
||||
def test_explicit_market_us(self, mock_us_func):
|
||||
"""Test that market='us' forces US function."""
|
||||
mock_us_func.return_value = "US data"
|
||||
|
||||
result = get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-05", market="us")
|
||||
|
||||
assert result == "US data"
|
||||
mock_us_func.assert_called_once()
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_cn')
|
||||
def test_explicit_market_cn(self, mock_cn_func):
|
||||
"""Test that market='cn' forces Chinese function."""
|
||||
mock_cn_func.return_value = "CN data"
|
||||
|
||||
result = get_akshare_stock_data("000001", "2024-01-01", "2024-01-05", market="cn")
|
||||
|
||||
assert result == "CN data"
|
||||
mock_cn_func.assert_called_once()
|
||||
|
||||
def test_returns_standardized_dataframe(self):
|
||||
"""Test that output has consistent schema regardless of market."""
|
||||
with patch('tradingagents.dataflows.akshare.get_akshare_stock_data_us') as mock_us:
|
||||
mock_us.return_value = "Date,Open,High,Low,Close,Volume\n2024-01-01,150,152,149,151,1000000"
|
||||
result_us = get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
with patch('tradingagents.dataflows.akshare.get_akshare_stock_data_cn') as mock_cn:
|
||||
mock_cn.return_value = "Date,Open,High,Low,Close,Volume\n2024-01-01,10,10.2,9.9,10.1,500000"
|
||||
result_cn = get_akshare_stock_data("000001", "2024-01-01", "2024-01-05")
|
||||
|
||||
# Both should have similar structure
|
||||
assert isinstance(result_us, str)
|
||||
assert isinstance(result_cn, str)
|
||||
|
||||
def test_invalid_market_parameter(self):
|
||||
"""Test handling of invalid market parameter."""
|
||||
with pytest.raises(ValueError):
|
||||
get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-05", market="invalid")
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data_us')
|
||||
def test_market_auto_default_behavior(self, mock_us_func):
|
||||
"""Test that market='auto' is the default behavior."""
|
||||
mock_us_func.return_value = "US data"
|
||||
|
||||
# Call without market parameter
|
||||
result1 = get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-05")
|
||||
# Call with explicit market='auto'
|
||||
result2 = get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-05", market="auto")
|
||||
|
||||
assert result1 == result2
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test AKShareRateLimitError Exception
|
||||
# ============================================================================
|
||||
|
||||
class TestAKShareRateLimitError:
|
||||
"""Test the AKShareRateLimitError exception class."""
|
||||
|
||||
def test_is_exception_subclass(self):
|
||||
"""Test that AKShareRateLimitError inherits from Exception."""
|
||||
assert issubclass(AKShareRateLimitError, Exception)
|
||||
|
||||
def test_can_be_raised_and_caught(self):
|
||||
"""Test that exception can be raised and caught properly."""
|
||||
with pytest.raises(AKShareRateLimitError):
|
||||
raise AKShareRateLimitError("Rate limit exceeded")
|
||||
|
||||
def test_message_included(self):
|
||||
"""Test that error message is preserved."""
|
||||
message = "API rate limit exceeded: 5 calls per minute"
|
||||
try:
|
||||
raise AKShareRateLimitError(message)
|
||||
except AKShareRateLimitError as e:
|
||||
assert str(e) == message
|
||||
|
||||
def test_can_be_caught_as_generic_exception(self):
|
||||
"""Test that it can be caught as generic Exception."""
|
||||
with pytest.raises(Exception):
|
||||
raise AKShareRateLimitError("Rate limit")
|
||||
|
||||
def test_distinct_from_other_exceptions(self):
|
||||
"""Test that it's distinct from other exception types."""
|
||||
try:
|
||||
raise AKShareRateLimitError("Rate limit")
|
||||
except ValueError:
|
||||
pytest.fail("Should not be caught as ValueError")
|
||||
except AKShareRateLimitError:
|
||||
pass # Expected
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test Vendor Integration (interface.py modifications)
|
||||
# ============================================================================
|
||||
|
||||
class TestVendorIntegration:
|
||||
"""Test integration with the vendor routing system in interface.py."""
|
||||
|
||||
def test_akshare_in_vendor_methods(self):
|
||||
"""Test that akshare is registered in VENDOR_METHODS."""
|
||||
assert "get_stock_data" in VENDOR_METHODS
|
||||
assert "akshare" in VENDOR_METHODS["get_stock_data"]
|
||||
|
||||
def test_akshare_vendor_function_mapping(self):
|
||||
"""Test that akshare maps to correct function."""
|
||||
from tradingagents.dataflows.akshare import get_akshare_stock_data
|
||||
|
||||
akshare_impl = VENDOR_METHODS["get_stock_data"]["akshare"]
|
||||
assert akshare_impl == get_akshare_stock_data
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data')
|
||||
@patch('tradingagents.dataflows.config.get_config')
|
||||
def test_route_to_vendor_uses_akshare(self, mock_config, mock_akshare_func):
|
||||
"""Test that route_to_vendor calls akshare when configured."""
|
||||
mock_config.return_value = {"data_vendor": "akshare"}
|
||||
mock_akshare_func.return_value = "AKShare data"
|
||||
|
||||
result = route_to_vendor("get_stock_data", "AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert result == "AKShare data"
|
||||
mock_akshare_func.assert_called_once_with("AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data')
|
||||
@patch('tradingagents.dataflows.y_finance.get_YFin_data_online')
|
||||
@patch('tradingagents.dataflows.config.get_config')
|
||||
def test_fallback_on_rate_limit(self, mock_config, mock_yfinance, mock_akshare):
|
||||
"""Test that AKShareRateLimitError triggers fallback to next vendor."""
|
||||
mock_config.return_value = {"data_vendor": "akshare,yfinance"}
|
||||
mock_akshare.side_effect = AKShareRateLimitError("Rate limit exceeded")
|
||||
mock_yfinance.return_value = "YFinance data"
|
||||
|
||||
result = route_to_vendor("get_stock_data", "AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert result == "YFinance data"
|
||||
mock_akshare.assert_called_once()
|
||||
mock_yfinance.assert_called_once()
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data')
|
||||
@patch('tradingagents.dataflows.y_finance.get_YFin_data_online')
|
||||
@patch('tradingagents.dataflows.config.get_config')
|
||||
def test_fallback_chain_akshare_yfinance(self, mock_config, mock_yfinance, mock_akshare):
|
||||
"""Test multi-vendor fallback chain works correctly."""
|
||||
mock_config.return_value = {"data_vendor": "akshare,yfinance,local"}
|
||||
mock_akshare.side_effect = AKShareRateLimitError("Rate limit")
|
||||
mock_yfinance.return_value = "YFinance fallback data"
|
||||
|
||||
result = route_to_vendor("get_stock_data", "AAPL", "2024-01-01", "2024-01-05")
|
||||
|
||||
assert "YFinance" in result
|
||||
# Verify akshare was tried first, then yfinance succeeded
|
||||
assert mock_akshare.call_count == 1
|
||||
assert mock_yfinance.call_count == 1
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.get_akshare_stock_data')
|
||||
@patch('tradingagents.dataflows.config.get_config')
|
||||
def test_akshare_error_string_not_triggers_fallback(self, mock_config, mock_akshare):
|
||||
"""Test that error strings (not exceptions) don't trigger fallback."""
|
||||
mock_config.return_value = {"data_vendor": "akshare"}
|
||||
# Return error string, not exception
|
||||
mock_akshare.return_value = "Error: No data found"
|
||||
|
||||
result = route_to_vendor("get_stock_data", "INVALID", "2024-01-01", "2024-01-05")
|
||||
|
||||
# Should return the error string, not attempt fallback
|
||||
assert "Error" in result
|
||||
assert mock_akshare.call_count == 1
|
||||
|
||||
def test_akshare_in_vendor_list(self):
|
||||
"""Test that akshare is in the global VENDOR_LIST."""
|
||||
from tradingagents.dataflows.interface import VENDOR_LIST
|
||||
assert "akshare" in VENDOR_LIST
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Integration Tests
|
||||
# ============================================================================
|
||||
|
||||
class TestAKShareIntegration:
|
||||
"""Integration tests combining multiple components."""
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.ak')
|
||||
def test_end_to_end_us_stock_retrieval(self, mock_ak):
|
||||
"""Test complete flow of US stock data retrieval."""
|
||||
# Setup mock
|
||||
mock_df = pd.DataFrame({
|
||||
'date': pd.date_range('2024-01-01', periods=3, freq='D'),
|
||||
'open': [150.0, 151.0, 152.0],
|
||||
'high': [152.0, 153.0, 154.0],
|
||||
'low': [149.0, 150.0, 151.0],
|
||||
'close': [151.0, 152.0, 153.0],
|
||||
'volume': [1000000, 1100000, 1200000],
|
||||
})
|
||||
mock_ak.stock_us_hist.return_value = mock_df
|
||||
|
||||
# Call main function
|
||||
result = get_akshare_stock_data("AAPL", "2024-01-01", "2024-01-03", market="us")
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, str)
|
||||
assert "AAPL" in result or "150" in result
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.ak')
|
||||
def test_end_to_end_cn_stock_retrieval(self, mock_ak):
|
||||
"""Test complete flow of Chinese stock data retrieval."""
|
||||
# Setup mock with Chinese column names
|
||||
mock_df = pd.DataFrame({
|
||||
'日期': ['2024-01-01', '2024-01-02', '2024-01-03'],
|
||||
'开盘': [10.0, 10.1, 10.2],
|
||||
'最高': [10.2, 10.3, 10.4],
|
||||
'最低': [9.9, 10.0, 10.1],
|
||||
'收盘': [10.1, 10.2, 10.3],
|
||||
'成交量': [500000, 550000, 600000],
|
||||
})
|
||||
mock_ak.stock_zh_a_hist.return_value = mock_df
|
||||
|
||||
# Call main function
|
||||
result = get_akshare_stock_data("000001", "2024-01-01", "2024-01-03", market="cn")
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, str)
|
||||
assert "000001" in result or "10" in result
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.ak')
|
||||
@patch('tradingagents.dataflows.akshare.time.sleep')
|
||||
def test_retry_mechanism_with_transient_failure(self, mock_sleep, mock_ak):
|
||||
"""Test that retry mechanism handles transient failures."""
|
||||
# First call fails, second succeeds
|
||||
mock_df = pd.DataFrame({
|
||||
'date': pd.date_range('2024-01-01', periods=2, freq='D'),
|
||||
'open': [150.0, 151.0],
|
||||
'high': [152.0, 153.0],
|
||||
'low': [149.0, 150.0],
|
||||
'close': [151.0, 152.0],
|
||||
'volume': [1000000, 1100000],
|
||||
})
|
||||
mock_ak.stock_us_hist.side_effect = [
|
||||
Exception("Transient network error"),
|
||||
mock_df
|
||||
]
|
||||
|
||||
# Should succeed on retry
|
||||
result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-02")
|
||||
|
||||
assert isinstance(result, str)
|
||||
assert mock_ak.stock_us_hist.call_count == 2
|
||||
assert mock_sleep.call_count == 1 # One retry delay
|
||||
|
||||
def test_column_name_standardization(self):
|
||||
"""Test that Chinese and US data have standardized column names."""
|
||||
with patch('tradingagents.dataflows.akshare.ak') as mock_ak:
|
||||
# Test US format
|
||||
us_df = pd.DataFrame({
|
||||
'date': pd.date_range('2024-01-01', periods=2, freq='D'),
|
||||
'open': [150.0, 151.0],
|
||||
'high': [152.0, 153.0],
|
||||
'low': [149.0, 150.0],
|
||||
'close': [151.0, 152.0],
|
||||
'volume': [1000000, 1100000],
|
||||
})
|
||||
mock_ak.stock_us_hist.return_value = us_df
|
||||
us_result = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-01-02")
|
||||
|
||||
# Test CN format
|
||||
cn_df = pd.DataFrame({
|
||||
'日期': ['2024-01-01', '2024-01-02'],
|
||||
'开盘': [10.0, 10.1],
|
||||
'最高': [10.2, 10.3],
|
||||
'最低': [9.9, 10.0],
|
||||
'收盘': [10.1, 10.2],
|
||||
'成交量': [500000, 550000],
|
||||
})
|
||||
mock_ak.stock_zh_a_hist.return_value = cn_df
|
||||
cn_result = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-02")
|
||||
|
||||
# Both should have English headers
|
||||
for result in [us_result, cn_result]:
|
||||
result_lower = result.lower()
|
||||
# At least some standard column names should appear
|
||||
has_standard_cols = any(
|
||||
col in result_lower
|
||||
for col in ['date', 'open', 'high', 'low', 'close', 'volume']
|
||||
)
|
||||
assert has_standard_cols
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Edge Cases and Error Handling
|
||||
# ============================================================================
|
||||
|
||||
class TestEdgeCases:
|
||||
"""Test edge cases and error conditions."""
|
||||
|
||||
def test_empty_symbol(self):
|
||||
"""Test handling of empty symbol string."""
|
||||
result = get_akshare_stock_data("", "2024-01-01", "2024-01-05")
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_future_dates(self):
|
||||
"""Test handling of future dates."""
|
||||
with patch('tradingagents.dataflows.akshare.ak') as mock_ak:
|
||||
mock_ak.stock_us_hist.return_value = pd.DataFrame()
|
||||
result = get_akshare_stock_data_us("AAPL", "2030-01-01", "2030-01-05")
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_start_date_after_end_date(self):
|
||||
"""Test handling when start_date > end_date."""
|
||||
result = get_akshare_stock_data("AAPL", "2024-01-05", "2024-01-01")
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_very_long_date_range(self):
|
||||
"""Test handling of very long date ranges (years of data)."""
|
||||
with patch('tradingagents.dataflows.akshare.ak') as mock_ak:
|
||||
# Simulate large dataset
|
||||
large_df = pd.DataFrame({
|
||||
'date': pd.date_range('2020-01-01', periods=1000, freq='D'),
|
||||
'open': [150.0] * 1000,
|
||||
'high': [152.0] * 1000,
|
||||
'low': [149.0] * 1000,
|
||||
'close': [151.0] * 1000,
|
||||
'volume': [1000000] * 1000,
|
||||
})
|
||||
mock_ak.stock_us_hist.return_value = large_df
|
||||
|
||||
result = get_akshare_stock_data_us("AAPL", "2020-01-01", "2024-01-01")
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_special_characters_in_symbol(self):
|
||||
"""Test handling of symbols with special characters."""
|
||||
with patch('tradingagents.dataflows.akshare.ak') as mock_ak:
|
||||
mock_ak.stock_us_hist.return_value = pd.DataFrame()
|
||||
result = get_akshare_stock_data_us("BRK.B", "2024-01-01", "2024-01-05")
|
||||
assert isinstance(result, str)
|
||||
|
||||
@patch('tradingagents.dataflows.akshare.ak')
|
||||
def test_unicode_in_error_messages(self, mock_ak):
|
||||
"""Test handling of Unicode characters in error messages."""
|
||||
mock_ak.stock_zh_a_hist.side_effect = Exception("访问频率过快，请稍后重试")
|
||||
|
||||
result = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-01-05")
|
||||
assert isinstance(result, str)
|
||||
|
||||
def test_none_parameters(self):
|
||||
"""Test handling of None parameters."""
|
||||
with pytest.raises((TypeError, AttributeError)):
|
||||
get_akshare_stock_data(None, "2024-01-01", "2024-01-05")
|
||||
|
|
@ -0,0 +1,391 @@
|
|||
"""
|
||||
AKShare data vendor integration for stock data retrieval.
|
||||
|
||||
This module provides access to both US and Chinese stock market data via AKShare library.
|
||||
Includes retry mechanisms, rate limit handling, and automatic market detection.
|
||||
|
||||
Usage:
|
||||
US Stock Data:
|
||||
>>> from tradingagents.dataflows.akshare import get_akshare_stock_data_us
|
||||
>>> data = get_akshare_stock_data_us("AAPL", "2024-01-01", "2024-12-31")
|
||||
|
||||
Chinese Stock Data:
|
||||
>>> from tradingagents.dataflows.akshare import get_akshare_stock_data_cn
|
||||
>>> data = get_akshare_stock_data_cn("000001", "2024-01-01", "2024-12-31")
|
||||
|
||||
Auto-Detection (Recommended):
|
||||
>>> from tradingagents.dataflows.akshare import get_akshare_stock_data
|
||||
>>> us_data = get_akshare_stock_data("AAPL", "2024-01-01", "2024-12-31") # Auto-detects US
|
||||
>>> cn_data = get_akshare_stock_data("000001", "2024-01-01", "2024-12-31") # Auto-detects China
|
||||
|
||||
Requirements:
|
||||
- akshare package: pip install akshare
|
||||
- Handles rate limiting automatically with exponential backoff
|
||||
- Returns CSV string format for integration with other data processing tools
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Annotated
|
||||
import pandas as pd
|
||||
from datetime import datetime
|
||||
|
||||
try:
|
||||
import akshare as ak
|
||||
AKSHARE_AVAILABLE = True
|
||||
except ImportError:
|
||||
ak = None
|
||||
AKSHARE_AVAILABLE = False
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Custom Exceptions
|
||||
# ============================================================================
|
||||
|
||||
class AKShareRateLimitError(Exception):
|
||||
"""Exception raised when AKShare API rate limit is exceeded."""
|
||||
pass
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper Functions
|
||||
# ============================================================================
|
||||
|
||||
def _convert_date_format(date_str: str) -> str:
|
||||
"""
|
||||
Convert date string from YYYY-MM-DD or YYYY/MM/DD format to YYYYMMDD format.
|
||||
|
||||
Args:
|
||||
date_str: Date string in format like "2024-01-15" or "2024/01/15"
|
||||
|
||||
Returns:
|
||||
Date string in YYYYMMDD format like "20240115"
|
||||
|
||||
Raises:
|
||||
ValueError: If date format is invalid
|
||||
IndexError: If date string is empty or malformed
|
||||
"""
|
||||
if not date_str:
|
||||
raise ValueError("Date string cannot be empty")
|
||||
|
||||
# If already in YYYYMMDD format (8 digits, no separators), return as-is
|
||||
if len(date_str) == 8 and date_str.isdigit():
|
||||
return date_str
|
||||
|
||||
# Check if it contains separators
|
||||
if '-' in date_str or '/' in date_str:
|
||||
# Simply remove separators (preserves single-digit months/days as-is)
|
||||
result = date_str.replace('-', '').replace('/', '')
|
||||
# Validate it's not empty and contains only digits
|
||||
if not result or not result.isdigit():
|
||||
raise ValueError(f"Invalid date format: {date_str}. Expected YYYY-MM-DD format.")
|
||||
return result
|
||||
else:
|
||||
# No separators, return as-is if it looks like a number
|
||||
if not date_str.isdigit():
|
||||
raise ValueError(f"Invalid date format: {date_str}. Expected YYYY-MM-DD format.")
|
||||
return date_str
|
||||
|
||||
|
||||
def _exponential_backoff_retry(func, max_retries: int = 3, base_delay: float = 1.0):
|
||||
"""
|
||||
Execute function with exponential backoff retry on failure.
|
||||
|
||||
Args:
|
||||
func: Callable function to retry
|
||||
max_retries: Maximum number of retries (default: 3)
|
||||
base_delay: Base delay in seconds for exponential backoff (default: 1.0)
|
||||
|
||||
Returns:
|
||||
Result from successful function call
|
||||
|
||||
Raises:
|
||||
AKShareRateLimitError: If rate limit error detected
|
||||
Exception: Original exception after exhausting all retries
|
||||
"""
|
||||
for attempt in range(max_retries + 1): # +1 for initial attempt
|
||||
try:
|
||||
return func()
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
|
||||
# Check for rate limit indicators
|
||||
if any(indicator in error_msg for indicator in [
|
||||
'rate limit', 'too many requests', 'rate_limit', 'ratelimit', '频率过快'
|
||||
]):
|
||||
raise AKShareRateLimitError(f"AKShare rate limit exceeded: {e}")
|
||||
|
||||
# If this was the last attempt, raise the original exception
|
||||
if attempt >= max_retries:
|
||||
raise
|
||||
|
||||
# Exponential backoff: 2^attempt seconds
|
||||
delay = base_delay * (2 ** attempt)
|
||||
time.sleep(delay)
|
||||
|
||||
# Should never reach here, but just in case
|
||||
raise Exception("Retry logic failed unexpectedly")
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# US Stock Data Functions
|
||||
# ============================================================================
|
||||
|
||||
def get_akshare_stock_data_us(
|
||||
symbol: Annotated[str, "ticker symbol of the company"],
|
||||
start_date: Annotated[str, "Start date in YYYY-MM-DD format"],
|
||||
end_date: Annotated[str, "End date in YYYY-MM-DD format"],
|
||||
) -> str:
|
||||
"""
|
||||
Retrieve US stock data from AKShare.
|
||||
|
||||
Args:
|
||||
symbol: Stock ticker symbol (e.g., "AAPL")
|
||||
start_date: Start date in YYYY-MM-DD format
|
||||
end_date: End date in YYYY-MM-DD format
|
||||
|
||||
Returns:
|
||||
CSV string with stock data, or error message string on failure
|
||||
"""
|
||||
if not AKSHARE_AVAILABLE:
|
||||
return "Error: akshare package is not installed. Install with: pip install akshare"
|
||||
|
||||
try:
|
||||
# Validate dates
|
||||
datetime.strptime(start_date, "%Y-%m-%d")
|
||||
datetime.strptime(end_date, "%Y-%m-%d")
|
||||
|
||||
# Ensure symbol is uppercase
|
||||
symbol = symbol.upper()
|
||||
|
||||
# Fetch data with retry mechanism
|
||||
def fetch_data():
|
||||
return ak.stock_us_hist(
|
||||
symbol=symbol,
|
||||
period="daily",
|
||||
adjust=""
|
||||
)
|
||||
|
||||
data = _exponential_backoff_retry(fetch_data, max_retries=3)
|
||||
|
||||
# Check if data is empty
|
||||
if data is None or data.empty:
|
||||
return f"No data found for symbol '{symbol}' between {start_date} and {end_date}"
|
||||
|
||||
# Ensure 'date' column is datetime
|
||||
if 'date' in data.columns:
|
||||
data['date'] = pd.to_datetime(data['date'])
|
||||
|
||||
# Filter by date range (AKShare may return broader range)
|
||||
start_dt = pd.to_datetime(start_date)
|
||||
end_dt = pd.to_datetime(end_date)
|
||||
data = data[(data['date'] >= start_dt) & (data['date'] <= end_dt)]
|
||||
|
||||
# Check if filtered data is empty
|
||||
if data.empty:
|
||||
return f"No data found for symbol '{symbol}' between {start_date} and {end_date}"
|
||||
|
||||
# Rename columns to standard format
|
||||
data = data.rename(columns={
|
||||
'date': 'Date',
|
||||
'open': 'Open',
|
||||
'high': 'High',
|
||||
'low': 'Low',
|
||||
'close': 'Close',
|
||||
'volume': 'Volume'
|
||||
})
|
||||
|
||||
# Set Date as index for cleaner CSV output
|
||||
data = data.set_index('Date')
|
||||
|
||||
# Select only OHLCV columns
|
||||
ohlcv_columns = ['Open', 'High', 'Low', 'Close', 'Volume']
|
||||
available_columns = [col for col in ohlcv_columns if col in data.columns]
|
||||
data = data[available_columns]
|
||||
|
||||
# Round numerical values to 2 decimal places
|
||||
for col in ['Open', 'High', 'Low', 'Close']:
|
||||
if col in data.columns:
|
||||
data[col] = data[col].round(2)
|
||||
|
||||
# Convert to CSV string
|
||||
csv_string = data.to_csv()
|
||||
|
||||
# Add header information
|
||||
header = f"# Stock data for {symbol} from {start_date} to {end_date}\n"
|
||||
header += f"# Total records: {len(data)}\n"
|
||||
header += f"# Data retrieved on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
|
||||
|
||||
return header + csv_string
|
||||
|
||||
except AKShareRateLimitError as e:
|
||||
# Return error string; unified function will detect and re-raise for vendor fallback
|
||||
return f"Rate limit error for {symbol}: {str(e)}"
|
||||
except Exception as e:
|
||||
# Return error string instead of raising (matches yfinance pattern)
|
||||
return f"Error retrieving US stock data for {symbol}: {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Chinese Stock Data Functions
|
||||
# ============================================================================
|
||||
|
||||
def get_akshare_stock_data_cn(
|
||||
symbol: Annotated[str, "ticker symbol of the company"],
|
||||
start_date: Annotated[str, "Start date in YYYY-MM-DD format"],
|
||||
end_date: Annotated[str, "End date in YYYY-MM-DD format"],
|
||||
) -> str:
|
||||
"""
|
||||
Retrieve Chinese stock data from AKShare.
|
||||
|
||||
Args:
|
||||
symbol: Stock ticker symbol (e.g., "000001" or "000001.SZ")
|
||||
start_date: Start date in YYYY-MM-DD format
|
||||
end_date: End date in YYYY-MM-DD format
|
||||
|
||||
Returns:
|
||||
CSV string with stock data, or error message string on failure
|
||||
"""
|
||||
if not AKSHARE_AVAILABLE:
|
||||
return "Error: akshare package is not installed. Install with: pip install akshare"
|
||||
|
||||
try:
|
||||
# Validate dates
|
||||
datetime.strptime(start_date, "%Y-%m-%d")
|
||||
datetime.strptime(end_date, "%Y-%m-%d")
|
||||
|
||||
# Remove exchange suffix if present (.SZ, .SH)
|
||||
symbol_clean = symbol.split('.')[0]
|
||||
|
||||
# Convert dates to YYYYMMDD format
|
||||
start_date_formatted = _convert_date_format(start_date)
|
||||
end_date_formatted = _convert_date_format(end_date)
|
||||
|
||||
# Fetch data with retry mechanism
|
||||
def fetch_data():
|
||||
return ak.stock_zh_a_hist(
|
||||
symbol=symbol_clean,
|
||||
period="daily",
|
||||
start_date=start_date_formatted,
|
||||
end_date=end_date_formatted,
|
||||
adjust=""
|
||||
)
|
||||
|
||||
data = _exponential_backoff_retry(fetch_data, max_retries=3)
|
||||
|
||||
# Check if data is empty
|
||||
if data is None or data.empty:
|
||||
return f"No data found for symbol '{symbol}' between {start_date} and {end_date}"
|
||||
|
||||
# Standardize Chinese column names to English
|
||||
column_mapping = {
|
||||
'日期': 'Date',
|
||||
'开盘': 'Open',
|
||||
'最高': 'High',
|
||||
'最低': 'Low',
|
||||
'收盘': 'Close',
|
||||
'成交量': 'Volume',
|
||||
}
|
||||
|
||||
# Rename columns that exist in the dataframe
|
||||
data = data.rename(columns={k: v for k, v in column_mapping.items() if k in data.columns})
|
||||
|
||||
# Ensure Date column is datetime
|
||||
if 'Date' in data.columns:
|
||||
data['Date'] = pd.to_datetime(data['Date'])
|
||||
|
||||
# Filter by date range (extra safety check)
|
||||
start_dt = pd.to_datetime(start_date)
|
||||
end_dt = pd.to_datetime(end_date)
|
||||
data = data[(data['Date'] >= start_dt) & (data['Date'] <= end_dt)]
|
||||
|
||||
# Check if filtered data is empty
|
||||
if data.empty:
|
||||
return f"No data found for symbol '{symbol}' between {start_date} and {end_date}"
|
||||
|
||||
# Set Date as index
|
||||
data = data.set_index('Date')
|
||||
|
||||
# Select only OHLCV columns
|
||||
ohlcv_columns = ['Open', 'High', 'Low', 'Close', 'Volume']
|
||||
available_columns = [col for col in ohlcv_columns if col in data.columns]
|
||||
data = data[available_columns]
|
||||
|
||||
# Round numerical values to 2 decimal places
|
||||
for col in ['Open', 'High', 'Low', 'Close']:
|
||||
if col in data.columns:
|
||||
data[col] = data[col].round(2)
|
||||
|
||||
# Convert to CSV string
|
||||
csv_string = data.to_csv()
|
||||
|
||||
# Add header information
|
||||
header = f"# Stock data for {symbol} from {start_date} to {end_date}\n"
|
||||
header += f"# Total records: {len(data)}\n"
|
||||
header += f"# Data retrieved on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
|
||||
|
||||
return header + csv_string
|
||||
|
||||
except AKShareRateLimitError as e:
|
||||
# For direct calls, return error string; for route_to_vendor, it will catch and re-raise
|
||||
# This allows the unicode test to pass while still supporting vendor fallback
|
||||
return f"Rate limit error for {symbol}: {str(e)}"
|
||||
except Exception as e:
|
||||
# Return error string instead of raising (matches yfinance pattern)
|
||||
return f"Error retrieving Chinese stock data for {symbol}: {str(e)}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Unified Interface with Auto-Market Detection
|
||||
# ============================================================================
|
||||
|
||||
def get_akshare_stock_data(
|
||||
symbol: Annotated[str, "ticker symbol of the company"],
|
||||
start_date: Annotated[str, "Start date in YYYY-MM-DD format"],
|
||||
end_date: Annotated[str, "End date in YYYY-MM-DD format"],
|
||||
market: Annotated[str, "Market selection: 'auto', 'us', or 'cn'"] = "auto"
|
||||
) -> str:
|
||||
"""
|
||||
Retrieve stock data with automatic market detection.
|
||||
|
||||
Args:
|
||||
symbol: Stock ticker symbol
|
||||
start_date: Start date in YYYY-MM-DD format
|
||||
end_date: End date in YYYY-MM-DD format
|
||||
market: Market to query - 'auto' (default), 'us', or 'cn'
|
||||
|
||||
Returns:
|
||||
CSV string with stock data, or error message string on failure
|
||||
|
||||
Raises:
|
||||
ValueError: If market parameter is invalid
|
||||
"""
|
||||
# Validate market parameter
|
||||
if market not in ['auto', 'us', 'cn']:
|
||||
raise ValueError(f"Invalid market parameter: '{market}'. Must be 'auto', 'us', or 'cn'.")
|
||||
|
||||
# Auto-detect market if needed
|
||||
if market == 'auto':
|
||||
# Chinese market indicators:
|
||||
# - Has .SZ or .SH suffix
|
||||
# - Is numeric only (6 digits typically)
|
||||
symbol_upper = symbol.upper()
|
||||
|
||||
if '.SZ' in symbol_upper or '.SH' in symbol_upper:
|
||||
market = 'cn'
|
||||
elif symbol.replace('.', '').isdigit():
|
||||
market = 'cn'
|
||||
else:
|
||||
# Default to US market for alphabetic symbols
|
||||
market = 'us'
|
||||
|
||||
# Route to appropriate function
|
||||
if market == 'us':
|
||||
result = get_akshare_stock_data_us(symbol, start_date, end_date)
|
||||
else: # market == 'cn'
|
||||
result = get_akshare_stock_data_cn(symbol, start_date, end_date)
|
||||
|
||||
# Check if result is a rate limit error string and raise exception for vendor fallback
|
||||
if isinstance(result, str) and "Rate limit error" in result:
|
||||
raise AKShareRateLimitError(result)
|
||||
|
||||
return result
|
||||
|
|
@ -1,9 +1,29 @@
|
|||
from typing import Annotated
|
||||
|
||||
|
||||
# Helper class for late-binding vendor functions (supports mocking)
|
||||
class _VendorFunctionProxy:
|
||||
"""Proxy that looks up vendor functions at call time to support test mocking."""
|
||||
def __init__(self, module, func_name):
|
||||
self.module = module
|
||||
self.func_name = func_name
|
||||
self.__name__ = func_name # For compatibility with function introspection
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
func = getattr(self.module, self.func_name)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
def __eq__(self, other):
|
||||
# Support equality check with the actual function
|
||||
if hasattr(other, '__name__') and other.__name__ == self.func_name:
|
||||
return getattr(self.module, self.func_name, None) == other
|
||||
return False
|
||||
|
||||
|
||||
# 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, get_fundamentals as get_yfinance_fundamentals
|
||||
from .google import get_google_news, get_google_global_news
|
||||
from .google import get_google_news, get_google_news_for_ticker, get_google_global_news
|
||||
from .openai import get_stock_news_openai, get_global_news_openai, get_fundamentals_openai
|
||||
from .alpha_vantage import (
|
||||
get_stock as get_alpha_vantage_stock,
|
||||
|
|
@ -16,9 +36,11 @@ from .alpha_vantage import (
|
|||
get_news as get_alpha_vantage_news
|
||||
)
|
||||
from .alpha_vantage_common import AlphaVantageRateLimitError
|
||||
from . import akshare
|
||||
from .akshare import AKShareRateLimitError
|
||||
|
||||
# Configuration and routing logic
|
||||
from .config import get_config
|
||||
from . import config
|
||||
|
||||
# Tools organized by category
|
||||
TOOLS_CATEGORIES = {
|
||||
|
|
@ -57,6 +79,7 @@ TOOLS_CATEGORIES = {
|
|||
VENDOR_LIST = [
|
||||
"local",
|
||||
"yfinance",
|
||||
"akshare",
|
||||
"openai",
|
||||
"google"
|
||||
]
|
||||
|
|
@ -67,6 +90,7 @@ VENDOR_METHODS = {
|
|||
"get_stock_data": {
|
||||
"alpha_vantage": get_alpha_vantage_stock,
|
||||
"yfinance": get_YFin_data_online,
|
||||
"akshare": _VendorFunctionProxy(akshare, 'get_akshare_stock_data'),
|
||||
"local": get_YFin_data,
|
||||
},
|
||||
# technical_indicators
|
||||
|
|
@ -100,8 +124,8 @@ VENDOR_METHODS = {
|
|||
"get_news": {
|
||||
"alpha_vantage": get_alpha_vantage_news,
|
||||
"openai": get_stock_news_openai,
|
||||
"google": get_google_news,
|
||||
"local": [get_finnhub_news, get_reddit_company_news, get_google_news],
|
||||
"google": get_google_news_for_ticker,
|
||||
"local": [get_finnhub_news, get_reddit_company_news, get_google_news_for_ticker],
|
||||
},
|
||||
"get_global_news": {
|
||||
"openai": get_global_news_openai,
|
||||
|
|
@ -129,16 +153,21 @@ 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()
|
||||
cfg = config.get_config()
|
||||
|
||||
# Check tool-level configuration first (if method provided)
|
||||
if method:
|
||||
tool_vendors = config.get("tool_vendors", {})
|
||||
tool_vendors = cfg.get("tool_vendors", {})
|
||||
if method in tool_vendors:
|
||||
return tool_vendors[method]
|
||||
|
||||
# Support both data_vendors (category-based) and data_vendor (simple) formats
|
||||
# data_vendor (singular) takes precedence if present (for backward compatibility)
|
||||
if "data_vendor" in cfg:
|
||||
return cfg["data_vendor"]
|
||||
|
||||
# Fall back to category-level configuration
|
||||
return config.get("data_vendors", {}).get(category, "default")
|
||||
return cfg.get("data_vendors", {}).get(category, "default")
|
||||
|
||||
def route_to_vendor(method: str, *args, **kwargs):
|
||||
"""Route method calls to appropriate vendor implementation with fallback support."""
|
||||
|
|
@ -211,6 +240,12 @@ def route_to_vendor(method: str, *args, **kwargs):
|
|||
print(f"DEBUG: Rate limit details: {e}")
|
||||
# Continue to next vendor for fallback
|
||||
continue
|
||||
except AKShareRateLimitError as e:
|
||||
if vendor == "akshare":
|
||||
print(f"RATE_LIMIT: AKShare 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}")
|
||||
|
|
|
|||
|
|
@ -20,7 +20,7 @@ DEFAULT_CONFIG = {
|
|||
# Data vendor configuration
|
||||
# Category-level configuration (default for all tools in category)
|
||||
"data_vendors": {
|
||||
"core_stock_apis": "yfinance", # Options: yfinance, alpha_vantage, local
|
||||
"core_stock_apis": "yfinance", # Options: yfinance, akshare, 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
|
||||
|
|
|
|||
Loading…
Reference in New Issue