from typing import List, Dict, Optional from custom_types import ( Code, Prompt, AppType, File, Space, Tutorial, App, WebApp, GradioApp, StreamlitApp, ReactApp, Code, ) from prompts import ( createLlamaPrompt, createSpace, isPythonOrGradioAppPrompt, isReactAppPrompt, isStreamlitAppPrompt, getWebApp, getGradioApp, getReactApp, getStreamlitApp, parseTutorial, generateFiles, ) from huggingface_hub import InferenceClient class Agent: def __init__(self, prompts: Dict[str, any]): self.prompts = prompts self.client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1") def process(self, user_input: str) -> str: """ Processes the user's input and generates code. """ # Parse the user's input app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial = self.parse_input(user_input) # Generate a prompt for the Llama model prompt = self.prompts["createLlamaPrompt"]( app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial ) # Generate code using the Llama model code = self.generate_code(prompt) # Generate files for the application files = self.prompts["generateFiles"]( app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial ) # Return the generated code and files return f"Code: {code}\nFiles: {files}" def parse_input(self, user_input: str) -> tuple: """ Parses the user's input and extracts the relevant information. """ # Extract the app type app_type = self.extract_app_type(user_input) # Extract the app name app_name = self.extract_app_name(user_input) # Extract the app description app_description = self.extract_app_description(user_input) # Extract the app features app_features = self.extract_app_features(user_input) # Extract the app dependencies app_dependencies = self.extract_app_dependencies(user_input) # Extract the app space app_space = self.extract_app_space(user_input) # Extract the app tutorial app_tutorial = self.extract_app_tutorial(user_input) return app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial def extract_app_type(self, user_input: str) -> AppType: """ Extracts the app type from the user's input. """ # Check if the user specified a specific app type if "web app" in user_input: return AppType.WEB_APP elif "gradio app" in user_input: return AppType.GRADIO_APP elif "streamlit app" in user_input: return AppType.STREAMLIT_APP elif "react app" in user_input: return AppType.REACT_APP # Otherwise, assume the user wants a web app return AppType.WEB_APP def extract_app_name(self, user_input: str) -> str: """ Extracts the app name from the user's input. """ # Find the substring "app name is:" start_index = user_input.find("app name is:") + len("app name is:") # Find the end of the app name end_index = user_input.find(".", start_index) # Extract the app name app_name = user_input[start_index:end_index].strip() return app_name def extract_app_description(self, user_input: str) -> str: """ Extracts the app description from the user's input. """ # Find the substring "app description is:" start_index = user_input.find("app description is:") + len("app description is:") # Find the end of the app description end_index = user_input.find(".", start_index) # Extract the app description app_description = user_input[start_index:end_index].strip() return app_description def extract_app_features(self, user_input: str) -> List[str]: """ Extracts the app features from the user's input. """ # Find the substring "app features are:" start_index = user_input.find("app features are:") + len("app features are:") # Find the end of the app features end_index = user_input.find(".", start_index) # Extract the app features app_features_str = user_input[start_index:end_index].strip() # Split the app features string into a list app_features = app_features_str.split(", ") return app_features def extract_app_dependencies(self, user_input: str) -> List[str]: """ Extracts the app dependencies from the user's input. """ # Find the substring "app dependencies are:" start_index = user_input.find("app dependencies are:") + len("app dependencies are:") # Find the end of the app dependencies end_index = user_input.find(".", start_index) # Extract the app dependencies app_dependencies_str = user_input[start_index:end_index].strip() # Split the app dependencies string into a list app_dependencies = app_dependencies_str.split(", ") return app_dependencies def extract_app_space(self, user_input: str) -> Optional[Space]: """ Extracts the app space from the user's input. """ # Find the substring "app space is:" start_index = user_input.find("app space is:") + len("app space is:") # Find the end of the app space end_index = user_input.find(".", start_index) # Extract the app space app_space_str = user_input[start_index:end_index].strip() # Create a Space object app_space = Space(space=app_space_str) return app_space def extract_app_tutorial(self, user_input: str) -> Optional[Tutorial]: """ Extracts the app tutorial from the user's input. """ # Find the substring "app tutorial is:" start_index = user_input.find("app tutorial is:") + len("app tutorial is:") # Find the end of the app tutorial end_index = user_input.find(".", start_index) # Extract the app tutorial app_tutorial_str = user_input[start_index:end_index].strip() # Create a Tutorial object app_tutorial = Tutorial(tutorial=app_tutorial_str) return app_tutorial def generate_code(self, prompt: Prompt) -> Code: """ Generates code using the Llama model. """ # Send the prompt to the Llama model response = self.client(prompt.prompt) # Extract the generated code code = response["generated_text"] code = code.replace("```", "") code = code.replace("```", "") # Create a Code object code = Code(code=code, language="python") return code def generate_files(self, app_type: AppType, app_name: str, app_description: str, app_features: List[str], app_dependencies: List[str], app_space: Optional[Space] = None, app_tutorial: Optional[Tutorial] = None) -> List[File]: """ Generates files for the application. """ # Generate files based on the app type files = self.prompts["generateFiles"]( app_type, app_name, app_description, app_features, app_dependencies, app_space, app_tutorial ) return files def main(): """ Main function for the application. """ # Create an agent agent = Agent( prompts={ "createLlamaPrompt": createLlamaPrompt, "createSpace": createSpace, "isPythonOrGradioAppPrompt": isPythonOrGradioAppPrompt, "isReactAppPrompt": isReactAppPrompt, "isStreamlitAppPrompt": isStreamlitAppPrompt, "getWebApp": getWebApp, "getGradioApp": getGradioApp, "getReactApp": getReactApp, "getStreamlitApp": getStreamlitApp, "parseTutorial": parseTutorial, "generateFiles": generateFiles, } ) # Get user input user_input = input("Enter your request: ") # Process the user's input response = agent.process(user_input) # Print the response print(response) if __name__ == "__main__": main()