fix: added conversation validation in chat

antoninoLorenzo · Feb 23, 2025 · 10443fc · 10443fc
1 parent 1d5bcc0
commit 10443fc
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/src/chat/service.py b/src/chat/service.py
@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 from src.agent import Agent, init_default_architecture
 from src.chat.model import AGENT_SETTINGS
@@ -30,7 +30,10 @@ def new_conversation(self, name: str) -> Conversation:
         )
         return self.__memory[new_conversation_id]
 
-    def get_conversation(self, conversation_id: int) -> Conversation:
+    def get_conversation(self, conversation_id: int) -> Optional[Conversation]:
+        """
+        :return: a Conversation or None
+        """
         return self.__memory[conversation_id]
 
     def rename_conversation(self, conversation_id: int, new_name: str):

diff --git a/src/chat/view.py b/src/chat/view.py
@@ -92,6 +92,8 @@ async def query(
         raise HTTPException(status_code=400, detail='expected {"query": str}')
 
     conversation = conversation_service.get_conversation(conversation_id)
+    if not conversation:
+        raise HTTPException(status_code=404, detail='conversation not found')
+
     conversation += Message(role=Role.USER, content=usr_query)
-
     return StreamingResponse(query_generator(agent, conversation))