Repository: Quansight/ragna (python)

Issue:
don't return the full chat object on the prepare and answer endpoints of the REST API
### Feature description

I've noticed that the `/chats/{id}/answer` endpoint returns both the latest `message` and the full `chat`:

https://github.com/Quansight/ragna/blob/1e256e935e2cb851d686882afd5377002ca011aa/ragna/_api/core.py#L210-L227

I thought that maybe the UI needed all that state about the chat for some distributed system reason, but it looks like it does not:

https://github.com/Quansight/ragna/blob/1e256e935e2cb851d686882afd5377002ca011aa/ragna/_ui/central_view.py#L361-L369

Assuming there is a need to maintain API backward compatibility, could we add a `/chats/{id}/quick_answer` endpoint to Ragna that only returns a `schemas.Message` object as opposed to a heavyweight `schemas.MessageOutput` object?

### Value and/or benefit

Every call to `/chats/{id}/answer` appends another message to the chat. This means that each time the endpoint is called, the amount of data that needs to be sent over the wire increases. Over time, this O(m) growth makes getting answers more and more sluggish. Ideally, the `/chats/{id}/answer` endpoint would be O(1) from a messages in the chat perspective. The proposed `/chats/{id}/quick_answer` endpoint would act in that constant time/space manner.

### Anything else?

_No response_

Pull request description:
don't return full chat as part of message output
Closes #242.

Patch:
diff --git a/ragna/deploy/_api/core.py b/ragna/deploy/_api/core.py
index af2a4e0..6dfa7bf 100644
--- a/ragna/deploy/_api/core.py
+++ b/ragna/deploy/_api/core.py
@@ -208,25 +208,24 @@ def app(config: Config) -> FastAPI:
             return database.get_chat(session, user=user, id=id)
 
     @app.post("/chats/{id}/prepare")
-    async def prepare_chat(
-        user: UserDependency, id: uuid.UUID
-    ) -> schemas.MessageOutput:
+    async def prepare_chat(user: UserDependency, id: uuid.UUID) -> schemas.Message:
         with get_session() as session:
             chat = database.get_chat(session, user=user, id=id)
 
             core_chat = schema_to_core_chat(session, user=user, chat=chat)
+
             welcome = schemas.Message.from_core(await core_chat.prepare())
+
             chat.prepared = True
             chat.messages.append(welcome)
-
             database.update_chat(session, user=user, chat=chat)
 
-            return schemas.MessageOutput(message=welcome, chat=chat)
+            return welcome
 
     @app.post("/chats/{id}/answer")
     async def answer(
         user: UserDependency, id: uuid.UUID, prompt: str
-    ) -> schemas.MessageOutput:
+    ) -> schemas.Message:
         with get_session() as session:
             chat = database.get_chat(session, user=user, id=id)
             chat.messages.append(
@@ -236,11 +235,11 @@ def app(config: Config) -> FastAPI:
             core_chat = schema_to_core_chat(session, user=user, chat=chat)
 
             answer = schemas.Message.from_core(await core_chat.answer(prompt))
-            chat.messages.append(answer)
 
+            chat.messages.append(answer)
             database.update_chat(session, user=user, chat=chat)
 
-            return schemas.MessageOutput(message=answer, chat=chat)
+            return answer
 
     @app.delete("/chats/{id}")
     async def delete_chat(user: UserDependency, id: uuid.UUID) -> None:
diff --git a/ragna/deploy/_api/schemas.py b/ragna/deploy/_api/schemas.py
index 9edc260..b64c2c3 100644
--- a/ragna/deploy/_api/schemas.py
+++ b/ragna/deploy/_api/schemas.py
@@ -78,8 +78,3 @@ class Chat(BaseModel):
     metadata: ChatMetadata
     messages: list[Message] = Field(default_factory=list)
     prepared: bool = False
-
-
-class MessageOutput(BaseModel):
-    message: Message
-    chat: Chat
diff --git a/ragna/deploy/_ui/api_wrapper.py b/ragna/deploy/_ui/api_wrapper.py
index 647be9b..06aa718 100644
--- a/ragna/deploy/_ui/api_wrapper.py
+++ b/ragna/deploy/_ui/api_wrapper.py
@@ -63,7 +63,7 @@ class ApiWrapper(param.Parameterized):
         return json_data
 
     async def answer(self, chat_id, prompt):
-        json_data = (
+        return self.improve_message(
             (
                 await self.client.post(
                     f"/chats/{chat_id}/answer",
@@ -75,13 +75,6 @@ class ApiWrapper(param.Parameterized):
             .json()
         )
 
-        json_data["message"] = self.improve_message(json_data["message"])
-        json_data["chat"]["messages"] = [
-            self.improve_message(msg) for msg in json_data["chat"]["messages"]
-        ]
-
-        return json_data
-
     async def get_components(self):
         return (await self.client.get("/components")).raise_for_status().json()
 
diff --git a/ragna/deploy/_ui/central_view.py b/ragna/deploy/_ui/central_view.py
index 2ea9e88..aa460c1 100644
--- a/ragna/deploy/_ui/central_view.py
+++ b/ragna/deploy/_ui/central_view.py
@@ -360,12 +360,12 @@ class CentralView(pn.viewable.Viewer):
         try:
             answer = await self.api_wrapper.answer(self.current_chat["id"], contents)
 
-            self.current_chat["messages"].append(answer["message"])
+            self.current_chat["messages"].append(answer)
 
             yield {
                 "user": "Ragna",
                 "avatar": "🤖",
-                "value": answer["message"]["content"],
+                "value": answer["content"],
             }
         except Exception as e:
             print(e)


Test patch:
diff --git a/tests/deploy/api/test_e2e.py b/tests/deploy/api/test_e2e.py
index 3e09263..a1c69da 100644
--- a/tests/deploy/api/test_e2e.py
+++ b/tests/deploy/api/test_e2e.py
@@ -99,25 +99,27 @@ def check_api(config):
         assert client.get("/chats").raise_for_status().json() == [chat]
         assert client.get(f"/chats/{chat['id']}").raise_for_status().json() == chat
 
-        json = client.post(f"/chats/{chat['id']}/prepare").raise_for_status().json()
-        message, chat = json["message"], json["chat"]
+        message = client.post(f"/chats/{chat['id']}/prepare").raise_for_status().json()
         assert message["role"] == "system"
         assert message["sources"] == []
+
+        chat = client.get(f"/chats/{chat['id']}").raise_for_status().json()
         assert chat["prepared"]
         assert len(chat["messages"]) == 1
         assert chat["messages"][-1] == message
 
         prompt = "?"
-        json = (
+        message = (
             client.post(f"/chats/{chat['id']}/answer", params={"prompt": prompt})
             .raise_for_status()
             .json()
         )
-        message, chat = json["message"], json["chat"]
         assert message["role"] == "assistant"
         assert {source["document"]["name"] for source in message["sources"]} == {
             document_path.name
         }
+
+        chat = client.get(f"/chats/{chat['id']}").raise_for_status().json()
         assert len(chat["messages"]) == 3
         assert (
             chat["messages"][-2]["role"] == "user"