Update app.py

app.py

@@ -13,9 +13,7 @@ import json
 now = datetime.now()
 date_time_str = now.strftime("%Y-%m-%d %H:%M:%S")
 
-client = InferenceClient(
-    "mistralai/Mixtral-8x7B-Instruct-v0.1"
-)
+client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
 # --- Set up logging ---
 logging.basicConfig(
@@ -24,16 +22,14 @@ logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(message)s",
 )
 
-agents =[
+agents = [
     "WEB_DEV",
     "AI_SYSTEM_PROMPT",
     "PYTHON_CODE_DEV"
 ]
-############################################
 
 VERBOSE = True
 MAX_HISTORY = 5
-#MODEL = "gpt-3.5-turbo"  # "gpt-4"
 
 PREFIX = """
 {date_time_str}
@@ -80,13 +76,12 @@ thought:
 """
 
 def format_prompt(message, history, max_history_turns=2):
-  prompt = " "
-  # Keep only the last 'max_history_turns' turns
-  for user_prompt, bot_response in history[-max_history_turns:]:
-    prompt += f"[INST] {user_prompt} [/  "
-    prompt += f" {bot_response}"
-  prompt += f"[INST] {message} [/  "
-  return prompt
+    prompt = " "
+    # Keep only the last 'max_history_turns' turns
+    for user_prompt, bot_response in history[-max_history_turns:]:
+        prompt += f"[INST] {user_prompt} [/INST] {bot_response} "
+    prompt += f"[INST] {message} [/INST] "
+    return prompt
 
 def run_gpt(
     prompt_template,
@@ -95,9 +90,9 @@ def run_gpt(
     purpose,
     **prompt_kwargs,
 ):
-    seed = random.randint(1,1111111111111111)
+    seed = random.randint(1, 1111111111111111)
     logging.info(f"Seed: {seed}")  # Log the seed
-    
+
     content = PREFIX.format(
         date_time_str=date_time_str,
         purpose=purpose,
@@ -105,22 +100,31 @@ def run_gpt(
     ) + prompt_template.format(**prompt_kwargs)
     if VERBOSE:
         logging.info(LOG_PROMPT.format(content))  # Log the prompt
-    
+
     resp = client.text_generation(content, max_new_tokens=max_tokens, stop_sequences=stop_tokens, temperature=0.7, top_p=0.8, repetition_penalty=1.5)
     if VERBOSE:
         logging.info(LOG_RESPONSE.format(resp))  # Log the response
     return resp
 
-def generate( prompt, history, agent_name=agents[0], sys_prompt="", temperature=0.7, max_new_tokens=2048, top_p=0.8, repetition_penalty=1.5, ):
+def generate(
+    prompt,
+    history,
+    agent_name=agents[0],
+    sys_prompt="",
+    temperature=0.7,
+    max_new_tokens=2048,
+    top_p=0.8,
+    repetition_penalty=1.5,
+):
     content = PREFIX.format(
         date_time_str=date_time_str,
         purpose=purpose,
         safe_search=safe_search,
-    ) + prompt_template.format(**prompt_kwargs)
+    ) + prompt
     if VERBOSE:
         logging.info(LOG_PROMPT.format(content))  # Log the prompt
-    
-    stream = client.text_generation(content, **generate_kwargs, stream=True, details=True, return_full_text=False)
+
+    stream = client.text_generation(content, stream=True, details=True, return_full_text=False, temperature=temperature, top_p=top_p, repetition_penalty=repetition_penalty, max_new_tokens=max_new_tokens)
     resp = ""
     for response in stream:
         resp += response.token.text
@@ -129,7 +133,6 @@ def generate( prompt, history, agent_name=agents[0], sys_prompt="", temperature=
         logging.info(LOG_RESPONSE.format(resp))  # Log the response
     return resp
 
-
 def compress_history(purpose, task, history, directory):
     resp = run_gpt(
         COMPRESS_HISTORY_PROMPT,
@@ -141,19 +144,13 @@ def compress_history(purpose, task, history, directory):
     )
     history = "observation: {}\n".format(resp)
     return history
-    
+
 def call_search(purpose, task, history, directory, action_input):
     logging.info(f"CALLING SEARCH: {action_input}")
     try:
-
         if "http" in action_input:
-            if "<" in action_input:
-                action_input = action_input.strip("<")
-            if ">" in action_input:
-                action_input = action_input.strip(">")
-            
+            action_input = action_input.strip("<>").strip()
             response = i_s(action_input)
-            #response = google(search_return)
             logging.info(f"Search Result: {response}")
             history += "observation: search result is: {}\n".format(response)
         else:
@@ -166,13 +163,13 @@ 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:"],
+        stop_tokens=["observation:", "task:", "action:", "thought:"],
         max_tokens=32000,
         purpose=purpose,
         task=task,
         history=history,
     )
-    lines = resp.strip().strip("\n").split("\n")
+    lines = resp.strip().split("\n")
     for line in lines:
         if line == "":
             continue
@@ -180,7 +177,6 @@ def call_main(purpose, task, history, directory, action_input):
             history += "{}\n".format(line)
             logging.info(f"Thought: {line}")
         elif line.startswith("action: "):
-            
             action_name, action_input = parse_action(line)
             logging.info(f"Action: {action_name} - {action_input}")
             history += "{}\n".format(line)
@@ -192,13 +188,8 @@ def call_main(purpose, task, history, directory, action_input):
         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
-            #assert False, "unknown action: {}".format(line)
     return "MAIN", None, history, task
 
-
 def call_set_task(purpose, task, history, directory, action_input):
     logging.info(f"CALLING SET_TASK: {action_input}")
     task = run_gpt(
@@ -222,27 +213,26 @@ NAME_TO_FUNC = {
     "UPDATE-TASK": call_set_task,
     "SEARCH": call_search,
     "COMPLETE": end_fn,
-
 }
 
 def run_action(purpose, task, history, directory, action_name, action_input):
     logging.info(f"RUNNING ACTION: {action_name} - {action_input}")
     try:
         if "RESPONSE" in action_name or "COMPLETE" in action_name:
-            action_name="COMPLETE"
-            task="END"
+            action_name = "COMPLETE"
+            task = "END"
             return action_name, "COMPLETE", history, task
-    
+
         # compress the history when it is long
         if len(history.split("\n")) > MAX_HISTORY:
             logging.info("COMPRESSING HISTORY")
             history = compress_history(purpose, task, history, directory)
-        if not action_name in NAME_TO_FUNC:
-            action_name="MAIN"
-        if action_name == "" or action_name == None:
-            action_name="MAIN"
+        if action_name not in NAME_TO_FUNC:
+            action_name = "MAIN"
+        if action_name == "" or action_name is None:
+            action_name = "MAIN"
         assert action_name in NAME_TO_FUNC
-    
+
         logging.info(f"RUN: {action_name} - {action_input}")
         return NAME_TO_FUNC[action_name](purpose, task, history, directory, action_input)
     except Exception as e:
@@ -250,14 +240,11 @@ def run_action(purpose, task, history, directory, action_name, action_input):
         logging.error(f"Error in run_action: {e}")
         return "MAIN", None, history, task
 
-def run(purpose,history):
-    
-    #print(purpose)
-    #print(hist)
-    task=None
-    directory="./"
+def run(purpose, history):
+    task = None
+    directory = "./"
     if history:
-        history=str(history).strip("[]")
+        history = str(history).strip("[]")
     if not history:
         history = ""
     
@@ -280,23 +267,19 @@ def run(purpose,history):
             action_input,
         )
         yield (history)
-        #yield ("",[(purpose,history)])
         if task == "END":
             return (history)
-            #return ("", [(purpose,history)])
-
 
-
-################################################
-
-def format_prompt(message, history, max_history_turns=5):
-  prompt = " "
-  # Keep only the last 'max_history_turns' turns
-  for user_prompt, bot_response in history[-max_history_turns:]:
-    prompt += f"[INST] {user_prompt} [/  "
-    prompt += f" {bot_response}"
-  prompt += f"[INST] {message} [/  "
-  return prompt
+def parse_action(line):
+    """Parse the action line to get the action name and input."""
+    parts = line.split(":", 1)
+    if len(parts) == 2:
+        action_name = parts[0].replace("action", "").strip()
+        action_input = parts[1].strip()
+    else:
+        action_name = parts[0].replace("action", "").strip()
+        action_input = ""
+    return action_name, action_input
 
 def main():
     with gr.Blocks() as demo:
@@ -305,8 +288,7 @@ def main():
 
         # Chat Interface
         chatbot = gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel")
-        #chatbot.load(examples)
-
+        
         # Input Components
         message = gr.Textbox(label="Enter your message", placeholder="Ask me anything!")
         purpose = gr.Textbox(label="Purpose", placeholder="What is the purpose of this interaction?")
@@ -331,4 +313,18 @@ def main():
             history = gr.State([])
             examples = [
                 ["What is the purpose of this AI agent?", "I am designed to assist with no-code development tasks."],
-                ["Can you help me generate a Python function to calculate the factorial of a number?", "
+                ["Can you help me generate a Python function to calculate the factorial of a number?", "Sure! Here is a Python function to calculate the factorial of a number:"],
+            ]
+
+        def chat(purpose, message, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, history):
+            prompt = format_prompt(message, history)
+            response = generate(prompt, history, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            history.append((message, response))
+            return history, history
+
+        submit_button.click(chat, inputs=[purpose, message, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty, history], outputs=[chatbot, history])
+
+    demo.launch()
+
+if __name__ == "__main__":
+    main()