From ac86cd3b4a9e272deda81b1adf78ef921d15b050 Mon Sep 17 00:00:00 2001 From: nexy7574 Date: Fri, 31 May 2024 18:30:04 +0100 Subject: [PATCH] Fix default server --- src/cogs/ollama.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py index 69a00ef..f54bd8e 100644 --- a/src/cogs/ollama.py +++ b/src/cogs/ollama.py @@ -369,6 +369,8 @@ class ChatHistory: SERVER_KEYS = list(CONFIG["ollama"].keys()) +SERVER_KEYS_AUTOCOMPLETE = SERVER_KEYS.copy() +SERVER_KEYS_AUTOCOMPLETE.remove("order") class OllamaGetPrompt(discord.ui.Modal): @@ -444,7 +446,8 @@ class Ollama(commands.Cog): self.servers = { server: asyncio.Lock() for server in CONFIG["ollama"] } - if CONFIG["ollama"].get("order"): + self.servers.pop("order", None) + if CONFIG["ollama"].pop("order"): self.servers = {} for key in CONFIG["ollama"]["order"]: self.servers[key] = asyncio.Lock() @@ -497,7 +500,7 @@ class Ollama(commands.Cog): ), ], server: typing.Annotated[ - str, discord.Option(str, "The server to use for ollama.", default="next", choices=SERVER_KEYS) + str, discord.Option(str, "The server to use for ollama.", default="next", choices=SERVER_KEYS_AUTOCOMPLETE) ], context: typing.Annotated[ str, discord.Option(str, "The context key of a previous ollama response to use as context.", default=None) @@ -513,6 +516,8 @@ class Ollama(commands.Cog): discord.Option(discord.Attachment, "An image to feed into ollama. Only works with llava.", default=None), ], ): + if server == "next": + server = tuple(self.servers.keys())[0] if not SERVER_KEYS: return await ctx.respond("No servers available. Please try again later.") system_query = None