ryan
746aaa4844
Postgres had only primary keys on these tables, so loading a conversation, listing conversations, checking for a first message on every send, and fetching memories/artifacts were sequential scans. Migration generated with aerich; index-only apart from the no-op JSONB re-casts aerich emits on every run (as in migrations 16 and 18), which were checked not to rewrite tables. Upgrade and downgrade verified on a throwaway pgvector database. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
74 lines
2.3 KiB
Python
74 lines
2.3 KiB
Python
import enum
|
|
|
|
from tortoise import fields
|
|
from tortoise.contrib.pydantic import (
|
|
pydantic_model_creator,
|
|
pydantic_queryset_creator,
|
|
)
|
|
from tortoise.models import Model
|
|
|
|
|
|
class Speaker(enum.Enum):
|
|
USER = "user"
|
|
SIMBA = "simba"
|
|
TOOL = "tool"
|
|
|
|
|
|
class Conversation(Model):
|
|
id = fields.UUIDField(primary_key=True)
|
|
name = fields.CharField(max_length=255)
|
|
created_at = fields.DatetimeField(auto_now_add=True)
|
|
updated_at = fields.DatetimeField(auto_now=True)
|
|
user: fields.ForeignKeyRelation = fields.ForeignKeyField(
|
|
"models.User", related_name="conversations", null=True
|
|
)
|
|
channel = fields.CharField(max_length=20, default="web", null=True)
|
|
|
|
class Meta:
|
|
table = "conversations"
|
|
# Conversation list: a user's conversations, most recent first.
|
|
indexes = (("user_id", "updated_at"),)
|
|
|
|
|
|
class ConversationMessage(Model):
|
|
id = fields.UUIDField(primary_key=True)
|
|
text = fields.TextField()
|
|
conversation = fields.ForeignKeyField(
|
|
"models.Conversation", related_name="messages"
|
|
)
|
|
created_at = fields.DatetimeField(auto_now_add=True)
|
|
speaker = fields.CharEnumField(enum_type=Speaker, max_length=10)
|
|
image_key = fields.CharField(max_length=512, null=True, default=None)
|
|
message_handle = fields.CharField(max_length=255, null=True, default=None)
|
|
|
|
class Meta:
|
|
table = "conversation_messages"
|
|
# Loading a conversation's history in order.
|
|
indexes = (("conversation_id", "created_at"),)
|
|
|
|
|
|
class UserMemory(Model):
|
|
id = fields.UUIDField(primary_key=True)
|
|
user = fields.ForeignKeyField("models.User", related_name="memories", db_index=True)
|
|
content = fields.TextField()
|
|
created_at = fields.DatetimeField(auto_now_add=True)
|
|
updated_at = fields.DatetimeField(auto_now=True)
|
|
|
|
class Meta:
|
|
table = "user_memories"
|
|
|
|
|
|
PydConversationMessage = pydantic_model_creator(ConversationMessage)
|
|
PydConversation = pydantic_model_creator(
|
|
Conversation, name="Conversation", allow_cycles=True, exclude=("user",)
|
|
)
|
|
PydConversationWithMessages = pydantic_model_creator(
|
|
Conversation,
|
|
name="ConversationWithMessages",
|
|
allow_cycles=True,
|
|
exclude=("user",),
|
|
include=("messages",),
|
|
)
|
|
PydListConversation = pydantic_queryset_creator(Conversation)
|
|
PydListConversationMessage = pydantic_queryset_creator(ConversationMessage)
|