70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
|
"""Message model."""
|
||
|
|
|
||
|
|
from typing import Optional
|
||
|
|
from sqlalchemy import String, Integer, Text, Enum, JSON
|
||
|
|
from sqlalchemy.orm import Mapped, mapped_column
|
||
|
|
import enum
|
||
|
|
|
||
|
|
from ..db.base import BaseModel
|
||
|
|
|
||
|
|
|
||
|
|
class MessageRole(str, enum.Enum):
|
||
|
|
"""Message role enumeration."""
|
||
|
|
USER = "user"
|
||
|
|
ASSISTANT = "assistant"
|
||
|
|
SYSTEM = "system"
|
||
|
|
|
||
|
|
|
||
|
|
class MessageType(str, enum.Enum):
|
||
|
|
"""Message type enumeration."""
|
||
|
|
TEXT = "text"
|
||
|
|
IMAGE = "image"
|
||
|
|
FILE = "file"
|
||
|
|
AUDIO = "audio"
|
||
|
|
|
||
|
|
|
||
|
|
class Message(BaseModel):
|
||
|
|
"""Message model."""
|
||
|
|
|
||
|
|
__tablename__ = "messages"
|
||
|
|
|
||
|
|
conversation_id: Mapped[int] = mapped_column(Integer, nullable=False) # Removed ForeignKey("conversations.id")
|
||
|
|
role: Mapped[MessageRole] = mapped_column(Enum(MessageRole), nullable=False)
|
||
|
|
content: Mapped[str] = mapped_column(Text, nullable=False)
|
||
|
|
message_type: Mapped[MessageType] = mapped_column(Enum(MessageType), default=MessageType.TEXT, nullable=False)
|
||
|
|
message_metadata: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True) # Store additional data like file info, tokens used, etc.
|
||
|
|
|
||
|
|
# For knowledge base context
|
||
|
|
context_documents: Mapped[Optional[dict]] = mapped_column(JSON, nullable=True) # Store retrieved document references
|
||
|
|
|
||
|
|
# Token usage tracking
|
||
|
|
prompt_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
|
||
|
|
completion_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
|
||
|
|
total_tokens: Mapped[Optional[int]] = mapped_column(Integer, nullable=True)
|
||
|
|
|
||
|
|
# Relationships removed to eliminate foreign key constraints
|
||
|
|
|
||
|
|
def __repr__(self):
|
||
|
|
content_preview = self.content[:50] + "..." if len(self.content) > 50 else self.content
|
||
|
|
return f"<Message(id={self.id}, role='{self.role}', content='{content_preview}')>"
|
||
|
|
|
||
|
|
def to_dict(self, include_metadata=True):
|
||
|
|
"""Convert to dictionary."""
|
||
|
|
data = super().to_dict()
|
||
|
|
if not include_metadata:
|
||
|
|
data.pop('message_metadata', None)
|
||
|
|
data.pop('context_documents', None)
|
||
|
|
data.pop('prompt_tokens', None)
|
||
|
|
data.pop('completion_tokens', None)
|
||
|
|
data.pop('total_tokens', None)
|
||
|
|
return data
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_from_user(self):
|
||
|
|
"""Check if message is from user."""
|
||
|
|
return self.role == MessageRole.USER
|
||
|
|
|
||
|
|
@property
|
||
|
|
def is_from_assistant(self):
|
||
|
|
"""Check if message is from assistant."""
|
||
|
|
return self.role == MessageRole.ASSISTANT
|