Update app.py

app.py

@@ -17,30 +17,51 @@ class InvalidActionError(Exception):
     """Raised when an invalid action is provided."""
     pass
 
+
 class InvalidInputError(Exception):
     """Raised when invalid input is provided for an action."""
     pass
 
+
 class CodeGenerationError(Exception):
     """Raised when code generation fails."""
     pass
 
+
 class AppTestingError(Exception):
     """Raised when app testing fails."""
     pass
 
+
 class WorkspaceExplorerError(Exception):
     """Raised when workspace exploration fails."""
     pass
 
+
 class PromptManagementError(Exception):
     """Raised when prompt management fails."""
     pass
 
+
 class SearchError(Exception):
     """Raised when search 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
+
 # --- AI Agent Class ---
 class AIAgent:
     def __init__(self):
@@ -48,9 +69,9 @@ class AIAgent:
         self.tools = {
             "SEARCH": self.search,
             "CODEGEN": self.code_generation,
-            "REFINE-CODE": self.refine_code,  # Implemented within the class
-            "TEST-CODE": self.test_code,  # Implemented within the class
-            "INTEGRATE-CODE": self.integrate_code,  # Implemented within the class
+            "REFINE-CODE": self.refine_code,  # Use internal function
+            "TEST-CODE": self.test_code,  # Use internal function
+            "INTEGRATE-CODE": self.integrate_code,  # Use internal function
             "TEST-APP": self.test_app,
             "GENERATE-REPORT": self.generate_report,
             "WORKSPACE-EXPLORER": self.workspace_explorer,
@@ -89,17 +110,17 @@ class AIAgent:
     def code_generation(self, snippet: str) -> str:
         """Generates code based on the provided snippet or description."""
         try:
-            generated_text = self.code_generator(snippet, max_length=500, num_return_sequences=1)[0]['generated_text']
+            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}")
 
     # --- Code Refinement Functionality ---
-    def refine_code(self, file_path: str) -> str:
-        """Refines the code in the specified file."""
+    def refine_code(self, code: str) -> str:
+        """Refines the provided code string."""
         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:
@@ -108,13 +129,15 @@ class AIAgent:
             raise CodeRefinementError(f"Error during code refinement: {e}")
 
     # --- Code Testing Functionality ---
-    def test_code(self, file_path: str) -> str:
-        """Tests the code in the specified file."""
+    def test_code(self, code: str) -> str:
+        """Tests the provided code string using pylint."""
         try:
-            results = lint.run(file_path, do_exit=False)
-            report = StringIO()
-            results.report(report=report)
-            return report.getvalue()
+            # Use pylint to lint the code
+            lint_output = StringIO()
+            sys.stdout = lint_output
+            lint.Run(code.split('\n'), do_exit=False)
+            sys.stdout = sys.__stdout__
+            return lint_output.getvalue()
         except Exception as e:
             raise CodeTestingError(f"Error during code testing: {e}")
 
@@ -122,12 +145,11 @@ class AIAgent:
     def integrate_code(self, file_path: str, code_snippet: str) -> str:
         """Integrates the code snippet into the specified file."""
         try:
-            with open(file_path, 'r') as f:
-                existing_code = f.read()
-            integrated_code = existing_code + "\n" + code_snippet
-            with open(file_path, 'w') as f:
-                f.write(integrated_code)
-            return f"Code snippet integrated into {file_path} successfully."
+            # For simplicity, we'll just append the code snippet to the file
+            # In a real scenario, you'd need more sophisticated logic
+            with open(file_path, 'a') as f:
+                f.write(code_snippet)
+            return f"Code snippet integrated into {file_path}"
         except Exception as e:
             raise CodeIntegrationError(f"Error during code integration: {e}")
 
@@ -185,11 +207,11 @@ class AIAgent:
             elif action == "CODEGEN":
                 return "Provide a code snippet to generate code from, or describe what you want the code to do."
             elif action == "REFINE-CODE":
-                return "Provide the file path of the code to refine."
+                return "Provide the code to refine."
             elif action == "TEST-CODE":
-                return "Provide the file path of the code to test."
+                return "Provide the code to test."
             elif action == "INTEGRATE-CODE":
-                return "Provide the file path and code snippet to integrate."
+                return "Provide the file path and code snippet to integrate. For example: /path/to/your/file.py \"\"\"print('Hello, World!')\"\"\""
             elif action == "TEST-APP":
                 return "Test the application."
             elif action == "GENERATE-REPORT":
@@ -289,40 +311,66 @@ class AIAgent:
             if action in self.tools:
                 if args:
                     try:
-                        self.task_history.append({
-                            "action": action,
-                            "input": " ".join(args),
-                            "output": self.tools[action](" ".join(args))
-                        })
-                        print(f"Action: {action}\nInput: {' '.join(args)}\nOutput: {self.tools[action](' '.join(args))}")
+                        self.task_history.append(
+                            {
+                                "action": action,
+                                "input": " ".join(args),
+                                "output": self.tools[action](" ".join(args)),
+                            }
+                        )
+                        print(
+                            f"Action: {action}\nInput: {' '.join(args)}\nOutput: {self.tools[action](' '.join(args))}"
+                        )
                     except Exception as e:
-                        self.task_history.append({
-                            "action": action,
-                            "input": " ".join(args),
-                            "output": f"Error: {e}"
-                        })
-                        print(f"Action: {action}\nInput: {' '.join(args)}\nOutput: Error: {e}")
+                        self.task_history.append(
+                            {
+                                "action": action,
+                                "input": " ".join(args),
+                                "output": f"Error: {e}",
+                            }
+                        )
+                        print(
+                            f"Action: {action}\nInput: {' '.join(args)}\nOutput: Error: {e}"
+                        )
                 else:
                     try:
-                        self.task_history.append({
-                            "action": action,
-                            "input": None,
-                            "output": self.tools[action]()
-                        })
-                        print(f"Action: {action}\nInput: None\nOutput: {self.tools[action]()}")
+                        self.task_history.append(
+                            {
+                                "action": action,
+                                "input": None,
+                                "output": self.tools[action](),
+                            }
+                        )
+                        print(
+                            f"Action: {action}\nInput: None\nOutput: {self.tools[action]()}"
+                        )
                     except Exception as e:
-                        self.task_history.append({
-                            "action": action,
-                            "input": None,
-                            "output": f"Error: {e}"
-                        })
-                        print(f"Action: {action}\nInput: None\nOutput: Error: {e}")
+                        self.task_history.append(
+                            {
+                                "action": action,
+                                "input": None,
+                                "output": f"Error: {e}",
+                            }
+                        )
+                        print(
+                            f"Action: {action}\nInput: None\nOutput: Error: {e}"
+                        )
             else:
-                raise InvalidActionError("Invalid action. Please choose a valid action from the list of tools.")
-        except (InvalidActionError, InvalidInputError, CodeGenerationError, 
-                CodeRefinementError, CodeTestingError, CodeIntegrationError, 
-                AppTestingError, WorkspaceExplorerError, PromptManagementError, 
-                SearchError) as e:
+                raise InvalidActionError(
+                    "Invalid action. Please choose a valid action from the list of tools."
+                )
+        except (
+            InvalidActionError,
+            InvalidInputError,
+            CodeGenerationError,
+            CodeRefinementError,
+            CodeTestingError,
+            CodeIntegrationError,
+            AppTestingError,
+            WorkspaceExplorerError,
+            PromptManagementError,
+            SearchError,
+        ) as e:
             print(f"Error: {e}")
 
     # --- Main Loop of the Agent ---
@@ -332,31 +380,34 @@ class AIAgent:
             input_str = input("Enter a command for the AI Agent: ")
             self.handle_input(input_str)
 
+
 # --- Streamlit Integration ---
-if __name__ == '__main__':
+if __name__ == "__main__":
     agent = AIAgent()
     st.title("AI Agent")
     st.write("Enter a command for the AI Agent:")
     input_str = st.text_input("")
-    agent.handle_input(input_str)
-    agent.run()
+    if input_str:
+        agent.handle_input(input_str)
+    # agent.run()  # Removed to avoid conflict with Streamlit
 
 # --- Gradio Integration ---
 def gradio_interface(input_text):
     """Gradio interface function."""
     try:
         agent.handle_input(input_text)
-        output = agent.task_history[-1]['output']  # Get the latest output
+        output = agent.task_history[-1]["output"]  # Get the latest output
         return output
     except Exception as e:
         return f"Error: {e}"
 
+
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Textbox(label="Enter Command"),
     outputs=gr.Textbox(label="Output"),
     title="AI Agent",
-    description="Interact with the AI Agent."
+    description="Interact with the AI Agent.",
 )
 
 iface.launch()