Allow requiring GPU server

This commit is contained in:
Nexus 2024-07-28 23:50:29 +01:00
parent adfbeb815f
commit 1be3c51160
Signed by: nex
GPG key ID: 0FA334385D0B689F

View file

@ -123,6 +123,11 @@ class AIModule(niobot.Module):
@niobot.command("ollama.chat", greedy=True)
async def chat(self, ctx: niobot.Context):
"""Chat with the model."""
if "--gpu" in ctx.args:
ctx.args.remove("--gpu")
gpu_only = True
else:
gpu_only = False
try:
message = " ".join(ctx.args)
users = self.read_users()
@ -131,7 +136,7 @@ class AIModule(niobot.Module):
return
model = users[ctx.message.sender]
res = await ctx.respond("Finding server...")
server = await self.find_server()
server = await self.find_server(gpu_only=gpu_only)
if not server:
await res.edit(content="No servers available.")
return