From f8e3b5368f1fb33e0da474cbe4be6de81e0aea8e Mon Sep 17 00:00:00 2001 From: Piyush Jain Date: Mon, 10 Apr 2023 21:33:31 -0700 Subject: [PATCH] Make provider call async --- packages/jupyter-ai/jupyter_ai/handlers.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/jupyter-ai/jupyter_ai/handlers.py b/packages/jupyter-ai/jupyter_ai/handlers.py index 2a728261d..6ef431ace 100644 --- a/packages/jupyter-ai/jupyter_ai/handlers.py +++ b/packages/jupyter-ai/jupyter_ai/handlers.py @@ -207,7 +207,7 @@ def broadcast_message(self, message: any, exclude_current_user: Optional[bool] = if client: client.write_message(message) - def on_message(self, message): + async def on_message(self, message): self.log.debug("Message recieved: %s", message) try: @@ -226,7 +226,7 @@ def on_message(self, message): self.broadcast_message(message=data, exclude_current_user=True) # process the message - response = self.chat_provider.predict(input=message.content) + response = await ensure_async(self.chat_provider.apredict(input=message.content)) response = AIMessage( content=response