Update app.py

app.py

@@ -5,7 +5,6 @@ import time
 from typing import Dict, List, Tuple
 from datetime import datetime
 import logging
-
 import gradio as gr
 from huggingface_hub import InferenceClient
 from safe_search import safe_search
@@ -19,7 +18,6 @@ TEMPERATURE = 0.7
 TOP_P = 0.8
 REPETITION_PENALTY = 1.5
 MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-
 API_KEY = os.getenv("HUGGINGFACE_API_KEY")
 
 # --- Logging Setup ---
@@ -44,62 +42,54 @@ PREFIX = """
 Purpose: {purpose}
 Safe Search: {safe_search}
 """
-
 LOG_PROMPT = """
 PROMPT: {content}
 """
-
 LOG_RESPONSE = """
 RESPONSE: {resp}
 """
-
 COMPRESS_HISTORY_PROMPT = """
 You are a helpful AI assistant. Your task is to compress the following history into a summary that is no longer than 512 tokens.
-History:
-{history}
+History: {history}
 """
-
 ACTION_PROMPT = """
 You are a helpful AI assistant. You are working on the task: {task}
-Your current history is:
-{history}
+Your current history is: {history}
 What is your next thought?
-thought: 
+thought:
 What is your next action?
-action: 
+action:
 """
-
 TASK_PROMPT = """
-You are a helpful AI assistant. Your current history is:
-{history}
+You are a helpful AI assistant. Your current history is: {history}
 What is the next task?
-task: 
+task:
 """
-
 UNDERSTAND_TEST_RESULTS_PROMPT = """
-You are a helpful AI assistant. The test results are:
-{test_results}
+You are a helpful AI assistant. The test results are: {test_results}
 What do you want to know about the test results?
-thought: 
+thought:
 """
 
 # --- Functions ---
 def format_prompt(message: str, history: List[Tuple[str, str]], max_history_turns: int = 2) -> str:
+    """Formats the prompt for the LLM, including the message and recent history."""
     prompt = " "
     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_llm(
     prompt_template: str,
     stop_tokens: List[str],
     purpose: str,
-    **prompt_kwargs: Dict
+    **prompt_kwargs: Dict,
 ) -> str:
+    """Runs the LLM with the given prompt template, stop tokens, and purpose."""
     seed = random.randint(1, 1111111111111111)
     logging.info(f"Seed: {seed}")
