Added class, method and function doc strings

scripts/ai_config.py

@@ -3,7 +3,25 @@ import data
 
 
 class AIConfig:
-    def __init__(self, ai_name="", ai_role="", ai_goals=[]):
+    """
+    A class object that contains the configuration information for the AI
+      
+    Attributes:
+        ai_name (str): The name of the AI.
+        ai_role (str): The description of the AI's role.
+        ai_goals (list): The list of objectives the AI is supposed to complete.
+    """
+    def __init__(self, ai_name:str="", ai_role:str="", ai_goals:list=[]) -> None:
+        """
+        Initialize a class instance
+        
+        Parameters:
+            ai_name (str): The name of the AI.
+            ai_role (str): The description of the AI's role.
+            ai_goals (list): The list of objectives the AI is supposed to complete.
+        Returns:
+            None
+        """
         self.ai_name = ai_name
         self.ai_role = ai_role
         self.ai_goals = ai_goals
@@ -12,8 +30,18 @@ class AIConfig:
     SAVE_FILE = "../ai_settings.yaml"
 
     @classmethod
-    def load(cls, config_file=SAVE_FILE):
-        # Load variables from yaml file if it exists
+    def load(cls:object, config_file:str=SAVE_FILE) -> object:
+        """
+        Returns class object with parameters (ai_name, ai_role, ai_goals) loaded from yaml file if yaml file exists, 
+        else returns class with no parameters.
+  
+        Parameters:
+           cls (class object): An AIConfig Class object.
+           config_file (int): The path to the config yaml file. DEFAULT: "../ai_settings.yaml"
+           
+        Returns:
+            cls (object): A instance of given cls object  
+        """
         try:
             with open(config_file) as file:
                 config_params = yaml.load(file, Loader=yaml.FullLoader)
@@ -26,12 +54,30 @@ class AIConfig:
 
         return cls(ai_name, ai_role, ai_goals)
 
-    def save(self, config_file=SAVE_FILE):
+    def save(self, config_file:str=SAVE_FILE) -> None:
+        """
+        Saves the class parameters to the specified file yaml file path as a yaml file.
+  
+        Parameters: 
+            config_file(str): The path to the config yaml file. DEFAULT: "../ai_settings.yaml"
+           
+        Returns:
+            None 
+        """
         config = {"ai_name": self.ai_name, "ai_role": self.ai_role, "ai_goals": self.ai_goals}
         with open(config_file, "w") as file:
             yaml.dump(config, file)
 
-    def construct_full_prompt(self):
+    def construct_full_prompt(self) -> str:
+        """
+        Returns a prompt to the user with the class information in an organized fashion.
+  
+        Parameters: 
+            None
+           
+        Returns:
+            full_prompt (str): A string containing the intitial prompt for the user including the ai_name, ai_role and ai_goals.
+        """
         prompt_start = """Your decisions must always be made independently without seeking user assistance. Play to your strengths as an LLM and pursue simple strategies with no legal complications."""
 
         # Construct full prompt
@@ -41,3 +87,7 @@ class AIConfig:
 
         full_prompt += f"\n\n{data.load_prompt()}"
         return full_prompt
+
+if __name__ == "__main__":
+    import doctest
+    doctest.testmod()

scripts/ai_functions.py

@@ -5,9 +5,16 @@ from call_ai_function import call_ai_function
 from json_parser import fix_and_parse_json
 cfg = Config()
 
-# Evaluating code
 
 def evaluate_code(code: str) -> List[str]:
+    """
+    A function that takes in a string and returns a response from create chat completion api call.
+      
+    Parameters:
+        code (str): Code to be evaluated.
+    Returns:
+        A result string from create chat completion. A list of suggestions to improve the code.
+    """
     function_string = "def analyze_code(code: str) -> List[str]:"
     args = [code]
     description_string = """Analyzes the given code and returns a list of suggestions for improvements."""
@@ -17,9 +24,16 @@ def evaluate_code(code: str) -> List[str]:
     return result_string
 
 
-# Improving code
-
 def improve_code(suggestions: List[str], code: str) -> str:
+    """
+    A function that takes in code and suggestions and returns a response from create chat completion api call. 
+      
+    Parameters:
+        suggestions (List): A list of suggestions around what needs to be improved.
+        code (str): Code to be improved.
+    Returns:
+        A result string from create chat completion. Improved code in response.
+    """
     function_string = (
         "def generate_improved_code(suggestions: List[str], code: str) -> str:"
     )
@@ -30,10 +44,16 @@ def improve_code(suggestions: List[str], code: str) -> str:
     return result_string
 
 
-# Writing tests
-
-
 def write_tests(code: str, focus: List[str]) -> str:
+    """
+    A function that takes in code and focus topics and returns a response from create chat completion api call.
+      
+    Parameters:
+        focus (List): A list of suggestions around what needs to be improved.
+        code (str): Code for test cases to be generated against.
+    Returns:
+        A result string from create chat completion. Test cases for the submitted code in response.
+    """
     function_string = (
         "def create_test_cases(code: str, focus: Optional[str] = None) -> str:"
     )