Update app.py

app.py

@@ -7,6 +7,8 @@ from safe_search import safe_search
 from i_search import google
 from i_search import i_search as i_s
 from datetime import datetime
+import logging
+
 now = datetime.now()
 date_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
 
@@ -14,8 +16,14 @@ client = InferenceClient(
     "mistralai/Mixtral-8x7B-Instruct-v0.1"
 )
 
-############################################
+# --- Set up logging ---
+logging.basicConfig(
+    filename="app.log",  # Name of the log file
+    level=logging.INFO,  # Set the logging level (INFO, DEBUG, etc.)
+    format="%(asctime)s - %(levelname)s - %(message)s",
+)
 
+############################################
 
 VERBOSE = True
 MAX_HISTORY = 100
@@ -72,16 +80,15 @@ What do you want to know about the test results?
 thought: 
 """
 
-def format_prompt(message, history):
+def format_prompt(message, history, max_history_turns=5):
   prompt = "<s>"
-  for user_prompt, bot_response in history:
+  # Keep only the last 'max_history_turns' turns
+  for user_prompt, bot_response in history[-max_history_turns:]:
     prompt += f"[INST] {user_prompt} [/INST]"
     prompt += f" {bot_response}</s> "
   prompt += f"[INST] {message} [/INST]"
   return prompt
 
-
-
 def run_gpt(
     prompt_template,
     stop_tokens,
@@ -90,7 +97,7 @@ def run_gpt(
     **prompt_kwargs,
 ):
     seed = random.randint(1,1111111111111111)
-    print (seed)
+    logging.info(f"Seed: {seed}")  # Log the seed
     generate_kwargs = dict(
         temperature=1.0,
         max_new_tokens=2096,
@@ -100,26 +107,21 @@ def run_gpt(
         seed=seed,
     )
 
-    
     content = PREFIX.format(
         date_time_str=date_time_str,
         purpose=purpose,
         safe_search=safe_search,
     ) + prompt_template.format(**prompt_kwargs)
     if VERBOSE:
-        print(LOG_PROMPT.format(content))
-    
+        logging.info(LOG_PROMPT.format(content))  # Log the prompt
     
-    #formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
-    #formatted_prompt = format_prompt(f'{content}', history)
-
     stream = client.text_generation(content, **generate_kwargs, stream=True, details=True, return_full_text=False)
     resp = ""
     for response in stream:
         resp += response.token.text
 
     if VERBOSE:
-        print(LOG_RESPONSE.format(resp))
+        logging.info(LOG_RESPONSE.format(resp))  # Log the response
     return resp
 
 
@@ -136,7 +138,7 @@ def compress_history(purpose, task, history, directory):
     return history
     
 def call_search(purpose, task, history, directory, action_input):
-    print("CALLING SEARCH")
+    logging.info(f"CALLING SEARCH: {action_input}")
     try:
         
         if "http" in action_input:
@@ -147,7 +149,7 @@ def call_search(purpose, task, history, directory, action_input):
             
             response = i_s(action_input)
             #response = google(search_return)
-            print(response)
+            logging.info(f"Search Result: {response}")
             history += "observation: search result is: {}\n".format(response)
         else:
             history += "observation: I need to provide a valid URL to 'action: SEARCH action_input=https://URL'\n"
@@ -156,6 +158,7 @@ def call_search(purpose, task, history, directory, action_input):
     return "MAIN", None, history, task
 
 def call_main(purpose, task, history, directory, action_input):
+    logging.info(f"CALLING MAIN: {action_input}")
     resp = run_gpt(
         ACTION_PROMPT,
         stop_tokens=["observation:", "task:", "action:","thought:"],
@@ -170,12 +173,11 @@ def call_main(purpose, task, history, directory, action_input):
             continue
         if line.startswith("thought: "):
             history += "{}\n".format(line)
+            logging.info(f"Thought: {line}")
         elif line.startswith("action: "):
             
             action_name, action_input = parse_action(line)
-            print (f'ACTION_NAME :: {action_name}')
-            print (f'ACTION_INPUT :: {action_input}')
-            
+            logging.info(f"Action: {action_name} - {action_input}")
             history += "{}\n".format(line)
             if "COMPLETE" in action_name or "COMPLETE" in action_input:
                 task = "END"
@@ -184,6 +186,7 @@ def call_main(purpose, task, history, directory, action_input):
                 return action_name, action_input, history, task
         else:
             history += "{}\n".format(line)
+            logging.info(f"Other Output: {line}")
             #history += "observation: the following command did not produce any useful output: '{}', I need to check the commands syntax, or use a different command\n".format(line)
             
             #return action_name, action_input, history, task
@@ -192,6 +195,7 @@ def call_main(purpose, task, history, directory, action_input):
 
 
 def call_set_task(purpose, task, history, directory, action_input):
+    logging.info(f"CALLING SET_TASK: {action_input}")
     task = run_gpt(
         TASK_PROMPT,
         stop_tokens=[],
@@ -204,6 +208,7 @@ def call_set_task(purpose, task, history, directory, action_input):
     return "MAIN", None, history, task
 
 def end_fn(purpose, task, history, directory, action_input):
+    logging.info(f"CALLING END_FN: {action_input}")
     task = "END"
     return "COMPLETE", "COMPLETE", history, task
 
@@ -216,7 +221,7 @@ NAME_TO_FUNC = {
 }
 
 def run_action(purpose, task, history, directory, action_name, action_input):
-    print(f'action_name::{action_name}')
+    logging.info(f"RUNNING ACTION: {action_name} - {action_input}")
     try:
         if "RESPONSE" in action_name or "COMPLETE" in action_name:
             action_name="COMPLETE"
@@ -225,8 +230,7 @@ def run_action(purpose, task, history, directory, action_name, action_input):
     
         # compress the history when it is long
         if len(history.split("\n")) > MAX_HISTORY:
-            if VERBOSE:
-                print("COMPRESSING HISTORY")
+            logging.info("COMPRESSING HISTORY")
             history = compress_history(purpose, task, history, directory)
         if not action_name in NAME_TO_FUNC:
             action_name="MAIN"
@@ -234,11 +238,11 @@ def run_action(purpose, task, history, directory, action_name, action_input):
             action_name="MAIN"
         assert action_name in NAME_TO_FUNC
     
-        print("RUN: ", action_name, action_input)
+        logging.info(f"RUN: {action_name} - {action_input}")
         return NAME_TO_FUNC[action_name](purpose, task, history, directory, action_input)
     except Exception as e:
         history += "observation: the previous command did not produce any useful output, I need to check the commands syntax, or use a different command\n"
-
+        logging.error(f"Error in run_action: {e}")
         return "MAIN", None, history, task
 
 def run(purpose,history):
@@ -255,14 +259,12 @@ def run(purpose,history):
     action_name = "UPDATE-TASK" if task is None else "MAIN"
     action_input = None
     while True:
-        print("")
-        print("")
-        print("---")
-        print("purpose:", purpose)
-        print("task:", task)
-        print("---")
-        print(history)
-        print("---")
+        logging.info(f"---")
+        logging.info(f"Purpose: {purpose}")
+        logging.info(f"Task: {task}")
+        logging.info(f"---")
+        logging.info(f"History: {history}")
+        logging.info(f"---")
 
         action_name, action_input, history, task = run_action(
             purpose,
@@ -282,9 +284,10 @@ def run(purpose,history):
 
 ################################################
 
-def format_prompt(message, history):
+def format_prompt(message, history, max_history_turns=5):
   prompt = "<s>"
-  for user_prompt, bot_response in history:
+  # Keep only the last 'max_history_turns' turns
+  for user_prompt, bot_response in history[-max_history_turns:]:
     prompt += f"[INST] {user_prompt} [/INST]"
     prompt += f" {bot_response}</s> "
   prompt += f"[INST] {message} [/INST]"
@@ -322,6 +325,7 @@ def generate(
     )
 
     formatted_prompt = format_prompt(prompt, history, max_history_turns=5)  # Truncated history
+    logging.info(f"Formatted Prompt: {formatted_prompt}")
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
@@ -331,7 +335,6 @@ def generate(
     return output
 
 
-
 additional_inputs=[
     gr.Dropdown(
         label="Agents",