130 lines
No EOL
6.9 KiB
Python
130 lines
No EOL
6.9 KiB
Python
"""
|
|
Chat model classes for the chat system.
|
|
"""
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Dict, Any, Optional
|
|
from datetime import datetime
|
|
import uuid
|
|
|
|
from modules.shared.attributeUtils import Label, BaseModelWithUI
|
|
|
|
|
|
# WORKFLOW MODELS
|
|
|
|
class ChatContent(BaseModelWithUI):
|
|
"""Data model for chat content"""
|
|
sequenceNr: int = Field(description="Sequence number of the content")
|
|
name: str = Field(description="Name of the content")
|
|
data: str = Field(description="The actual content data")
|
|
mimeType: str = Field(description="MIME type of the content")
|
|
metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
|
|
|
|
class ChatDocument(BaseModelWithUI):
|
|
"""Data model for a chat document"""
|
|
id: str = Field(description="Primary key")
|
|
fileId: int = Field(description="Foreign key to file")
|
|
filename: str = Field(description="Name of the file")
|
|
fileSize: int = Field(description="Size of the file")
|
|
mimeType: str = Field(description="MIME type of the file")
|
|
contents: List[ChatContent] = Field(default_factory=list, description="List of chat contents")
|
|
|
|
class ChatStat(BaseModelWithUI):
|
|
"""Data model for chat statistics"""
|
|
id: str = Field(description="Primary key")
|
|
processingTime: Optional[float] = Field(None, description="Processing time in seconds")
|
|
tokenCount: Optional[int] = Field(None, description="Number of tokens processed")
|
|
bytesSent: Optional[int] = Field(None, description="Number of bytes sent")
|
|
bytesReceived: Optional[int] = Field(None, description="Number of bytes received")
|
|
|
|
class ChatLog(BaseModelWithUI):
|
|
"""Data model for a chat log"""
|
|
id: str = Field(description="Primary key")
|
|
workflowId: str = Field(description="Foreign key to workflow")
|
|
message: str = Field(description="Log message")
|
|
type: str = Field(description="Type of log entry")
|
|
timestamp: str = Field(description="Timestamp of the log entry")
|
|
agentName: str = Field(description="Name of the agent")
|
|
status: str = Field(description="Status of the log entry")
|
|
progress: Optional[int] = Field(None, description="Progress percentage")
|
|
|
|
class ChatMessage(BaseModelWithUI):
|
|
"""Data model for a chat message"""
|
|
id: str = Field(description="Primary key")
|
|
workflowId: str = Field(description="Foreign key to workflow")
|
|
parentMessageId: Optional[str] = Field(None, description="Parent message ID for threading")
|
|
agentName: Optional[str] = Field(None, description="Name of the agent")
|
|
documents: List[ChatDocument] = Field(default_factory=list, description="Associated documents")
|
|
message: Optional[str] = Field(None, description="Message content")
|
|
role: str = Field(description="Role of the message sender")
|
|
status: str = Field(description="Status of the message")
|
|
sequenceNr: int = Field(description="Sequence number of the message")
|
|
startedAt: str = Field(description="When the message processing started")
|
|
finishedAt: Optional[str] = Field(None, description="When the message processing finished")
|
|
stats: Optional[ChatStat] = Field(None, description="Statistics for this message")
|
|
|
|
class ChatWorkflow(BaseModelWithUI):
|
|
"""Data model for a chat workflow"""
|
|
id: str = Field(description="Primary key")
|
|
mandateId: str = Field(description="ID of the mandate this workflow belongs to")
|
|
status: str = Field(description="Current status of the workflow")
|
|
name: Optional[str] = Field(None, description="Name of the workflow")
|
|
currentRound: int = Field(description="Current round number")
|
|
lastActivity: str = Field(description="Timestamp of last activity")
|
|
startedAt: str = Field(description="When the workflow started")
|
|
logs: List[ChatLog] = Field(default_factory=list, description="Workflow logs")
|
|
messages: List[ChatMessage] = Field(default_factory=list, description="Messages in the workflow")
|
|
stats: Optional[ChatStat] = Field(None, description="Workflow statistics")
|
|
|
|
label: Label = Field(
|
|
default=Label(default="Chat Workflow", translations={"en": "Chat Workflow", "fr": "Flux de travail de chat"}),
|
|
description="Label for the class"
|
|
)
|
|
|
|
fieldLabels: Dict[str, Label] = {
|
|
"id": Label(default="ID", translations={}),
|
|
"mandateId": Label(default="Mandate ID", translations={"en": "Mandate ID", "fr": "ID du mandat"}),
|
|
"status": Label(default="Status", translations={"en": "Status", "fr": "Statut"}),
|
|
"name": Label(default="Name", translations={"en": "Name", "fr": "Nom"}),
|
|
"currentRound": Label(default="Current Round", translations={"en": "Current Round", "fr": "Tour actuel"}),
|
|
"lastActivity": Label(default="Last Activity", translations={"en": "Last Activity", "fr": "Dernière activité"}),
|
|
"startedAt": Label(default="Started At", translations={"en": "Started At", "fr": "Démarré le"}),
|
|
"logs": Label(default="Logs", translations={"en": "Logs", "fr": "Journaux"}),
|
|
"messages": Label(default="Messages", translations={"en": "Messages", "fr": "Messages"}),
|
|
"stats": Label(default="Statistics", translations={"en": "Statistics", "fr": "Statistiques"})
|
|
}
|
|
|
|
# AGENT AND TASK MODELS
|
|
|
|
class Agent(BaseModelWithUI):
|
|
"""Data model for an agent"""
|
|
id: str = Field(description="Primary key")
|
|
name: str = Field(description="Name of the agent")
|
|
description: str = Field(description="Description of the agent")
|
|
capabilities: List[str] = Field(default_factory=list, description="List of agent capabilities")
|
|
|
|
class AgentResponse(BaseModelWithUI):
|
|
"""Data model for an agent response"""
|
|
response: str = Field(description="Response content from the agent")
|
|
documents: List[ChatDocument] = Field(default_factory=list, description="Documents associated with the response")
|
|
|
|
class TaskItem(BaseModelWithUI):
|
|
"""Data model for a task item"""
|
|
sequenceNr: int = Field(description="Sequence number of the task")
|
|
agentName: str = Field(description="Name of the agent assigned to this task")
|
|
prompt: str = Field(description="Prompt for the task")
|
|
userLanguage: str = Field(description="User's preferred language")
|
|
filesInput: List[str] = Field(default_factory=list, description="Input files (format: filename;[documentId])")
|
|
filesOutput: List[str] = Field(default_factory=list, description="Output files (format: filename)")
|
|
|
|
class TaskPlan(BaseModelWithUI):
|
|
"""Data model for a task plan"""
|
|
fileList: List[str] = Field(default_factory=list, description="List of files (format: filename)")
|
|
taskItems: List[TaskItem] = Field(default_factory=list, description="List of task items in the plan")
|
|
userLanguage: str = Field(description="User's preferred language")
|
|
userResponse: str = Field(description="User's response or feedback")
|
|
|
|
class UserInputRequest(BaseModelWithUI):
|
|
"""Data model for a user input request"""
|
|
prompt: str = Field(description="Prompt for the user")
|
|
listFileId: List[int] = Field(default_factory=list, description="List of file IDs") |