added chat modified, created, and title fields

inference/server/alembic/versions/2023_03_12_1735-af8668441ff0_added_created_and_title_to_chat.py

@@ -0,0 +1,32 @@
+"""added created and title to chat
+
+Revision ID: af8668441ff0
+Revises: 486fe9e7fb84
+Create Date: 2023-03-12 17:35:27.837252
+
+"""
+import sqlalchemy as sa
+import sqlmodel
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "af8668441ff0"
+down_revision = "486fe9e7fb84"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("chat", sa.Column("created_at", sa.DateTime(), nullable=False))
+    op.add_column("chat", sa.Column("title", sqlmodel.sql.sqltypes.AutoString(), nullable=True))
+    op.create_index(op.f("ix_chat_created_at"), "chat", ["created_at"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_chat_created_at"), table_name="chat")
+    op.drop_column("chat", "title")
+    op.drop_column("chat", "created_at")
+    # ### end Alembic commands ###

inference/server/alembic/versions/2023_03_12_1736-b247f202e522_added_modified_to_chat.py

@@ -0,0 +1,29 @@
+"""added modified to chat
+
+Revision ID: b247f202e522
+Revises: af8668441ff0
+Create Date: 2023-03-12 17:36:53.061631
+
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b247f202e522"
+down_revision = "af8668441ff0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column("chat", sa.Column("modified_at", sa.DateTime(), nullable=False))
+    op.create_index(op.f("ix_chat_modified_at"), "chat", ["modified_at"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_chat_modified_at"), table_name="chat")
+    op.drop_column("chat", "modified_at")
+    # ### end Alembic commands ###

inference/server/oasst_inference_server/models/chat.py

@@ -64,14 +64,26 @@ class DbChat(SQLModel, table=True):
     id: str = Field(default_factory=lambda: str(uuid4()), primary_key=True)
 
     user_id: str = Field(foreign_key="user.id", index=True)
+    created_at: datetime.datetime = Field(default_factory=datetime.datetime.utcnow, index=True)
+    modified_at: datetime.datetime = Field(default_factory=datetime.datetime.utcnow, index=True)
+    title: str | None = Field(None)
 
     messages: list[DbMessage] = Relationship(back_populates="chat")
 
     def to_list_read(self) -> chat_schema.ChatListRead:
-        return chat_schema.ChatListRead(id=self.id)
+        return chat_schema.ChatListRead(
+            id=self.id,
+            created_at=self.created_at,
+            title=self.title,
+        )
 
     def to_read(self) -> chat_schema.ChatRead:
-        return chat_schema.ChatRead(id=self.id, messages=[m.to_read() for m in self.messages])
+        return chat_schema.ChatRead(
+            id=self.id,
+            created_at=self.created_at,
+            title=self.title,
+            messages=[m.to_read() for m in self.messages],
+        )
 
     def get_msg_dict(self) -> dict[str, DbMessage]:
         return {m.id: m for m in self.messages}

inference/server/oasst_inference_server/routes/chats.py

@@ -55,6 +55,7 @@ async def create_message(
     """Allows the client to stream the results of a request."""
 
     try:
+        ucr: UserChatRepository
         async with deps.manual_user_chat_repository(user_id) as ucr:
             prompter_message = await ucr.add_prompter_message(
                 chat_id=chat_id, parent_id=request.parent_id, content=request.content

inference/server/oasst_inference_server/schemas/chat.py

@@ -1,3 +1,4 @@
+import datetime
 from typing import Annotated, Literal, Union
 
 import pydantic
@@ -58,10 +59,11 @@ class CreateChatRequest(pydantic.BaseModel):
 
 class ChatListRead(pydantic.BaseModel):
     id: str
+    created_at: datetime.datetime
+    title: str | None
 
 
-class ChatRead(pydantic.BaseModel):
-    id: str
+class ChatRead(ChatListRead):
     messages: list[inference.MessageRead]
 
 

inference/server/oasst_inference_server/user_chat_repository.py

@@ -72,6 +72,7 @@ async def add_prompter_message(self, chat_id: str, parent_id: str | None, conten
         if parent_id is None:
             if len(chat.messages) > 0:
                 raise fastapi.HTTPException(status_code=400, detail="Trying to add first message to non-empty chat")
+            chat.title = content
         else:
             msg_dict = chat.get_msg_dict()
             if parent_id not in msg_dict:
@@ -87,6 +88,7 @@ async def add_prompter_message(self, chat_id: str, parent_id: str | None, conten
             content=content,
         )
         self.session.add(message)
+        chat.modified_at = message.created_at
 
         await self.session.commit()
         logger.debug(f"Added prompter message {len(content)=} to chat {chat_id}")