145 lines
4.9 KiB
Python
145 lines
4.9 KiB
Python
"""
|
|
数据库模型定义
|
|
定义工作流相关的所有数据表结构
|
|
"""
|
|
|
|
from datetime import datetime
|
|
from sqlalchemy import Column, String, Text, DateTime, ForeignKey, Integer, Boolean
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.orm import relationship
|
|
|
|
# 创建基类
|
|
Base = declarative_base()
|
|
|
|
|
|
class Workflow(Base):
|
|
"""工作流模型"""
|
|
|
|
__tablename__ = "workflows"
|
|
|
|
name = Column(String, primary_key=True)
|
|
base_name = Column(String, nullable=False)
|
|
version = Column(String, nullable=False)
|
|
workflow_json = Column(Text, nullable=False)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"name": self.name,
|
|
"base_name": self.base_name,
|
|
"version": self.version,
|
|
"workflow_json": self.workflow_json,
|
|
}
|
|
|
|
|
|
class WorkflowRun(Base):
|
|
"""工作流运行记录模型"""
|
|
|
|
__tablename__ = "workflow_run"
|
|
|
|
id = Column(String, primary_key=True)
|
|
workflow_name = Column(String, nullable=False)
|
|
prompt_id = Column(String)
|
|
client_id = Column(String)
|
|
status = Column(String, nullable=False, default="pending")
|
|
server_url = Column(String)
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
started_at = Column(DateTime)
|
|
completed_at = Column(DateTime)
|
|
error_message = Column(Text)
|
|
workflow_json = Column(Text, nullable=False)
|
|
api_spec = Column(Text, nullable=False)
|
|
request_data = Column(Text, nullable=False)
|
|
result = Column(Text)
|
|
|
|
# 关联关系
|
|
nodes = relationship(
|
|
"WorkflowRunNode", back_populates="workflow_run", cascade="all, delete-orphan"
|
|
)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"id": self.id,
|
|
"workflow_name": self.workflow_name,
|
|
"prompt_id": self.prompt_id,
|
|
"client_id": self.client_id,
|
|
"status": self.status,
|
|
"server_url": self.server_url,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
"started_at": self.started_at.isoformat() if self.started_at else None,
|
|
"completed_at": (
|
|
self.completed_at.isoformat() if self.completed_at else None
|
|
),
|
|
"error_message": self.error_message,
|
|
"workflow_json": self.workflow_json,
|
|
"api_spec": self.api_spec,
|
|
"request_data": self.request_data,
|
|
"result": self.result,
|
|
}
|
|
|
|
|
|
class WorkflowRunNode(Base):
|
|
"""工作流运行节点模型"""
|
|
|
|
__tablename__ = "workflow_run_nodes"
|
|
|
|
id = Column(String, primary_key=True)
|
|
workflow_run_id = Column(String, ForeignKey("workflow_run.id"), nullable=False)
|
|
node_id = Column(String, nullable=False)
|
|
node_type = Column(String, nullable=False)
|
|
status = Column(String, nullable=False, default="pending")
|
|
started_at = Column(DateTime)
|
|
completed_at = Column(DateTime)
|
|
output_data = Column(Text)
|
|
error_message = Column(Text)
|
|
|
|
# 关联关系
|
|
workflow_run = relationship("WorkflowRun", back_populates="nodes")
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"id": self.id,
|
|
"workflow_run_id": self.workflow_run_id,
|
|
"node_id": self.node_id,
|
|
"node_type": self.node_type,
|
|
"status": self.status,
|
|
"started_at": self.started_at.isoformat() if self.started_at else None,
|
|
"completed_at": (
|
|
self.completed_at.isoformat() if self.completed_at else None
|
|
),
|
|
"output_data": self.output_data,
|
|
"error_message": self.error_message,
|
|
}
|
|
|
|
|
|
class ComfyUIServer(Base):
|
|
"""ComfyUI 服务器模型"""
|
|
|
|
__tablename__ = "comfyui_servers"
|
|
|
|
name = Column(String, primary_key=True)
|
|
http_url = Column(String, nullable=False)
|
|
ws_url = Column(String, nullable=False)
|
|
status = Column(String, nullable=False, default="offline")
|
|
last_health_check = Column(DateTime)
|
|
current_tasks = Column(Integer, default=0)
|
|
max_concurrent_tasks = Column(Integer, default=1)
|
|
capabilities = Column(Text) # JSON 字符串
|
|
server_metadata = Column(Text) # JSON 字符串
|
|
created_at = Column(DateTime, default=datetime.utcnow)
|
|
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
|
|
|
|
def to_dict(self) -> dict:
|
|
return {
|
|
"name": self.name,
|
|
"http_url": self.http_url,
|
|
"ws_url": self.ws_url,
|
|
"status": self.status,
|
|
"last_health_check": self.last_health_check.isoformat() if self.last_health_check else None,
|
|
"current_tasks": self.current_tasks,
|
|
"max_concurrent_tasks": self.max_concurrent_tasks,
|
|
"capabilities": self.capabilities,
|
|
"server_metadata": self.server_metadata,
|
|
"created_at": self.created_at.isoformat() if self.created_at else None,
|
|
"updated_at": self.updated_at.isoformat() if self.updated_at else None,
|
|
}
|