From 42215bff2aa7a7cbad42c1c7082033fdb42f4cca Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Thu, 18 Apr 2024 00:43:37 +0100
Subject: [PATCH] use a lock for AskAI

---
 src/cogs/ollama.py | 47 +++++++++++++++++++++++++++-------------------
 1 file changed, 28 insertions(+), 19 deletions(-)

diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py
index 48daf0c..b5fa8dd 100644
--- a/src/cogs/ollama.py
+++ b/src/cogs/ollama.py
@@ -447,6 +447,7 @@ class Ollama(commands.Cog):
         self.last_server = 0
         self.contexts = {}
         self.history = ChatHistory()
+        self.lock = asyncio.Lock()
 
     def next_server(self, increment: bool = True) -> str:
         """Returns the next server key."""
@@ -884,30 +885,38 @@ class Ollama(commands.Cog):
                         "Downloading orca-mini:3b on server %r - %s (%.2f%%)", server, handler.status, handler.percent
                     )
 
+        if self.lock.locked():
+            await ctx.respond("Waiting for server to be free...")
+            async with self.lock:
+                await ctx.delete(delay=0.1)
         messages = self.history.get_history(thread)
-        embed = discord.Embed(description="*waiting...*")
+        embed = discord.Embed(description="*Waking Ollama up...*")
         embed.add_field(
             name="\u200b",
             value="[in reply to message from {0.author.display_name}]({0.jump_url})".format(message),
         )
-        await ctx.respond(embed=embed)
-        last_edit = time.time()
-        with client.new_chat("orca-mini:3b", messages) as handler:
-            async for ln in handler:
-                done = ln.get("done") is True
-                embed.description = handler.result
-                if len(embed.description) >= 4096:
-                    break
-                if len(embed.description) >= 3250:
-                    embed.colour = discord.Color.gold()
-                    embed.set_footer(text="Warning: {:,}/4096 characters.".format(len(embed.description)))
-                else:
-                    embed.colour = discord.Color.blurple()
-                    embed.set_footer(text="Using server %r" % server, icon_url=CONFIG["ollama"][server].get("icon_url"))
-                if time.time() >= (last_edit + 5.1) or done:
-                    await ctx.edit(embed=embed)
-                if done:
-                    break
+        async with self.lock:
+            await ctx.respond(embed=embed)
+            last_edit = time.time()
+            with client.new_chat("orca-mini:3b", messages) as handler:
+                async for ln in handler:
+                    done = ln.get("done") is True
+                    embed.description = handler.result
+                    if len(embed.description) >= 4096:
+                        break
+                    if len(embed.description) >= 3250:
+                        embed.colour = discord.Color.gold()
+                        embed.set_footer(text="Warning: {:,}/4096 characters.".format(len(embed.description)))
+                    else:
+                        embed.colour = discord.Color.blurple()
+                        embed.set_footer(text="Using server %r" % server, icon_url=CONFIG["ollama"][server].get("icon_url"))
+                    if time.time() >= (last_edit + 5.1) or done is True:
+                        await ctx.edit(embed=embed)
+                        last_edit = time.time()
+                    if done:
+                        break
+        embed.colour = discord.Colour.dark_theme()
+        return await ctx.edit(embed=embed)
 
 
 def setup(bot):