huibao/backend/app/models.py

"""
数据库模型定义
"""
from sqlalchemy import Column, Integer, String, Text, DateTime, Boolean, ForeignKey, Enum
from sqlalchemy.orm import relationship
from datetime import datetime, timezone
import enum

from app.database import Base


def utc_now():
    """获取当前UTC时间 (不带时区，适配PostgreSQL)"""
    # asyncpg 要求 TIMESTAMP WITHOUT TIME ZONE 字段必须传入 naive datetime
    return datetime.now(timezone.utc).replace(tzinfo=None)


class AnalysisStatus(enum.Enum):
    """分析状态枚举"""
    PENDING = "pending"          # 待分析
    EXTRACTING = "extracting"    # 提取中
    ANALYZING = "analyzing"      # 分析中
    SUMMARIZING = "summarizing"  # 汇总中
    COMPLETED = "completed"      # 已完成
    FAILED = "failed"            # 失败


class Company(Base):
    """上市公司表"""
    __tablename__ = "companies"
    
    id = Column(Integer, primary_key=True, index=True)
    stock_code = Column(String(10), unique=True, index=True, nullable=False, comment="股票代码")
    company_name = Column(String(100), nullable=False, comment="公司名称")
    short_name = Column(String(50), comment="公司简称")
    industry = Column(String(50), comment="所属行业")
    org_id = Column(String(50), comment="巨潮机构ID")
    is_active = Column(Boolean, default=True, comment="是否启用监控")
    created_at = Column(DateTime, default=utc_now, comment="创建时间")
    updated_at = Column(DateTime, default=utc_now, onupdate=utc_now, comment="更新时间")
    
    # 关联报告
    reports = relationship("Report", back_populates="company", cascade="all, delete-orphan")
    
    def __repr__(self):
        return f"<Company {self.stock_code} {self.short_name}>"


class Report(Base):
    """年报/半年报表"""
    __tablename__ = "reports"
    
    id = Column(Integer, primary_key=True, index=True)
    company_id = Column(Integer, ForeignKey("companies.id"), nullable=False)
    
    # 报告基本信息
    title = Column(String(500), nullable=False, comment="报告标题")
    report_type = Column(String(50), comment="报告类型: 年度报告/半年度报告")
    report_year = Column(Integer, comment="报告年份")
    report_period = Column(String(20), comment="报告期间: 年报/半年报")
    
    # 巨潮信息
    announcement_id = Column(String(100), unique=True, comment="公告ID")
    announcement_time = Column(DateTime, comment="公告时间")
    pdf_url = Column(String(500), comment="PDF下载链接")
    
    # 本地存储
    local_path = Column(String(500), comment="本地存储路径")
    file_size = Column(Integer, comment="文件大小(字节)")
    
    # 状态
    is_downloaded = Column(Boolean, default=False, comment="是否已下载")
    is_extracted = Column(Boolean, default=False, comment="是否已提取内容")
    is_analyzed = Column(Boolean, default=False, comment="是否已分析")
    analysis_status = Column(String(20), default=AnalysisStatus.PENDING.value, comment="分析状态")
    
    created_at = Column(DateTime, default=utc_now, comment="创建时间")
    updated_at = Column(DateTime, default=utc_now, onupdate=utc_now, comment="更新时间")
    
    # 关联
    company = relationship("Company", back_populates="reports")
    extracted_contents = relationship("ExtractedContent", back_populates="report", cascade="all, delete-orphan")
    analysis_results = relationship("AnalysisResult", back_populates="report", cascade="all, delete-orphan")
    
    def __repr__(self):
        return f"<Report {self.title}>"


class ExtractedContent(Base):
    """提取的内容表"""
    __tablename__ = "extracted_contents"
    
    id = Column(Integer, primary_key=True, index=True)
    report_id = Column(Integer, ForeignKey("reports.id"), nullable=False)
    
    section_name = Column(String(100), comment="章节名称")
    section_keyword = Column(String(50), comment="匹配的关键词")
    content = Column(Text, comment="提取的内容")
    page_start = Column(Integer, comment="起始页码")
    page_end = Column(Integer, comment="结束页码")
    char_count = Column(Integer, comment="字符数")
    
    created_at = Column(DateTime, default=utc_now, comment="创建时间")
    
    # 关联
    report = relationship("Report", back_populates="extracted_contents")
    
    def __repr__(self):
        return f"<ExtractedContent {self.section_name}>"


class AnalysisResult(Base):
    """AI分析结果表"""
    __tablename__ = "analysis_results"
    
    id = Column(Integer, primary_key=True, index=True)
    report_id = Column(Integer, ForeignKey("reports.id"), nullable=False)
    
    # 分析信息
    analysis_type = Column(String(50), comment="分析类型: section/summary")
    section_name = Column(String(100), comment="分析的章节名称(如果是章节分析)")
    ai_model = Column(String(50), comment="使用的AI模型")
    
    # 分析内容
    prompt = Column(Text, comment="发送给AI的提示词")
    response = Column(Text, comment="AI的原始响应")
    summary = Column(Text, comment="提取的摘要")
    key_points = Column(Text, comment="关键要点(JSON格式)")
    
    # 状态
    is_final = Column(Boolean, default=False, comment="是否为最终汇总结果")
    token_count = Column(Integer, comment="消耗的token数")
    
    created_at = Column(DateTime, default=utc_now, comment="创建时间")
    
    # 关联
    report = relationship("Report", back_populates="analysis_results")
    
    def __repr__(self):
        return f"<AnalysisResult {self.analysis_type} for Report {self.report_id}>"


class TaskLog(Base):
    """任务日志表"""
    __tablename__ = "task_logs"
    
    id = Column(Integer, primary_key=True, index=True)
    
    task_type = Column(String(50), comment="任务类型: crawl/extract/analyze")
    task_name = Column(String(200), comment="任务名称")
    status = Column(String(20), comment="状态: started/running/completed/failed")
    message = Column(Text, comment="日志消息")
    error = Column(Text, comment="错误信息")
    
    started_at = Column(DateTime, default=utc_now, comment="开始时间")
    completed_at = Column(DateTime, comment="完成时间")
    
    def __repr__(self):
        return f"<TaskLog {self.task_type} {self.status}>"