From 921630e53078f46e53bc337c87ba26a019bf5221 Mon Sep 17 00:00:00 2001
From: M <mailinator@mailinator.com>
Date: Mon, 7 Jul 2025 19:00:08 +0200
Subject: [PATCH] Remove thinking from the output

---
 knowledge_storm/lm.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/knowledge_storm/lm.py b/knowledge_storm/lm.py
index d5836f72..37620b01 100644
--- a/knowledge_storm/lm.py
+++ b/knowledge_storm/lm.py
@@ -4,6 +4,7 @@
 import logging
 import os
 import random
+import re
 import requests
 import threading
 from typing import Optional, Literal, Any
@@ -265,6 +266,10 @@ def __call__(self, prompt=None, messages=None, **kwargs):
         )
         self.history.append(entry)
 
+        # remove everything within <think></think>
+        outputs = [re.sub("<think>.*?(</think>|$)", " ", output, flags=re.DOTALL)
+            for output in outputs]
+
         return outputs