├── docs ├── usage.md └── assets │ ├── demo.gif │ ├── squema.jpg │ ├── mads-logo.png │ ├── demo-tera-systems.mp4 │ └── MADS__Multi-agents-for-data-science.pdf ├── tests ├── __init__.py ├── test_agent_manager.py └── test_configs.py ├── src └── mads │ ├── __init__.py │ ├── config.py │ ├── chat_manager.py │ └── agent_manager.py ├── setup.py ├── .gitattributes ├── LICENSE ├── pyproject.toml ├── .gitignore └── README.md /docs/usage.md: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/mads/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | setup() -------------------------------------------------------------------------------- /docs/assets/demo.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AiFlowSolutions/MADS/HEAD/docs/assets/demo.gif -------------------------------------------------------------------------------- /docs/assets/squema.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AiFlowSolutions/MADS/HEAD/docs/assets/squema.jpg -------------------------------------------------------------------------------- /docs/assets/mads-logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AiFlowSolutions/MADS/HEAD/docs/assets/mads-logo.png -------------------------------------------------------------------------------- /docs/assets/demo-tera-systems.mp4: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AiFlowSolutions/MADS/HEAD/docs/assets/demo-tera-systems.mp4 -------------------------------------------------------------------------------- /docs/assets/MADS__Multi-agents-for-data-science.pdf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/AiFlowSolutions/MADS/HEAD/docs/assets/MADS__Multi-agents-for-data-science.pdf -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | *.joblib filter=lfs diff=lfs merge=lfs -text 2 | *.pkl filter=lfs diff=lfs merge=lfs -text 3 | tasks/kc_house_files_generated/house_price_model.joblib filter=lfs diff=lfs merge=lfs -text 4 | -------------------------------------------------------------------------------- /tests/test_agent_manager.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | from src.mads.agent_manager import AgentManager 4 | 5 | def test_data_initialization(): 6 | 7 | llm_config = {"param1": "value1"} # random configuration 8 | data_name_with_csv = "data.csv" 9 | agent_manager_with_csv = AgentManager(llm_config, data_name_with_csv) 10 | assert agent_manager_with_csv.data == "data" -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 AiFlowSolutions 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=42", "wheel"] 3 | build-backend = "setuptools.build_meta" 4 | 5 | [project] 6 | name = "pymads" 7 | version = "0.0.6" 8 | dependencies = [ 9 | "annotated-types>=0.7.0", 10 | "anyio>=4.4.0", 11 | "certifi>=2024.6.2", 12 | "charset-normalizer>=3.3.2", 13 | "colorama>=0.4.6", 14 | "diskcache>=5.6.3", 15 | "distro>=1.9.0", 16 | "docker>=7.1.0", 17 | "FLAML>=2.1.2", 18 | "h11>=0.14.0", 19 | "httpcore>=1.0.5", 20 | "httpx>=0.27.0", 21 | "idna>=3.7", 22 | "numpy>=1.26.4", 23 | "openai>=1.31.0", 24 | "packaging>=24.0", 25 | "pandas>=2.2.2", 26 | "pyautogen>=0.2.28", 27 | "pydantic>=2.7.3", 28 | "pydantic_core>=2.18.4", 29 | "python-dateutil>=2.9.0.post0", 30 | "python-dotenv>=1.0.1", 31 | "pytz>=2024.1", 32 | "regex>=2024.5.15", 33 | "requests>=2.32.3", 34 | "six>=1.16.0", 35 | "sniffio>=1.3.1", 36 | "termcolor>=2.4.0", 37 | "tiktoken>=0.7.0", 38 | "tqdm>=4.66.4", 39 | "typing_extensions>=4.12.1", 40 | "tzdata>=2024.1", 41 | "urllib3>=2.2.1" 42 | ] 43 | authors = [ 44 | { name="Diogo Pedrosa", email="diogofranciscop@hotmail.com"}, 45 | { name="Iago Gaspar", email="wutsuperwut@gmail.com" }, 46 | 47 | ] 48 | description = "A framework that uses multi-agents to enable users to perform a systematic data science pipeline with just two inputs." 49 | readme = "README.md" 50 | requires-python = ">=3.11.7" 51 | classifiers = [ 52 | "Programming Language :: Python :: 3", 53 | "License :: OSI Approved :: MIT License", 54 | "Operating System :: OS Independent", 55 | ] 56 | 57 | [project.urls] 58 | Homepage = "https://github.com/AiFlowSolutions/MADS" 59 | Issues = "https://github.com/AiFlowSolutions/MADS/issues" 60 | 61 | 62 | [tool.pytest] 63 | addopts = "--strict-markers" 64 | testpaths = ["tests"] 65 | -------------------------------------------------------------------------------- /tests/test_configs.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import sys 3 | import os 4 | 5 | # Add the project's root directory to the Python path 6 | sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) 7 | 8 | from src.mads.config import configure_llm, is_termination_msg_chat, create_task_folders 9 | 10 | 11 | # Tests for the configure_llm function 12 | 13 | def test_configure_llm_llama3(): 14 | llm_name = "llama3-70b-8192" 15 | api_key = "your_api_key" 16 | expected_config = [{ 17 | "model": llm_name, 18 | "base_url": "https://api.groq.com/openai/v1", 19 | "api_key": api_key, 20 | "temperature": 0, 21 | "cache_seed": None, 22 | }] 23 | assert configure_llm(llm_name, api_key) == expected_config 24 | 25 | def test_configure_llm_gpt_3_5_turbo(): 26 | llm_name = "gpt-3.5-turbo-0125" 27 | api_key = "your_api_key" 28 | expected_config = [{ 29 | "model": llm_name, 30 | "api_key": api_key, 31 | "temperature": 0, 32 | "cache_seed": None, 33 | }] 34 | assert configure_llm(llm_name, api_key) == expected_config 35 | 36 | def test_configure_llm_invalid_name(): 37 | with pytest.raises(ValueError): 38 | configure_llm("invalid_llm_name", "your_api_key") 39 | 40 | 41 | # Tests for the Termination message function 42 | 43 | def test_is_termination_msg_chat_empty(): 44 | message = "" 45 | assert is_termination_msg_chat(message) == True 46 | 47 | def test_is_termination_msg_chat_terminate(): 48 | message = "This is a message TERMINATE" 49 | assert is_termination_msg_chat(message) == True 50 | 51 | def test_is_termination_msg_chat_not_terminate(): 52 | message = "This is not a termination message" 53 | assert is_termination_msg_chat(message) == False 54 | 55 | def test_is_termination_msg_chat_dict(): 56 | message = {"content": "TERMINATE"} 57 | assert is_termination_msg_chat(message) == True 58 | 59 | def test_is_termination_msg_chat_dict_empty_content(): 60 | message = {"content": ""} 61 | assert is_termination_msg_chat(message) == True 62 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # poetry 98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 102 | #poetry.lock 103 | 104 | # pdm 105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 106 | #pdm.lock 107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 108 | # in version control. 109 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 110 | .pdm.toml 111 | .pdm-python 112 | .pdm-build/ 113 | 114 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 115 | __pypackages__/ 116 | 117 | # Celery stuff 118 | celerybeat-schedule 119 | celerybeat.pid 120 | 121 | # SageMath parsed files 122 | *.sage.py 123 | 124 | # Environments 125 | .env 126 | .venv 127 | env/ 128 | venv/ 129 | ENV/ 130 | env.bak/ 131 | venv.bak/ 132 | 133 | # Spyder project settings 134 | .spyderproject 135 | .spyproject 136 | 137 | # Rope project settings 138 | .ropeproject 139 | 140 | # mkdocs documentation 141 | /site 142 | 143 | # mypy 144 | .mypy_cache/ 145 | .dmypy.json 146 | dmypy.json 147 | 148 | # Pyre type checker 149 | .pyre/ 150 | 151 | # pytype static type analyzer 152 | .pytype/ 153 | 154 | # Cython debug symbols 155 | cython_debug/ 156 | 157 | # PyCharm 158 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 159 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 160 | # and can be added to the global gitignore or merged into this file. For a more nuclear 161 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 162 | #.idea/ 163 | -------------------------------------------------------------------------------- /src/mads/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | from dotenv import load_dotenv 3 | 4 | load_dotenv() 5 | 6 | 7 | STANDARD_CODE = """In the following cases, suggest python code (in a python coding block) or 8 | shell script (in a sh coding block) for the user to execute.\n 9 | 1. When you need to collect info, use the code to output the info you 10 | need, for example, browse or search the web, download/read a file, print the 11 | content of a webpage or a file, get the current date/time, check the 12 | operating system. After sufficient info is printed and the task is ready to 13 | be solved based on your language skill, you can solve the task by yourself.\n 14 | 2. When you need to perform some task with code, use the code to perform 15 | the task and output the result. Finish the task smartly.\n 16 | Don't write code that make plots or visualizations.\n 17 | After writing code cells, always let the user run it.\n 18 | The user can't replace anything from the code you give to him. So don not suggest the user to replace code.""" 19 | 20 | 21 | def configure_llm(llm_name: str, api_key:str) -> dict: 22 | """ 23 | Choose wihch LLM configuration the agents pipeline will have. 24 | 25 | Args: 26 | name (str): name of the LLM the user want to use. 27 | Returns: 28 | llm_config (dict): a dictionary with the LLM configuraÇtion for the LLM the user have choosed. 29 | """ 30 | if llm_name == "llama3-70b-8192": 31 | llm_config=[{ 32 | "model": llm_name, 33 | "base_url": "https://api.groq.com/openai/v1", 34 | "api_key": api_key, 35 | "temperature": 0, 36 | "cache_seed": None, 37 | }] 38 | elif llm_name == "gpt-3.5-turbo-0125": 39 | llm_config = [{ 40 | "model": llm_name, 41 | "api_key": api_key, 42 | "temperature": 0, 43 | "cache_seed": None, 44 | }] 45 | else: 46 | raise ValueError("The specified LLM name doesn't exist in the existing configurations. Please choose another or modify the function in 'config.py'.") 47 | 48 | return llm_config 49 | 50 | def is_termination_msg_chat(message:str) -> bool: 51 | """ 52 | Checks if the message from the last agent was a termination message or not. 53 | 54 | Args: 55 | message (str): last message from the agent 56 | Return: 57 | state (boolean): true if the message is empty or if ends with TERMINATE, False otherwise 58 | """ 59 | 60 | if isinstance(message, dict): 61 | message = message.get("content", "") 62 | 63 | termination_keywords = "TERMINATE" 64 | state = message.rstrip().endswith(termination_keywords) or not message.rstrip() 65 | return state 66 | 67 | def tasks(user_problem:str) -> list: 68 | """ 69 | Initialize the general tasks for each agent. 70 | 71 | Args: 72 | user_problem (str): A string that the user have inputed defining the problem he wants to solve 73 | Return: 74 | taks_prompts (list): A list containing the tasks for each agent, updated with the user problem. 75 | """ 76 | 77 | task_prompts = [ 78 | f"This is my problem: {user_problem}. What is the Machine learning problem?", 79 | "I want to perform data analysis", 80 | "Based on the relevant insights, identify the most relevant machine learning model to use.", 81 | "Base on the relevant insights, make the necessary transformations to the data, separate the data by features and target based on the problem type and split the dataset.", 82 | "Fit the training data, make predictions, and evaluate them.", 83 | "Compile a detailed report with insights from other agents.", 84 | ] 85 | return task_prompts 86 | 87 | def create_task_folders(): 88 | """ 89 | Create a folder structure for tasks with datasets and generated files. 90 | """ 91 | try: 92 | # Create the main 'tasks' folder 93 | os.makedirs('tasks', exist_ok=True) 94 | 95 | # Create the 'datasets' folder inside 'tasks' 96 | os.makedirs(os.path.join('tasks', 'datasets'), exist_ok=True) 97 | 98 | # Create the 'generated_files' folder inside 'tasks' 99 | os.makedirs(os.path.join('tasks', 'generated_files'), exist_ok=True) 100 | 101 | print("Folders created successfully.") 102 | except Exception as e: 103 | print(f"An error occurred: {str(e)}") 104 | 105 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 |
2 | 3 | 4 | **MADS - Multi-Agents for Data Science!** 5 | **Our goal is to enable everyone to apply machine learning with just two inputs!** 6 |

