''' Author: caoziyuan ziyuan.cao@zhuying.com Date: 2025-12-22 17:42:10 LastEditors: caoziyuan ziyuan.cao@zhuying.com LastEditTime: 2025-12-22 18:03:28 FilePath: \naming-backend\app\api\v1\module_application\ai\model.py Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE ''' # -*- coding: utf-8 -*- from sqlalchemy import JSON, String, Integer, Text, ForeignKey from sqlalchemy.orm import Mapped, mapped_column, relationship from app.core.base_model import ModelMixin, UserMixin class McpModel(ModelMixin, UserMixin): """ MCP 服务器表 MCP类型: - 0: stdio (标准输入输出) - 1: sse (Server-Sent Events) """ __tablename__: str = 'app_ai_mcp' __table_args__: dict[str, str] = ({'comment': 'MCP 服务器表'}) __loader_options__: list[str] = ["created_by", "updated_by"] name: Mapped[str] = mapped_column(String(50), comment='MCP 名称') type: Mapped[int] = mapped_column(Integer, default=0, comment='MCP 类型(0:stdio 1:sse)') url: Mapped[str | None] = mapped_column(String(255), default=None, comment='远程 SSE 地址') command: Mapped[str | None] = mapped_column(String(255), default=None, comment='MCP 命令') args: Mapped[str | None] = mapped_column(String(255), default=None, comment='MCP 命令参数') env: Mapped[dict[str, str] | None] = mapped_column(JSON(), default=None, comment='MCP 环境变量') class AIProviderModel(ModelMixin, UserMixin): """ AI供应商配置表 存储AI服务供应商的接口地址和API Key """ __tablename__: str = 'app_ai_provider' __table_args__: dict[str, str] = ({'comment': 'AI供应商配置表'}) __loader_options__: list[str] = ["created_by", "updated_by"] name: Mapped[str] = mapped_column(String(50), nullable=False, comment='供应商名称') provider_type: Mapped[str] = mapped_column(String(50), nullable=False, comment='供应商类型(openai/deepseek/anthropic/gemini/qwen等)') base_url: Mapped[str] = mapped_column(String(255), nullable=False, comment='接口地址BaseURL') api_key: Mapped[str] = mapped_column(String(255), nullable=False, comment='API Key') is_default: Mapped[int] = mapped_column(Integer, default=0, comment='是否默认供应商(0:否 1:是)') class EmbeddingConfigModel(ModelMixin, UserMixin): """ 知识库向量化配置表 支持本地或远程向量化服务 """ __tablename__: str = 'app_ai_embedding_config' __table_args__: dict[str, str] = ({'comment': '知识库向量化配置表'}) __loader_options__: list[str] = ["created_by", "updated_by"] name: Mapped[str] = mapped_column(String(50), nullable=False, comment='配置名称') embedding_type: Mapped[int] = mapped_column(Integer, default=0, comment='向量化类型(0:本地 1:远程)') model_name: Mapped[str] = mapped_column(String(100), nullable=False, comment='Embedding模型名称') base_url: Mapped[str | None] = mapped_column(String(255), default=None, comment='远程接口地址(远程模式必填)') api_key: Mapped[str | None] = mapped_column(String(255), default=None, comment='远程API Key(远程模式必填)') is_default: Mapped[int] = mapped_column(Integer, default=0, comment='是否默认配置(0:否 1:是)') class KnowledgeBaseModel(ModelMixin, UserMixin): """ 知识库表 存储知识库信息，关联向量化配置 """ __tablename__: str = 'app_ai_knowledge_base' __table_args__: dict[str, str] = ({'comment': '知识库表'}) __loader_options__: list[str] = ["created_by", "updated_by", "embedding_config"] name: Mapped[str] = mapped_column(String(100), nullable=False, comment='知识库名称') embedding_config_id: Mapped[int | None] = mapped_column( Integer, ForeignKey('app_ai_embedding_config.id', ondelete="SET NULL", onupdate="CASCADE"), default=None, nullable=True, index=True, comment='向量化配置ID' ) collection_name: Mapped[str] = mapped_column(String(100), nullable=False, comment='ChromaDB集合名称') document_count: Mapped[int] = mapped_column(Integer, default=0, comment='文档数量') vector_count: Mapped[int] = mapped_column(Integer, default=0, comment='向量数量') kb_status: Mapped[int] = mapped_column(Integer, default=0, comment='知识库状态(0:待处理 1:处理中 2:已完成 3:处理失败)') error_message: Mapped[str | None] = mapped_column(Text, default=None, comment='错误信息') file_paths: Mapped[list[str] | None] = mapped_column(JSON, default=None, comment='文件路径列表') # 关联关系 embedding_config: Mapped["EmbeddingConfigModel | None"] = relationship( "EmbeddingConfigModel", lazy="selectin", foreign_keys=[embedding_config_id], uselist=False ) class AIModelConfigModel(ModelMixin, UserMixin): """ AI模型配置表 存储不同类型AI模型的配置信息 模型类型: - enterprise_naming(企业起名), enterprise_renaming(企业改名), enterprise_scoring(企业测名), enterprise_scoring_trial(企业测名试用) - personal_naming(个人起名), personal_renaming(个人改名), personal_scoring(个人测名), personal_scoring_trial(个人测名试用) """ __tablename__: str = 'app_ai_model_config' __table_args__: dict[str, str] = ({'comment': 'AI模型配置表'}) __loader_options__: list[str] = ["created_by", "updated_by", "provider", "knowledge_bases"] model_type: Mapped[str] = mapped_column(String(50), nullable=False, unique=True, comment='模型类型(enterprise_naming/enterprise_renaming/enterprise_scoring/enterprise_scoring_trial/personal_naming/personal_renaming/personal_scoring/personal_scoring_trial)') model_name: Mapped[str | None] = mapped_column(String(100), default=None, comment='使用的模型名称') provider_id: Mapped[int | None] = mapped_column( Integer, ForeignKey('app_ai_provider.id', ondelete="SET NULL", onupdate="CASCADE"), default=None, nullable=True, index=True, comment='AI供应商ID' ) system_prompt: Mapped[str | None] = mapped_column(Text, default=None, comment='系统提示词') temperature: Mapped[float] = mapped_column(default=1.0, comment='模型温度(0-2)') knowledge_base_ids: Mapped[list[int] | None] = mapped_column(JSON, default=None, comment='关联的知识库ID列表') # 关联关系 provider: Mapped["AIProviderModel | None"] = relationship( "AIProviderModel", lazy="selectin", foreign_keys=[provider_id], uselist=False ) class AIModelTrainingMessageModel(ModelMixin, UserMixin): """ AI模型训练对话记录表 存储训练对话的历史记录 """ __tablename__: str = 'app_ai_model_training_message' __table_args__: dict[str, str] = ({'comment': 'AI模型训练对话记录表'}) __loader_options__: list[str] = ["created_by", "updated_by"] model_config_id: Mapped[int] = mapped_column( Integer, ForeignKey('app_ai_model_config.id', ondelete="CASCADE", onupdate="CASCADE"), nullable=False, index=True, comment='模型配置ID' ) role: Mapped[str] = mapped_column(String(20), nullable=False, comment='角色(user/assistant)') content: Mapped[str] = mapped_column(Text, nullable=False, comment='消息内容')