From 19101f29af98fd8a88b2fe4296d14b21f1ed44ba Mon Sep 17 00:00:00 2001
From: projectmoon <projectmoon@agnos.is>
Date: Sat, 23 Nov 2024 12:19:23 +0100
Subject: [PATCH] thinking filter: more reliable output

---
 .gitignore   |  1 +
 CHANGELOG.md |  3 +++
 thinking.py  | 22 +++++++++++++++++++++-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 3d59aed..b65bb2a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@ chroma/
 chromatest.py
 env/
 gpu_layer_scaler_rocm.py
+exif.py
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7f69e6b..cd6f336 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -179,6 +179,9 @@
 
 # Collapsible Thought Filter
 
+**0.2.0:**
+ - Fix issue with output disappearing.
+
 **0.1.0:**
  - Initial release.
 
diff --git a/thinking.py b/thinking.py
index e15c487..47d37e8 100644
--- a/thinking.py
+++ b/thinking.py
@@ -2,7 +2,7 @@
 title: Collapsible Thought Filter
 author: projectmoon
 author_url: https://git.agnos.is/projectmoon/open-webui-filters
-version: 0.1.0
+version: 0.2.0
 license: AGPL-3.0+, MIT
 required_open_webui_version: 0.3.32
 """
@@ -87,6 +87,10 @@ class Filter:
         tag = self.valves.thought_tag
         return f"<{tag}>(.*?)</{tag}>"
 
+    def _create_output_regex(self) -> str:
+        tag = self.valves.output_tag
+        return f"<{tag}>(.*?)</{tag}>"
+
     def _create_thought_tag_deletion_regex(self) -> str:
         tag = self.valves.thought_tag
         return "</?{{THINK}}>[\s\S]*?</{{THINK}}>".replace("{{THINK}}", tag)
@@ -101,9 +105,18 @@ class Filter:
 
         # collapsible thinking process section
         thought_regex = self._create_thought_regex()
+        output_regex = self._create_output_regex()
         reply = messages[-1]["content"]
+
         thoughts = re.findall(thought_regex, reply, re.DOTALL)
         thoughts = "\n".join(thoughts).strip()
+
+        output = re.findall(output_regex, reply, re.DOTALL)
+        output = "\n".join(output).strip()
+
+        print(thoughts)
+        print(output)
+
         enclosure = THOUGHT_ENCLOSURE.replace("{{THOUGHT_TITLE}}", self.valves.thought_title)
         enclosure = enclosure.replace("{{THOUGHTS}}", thoughts).strip()
 
@@ -116,6 +129,13 @@ class Filter:
         reply = reply.replace(f"<{self.valves.output_tag}>", "", 1)
         reply = reply.replace(f"</{self.valves.output_tag}>", "", 1)
 
+        # because some models do not close the output tag, we prefer
+        # using the captured output via regex, but if that does not
+        # work, we use whatever's left over as the output (which is
+        # already set).
+        if output is not None and len(output) > 0:
+            reply = output
+
         # prevents empty thought process blocks when filter used with
         # malformed LLM output.
         if len(enclosure) > 0: