MODIFIED: add one more agent in ml assistant

agents/ml_assistant/agents.py

@@ -51,13 +51,28 @@ def model_recommendation_agent(self):
 
     def starter_code_agent(self):
         return Agent(
-    role='Starter_Code_Generator_Agent',
+        role='Starter_Code_Generator_Agent',
         goal="""generate starter Python code for the project, including data loading, 
                 model definition, and a basic training loop, based on findings from the problem definitions,
                 data assessment and model recommendation""",
         backstory="""You are a code wizard, able to generate starter code templates that users 
                 can customize for their projects. Your goal is to give users a head start in their coding efforts.""",
         verbose=True,
         allow_delegation=False,
-        llm=self.llm,
-        )
+        llm=self.llm
+                )
+
+    def summarization_agent(self):
+        return Agent(
+            role='Summarization_Agent',
+            goal="""Summarize findings from each of the previous steps of the ML discovery process.
+            Include all findings from the problem definitions, data assessment and model recommendation 
+            and all code provided from the starter code generator.
+            """,
+            backstory="""You are a seasoned data scientist, able to break down machine learning problems for
+            less experienced practitioners, provide valuable insight into the problem and why certain ML models
+            are appropriate, and write good, simple code to help get started on solving the problem.
+            """,
+            verbose=True,
+            allow_delegation=False,
+            llm=self.llm)

agents/ml_assistant/crew.py

@@ -17,11 +17,13 @@ def run_ml_crew(file_path, user_question, model="llama3-70b-8192"):
     data_assessment_agent = agents.data_assessment_agent()
     model_recommendation_agent = agents.model_recommendation_agent()
     starter_code_agent = agents.starter_code_agent()
+    summarize_agent = agents.summarization_agent()
 
     task_define_problem = tasks.task_define_problem(problem_definition_agent)
     task_assess_data = tasks.task_assess_data(data_assessment_agent)
     task_recommend_model = tasks.task_recommend_model(model_recommendation_agent)
     task_generate_code = tasks.task_generate_code(starter_code_agent)
+    task_summarize =  tasks.task_summarize(summarize_agent)
 
     # Format the input data for agents
     input_data = {
@@ -32,8 +34,8 @@ def run_ml_crew(file_path, user_question, model="llama3-70b-8192"):
 
     # Initialize and run the crew
     ml_crew = Crew(
-        agents=[problem_definition_agent, data_assessment_agent, model_recommendation_agent, starter_code_agent],
-        tasks=[task_define_problem, task_assess_data, task_recommend_model, task_generate_code],
+        agents=[problem_definition_agent, data_assessment_agent, model_recommendation_agent, starter_code_agent,summarize_agent],
+        tasks=[task_define_problem, task_assess_data, task_recommend_model, task_generate_code,task_summarize],
         verbose=True
     )
 

agents/ml_assistant/tasks.py

@@ -40,4 +40,14 @@ def task_generate_code(self,agent):
             including snippets for package import, data handling, model definition, and training. """,
             agent=agent,
             expected_output="Python code snippets for package import, data handling, model definition, and training, tailored to the user's project, plus a brief summary of the problem and model recommendations."
+        )
+
+    def task_summarize(self,agent):
+        return Task(
+            description="""
+            Summarize the results of the problem definition, data assessment, model recommendation and starter code generator.
+            Don't forget to share the entirety of the starter code!
+            """,
+            agent=agent,
+            expected_output = "Summary of the overall use case"
         )