qzl
66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
from __future__ import annotations
|
|
|
|
from decimal import Decimal
|
|
import uuid
|
|
from enum import Enum
|
|
|
|
from sqlalchemy import (
|
|
JSON,
|
|
Enum as SqlEnum,
|
|
ForeignKey,
|
|
Integer,
|
|
Numeric,
|
|
String,
|
|
Text,
|
|
UniqueConstraint,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import JSONB, UUID
|
|
from sqlalchemy.orm import Mapped, mapped_column
|
|
|
|
from core.db.base import Base, SoftDeleteMixin, TimestampMixin
|
|
|
|
|
|
class AgentChatMessageRole(str, Enum):
|
|
USER = "user"
|
|
ASSISTANT = "assistant"
|
|
SYSTEM = "system"
|
|
TOOL = "tool"
|
|
|
|
|
|
class AgentChatMessage(TimestampMixin, SoftDeleteMixin, Base):
|
|
__tablename__: str = "messages"
|
|
__table_args__: tuple[UniqueConstraint] = (
|
|
UniqueConstraint("session_id", "seq", name="uq_messages_session_seq"),
|
|
)
|
|
|
|
id: Mapped[uuid.UUID] = mapped_column(
|
|
UUID(as_uuid=True), primary_key=True, default=uuid.uuid4
|
|
)
|
|
session_id: Mapped[uuid.UUID] = mapped_column(
|
|
UUID(as_uuid=True),
|
|
ForeignKey("sessions.id", ondelete="CASCADE"),
|
|
nullable=False,
|
|
index=True,
|
|
)
|
|
seq: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
role: Mapped[AgentChatMessageRole] = mapped_column(
|
|
SqlEnum(
|
|
AgentChatMessageRole,
|
|
name="agent_chat_message_role",
|
|
native_enum=False,
|
|
values_callable=lambda enum_cls: [item.value for item in enum_cls],
|
|
),
|
|
nullable=False,
|
|
)
|
|
content: Mapped[str] = mapped_column(Text, nullable=False)
|
|
model_code: Mapped[str | None] = mapped_column(String(50), nullable=True)
|
|
tool_name: Mapped[str | None] = mapped_column(String(100), nullable=True)
|
|
input_tokens: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
output_tokens: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
cost: Mapped[Decimal] = mapped_column(Numeric(12, 6), nullable=False, default=0)
|
|
currency: Mapped[str] = mapped_column(String(3), nullable=False, default="USD")
|
|
latency_ms: Mapped[int | None] = mapped_column(Integer, nullable=True)
|
|
metadata_json: Mapped[dict[str, object] | None] = mapped_column(
|
|
"metadata", JSON().with_variant(JSONB, "postgresql"), nullable=True
|
|
)
|