Merge pull request #932 from lucagobbi/fix_message_endpoint

Run stray in threadpool to allow tool execution on http message endpoint

core/cat/looking_glass/stray_cat.py

@@ -469,17 +469,22 @@ async def __call__(self, message_dict):
 
         return final_output
 
-    def run(self, user_message_json):
+    def run(self, user_message_json, return_message=False):
         try:
             cat_message = self.loop.run_until_complete(self.__call__(user_message_json))
-            # send message back to client
-            self.send_chat_message(cat_message)
+            if return_message:
+                # return the message for HTTP usage
+                return cat_message
+            else:
+                # send message back to client via WS
+                self.send_chat_message(cat_message)
         except Exception as e:
-            # Log any unexpected errors
             log.error(e)
             traceback.print_exc()
-            # Send error as websocket message
-            self.send_error(e)
+            if return_message:
+                return {"error": str(e)}
+            else:
+                self.send_error(e)
 
     def classify(
         self, sentence: str, labels: List[str] | Dict[str, List[str]]

core/cat/routes/base.py

@@ -1,4 +1,5 @@
 from fastapi import APIRouter, Depends, Body
+from fastapi.concurrency import run_in_threadpool
 from typing import Dict
 import tomli
 from cat.auth.permissions import AuthPermission, AuthResource
@@ -27,5 +28,6 @@ async def message_with_cat(
     stray=Depends(HTTPAuth(AuthResource.CONVERSATION, AuthPermission.WRITE)),
 ) -> Dict:
     """Get a response from the Cat"""
-    answer = await stray({"user_id": stray.user_id, **payload})
+    user_message_json = {"user_id": stray.user_id, **payload}
+    answer = await run_in_threadpool(stray.run, user_message_json, True)
     return answer

core/cat/routes/websocket.py

@@ -22,7 +22,7 @@ async def receive_message(websocket: WebSocket, stray: StrayCat):
         user_message["user_id"] = stray.user_id
 
         # Run the `stray` object's method in a threadpool since it might be a CPU-bound operation.
-        await run_in_threadpool(stray.run, user_message)
+        await run_in_threadpool(stray.run, user_message, return_message=False)
 
 
 @router.websocket("/ws")