TradingAgents/tradingagents/agents/discovery/models.py

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"),
        )