Update agent.py

agent.py

@@ -1,5 +1,5 @@
 from typing import List, Dict, Optional
-from types import (
+from custom_types import (
     Code,
     Prompt,
     AppType,
@@ -31,15 +31,10 @@ 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"
-        )
+        self.client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
 
     def process(self, user_input: str) -> str:
-        """
-        Processes the user's input and generates code.
-        """
-
+        """ 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)
 
@@ -60,38 +55,26 @@ class Agent:
         return f"Code: {code}\nFiles: {files}"
 
     def parse_input(self, user_input: str) -> tuple:
-        """
-        Parses the user's input and extracts the relevant information.
-        """
-
+        """ 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.
-        """
-
+        """ 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
@@ -101,144 +84,124 @@ class Agent:
             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.
-        """
-
+        """ 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.
-        """
-
+        """ 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.
-        """
-
+        """ 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.
-        """
-
+        """ 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.
-        """
-
+        """ 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.
-        """
-
+        """ 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.
-        """
-
+        """ 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.
-        """
-
+        """ 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()