""" 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): """Unified model for action results with workflow state management""" # Core result fields 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") error: Optional[str] = Field(None, description="Error message if any") # Action identification actionId: Optional[str] = Field(None, description="ID of the action that produced this result") actionMethod: Optional[str] = Field(None, description="Method of the action that produced this result") actionName: Optional[str] = Field(None, description="Name of the action that produced this result") # Document handling documents: List[str] = Field(default_factory=list, description="List of document references") resultLabel: Optional[str] = Field(None, description="Label for the result") # Validation and workflow state validation: Dict[str, Any] = Field(default_factory=dict, description="Validation information") is_retry: bool = Field(default=False, description="Whether this is a retry attempt") previous_error: Optional[str] = Field(None, description="Previous error message for retries") applied_improvements: List[str] = Field(default_factory=list, description="Improvements applied for retry") @classmethod def success(cls, documents: List[str] = None, resultLabel: str = None, data: Dict[str, Any] = None, actionId: str = None, actionMethod: str = None, actionName: str = None) -> 'ActionResult': """Create a successful action result""" return cls( success=True, data=data or {}, documents=documents or [], resultLabel=resultLabel, actionId=actionId, actionMethod=actionMethod, actionName=actionName ) @classmethod def failure(cls, error: str, data: Dict[str, Any] = None, actionId: str = None, actionMethod: str = None, actionName: str = None) -> 'ActionResult': """Create a failed action result""" return cls( success=False, data=data or {}, error=error, actionId=actionId, actionMethod=actionMethod, actionName=actionName ) @classmethod def retry(cls, previous_result: 'ActionResult', improvements: List[str] = None) -> 'ActionResult': """Create a retry action result based on a previous result""" return cls( success=previous_result.success, data=previous_result.data, metadata=previous_result.metadata, validation=previous_result.validation, error=previous_result.error, documents=previous_result.documents, resultLabel=previous_result.resultLabel, actionId=previous_result.actionId, actionMethod=previous_result.actionMethod, actionName=previous_result.actionName, is_retry=True, previous_error=previous_result.error, applied_improvements=improvements or [] ) # 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"}, "documents": {"en": "Documents", "fr": "Documents"}, "resultLabel": {"en": "Result Label", "fr": "Étiquette du résultat"}, "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"}, "is_retry": {"en": "Is Retry", "fr": "Est une nouvelle tentative"}, "previous_error": {"en": "Previous Error", "fr": "Erreur précédente"}, "applied_improvements": {"en": "Applied Improvements", "fr": "Améliorations appliquées"} } ) # ===== 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") # NEW: Optional document format specification expectedDocumentFormats: Optional[List[Dict[str, str]]] = Field(None, description="Expected document formats (optional)") 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"} } ) # ====== WORKFLOW SUPPORT MODELS (for managerChat.py compatibility) ====== class TaskStep(BaseModel, ModelMixin): id: str objective: str dependencies: Optional[list[str]] = [] success_criteria: Optional[list[str]] = [] estimated_complexity: Optional[str] = None class TaskContext(BaseModel, ModelMixin): task_step: TaskStep workflow: Optional['ChatWorkflow'] = None workflow_id: Optional[str] = None available_documents: Optional[list[str]] = [] previous_results: Optional[list[str]] = [] improvements: Optional[list[str]] = [] retry_count: Optional[int] = 0 previous_action_results: Optional[list] = [] previous_review_result: Optional[dict] = None is_regeneration: Optional[bool] = False failure_patterns: Optional[list[str]] = [] failed_actions: Optional[list] = [] successful_actions: Optional[list] = [] class ReviewContext(BaseModel, ModelMixin): task_step: TaskStep task_actions: Optional[list] = [] action_results: Optional[list] = [] step_result: Optional[dict] = {} workflow_id: Optional[str] = None previous_results: Optional[list[str]] = [] class ReviewResult(BaseModel, ModelMixin): status: str reason: Optional[str] = None improvements: Optional[list[str]] = [] quality_score: Optional[int] = 5 missing_outputs: Optional[list[str]] = [] met_criteria: Optional[list[str]] = [] unmet_criteria: Optional[list[str]] = [] confidence: Optional[float] = 0.5 class TaskPlan(BaseModel, ModelMixin): overview: str tasks: list[TaskStep] class WorkflowResult(BaseModel, ModelMixin): status: str completed_tasks: int total_tasks: int execution_time: float final_results_count: int error: Optional[str] = None phase: Optional[str] = None