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 | |
| 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): | |
| # --- Initialize Tools and Attributes --- | |
| self.tools = { | |
| "SEARCH": self.search, | |
| "CODEGEN": self.code_generation, | |
| "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, | |
| "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 = None # Initialize code generator later | |
| self.available_models = [ | |
| "gpt2", | |
| "facebook/bart-large-cnn", | |
| "google/flan-t5-xl", | |
| "bigscience/T0_3B", | |
| ] # Add more as needed | |
| self.selected_model = "gpt2" # Default model | |
| # --- 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: | |
| if not self.code_generator: | |
| self.code_generator = pipeline( | |
| 'text-generation', model=self.selected_model | |
| ) | |
| 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, code: str) -> str: | |
| """Refines the provided code string.""" | |
| try: | |
| 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, code: str) -> str: | |
| """Tests the provided code string using pylint.""" | |
| try: | |
| # 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}") | |
| # --- Code Integration Functionality --- | |
| def integrate_code(self, file_path: str, code_snippet: str) -> str: | |
| """Integrates the code snippet into the specified file.""" | |
| try: | |
| # 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}") | |
| # --- 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 code to refine." | |
| elif action == "TEST-CODE": | |
| return "Provide the code to test." | |
| elif action == "INTEGRATE-CODE": | |
| 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": | |
| 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.set_page_config( | |
| page_title="AI Agent", | |
| page_icon="🤖", | |
| layout="wide", | |
| initial_sidebar_state="expanded", | |
| ) | |
| # --- Tabbed Navigation --- | |
| tabs = st.tabs(["Agent Generation", "Chat App"]) | |
| with tabs[0]: | |
| st.title("AI Agent Generation") | |
| st.sidebar.title("Agent Settings") | |
| # --- Command Dropdown --- | |
| command_options = [ | |
| "SEARCH", | |
| "CODEGEN", | |
| "REFINE-CODE", | |
| "TEST-CODE", | |
| "INTEGRATE-CODE", | |
| "TEST-APP", | |
| "GENERATE-REPORT", | |
| "WORKSPACE-EXPLORER", | |
| "ADD_PROMPT", | |
| "ACTION_PROMPT", | |
| "COMPRESS_HISTORY_PROMPT", | |
| "LOG_PROMPT", | |
| "LOG_RESPONSE", | |
| "MODIFY_PROMPT", | |
| "PREFIX", | |
| "SEARCH_QUERY", | |
| "READ_PROMPT", | |
| "TASK_PROMPT", | |
| "UNDERSTAND_TEST_RESULTS_PROMPT", | |
| ] | |
| selected_command = st.sidebar.selectbox("Command", command_options) | |
| # --- Model Dropdown --- | |
| selected_model = st.sidebar.selectbox( | |
| "Model", agent.available_models, index=agent.available_models.index(agent.selected_model) | |
| ) | |
| agent.selected_model = selected_model | |
| # --- Input Field --- | |
| input_str = st.text_input(f"Enter input for {selected_command}:") | |
| # --- Execute Command --- | |
| if st.button("Execute"): | |
| if input_str: | |
| agent.handle_input(f"{selected_command} {input_str}") | |
| st.write(f"Output: {agent.task_history[-1]['output']}") | |
| # --- Task History --- | |
| st.subheader("Task History") | |
| for task in agent.task_history: | |
| st.write(f"**Action:** {task['action']}") | |
| st.write(f"**Input:** {task['input']}") | |
| st.write(f"**Output:** {task['output']}") | |
| # --- Workspace Explorer --- | |
| st.subheader("Workspace Explorer") | |
| with st.expander("Explore Workspace"): | |
| try: | |
| workspace_output = agent.workspace_explorer() | |
| st.write(workspace_output) | |
| except WorkspaceExplorerError as e: | |
| st.error(f"Error exploring workspace: {e}") | |
| with tabs[1]: | |
| st.title("Chat App") | |
| st.sidebar.title("Chat Settings") | |
| # --- Model Dropdown --- | |
| selected_chat_model = st.sidebar.selectbox( | |
| "Model", agent.available_models, index=agent.available_models.index(agent.selected_model) | |
| ) | |
| agent.selected_model = selected_chat_model | |
| # --- Chat History --- | |
| chat_history = st.empty() | |
| chat_input = st.text_input("Enter your message:") | |
| # --- Chat Input Handling --- | |
| if st.button("Send"): | |
| if chat_input: | |
| user_message = chat_input | |
| chat_history.markdown(f"**You:** {user_message}") | |
| chat_input.empty() | |
| # --- Generate Response --- | |
| try: | |
| if not agent.code_generator: | |
| agent.code_generator = pipeline( | |
| "text-generation", model=agent.selected_model | |
| ) | |
| response = agent.code_generator( | |
| user_message, max_length=500, num_return_sequences=1 | |
| )[0]["generated_text"] | |
| chat_history.markdown(f"**Agent:** {response}") | |
| except Exception as e: | |
| chat_history.markdown(f"**Agent:** Error: {e}") | |
| # --- 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() |