From 086dffd88c228223cb3c04346898a4675779de5d Mon Sep 17 00:00:00 2001
From: nexy7574 <git@nexy7574.co.uk>
Date: Sun, 14 Apr 2024 17:23:05 +0100
Subject: [PATCH] Log download progress to console

---
 src/cogs/ollama.py | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py
index 7bdb87c..8d1933f 100644
--- a/src/cogs/ollama.py
+++ b/src/cogs/ollama.py
@@ -82,6 +82,13 @@ class OllamaDownloadHandler:
     def abort(self):
         self._abort.set()
 
+    def __enter__(self):
+        self._abort.clear()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._abort.set()
+
     @property
     def percent(self) -> float:
         return round((self._completed / self._total) * 100, 2)
@@ -946,7 +953,9 @@ class Ollama(commands.Cog):
 
         client = OllamaClient(CONFIG["ollama"][server]["base_url"])
         if not await client.has_model_named("orca-mini", "3b"):
-            await client.download_model("orca-mini", "3b")
+            with client.download_model("orca-mini", "3b") as handler:
+                async for _ in handler:
+                    self.log.info("Downloading orca-mini:3b on server %r - %.2f%", server, handler.percent)
 
         messages = self.history.get_history(thread)
         embed = discord.Embed(description=">>> ")