nex/college-bot-v2
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Activity Actions

fix next_server
All checks were successful
Build and Publish Jimmy.2 / build_and_publish (push) Successful in 7s
Details

Browse source
This commit is contained in:
Nexus 2024-05-31 18:00:56 +01:00
parent 34a3cc622e
commit 70afac54ca
Signed by: nex
GPG key ID: 0FA334385D0B689F
1 changed files with 13 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
src/cogs/ollama.py
View file

@ -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)