Inference: Associate chats with user IDs (LAION-AI#1826)

Closes LAION-AI#1788.

Author: olliestanley

inference/server/alembic/versions/2023_02_23_2245-b74c66553643_add_user_id_to_chats_table.py

@@ -0,0 +1,24 @@
+"""Add user id to chats table
+
+Revision ID: b74c66553643
+Revises: b365a18db6fd
+Create Date: 2023-02-23 22:45:19.946188
+"""
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "b74c66553643"
+down_revision = "b365a18db6fd"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column("chat", sa.Column("user_id", sa.String(), sa.ForeignKey("user.id"), nullable=True))
+    op.create_index(op.f("ix_chat_user_id"), "chat", ["user_id"], unique=False)
+
+
+def downgrade() -> None:
+    op.drop_index(op.f("ix_chat_user_id"), table_name="chat")
+    op.drop_column("chat", "user_id")

inference/server/main.py

@@ -182,8 +182,9 @@ async def list_chats(
 ) -> interface.ListChatsResponse:
     """Lists all chats."""
     logger.info("Listing all chats.")
-    chats = cr.get_chat_list()
-    return interface.ListChatsResponse(chats=chats)
+    chats = cr.get_chats(user_id)
+    chats_list = [chat.to_list_read() for chat in chats]
+    return interface.ListChatsResponse(chats=chats_list)
 
 
 @app.post("/chat")
@@ -194,7 +195,7 @@ async def create_chat(
 ) -> interface.ChatListRead:
     """Allows a client to create a new chat."""
     logger.info(f"Received {request=}")
-    chat = cr.create_chat()
+    chat = cr.create_chat(user_id)
     return chat.to_list_read()
 
 
@@ -206,6 +207,9 @@ async def get_chat(
 ) -> interface.ChatRead:
     """Allows a client to get the current state of a chat."""
     chat = cr.get_chat_by_id(id)
+    # currently, user_id will be None if server auth is disabled
+    if user_id and chat.user_id != user_id:
+        raise HTTPException(status_code=404, detail="Chat not found")
     return chat.to_read()
 
 

inference/server/oasst_inference_server/__init__.py

@@ -11,6 +11,18 @@
 oauth2_scheme = APIKeyCookie(name=settings.auth_cookie_name)
 
 
+def derive_key() -> bytes:
+    """Derive a key from the auth secret."""
+    hkdf = HKDF(
+        algorithm=hashes.SHA256(),
+        length=settings.auth_length,
+        salt=settings.auth_salt,
+        info=settings.auth_info,
+    )
+    key = hkdf.derive(settings.auth_secret)
+    return key
+
+
 def create_access_token(data: dict) -> str:
     """Create encoded JSON Web Token (JWT) using the given data."""
     expires_delta = timedelta(minutes=settings.auth_access_token_expire_minutes)
@@ -19,32 +31,20 @@ def create_access_token(data: dict) -> str:
     to_encode.update({"exp": expire})
 
     # Generate a key from the auth secret
-    hkdf = HKDF(
-        algorithm=hashes.SHA256(),
-        length=settings.auth_length,
-        salt=settings.auth_salt,
-        info=settings.auth_info,
-    )
-    key = hkdf.derive(settings.auth_secret)
+    key = derive_key()
 
     # Encrypt the payload using JWE
     token: bytes = jwe.encrypt(to_encode, key)
     return token.decode()
 
 
 def get_current_user_id(token: str = Security(oauth2_scheme)) -> str | None:
+    """Decode the current user JWT token and return the payload."""
     if not settings.use_auth:
         return None
 
-    """Decode the current user JWT token and return the payload."""
     # Generate a key from the auth secret
-    hkdf = HKDF(
-        algorithm=hashes.SHA256(),
-        length=settings.auth_length,
-        salt=settings.auth_salt,
-        info=settings.auth_info,
-    )
-    key = hkdf.derive(settings.auth_secret)
+    key = derive_key()
 
     # Decrypt the JWE token
     try:

inference/server/oasst_inference_server/auth.py

@@ -20,8 +20,11 @@ def maybe_commit(self) -> None:
         if self.do_commit:
             self.session.commit()
 
-    def get_chats(self) -> list[models.DbChat]:
-        return self.session.exec(sqlmodel.select(models.DbChat)).all()
+    def get_chats(self, user_id: str | None = None) -> list[models.DbChat]:
+        query = sqlmodel.select(models.DbChat)
+        if user_id:
+            query = query.filter(models.DbChat.user_id == user_id)
+        return self.session.exec(query).all()
 
     def get_pending_chats(self) -> list[models.DbChat]:
         return self.session.exec(
@@ -31,10 +34,6 @@ def get_pending_chats(self) -> list[models.DbChat]:
             )
         ).all()
 
-    def get_chat_list(self) -> list[interface.ChatListRead]:
-        chats = self.get_chats()
-        return [chat.to_read() for chat in chats]
-
     def get_prompter_message_by_id(self, message_id: str, for_update=False) -> models.DbMessage:
         query = sqlmodel.select(models.DbMessage).where(
             models.DbMessage.id == message_id, models.DbMessage.role == "prompter"
@@ -53,15 +52,17 @@ def get_assistant_message_by_id(self, message_id: str, for_update=False) -> mode
         message = self.session.exec(query).one()
         return message
 
-    def get_chat_by_id(self, chat_id: str, for_update=False) -> models.DbChat:
+    def get_chat_by_id(self, chat_id: str, user_id: str | None = None, for_update=False) -> models.DbChat:
         query = sqlmodel.select(models.DbChat).where(models.DbChat.id == chat_id)
+        if user_id:
+            query = query.where(models.DbChat.user_id == user_id)
         if for_update:
             query = query.with_for_update()
         chat = self.session.exec(query).one()
         return chat
 
-    def create_chat(self) -> models.DbChat:
-        chat = models.DbChat()
+    def create_chat(self, user_id: str) -> models.DbChat:
+        chat = models.DbChat(user_id=user_id)
         self.session.add(chat)
         self.maybe_commit()
         return chat

inference/server/oasst_inference_server/chat_repository.py

@@ -1,18 +1,27 @@
 import fastapi
+from fastapi import Depends
 from loguru import logger
-from oasst_inference_server import deps, interface, queueing
+from oasst_inference_server import auth, deps, interface, queueing
 from oasst_shared.schemas import inference
+from sqlalchemy.exc import NoResultFound
 from sse_starlette.sse import EventSourceResponse
 
 
 async def handle_create_message(
     chat_id: str,
     message_request: interface.MessageRequest,
     fastapi_request: fastapi.Request,
+    user_id: str = Depends(auth.get_current_user_id),
 ) -> EventSourceResponse:
     """Allows the client to stream the results of a request."""
 
     with deps.manual_chat_repository() as cr:
+        try:
+            # Ensure the user can access the chat
+            cr.get_chat_by_id(chat_id, user_id=user_id)
+        except NoResultFound:
+            return fastapi.Response(status_code=404)
+
         try:
             prompter_message = cr.add_prompter_message(
                 chat_id=chat_id, parent_id=message_request.parent_id, content=message_request.content

inference/server/oasst_inference_server/client_handler.py

@@ -45,6 +45,8 @@ 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)
+
     messages: list[DbMessage] = Relationship(back_populates="chat")
 
     def to_list_read(self) -> interface.ChatListRead: