Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import subprocess
-import streamlit as st
+import streamlit as st  # For Streamlit
+import gradio as gr  # For Gradio
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import black
 from pylint import lint
@@ -12,7 +13,16 @@ import requests
 from bs4 import BeautifulSoup
 from typing import List, Dict, Optional
 
-# Define custom exceptions for better error handling
+from utils.code_utils import (
+    refine_code,
+    test_code,
+    integrate_code,
+    CodeRefinementError,
+    CodeTestingError,
+    CodeIntegrationError,
+)
+
+# --- Define custom exceptions for better error handling ---
 class InvalidActionError(Exception):
     """Raised when an invalid action is provided."""
     pass
@@ -25,18 +35,6 @@ class CodeGenerationError(Exception):
     """Raised when code generation fails."""
     pass
 
-class CodeRefinementError(Exception):
-    """Raised when code refinement fails."""
-    pass
-
-class CodeTestingError(Exception):
-    """Raised when code testing fails."""
-    pass
-
-class CodeIntegrationError(Exception):
-    """Raised when code integration fails."""
-    pass
-
 class AppTestingError(Exception):
     """Raised when app testing fails."""
     pass
@@ -53,14 +51,16 @@ class SearchError(Exception):
     """Raised when search fails."""
     pass
 
+# --- Define a class for the AI Agent ---
 class AIAgent:
     def __init__(self):
