ai-chat-system/models_v2.py

"""
AI对话系统 v2.0.0 - 数据库模型重构
支持：大模型池、Agent管理、渠道独立绑定、思考功能开关
"""
from sqlalchemy import create_engine, Column, Integer, String, Text, Boolean, DateTime, ForeignKey, JSON, Float
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from datetime import datetime
import os

DATABASE_URL = os.environ.get('DATABASE_URL', 'sqlite:///./ai_chat_v2.db')

engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False})
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()


# ==================== 大模型池 ====================

class LLMProvider(Base):
    """大模型提供商/池配置"""
    __tablename__ = 'llm_providers'
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, index=True)  # 名称标识，如 "DeepSeek", "本地LLM"
    api_base = Column(String(500))  # API地址
    api_key = Column(String(500))  # API密钥
    models = Column(JSON, default=list)  # 可用模型列表 [{"id": "xxx", "name": "xxx"}]
    default_model = Column(String(100))  # 默认模型
    
    # 思考功能支持
    supports_thinking = Column(Boolean, default=False)  # 是否原生支持思考
    thinking_model = Column(String(100), nullable=True)  # 思考模式模型名（如有单独模型）
    
    # 配额和限制
    max_tokens = Column(Integer, default=4096)
    temperature = Column(Float, default=0.7)
    
    # 状态
    is_active = Column(Boolean, default=True)
    priority = Column(Integer, default=0)  # 优先级，用于自动选择
    description = Column(Text, nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关联
    agents = relationship("Agent", back_populates="llm_provider")


# ==================== Agent ====================

class Agent(Base):
    """智能体配置"""
    __tablename__ = 'agents'
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100), unique=True, index=True)  # Agent名称
    display_name = Column(String(100))  # 显示名称
    
    # 工具配置
    tools = Column(JSON, default=list)  # 可用工具列表 ["search", "calculator", ...]
    
    # 大模型配置
    llm_provider_id = Column(Integer, ForeignKey('llm_providers.id'))
    model_override = Column(String(100), nullable=True)  # 覆盖Provider默认模型
    
    # 系统设定
    system_prompt = Column(Text, default="你是一个有用的AI助手。")  # 系统提示词
    
    # 思考功能
    enable_thinking = Column(Boolean, default=True)  # 是否启用思考
    thinking_prompt = Column(Text, nullable=True)  # 思考提示词模板
    thinking_prefix = Column(String(50), default="<think>")  # 思考内容前缀标识
    thinking_suffix = Column(String(50), default="</think>")  # 思考内容后缀标识
    
    # 其他配置
    max_history = Column(Integer, default=20)  # 最大历史消息数
    temperature_override = Column(Float, nullable=True)  # 覆盖温度
    
    # 状态
    is_active = Column(Boolean, default=True)
    is_default = Column(Boolean, default=False)  # 是否为默认Agent
    description = Column(Text, nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关联
    llm_provider = relationship("LLMProvider", back_populates="agents")
    channel_mappings = relationship("ChannelAgentMapping", back_populates="agent")


# ==================== 渠道 ====================

class Channel(Base):
    """渠道配置 - 网页端、Matrix端等"""
    __tablename__ = 'channels'
    
    id = Column(Integer, primary_key=True, index=True)
    channel_type = Column(String(20), index=True)  # web, matrix, telegram, etc.
    name = Column(String(100))  # 渠道名称，如 "网页端主入口", "Matrix AI Bot"
    
    # 渠道配置（JSON）
    config = Column(JSON, default=dict)
    # Matrix示例: {"homeserver": "xxx", "username": "xxx", "password": "xxx"}
    # Web示例: {"require_login": false, "session_timeout": 3600}
    
    # 状态
    is_active = Column(Boolean, default=True)
    is_primary = Column(Boolean, default=False)  # 是否为主要渠道
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    
    # 关联
    agent_mappings = relationship("ChannelAgentMapping", back_populates="channel", order_by="ChannelAgentMapping.priority")
    conversations = relationship("Conversation", back_populates="channel")


# ==================== 渠道-Agent映射 ====================

class ChannelAgentMapping(Base):
    """渠道与Agent的映射关系"""
    __tablename__ = 'channel_agent_mappings'
    
    id = Column(Integer, primary_key=True, index=True)
    channel_id = Column(Integer, ForeignKey('channels.id'))
    agent_id = Column(Integer, ForeignKey('agents.id'))
    
    # 映射配置
    priority = Column(Integer, default=0)  # 优先级（数字越小优先级越高）
    mode = Column(String(20), default='single')  # single(单Agent), round_robin(轮询), fallback(备用)
    weight = Column(Integer, default=1)  # 轮询权重
    
    # 条件配置（可选）
    conditions = Column(JSON, nullable=True)  # 条件触发，如 {"user_type": "vip"}
    
    is_active = Column(Boolean, default=True)
    created_at = Column(DateTime, default=datetime.utcnow)
    
    # 关联
    channel = relationship("Channel", back_populates="agent_mappings")
    agent = relationship("Agent", back_populates="channel_mappings")


# ==================== 用户和会话（保留原有结构，增加渠道关联） ====================

class User(Base):
    """用户表"""
    __tablename__ = 'users'
    
    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(String(100), unique=True, index=True)
    display_name = Column(String(100))
    user_type = Column(String(20), default='web')  # web, matrix, telegram
    matrix_user_id = Column(String(200), nullable=True)
    
    # 用户级别（用于Agent选择条件）
    user_level = Column(String(20), default='normal')  # normal, vip, admin
    
    created_at = Column(DateTime, default=datetime.utcnow)
    last_active_at = Column(DateTime, default=datetime.utcnow)
    is_active = Column(Boolean, default=True)
    
    conversations = relationship("Conversation", back_populates="user")


class Conversation(Base):
    """对话会话"""
    __tablename__ = 'conversations'
    
    id = Column(Integer, primary_key=True, index=True)
    conversation_id = Column(String(100), unique=True, index=True)
    user_id = Column(Integer, ForeignKey('users.id'))
    channel_id = Column(Integer, ForeignKey('channels.id'), nullable=True)  # 新增渠道关联
    
    title = Column(String(200), nullable=True)
    
    # 当前使用的Agent（会话可切换Agent）
    current_agent_id = Column(Integer, ForeignKey('agents.id'), nullable=True)
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    is_active = Column(Boolean, default=True)
    extra_data = Column(JSON, nullable=True)
    
    user = relationship("User", back_populates="conversations")
    channel = relationship("Channel", back_populates="conversations")
    messages = relationship("Message", back_populates="conversation")


class Message(Base):
    """消息表"""
    __tablename__ = 'messages'
    
    id = Column(Integer, primary_key=True, index=True)
    conversation_id = Column(Integer, ForeignKey('conversations.id'))
    role = Column(String(20))  # user, assistant, system, thinking(思考过程)
    content = Column(Text)
    source = Column(String(20))  # web, matrix
    
    # 思考内容（如果启用了思考功能）
    thinking_content = Column(Text, nullable=True)  # 思考过程
    
    # Agent信息
    agent_id = Column(Integer, ForeignKey('agents.id'), nullable=True)  # 生成此消息的Agent
    model_used = Column(String(100), nullable=True)  # 使用的模型
    
    created_at = Column(DateTime, default=datetime.utcnow)
    extra_data = Column(JSON, nullable=True)
    
    conversation = relationship("Conversation", back_populates="messages")


# ==================== Matrix房间映射（保留） ====================

class MatrixRoomMapping(Base):
    """Matrix房间映射"""
    __tablename__ = 'matrix_room_mapping'
    
    id = Column(Integer, primary_key=True, index=True)
    room_id = Column(String(200), unique=True, index=True)
    user_id = Column(Integer, ForeignKey('users.id'))
    conversation_id = Column(Integer, ForeignKey('conversations.id'))
    channel_id = Column(Integer, ForeignKey('channels.id'), nullable=True)
    created_at = Column(DateTime, default=datetime.utcnow)


# ==================== 搜索工具配置 ====================

class ToolConfig(Base):
    """工具配置（通用，支持搜索、计算器、代码执行等）"""
    __tablename__ = 'tool_configs'
    
    id = Column(Integer, primary_key=True, index=True)
    name = Column(String(100))  # 工具名称，如 "Tavily Search"、"Calculator"
    tool_type = Column(String(50), index=True)  # 工具类型：search, calculator, code_runner, image_gen, etc.
    provider = Column(String(50), nullable=True)  # 提供商（可选）：tavily, google, wolfram, etc.
    
    # API配置（JSON，不同工具可能有不同配置）
    config = Column(JSON, default=dict)
    # search示例: {"api_key": "xxx", "max_results": 5, "search_depth": "basic"}
    # calculator示例: {"api_base": "xxx"}
    
    # 状态
    is_active = Column(Boolean, default=True)
    is_default = Column(Boolean, default=False)  # 是否为该类型的默认工具
    
    # 统计
    total_calls = Column(Integer, default=0)  # 总调用次数
    success_calls = Column(Integer, default=0)  # 成功次数
    failed_calls = Column(Integer, default=0)  # 失败次数
    
    created_at = Column(DateTime, default=datetime.utcnow)
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


class ToolUsageLog(Base):
    """工具使用日志"""
    __tablename__ = 'tool_usage_logs'
    
    id = Column(Integer, primary_key=True, index=True)
    tool_id = Column(Integer, ForeignKey('tool_configs.id'))
    tool_type = Column(String(50), index=True)
    
    # 调用信息
    query = Column(Text)  # 调用参数/查询内容
    success = Column(Boolean, default=True)
    error_message = Column(Text, nullable=True)
    result_summary = Column(Text, nullable=True)  # 结果摘要
    
    # 关联信息
    conversation_id = Column(String(100), nullable=True)
    agent_id = Column(Integer, nullable=True)
    user_id = Column(String(100), nullable=True)
    
    # 性能
    duration_ms = Column(Integer, nullable=True)  # 调用耗时（毫秒）
    
    called_at = Column(DateTime, default=datetime.utcnow)


# ==================== 系统配置（保留） ====================

class SystemConfig(Base):
    """系统配置表"""
    __tablename__ = 'system_config'
    
    id = Column(Integer, primary_key=True, index=True)
    key = Column(String(100), unique=True, index=True)
    value = Column(Text)
    description = Column(String(500))
    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)


