253 lines
8.6 KiB
Python
253 lines
8.6 KiB
Python
from sqlalchemy import (
|
||
Column, Integer, String, Text, DateTime, Boolean, Float, ForeignKey
|
||
)
|
||
from sqlalchemy.orm import relationship, declarative_base
|
||
|
||
from db.session import Base
|
||
from models.base import CoreModel
|
||
from models.user import DjangoUser # 确保导入 DjangoUser
|
||
|
||
# 状态选择类（示例）
|
||
class CommonStatus:
|
||
DISABLED = 0
|
||
ENABLED = 1
|
||
|
||
@staticmethod
|
||
def choices():
|
||
return [(0, '禁用'), (1, '启用')]
|
||
|
||
|
||
# 平台选择类（示例）
|
||
class PlatformChoices:
|
||
OPENAI = 'openai'
|
||
ALIMNS = 'alimns'
|
||
|
||
@staticmethod
|
||
def choices():
|
||
return [('openai', 'OpenAI'), ('alimns', '阿里云MNS')]
|
||
|
||
|
||
# 消息类型选择类（示例）
|
||
class MessageType:
|
||
SYSTEM = "system" # 系统指令
|
||
USER = "user" # 用户消息
|
||
ASSISTANT = "assistant" # 助手回复
|
||
FUNCTION = "function" # 函数返回结果
|
||
|
||
@staticmethod
|
||
def choices():
|
||
"""返回可用的消息角色选项"""
|
||
return [
|
||
(MessageType.SYSTEM, "系统"),
|
||
(MessageType.USER, "用户"),
|
||
(MessageType.ASSISTANT, "助手"),
|
||
(MessageType.FUNCTION, "函数")
|
||
]
|
||
|
||
|
||
class MessageContentType:
|
||
"""消息内容类型"""
|
||
TEXT = "text"
|
||
FUNCTION_CALL = "function_call"
|
||
|
||
@staticmethod
|
||
def choices():
|
||
"""返回可用的内容类型选项"""
|
||
return [
|
||
(MessageContentType.TEXT, "文本"),
|
||
(MessageContentType.FUNCTION_CALL, "函数调用")
|
||
]
|
||
|
||
# AI API 密钥表
|
||
class AIApiKey(CoreModel):
|
||
__tablename__ = 'ai_api_key'
|
||
|
||
name = Column(String(255), nullable=False)
|
||
platform = Column(String(100), nullable=False)
|
||
api_key = Column(String(255), nullable=False)
|
||
url = Column(String(255), nullable=True)
|
||
status = Column(Integer, default=CommonStatus.DISABLED)
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
|
||
# AI 模型表
|
||
class AIModel(CoreModel):
|
||
__tablename__ = 'ai_model'
|
||
|
||
name = Column(String(64), nullable=False)
|
||
sort = Column(Integer, default=0)
|
||
status = Column(Integer, default=CommonStatus.DISABLED)
|
||
key_id = Column(Integer, ForeignKey('ai_api_key.id'), nullable=False)
|
||
model_type = Column(String(32), nullable=True)
|
||
platform = Column(String(32), nullable=False)
|
||
model = Column(String(64), nullable=False)
|
||
temperature = Column(Float, nullable=True)
|
||
max_tokens = Column(Integer, nullable=True)
|
||
max_contexts = Column(Integer, nullable=True)
|
||
|
||
key = relationship('AIApiKey', backref='models')
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
|
||
# AI 工具表
|
||
class Tool(CoreModel):
|
||
__tablename__ = 'ai_tool'
|
||
|
||
name = Column(String(128), nullable=False)
|
||
description = Column(String(256), nullable=True)
|
||
status = Column(Integer, default=0)
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
|
||
# AI 知识库表
|
||
class Knowledge(CoreModel):
|
||
__tablename__ = 'ai_knowledge'
|
||
|
||
name = Column(String(255), nullable=False)
|
||
description = Column(Text, nullable=True)
|
||
embedding_model_id = Column(Integer, ForeignKey('ai_model.id'), nullable=False)
|
||
embedding_model = Column(String(32), nullable=False)
|
||
top_k = Column(Integer, default=0)
|
||
similarity_threshold = Column(Float, nullable=False)
|
||
status = Column(Integer, default=CommonStatus.DISABLED)
|
||
|
||
embedding_model_rel = relationship('AIModel', backref='knowledges')
|
||
documents = relationship('KnowledgeDocument', backref='knowledge', cascade='all, delete-orphan')
|
||
segments = relationship('KnowledgeSegment', backref='knowledge', cascade='all, delete-orphan')
|
||
roles = relationship('ChatRole', secondary='ai_chat_role_knowledge', backref='knowledges')
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
|
||
# AI 知识库文档表
|
||
class KnowledgeDocument(CoreModel):
|
||
__tablename__ = 'ai_knowledge_document'
|
||
|
||
knowledge_id = Column(Integer, ForeignKey('ai_knowledge.id'), nullable=False)
|
||
name = Column(String(255), nullable=False)
|
||
url = Column(String(1024), nullable=False)
|
||
content = Column(Text, nullable=False)
|
||
content_length = Column(Integer, nullable=False)
|
||
tokens = Column(Integer, nullable=False)
|
||
segment_max_tokens = Column(Integer, nullable=False)
|
||
retrieval_count = Column(Integer, default=0)
|
||
status = Column(Integer, default=CommonStatus.DISABLED)
|
||
|
||
segments = relationship('KnowledgeSegment', backref='document', cascade='all, delete-orphan')
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
|
||
# AI 知识库分段表
|
||
class KnowledgeSegment(CoreModel):
|
||
__tablename__ = 'ai_knowledge_segment'
|
||
|
||
knowledge_id = Column(Integer, ForeignKey('ai_knowledge.id'), nullable=False)
|
||
document_id = Column(Integer, ForeignKey('ai_knowledge_document.id'), nullable=False)
|
||
content = Column(Text, nullable=False)
|
||
content_length = Column(Integer, nullable=False)
|
||
tokens = Column(Integer, nullable=False)
|
||
vector_id = Column(String(100), nullable=True)
|
||
retrieval_count = Column(Integer, default=0)
|
||
status = Column(Integer, default=CommonStatus.DISABLED)
|
||
|
||
def __str__(self):
|
||
return f"Segment {self.id}"
|
||
|
||
|
||
# AI 聊天角色表
|
||
class ChatRole(CoreModel):
|
||
__tablename__ = 'ai_chat_role'
|
||
|
||
name = Column(String(128), nullable=False)
|
||
avatar = Column(String(256), nullable=False)
|
||
description = Column(String(256), nullable=True)
|
||
status = Column(Integer, default=CommonStatus.DISABLED)
|
||
sort = Column(Integer, default=0)
|
||
public_status = Column(Boolean, default=False)
|
||
category = Column(String(32), nullable=True)
|
||
model_id = Column(Integer, ForeignKey('ai_model.id'), nullable=False)
|
||
system_message = Column(String(1024), nullable=True)
|
||
user_id = Column(
|
||
Integer,
|
||
ForeignKey('system_users.id'), # 假设DjangoUser表名是system_users
|
||
nullable=True # 允许为空（如匿名角色）
|
||
)
|
||
user = relationship(DjangoUser, backref='chat_roles') # 正确：DjangoUser 已定义并导入
|
||
|
||
model = relationship('AIModel', backref='chat_roles')
|
||
tools = relationship('Tool', secondary='ai_chat_role_tool', backref='roles')
|
||
# conversations = relationship('ChatConversation', backref='role', cascade='all, delete-orphan')
|
||
# messages = relationship('ChatMessage', backref='role', cascade='all, delete-orphan')
|
||
|
||
def __str__(self):
|
||
return self.name
|
||
|
||
# AI 聊天对话表
|
||
class ChatConversation(CoreModel):
|
||
__tablename__ = 'ai_chat_conversation'
|
||
|
||
title = Column(String(256), nullable=False)
|
||
pinned = Column(Boolean, default=False)
|
||
pinned_time = Column(DateTime, nullable=True)
|
||
user_id = Column(Integer, ForeignKey('system_users.id'), nullable=True)
|
||
role_id = Column(Integer, ForeignKey('ai_chat_role.id'), nullable=True)
|
||
model_id = Column(Integer, ForeignKey('ai_model.id'), nullable=False)
|
||
model = Column(String(32), nullable=False)
|
||
system_message = Column(String(1024), nullable=True)
|
||
temperature = Column(Float, nullable=False)
|
||
max_tokens = Column(Integer, nullable=False)
|
||
max_contexts = Column(Integer, nullable=False)
|
||
user = relationship(DjangoUser, backref='conversations') # 正确：DjangoUser 已定义并导入
|
||
|
||
model_rel = relationship('AIModel', backref='conversations')
|
||
messages = relationship('ChatMessage', backref='conversation', cascade='all, delete-orphan')
|
||
|
||
def __str__(self):
|
||
return self.title
|
||
|
||
|
||
# AI 聊天消息表
|
||
class ChatMessage(CoreModel):
|
||
__tablename__ = 'ai_chat_message'
|
||
|
||
conversation_id = Column(Integer, ForeignKey('ai_chat_conversation.id'), nullable=False)
|
||
user_id = Column(Integer, ForeignKey('system_users.id'), nullable=True)
|
||
role_id = Column(Integer, ForeignKey('ai_chat_role.id'), nullable=True)
|
||
model = Column(String(32), nullable=False)
|
||
model_id = Column(Integer, ForeignKey('ai_model.id'), nullable=False)
|
||
type = Column(String(16), nullable=False)
|
||
reply_id = Column(Integer, nullable=True)
|
||
content = Column(String(2048), nullable=False)
|
||
use_context = Column(Boolean, default=False)
|
||
segment_ids = Column(String(2048), nullable=True)
|
||
|
||
user = relationship(DjangoUser, backref='messages') # 正确：DjangoUser 已定义并导入
|
||
model_rel = relationship('AIModel', backref='messages')
|
||
|
||
def __str__(self):
|
||
return self.content[:30]
|
||
|
||
|
||
# 聊天角色与知识库的关联表
|
||
class ChatRoleKnowledge(Base):
|
||
__tablename__ = 'ai_chat_role_knowledge'
|
||
|
||
chat_role_id = Column(Integer, ForeignKey('ai_chat_role.id'), primary_key=True)
|
||
knowledge_id = Column(Integer, ForeignKey('ai_knowledge.id'), primary_key=True)
|
||
|
||
|
||
# 聊天角色与工具的关联表
|
||
class ChatRoleTool(Base):
|
||
__tablename__ = 'ai_chat_role_tool'
|
||
|
||
chat_role_id = Column(Integer, ForeignKey('ai_chat_role.id'), primary_key=True)
|
||
tool_id = Column(Integer, ForeignKey('ai_tool.id'), primary_key=True) |