-
     content = PREFIX.format(
         date_time_str=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
         purpose=purpose,
@@ -107,13 +97,20 @@ def run_llm(
     ) + prompt_template.format(**prompt_kwargs)
     if VERBOSE:
         logging.info(LOG_PROMPT.format(content=content))
-
     client = InferenceClient(model=MODEL_NAME, token=API_KEY)
-    resp = client.text_generation(content, max_new_tokens=MAX_TOKENS, stop_sequences=stop_tokens, temperature=TEMPERATURE, top_p=TOP_P, repetition_penalty=REPETITION_PENALTY)
+    resp = client.text_generation(
+        content,
+        max_new_tokens=MAX_TOKENS,
+        stop_sequences=stop_tokens,
+        temperature=TEMPERATURE,
+        top_p=TOP_P,
+        repetition_penalty=REPETITION_PENALTY,
+    )
     if VERBOSE:
         logging.info(LOG_RESPONSE.format(resp=resp))
     return resp
 
+
 def generate(
     prompt: str,
     history: List[Tuple[str, str]],
@@ -124,6 +121,7 @@ def generate(
     top_p: float = TOP_P,
     repetition_penalty: float = REPETITION_PENALTY,
 ) -> str:
+    """Generates a response from the LLM based on the prompt, history, and other parameters."""
     content = PREFIX.format(
         date_time_str=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
         purpose=f"Generating response as {agent_name}",
@@ -131,121 +129,240 @@ def generate(
     ) + sys_prompt + "\n" + prompt
     if VERBOSE:
         logging.info(LOG_PROMPT.format(content=content))
-
     client = InferenceClient(model=MODEL_NAME, token=API_KEY)
-    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)
-    return "".join(chunk.text for chunk in stream)
+    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,
+    )
+    response = "".join(chunk.text for chunk in stream)
+    return response
+
 
 def main():
+    """Main function to launch the Gradio interface."""
     with gr.Blocks() as demo:
         gr.Markdown("## FragMixt: The No-Code Development Powerhouse")
         gr.Markdown("### Your AI-Powered Development Companion")
-
         with gr.Row():
             with gr.Column(scale=3):
-                chatbot = gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel")
-                message = gr.Textbox(label="Enter your message", placeholder="Ask me anything!")
+                chatbot = gr.Chatbot(
+                    show_label=False,
+                    show_share_button=False,
+                    show_copy_button=True,
+                    likeable=True,
+                    layout="panel",
+                )
+                message = gr.Textbox(
+                    label="Enter your message", placeholder="Ask me anything!"
+                )
                 submit_button = gr.Button(value="Send")
-            
             with gr.Column(scale=1):
-                purpose = gr.Textbox(label="Purpose", placeholder="What is the purpose of this interaction?")
-                agent_name = gr.Dropdown(label="Agents", choices=[s for s in agents], value=agents[0], interactive=True)
-                sys_prompt = gr.Textbox(label="System Prompt", max_lines=1, interactive=True)
-                temperature = gr.Slider(label="Temperature", value=TEMPERATURE, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Higher values produce more diverse outputs")
-                max_new_tokens = gr.Slider(label="Max new tokens", value=MAX_TOKENS, minimum=0, maximum=1048*10, step=64, interactive=True, info="The maximum numbers of new tokens")
-                top_p = gr.Slider(label="Top-p (nucleus sampling)", value=TOP_P, minimum=0.0, maximum=1, step=0.05, interactive=True, info="Higher values sample more low-probability tokens")
-                repetition_penalty = gr.Slider(label="Repetition penalty", value=REPETITION_PENALTY, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Penalize repeated tokens")
-
+                purpose = gr.Textbox(
+                    label="Purpose", placeholder="What is the purpose of this interaction?"
+                )
+                agent_name = gr.Dropdown(
+                    label="Agents",
+                    choices=[s for s in agents],
+                    value=agents[0],
+                    interactive=True,
+                )
+                sys_prompt = gr.Textbox(
+                    label="System Prompt", max_lines=1, interactive=True
+                )
+                temperature = gr.Slider(
+                    label="Temperature",
+                    value=TEMPERATURE,
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    interactive=True,
+                    info="Higher values produce more diverse outputs",
+                )
+                max_new_tokens = gr.Slider(
+                    label="Max new tokens",
+                    value=MAX_TOKENS,
+                    minimum=0,
+                    maximum=1048 * 10,
+                    step=64,
+                    interactive=True,
+                    info="The maximum numbers of new tokens",
+                )
+                top_p = gr.Slider(
+                    label="Top-p (nucleus sampling)",
+                    value=TOP_P,
+                    minimum=0.0,
+                    maximum=1,
+                    step=0.05,
+                    interactive=True,
+                    info="Higher values sample more low-probability tokens",
+                )
+                repetition_penalty = gr.Slider(
+                    label="Repetition penalty",
+                    value=REPETITION_PENALTY,
+                    minimum=1.0,
+                    maximum=2.0,
+                    step=0.05,
+                    interactive=True,
+                    info="Penalize repeated tokens",
+                )
         with gr.Tabs():
             with gr.TabItem("Project Explorer"):
-                project_path = gr.Textbox(label="Project Path", placeholder="/home/user/app/current_project")
+                project_path = gr.Textbox(
+                    label="Project Path", placeholder="/home/user/app/current_project"
+                )
                 explore_button = gr.Button(value="Explore")
                 project_output = gr.Textbox(label="File Tree", lines=20)
-
             with gr.TabItem("Code Editor"):
                 code_editor = gr.Code(label="Code Editor", language="python")
                 run_code_button = gr.Button(value="Run Code")
                 code_output = gr.Textbox(label="Code Output", lines=10)
-
             with gr.TabItem("File Management"):
-                file_list = gr.Dropdown(label="Select File", choices=[], interactive=True)
+                file_list = gr.Dropdown(
+                    label="Select File", choices=[], interactive=True
+                )
                 file_content = gr.Textbox(label="File Content", lines=20)
                 save_file_button = gr.Button(value="Save File")
                 create_file_button = gr.Button(value="Create New File")
                 delete_file_button = gr.Button(value="Delete File")
-
         history = gr.State([])
 
-        def chat(purpose: str, message: str, agent_name: str, sys_prompt: str, temperature: float, max_new_tokens: int, top_p: float, repetition_penalty: float, history: List[Tuple[str, str]]) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
+        def chat(
+            purpose: str,
+            message: str,
+            agent_name: str,
+            sys_prompt: str,
+            temperature: float,
+            max_new_tokens: int,
+            top_p: float,
+            repetition_penalty: float,
+            history: List[Tuple[str, str]],
+        ) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
+            """Handles the chat interaction, generating responses and updating history."""
             prompt = format_prompt(message, history)
-            response = generate(prompt, history, agent_name, sys_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+            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])
+        submit_button.click(
+            chat,
+            inputs=[
+                purpose,
+                message,
+                agent_name,
+                sys_prompt,
+                temperature,
+                max_new_tokens,
+                top_p,
+                repetition_penalty,
+                history,
+            ],
+            outputs=[chatbot, history],
+        )
 
         def explore_project(project_path: str) -> str:
+            """Explores the project directory and displays the file tree."""
             try:
                 tree = subprocess.check_output(["tree", project_path]).decode("utf-8")
                 return tree
             except Exception as e:
                 return f"Error exploring project: {e}"
 
-        explore_button.click(explore_project, inputs=[project_path], outputs=[project_output])
+        explore_button.click(
+            explore_project, inputs=[project_path], outputs=[project_output]
+        )
 
         def run_code(code: str) -> str:
+            """Executes the Python code in the code editor and returns the output."""
             try:
                 exec_globals = {}
                 exec(code, exec_globals)
-                output = exec_globals.get('__builtins__', {}).get('print', print)
+                output = exec_globals.get("__builtins__", {}).get("print", print)
                 return str(output)
             except Exception as e:
                 return f"Error running code: {e}"
 
-        run_code_button.click(run_code, inputs=[code_editor], outputs=[code_output])
+        run_code_button.click(
+            run_code, inputs=[code_editor], outputs=[code_output]
+        )
 
         def load_file_list(project_path: str) -> List[str]:
+            """Loads the list of files in the project directory."""
             try:
-                return [f for f in os.listdir(project_path) if os.path.isfile(os.path.join(project_path, f))]
+                return [
+                    f
+                    for f in os.listdir(project_path)
+                    if os.path.isfile(os.path.join(project_path, f))
+                ]
             except Exception as e:
                 return [f"Error loading file list: {e}"]
 
         def load_file_content(project_path: str, file_name: str) -> str:
+            """Loads the content of the selected file."""
             try:
-                with open(os.path.join(project_path, file_name), 'r') as file:
+                with open(os.path.join(project_path, file_name), "r") as file:
                     return file.read()
             except Exception as e:
                 return f"Error loading file content: {e}"
 
         def save_file(project_path: str, file_name: str, content: str) -> str:
+            """Saves the content to the selected file."""
             try:
-                with open(os.path.join(project_path, file_name), 'w') as file:
+                with open(os.path.join(project_path, file_name), "w") as file:
                     file.write(content)
                 return f"File {file_name} saved successfully."
             except Exception as e:
                 return f"Error saving file: {e}"
 
         def create_file(project_path: str, file_name: str) -> str:
+            """Creates a new file in the project directory."""
             try:
-                open(os.path.join(project_path, file_name), 'a').close()
+                open(os.path.join(project_path, file_name), "a").close()
                 return f"File {file_name} created successfully."
             except Exception as e:
                 return f"Error creating file: {e}"
 
         def delete_file(project_path: str, file_name: str) -> str:
+            """Deletes the selected file from the project directory."""
             try:
                 os.remove(os.path.join(project_path, file_name))
                 return f"File {file_name} deleted successfully."
             except Exception as e:
                 return f"Error deleting file: {e}"
 
-        project_path.change(load_file_list, inputs=[project_path], outputs=[file_list])
-        file_list.change(load_file_content, inputs=[project_path, file_list], outputs=[file_content])
-        save_file_button.click(save_file, inputs=[project_path, file_list, file_content], outputs=[gr.Textbox()])
-        create_file_button.click(create_file, inputs=[project_path, gr.Textbox(label="New File Name")], outputs=[gr.Textbox()])
-        delete_file_button.click(delete_file, inputs=[project_path, file_list], outputs=[gr.Textbox()])
-
+        project_path.change(
+            load_file_list, inputs=[project_path], outputs=[file_list]
+        )
+        file_list.change(
+            load_file_content, inputs=[project_path, file_list], outputs=[file_content]
+        )
+        save_file_button.click(
+            save_file, inputs=[project_path, file_list, file_content], outputs=[gr.Textbox()]
+        )
+        create_file_button.click(
+            create_file,
+            inputs=[project_path, gr.Textbox(label="New File Name")],
+            outputs=[gr.Textbox()],
+        )
+        delete_file_button.click(
+            delete_file, inputs=[project_path, file_list], outputs=[gr.Textbox()]
+        )
     demo.launch()
 
+
 if __name__ == "__main__":
     main()