def init_db():
    """初始化数据库"""
    Base.metadata.create_all(bind=engine)


def get_db():
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


# ==================== 初始化默认数据 ====================

def init_default_data():
    """初始化默认数据：大模型池、Agent、渠道"""
    db = SessionLocal()
    try:
        # 检查是否已有数据
        if db.query(LLMProvider).count() > 0:
            return
        
        # 1. 创建默认大模型池
        default_llm = LLMProvider(
            name="本地LLM Proxy",
            api_base="http://192.168.2.17:19007/v1",
            api_key="xxxx",
            models=[
                {"id": "auto", "name": "auto (自动选择)"},
                {"id": "qwen3.5-4b", "name": "qwen3.5-4b"},
                {"id": "dsv32", "name": "dsv32"},
                {"id": "glm-4", "name": "glm-4"}
            ],
            default_model="auto",
            supports_thinking=False,
            is_active=True,
            priority=0,
            description="本地LLM代理服务"
        )
        db.add(default_llm)
        
        # 2. 创建默认Agent
        default_agent = Agent(
            name="default",
            display_name="默认助手",
            llm_provider_id=1,
            system_prompt="你是一个有用的AI助手。",
            enable_thinking=True,
            thinking_prompt="请先仔细思考这个问题，然后给出回答。",
            is_active=True,
            is_default=True,
            description="默认AI助手"
        )
        db.add(default_agent)
        
        # 3. 创建渠道
        web_channel = Channel(
            channel_type="web",
            name="网页端",
            config={"require_login": False},
            is_active=True,
            is_primary=True
        )
        db.add(web_channel)
        
        matrix_channel = Channel(
            channel_type="matrix",
            name="Matrix AI Bot",
            config={
                "homeserver": "http://matrix.tphai.com",
                "username": "@tester:matrix.tphai.com",
                "password": "tester12345@!"
            },
            is_active=True
        )
        db.add(matrix_channel)
        
        db.commit()
        
        # 4. 创建渠道-Agent映射
        web_mapping = ChannelAgentMapping(
            channel_id=1,
            agent_id=1,
            priority=0,
            mode='single'
        )
        db.add(web_mapping)
        
        matrix_mapping = ChannelAgentMapping(
            channel_id=2,
            agent_id=1,
            priority=0,
            mode='single'
        )
        db.add(matrix_mapping)
        
        # 5. 创建默认搜索工具配置
        search_config = SearchToolConfig(
            name="Tavily Search",
            provider="tavily",
            api_key="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",
            max_results=5,
            search_depth="basic",
            is_active=True,
            is_default=True
        )
        db.add(search_config)
        
        db.commit()
        print("默认数据初始化完成")
        
    finally:
        db.close()