456 lines
21 KiB
Python
456 lines
21 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, UTC
|
|
import uuid
|
|
from enum import Enum
|
|
|
|
from modules.shared.attributeUtils import register_model_labels, ModelMixin
|
|
|
|
# ===== Method Models =====
|
|
|
|
class ActionResult(BaseModel, ModelMixin):
|
|
"""Model for action results from a methods action"""
|
|
success: bool = Field(description="Whether the method execution was successful")
|
|
data: Dict[str, Any] = Field(description="Result data")
|
|
metadata: Dict[str, Any] = Field(default_factory=dict, description="Additional metadata")
|
|
validation: List[str] = Field(default_factory=list, description="Validation messages")
|
|
error: Optional[str] = Field(None, description="Error message if any")
|
|
|
|
# Register labels for ActionResult
|
|
register_model_labels(
|
|
"ActionResult",
|
|
{"en": "Action Result", "fr": "Résultat de l'action"},
|
|
{
|
|
"success": {"en": "Success", "fr": "Succès"},
|
|
"data": {"en": "Data", "fr": "Données"},
|
|
"metadata": {"en": "Metadata", "fr": "Métadonnées"},
|
|
"validation": {"en": "Validation", "fr": "Validation"},
|
|
"error": {"en": "Error", "fr": "Erreur"}
|
|
}
|
|
)
|
|
|
|
# ===== Base Enums and Simple Models =====
|
|
|
|
class TaskStatus(str, Enum):
|
|
"""Task status enumeration"""
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
CANCELLED = "cancelled"
|
|
|
|
# Register labels for TaskStatus
|
|
register_model_labels(
|
|
"TaskStatus",
|
|
{"en": "Task Status", "fr": "Statut de la tâche"},
|
|
{
|
|
"PENDING": {"en": "Pending", "fr": "En attente"},
|
|
"RUNNING": {"en": "Running", "fr": "En cours"},
|
|
"COMPLETED": {"en": "Completed", "fr": "Terminé"},
|
|
"FAILED": {"en": "Failed", "fr": "Échec"},
|
|
"CANCELLED": {"en": "Cancelled", "fr": "Annulé"},
|
|
"ROLLED_BACK": {"en": "Rolled Back", "fr": "Annulé"}
|
|
}
|
|
)
|
|
|
|
class UserInputRequest(BaseModel, ModelMixin):
|
|
"""Data model for a user input request"""
|
|
prompt: str = Field(description="Prompt for the user")
|
|
listFileId: List[str] = Field(default_factory=list, description="List of file IDs")
|
|
userLanguage: str = Field(default="en", description="User's preferred language")
|
|
|
|
# Register labels for UserInputRequest
|
|
register_model_labels(
|
|
"UserInputRequest",
|
|
{"en": "User Input Request", "fr": "Demande de saisie utilisateur"},
|
|
{
|
|
"prompt": {"en": "Prompt", "fr": "Invite"},
|
|
"listFileId": {"en": "File IDs", "fr": "IDs des fichiers"},
|
|
"userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"}
|
|
}
|
|
)
|
|
|
|
# ===== Content Models =====
|
|
|
|
class ContentMetadata(BaseModel, ModelMixin):
|
|
"""Metadata for content items"""
|
|
size: int = Field(description="Content size in bytes")
|
|
pages: Optional[int] = Field(None, description="Number of pages for multi-page content")
|
|
error: Optional[str] = Field(None, description="Processing error if any")
|
|
width: Optional[int] = Field(None, description="Width in pixels for images/videos")
|
|
height: Optional[int] = Field(None, description="Height in pixels for images/videos")
|
|
colorMode: Optional[str] = Field(None, description="Color mode (e.g., RGB, CMYK, grayscale)")
|
|
fps: Optional[float] = Field(None, description="Frames per second for videos")
|
|
durationSec: Optional[float] = Field(None, description="Duration in seconds for videos/audio")
|
|
mimeType: str = Field(description="MIME type of the content")
|
|
base64Encoded: bool = Field(description="Whether the data is base64 encoded")
|
|
|
|
# Register labels for ContentMetadata
|
|
register_model_labels(
|
|
"ContentMetadata",
|
|
{"en": "Content Metadata", "fr": "Métadonnées du contenu"},
|
|
{
|
|
"size": {"en": "Size", "fr": "Taille"},
|
|
"pages": {"en": "Pages", "fr": "Pages"},
|
|
"error": {"en": "Error", "fr": "Erreur"},
|
|
"width": {"en": "Width", "fr": "Largeur"},
|
|
"height": {"en": "Height", "fr": "Hauteur"},
|
|
"colorMode": {"en": "Color Mode", "fr": "Mode de couleur"},
|
|
"fps": {"en": "FPS", "fr": "IPS"},
|
|
"durationSec": {"en": "Duration", "fr": "Durée"},
|
|
"mimeType": {"en": "MIME Type", "fr": "Type MIME"},
|
|
"base64Encoded": {"en": "Base64 Encoded", "fr": "Encodé en Base64"}
|
|
}
|
|
)
|
|
|
|
class ContentItem(BaseModel, ModelMixin):
|
|
"""Individual content item from a document"""
|
|
label: str = Field(description="Content label (e.g., tab name, tag name)")
|
|
data: str = Field(description="Extracted text content")
|
|
metadata: ContentMetadata = Field(description="Content metadata")
|
|
|
|
# Register labels for ContentItem
|
|
register_model_labels(
|
|
"ContentItem",
|
|
{"en": "Content Item", "fr": "Élément de contenu"},
|
|
{
|
|
"label": {"en": "Label", "fr": "Étiquette"},
|
|
"data": {"en": "Data", "fr": "Données"},
|
|
"metadata": {"en": "Metadata", "fr": "Métadonnées"}
|
|
}
|
|
)
|
|
|
|
class ChatDocument(BaseModel, ModelMixin):
|
|
"""Data model for a chat document"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
|
|
fileId: str = 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")
|
|
|
|
# Register labels for ChatDocument
|
|
register_model_labels(
|
|
"ChatDocument",
|
|
{"en": "Chat Document", "fr": "Document de chat"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"fileId": {"en": "File ID", "fr": "ID du fichier"},
|
|
"filename": {"en": "Filename", "fr": "Nom de fichier"},
|
|
"fileSize": {"en": "File Size", "fr": "Taille du fichier"},
|
|
"mimeType": {"en": "MIME Type", "fr": "Type MIME"}
|
|
}
|
|
)
|
|
|
|
class DocumentExchange(BaseModel, ModelMixin):
|
|
"""Data model for document exchange between AI actions"""
|
|
documentsLabel: str = Field(description="Label for the set of documents")
|
|
documents: List[str] = Field(default_factory=list, description="List of document references")
|
|
|
|
# Register labels for DocumentExchange
|
|
register_model_labels(
|
|
"DocumentExchange",
|
|
{"en": "Document Exchange", "fr": "Échange de documents"},
|
|
{
|
|
"documentsLabel": {"en": "Documents Label", "fr": "Label des documents"},
|
|
"documents": {"en": "Documents", "fr": "Documents"}
|
|
}
|
|
)
|
|
|
|
class ExtractedContent(BaseModel, ModelMixin):
|
|
"""Data model for extracted content"""
|
|
id: str = Field(description="Reference to source ChatDocument")
|
|
contents: List[ContentItem] = Field(default_factory=list, description="List of content items")
|
|
|
|
# Register labels for ExtractedContent
|
|
register_model_labels(
|
|
"ExtractedContent",
|
|
{"en": "Extracted Content", "fr": "Contenu extrait"},
|
|
{
|
|
"objectId": {"en": "Object ID", "fr": "ID de l'objet"},
|
|
"objectType": {"en": "Object Type", "fr": "Type d'objet"},
|
|
"contents": {"en": "Contents", "fr": "Contenus"}
|
|
}
|
|
)
|
|
|
|
# ===== Task Models =====
|
|
|
|
class TaskAction(BaseModel, ModelMixin):
|
|
"""Model for task actions"""
|
|
id: str = Field(..., description="Action ID")
|
|
execMethod: str = Field(..., description="Method to execute")
|
|
execAction: str = Field(..., description="Action to perform")
|
|
execParameters: Dict[str, Any] = Field(default_factory=dict, description="Action parameters")
|
|
execResultLabel: Optional[str] = Field(None, description="Label for the set of result documents")
|
|
status: TaskStatus = Field(default=TaskStatus.PENDING, description="Action status")
|
|
error: Optional[str] = Field(None, description="Error message if action failed")
|
|
retryCount: int = Field(default=0, description="Number of retries attempted")
|
|
retryMax: int = Field(default=3, description="Maximum number of retries")
|
|
processingTime: Optional[float] = Field(None, description="Processing time in seconds")
|
|
timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC), description="When the action was executed")
|
|
result: Optional[str] = Field(None, description="Result of the action")
|
|
resultDocuments: Optional[List[ChatDocument]] = Field(None, description="Result documents from the action")
|
|
|
|
def isSuccessful(self) -> bool:
|
|
"""Check if action was successful"""
|
|
return self.status == TaskStatus.COMPLETED
|
|
|
|
def hasError(self) -> bool:
|
|
"""Check if action has an error"""
|
|
return self.status == TaskStatus.FAILED
|
|
|
|
def getErrorMessage(self) -> Optional[str]:
|
|
"""Get error message if any"""
|
|
return self.error if self.hasError() else None
|
|
|
|
def setError(self, error: str) -> None:
|
|
"""Set action error"""
|
|
self.error = error
|
|
self.status = TaskStatus.FAILED
|
|
|
|
def setSuccess(self) -> None:
|
|
"""Set action as successful"""
|
|
self.status = TaskStatus.COMPLETED
|
|
self.error = None
|
|
|
|
# Register labels for TaskAction
|
|
register_model_labels(
|
|
"TaskAction",
|
|
{"en": "Task Action", "fr": "Action de tâche"},
|
|
{
|
|
"id": {"en": "Action ID", "fr": "ID de l'action"},
|
|
"execMethod": {"en": "Method", "fr": "Méthode"},
|
|
"execAction": {"en": "Action", "fr": "Action"},
|
|
"execParameters": {"en": "Parameters", "fr": "Paramètres"},
|
|
"execResultLabel": {"en": "Result Label", "fr": "Label du résultat"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"error": {"en": "Error", "fr": "Erreur"},
|
|
"retryCount": {"en": "Retry Count", "fr": "Nombre de tentatives"},
|
|
"retryMax": {"en": "Max Retries", "fr": "Tentatives max"},
|
|
"processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
|
|
"timestamp": {"en": "Timestamp", "fr": "Horodatage"},
|
|
"result": {"en": "Result", "fr": "Résultat"},
|
|
"resultDocuments": {"en": "Result Documents", "fr": "Documents de résultat"}
|
|
}
|
|
)
|
|
|
|
class TaskResult(BaseModel, ModelMixin):
|
|
"""Model for task results"""
|
|
taskId: str = Field(..., description="Task ID")
|
|
status: TaskStatus = Field(default=TaskStatus.PENDING, description="Task status")
|
|
success: bool = Field(..., description="Whether the task was successful")
|
|
feedback: Optional[str] = Field(None, description="Task feedback message")
|
|
error: Optional[str] = Field(None, description="Error message if task failed")
|
|
|
|
# Register labels for TaskResult
|
|
register_model_labels(
|
|
"TaskResult",
|
|
{"en": "Task Result", "fr": "Résultat de tâche"},
|
|
{
|
|
"taskId": {"en": "Task ID", "fr": "ID de la tâche"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"success": {"en": "Success", "fr": "Succès"},
|
|
"feedback": {"en": "Feedback", "fr": "Retour"},
|
|
"error": {"en": "Error", "fr": "Erreur"}
|
|
}
|
|
)
|
|
|
|
class TaskItem(BaseModel, ModelMixin):
|
|
"""Model for workflow tasks"""
|
|
id: str = Field(..., description="Task ID")
|
|
workflowId: str = Field(..., description="Workflow ID")
|
|
userInput: str = Field(..., description="User input that triggered the task")
|
|
status: TaskStatus = Field(default=TaskStatus.PENDING, description="Task status")
|
|
error: Optional[str] = Field(None, description="Error message if task failed")
|
|
startedAt: Optional[str] = Field(None, description="When the task started")
|
|
finishedAt: Optional[str] = Field(None, description="When the task finished")
|
|
actionList: List[TaskAction] = Field(default_factory=list, description="List of actions to execute")
|
|
retryCount: int = Field(default=0, description="Number of retries attempted")
|
|
retryMax: int = Field(default=3, description="Maximum number of retries")
|
|
rollbackOnFailure: bool = Field(default=True, description="Whether to rollback on failure")
|
|
dependencies: List[str] = Field(default_factory=list, description="List of task IDs this task depends on")
|
|
feedback: Optional[str] = Field(None, description="Task feedback message")
|
|
processingTime: Optional[float] = Field(None, description="Total processing time in seconds")
|
|
resultLabels: Optional[Dict[str, Any]] = Field(default_factory=dict, description="Map of result labels to their values")
|
|
|
|
def isSuccessful(self) -> bool:
|
|
"""Check if task was successful"""
|
|
return self.status == TaskStatus.COMPLETED
|
|
|
|
def hasError(self) -> bool:
|
|
"""Check if task has an error"""
|
|
return self.status == TaskStatus.FAILED
|
|
|
|
def getErrorMessage(self) -> Optional[str]:
|
|
"""Get error message if any"""
|
|
return self.error if self.hasError() else None
|
|
|
|
def getResultDocuments(self) -> List[ChatDocument]:
|
|
"""Get all documents from all successful actions"""
|
|
documents = []
|
|
for action in self.actionList:
|
|
if action.isSuccessful() and action.resultDocuments:
|
|
documents.extend(action.resultDocuments)
|
|
return documents
|
|
|
|
def getResultDocumentLabel(self) -> Optional[str]:
|
|
"""Get the label for the result documents"""
|
|
for action in self.actionList:
|
|
if action.isSuccessful() and action.execResultLabel:
|
|
return action.execResultLabel
|
|
return None
|
|
|
|
def getResultLabel(self, label: str) -> Optional[Any]:
|
|
"""Get value for a specific result label"""
|
|
return self.resultLabels.get(label) if self.resultLabels else None
|
|
|
|
# Register labels for TaskItem
|
|
register_model_labels(
|
|
"TaskItem",
|
|
{"en": "Task", "fr": "Tâche"},
|
|
{
|
|
"id": {"en": "Task ID", "fr": "ID de la tâche"},
|
|
"workflowId": {"en": "Workflow ID", "fr": "ID du workflow"},
|
|
"userInput": {"en": "User Input", "fr": "Entrée utilisateur"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"error": {"en": "Error", "fr": "Erreur"},
|
|
"startedAt": {"en": "Started At", "fr": "Démarré à"},
|
|
"finishedAt": {"en": "Finished At", "fr": "Terminé à"},
|
|
"actionList": {"en": "Actions", "fr": "Actions"},
|
|
"retryCount": {"en": "Retry Count", "fr": "Nombre de tentatives"},
|
|
"retryMax": {"en": "Max Retries", "fr": "Tentatives max"},
|
|
"rollbackOnFailure": {"en": "Rollback On Failure", "fr": "Annuler en cas d'échec"},
|
|
"dependencies": {"en": "Dependencies", "fr": "Dépendances"},
|
|
"feedback": {"en": "Feedback", "fr": "Retour"},
|
|
"processingTime": {"en": "Processing Time", "fr": "Temps de traitement"}
|
|
}
|
|
)
|
|
|
|
class ChatStat(BaseModel, ModelMixin):
|
|
"""Data model for chat statistics"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), 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")
|
|
successRate: Optional[float] = Field(None, description="Success rate of operations")
|
|
errorCount: Optional[int] = Field(None, description="Number of errors encountered")
|
|
|
|
# Register labels for ChatStat
|
|
register_model_labels(
|
|
"ChatStat",
|
|
{"en": "Chat Statistics", "fr": "Statistiques de chat"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
|
|
"tokenCount": {"en": "Token Count", "fr": "Nombre de tokens"},
|
|
"bytesSent": {"en": "Bytes Sent", "fr": "Octets envoyés"},
|
|
"bytesReceived": {"en": "Bytes Received", "fr": "Octets reçus"},
|
|
"successRate": {"en": "Success Rate", "fr": "Taux de succès"},
|
|
"errorCount": {"en": "Error Count", "fr": "Nombre d'erreurs"}
|
|
}
|
|
)
|
|
|
|
class ChatLog(BaseModel, ModelMixin):
|
|
"""Data model for a chat log"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), 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")
|
|
status: str = Field(description="Status of the log entry")
|
|
progress: Optional[int] = Field(None, description="Progress percentage")
|
|
performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")
|
|
|
|
# Register labels for ChatLog
|
|
register_model_labels(
|
|
"ChatLog",
|
|
{"en": "Chat Log", "fr": "Journal de chat"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
|
|
"message": {"en": "Message", "fr": "Message"},
|
|
"type": {"en": "Type", "fr": "Type"},
|
|
"timestamp": {"en": "Timestamp", "fr": "Horodatage"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"progress": {"en": "Progress", "fr": "Progression"},
|
|
"performance": {"en": "Performance", "fr": "Performance"}
|
|
}
|
|
)
|
|
|
|
class ChatMessage(BaseModel, ModelMixin):
|
|
"""Data model for a chat message"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
|
|
workflowId: str = Field(description="Foreign key to workflow")
|
|
parentMessageId: Optional[str] = Field(None, description="Parent message ID for threading")
|
|
documents: List[ChatDocument] = Field(default_factory=list, description="Associated documents")
|
|
documentsLabel: Optional[str] = Field(None, description="Label for the set of 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 (first, step, last)")
|
|
sequenceNr: int = Field(description="Sequence number of the message (set automatically)")
|
|
publishedAt: str = Field(description="When the message was published")
|
|
stats: Optional[ChatStat] = Field(None, description="Statistics for this message")
|
|
success: Optional[bool] = Field(None, description="Whether the message processing was successful")
|
|
actionId: Optional[str] = Field(None, description="ID of the action that produced this message")
|
|
actionMethod: Optional[str] = Field(None, description="Method of the action that produced this message")
|
|
actionName: Optional[str] = Field(None, description="Name of the action that produced this message")
|
|
|
|
# Register labels for ChatMessage
|
|
register_model_labels(
|
|
"ChatMessage",
|
|
{"en": "Chat Message", "fr": "Message de chat"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
|
|
"parentMessageId": {"en": "Parent Message ID", "fr": "ID du message parent"},
|
|
"documents": {"en": "Documents", "fr": "Documents"},
|
|
"documentsLabel": {"en": "Documents Label", "fr": "Label des documents"},
|
|
"message": {"en": "Message", "fr": "Message"},
|
|
"role": {"en": "Role", "fr": "Rôle"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"sequenceNr": {"en": "Sequence Number", "fr": "Numéro de séquence"},
|
|
"publishedAt": {"en": "Published At", "fr": "Publié le"},
|
|
"stats": {"en": "Statistics", "fr": "Statistiques"},
|
|
"success": {"en": "Success", "fr": "Succès"},
|
|
"actionId": {"en": "Action ID", "fr": "ID de l'action"},
|
|
"actionMethod": {"en": "Action Method", "fr": "Méthode de l'action"},
|
|
"actionName": {"en": "Action Name", "fr": "Nom de l'action"}
|
|
}
|
|
)
|
|
|
|
class ChatWorkflow(BaseModel, ModelMixin):
|
|
"""Data model for a chat workflow"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), 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")
|
|
tasks: List[TaskItem] = Field(default_factory=list, description="List of tasks in the workflow")
|
|
|
|
# Register labels for ChatWorkflow
|
|
register_model_labels(
|
|
"ChatWorkflow",
|
|
{"en": "Chat Workflow", "fr": "Flux de travail de chat"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"mandateId": {"en": "Mandate ID", "fr": "ID du mandat"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"name": {"en": "Name", "fr": "Nom"},
|
|
"currentRound": {"en": "Current Round", "fr": "Tour actuel"},
|
|
"lastActivity": {"en": "Last Activity", "fr": "Dernière activité"},
|
|
"startedAt": {"en": "Started At", "fr": "Démarré le"},
|
|
"logs": {"en": "Logs", "fr": "Journaux"},
|
|
"messages": {"en": "Messages", "fr": "Messages"},
|
|
"stats": {"en": "Statistics", "fr": "Statistiques"},
|
|
"tasks": {"en": "Tasks", "fr": "Tâches"}
|
|
}
|
|
)
|