Add status command

app/modules/ai.py

@@ -5,10 +5,10 @@ from contextlib import asynccontextmanager
 
 import httpx
 import typing
+
 from ollama import AsyncClient
 
 import niobot
-import tomllib
 
 if typing.TYPE_CHECKING:
     from ..main import TortoiseIntegratedBot
@@ -155,4 +155,45 @@ class AIModule(niobot.Module):
             await res.edit(content=response["message"]["content"])
         except Exception as e:
             logging.exception(e)
-            await res.edit(content="An error occurred.")
+            await ctx.respond(content="An error occurred.")
+
+    @niobot.command("ollama.status")
+    async def status(self, ctx: niobot.Context, gpu_only: bool = False):
+        """Checks which servers are online."""
+        lines: dict[str, dict[str, str | None | bool]] = {}
+        for name, cfg in self.bot.cfg["ollama"].items():
+            lines[name] = {
+                "url": cfg["url"],
+                "gpu": cfg["gpu"],
+                "online": None
+            }
+
+        emojis = {
+            True: "\N{white heavy check mark}",
+            False: "\N{cross mark}",
+            None: "\N{hourglass with flowing sand}"
+        }
+
+        def get_lines():
+            ln = []
+            for _n, _d in lines.items():
+                if gpu_only and _d["gpu"] is False:
+                    continue
+                ln.append(f"* **{_n}**: {emojis[_d['online']]}")
+            return "\n".join(ln)
+
+        response = await ctx.respond(get_lines())
+
+        for name, cfg in self.bot.cfg["ollama"].items():
+            url = cfg["url"]
+            gpu = cfg["gpu"]
+            if gpu_only and not gpu:
+                continue
+            async with ollama_client(url) as client:
+                try:
+                    await client.ps()
+                except (httpx.HTTPError, ConnectionError):
+                    lines[name]["online"] = False
+                else:
+                    lines[name]["online"] = True
+                await response.edit(content=get_lines())