105 lines
3.6 KiB
Python
105 lines
3.6 KiB
Python
"""
|
|
SQLAlchemy ORM 数据模型
|
|
对应数据库表结构 (基于 DATABASE_DESIGN.md)
|
|
"""
|
|
|
|
import uuid
|
|
from datetime import datetime, timezone
|
|
from enum import Enum as PyEnum
|
|
from sqlalchemy import Column, String, Text, DateTime, Integer, Enum, Index, ForeignKey
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
from sqlalchemy.orm import relationship
|
|
|
|
from .database import Base
|
|
|
|
|
|
class TaskStatus(str, PyEnum):
|
|
"""任务状态枚举"""
|
|
GENERATING = "generating" # 生成中 - TaskProcess 生成阶段
|
|
EXECUTING = "executing" # 执行中 - 任务执行阶段
|
|
STOPPED = "stopped" # 已停止 - 用户手动停止执行
|
|
COMPLETED = "completed" # 已完成 - 任务正常完成
|
|
|
|
|
|
def utc_now():
|
|
"""获取当前 UTC 时间"""
|
|
return datetime.now(timezone.utc)
|
|
|
|
|
|
class MultiAgentTask(Base):
|
|
"""多智能体任务记录模型"""
|
|
__tablename__ = "multi_agent_tasks"
|
|
|
|
task_id = Column(String(64), primary_key=True)
|
|
user_id = Column(String(64), nullable=False, index=True)
|
|
query = Column(Text, nullable=False)
|
|
agents_info = Column(JSONB, nullable=False)
|
|
task_outline = Column(JSONB)
|
|
assigned_agents = Column(JSONB)
|
|
agent_scores = Column(JSONB)
|
|
result = Column(JSONB)
|
|
status = Column(
|
|
Enum(TaskStatus, name="task_status_enum", create_type=False),
|
|
default=TaskStatus.GENERATING,
|
|
nullable=False
|
|
)
|
|
execution_count = Column(Integer, default=0, nullable=False)
|
|
generation_id = Column(String(64))
|
|
execution_id = Column(String(64))
|
|
rehearsal_log = Column(JSONB)
|
|
branches = Column(JSONB) # 任务大纲探索分支数据
|
|
created_at = Column(DateTime(timezone=True), default=utc_now)
|
|
updated_at = Column(DateTime(timezone=True), default=utc_now, onupdate=utc_now)
|
|
|
|
__table_args__ = (
|
|
Index("idx_multi_agent_tasks_status", "status"),
|
|
Index("idx_multi_agent_tasks_generation_id", "generation_id"),
|
|
Index("idx_multi_agent_tasks_execution_id", "execution_id"),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
"""转换为字典"""
|
|
return {
|
|
"task_id": self.task_id,
|
|
"user_id": self.user_id,
|
|
"query": self.query,
|
|
"agents_info": self.agents_info,
|
|
"task_outline": self.task_outline,
|
|
"assigned_agents": self.assigned_agents,
|
|
"agent_scores": self.agent_scores,
|
|
"result": self.result,
|
|
"status": self.status.value if self.status else None,
|
|
"execution_count": self.execution_count,
|
|
"generation_id": self.generation_id,
|
|
"execution_id": self.execution_id,
|
|
"rehearsal_log": self.rehearsal_log,
|
|
"branches": self.branches,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
"updated_at": self.updated_at.isoformat() if self.updated_at else None,
|
|
}
|
|
|
|
|
|
class UserAgent(Base):
|
|
"""用户保存的智能体配置模型 (可选表)"""
|
|
__tablename__ = "user_agents"
|
|
|
|
id = Column(String(64), primary_key=True)
|
|
user_id = Column(String(64), nullable=False, index=True)
|
|
agent_name = Column(String(100), nullable=False)
|
|
agent_config = Column(JSONB, nullable=False)
|
|
created_at = Column(DateTime(timezone=True), default=utc_now)
|
|
|
|
__table_args__ = (
|
|
Index("idx_user_agents_user_created", "user_id", "created_at"),
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
"""转换为字典"""
|
|
return {
|
|
"id": self.id,
|
|
"user_id": self.user_id,
|
|
"agent_name": self.agent_name,
|
|
"agent_config": self.agent_config,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
}
|