fixed bug

app_modules/llm_inference.py

@@ -69,41 +69,40 @@ class LLMInference(metaclass=abc.ABCMeta):
         return result
 
     def _execute_chain(self, chain, inputs, q, sh):
-        self.llm_loader.lock.acquire()
-        try:
-            q.put(chain(inputs, callbacks=[sh]))
-        finally:
-            # Release the lock
-            self.llm_loader.lock.release()
+        q.put(chain(inputs, callbacks=[sh]))
 
     def _run_chain(self, chain, inputs, streaming_handler):
-        que = Queue()
-
-        t = Thread(
-            target=self._execute_chain,
-            args=(chain, inputs, que, streaming_handler),
-        )
-        t.start()
+        self.llm_loader.lock.acquire()
+        try:
+            que = Queue()
 
-        if self.llm_loader.streamer is not None and isinstance(
-            self.llm_loader.streamer, TextIteratorStreamer
-        ):
-            count = (
-                2
-                if "chat_history" in inputs and len(inputs.get("chat_history")) > 0
-                else 1
+            t = Thread(
+                target=self._execute_chain,
+                args=(chain, inputs, que, streaming_handler),
             )
-
-            while count > 0:
-                try:
-                    for token in self.llm_loader.streamer:
-                        streaming_handler.on_llm_new_token(token)
-
-                    self.llm_loader.streamer.reset()
-                    count -= 1
-                except Exception:
-                    print("nothing generated yet - retry in 0.5s")
-                    time.sleep(0.5)
-
-        t.join()
-        return que.get()
+            t.start()
+
+            if self.llm_loader.streamer is not None and isinstance(
+                self.llm_loader.streamer, TextIteratorStreamer
+            ):
+                count = (
+                    2
+                    if "chat_history" in inputs and len(inputs.get("chat_history")) > 0
+                    else 1
+                )
+
+                while count > 0:
+                    try:
+                        for token in self.llm_loader.streamer:
+                            streaming_handler.on_llm_new_token(token)
+
+                        self.llm_loader.streamer.reset()
+                        count -= 1
+                    except Exception:
+                        print("nothing generated yet - retry in 0.5s")
+                        time.sleep(0.5)
+
+            t.join()
+            return que.get()
+        finally:
+            self.llm_loader.lock.release()