Update app.py

app.py

@@ -6,21 +6,11 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import black
 from pylint import lint
 from io import StringIO
-import openai
 import sys
 from datetime import datetime
 import requests
 from bs4 import BeautifulSoup
 from typing import List, Dict, Optional
-from utils import refine_code
-
-# Individual imports (assuming files are in the same directory as app.py)
-from utils.refine_code import refine_code  # Import the function
-from utils.refine_code import CodeRefinementError  # Import the class
-import test_code
-import integrate_code
-from test_code import CodeTestingError
-from integrate_code import CodeIntegrationError
 
 # --- Custom Exceptions for Enhanced Error Handling ---
 class InvalidActionError(Exception):
@@ -58,9 +48,9 @@ class AIAgent:
         self.tools = {
             "SEARCH": self.search,
             "CODEGEN": self.code_generation,
-            "REFINE-CODE": refine_code,  # Use external function
-            "TEST-CODE": test_code,  # Use external function
-            "INTEGRATE-CODE": integrate_code,  # Use external function
+            "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
             "TEST-APP": self.test_app,
             "GENERATE-REPORT": self.generate_report,
             "WORKSPACE-EXPLORER": self.workspace_explorer,
@@ -81,7 +71,6 @@ class AIAgent:
         self.search_engine_url: str = "https://www.google.com/search?q="  # Default search engine
         self.prompts: List[str] = []  # Store prompts for future use
         self.code_generator = pipeline('text-generation', model='gpt2')  # Initialize code generator
-        self.openai_api_key = os.getenv("OPENAI_API_KEY")  # Get OpenAI API Key from Environment Variable
 
     # --- Search Functionality ---
     def search(self, query: str) -> List[str]:
@@ -100,25 +89,48 @@ class AIAgent:
     def code_generation(self, snippet: str) -> str:
         """Generates code based on the provided snippet or description."""
         try:
-            # Use OpenAI's GPT-3 for more advanced code generation
-            if self.openai_api_key:
-                openai.api_key = self.openai_api_key
-                response = openai.Completion.create(
-                    engine="text-davinci-003",
-                    prompt=f"```python\n{snippet}\n```\nGenerate Python code based on this snippet or description:",
-                    max_tokens=500,
-                    temperature=0.7,
-                    top_p=1,
-                    frequency_penalty=0,
-                    presence_penalty=0
-                )
-                return response.choices[0].text
-            else:
-                generated_text = self.code_generator(snippet, max_length=500, num_return_sequences=1)[0]['generated_text']
-                return 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."""
+        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}")
+
+    # --- Code Testing Functionality ---
+    def test_code(self, file_path: str) -> str:
+        """Tests the code in the specified file."""
+        try:
+            results = lint.run(file_path, do_exit=False)
+            report = StringIO()
+            results.report(report=report)
+            return report.getvalue()
+        except Exception as e:
+            raise CodeTestingError(f"Error during code testing: {e}")
+
+    # --- Code Integration Functionality ---
+    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."
+        except Exception as e:
+            raise CodeIntegrationError(f"Error during code integration: {e}")
+
     # --- App Testing Functionality ---
     def test_app(self) -> str:
         """Tests the functionality of the app."""
@@ -347,4 +359,4 @@ iface = gr.Interface(
     description="Interact with the AI Agent."
 )
 
-iface.launch()
+iface.launch()