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

use a lock for AskAI

Browse source
This commit is contained in:
Nexus 2024-04-18 00:43:37 +01:00
parent 093991ad70
commit 42215bff2a
Signed by: nex
GPG key ID: 0FA334385D0B689F
1 changed files with 28 additions and 19 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

47
src/cogs/ollama.py
View file

@ -447,6 +447,7 @@ class Ollama(commands.Cog):
self.last_server = 0
self.contexts = {}
self.history = ChatHistory()
self.lock = asyncio.Lock()
def next_server(self, increment: bool = True) -> str:
"""Returns the next server key."""
@ -884,30 +885,38 @@ class Ollama(commands.Cog):
"Downloading orca-mini:3b on server %r - %s (%.2f%%)", server, handler.status, handler.percent
)
if self.lock.locked():
await ctx.respond("Waiting for server to be free...")
async with self.lock:
await ctx.delete(delay=0.1)
messages = self.history.get_history(thread)
embed = discord.Embed(description="*waiting...*")
embed = discord.Embed(description="*Waking Ollama up...*")
embed.add_field(
name="\u200b",
value="[in reply to message from {0.author.display_name}]({0.jump_url})".format(message),
)
await ctx.respond(embed=embed)
last_edit = time.time()
with client.new_chat("orca-mini:3b", messages) as handler:
async for ln in handler:
done = ln.get("done") is True
embed.description = handler.result
if len(embed.description) >= 4096:
break
if len(embed.description) >= 3250:
embed.colour = discord.Color.gold()
embed.set_footer(text="Warning: {:,}/4096 characters.".format(len(embed.description)))
else:
embed.colour = discord.Color.blurple()
embed.set_footer(text="Using server %r" % server, icon_url=CONFIG["ollama"][server].get("icon_url"))
if time.time() >= (last_edit + 5.1) or done:
await ctx.edit(embed=embed)
if done:
break
async with self.lock:
await ctx.respond(embed=embed)
last_edit = time.time()
with client.new_chat("orca-mini:3b", messages) as handler:
async for ln in handler:
done = ln.get("done") is True
embed.description = handler.result
if len(embed.description) >= 4096:
break
if len(embed.description) >= 3250:
embed.colour = discord.Color.gold()
embed.set_footer(text="Warning: {:,}/4096 characters.".format(len(embed.description)))
else:
embed.colour = discord.Color.blurple()
embed.set_footer(text="Using server %r" % server, icon_url=CONFIG["ollama"][server].get("icon_url"))
if time.time() >= (last_edit + 5.1) or done is True:
await ctx.edit(embed=embed)
last_edit = time.time()
if done:
break
embed.colour = discord.Colour.dark_theme()
return await ctx.edit(embed=embed)
def setup(bot):