375 lines
17 KiB
Python
375 lines
17 KiB
Python
"""
|
|
Chat model classes for the chat system.
|
|
"""
|
|
|
|
from pydantic import BaseModel, Field
|
|
from typing import List, Dict, Any, Optional, Union
|
|
from datetime import datetime
|
|
import uuid
|
|
|
|
from modules.shared.attributeUtils import register_model_labels, ModelMixin
|
|
|
|
|
|
# USER MODELS
|
|
|
|
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"}
|
|
}
|
|
)
|
|
|
|
# WORKFLOW MODELS
|
|
|
|
class ChatContent(BaseModel, ModelMixin):
|
|
"""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")
|
|
# Register labels for ChatContent
|
|
register_model_labels(
|
|
"ChatContent",
|
|
{"en": "Chat Content", "fr": "Contenu de chat"},
|
|
{
|
|
"sequenceNr": {"en": "Sequence Number", "fr": "Numéro de séquence"},
|
|
"name": {"en": "Name", "fr": "Nom"},
|
|
"data": {"en": "Data", "fr": "Données"},
|
|
"mimeType": {"en": "MIME Type", "fr": "Type MIME"},
|
|
"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")
|
|
contents: List[ChatContent] = Field(default_factory=list, description="List of chat contents")
|
|
# 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"},
|
|
"contents": {"en": "Contents", "fr": "Contenus"}
|
|
}
|
|
)
|
|
|
|
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")
|
|
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")
|
|
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"},
|
|
"agentName": {"en": "Agent Name", "fr": "Nom de l'agent"},
|
|
"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")
|
|
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")
|
|
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"},
|
|
"agentName": {"en": "Agent Name", "fr": "Nom de l'agent"},
|
|
"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"},
|
|
"startedAt": {"en": "Started At", "fr": "Démarré le"},
|
|
"finishedAt": {"en": "Finished At", "fr": "Terminé le"},
|
|
"stats": {"en": "Statistics", "fr": "Statistiques"},
|
|
"success": {"en": "Success", "fr": "Succès"}
|
|
}
|
|
)
|
|
|
|
class AgentTask(BaseModel, ModelMixin):
|
|
"""Data model for a task"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Primary key")
|
|
workflowId: str = Field(description="Foreign key to workflow")
|
|
agentName: str = Field(description="Name of the agent assigned to this task")
|
|
status: str = Field(description="Current status of the task")
|
|
progress: float = Field(description="Task progress (0-100)")
|
|
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")
|
|
filesOutput: List[str] = Field(default_factory=list, description="Output files")
|
|
result: Optional[ChatMessage] = Field(None, description="Task result message")
|
|
error: Optional[str] = Field(None, description="Error message if failed")
|
|
startedAt: str = Field(description="When the task started")
|
|
finishedAt: Optional[str] = Field(None, description="When the task finished")
|
|
performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")
|
|
# Register labels for AgentTask
|
|
register_model_labels(
|
|
"AgentTask",
|
|
{"en": "Task", "fr": "Tâche"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"workflowId": {"en": "Workflow ID", "fr": "ID du flux de travail"},
|
|
"agentName": {"en": "Agent Name", "fr": "Nom de l'agent"},
|
|
"status": {"en": "Status", "fr": "Statut"},
|
|
"progress": {"en": "Progress", "fr": "Progression"},
|
|
"prompt": {"en": "Prompt", "fr": "Invite"},
|
|
"userLanguage": {"en": "User Language", "fr": "Langue de l'utilisateur"},
|
|
"filesInput": {"en": "Input Files", "fr": "Fichiers d'entrée"},
|
|
"filesOutput": {"en": "Output Files", "fr": "Fichiers de sortie"},
|
|
"result": {"en": "Result", "fr": "Résultat"},
|
|
"error": {"en": "Error", "fr": "Erreur"},
|
|
"startedAt": {"en": "Started At", "fr": "Démarré le"},
|
|
"finishedAt": {"en": "Finished At", "fr": "Terminé le"},
|
|
"performance": {"en": "Performance", "fr": "Performance"}
|
|
}
|
|
)
|
|
|
|
class Agent(BaseModel, ModelMixin):
|
|
"""Data model for an agent"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), 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")
|
|
performance: Optional[Dict[str, Any]] = Field(None, description="Performance metrics")
|
|
# Register labels for Agent
|
|
register_model_labels(
|
|
"Agent",
|
|
{"en": "Agent", "fr": "Agent"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"name": {"en": "Name", "fr": "Nom"},
|
|
"description": {"en": "Description", "fr": "Description"},
|
|
"capabilities": {"en": "Capabilities", "fr": "Capacités"},
|
|
"performance": {"en": "Performance", "fr": "Performance"}
|
|
}
|
|
)
|
|
|
|
# WORKFLOW MODELS
|
|
|
|
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[AgentTask] = 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"}
|
|
}
|
|
)
|
|
|
|
# DOCUMENT MODELS
|
|
|
|
class DocumentExtraction(BaseModel, ModelMixin):
|
|
"""Data model for document extraction history"""
|
|
timestamp: str = Field(description="Timestamp of extraction")
|
|
type: str = Field(description="Type of document")
|
|
sections: List[str] = Field(default_factory=list, description="Extracted sections")
|
|
metadata: Dict[str, Any] = Field(default_factory=dict, description="Extraction metadata")
|
|
|
|
# Register labels for DocumentExtraction
|
|
register_model_labels(
|
|
"DocumentExtraction",
|
|
{"en": "Document Extraction", "fr": "Extraction de document"},
|
|
{
|
|
"timestamp": {"en": "Timestamp", "fr": "Horodatage"},
|
|
"type": {"en": "Type", "fr": "Type"},
|
|
"sections": {"en": "Sections", "fr": "Sections"},
|
|
"metadata": {"en": "Metadata", "fr": "Métadonnées"}
|
|
}
|
|
)
|
|
|
|
class DocumentContext(BaseModel, ModelMixin):
|
|
"""Data model for document context"""
|
|
id: str = Field(description="Document ID")
|
|
extractionHistory: List[DocumentExtraction] = Field(default_factory=list, description="History of extractions")
|
|
relevantSections: List[str] = Field(default_factory=list, description="Relevant sections")
|
|
processingStatus: Dict[str, str] = Field(default_factory=dict, description="Processing status")
|
|
|
|
# Register labels for DocumentContext
|
|
register_model_labels(
|
|
"DocumentContext",
|
|
{"en": "Document Context", "fr": "Contexte de document"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"extractionHistory": {"en": "Extraction History", "fr": "Historique d'extraction"},
|
|
"relevantSections": {"en": "Relevant Sections", "fr": "Sections pertinentes"},
|
|
"processingStatus": {"en": "Processing Status", "fr": "Statut de traitement"}
|
|
}
|
|
)
|
|
|
|
class DocumentMetadata(BaseModel, ModelMixin):
|
|
"""Data model for document metadata"""
|
|
type: str = Field(description="Document type")
|
|
format: str = Field(description="Document format")
|
|
size: int = Field(description="Document size in bytes")
|
|
pages: Optional[int] = Field(None, description="Number of pages")
|
|
sections: Optional[List[str]] = Field(None, description="Document sections")
|
|
error: Optional[str] = Field(None, description="Processing error if any")
|
|
|
|
# Register labels for DocumentMetadata
|
|
register_model_labels(
|
|
"DocumentMetadata",
|
|
{"en": "Document Metadata", "fr": "Métadonnées de document"},
|
|
{
|
|
"type": {"en": "Type", "fr": "Type"},
|
|
"format": {"en": "Format", "fr": "Format"},
|
|
"size": {"en": "Size", "fr": "Taille"},
|
|
"pages": {"en": "Pages", "fr": "Pages"},
|
|
"sections": {"en": "Sections", "fr": "Sections"},
|
|
"error": {"en": "Error", "fr": "Erreur"}
|
|
}
|
|
)
|
|
|
|
class ImageData(BaseModel, ModelMixin):
|
|
"""Data model for image data"""
|
|
data: str = Field(description="Base64 encoded image data")
|
|
format: str = Field(description="Image format")
|
|
page: Optional[int] = Field(None, description="Page number if from a multi-page document")
|
|
index: Optional[int] = Field(None, description="Image index in the document")
|
|
|
|
# Register labels for ImageData
|
|
register_model_labels(
|
|
"ImageData",
|
|
{"en": "Image Data", "fr": "Données d'image"},
|
|
{
|
|
"data": {"en": "Image Data", "fr": "Données d'image"},
|
|
"format": {"en": "Format", "fr": "Format"},
|
|
"page": {"en": "Page", "fr": "Page"},
|
|
"index": {"en": "Index", "fr": "Index"}
|
|
}
|
|
)
|
|
|
|
class DocumentContent(BaseModel, ModelMixin):
|
|
"""Data model for document content"""
|
|
text: Optional[str] = Field(None, description="Extracted text content")
|
|
data: Optional[Dict[str, Any]] = Field(None, description="Structured data content")
|
|
images: Optional[List[ImageData]] = Field(None, description="Extracted images")
|
|
metadata: DocumentMetadata = Field(description="Document metadata")
|
|
|
|
# Register labels for DocumentContent
|
|
register_model_labels(
|
|
"DocumentContent",
|
|
{"en": "Document Content", "fr": "Contenu de document"},
|
|
{
|
|
"text": {"en": "Text", "fr": "Texte"},
|
|
"data": {"en": "Data", "fr": "Données"},
|
|
"images": {"en": "Images", "fr": "Images"},
|
|
"metadata": {"en": "Metadata", "fr": "Métadonnées"}
|
|
}
|
|
)
|
|
|
|
class ProcessedDocument(BaseModel, ModelMixin):
|
|
"""Data model for processed document"""
|
|
id: str = Field(default_factory=lambda: str(uuid.uuid4()), description="Document ID")
|
|
name: str = Field(description="Document name")
|
|
contentType: str = Field(description="Content type")
|
|
content: DocumentContent = Field(description="Document content")
|
|
context: Optional[DocumentContext] = Field(None, description="Document context")
|
|
|
|
# Register labels for ProcessedDocument
|
|
register_model_labels(
|
|
"ProcessedDocument",
|
|
{"en": "Processed Document", "fr": "Document traité"},
|
|
{
|
|
"id": {"en": "ID", "fr": "ID"},
|
|
"name": {"en": "Name", "fr": "Nom"},
|
|
"contentType": {"en": "Content Type", "fr": "Type de contenu"},
|
|
"content": {"en": "Content", "fr": "Contenu"},
|
|
"context": {"en": "Context", "fr": "Contexte"}
|
|
}
|
|
)
|