clean up code

app_modules/llm_inference.py

@@ -5,7 +5,6 @@ import urllib
 from queue import Queue
 from threading import Thread
 
-from langchain.callbacks.tracers import LangChainTracer
 from langchain.chains.base import Chain
 
 from app_modules.llm_loader import LLMLoader, TextIteratorStreamer
@@ -24,12 +23,8 @@ class LLMInference(metaclass=abc.ABCMeta):
     def create_chain(self) -> Chain:
         pass
 
-    def get_chain(self, tracing: bool = False) -> Chain:
+    def get_chain(self) -> Chain:
         if self.chain is None:
-            if tracing:
-                tracer = LangChainTracer()
-                tracer.load_default_session()
-
             self.chain = self.create_chain()
 
         return self.chain
@@ -39,7 +34,6 @@ class LLMInference(metaclass=abc.ABCMeta):
         inputs,
         streaming_handler,
         q: Queue = None,
-        tracing: bool = False,
         testing: bool = False,
     ):
         print(inputs)
@@ -49,7 +43,7 @@ class LLMInference(metaclass=abc.ABCMeta):
         try:
             self.llm_loader.streamer.reset(q)
 
-            chain = self.get_chain(tracing)
+            chain = self.get_chain()
             result = (
                 self._run_chain(chain, inputs, streaming_handler, testing)
                 if streaming_handler is not None
@@ -84,7 +78,7 @@ class LLMInference(metaclass=abc.ABCMeta):
         )
         t.start()
 
-        if self.llm_loader.streamer.for_huggingface and not testing:
+        if self.llm_loader.streamer.for_huggingface:
             count = (
                 2
                 if "chat_history" in inputs and len(inputs.get("chat_history")) > 0
@@ -94,12 +88,14 @@ class LLMInference(metaclass=abc.ABCMeta):
             while count > 0:
                 try:
                     for token in self.llm_loader.streamer:
-                        streaming_handler.on_llm_new_token(token)
+                        if not testing:
+                            streaming_handler.on_llm_new_token(token)
 
                     self.llm_loader.streamer.reset()
                     count -= 1
                 except Exception:
-                    print("nothing generated yet - retry in 0.5s")
+                    if not testing:
+                        print("nothing generated yet - retry in 0.5s")
                     time.sleep(0.5)
 
         t.join()

test.py

@@ -72,7 +72,6 @@ while True:
         {"question": query, "chat_history": chat_history},
         custom_handler,
         None,
-        False,
         True,
     )
     end = timer()