+        # --- Initialize tools and attributes ---
         self.tools = {
             "SEARCH": self.search,
             "CODEGEN": self.code_generation,
-            "REFINE-CODE": self.refine_code,
-            "TEST-CODE": self.test_code,
-            "INTEGRATE-CODE": self.integrate_code,
+            "REFINE-CODE": refine_code,  # Use external function
+            "TEST-CODE": test_code,  # Use external function
+            "INTEGRATE-CODE": integrate_code,  # Use external function
             "TEST-APP": self.test_app,
             "GENERATE-REPORT": self.generate_report,
             "WORKSPACE-EXPLORER": self.workspace_explorer,
@@ -82,10 +82,9 @@ class AIAgent:
         self.prompts: List[str] = []  # Store prompts for future use
         self.code_generator = pipeline('text-generation', model='gpt2')  # Initialize code generator
 
+    # --- Implement search functionality ---
     def search(self, query: str) -> List[str]:
-        """
-        Performs a web search using the specified search engine.
-        """
+        """Performs a web search using the specified search engine."""
         search_url = self.search_engine_url + query
         try:
             response = requests.get(search_url)
@@ -96,68 +95,27 @@ class AIAgent:
         except requests.exceptions.RequestException as e:
             raise SearchError(f"Error during search: {e}")
 
+    # --- Implement code generation functionality ---
     def code_generation(self, snippet: str) -> str:
-        """
-        Generates code based on the provided snippet.
-        """
+        """Generates code based on the provided snippet."""
         try:
             generated_text = self.code_generator(snippet, max_length=500, num_return_sequences=1)[0]['generated_text']
             return generated_text
         except Exception as e:
             raise CodeGenerationError(f"Error during code generation: {e}")
 
-    def refine_code(self, file_path: str) -> str:
-        """
-        Refines the code in the specified file.
-        """
-        try:
-            with open(file_path, 'r') as f:
-                code = f.read()
-            refined_code = black.format_str(code, mode=black.FileMode())
-            return refined_code
-        except black.InvalidInput:
-            raise CodeRefinementError("Error: Invalid code input for black formatting.")
-        except Exception as e:
-            raise CodeRefinementError(f"Error during code refinement: {e}")
-
-    def test_code(self, file_path: str) -> str:
-        """
-        Tests the code in the specified file.
-        """
-        try:
-            with open(file_path, 'r') as f:
-                code = f.read()
-            output = StringIO()
-            lint.run(code, output=output)
-            return output.getvalue()
-        except Exception as e:
-            raise CodeTestingError(f"Error during code testing: {e}")
-
-    def integrate_code(self, file_path: str, code_snippet: str) -> str:
-        """
-        Integrates the code snippet into the specified file.
-        """
-        try:
-            with open(file_path, 'a') as f:
-                f.write(code_snippet)
-            return "Code integrated successfully."
-        except Exception as e:
-            raise CodeIntegrationError(f"Error during code integration: {e}")
-
+    # --- Implement app testing functionality ---
     def test_app(self) -> str:
-        """
-        Tests the functionality of the app.
-        """
+        """Tests the functionality of the app."""
         try:
             subprocess.run(['streamlit', 'run', 'app.py'], check=True)
             return "App tested successfully."
         except subprocess.CalledProcessError as e:
             raise AppTestingError(f"Error during app testing: {e}")
 
+    # --- Implement report generation functionality ---
     def generate_report(self) -> str:
-        """
-        Generates a report based on the task history.
-        """
+        """Generates a report based on the task history."""
         report = f"## Task Report: {self.current_task}\n\n"
         for task in self.task_history:
             report += f"**Action:** {task['action']}\n"
@@ -165,10 +123,9 @@ class AIAgent:
             report += f"**Output:** {task['output']}\n\n"
         return report
 
+    # --- Implement workspace exploration functionality ---
     def workspace_explorer(self) -> str:
-        """
-        Provides a workspace explorer functionality.
-        """
+        """Provides a workspace explorer functionality."""
         try:
             current_directory = os.getcwd()
             directories = []
@@ -183,20 +140,18 @@ class AIAgent:
         except Exception as e:
             raise WorkspaceExplorerError(f"Error during workspace exploration: {e}")
 
+    # --- Implement prompt management functionality ---
     def add_prompt(self, prompt: str) -> str:
-        """
-        Adds a new prompt to the agent's knowledge base.
-        """
+        """Adds a new prompt to the agent's knowledge base."""
         try:
             self.prompts.append(prompt)
             return f"Prompt '{prompt}' added successfully."
         except Exception as e:
             raise PromptManagementError(f"Error adding prompt: {e}")
 
+    # --- Implement prompt generation functionality ---
     def action_prompt(self, action: str) -> str:
-        """
-        Provides a prompt for a specific action.
-        """
+        """Provides a prompt for a specific action."""
         try:
             if action == "SEARCH":
                 return "What do you want to search for?"
@@ -241,29 +196,25 @@ class AIAgent:
         except InvalidActionError as e:
             raise e
 
+    # --- Implement prompt generation functionality ---
     def compress_history_prompt(self) -> str:
-        """
-        Provides a prompt to compress the task history.
-        """
+        """Provides a prompt to compress the task history."""
         return "Do you want to compress the task history?"
 
+    # --- Implement prompt generation functionality ---
     def log_prompt(self) -> str:
-        """
-        Provides a prompt to log a specific event.
-        """
+        """Provides a prompt to log a specific event."""
         return "What event do you want to log?"
 
+    # --- Implement logging functionality ---
     def log_response(self, event: str) -> str:
-        """
-        Logs the specified event.
-        """
+        """Logs the specified event."""
         print(f"Event logged: {event}")
         return "Event logged successfully."
 
+    # --- Implement prompt modification functionality ---
     def modify_prompt(self, prompt: str) -> str:
-        """
-        Modifies an existing prompt.
-        """
+        """Modifies an existing prompt."""
         try:
             # Find the prompt to modify
             # Update the prompt
@@ -271,46 +222,42 @@ class AIAgent:
         except Exception as e:
             raise PromptManagementError(f"Error modifying prompt: {e}")
 
+    # --- Implement prefix functionality ---
     def prefix(self, text: str) -> str:
-        """
-        Adds a prefix to the provided text.
-        """
+        """Adds a prefix to the provided text."""
         return f"PREFIX: {text}"
 
+    # --- Implement search query generation functionality ---
     def search_query(self, query: str) -> str:
-        """
-        Provides a search query for the specified topic.
-        """
+        """Provides a search query for the specified topic."""
         return f"Search query: {query}"
 
+    # --- Implement file reading functionality ---
     def read_prompt(self, file_path: str) -> str:
-        """
-        Provides a prompt to read the contents of a file.
-        """
+        """Provides a prompt to read the contents of a file."""
         try:
             with open(file_path, 'r') as f:
                 contents = f.read()
             return contents
+        except FileNotFoundError:
+            raise InvalidInputError(f"Error: File not found: {file_path}")
         except Exception as e:
             raise InvalidInputError(f"Error reading file: {e}")
 
+    # --- Implement task prompt generation functionality ---
     def task_prompt(self) -> str:
-        """
-        Provides a prompt to start a new task.
-        """
+        """Provides a prompt to start a new task."""
         return "What task do you want to start?"
 
+    # --- Implement test results understanding prompt generation functionality ---
     def understand_test_results_prompt(self) -> str:
-        """
-        Provides a prompt to understand the test results.
-        """
+        """Provides a prompt to understand the test results."""
         return "What do you want to know about the test results?"
 
+    # --- Implement input handling functionality ---
     def handle_input(self, input_str: str):
-        """
-        Handles user input and executes the corresponding action.
-        """
-        if input_str:
+        """Handles user input and executes the corresponding action."""
+        try:
             action, *args = input_str.split()
             if action in self.tools:
                 if args:
@@ -328,22 +275,20 @@ class AIAgent:
                 print(f"Action: {action}\nInput: {' '.join(args)}\nOutput: {self.tools[action](' '.join(args))}")
             else:
                 raise InvalidActionError("Invalid action. Please choose a valid action from the list of tools.")
-        else:
-            # Handle empty input (e.g., display a message or do nothing)
-            print("Please provide input.")
-        except (InvalidActionError, InvalidInputError, CodeGenerationError, CodeRefinementError,
-                CodeTestingError, CodeIntegrationError, AppTestingError, WorkspaceExplorerError,
-                PromptManagementError, SearchError) as e: 
+        except (InvalidActionError, InvalidInputError, CodeGenerationError, 
+                CodeRefinementError, CodeTestingError, CodeIntegrationError, 
+                AppTestingError, WorkspaceExplorerError, PromptManagementError, 
+                SearchError) as e:
             print(f"Error: {e}")
 
+    # --- Implement the main loop of the agent ---
     def run(self):
-        """
-        Runs the agent continuously, waiting for user input.
-        """
+        """Runs the agent continuously, waiting for user input."""
         while True:
             input_str = input("Enter a command for the AI Agent: ")
             self.handle_input(input_str)
 
+# --- Streamlit Integration ---
 if __name__ == '__main__':
     agent = AIAgent()
     st.title("AI Agent")