Add status command
This commit is contained in:
parent
3e110097bf
commit
c4f8f94ef9
1 changed files with 43 additions and 2 deletions
|
@ -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())
|
||||
|
|
Loading…
Reference in a new issue