Update llm_response_generator.py

llm_response_generator.py

@@ -1,7 +1,6 @@
 import os
 from langchain_community.llms import HuggingFaceHub
 from langchain_community.llms import OpenAI
-# from langchain.llms import HuggingFaceHub, OpenAI
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
 import warnings
@@ -9,17 +8,19 @@ import warnings
 warnings.filterwarnings("ignore")
 
 class LLLResponseGenerator():
-
     def __init__(self):
-        print("initialized")
+        self.context = "You are a mental health supporting non-medical assistant. DO NOT PROVIDE any medical advice with conviction."
+        self.conversation_history = []
 
+    def update_context(self, user_text):
+        self.conversation_history.append(user_text)
+        self.context = "\n".join(self.conversation_history)
 
     def llm_inference(
         self,
         model_type: str,
         question: str,
         prompt_template: str,
-        context: str,
         ai_tone: str,
         questionnaire: str,
         user_text: str,
@@ -37,7 +38,6 @@ class LLLResponseGenerator():
             model_str: Denotes the LLM vendor's name. Can be either 'huggingface' or 'openai'
             question: The question to be asked to the LLM.
             prompt_template: The prompt template itself.
-            context: Instructions for the LLM.
             ai_tone: Can be either empathy, encouragement or suggest medical help.
             questionnaire: Can be either depression, anxiety or adhd.
             user_text: Response given by the user.
@@ -65,13 +65,12 @@ class LLLResponseGenerator():
         )
 
         if model_type == "openai":
-            # https://api.python.langchain.com/en/stable/llms/langchain.llms.openai.OpenAI.html#langchain.llms.openai.OpenAI
             llm = OpenAI(
                 model_name=openai_model_name, temperature=temperature, max_tokens=max_length
             )
             llm_chain = LLMChain(prompt=prompt, llm=llm)
             return llm_chain.run(
-                context=context,
+                context=self.context,
                 ai_tone=ai_tone,
                 questionnaire=questionnaire,
                 question=question,
@@ -79,15 +78,14 @@ class LLLResponseGenerator():
             )
 
         elif model_type == "huggingface":
-            # https://python.langchain.com/docs/integrations/llms/huggingface_hub
             llm = HuggingFaceHub(
                 repo_id=hf_repo_id,
                 model_kwargs={"temperature": temperature, "max_length": max_length},
             )
 
             llm_chain = LLMChain(prompt=prompt, llm=llm)
-            response =  llm_chain.run(
-                context=context,
+            response = llm_chain.run(
+                context=self.context,
                 ai_tone=ai_tone,
                 questionnaire=questionnaire,
                 question=question,
@@ -108,8 +106,6 @@ if __name__ == "__main__":
     # Please ensure you have a .env file available with 'HUGGINGFACEHUB_API_TOKEN' and 'OPENAI_API_KEY' values.
     HUGGINGFACEHUB_API_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')
 
-    context = "You are a mental health supporting non-medical assistant. DO NOT PROVIDE any medical advice with conviction."
-
     ai_tone = "EMPATHY"
     questionnaire = "ADHD"
     question = (
@@ -136,17 +132,32 @@ if __name__ == "__main__":
 
     model = LLLResponseGenerator()
 
-
-    llm_response = model.llm_inference(
+    # Initial prompt
+    print("Bot:", model.llm_inference(
         model_type="huggingface",
         question=question,
         prompt_template=template,
-        context=context,
         ai_tone=ai_tone,
         questionnaire=questionnaire,
         user_text=user_text,
         temperature=temperature,
         max_length=max_length,
-        )
-
-    print(llm_response)
+        ))
+
+    while True:
+        user_input = input("You: ")
+        if user_input.lower() == "exit":
+            break
+
+        model.update_context(user_input)
+
+        print("Bot:", model.llm_inference(
+            model_type="huggingface",
+            question=question,
+            prompt_template=template,
+            ai_tone=ai_tone,
+            questionnaire=questionnaire,
+            user_text=user_input,
+            temperature=temperature,
+            max_length=max_length,
+        ))