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

Allow cancelling a download

Browse source
This commit is contained in:
Nexus 2024-01-11 14:58:00 +00:00
parent a8f15f3ea6
commit 0193ab25ad
1 changed files with 9 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/cogs/ollama.py
View file

@ -276,8 +276,15 @@ class Ollama(commands.Cog):
) )
embed.set_footer(text="Unable to continue.") embed.set_footer(text="Unable to continue.")
return await ctx.edit(embed=embed) return await ctx.edit(embed=embed)
view = OllamaView(ctx)
async for line in self.ollama_stream(response.content): async for line in self.ollama_stream(response.content):
if view.cancel.is_set():
embed = discord.Embed(
title="Download cancelled.",
colour=discord.Colour.red(),
timestamp=discord.utils.utcnow()
)
return await ctx.edit(embed=embed, view=None)
if time.time() >= (last_update + 5.1): if time.time() >= (last_update + 5.1):
if line.get("total") is not None and line.get("completed") is not None: if line.get("total") is not None and line.get("completed") is not None:
percent = (line["completed"] / line["total"]) * 100 percent = (line["completed"] / line["total"]) * 100
@ -285,7 +292,7 @@ class Ollama(commands.Cog):
percent = 50.0 percent = 50.0
embed.fields[0].value = progress_bar(percent, line["status"]) embed.fields[0].value = progress_bar(percent, line["status"])
await ctx.edit(embed=embed) await ctx.edit(embed=embed, view=view)
last_update = time.time() last_update = time.time()
else: else:
self.log.debug("Model %r already exists on server.", model) self.log.debug("Model %r already exists on server.", model)