gateway/modules/interfaces/serviceChatModel.py

"""
Chat model classes for the chat system.
"""

from pydantic import BaseModel, Field
from typing import List, Dict, Any, Optional, Union
from datetime import datetime, UTC
import uuid
from enum import Enum

from modules.shared.attributeUtils import register_model_labels, ModelMixin

# ===== Base Enums and Simple Models =====

class TaskStatus(str, Enum):
    """Task status enumeration"""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    ROLLED_BACK = "rolled_back"

# 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")

# 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"}
    }
)

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")
    mimeType: str = Field(description="MIME type of the content")
    base64Encoded: bool = Field(description="Whether the data is base64 encoded")
    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 TaskDocument(BaseModel, ModelMixin):
    """Data model for a task document"""
    id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
    data: str = Field(description="Base64 encoded file data")
    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 TaskDocument
register_model_labels(
    "TaskDocument",
    {"en": "Task Document", "fr": "Document de tâche"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "filename": {"en": "Filename", "fr": "Nom de fichier"},
        "fileSize": {"en": "File Size", "fr": "Taille du fichier"},
        "mimeType": {"en": "MIME Type", "fr": "Type MIME"},
        "data": {"en": "Data", "fr": "Données"}
    }
)

class ExtractedContent(BaseModel, ModelMixin):
    """Data model for extracted content"""
    objectId: str = Field(description="Reference to source document")
    objectType: str = Field(description="Type of source object ('ChatDocument' or 'TaskDocument')")
    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 TaskItem(BaseModel, ModelMixin):
    """Model for tasks"""
    id: str = Field(..., description="Unique task identifier")
    workflowId: str = Field(..., description="Associated workflow ID")
    status: TaskStatus = Field(default=TaskStatus.PENDING, description="Current task status")
    error: Optional[str] = Field(None, description="Error message if task failed")
    startedAt: Optional[datetime] = Field(None, description="Task start timestamp")
    finishedAt: Optional[datetime] = Field(None, description="Task completion timestamp")
    actionList: List[Dict[str, Any]] = Field(default_factory=list, description="List of actions to execute")
    documentsOutput: List[Dict[str, Any]] = Field(default_factory=list, description="Output documents")
    retryCount: int = Field(default=0, description="Number of retry attempts")
    retryMax: int = Field(default=3, description="Maximum number of retry attempts")
    rollbackOnFailure: bool = Field(default=True, description="Whether to rollback on failure")
    dependencies: List[str] = Field(default_factory=list, description="List of dependent task IDs")
    feedback: Optional[Dict[str, Any]] = Field(None, description="Task feedback data")

    def isCompleted(self) -> bool:
        """Check if task is completed"""
        return self.status == TaskStatus.COMPLETED

    def isFailed(self) -> bool:
        """Check if task has failed"""
        return self.status == TaskStatus.FAILED

    def canRetry(self) -> bool:
        """Check if task can be retried"""
        return self.retryCount < self.retryMax

    def start(self) -> None:
        """Start the task"""
        self.status = TaskStatus.RUNNING
        self.startedAt = datetime.now(UTC)

    def complete(self) -> None:
        """Mark task as completed"""
        self.status = TaskStatus.COMPLETED
        self.finishedAt = datetime.now(UTC)

    def fail(self, error: str) -> None:
        """Mark task as failed"""
        self.status = TaskStatus.FAILED
        self.error = error
        self.finishedAt = datetime.now(UTC)

    def cancel(self) -> None:
        """Cancel the task"""
        self.status = TaskStatus.CANCELLED
        self.finishedAt = datetime.now(UTC)

    def rollback(self) -> None:
        """Mark task as rolled back"""
        self.status = TaskStatus.ROLLED_BACK
        self.finishedAt = datetime.now(UTC)

    def incrementRetry(self) -> None:
        """Increment retry count"""
        self.retryCount += 1

    def addDependency(self, taskId: str) -> None:
        """Add a task dependency"""
        if taskId not in self.dependencies:
            self.dependencies.append(taskId)

    def removeDependency(self, taskId: str) -> None:
        """Remove a task dependency"""
        if taskId in self.dependencies:
            self.dependencies.remove(taskId)

    def addAction(self, action: Dict[str, Any]) -> None:
        """Add an action to the task"""
        self.actionList.append(action)

    def addDocumentOutput(self, document: Dict[str, Any]) -> None:
        """Add an output document"""
        self.documentsOutput.append(document)

    def setFeedback(self, feedback: Dict[str, Any]) -> None:
        """Set task feedback"""
        self.feedback = feedback

# Register labels for TaskItem
register_model_labels(
    "TaskItem",
    {"en": "Task", "fr": "Tâche"},
    {
        "id": {"en": "ID", "fr": "ID"},
        "workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
        "status": {"en": "Status", "fr": "Statut"},
        "error": {"en": "Error", "fr": "Erreur"},
        "startedAt": {"en": "Started At", "fr": "Démarré le"},
        "finishedAt": {"en": "Finished At", "fr": "Terminé le"},
        "actionList": {"en": "Action List", "fr": "Liste d'actions"},
        "documentsOutput": {"en": "Output Documents", "fr": "Documents de sortie"},
        "retryCount": {"en": "Retry Count", "fr": "Nombre de tentatives"},
        "retryMax": {"en": "Max Retries", "fr": "Tentatives maximales"},
        "rollbackOnFailure": {"en": "Rollback on Failure", "fr": "Annulation en cas d'échec"},
        "dependencies": {"en": "Dependencies", "fr": "Dépendances"},
        "feedback": {"en": "Task Feedback", "fr": "Retour sur la tâche"}
    }
)

class TaskResult(BaseModel, ModelMixin):
    """Model for task execution results"""
    taskId: str = Field(..., description="ID of the task this result belongs to")
    status: TaskStatus = Field(..., description="Result status")
    success: bool = Field(..., description="Whether the task was successful")
    error: Optional[str] = Field(None, description="Error message if task failed")
    data: Optional[Dict[str, Any]] = Field(None, description="Result data")
    documents: List[ChatDocument] = Field(default_factory=list, description="Output documents")
    feedback: Optional[str] = Field(None, description="Task feedback message")
    processingTime: Optional[float] = Field(None, description="Processing time in seconds")
    timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC), description="When the result was created")

    def isSuccessful(self) -> bool:
        """Check if result indicates success"""
        return self.success and self.status == TaskStatus.COMPLETED

    def hasError(self) -> bool:
        """Check if result has an error"""
        return not self.success or self.status == TaskStatus.FAILED

    def getErrorMessage(self) -> Optional[str]:
        """Get error message if any"""
        return self.error if self.hasError() else None

# Register labels for TaskResult
register_model_labels(
    "TaskResult",
    {"en": "Task Result", "fr": "Résultat de la tâche"},
    {
        "taskId": {"en": "Task ID", "fr": "ID de la tâche"},
        "status": {"en": "Status", "fr": "Statut"},
        "success": {"en": "Success", "fr": "Succès"},
        "error": {"en": "Error", "fr": "Erreur"},
        "data": {"en": "Data", "fr": "Données"},
        "documents": {"en": "Documents", "fr": "Documents"},
        "feedback": {"en": "Feedback", "fr": "Retour"},
        "processingTime": {"en": "Processing Time", "fr": "Temps de traitement"},
        "timestamp": {"en": "Timestamp", "fr": "Horodatage"}
    }
)

# ===== Message and Workflow Models =====

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")
    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")

# 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"},
        "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"}
    }
)

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"}
    }
)