diff options
Diffstat (limited to 'shared')
| -rw-r--r-- | shared/src/shared/models.py | 23 | ||||
| -rw-r--r-- | shared/src/shared/sentiment_models.py | 44 | ||||
| -rw-r--r-- | shared/tests/test_sentiment_models.py | 114 |
3 files changed, 181 insertions, 0 deletions
diff --git a/shared/src/shared/models.py b/shared/src/shared/models.py index 70820b5..a436c03 100644 --- a/shared/src/shared/models.py +++ b/shared/src/shared/models.py @@ -74,3 +74,26 @@ class Position(BaseModel): @property def unrealized_pnl(self) -> Decimal: return self.quantity * (self.current_price - self.avg_entry_price) + + +class NewsCategory(str, Enum): + POLICY = "policy" + EARNINGS = "earnings" + MACRO = "macro" + SOCIAL = "social" + FILING = "filing" + FED = "fed" + + +class NewsItem(BaseModel): + id: str = Field(default_factory=lambda: str(uuid.uuid4())) + source: str + headline: str + summary: Optional[str] = None + url: Optional[str] = None + published_at: datetime + symbols: list[str] = [] + sentiment: float + category: NewsCategory + raw_data: dict = {} + created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc)) diff --git a/shared/src/shared/sentiment_models.py b/shared/src/shared/sentiment_models.py new file mode 100644 index 0000000..a009601 --- /dev/null +++ b/shared/src/shared/sentiment_models.py @@ -0,0 +1,44 @@ +"""Sentiment scoring and stock selection models.""" + +from datetime import datetime +from typing import Optional + +from pydantic import BaseModel + +from shared.models import OrderSide + + +class SymbolScore(BaseModel): + symbol: str + news_score: float + news_count: int + social_score: float + policy_score: float + filing_score: float + composite: float + updated_at: datetime + + +class MarketSentiment(BaseModel): + fear_greed: int + fear_greed_label: str + vix: Optional[float] = None + fed_stance: str + market_regime: str + updated_at: datetime + + +class SelectedStock(BaseModel): + symbol: str + side: OrderSide + conviction: float + reason: str + key_news: list[str] + + +class Candidate(BaseModel): + symbol: str + source: str + direction: Optional[OrderSide] = None + score: float + reason: str diff --git a/shared/tests/test_sentiment_models.py b/shared/tests/test_sentiment_models.py new file mode 100644 index 0000000..74f1acd --- /dev/null +++ b/shared/tests/test_sentiment_models.py @@ -0,0 +1,114 @@ +"""Tests for news and sentiment models.""" + +import pytest +from datetime import datetime, timezone + +from shared.models import NewsCategory, NewsItem, OrderSide +from shared.sentiment_models import SymbolScore, MarketSentiment, SelectedStock, Candidate + + +def test_news_item_defaults(): + item = NewsItem( + source="finnhub", + headline="Test headline", + published_at=datetime(2026, 4, 2, tzinfo=timezone.utc), + sentiment=0.5, + category=NewsCategory.MACRO, + ) + assert item.id + assert item.symbols == [] + assert item.summary is None + assert item.raw_data == {} + assert item.created_at is not None + + +def test_news_item_with_symbols(): + item = NewsItem( + source="rss", + headline="AAPL earnings beat", + published_at=datetime(2026, 4, 2, tzinfo=timezone.utc), + sentiment=0.8, + category=NewsCategory.EARNINGS, + symbols=["AAPL"], + ) + assert item.symbols == ["AAPL"] + assert item.category == NewsCategory.EARNINGS + + +def test_news_category_values(): + assert NewsCategory.POLICY == "policy" + assert NewsCategory.EARNINGS == "earnings" + assert NewsCategory.MACRO == "macro" + assert NewsCategory.SOCIAL == "social" + assert NewsCategory.FILING == "filing" + assert NewsCategory.FED == "fed" + + +def test_symbol_score(): + score = SymbolScore( + symbol="AAPL", + news_score=0.5, + news_count=10, + social_score=0.3, + policy_score=0.0, + filing_score=0.2, + composite=0.3, + updated_at=datetime(2026, 4, 2, tzinfo=timezone.utc), + ) + assert score.symbol == "AAPL" + assert score.composite == 0.3 + + +def test_market_sentiment(): + ms = MarketSentiment( + fear_greed=25, + fear_greed_label="Extreme Fear", + vix=32.5, + fed_stance="hawkish", + market_regime="risk_off", + updated_at=datetime(2026, 4, 2, tzinfo=timezone.utc), + ) + assert ms.market_regime == "risk_off" + assert ms.vix == 32.5 + + +def test_market_sentiment_no_vix(): + ms = MarketSentiment( + fear_greed=50, + fear_greed_label="Neutral", + fed_stance="neutral", + market_regime="neutral", + updated_at=datetime(2026, 4, 2, tzinfo=timezone.utc), + ) + assert ms.vix is None + + +def test_selected_stock(): + ss = SelectedStock( + symbol="NVDA", + side=OrderSide.BUY, + conviction=0.85, + reason="CHIPS Act expansion", + key_news=["Trump signs CHIPS Act expansion"], + ) + assert ss.conviction == 0.85 + assert len(ss.key_news) == 1 + + +def test_candidate(): + c = Candidate( + symbol="TSLA", + source="sentiment", + direction=OrderSide.BUY, + score=0.75, + reason="High social buzz", + ) + assert c.direction == OrderSide.BUY + + c2 = Candidate( + symbol="XOM", + source="llm", + score=0.6, + reason="Oil price surge", + ) + assert c2.direction is None |
