optim v2: separated response for each component, mod becomes a classifier

allofresh_chatbot.py

@@ -1,6 +1,7 @@
 import os
 from dotenv import load_dotenv
 import re
+from loguru import logger
 
 from langchain import PromptTemplate, LLMChain
 from langchain.agents import initialize_agent, Tool
@@ -13,7 +14,7 @@ from langchain.llms import AzureOpenAI
 from langchain.prompts import PromptTemplate
 
 from utils import lctool_search_allo_api, cut_dialogue_history
-from prompts.mod_prompt import MOD_PROMPT, FALLBACK_MESSAGE
+from prompts.mod_prompt import MOD_PROMPT, FALLBACK_MESSAGE, MOD_PROMPT_OPTIM_v2
 from prompts.ans_prompt import ANS_PREFIX, ANS_FORMAT_INSTRUCTIONS, ANS_SUFFIX, ANS_CHAIN_PROMPT
 from prompts.reco_prompt import RECO_PREFIX, RECO_FORMAT_INSTRUCTIONS, RECO_SUFFIX, NO_RECO_OUTPUT
 
@@ -69,7 +70,7 @@ class AllofreshChatbot():
     
     def init_mod_chain(self):
         mod_prompt = PromptTemplate(
-            template=MOD_PROMPT,
+            template=MOD_PROMPT_OPTIM_v2,
             input_variables=["input"]
         )
 
@@ -97,7 +98,6 @@ class AllofreshChatbot():
             self.llms["gpt-4"],
             agent="conversational-react-description",
             verbose=self.debug,
-            memory=self.ans_memory,
             return_intermediate_steps=True,
             agent_kwargs={
                 'prefix': ANS_PREFIX, 
@@ -114,24 +114,6 @@ class AllofreshChatbot():
 
         # Define the first LLM chain with the shared AzureOpenAI object and prompt template
         return LLMChain(llm=self.llms["gpt-4"], prompt=ans_prompt)
-    
-    def ans_pipeline(self, text, debug_cost=False, keep_last_n_words=500):
-        try:
-            self.ans_agent.memory.buffer = cut_dialogue_history(self.ans_agent.memory.buffer, keep_last_n_words=keep_last_n_words)
-        except:
-            pass
-        finally:
-            with get_openai_callback() as openai_cb:
-                res = self.ans_agent({"input": text.strip()})
-                response = res['output'].replace("\\", "/")
-                
-                if debug_cost:
-                    print(f"Total Tokens: {openai_cb.total_tokens}")
-                    print(f"Prompt Tokens: {openai_cb.prompt_tokens}")
-                    print(f"Completion Tokens: {openai_cb.completion_tokens}")
-                    print(f"Total Cost (USD): ${openai_cb.total_cost}")
-
-                return response
         
     def init_reco_agent(self):
         reco_tools = [
@@ -195,10 +177,26 @@ class AllofreshChatbot():
             return self.ans_chain.run({"input": query, "chat_history": str(chat_history)})
         return FALLBACK_MESSAGE
     
-    def reco_optim_v1(self, chat_history): 
+    def answer_optim_v2(self, query, chat_history):
         """
         We plugged off the tools from the 'answering' component and replaced it with a simple chain
         """
+        # moderate
+        mod_verdict = self.mod_chain.run({"input": query})
+        llm_input = {"input": query, "chat_history": str(chat_history)}
+
+        logger.info(f"mod verdict: {mod_verdict}")
+        # if no need to access knowledge base
+        if mod_verdict == "ANS_CHAIN":
+            # answer question
+            return self.ans_chain.run(llm_input)
+        # if need to access knowledge base
+        elif mod_verdict == "ANS_AGENT":
+            res = self.ans_agent(llm_input)
+            return res['output'].replace("\\", "/")
+        return FALLBACK_MESSAGE
+
+    def reco_optim_v1(self, chat_history): 
         reco = self.reco_agent.run({"input": chat_history})
         # filter out reco (str) to only contain alphabeticals
         return reco if reco != NO_RECO_OUTPUT else None

app.py

@@ -40,14 +40,14 @@ def predict_answer(input, history):
     """
     Answering component
     """
-    answer = allo_chatbot.answer_optim_v1(input, preproc_history(history))
+    answer = allo_chatbot.answer_optim_v2(input, preproc_history(history))
 
     history.append({'role': 'user', 'content': None})
     history.append({'role': 'assistant', 'content': answer})
 
     return fetch_messages(history), history
 
-def predict_reco(input, history):
+def predict_reco(history):
     """
     Reco component
     """
@@ -73,7 +73,7 @@ with gr.Blocks() as app:
     ).success(
         predict_answer, [txt, state], [chatbot, state]
     ).success(
-        predict_reco, [txt, state], [chatbot, state]
+        predict_reco, [state], [chatbot, state]
     )
 
 app.queue(concurrency_count=4)

prompts/mod_prompt.py

@@ -15,7 +15,7 @@ NOTES:
 
 ONLY ANSWER with either [True, False]
 """
-MOD_PROMPT_OPTIM_v1 = """
+MOD_PROMPT_OPTIM_v2 = """
 You are MODERATOR. 
 You are to classify what the next Chatbot will be doing. The chatbot will assist in supermarket shopping and requires validation for any information relating to the supermarket