Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
-import streamlit as st  # For Streamlit
-import gradio as gr  # For Gradio
+import streamlit as st
+import gradio as gr
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 import black
 from pylint import lint
@@ -22,7 +22,7 @@ from utils.code_utils import (
     CodeIntegrationError,
 )
 
-# --- Define custom exceptions for better error handling ---
+# --- Custom Exceptions for Enhanced Error Handling ---
 class InvalidActionError(Exception):
     """Raised when an invalid action is provided."""
     pass
@@ -51,10 +51,10 @@ class SearchError(Exception):
     """Raised when search fails."""
     pass
 
-# --- Define a class for the AI Agent ---
+# --- AI Agent Class ---
 class AIAgent:
     def __init__(self):
-        # --- Initialize tools and attributes ---
+        # --- Initialize Tools and Attributes ---
         self.tools = {
             "SEARCH": self.search,
             "CODEGEN": self.code_generation,
@@ -81,8 +81,9 @@ 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
 
-    # --- Implement search functionality ---
+    # --- Search Functionality ---
     def search(self, query: str) -> List[str]:
         """Performs a web search using the specified search engine."""
         search_url = self.search_engine_url + query
@@ -95,16 +96,30 @@ class AIAgent:
         except requests.exceptions.RequestException as e:
             raise SearchError(f"Error during search: {e}")
 
-    # --- Implement code generation functionality ---
+    # --- Code Generation Functionality ---
     def code_generation(self, snippet: str) -> str:
-        """Generates code based on the provided snippet."""
+        """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']
-            return generated_text
+            # 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
         except Exception as e:
             raise CodeGenerationError(f"Error during code generation: {e}")
 
-    # --- Implement app testing functionality ---
+    # --- App Testing Functionality ---
     def test_app(self) -> str:
         """Tests the functionality of the app."""
         try:
@@ -113,7 +128,7 @@ class AIAgent:
         except subprocess.CalledProcessError as e:
             raise AppTestingError(f"Error during app testing: {e}")
 
-    # --- Implement report generation functionality ---
+    # --- Report Generation Functionality ---
     def generate_report(self) -> str:
         """Generates a report based on the task history."""
         report = f"## Task Report: {self.current_task}\n\n"
@@ -123,7 +138,7 @@ class AIAgent:
             report += f"**Output:** {task['output']}\n\n"
         return report
 
-    # --- Implement workspace exploration functionality ---
+    # --- Workspace Exploration Functionality ---
     def workspace_explorer(self) -> str:
         """Provides a workspace explorer functionality."""
         try:
@@ -140,7 +155,7 @@ class AIAgent:
         except Exception as e:
             raise WorkspaceExplorerError(f"Error during workspace exploration: {e}")
 
-    # --- Implement prompt management functionality ---
+    # --- Prompt Management Functionality ---
     def add_prompt(self, prompt: str) -> str:
         """Adds a new prompt to the agent's knowledge base."""
         try:
@@ -149,14 +164,14 @@ class AIAgent:
         except Exception as e:
             raise PromptManagementError(f"Error adding prompt: {e}")
 
-    # --- Implement prompt generation functionality ---
+    # --- Prompt Generation Functionality ---
     def action_prompt(self, action: str) -> str:
         """Provides a prompt for a specific action."""
         try:
             if action == "SEARCH":
                 return "What do you want to search for?"
             elif action == "CODEGEN":
-                return "Provide a code snippet to generate code from."
+                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."
             elif action == "TEST-CODE":
@@ -196,23 +211,23 @@ class AIAgent:
         except InvalidActionError as e:
             raise e
 
-    # --- Implement prompt generation functionality ---
+    # --- Prompt Generation Functionality ---
     def compress_history_prompt(self) -> str:
         """Provides a prompt to compress the task history."""
         return "Do you want to compress the task history?"
 
-    # --- Implement prompt generation functionality ---
+    # --- Prompt Generation Functionality ---
     def log_prompt(self) -> str:
         """Provides a prompt to log a specific event."""
         return "What event do you want to log?"
 
-    # --- Implement logging functionality ---
+    # --- Logging Functionality ---
     def log_response(self, event: str) -> str:
         """Logs the specified event."""
         print(f"Event logged: {event}")
         return "Event logged successfully."
 
-    # --- Implement prompt modification functionality ---
+    # --- Prompt Modification Functionality ---
     def modify_prompt(self, prompt: str) -> str:
         """Modifies an existing prompt."""
         try:
@@ -222,17 +237,17 @@ class AIAgent:
         except Exception as e:
             raise PromptManagementError(f"Error modifying prompt: {e}")
 
-    # --- Implement prefix functionality ---
+    # --- Prefix Functionality ---
     def prefix(self, text: str) -> str:
         """Adds a prefix to the provided text."""
         return f"PREFIX: {text}"
 
-    # --- Implement search query generation functionality ---
+    # --- Search Query Generation Functionality ---
     def search_query(self, query: str) -> str:
         """Provides a search query for the specified topic."""
         return f"Search query: {query}"
 
-    # --- Implement file reading functionality ---
+    # --- File Reading Functionality ---
     def read_prompt(self, file_path: str) -> str:
         """Provides a prompt to read the contents of a file."""
         try:
@@ -244,35 +259,52 @@ class AIAgent:
         except Exception as e:
             raise InvalidInputError(f"Error reading file: {e}")
 
-    # --- Implement task prompt generation functionality ---
+    # --- Task Prompt Generation Functionality ---
     def task_prompt(self) -> str:
         """Provides a prompt to start a new task."""
         return "What task do you want to start?"
 
-    # --- Implement test results understanding prompt generation functionality ---
+    # --- Test Results Understanding Prompt Generation Functionality ---
     def understand_test_results_prompt(self) -> str:
         """Provides a prompt to understand the test results."""
         return "What do you want to know about the test results?"
 
-    # --- Implement input handling functionality ---
+    # --- Input Handling Functionality ---
     def handle_input(self, input_str: str):
         """Handles user input and executes the corresponding action."""
         try:
             action, *args = input_str.split()
             if action in self.tools:
                 if args:
-                    self.task_history.append({
-                        "action": action,
-                        "input": " ".join(args),
-                        "output": self.tools[action](" ".join(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))}")
+                    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}")
                 else:
-                    self.task_history.append({
-                        "action": action,
-                        "input": None,
-                        "output": self.tools[action]()
-                    })
-                print(f"Action: {action}\nInput: {' '.join(args)}\nOutput: {self.tools[action](' '.join(args))}")
+                    try:
+                        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}")
             else:
                 raise InvalidActionError("Invalid action. Please choose a valid action from the list of tools.")
         except (InvalidActionError, InvalidInputError, CodeGenerationError, 
@@ -281,7 +313,7 @@ class AIAgent:
                 SearchError) as e:
             print(f"Error: {e}")
 
-    # --- Implement the main loop of the agent ---
+    # --- Main Loop of the Agent ---
     def run(self):
         """Runs the agent continuously, waiting for user input."""
         while True:
@@ -295,4 +327,24 @@ if __name__ == '__main__':
     st.write("Enter a command for the AI Agent:")
     input_str = st.text_input("")
     agent.handle_input(input_str)
-    agent.run()
+    agent.run()
+
+# --- 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
+        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."
+)
+
+iface.launch()