diff --git a/agents/ml_assistant/agents.py b/agents/ml_assistant/agents.py index ee42380..25e3034 100644 --- a/agents/ml_assistant/agents.py +++ b/agents/ml_assistant/agents.py @@ -51,7 +51,7 @@ 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""", @@ -59,5 +59,20 @@ def starter_code_agent(self): 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, - ) \ No newline at end of file + 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) \ No newline at end of file diff --git a/agents/ml_assistant/crew.py b/agents/ml_assistant/crew.py index 493a970..376e26c 100644 --- a/agents/ml_assistant/crew.py +++ b/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 ) diff --git a/agents/ml_assistant/tasks.py b/agents/ml_assistant/tasks.py index ff3d278..3310c95 100644 --- a/agents/ml_assistant/tasks.py +++ b/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" ) \ No newline at end of file