7 | 8 | [[Our Website]](https://aiflowsolutions.github.io/site-agi-flow-solutions/) | [[Our Research Website]](https://aiflowsolutions.github.io/site-agi-flow-research-robotics/) | [[pré-print MADS paper]](https://aiflowsolutions.github.io/site-agi-flow-research-robotics/papers.html) 9 | 10 | [![PyPI](https://img.shields.io/pypi/v/pymads)](https://pypi.org/project/pymads/) 11 | 12 |

13 |
14 | 15 | ## Table of contents 16 | 17 | - [Overview](#overview) 18 | - [Roadmap](#roadmap) 19 | - [Installation](#installation) 20 | - [Setting up MADS](#setting-up-mads) 21 | - [Contribuition](#contribuition) 22 | - [Issue Reporting](#issue-reporting) 23 | - [Contact us](#contact-us) 24 | - [License](#license) 25 | 26 | 27 | ## Overview 28 | MADS is a project aimed at creating a platform where users can update a dataset. Our agents will then execute all the necessary steps in the data science pipeline. The user simply needs to define the goal of the project, and our agents will handle the rest. In the end, the user will have access to a trained model, to the predictions and a report that includes insights from each agent. 29 | 30 | ### Squema 31 | ![Squema](docs/assets/squema.jpg) 32 | 33 | ### Demo 34 | This is a video demo of a plataform where we apply the MADS library. 35 | ![Demo](docs/assets/demo.gif) 36 | 37 | ## Roadmap 38 | - Implement a Reiforcment Learning Agent to improve the current prompts. 39 | - Create tools for the agents to use them. 40 | - Example 1: For the model-building agent, provide a tool like one from Nixtla's library to improve forecasting for time series problems. 41 | - Example 2: For the model consultant agent, create a tool to optimize the selected model. 42 | - Introduce a new agent to interact with the data analyst and generate visualizations useful for the final report. 43 | 44 | ## Installation 45 | MADS requires Python >=3.11.7 installed on your system. We recommend setting up a virtual environment before starting to work with MADS: 46 | - Create a virtual environment: `python -m venv .venv` 47 | - Activate your virtual environment: `.\.venv\Scripts\activate` 48 | 49 | To begin using our library, simply install it via pip: 50 | `pip install pymads` 51 | 52 | ## Setting Up MADS 53 | Before running the bellow script create a `.env` file and place your API Key inside it: `GROQ_API_KEY="your_api_key"` or `OPENAI_API_KEY="your_api_key"`. 54 | Although by default, Docker usage for code execution is set to false, some users have encountered issues with agents still attempting to use the Docker engine. To ensure that the agents do not look for Docker, place `AUTOGEN_USE_DOCKER="False"` inside your `.env` file. 55 | 56 | ```python 57 | import os 58 | from mads.chat_manager import ChatManager 59 | from mads.config import configure_llm, create_task_folders 60 | 61 | # Check if the 'tasks' directory exists, if not, create it. 62 | if not os.path.exists('tasks'): 63 | create_task_folders() 64 | 65 | # After creating the folder, upload the dataset (a csv) you wish to test into tasks/datasets. 66 | # Note: If there is no dataset in tasks, the pipeline will proceed but yield no results. 67 | 68 | # Set your API keys here. 69 | # These keys may be from Groq or OpenAI; current configurations are better suited for Groq. 70 | 71 | GROQ_API_KEY = os.getenv("GROQ_API_KEY") 72 | # OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") 73 | 74 | # Select the model to use (tested options include 'llama3-70b-8192' and 'gpt-3.5-turbo-0125'). 75 | # To test additional models, modify the config.py script accordingly. 76 | model = configure_llm("llama3-70b-8192", GROQ_API_KEY) 77 | 78 | # Define the supervised ML problem that you intend to solve. 79 | problem = "I want to predict wine quality" 80 | 81 | # Specify the filename of your data 82 | dataset = "winequality-red.csv" 83 | 84 | # Choose the agents to deploy. 85 | # Here, we select all six available agents. 86 | agents = [1, 2, 3, 4, 5, 6] 87 | 88 | # Configure the ChatManager class. 89 | chat_manager = ChatManager(dataset, problem, model, agents) 90 | 91 | # Begin the chat sessions. 92 | chat_results = chat_manager.initiate_chats() 93 | ``` 94 | 95 | ## Contribuition 96 | MADS is open-source and we welcome contributions. If you're looking to contribute, please: 97 | 98 | - Fork the repository. 99 | - Create a new branch for your feature. 100 | - Add your feature or improvement. 101 | - Send a pull request. 102 | 103 | ## Issue Reporting 104 | 105 | If you encounter any problems while using our library, please don't hesitate to [create an issue](https://github.com/AiFlowSolutions/MADS/issues) on GitHub. When reporting issues, please provide as much detail as possible, including: 106 | 107 | - Steps to reproduce the issue 108 | - Expected behavior 109 | - Actual behavior 110 | - Any error messages or stack traces 111 | 112 | Your feedback is valuable to us and will help improve the library for everyone. Thank you for contributing to the project! 113 | 114 | ## Contact Us 115 | 116 | If you have any questions, suggestions, or feedback regarding MADS, please feel free to reach out to us: 117 | 118 | - **Company Email**: [info@aiflowsolutions.eu](mailto:info@aiflowsolutions.eu) 119 | 120 | ##### Main Contributors 121 | 122 | - **Email**: [wutsuperwut@gmail.com](mailto:wutsuperwut@gmail.com) 123 | - **Email**: [diogofranciscop@hotmail.com](mailto:diogofranciscop@hotmail.com) 124 | 125 | We are committed to improving MADS based on your input and look forward to hearing from you! 126 | 127 | ## Ackonwledgments 128 | A heartfelt thank you to all the contributors of the autogen framework. Your dedication and hard work have been instrumental in making this project possible. We deeply appreciate the entire community's support and involvement. 129 | 130 | ## License 131 | MADS is released under the MIT License. 132 | -------------------------------------------------------------------------------- /src/mads/chat_manager.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from mads.agent_manager import AgentManager 3 | from mads.config import tasks 4 | 5 | # Configure logging 6 | logging.basicConfig(level=logging.ERROR, format='%(asctime)s - %(levelname)s - %(message)s') 7 | 8 | class ChatManager: 9 | """ 10 | This class manages the initiation of chats for a sequence of data science tasks. 11 | 12 | Attributes: 13 | task (function): A function that returns the user's problem. 14 | dataset_path (str): The datasets needed for the task. 15 | agent_manager (AgentManager): An instance of the AgentManager class. 16 | user (ConversableAgent): The user agent. 17 | include_chats (list): A list of chat IDs to include in the conversation (i.e. which agents to include). 18 | """ 19 | def __init__(self, dataset_name, user_problem, model, include_chats=None): 20 | try: 21 | self.task = tasks(user_problem) 22 | self.dataset_path = f'datasets/{dataset_name}' 23 | self.agent_manager = AgentManager(llm_config=model, data_name=dataset_name) 24 | self.user = self.agent_manager.user_agent() 25 | self.include_chats = include_chats if include_chats is not None else list(range(1, 7)) 26 | except Exception as e: 27 | logging.error("Error during initialization", exc_info=True) 28 | raise 29 | 30 | def initiate_chats(self): 31 | try: 32 | carryover_msg = f"The needed data is in {self.dataset_path}" 33 | 34 | chat_configurations = [ 35 | self._build_chat_config( 36 | chat_id, 37 | { 38 | "summary_method": "reflection_with_llm", 39 | **({"carryover": carryover_msg} if chat_id in [1, 2, 3] else {}) 40 | }, 41 | ) 42 | for chat_id in self.include_chats 43 | ] 44 | 45 | chat_results = self.user.initiate_chats(chat_configurations) 46 | return chat_results 47 | except Exception as e: 48 | logging.error("Error during chat initiation", exc_info=True) 49 | raise 50 | 51 | def _get_summary_prompt(self): 52 | try: 53 | return { 54 | 1: """ 55 | Summarize the content and format summary EXACTLY as follows: 56 | --- 57 | *Data set name*: 58 | `Acme Corp` 59 | --- 60 | *User problem*: 61 | `Regression, Classification or Time-series` 62 | --- 63 | *Target variable* 64 | `write here the target variable` 65 | --- 66 | *Indications on how to write the code to read data* 67 | `Write the code that you read the data` 68 | """, 69 | 2: """ 70 | Summarize the content and format summary EXACTLY as follows: 71 | --- 72 | *Problem*: 73 | `Write the machine learning problem` 74 | --- 75 | *Correlations*: 76 | `Relevant correlations` 77 | --- 78 | *Columns name*: 79 | `Columns: column1, column2...` 80 | --- 81 | *Relevant insights*: 82 | `Useful insights found for the next agents` 83 | --- 84 | """, 85 | 3: """ 86 | Summarize the content and format summary EXACTLY as follows: 87 | --- 88 | *Problem*: 89 | `problem written here` 90 | --- 91 | *Machine learning model to use*: 92 | `ML model` 93 | --- 94 | *Explanation and alternatives*: 95 | `Explanation of why that model was chosen` 96 | --- 97 | """, 98 | 4: """ 99 | Summarize the content and format summary EXACTLY as follows: 100 | --- 101 | *Transformations*: 102 | `Transformations you've done to the data` 103 | --- 104 | *Splitting*: 105 | `The split you've done and where you save the data` 106 | --- 107 | **Read the data** 108 | `pd.read_csv('X_train.csv')` 109 | `pd.read_csv('y_train.csv')` 110 | `pd.read_csv('X_test.csv')` 111 | `pd.read_csv('y_test.csv')` 112 | """, 113 | 5: """ 114 | Summarize the content and format summary EXACTLY as follows: 115 | --- 116 | *ML model used*: 117 | `ML model` 118 | --- 119 | *Place where you saved predictions*: 120 | `acmecorp.com` 121 | --- 122 | *Results of the evaluations*: 123 | `Metric: result` 124 | --- 125 | """, 126 | 6: """ 127 | Summarize the content and format summary EXACTLY as follows: 128 | --- 129 | *Data set name*: 130 | `name of the dataset` 131 | --- 132 | *User problem*: 133 | `problem of the user` 134 | --- 135 | *Target variable* 136 | `the variable we're trying to predict` 137 | --- 138 | **Correlations:** 139 | 'Correlations found' 140 | --- 141 | **Columns:** 142 | 'the columns' 143 | --- 144 | **Relevant Insights:** 145 | 'The relevant insights' 146 | --- 147 | *Machine learning model to use*: 148 | Random Forest Regressor 149 | --- 150 | *Explanation and alternatives*: 151 | 'Explanations of the machine learning model used' 152 | --- 153 | *Transformations*: 154 | `the transformations made to the data` 155 | --- 156 | *Splitting*: 157 | `The splitting done to the data` 158 | --- 159 | *Results of the evaluations*: 160 | `results of the metrics` 161 | --- 162 | """ 163 | } 164 | except Exception as e: 165 | logging.error("Error in _get_summary_prompt", exc_info=True) 166 | raise 167 | 168 | def _build_chat_config(self, chat_id, base_config): 169 | try: 170 | config = base_config.copy() 171 | summary_args = self._get_summary_prompt() 172 | chat_methods = self._get_chat_methods() 173 | 174 | config.update({ 175 | "chat_id": chat_id, 176 | "recipient": chat_methods[chat_id](), 177 | "message": self.task[chat_id - 1], 178 | "summary_args": {"summary_prompt":summary_args[chat_id]} 179 | }) 180 | 181 | if chat_id == 1: 182 | config["clear_history"] = True 183 | config["silent"] = False 184 | 185 | return config 186 | except Exception as e: 187 | logging.error(f"Error in _build_chat_config for chat_id {chat_id}", exc_info=True) 188 | raise 189 | 190 | def _get_chat_methods(self): 191 | try: 192 | return { 193 | 1: self.agent_manager.problem_definer_agent, 194 | 2: self.agent_manager.data_analyst_agent, 195 | 3: self.agent_manager.model_consultor_agent, 196 | 4: self.agent_manager.feature_engineer_agent, 197 | 5: self.agent_manager.model_builder_agent, 198 | 6: self.agent_manager.report_builder_agent, 199 | } 200 | except Exception as e: 201 | logging.error("Error in _get_chat_methods", exc_info=True) 202 | raise 203 | -------------------------------------------------------------------------------- /src/mads/agent_manager.py: -------------------------------------------------------------------------------- 1 | from autogen import ConversableAgent 2 | from mads.config import is_termination_msg_chat, STANDARD_CODE 3 | 4 | class AgentManager: 5 | """ 6 | This class manages the creation of different types of agents for a data science pipeline. 7 | 8 | Attributes: 9 | llm_config (dict): Configuration for the low-level model. 10 | data (str): The name of the dataset. 11 | """ 12 | def __init__(self, llm_config, data_name): 13 | self.llm_config = {"config_list": llm_config} 14 | self.data = data_name.replace('.csv', '') 15 | self.costar_instructions = { 16 | "User": """ 17 | # CONTEXT # 18 | You are an automated code execution service, responsible for running code snippets submitted by agents. 19 | 20 | # OBJECTIVE # 21 | Your sole function is to execute the provided code exactly as it is, including necessary library installations, and report any output or errors generated by the code. 22 | 23 | # STYLE # 24 | Interactions should be minimalistic and strictly pertain to the outcome of the code execution. Omit extraneous information such as verbose confirmation of successful operations. 25 | 26 | # TONE # 27 | The tone should be neutral and functional, aimed at delivering clear and essential feedback concerning the execution results. 28 | 29 | # AUDIENCE # 30 | Your primary users are agents who depend on the accurate execution of code and clear reporting of outcomes to progress in their work. 31 | 32 | # RESPONSE # 33 | You will run the supplied code and communicate only the outputs or errors that arise from the execution. No additional commentary or confirmations will be provided unless they directly relate to the code's functioning. 34 | """, 35 | "Problem Definer": """ 36 | # CONTEXT # 37 | The user is engaged in machine learning and is trying to solve a specific type of problem. 38 | They require assistance in correctly identifying whether their task falls under classification, regression, time-series, or another category within machine learning. 39 | 40 | # OBJECTIVE # 41 | Clear define the problem the user is trying to solve: 42 | - Identify how you should read the data to read it in a csv. Check needed delimiters argument, for a cleaning reading. This is a important step! 43 | - Identify the type of task (regression, classification, clustering or time-series). 44 | - Identify which variables are the features, and which are the target. 45 | When tou have these 3 points. Write TERMINATE 46 | 47 | # STYLE # 48 | The response should be clear and concise, focusing solely on pinpointing the precise nature of the ML problem based on the details shared by the user. 49 | Avoid technical jargon that pertains to preprocessing, built models and visualization since you are not providing instructions on those tasks. 50 | Avoid visual representation. 51 | 52 | # TONE # 53 | The tone should be professional and informative, demonstrating expertise in machine learning concepts to foster trust and authority. 54 | 55 | # AUDIENCE # 56 | The primary audience is individuals or entities involved in a machine learning project who possess a technical background and need expert validation of their problem type. 57 | 58 | # RESPONSE # 59 | You have three types of responses: 60 | RESPONSE 1: 61 | If you need to see the data, write python code. 62 | Use the print function to see what you want. 63 | You're script should be in one block. 64 | You should just writte python code in this responses. 65 | RESPONSE 2: 66 | Analyse the output that the user gave you, and responde to the OBJETIVES 67 | RESPONSE 3: 68 | When the user replies with exitcode: 0 (execution succeeded) write TERMINATE 69 | """, 70 | "Data Analyst":""" 71 | # CONTEXT # 72 | You are part of a sequence of agents and you are resposible to give informations about the data to the next agent, a feature engineer agent. 73 | 74 | # OBJECTIVE # 75 | Execute each task one at time: 76 | Check the type of data of each variable and the range of values it can take. 77 | For numerical variables, calculate descriptive statistics. 78 | Calculate correlation coefficients. 79 | Analyse the correlations. 80 | Write TERMINATE 81 | 82 | # STYLE # 83 | Avoid technical jargon that pertains to preprocessing, built models and visualization since you are not providing instructions on those tasks. 84 | Avoid visual representation. 85 | 86 | # TONE # 87 | Maintain an advisory and supportive tone throughout the consultation process, ensuring that the feature engineer feels guided and well-informed about handling their dataset. 88 | 89 | # AUDIENCE # 90 | The primary audience is a feature engineer, so tailor your response to someone with knowledge in feature engineering but who may require analytical expertise. 91 | 92 | # RESPONSE # 93 | You have three types of responses: 94 | RESPONSE 1: 95 | When you need to get informations from the data, write python code. 96 | You're scripts should always be in one block. 97 | You should retrain from plots and avoid visualizations. 98 | You should just writte python code in this responses. 99 | RESPONSE 2: 100 | Analyse the output of the code the user have runed. 101 | RESPONSE 3: 102 | After analysing the correlations write TERMINATE 103 | """, 104 | "Model Consultor":""" 105 | # CONTEXT # 106 | The user requires expertise in selecting the most appropriate statistical or machine learning model for their specific data problem. 107 | As a specialist, you'rr role is to advise on the optimal choice of model while explicitly excluding the use of the Prophet model from my recommendations. 108 | 109 | # OBJECTIVE # 110 | To provide the user with a clear recommendation on the best-suited model that aligns with their data and predictive requirements, taking into consideration all relevant factors except the use of the Prophet model. 111 | 112 | 113 | # STYLE # 114 | Advice should be succinct and focused, directly addressing the criteria and rationale behind the selection of a particular model. The guidance will be purely textual without any code examples or visual elements. 115 | 116 | # TONE # 117 | The tone should be informative and authoritative, instilling confidence in the user regarding the recommended model's suitability for their needs. 118 | 119 | # AUDIENCE # 120 | The intended audience is a user who may range from being a novice to an experienced data science practitioner. 121 | They are seeking expert advice on model selection to inform their work. 122 | 123 | # RESPONSE # 124 | Your response should be short and concise. 125 | When you have decided which is the best model to use write 'TERMINATE' 126 | """, 127 | "Feature Engineer":""" 128 | # CONTEXT # 129 | You are specializing in data preparation for machine learning applications. 130 | This crucial step involves a series of tasks aimed at ensuring the data is ready for modeling. 131 | 132 | # OBJECTIVE # 133 | Execute each task one at time: 134 | Fill missing values. 135 | For time-series data, create lag features, rolling mean or rolling standard deviation. 136 | Split the dataset into training and test subsets. 137 | Save them in X_test.csv, X_train.csv, y_test.csv and y_train.csv. 138 | 139 | # STYLE # 140 | The response should be technical and instructional, providing clear guidelines for the data preprocessing workflow required prior to ML modeling. 141 | 142 | # TONE # 143 | The tone of the response should be informative and precise, maintaining a professional demeanor suitable for data science practitioners looking for guidance in their workflow. 144 | 145 | # AUDIENCE # 146 | The target audience includes data scientists, machine learning engineers, and other professionals working with data who require a systematic approach to preparing their datasets. 147 | 148 | # RESPONSE # 149 | You have three types of responses: 150 | RESPONSE 1: 151 | To solve your tasks, write python code. 152 | You're scripts should always be in one block. 153 | You should just writte python code in this responses. 154 | RESPONSE 2: 155 | Analyse the output of the code the user have runed. 156 | RESPONSE 3: 157 | After the user responds to you that the dara was saved sussfully and with a exitcode: 0 (execution succeeded) write TERMINATE 158 | """, 159 | "Model Builder": 160 | f""" 161 | # CONTEXT # 162 | As a machine learning engineer, you are expected to handle model training and evaluation comprehensively. 163 | This encompasses various responsibilities, including analyzing datasets and applying machine learning algorithms. 164 | 165 | # OBJECTIVE # 166 | Train the machine learning model with `X_train.csv` and `y_train.csv` files. 167 | Make predictions with `X_test.csv`. 168 | Evaluate using `y_test.csv`. (with RMSE and MAE for regression and time-series and with Accuracy for classification), and save the necessary outputs 169 | You should save the machine learnig model as ML_{self.data}.pkl file, in a directory named 'generated_files'. 170 | You should save the predictions as pred_{self.data}.csv file, in a directory named 'generated_files'. 171 | 172 | # STYLE # 173 | The instructions should be communicated with technical accuracy, offering a step-by-step approach for training and evaluating the ML model. 174 | The language used will be precise, catering to a professional audience well-versed in machine learning workflows. 175 | Retrain from asking the user for inputs. 176 | 177 | # TONE # 178 | Maintain an instructional but supportive tone throughout, ensuring clarity for users who are working on training and predicting with ML models. 179 | It should instill confidence in them to perform the required tasks effectively. 180 | 181 | # AUDIENCE # 182 | This explanation is intended for machine learning engineers, data scientists, and others in related fields who have a solid understanding of model development processes, from training to prediction. 183 | 184 | # RESPONSE # 185 | You have three types of responses: 186 | RESPONSE 1: 187 | Write python code to solve the task. 188 | You're scripts should always be in one block. 189 | You should just writte python code in this responses. 190 | RESPONSE 2: 191 | Analyse the output of the code the user have runed. 192 | RESPONSE 3: 193 | When the user replys to you with exitcode: 0 the model was saved sussefully write TERMINATE 194 | """, 195 | "Report Builder":f""" 196 | # CONTEXT # 197 | You are a report generator tasked with converting insights gathered by previous agents into an organized report. 198 | 199 | # OBJECTIVE # 200 | Write the problem of the user. 201 | Write the target variable. 202 | Write the columns. 203 | Write the correlations. 204 | Write the relevant insights. 205 | Write the machine learning model used. 206 | Write the explanations about the machine learning model used. 207 | Write the transformations done. 208 | Write the metrics of the evaluations of the predictions. 209 | Save everything in report_{self.data}.txt file, and then write a Python script that stores this file in a directory named 'generated_files'. 210 | 211 | # STYLE # 212 | The response should consist of clear and executable Python or shell code, ready for the user to run without any modifications required on their part. 213 | 214 | # TONE # 215 | Professional and instructional, ensuring the user understands that the provided code is complete and will perform the task as described when executed. 216 | 217 | # AUDIENCE # 218 | The end-user who requires a Python script to automate the process of saving reports but does not possess the necessary coding skills to modify or write the code themselves. 219 | 220 | # RESPONSE # 221 | You have three types of responses: 222 | RESPONSE 1: 223 | Write python code to solve the task. 224 | You're scripts should always be in one block. 225 | You should just writte python code in this responses. 226 | RESPONSE 2: 227 | When the user replys to you that the report was saved sucefully write TERMINATE 228 | """ 229 | } 230 | self.chain_of_thought_instructions = { 231 | 232 | } 233 | 234 | def _create_agent(self, name, system_message) -> ConversableAgent: 235 | """ 236 | Create a ConversableAgent with common configuration. 237 | The system_message gets formatted with STANDARD_CODE if necessary. 238 | 239 | Args: 240 | name (str): name of the agent 241 | system_message (str): instructions for the agent 242 | Return: 243 | A new agent that is a ConversableAgent object. 244 | """ 245 | return ConversableAgent( 246 | name=name, 247 | llm_config=self.llm_config, 248 | code_execution_config={ 249 | "last_n_messages": 2, 250 | "work_dir": "tasks", 251 | "use_docker": False, 252 | } if name == "User" else False, 253 | human_input_mode="NEVER", 254 | system_message=system_message.format(STANDARD_CODE=STANDARD_CODE), 255 | is_termination_msg=is_termination_msg_chat if name == "User" else None, 256 | ) 257 | 258 | def user_agent(self, method="costar"): 259 | description_key = "User" 260 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 261 | return self._create_agent(description_key, description) 262 | 263 | 264 | def problem_definer_agent(self,method="costar"): 265 | description_key = "Problem Definer" 266 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 267 | return self._create_agent(description_key, description) 268 | 269 | def data_analyst_agent(self,method="costar"): 270 | description_key = "Data Analyst" 271 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 272 | return self._create_agent(description_key, description) 273 | 274 | def model_consultor_agent(self,method="costar"): 275 | description_key = "Model Consultor" 276 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 277 | return self._create_agent(description_key, description) 278 | 279 | def feature_engineer_agent(self,method="costar"): 280 | description_key = "Feature Engineer" 281 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 282 | return self._create_agent(description_key, description) 283 | 284 | def model_builder_agent(self,method="costar"): 285 | description_key = "Model Builder" 286 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 287 | return self._create_agent(description_key, description) 288 | 289 | def report_builder_agent(self,method="costar"): 290 | description_key = "Report Builder" 291 | description = (self.costar_instructions if method == "costar" else self.chain_of_thought_instructions)[description_key] 292 | return self._create_agent(description_key, description) 293 | --------------------------------------------------------------------------------