forked from innovacion/Mayacontigo
124 lines
3.5 KiB
Python
124 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from datetime import UTC, datetime
|
|
|
|
from sqlalchemy import ForeignKey, select
|
|
from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP
|
|
from sqlalchemy.ext.asyncio import AsyncAttrs
|
|
from sqlalchemy.orm import (
|
|
DeclarativeBase,
|
|
Mapped,
|
|
MappedAsDataclass,
|
|
column_property,
|
|
mapped_column,
|
|
relationship,
|
|
)
|
|
|
|
|
|
class Base(DeclarativeBase, MappedAsDataclass, AsyncAttrs):
|
|
type_annotation_map = {
|
|
datetime: TIMESTAMP(timezone=True),
|
|
dict: JSONB,
|
|
}
|
|
|
|
|
|
class CommonMixin(MappedAsDataclass):
|
|
id: Mapped[int] = mapped_column(init=False, primary_key=True)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
init=False, default_factory=lambda: datetime.now(UTC)
|
|
)
|
|
updated_at: Mapped[datetime | None] = mapped_column(
|
|
init=False, default=None, onupdate=lambda: datetime.now(UTC)
|
|
)
|
|
active: Mapped[bool] = mapped_column(init=False, default=True)
|
|
|
|
|
|
class User(CommonMixin, Base):
|
|
__tablename__ = "users"
|
|
|
|
name: Mapped[str] = mapped_column(unique=True, index=True)
|
|
|
|
conversations: Mapped[list[Conversation]] = relationship(
|
|
init=False, back_populates="user"
|
|
)
|
|
|
|
|
|
class Assistant(CommonMixin, Base):
|
|
__tablename__ = "assistants"
|
|
|
|
name: Mapped[str] = mapped_column(unique=True, index=True)
|
|
system_prompt: Mapped[str]
|
|
|
|
conversations: Mapped[list[Conversation]] = relationship(
|
|
init=False, back_populates="assistant"
|
|
)
|
|
|
|
|
|
class Conversation(CommonMixin, Base):
|
|
__tablename__ = "conversations"
|
|
|
|
user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
|
|
assistant_id: Mapped[int] = mapped_column(ForeignKey("assistants.id"))
|
|
|
|
user: Mapped[User] = relationship(init=False, back_populates="conversations")
|
|
assistant: Mapped[Assistant] = relationship(
|
|
init=False, back_populates="conversations"
|
|
)
|
|
messages: Mapped[list[Message]] = relationship(
|
|
init=False, order_by="Message.created_at"
|
|
)
|
|
|
|
assistant_name = column_property(
|
|
select(Assistant.name).where(Assistant.id == assistant_id).scalar_subquery()
|
|
)
|
|
|
|
def add_message(
|
|
self,
|
|
role: str,
|
|
content: str | None = None,
|
|
tools: dict | None = None,
|
|
query_id: int | None = None,
|
|
):
|
|
self.messages.append(
|
|
Message(
|
|
role=role,
|
|
content=content,
|
|
tools=tools,
|
|
query_id=query_id,
|
|
conversation_id=self.id,
|
|
)
|
|
)
|
|
|
|
async def to_openai_format(self):
|
|
messages = await self.awaitable_attrs.messages
|
|
return [(await m.to_openai_format()) for m in messages]
|
|
|
|
|
|
class Message(CommonMixin, Base):
|
|
__tablename__ = "messages"
|
|
|
|
conversation_id: Mapped[int] = mapped_column(ForeignKey("conversations.id"))
|
|
|
|
role: Mapped[str]
|
|
content: Mapped[str | None] = mapped_column(default=None)
|
|
feedback: Mapped[bool | None] = mapped_column(default=None)
|
|
tools: Mapped[dict | None] = mapped_column(default=None)
|
|
query_id: Mapped[int | None] = mapped_column(default=None)
|
|
|
|
async def to_openai_format(self):
|
|
role = await self.awaitable_attrs.role
|
|
content = await self.awaitable_attrs.content
|
|
tools = await self.awaitable_attrs.tools
|
|
return {
|
|
"role": role,
|
|
"content": content,
|
|
**(tools or {}),
|
|
}
|
|
|
|
|
|
class Comment(CommonMixin, Base):
|
|
__tablename__ = "comments"
|
|
|
|
message_id: Mapped[int] = mapped_column(ForeignKey("messages.id"))
|
|
content: Mapped[str]
|