ai-LiaoWangweb-app/app/models.py

from . import db, login_manager
from flask_login import UserMixin
from werkzeug.security import generate_password_hash, check_password_hash
from datetime import datetime

class User(UserMixin, db.Model):
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(64), index=True, unique=True)
    password_hash = db.Column(db.String(512))

    # SSO 相关字段
    sso_sub = db.Column(db.String(256), unique=True, nullable=True)
    real_name = db.Column(db.String(100), nullable=True)
    roles = db.Column(db.Text, nullable=True)  # JSON 字符串，如 ["super_admin", "sam_sys_admin"]
    email = db.Column(db.String(120), nullable=True)
    phone = db.Column(db.String(30), nullable=True)
    avatar_url = db.Column(db.String(500), nullable=True)

    def set_password(self, password):
        self.password_hash = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password_hash, password)

@login_manager.user_loader
def load_user(id):
    return User.query.get(int(id))

class SpiderSource(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(100), nullable=False, unique=True)
    code_identifier = db.Column(db.String(100), nullable=False, unique=True) # e.g. 'baidusearch' or 'generic_custom'
    description = db.Column(db.Text)
    status = db.Column(db.String(20), default='active') # active, disabled
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    
    # Generic Spider Configuration
    type = db.Column(db.String(20), default='script') # 'script', 'generic'
    url = db.Column(db.String(500)) # Base URL for generic spider
    method = db.Column(db.String(10), default='GET')
    headers = db.Column(db.Text) # JSON string for headers
    params = db.Column(db.Text) # JSON string for default params
    search_param_key = db.Column(db.String(50), default='q') # The query parameter key for keyword
    selectors = db.Column(db.Text) # JSON string for parsing rules: list, title, link, abstract, source, cover
    
    # Pagination Configuration
    has_pagination = db.Column(db.Boolean, default=False)
    pagination_param = db.Column(db.String(50)) # e.g. 'pn'
    pagination_step = db.Column(db.Integer, default=10) # e.g. 10
    pagination_start = db.Column(db.Integer, default=0) # e.g. 0
    
    tasks = db.relationship('SpiderTask', backref='spider_source', lazy='dynamic')

class SpiderTask(db.Model):
    __tablename__ = 'collection_task'
    id = db.Column(db.Integer, primary_key=True)
    keyword = db.Column(db.String(100), nullable=False)
    spider_source_id = db.Column(db.Integer, db.ForeignKey('spider_source.id'), nullable=False)
    pages = db.Column(db.Integer, default=1) # Number of pages to crawl
    status = db.Column(db.String(20), default='pending') # pending, running, completed, failed
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    finished_at = db.Column(db.DateTime)
    results = db.relationship('SpiderResult', backref='task', lazy='dynamic')

class SpiderResult(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    task_id = db.Column(db.Integer, db.ForeignKey('collection_task.id'))
    title = db.Column(db.String(500))
    abstract = db.Column(db.Text)
    source = db.Column(db.String(200))
    cover = db.Column(db.String(500))
    link = db.Column(db.Text)
    published_at = db.Column(db.String(50)) # Date extracted from source
    has_deep_collection = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

class DeepCollection(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(500)) # Title from spider result
    url = db.Column(db.String(500), nullable=False, unique=True)
    content = db.Column(db.Text) # Extracted content
    summary = db.Column(db.Text) # AI Summary (optional)
    status = db.Column(db.String(20), default='pending') # pending, running, completed, failed
    progress = db.Column(db.Integer, default=0)
    progress_msg = db.Column(db.String(200), default='')
    error_msg = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)

class AIModel(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    name = db.Column(db.String(100), nullable=False) # e.g. "SiliconFlow Qwen"
    provider = db.Column(db.String(50), default='openai') # e.g. openai, siliconflow, etc.
    api_base = db.Column(db.String(500), nullable=False) # URL
    api_key = db.Column(db.String(500), nullable=False)
    model_name = db.Column(db.String(200), nullable=False) # e.g. "Qwen/Qwen3-Next-80B-A3B-Instruct"
    system_prompt = db.Column(db.Text)
    is_active = db.Column(db.Boolean, default=True)
    is_default = db.Column(db.Boolean, default=False)
    total_tokens = db.Column(db.BigInteger, default=0)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    logs = db.relationship('TokenUsageLog', backref='model', lazy='dynamic')

class TokenUsageLog(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    model_id = db.Column(db.Integer, db.ForeignKey('ai_model.id'))
    prompt_tokens = db.Column(db.Integer, default=0)
    completion_tokens = db.Column(db.Integer, default=0)
    total_tokens = db.Column(db.Integer, default=0)
    request_type = db.Column(db.String(50), default='test') # test, deep_collect
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

class AIConversation(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    title = db.Column(db.String(200), default='New Chat')
    user_id = db.Column(db.Integer, db.ForeignKey('user.id'))
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    messages = db.relationship('AIMessage', backref='conversation', lazy='dynamic', cascade='all, delete-orphan')

class AIMessage(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    conversation_id = db.Column(db.Integer, db.ForeignKey('ai_conversation.id'), nullable=False)
    role = db.Column(db.String(20), nullable=False) # user, assistant, system
    content = db.Column(db.Text)
    meta_data = db.Column(db.Text) # JSON string for charts, tool calls, etc.
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

CollectionTask = SpiderTask

class KnowledgeImportTask(db.Model):
    """样本中心批量入库任务本地记录。"""
    __tablename__ = 'knowledge_import_task'

    id = db.Column(db.Integer, primary_key=True)
    task_no = db.Column(db.String(64), unique=True, nullable=False, index=True)
    sample_task_id = db.Column(db.String(64), nullable=True, index=True)
    kb_id = db.Column(db.String(64), nullable=False, index=True)
    kb_name = db.Column(db.String(200), nullable=True)
    callback_url = db.Column(db.String(500), nullable=True)
    status = db.Column(db.String(20), default='pending', index=True)  # pending/processing/success/failed
    status_detail = db.Column(db.Text, nullable=True)  # JSON
    progress = db.Column(db.Text, nullable=True)  # JSON: {total, processed, succeeded, failed}
    poll_count = db.Column(db.Integer, default=0)
    last_poll_at = db.Column(db.DateTime, nullable=True)
    next_poll_at = db.Column(db.DateTime, nullable=True, index=True)
    error_message = db.Column(db.Text, nullable=True)
    created_by = db.Column(db.Integer, db.ForeignKey('user.id'), nullable=True)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)