Multi-Agent System for Science, Made by Cosmologists, Powered by AG2.
Try cmbagent in Colab and check our demo video on YouTube!
This is open-source research-ready software. Please star the repo ⭐ and cite Laverick et al (2024) to support our open-source work.
-
Check the demo notebooks.
-
Best perfmonces are obtained with top-scoring models. (We like Gemini, GPT-4o/4.5/o3, Claude. see our examples.)
-
Currently, RAG agents use
file_searchon OpenAI vector stores.
Note: Please fork and contribute to the repo. We give access to our top-tier OpenAI, Anthropic and Cloud organizations to our top contributors.
We emphasize that cmbagent is under active development and apologize for any bugs.
The backbone of cmbagent is AG2. Please star the AG2 repo ⭐ and cite Wu et al (2023)!
Cmbagent acts according to a Planning and Control strategy with no human-in-the-loop.
You give a task to solve, then:
Planning
- A plan is designed from a conversation between a planner and a plan reviewer.
- Once the number of feedbacks (reviews) is exhausted the plan is recorded in context and cmbagent switches to control.
Control
- The plan is executed step-by-step.
- Sub-tasks are handed over to a single agent in each step.
See Installation, and then in a Jupyter notebook, do:
The output of this session is here.
import os
from cmbagent import CMBAgent
cmbagent = CMBAgent(agent_llm_configs = {
'engineer': {
"model": "o3-mini-2025-01-31",
"reasoning_effort": "high", ## for gpt-4.5-preview-2025-02-27, gpt-4o-2024-11-20, gpt-4o-mini, etc, comment out this line.
"api_key": os.getenv("OPENAI_API_KEY"),
"api_type": "openai",
},
'researcher': {
"model": "gemini-2.0-pro-exp-02-05",
"api_key": os.getenv("GEMINI_API_KEY"),
"api_type": "google",
}})
task = """
Generate simulated stock market data that mimics the behavior of 500 stocks across various sectors (similar to the S&P 500) over a 2-year period.
Within this timeframe, incorporate a financial crisis lasting a few weeks in the middle of the period.
Your simulation should capture:
Sudden volatility spikes, market jumps, and heavy-tailed returns.
Periods of extreme uncertainty and rapid price changes.
Realistic correlations among stocks, particularly reflecting sector-based dependencies.
After generating the simulated data, apply the Black-Scholes Merton model to price call options—focusing specifically on the impact of varying strike prices.
Provide visualizations that illustrate:
The evolution of stock prices, highlighting the crisis period.
Volatility patterns and any market jumps.
The correlation structure among different sectors.
The pricing behavior of call options during both normal and crisis conditions.
Finally, analyze how the crisis impacts option pricing and discuss any limitations or insights regarding the model's performance under extreme market conditions.
"""
cmbagent.solve(task,
max_rounds=500, # set to a high number, this is the max number of total agent calls
shared_context = {'feedback_left': 1, # number of feedbacks on the plan, generally want to set to a low number, as this adds unnecessary complexity to the workflow.
'maximum_number_of_steps_in_plan': 5})Your outputs will be stored in the output directory.
To update a vector stores with local files in your CMBAGENT_DATA folder (see Getting the RAG data), for your RAG agents use:
cmbagent = CMBAgent(make_vector_stores=['name_of_agent'])Before you can use cmbagent, you need to set your OpenAI API key as an environment variable. Do this in a terminal, before launching Jupyter-lab.
For Unix-based systems (Linux, macOS), do:
export OPENAI_API_KEY="sk-..." ## mandatory for the RAG agents
export ANTHROPIC_API_KEY="sk-..." ## optional
export GEMINI_API_KEY="AI...." ## optional (paste in your bashrc or zshrc file, if possible.)
For Windows, use WSL and the same command.
By default, cmbagent uses models from oai/anthropic/google. If you want to pick different LLMs, just adapat agent_llm_configs as above, or the default_agent_llm_configs in utils.py.
These features are based on AG2.
- Multi-LLM multi-agent system
- Tool calls
- Retrieval Augmented Generation
- Structured output
- Swarm
- No human-in-the-loop
- Local code development and execution
Check our examples and try cmbagent in Colab!
@misc{Laverick:2024fyh,
author = "Laverick, Andrew and Surrao, Kristen and Zubeldia, Inigo and Bolliet, Boris and Cranmer, Miles and Lewis, Antony and Sherwin, Blake and Lesgourgues, Julien",
title = "{Multi-Agent System for Cosmological Parameter Analysis}",
eprint = "2412.00431",
archivePrefix = "arXiv",
primaryClass = "astro-ph.IM",
month = "11",
year = "2024"
}Before installing cmbagent, create a virtual environment (we use python3.12):
python3.12 -m venv /path/to/your/envs/cmbagent_env
source /path/to/your/envs/cmbagent_env/bin/activateThen, follow these steps:
Clone and install our package from GitHub.
git clone https://github.com/CMBAgents/cmbagent.git
cd cmbagent
pip install -e .That's it!
If you want to develop the package, you should also install in editable mode
git https://github.com/CMBAgents/ag2.git
cd ag2
git checkout before_gemini_modif
pip install -e .If you ignore this, it is OK, just note that cmbagent_data will go into your $HOME.
If you are a cosmologist, there is already some RAG files for you to play with. If you are not a cosmologist, just modify the code/documents so it does RAG on your own documents of interest. It's pretty straightforward.
Then, do this:
git clone https://github.com/CMBAgents/cmbagent_data.git
export CMBAGENT_DATA=/where/you/have/cloned/cmbagent_dataNote that you need to set the CMBAGENT_DATA environment variable accordingly before using cmbagent
in any future session. Maybe you want to add this to your .bashrc or .zshrc file, or in your activate script!
If you don't set this, the rag data will go into your home directory (you should see a folder cmbagent_data with same content as here.
The core of the code is located in cmbagent.py.
RAG agents can be found in rag_agents. You can make your own easily.
Apart from the RAG agents, we have assistant agents (like an engineer and a planner) and a code execution agent (executor).
All agents inherit from the BaseAgent class. You can find the definition of BaseAgent in the base_agent.py file.
Check the demos.
Task
"""
Provide a gif of the matter power spectrum varying z between 0 and 20, showing the scale where non-linear effects become important
"""
Result
(Note: after
This example illustrates how to skip the Planning step and enter directly into Control mode.
Task
task = r"""
Provide a gif of harmonic oscillator trajectories.
Instruction: make sure the code is optimized.
"""
cmbagent.solve(task,
max_rounds=500,
initial_agent='engineer',
shared_context = {'feedback_left': 0,
"number_of_steps_in_plan": 1,
'maximum_number_of_steps_in_plan': 1})
Result with o3-mini-2025-01-31 [high]
Set-up with:
cmbagent = CMBAgent(agent_llm_configs = {
'engineer': {
"model": "o3-mini-2025-01-31",
"reasoning_effort": "high"
"api_key": os.getenv("OPENAI_API_KEY"),
"api_type": "openai"}})
Result with gemini-2.0-pro-exp-02-05
Set-up with:
cmbagent = CMBAgent(agent_llm_configs = {
'engineer': {
"model": "gemini-2.0-pro-exp-02-05",
"api_key": os.getenv("GEMINI_API_KEY"),
"api_type": "google"}})
Result with claude-3-7-sonnet-20250219
Set-up with:
cmbagent = CMBAgent(agent_llm_configs = {
'engineer': {
"model": "claude-3-7-sonnet-20250219",
"api_key": os.getenv("ANTHROPIC_API_KEY"),
"api_type": "anthropic"}})
Not bad, but axis limits are questionable on the left panel.
Our project is funded by the Cambridge Centre for Data-Driven Discovery Accelerate Programme. We are grateful to Mark Sze for help with AG2.