✨ feat: added yaml step management

MANIFEST.in

@@ -0,0 +1 @@
+include agent/templates/*.yml

agent/cli/cli.py

@@ -81,9 +81,9 @@ def config(general):
 
 
 @cli.command()
-def go():
+def chat():
     """
-    go: start the working your ML project.
+    chat: start an interactive chat with LLM to work on your ML project.
     """
     configuration = Config()
     model = load_model()

agent/prompt/base.py

@@ -1,15 +1,37 @@
+from agent.types import Step
+
+from agent.templates import load_step
+
+
 def pmpt_sys_init(
         lang: str
 ) -> str:
     return f"""
     You are an Machine learning engineer, and you are currently working on
      an ML project using {lang} as the primary language.
     
-    You are asked to generate a script and the script file name based on the user requirements.
+    You are asked to generate a code script and the corresponding file name based on the user requirements.
     The output format should be:
     
     File Name: {{file name}}
     
     Code: {{code}}
     
     """
+
+
+def hint_step(
+        step: Step
+) -> str:
+    return f"""
+    You are currently working on step {step} of the project.
+    An ML project in our workflow contains multiple steps, and each step is a task that you need to complete.
+    
+    Now, you are currently on step {step.step}:{step.name}
+    This step is about {step.description}
+    
+    """
+
+
+if __name__ == "__main__":
+    print(hint_step(load_step("data_collection.yml")))

agent/templates/__init__.py

@@ -0,0 +1 @@
+from .utils import *

agent/templates/data_collection.yml

@@ -0,0 +1,3 @@
+step: 0
+name: "Data Collection"
+description: "Data collection in machine learning involves gathering, storing, and preparing high-quality, relevant data from various sources to train and evaluate machine learning models."

agent/templates/data_engineering.yml

@@ -0,0 +1,3 @@
+step: 1
+name: "Data Engineering"
+description: "Data engineering focuses on the practical application of data collection and involves processing, cleansing, and verifying the integrity of data used for analysis. This step ensures that the dataset is ready for training by handling missing data, encoding categorical features, normalizing numerical data, and partitioning datasets into training, validation, and test sets."

agent/templates/model_deployment.yml

@@ -0,0 +1,3 @@
+step: 5
+name: "Model Deployment"
+description: "Model deployment is the final step in the machine learning workflow where the trained model is deployed into a production environment. This can be a web server, a cloud-based application, or embedded systems, allowing the model to make predictions on new data. This step also includes monitoring the model's performance over time and updating it if necessary to maintain its effectiveness."

agent/templates/model_evaluation.yml

@@ -0,0 +1,3 @@
+step: 4
+name: "Model Evaluation"
+description: "After training, the model is evaluated to assess its performance. This typically includes using metrics such as accuracy, precision, recall, and F1-score for classification models, or mean squared error and R² for regression models. Evaluation helps in determining how well the model is likely to perform on unseen data."

agent/templates/model_selection.yml

@@ -0,0 +1,3 @@
+step: 2
+name: "Model Selection"
+description: "Model selection involves choosing the appropriate machine learning algorithms based on the problem type (e.g., regression, classification, clustering) and the nature of the data. This step may include comparing different models, tuning hyperparameters, and using techniques like cross-validation to ascertain the best model for the specific dataset and task."

agent/templates/model_training.yml

@@ -0,0 +1,3 @@
+step: 3
+name: "Model Training"
+description: "Model training is the process where the selected machine learning model is trained on the preprocessed data. This involves feeding the data into the model multiple times, adjusting model parameters (weights) to minimize a loss function, and using optimization techniques like gradient descent to find the best model parameters."

agent/templates/utils.py

@@ -0,0 +1,25 @@
+import os
+import yaml
+from pydantic import ValidationError
+
+from agent.types import Step
+
+
+def load_step(file_name: str) -> Step:
+    """
+    Load a step from a .yaml file.
+    :param file_name: the name of the configuration file.
+    :return:
+    """
+
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+    yml_path = os.path.join(dir_path, file_name)
+
+    with open(yml_path, 'r') as file:
+        data = yaml.safe_load(file)
+    try:
+        config = Step(**data)
+        return config
+    except ValidationError as e:
+        print(f"Error in loading step file: {e}")
+        raise

agent/types/__init__.py

@@ -1 +1,2 @@
 from .project import ProjectState
+from .step import Step

agent/types/step.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class Step(BaseModel):
+    step: int
+    name: str
+    description: str

agent/utils/__init__.py

@@ -1,4 +1,3 @@
 from .config import *
 from .system import *
-from .code import *
 from .chat import Chat

agent/utils/system.py

@@ -1,7 +1,9 @@
 import os
+import re
 import yaml
 from rich.console import Console
 
+from agent.utils import Config
 from agent.types import ProjectState
 from agent.const import CONFIG_PROJECT_FILE
 
@@ -103,3 +105,32 @@ def update_project_state(project_path: str, content_dict: dict = None):
         console.log(f"[green]File '{file_path}' updated successfully.")
     except IOError as error:
         console.log(f"[red]Updating the file '{file_path}' failed due to: {error}")
+
+
+def extract_and_save_file(input_text):
+    """
+    Extracts the file name and code block from a text formatted as specified,
+    then creates a local file with the file name and writes the code into it.
+
+    Args:
+    input_text (str): The text containing the file name and code block.
+
+    Returns:
+    str: The name of the file created.
+    """
+    file_name_match = re.search(r"File Name:\s*(.*?)\s*\n+\s*Code:", input_text, re.DOTALL)
+    if not file_name_match:
+        raise ValueError("File name not found in the text.")
+
+    file_name = file_name_match.group(1).strip()
+
+    project_file_path = os.path.join(Config().read()['project']['path'], file_name)
+    code_match = re.search(r"```(?:[a-zA-Z0-9]+)?\n(.*?)```", input_text, re.DOTALL)
+    if not code_match:
+        raise ValueError("Code block not found in the text.")
+
+    code = code_match.group(1).strip()
+    with open(project_file_path, 'w') as file:
+        file.write(code)
+
+    return project_file_path, code

setup.py

@@ -42,6 +42,10 @@ def readme():
             "mle=agent.cli.cli:cli",
         ]
     },
+    package_data={
+        'steps': ['agent/templates/*.yml']
+    },
+    zip_safe=False,
     include_package_data=True,
     install_requires=requirements,
     setup_requires=['setuptools>=38.6.0'],