From 70afac54ca4b71f39fd9aa0d65298f32dd05b7fd Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Fri, 31 May 2024 18:00:56 +0100
Subject: [PATCH] fix next_server

---
 src/cogs/ollama.py | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py
index 643fb67..7e3e473 100644
--- a/src/cogs/ollama.py
+++ b/src/cogs/ollama.py
@@ -449,15 +449,17 @@ class Ollama(commands.Cog):
             for key in CONFIG["ollama"]["order"]:
                 self.servers[key] = asyncio.Lock()
 
-    def next_server(self) -> str:
+    def next_server(self, tried: typing.Iterable[str] = None) -> str:
         """
         Returns the next server key.
 
+        :param tried: A list of keys already tried
         :returns: The key for the server
         :raises: RuntimeError - If no servers are available.
         """
+        tried = tried or set()
         for server_name, locked in self.servers.items():
-            if not locked.locked():
+            if locked.locked() is False and server_name not in tried:
                 return server_name
         raise RuntimeError("No servers available.")
 
@@ -593,11 +595,15 @@ class Ollama(commands.Cog):
             embed.set_footer(text="Using server %r" % server, icon_url=server_config.get("icon_url"))
             await ctx.respond(embed=embed)
             if not await self.check_server(server_config["base_url"]):
+                tried = {server}
                 for i in range(10):
                     try:
-                        server = self.next_server()
+                        server = self.next_server(tried)
                     except RuntimeError:
+                        tried.add(server)
                         continue
+                    finally:
+                        tried.add(server)
                     embed = discord.Embed(
                         title="Server was offline. Trying next server.",
                         description=f"Trying server {server}...",
@@ -709,8 +715,8 @@ class Ollama(commands.Cog):
             )
             embed.set_author(
                 name=model,
-                url="https://ollama.ai/library/" + model.split(":")[0],
-                icon_url="https://ollama.ai/public/ollama.png",
+                url="https://ollama.com/library/" + model.split(":")[0],
+                icon_url="https://ollama.com/public/ollama.png",
             )
             embed.add_field(
                 name="Prompt", value=">>> " + textwrap.shorten(query, width=1020, placeholder="..."), inline=False
@@ -885,10 +891,12 @@ class Ollama(commands.Cog):
         user_message = {"role": "user", "content": message.content}
         self.history.add_message(thread, "user", user_message["content"])
 
+        tried = set()
         for _ in range(10):
             server = self.next_server()
             if await self.check_server(CONFIG["ollama"][server]["base_url"]):
                 break
+            tried.add(server)
         else:
             return await ctx.respond("All servers are offline. Please try again later.", ephemeral=True)