diff --git a/api/models.py b/api/models.py index 78567f2..0cb2614 100644 --- a/api/models.py +++ b/api/models.py @@ -177,6 +177,8 @@ class MarketDataResponse(BaseModel): market_cap: Optional[float] = None pe_ratio: Optional[float] = None timestamp: datetime + provider: Optional[str] = None + feed_latency_ms: Optional[float] = None class TechnicalIndicatorsResponse(BaseModel): diff --git a/docs/user-guide/websocket-feeds.md b/docs/user-guide/websocket-feeds.md new file mode 100644 index 0000000..77f86a0 --- /dev/null +++ b/docs/user-guide/websocket-feeds.md @@ -0,0 +1,48 @@ +# WebSocket Market Data Feeds + +TradeGraph now supports low-latency market data by connecting to real-time vendor feeds. The financial analysis agent uses these feeds automatically when `include_market_data=True`, so you can mix fresh tick data with Yahoo Finance fundamentals in a single run. + +## Supported Providers + +| Provider | Asset Class | Notes | +| --- | --- | --- | +| Polygon | Equities | Requires `POLYGON_API_KEY`. Uses trades channel `T.*`. | +| Alpaca | Equities | Requires `ALPACA_API_KEY` and `ALPACA_API_SECRET`. Supports `ALPACA_DATA_FEED` (`iex` or `sip`). | +| Binance | Crypto | No authentication required. Streams trades via `symbol@trade`. | + +Use the environment variables (or settings) to select defaults: + +```env +DEFAULT_EQUITY_FEED_PROVIDER=polygon +DEFAULT_CRYPTO_FEED_PROVIDER=binance +POLYGON_API_KEY=... +ALPACA_API_KEY=... +ALPACA_API_SECRET=... +ALPACA_DATA_FEED=iex +WEBSOCKET_TIMEOUT_SECONDS=10 +``` + +You can override feeds per symbol at runtime with `asset_types` and `feed_overrides` fields on the `FinancialAnalysisAgent` input payload. + +## Monitoring Latency + +Each `MarketData` entry now reports two additional fields: + +``` +provider: Source of the trade (`polygon`, `alpaca`, `binance`, ...) +feed_latency_ms: Milliseconds between trade timestamp and ingestion +``` + +This helps dashboards surface stale feeds or connectivity issues. + +## Quick CLI Test + +The repository ships with `examples/websocket_feed_demo.py` for manual testing. + +```bash +python examples/websocket_feed_demo.py AAPL --asset-type equity \ + --polygon-api-key $POLYGON_API_KEY +``` + +Specify `--provider alpaca` or `--asset-type crypto` to route to other feeds. The script prints the decoded trade payload along with the calculated latency so you can verify credentials before running full analyses. +``` diff --git a/examples/websocket_feed_demo.py b/examples/websocket_feed_demo.py new file mode 100644 index 0000000..6904682 --- /dev/null +++ b/examples/websocket_feed_demo.py @@ -0,0 +1,103 @@ +"""Quick script to test streaming market data feeds.""" + +import argparse +import asyncio +import os +from pprint import pprint + +from tradegraph_financial_advisor.services.market_data import ( + MarketDataFeedConfig, + MarketDataWebSocketClient, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Test WebSocket market feeds") + parser.add_argument("symbol", help="Ticker symbol, e.g., AAPL or BTCUSDT") + parser.add_argument( + "--asset-type", + default="equity", + choices=["equity", "crypto"], + help="Asset class for routing to the right feed", + ) + parser.add_argument( + "--provider", + default=None, + help="Override provider (binance, polygon, alpaca)", + ) + parser.add_argument( + "--equity-provider", + default=os.getenv("DEFAULT_EQUITY_FEED_PROVIDER", "polygon"), + help="Default provider for equities", + ) + parser.add_argument( + "--crypto-provider", + default=os.getenv("DEFAULT_CRYPTO_FEED_PROVIDER", "binance"), + help="Default provider for crypto symbols", + ) + parser.add_argument( + "--polygon-api-key", + default=os.getenv("POLYGON_API_KEY"), + help="Polygon API key", + ) + parser.add_argument( + "--alpaca-api-key", + default=os.getenv("ALPACA_API_KEY"), + help="Alpaca API key", + ) + parser.add_argument( + "--alpaca-api-secret", + default=os.getenv("ALPACA_API_SECRET"), + help="Alpaca API secret", + ) + parser.add_argument( + "--alpaca-feed", + default=os.getenv("ALPACA_DATA_FEED", "iex"), + help="Alpaca feed (iex or sip)", + ) + parser.add_argument( + "--timeout", + type=int, + default=int(os.getenv("WEBSOCKET_TIMEOUT_SECONDS", "10")), + help="Seconds to wait for a trade message", + ) + return parser.parse_args() + + +def build_config(args: argparse.Namespace) -> MarketDataFeedConfig: + return MarketDataFeedConfig( + equity_provider=args.equity_provider, + crypto_provider=args.crypto_provider, + polygon_api_key=args.polygon_api_key, + alpaca_api_key=args.alpaca_api_key, + alpaca_api_secret=args.alpaca_api_secret, + alpaca_feed=args.alpaca_feed, + timeout_seconds=args.timeout, + ) + + +async def main() -> None: + args = parse_args() + config = build_config(args) + + async with MarketDataWebSocketClient(config=config) as client: + trade = await client.get_realtime_trade( + args.symbol, + asset_type=args.asset_type, + provider_override=args.provider, + ) + + pprint( + { + "symbol": trade.symbol, + "price": trade.price, + "size": trade.size, + "provider": trade.provider, + "timestamp": trade.timestamp.isoformat(), + "latency_ms": trade.latency_ms, + } + ) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/mkdocs.yml b/mkdocs.yml index e3b6f21..f19cd3e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -84,6 +84,7 @@ nav: - Python API: user-guide/python-api.md - Web Interface: user-guide/web-interface.md - Analysis Types: user-guide/analysis-types.md + - WebSocket Feeds: user-guide/websocket-feeds.md - API Reference: - Core Classes: api-reference/core.md - Agents: api-reference/agents.md @@ -113,4 +114,4 @@ extra: - icon: fontawesome/brands/python link: https://pypi.org/project/tradegraph-financial-advisor version: - provider: mike \ No newline at end of file + provider: mike diff --git a/src/tradegraph_financial_advisor/agents/financial_agent.py b/src/tradegraph_financial_advisor/agents/financial_agent.py index c54376c..43ef3cf 100644 --- a/src/tradegraph_financial_advisor/agents/financial_agent.py +++ b/src/tradegraph_financial_advisor/agents/financial_agent.py @@ -8,6 +8,11 @@ from .base_agent import BaseAgent from ..models.financial_data import CompanyFinancials, MarketData, TechnicalIndicators from ..config.settings import settings +from ..services.market_data import ( + MarketDataFeedConfig, + MarketDataWebSocketClient, + WebSocketFeedError, +) class FinancialAnalysisAgent(BaseAgent): @@ -18,16 +23,33 @@ def __init__(self, **kwargs): **kwargs, ) self.session: Optional[aiohttp.ClientSession] = None + self.market_data_stream: Optional[MarketDataWebSocketClient] = None async def start(self) -> None: await super().start() self.session = aiohttp.ClientSession( timeout=aiohttp.ClientTimeout(total=settings.analysis_timeout_seconds) ) + feed_config = MarketDataFeedConfig( + equity_provider=settings.default_equity_feed_provider, + crypto_provider=settings.default_crypto_feed_provider, + polygon_api_key=settings.polygon_api_key, + alpaca_api_key=settings.alpaca_api_key, + alpaca_api_secret=settings.alpaca_api_secret, + alpaca_feed=settings.alpaca_data_feed, + timeout_seconds=settings.websocket_timeout_seconds, + ) + self.market_data_stream = MarketDataWebSocketClient( + session=self.session, + config=feed_config, + ) async def stop(self) -> None: if self.session: await self.session.close() + if self.market_data_stream: + await self.market_data_stream.aclose() + self.market_data_stream = None await super().stop() async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]: @@ -35,6 +57,8 @@ async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]: include_financials = input_data.get("include_financials", True) include_technical = input_data.get("include_technical", True) include_market_data = input_data.get("include_market_data", True) + asset_types = input_data.get("asset_types", {}) # symbol -> equity/crypto + feed_overrides = input_data.get("feed_overrides", {}) # symbol -> provider name logger.info(f"Analyzing financial data for symbols: {symbols}") @@ -45,7 +69,11 @@ async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]: symbol_data = {} if include_market_data: - market_data = await self._get_market_data(symbol) + asset_type = asset_types.get(symbol, "equity") + feed_override = feed_overrides.get(symbol) + market_data = await self._get_market_data( + symbol, asset_type=asset_type, feed_override=feed_override + ) symbol_data["market_data"] = ( market_data.dict() if market_data else None ) @@ -73,32 +101,43 @@ async def execute(self, input_data: Dict[str, Any]) -> Dict[str, Any]: "analysis_timestamp": datetime.now().isoformat(), } - async def _get_market_data(self, symbol: str) -> Optional[MarketData]: + async def _get_market_data( + self, symbol: str, asset_type: str = "equity", feed_override: Optional[str] = None + ) -> Optional[MarketData]: try: - ticker = yf.Ticker(symbol) - info = ticker.info - history = ticker.history(period="1d") + if not self.market_data_stream: + raise RuntimeError("Market data stream is not initialized. Did you call start()?") - if history.empty: - return None + trade = await self.market_data_stream.get_realtime_trade( + symbol, asset_type=asset_type, provider_override=feed_override + ) - latest = history.iloc[-1] + ticker = yf.Ticker(symbol) + info = ticker.info + previous_close = info.get("previousClose") + change = trade.price - previous_close if previous_close else 0.0 + change_percent = ( + (change / previous_close) * 100 if previous_close else 0.0 + ) market_data = MarketData( symbol=symbol, - current_price=float(latest["Close"]), - change=float(latest["Close"] - latest["Open"]), - change_percent=float( - (latest["Close"] - latest["Open"]) / latest["Open"] * 100 - ), - volume=int(latest["Volume"]), + current_price=float(trade.price), + change=float(change), + change_percent=float(change_percent), + volume=int(trade.size or 0), market_cap=info.get("marketCap"), pe_ratio=info.get("trailingPE"), - timestamp=datetime.now(), + timestamp=trade.timestamp, + provider=trade.provider, + feed_latency_ms=trade.latency_ms, ) return market_data + except WebSocketFeedError as e: + logger.error(f"WebSocket feed error for {symbol}: {str(e)}") + return None except Exception as e: logger.error(f"Error fetching market data for {symbol}: {str(e)}") return None diff --git a/src/tradegraph_financial_advisor/config/settings.py b/src/tradegraph_financial_advisor/config/settings.py index f509ca9..4f99d1a 100644 --- a/src/tradegraph_financial_advisor/config/settings.py +++ b/src/tradegraph_financial_advisor/config/settings.py @@ -10,10 +10,21 @@ class Settings(BaseSettings): openai_api_key: str = Field("", env="OPENAI_API_KEY") alpha_vantage_api_key: Optional[str] = Field(None, env="ALPHA_VANTAGE_API_KEY") financial_data_api_key: Optional[str] = Field(None, env="FINANCIAL_DATA_API_KEY") + polygon_api_key: Optional[str] = Field(None, env="POLYGON_API_KEY") + alpaca_api_key: Optional[str] = Field(None, env="ALPACA_API_KEY") + alpaca_api_secret: Optional[str] = Field(None, env="ALPACA_API_SECRET") log_level: str = Field("INFO", env="LOG_LEVEL") max_concurrent_agents: int = Field(5, env="MAX_CONCURRENT_AGENTS") analysis_timeout_seconds: int = Field(30, env="ANALYSIS_TIMEOUT_SECONDS") + websocket_timeout_seconds: int = Field(10, env="WEBSOCKET_TIMEOUT_SECONDS") + default_equity_feed_provider: str = Field( + "polygon", env="DEFAULT_EQUITY_FEED_PROVIDER" + ) + default_crypto_feed_provider: str = Field( + "binance", env="DEFAULT_CRYPTO_FEED_PROVIDER" + ) + alpaca_data_feed: str = Field("iex", env="ALPACA_DATA_FEED") news_sources: List[str] = Field( default_factory=lambda: [ diff --git a/src/tradegraph_financial_advisor/models/financial_data.py b/src/tradegraph_financial_advisor/models/financial_data.py index a2b18aa..83d26a7 100644 --- a/src/tradegraph_financial_advisor/models/financial_data.py +++ b/src/tradegraph_financial_advisor/models/financial_data.py @@ -53,6 +53,8 @@ class MarketData(BaseModel): market_cap: Optional[float] = None pe_ratio: Optional[float] = None timestamp: datetime + provider: Optional[str] = None + feed_latency_ms: Optional[float] = None class TechnicalIndicators(BaseModel): diff --git a/src/tradegraph_financial_advisor/services/local_scraping_service.py b/src/tradegraph_financial_advisor/services/local_scraping_service.py index 864338f..6b5c02c 100644 --- a/src/tradegraph_financial_advisor/services/local_scraping_service.py +++ b/src/tradegraph_financial_advisor/services/local_scraping_service.py @@ -1,19 +1,44 @@ from typing import Any, Dict, List from loguru import logger -from ddgs import DDGS -from crawl4ai import AsyncWebCrawler + +try: + from crawl4ai import AsyncWebCrawler +except ModuleNotFoundError: # pragma: no cover - optional dependency + AsyncWebCrawler = None # type: ignore + logger.warning( + "crawl4ai package is not installed; LocalScrapingService scraping is disabled" + ) + +try: + from ddgs import DDGS +except ModuleNotFoundError: # pragma: no cover - optional dependency for scraping + DDGS = None # type: ignore + logger.warning( + "ddgs package is not installed; LocalScrapingService search methods are disabled" + ) from ..models.financial_data import NewsArticle from ..utils.helpers import generate_summary class LocalScrapingService: def __init__(self): - self.crawler = AsyncWebCrawler() + self.crawler = AsyncWebCrawler() if AsyncWebCrawler else None + + def _ensure_dependencies(self) -> None: + if AsyncWebCrawler is None: + raise RuntimeError( + "crawl4ai package is required for scraping but is not installed" + ) + if DDGS is None: + raise RuntimeError( + "ddgs package is required for scraping but is not installed" + ) async def search_and_scrape_news( self, symbols: List[str], max_articles_per_symbol: int = 5 ) -> List[NewsArticle]: all_articles = [] + self._ensure_dependencies() async with DDGS() as ddgs: for symbol in symbols: query = f"{symbol} stock news" @@ -53,6 +78,7 @@ async def search_and_scrape_financial_reports( ) -> List[Dict[str, Any]]: query = f"{company_symbol} {report_type} site:sec.gov" filings = [] + self._ensure_dependencies() async with DDGS() as ddgs: try: results = await ddgs.text( @@ -85,11 +111,13 @@ async def search_and_scrape_financial_reports( async def start(self): logger.info("LocalScrapingService started.") + self._ensure_dependencies() await self.crawler.start() async def stop(self): logger.info("LocalScrapingService stopped.") - await self.crawler.close() + if self.crawler: + await self.crawler.close() async def health_check(self) -> bool: # For now, we assume the service is healthy if it can be instantiated. diff --git a/src/tradegraph_financial_advisor/services/market_data/__init__.py b/src/tradegraph_financial_advisor/services/market_data/__init__.py new file mode 100644 index 0000000..3b68fe9 --- /dev/null +++ b/src/tradegraph_financial_advisor/services/market_data/__init__.py @@ -0,0 +1,15 @@ +"""Market data services.""" + +from .streaming import ( + MarketDataFeedConfig, + MarketDataWebSocketClient, + RealtimeTrade, + WebSocketFeedError, +) + +__all__ = [ + "MarketDataFeedConfig", + "MarketDataWebSocketClient", + "RealtimeTrade", + "WebSocketFeedError", +] diff --git a/src/tradegraph_financial_advisor/services/market_data/streaming.py b/src/tradegraph_financial_advisor/services/market_data/streaming.py new file mode 100644 index 0000000..3336ff3 --- /dev/null +++ b/src/tradegraph_financial_advisor/services/market_data/streaming.py @@ -0,0 +1,238 @@ +"""WebSocket-based market data feeds used across the project.""" + +from __future__ import annotations + +import asyncio +import json +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Any, Dict, Optional + +import aiohttp +from aiohttp import ClientSession, WSMsgType +from loguru import logger + + +def _utc_now() -> datetime: + return datetime.now(timezone.utc) + + +@dataclass +class RealtimeTrade: + """Represents the latest trade received from a WebSocket feed.""" + + symbol: str + price: float + size: Optional[float] + timestamp: datetime + provider: str + raw: Dict[str, Any] = field(default_factory=dict) + + @property + def latency_ms(self) -> float: + """Approximate feed latency in milliseconds.""" + return max((_utc_now() - self.timestamp).total_seconds() * 1000, 0.0) + + +@dataclass +class MarketDataFeedConfig: + """Configuration for selecting WebSocket feeds.""" + + equity_provider: str = "polygon" + crypto_provider: str = "binance" + polygon_api_key: Optional[str] = None + alpaca_api_key: Optional[str] = None + alpaca_api_secret: Optional[str] = None + alpaca_feed: str = "iex" + timeout_seconds: int = 8 + + def provider_for(self, asset_type: str, override: Optional[str] = None) -> str: + asset_type = asset_type.lower() + if override: + return override.lower() + if asset_type == "crypto": + return self.crypto_provider.lower() + return self.equity_provider.lower() + + +class WebSocketFeedError(RuntimeError): + """Raised when a WebSocket feed cannot be used.""" + + +class MarketDataWebSocketClient: + """Small helper that fetches single-trade snapshots from vendor feeds.""" + + def __init__( + self, + session: Optional[ClientSession] = None, + config: Optional[MarketDataFeedConfig] = None, + ) -> None: + self._session = session + self._owns_session = session is None + self.config = config or MarketDataFeedConfig() + + async def __aenter__(self) -> "MarketDataWebSocketClient": + await self._ensure_session() + return self + + async def __aexit__(self, *exc_info: Any) -> None: + await self.aclose() + + async def aclose(self) -> None: + if self._owns_session and self._session: + await self._session.close() + self._session = None + + async def _ensure_session(self) -> ClientSession: + if self._session is None or self._session.closed: + timeout = aiohttp.ClientTimeout(total=self.config.timeout_seconds + 2) + self._session = ClientSession(timeout=timeout) + self._owns_session = True + return self._session + + async def get_realtime_trade( + self, + symbol: str, + asset_type: str = "equity", + provider_override: Optional[str] = None, + ) -> RealtimeTrade: + """Fetch the next available trade for the requested symbol.""" + + provider = self.config.provider_for(asset_type, provider_override) + symbol = symbol.upper() + + if provider == "binance": + return await self._fetch_from_binance(symbol) + if provider == "polygon": + return await self._fetch_from_polygon(symbol) + if provider == "alpaca": + return await self._fetch_from_alpaca(symbol) + + raise WebSocketFeedError(f"Unsupported WebSocket provider: {provider}") + + async def _fetch_from_binance(self, symbol: str) -> RealtimeTrade: + session = await self._ensure_session() + stream_symbol = self._format_binance_symbol(symbol) + url = f"wss://stream.binance.com:9443/ws/{stream_symbol}@trade" + logger.debug(f"Connecting to Binance WebSocket for {stream_symbol}") + + try: + async with session.ws_connect(url, heartbeat=45) as ws: + payload = await self._next_json(ws) + trade_event = payload + if "data" in payload: # combined stream response + trade_event = payload["data"] + + event_time = datetime.fromtimestamp( + trade_event.get("E", trade_event.get("eventTime", 0)) / 1000, + tz=timezone.utc, + ) + return RealtimeTrade( + symbol=symbol, + price=float(trade_event["p"]), + size=float(trade_event.get("q") or 0.0), + timestamp=event_time, + provider="binance", + raw=trade_event, + ) + except asyncio.TimeoutError as exc: + raise WebSocketFeedError("Timed out waiting for Binance trade data") from exc + except aiohttp.ClientError as exc: + raise WebSocketFeedError(f"Binance WebSocket error: {exc}") from exc + + def _format_binance_symbol(self, symbol: str) -> str: + normalized = symbol.replace("/", "").lower() + if not normalized.endswith("usdt"): + normalized = f"{normalized}usdt" + return normalized + + async def _fetch_from_polygon(self, symbol: str) -> RealtimeTrade: + if not self.config.polygon_api_key: + raise WebSocketFeedError( + "POLYGON_API_KEY is required for Polygon WebSocket access" + ) + + session = await self._ensure_session() + url = "wss://socket.polygon.io/stocks" + logger.debug(f"Connecting to Polygon WebSocket for {symbol}") + + try: + async with session.ws_connect(url, heartbeat=30) as ws: + await ws.send_json({"action": "auth", "params": self.config.polygon_api_key}) + await ws.send_json({"action": "subscribe", "params": f"T.{symbol}"}) + + while True: + payload = await self._next_json(ws) + events = payload if isinstance(payload, list) else [payload] + for event in events: + if event.get("ev") != "T": + continue + event_time = datetime.fromtimestamp(event["t"] / 1_000_000_000, tz=timezone.utc) + return RealtimeTrade( + symbol=symbol, + price=float(event["p"]), + size=float(event.get("s") or 0.0), + timestamp=event_time, + provider="polygon", + raw=event, + ) + except asyncio.TimeoutError as exc: + raise WebSocketFeedError("Timed out waiting for Polygon trade data") from exc + except aiohttp.ClientError as exc: + raise WebSocketFeedError(f"Polygon WebSocket error: {exc}") from exc + + async def _fetch_from_alpaca(self, symbol: str) -> RealtimeTrade: + if not self.config.alpaca_api_key or not self.config.alpaca_api_secret: + raise WebSocketFeedError( + "ALPACA_API_KEY and ALPACA_API_SECRET are required for Alpaca WebSocket access" + ) + + feed = self.config.alpaca_feed + session = await self._ensure_session() + url = f"wss://stream.data.alpaca.markets/v2/{feed}" + logger.debug(f"Connecting to Alpaca {feed} feed for {symbol}") + + try: + async with session.ws_connect(url, heartbeat=30) as ws: + await ws.send_json( + { + "action": "auth", + "key": self.config.alpaca_api_key, + "secret": self.config.alpaca_api_secret, + } + ) + await ws.send_json({"action": "subscribe", "trades": [symbol], "quotes": [], "bars": []}) + + while True: + payload = await self._next_json(ws) + events = payload if isinstance(payload, list) else [payload] + for event in events: + if event.get("T") != "t": + continue + timestamp = event.get("t") + if isinstance(timestamp, str): + event_time = datetime.fromisoformat(timestamp.replace("Z", "+00:00")) + else: + event_time = datetime.fromtimestamp(timestamp, tz=timezone.utc) + return RealtimeTrade( + symbol=symbol, + price=float(event.get("p")), + size=float(event.get("s") or 0.0), + timestamp=event_time, + provider="alpaca", + raw=event, + ) + except asyncio.TimeoutError as exc: + raise WebSocketFeedError("Timed out waiting for Alpaca trade data") from exc + except aiohttp.ClientError as exc: + raise WebSocketFeedError(f"Alpaca WebSocket error: {exc}") from exc + + async def _next_json(self, ws: aiohttp.ClientWebSocketResponse) -> Any: + msg = await asyncio.wait_for(ws.receive(), timeout=self.config.timeout_seconds) + if msg.type == WSMsgType.TEXT: + return json.loads(msg.data) + if msg.type == WSMsgType.BINARY: + return json.loads(msg.data.decode("utf-8")) + if msg.type in (WSMsgType.CLOSED, WSMsgType.ERROR): + raise WebSocketFeedError("WebSocket connection closed unexpectedly") + return {} diff --git a/src/tradegraph_financial_advisor/visualization/charts.py b/src/tradegraph_financial_advisor/visualization/charts.py index 304d01a..0a752bf 100644 --- a/src/tradegraph_financial_advisor/visualization/charts.py +++ b/src/tradegraph_financial_advisor/visualization/charts.py @@ -1,4 +1,7 @@ -import plotly.graph_objects as go +try: # pragma: no cover - optional visualization dependency + import plotly.graph_objects as go +except ModuleNotFoundError: # pragma: no cover + go = None # type: ignore def create_portfolio_allocation_chart( @@ -11,6 +14,9 @@ def create_portfolio_allocation_chart( recommendations: List of recommendation dicts from analysis results output_path: Where to save the HTML file """ + if go is None: + raise RuntimeError("plotly is required for chart generation but is not installed") + symbols = [rec["symbol"] for rec in recommendations] allocations = [rec["allocation_percentage"] * 100 for rec in recommendations] diff --git a/tests/unit/test_agents.py b/tests/unit/test_agents.py index 4316a70..e21b812 100644 --- a/tests/unit/test_agents.py +++ b/tests/unit/test_agents.py @@ -1,6 +1,6 @@ import pytest -from unittest.mock import patch -from datetime import datetime +from unittest.mock import AsyncMock, patch +from datetime import datetime, timezone from tradegraph_financial_advisor.agents.base_agent import BaseAgent from tradegraph_financial_advisor.agents.news_agent import NewsReaderAgent @@ -15,6 +15,30 @@ NewsArticle, SentimentType, ) +from tradegraph_financial_advisor.services.market_data.streaming import RealtimeTrade + + +@pytest.fixture +def mock_realtime_trade(): + return RealtimeTrade( + symbol="AAPL", + price=195.25, + size=150, + timestamp=datetime.now(timezone.utc), + provider="polygon", + raw={}, + ) + + +@pytest.fixture +def patch_market_data_client(mock_realtime_trade): + with patch( + "tradegraph_financial_advisor.agents.financial_agent.MarketDataWebSocketClient" + ) as mock_client_cls: + client_instance = mock_client_cls.return_value + client_instance.get_realtime_trade = AsyncMock(return_value=mock_realtime_trade) + client_instance.aclose = AsyncMock() + yield client_instance class TestBaseAgent: @@ -169,7 +193,9 @@ async def test_financial_agent_initialization(self): assert "financial" in agent.description.lower() @pytest.mark.asyncio - async def test_execute_financial_analysis(self, mock_yfinance_ticker): + async def test_execute_financial_analysis( + self, mock_yfinance_ticker, patch_market_data_client + ): """Test financial analysis execution.""" agent = FinancialAnalysisAgent() @@ -196,12 +222,16 @@ async def test_execute_financial_analysis(self, mock_yfinance_ticker): await agent.stop() @pytest.mark.asyncio - async def test_market_data_extraction(self, mock_yfinance_ticker): + async def test_market_data_extraction( + self, mock_yfinance_ticker, patch_market_data_client + ): """Test market data extraction.""" agent = FinancialAnalysisAgent() with patch("yfinance.Ticker", return_value=mock_yfinance_ticker): + await agent.start() market_data = await agent._get_market_data("AAPL") + await agent.stop() assert market_data is not None assert market_data.symbol == "AAPL"