Spaces:
Sleeping
Sleeping
| import os | |
| import subprocess | |
| import streamlit as st | |
| import gradio as gr | |
| from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer | |
| import black | |
| from pylint import lint | |
| from io import StringIO | |
| import sys | |
| from datetime import datetime | |
| import requests | |
| from bs4 import BeautifulSoup | |
| from typing import List, Dict, Optional | |
| # --- Custom Exceptions for Enhanced Error Handling --- | |
| 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 | |
| # --- AI Agent Class --- | |
| class AIAgent: | |
| def __init__(self): | |
| # --- Initialize Tools and Attributes --- | |
| 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 | |
| "TEST-APP": self.test_app, | |
| "GENERATE-REPORT": self.generate_report, | |
| "WORKSPACE-EXPLORER": self.workspace_explorer, | |
| "ADD_PROMPT": self.add_prompt, | |
| "ACTION_PROMPT": self.action_prompt, | |
| "COMPRESS_HISTORY_PROMPT": self.compress_history_prompt, | |
| "LOG_PROMPT": self.log_prompt, | |
| "LOG_RESPONSE": self.log_response, | |
| "MODIFY_PROMPT": self.modify_prompt, | |
| "PREFIX": self.prefix, | |
| "SEARCH_QUERY": self.search_query, | |
| "READ_PROMPT": self.read_prompt, | |
| "TASK_PROMPT": self.task_prompt, | |
| "UNDERSTAND_TEST_RESULTS_PROMPT": self.understand_test_results_prompt, | |
| } | |
| self.task_history: List[Dict[str, str]] = [] | |
| self.current_task: Optional[str] = None | |
| 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 | |
| # --- 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 | |
| try: | |
| response = requests.get(search_url) | |
| response.raise_for_status() # Raise an exception for bad status codes | |
| soup = BeautifulSoup(response.content, 'html.parser') | |
| results = soup.find_all('a', href=True) | |
| return [result['href'] for result in results] | |
| except requests.exceptions.RequestException as e: | |
| raise SearchError(f"Error during search: {e}") | |
| # --- Code Generation Functionality --- | |
| 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'] | |
| 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.""" | |
| 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}") | |
| # --- 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" | |
| for task in self.task_history: | |
| report += f"**Action:** {task['action']}\n" | |
| report += f"**Input:** {task['input']}\n" | |
| report += f"**Output:** {task['output']}\n\n" | |
| return report | |
| # --- Workspace Exploration Functionality --- | |
| def workspace_explorer(self) -> str: | |
| """Provides a workspace explorer functionality.""" | |
| try: | |
| current_directory = os.getcwd() | |
| directories = [] | |
| files = [] | |
| for item in os.listdir(current_directory): | |
| item_path = os.path.join(current_directory, item) | |
| if os.path.isdir(item_path): | |
| directories.append(item) | |
| elif os.path.isfile(item_path): | |
| files.append(item) | |
| return f"**Directories:** {directories}\n**Files:** {files}" | |
| except Exception as e: | |
| raise WorkspaceExplorerError(f"Error during workspace exploration: {e}") | |
| # --- Prompt Management Functionality --- | |
| def add_prompt(self, prompt: str) -> str: | |
| """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}") | |
| # --- 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, 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": | |
| return "Provide the file path of the code to test." | |
| elif action == "INTEGRATE-CODE": | |
| return "Provide the file path and code snippet to integrate." | |
| elif action == "TEST-APP": | |
| return "Test the application." | |
| elif action == "GENERATE-REPORT": | |
| return "Generate a report based on the task history." | |
| elif action == "WORKSPACE-EXPLORER": | |
| return "Explore the current workspace." | |
| elif action == "ADD_PROMPT": | |
| return "Enter the new prompt to add." | |
| elif action == "ACTION_PROMPT": | |
| return "Enter the action to get a prompt for." | |
| elif action == "COMPRESS_HISTORY_PROMPT": | |
| return "Compress the task history." | |
| elif action == "LOG_PROMPT": | |
| return "Enter the event to log." | |
| elif action == "LOG_RESPONSE": | |
| return "Log the specified event." | |
| elif action == "MODIFY_PROMPT": | |
| return "Enter the prompt to modify." | |
| elif action == "PREFIX": | |
| return "Enter the text to add a prefix to." | |
| elif action == "SEARCH_QUERY": | |
| return "Enter the topic to generate a search query for." | |
| elif action == "READ_PROMPT": | |
| return "Enter the file path to read." | |
| elif action == "TASK_PROMPT": | |
| return "Enter the new task to start." | |
| elif action == "UNDERSTAND_TEST_RESULTS_PROMPT": | |
| return "Enter your question about the test results." | |
| else: | |
| raise InvalidActionError("Please provide a valid action.") | |
| except InvalidActionError as e: | |
| raise e | |
| # --- 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?" | |
| # --- Prompt Generation Functionality --- | |
| def log_prompt(self) -> str: | |
| """Provides a prompt to log a specific event.""" | |
| return "What event do you want to log?" | |
| # --- Logging Functionality --- | |
| def log_response(self, event: str) -> str: | |
| """Logs the specified event.""" | |
| print(f"Event logged: {event}") | |
| return "Event logged successfully." | |
| # --- Prompt Modification Functionality --- | |
| def modify_prompt(self, prompt: str) -> str: | |
| """Modifies an existing prompt.""" | |
| try: | |
| # Find the prompt to modify | |
| # Update the prompt | |
| return f"Prompt '{prompt}' modified successfully." | |
| except Exception as e: | |
| raise PromptManagementError(f"Error modifying prompt: {e}") | |
| # --- Prefix Functionality --- | |
| def prefix(self, text: str) -> str: | |
| """Adds a prefix to the provided text.""" | |
| return f"PREFIX: {text}" | |
| # --- Search Query Generation Functionality --- | |
| def search_query(self, query: str) -> str: | |
| """Provides a search query for the specified topic.""" | |
| return f"Search query: {query}" | |
| # --- File Reading Functionality --- | |
| def read_prompt(self, file_path: str) -> str: | |
| """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}") | |
| # --- 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?" | |
| # --- 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?" | |
| # --- 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: | |
| 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: | |
| 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, | |
| CodeRefinementError, CodeTestingError, CodeIntegrationError, | |
| AppTestingError, WorkspaceExplorerError, PromptManagementError, | |
| SearchError) as e: | |
| print(f"Error: {e}") | |
| # --- Main Loop of the Agent --- | |
| def run(self): | |
| """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") | |
| st.write("Enter a command for the AI Agent:") | |
| input_str = st.text_input("") | |
| agent.handle_input(input_str) | |
| 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() |