From 419ed8efe4577477c1ed3a7b51a426ca94f8084a Mon Sep 17 00:00:00 2001
From: nexy7574 <me@nexy7574.co.uk>
Date: Sat, 6 Jan 2024 22:11:20 +0000
Subject: [PATCH] Better log ollama

---
 .idea/workspace.xml | 19 +++++++++++++++----
 src/cogs/ollama.py  |  8 +++++++-
 2 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/.idea/workspace.xml b/.idea/workspace.xml
index b3fde07..27145cc 100644
--- a/.idea/workspace.xml
+++ b/.idea/workspace.xml
@@ -4,7 +4,9 @@
     <option name="autoReloadType" value="SELECTIVE" />
   </component>
   <component name="ChangeListManager">
-    <list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Fix stupid aiohttp base url req" />
+    <list default="true" id="aa6d62a8-d64d-4a60-a85f-8d9fa52b6b49" name="Changes" comment="Fix ollama">
+      <change beforePath="$PROJECT_DIR$/src/cogs/ollama.py" beforeDir="false" afterPath="$PROJECT_DIR$/src/cogs/ollama.py" afterDir="false" />
+    </list>
     <option name="SHOW_DIALOG" value="false" />
     <option name="HIGHLIGHT_CONFLICTS" value="true" />
     <option name="HIGHLIGHT_NON_ACTIVE_CHANGELIST" value="false" />
@@ -76,7 +78,7 @@
       <workItem from="1704296563227" duration="780000" />
       <workItem from="1704317484555" duration="1407000" />
       <workItem from="1704572901106" duration="125000" />
-      <workItem from="1704573038888" duration="5416000" />
+      <workItem from="1704573038888" duration="5748000" />
     </task>
     <task id="LOCAL-00001" summary="Update gitignore">
       <option name="closed" value="true" />
@@ -182,7 +184,15 @@
       <option name="project" value="LOCAL" />
       <updated>1704578291055</updated>
     </task>
-    <option name="localTasksCounter" value="14" />
+    <task id="LOCAL-00014" summary="Fix ollama">
+      <option name="closed" value="true" />
+      <created>1704578625479</created>
+      <option name="number" value="00014" />
+      <option name="presentableId" value="LOCAL-00014" />
+      <option name="project" value="LOCAL" />
+      <updated>1704578625479</updated>
+    </task>
+    <option name="localTasksCounter" value="15" />
     <servers />
   </component>
   <component name="TypeScriptGeneratedFilesManager">
@@ -203,6 +213,7 @@
     <MESSAGE value="FIx unicode" />
     <MESSAGE value="Remove top level import" />
     <MESSAGE value="Fix stupid aiohttp base url req" />
-    <option name="LAST_COMMIT_MESSAGE" value="Fix stupid aiohttp base url req" />
+    <MESSAGE value="Fix ollama" />
+    <option name="LAST_COMMIT_MESSAGE" value="Fix ollama" />
   </component>
 </project>
\ No newline at end of file
diff --git a/src/cogs/ollama.py b/src/cogs/ollama.py
index f50f07d..d6fea50 100644
--- a/src/cogs/ollama.py
+++ b/src/cogs/ollama.py
@@ -144,7 +144,9 @@ class Ollama(commands.Cog):
             await ctx.respond(embed=embed)
 
             try:
+                self.log.debug("Connecting to %r", server_config["base_url"])
                 async with session.post("/api/show", json={"name": model}) as resp:
+                    self.log.debug("%r responded.", server_config["base_url"])
                     if resp.status not in [404, 200]:
                         embed = discord.Embed(
                             url=resp.url,
@@ -166,6 +168,7 @@ class Ollama(commands.Cog):
                 return await ctx.edit(embed=embed)
 
             if resp.status == 404:
+                self.log.debug("Beginning download of %r", model)
                 def progress_bar(value: float, action: str = None):
                     bar = "\N{large green square}" * round(value / 10)
                     bar += "\N{white large square}" * (10 - len(bar))
@@ -207,6 +210,8 @@ class Ollama(commands.Cog):
                             embed.fields[0].value = progress_bar(percent, line["status"])
                             await ctx.edit(embed=embed)
                             last_update = time.time()
+            else:
+                self.log.debug("Model %r already exists on server.", model)
 
             embed = discord.Embed(
                 title="Generating response...",
@@ -215,6 +220,7 @@ class Ollama(commands.Cog):
                 timestamp=discord.utils.utcnow()
             )
             await ctx.edit(embed=embed)
+            self.log.debug("Beginning to generate response.")
             async with session.post(
                 "/api/generate",
                 json={
@@ -242,7 +248,7 @@ class Ollama(commands.Cog):
                         if line.get("done"):
                             embed.title = "Done!"
                             embed.color = discord.Color.green()
-                        embed.description += line["response"]
+                        embed.description += line["response"].strip()
                         embed.timestamp = discord.utils.utcnow()
                         if len(embed.description) >= 4096:
                             embed.description = embed.description[:4093] + "..."