From 38f47b96998ededa684512e15be595f15dcb000d Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Tue, 28 May 2024 00:57:18 +0100
Subject: [PATCH] Add ollama status checker

---
 src/cogs/ollama.py | 30 +++++++++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py
index e2a7555..a79fc3e 100644
--- a/src/cogs/ollama.py
+++ b/src/cogs/ollama.py
@@ -907,7 +907,7 @@ class Ollama(commands.Cog):
         messages = self.history.get_history(thread)
         embed = discord.Embed(description="*Waking Ollama up...*")
         self.log.debug("Acquiring lock")
-        async with self.lock:
+        async with self.servers[server]:
             await ctx.respond(embed=embed, ephemeral=True)
             last_edit = time.time()
             msg = None
@@ -937,6 +937,34 @@ class Ollama(commands.Cog):
         embed.colour = discord.Colour.dark_theme()
         return await msg.edit(embed=embed)
 
+    @commands.command()
+    async def ollama_status(self, ctx: commands.Context):
+        embed = discord.Embed(
+            title="Ollama server statuses:",
+            color=discord.Color.blurple(),
+            timestamp=discord.utils.utcnow(),
+        )
+        for server, lock in self.servers.items():
+            embed.add_field(
+                name=server,
+                value="\U000026a0\U0000fe0fIn use" if lock.locked() else "\U000023f3 Checking..."
+            )
+
+        msg = await ctx.reply(embed=embed)
+        last_edit = time.time()
+        for server in self.servers.items():
+            try:
+                ok = await self.check_server(CONFIG["ollama"][server]["base_url"])
+            except Exception as e:
+                ok = False
+                self.log.error("Error while checking server %r: %r", server, e)
+            for field in embed.fields:
+                if field.name == server:
+                    field.value = "\U00002705 Online" if ok else "\U0000274c Offline"
+                    break
+            await asyncio.sleep(1 - (time.time() - last_edit))
+            await msg.edit(embed=embed)
+
 
 def setup(bot):
     bot.add_cog(Ollama(bot))