56 lines
2.2 KiB
Python
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 |