31 lines
1.5 KiB
Python
31 lines
1.5 KiB
Python
from typing import Optional
|
|
|
|
from sqlalchemy import Index
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from letta.orm.custom_columns import ToolCallColumn
|
|
from letta.orm.mixins import AgentMixin, OrganizationMixin
|
|
from letta.orm.sqlalchemy_base import SqlalchemyBase
|
|
from letta.schemas.message import Message as PydanticMessage
|
|
from letta.schemas.openai.chat_completions import ToolCall
|
|
|
|
|
|
class Message(SqlalchemyBase, OrganizationMixin, AgentMixin):
|
|
"""Defines data model for storing Message objects"""
|
|
|
|
__tablename__ = "messages"
|
|
__table_args__ = (Index("ix_messages_agent_created_at", "agent_id", "created_at"),)
|
|
__pydantic_model__ = PydanticMessage
|
|
|
|
id: Mapped[str] = mapped_column(primary_key=True, doc="Unique message identifier")
|
|
role: Mapped[str] = mapped_column(doc="Message role (user/assistant/system/tool)")
|
|
text: Mapped[Optional[str]] = mapped_column(nullable=True, doc="Message content")
|
|
model: Mapped[Optional[str]] = mapped_column(nullable=True, doc="LLM model used")
|
|
name: Mapped[Optional[str]] = mapped_column(nullable=True, doc="Name for multi-agent scenarios")
|
|
tool_calls: Mapped[ToolCall] = mapped_column(ToolCallColumn, doc="Tool call information")
|
|
tool_call_id: Mapped[Optional[str]] = mapped_column(nullable=True, doc="ID of the tool call")
|
|
|
|
# Relationships
|
|
agent: Mapped["Agent"] = relationship("Agent", back_populates="messages", lazy="selectin")
|
|
organization: Mapped["Organization"] = relationship("Organization", back_populates="messages", lazy="selectin")
|