TradingAgents/app/domain/models.py

56 lines
2.2 KiB
Python

from datetime import date, datetime
from typing import List, Optional
from sqlmodel import Field, SQLModel, JSON, Column
import enum
class User(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
email: str = Field(unique=True, index=True)
username: str = Field(unique=True, index=True)
hashed_password: str
first_name: Optional[str] = None
last_name: Optional[str] = None
is_active: bool = Field(default=True)
is_superuser: bool = Field(default=False)
created_at: datetime = Field(default_factory=datetime.utcnow)
updated_at: datetime = Field(default_factory=datetime.utcnow, sa_column_kwargs={"onupdate": datetime.utcnow})
class UserProfile(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
user_id: int = Field(foreign_key="user.id", unique=True)
encrypted_openai_api_key: Optional[str] = None
default_ticker: str = Field(default="SPY")
preferred_research_depth: int = Field(default=3)
preferred_shallow_thinker: str = Field(default="gpt-4o-mini")
preferred_deep_thinker: str = Field(default="gpt-4o")
created_at: datetime = Field(default_factory=datetime.utcnow)
updated_at: datetime = Field(default_factory=datetime.utcnow, sa_column_kwargs={"onupdate": datetime.utcnow})
class AnalysisStatus(str, enum.Enum):
PENDING = "pending"
RUNNING = "running"
COMPLETED = "completed"
FAILED = "failed"
CANCELLED = "cancelled"
class AnalysisSession(SQLModel, table=True):
id: Optional[int] = Field(default=None, primary_key=True)
user_id: int = Field(foreign_key="user.id")
ticker: str
analysis_date: date
analysts_selected: List[str] = Field(sa_column=Column(JSON))
research_depth: int
llm_provider: str
backend_url: str
shallow_thinker: str
deep_thinker: str
status: AnalysisStatus = Field(default=AnalysisStatus.PENDING)
final_report: Optional[str] = None
error_message: Optional[str] = None
created_at: datetime = Field(default_factory=datetime.utcnow)
started_at: Optional[datetime] = None
completed_at: Optional[datetime] = None