diff --git a/autogen/version.py b/autogen/version.py
index ab14361953..3b1c89cd7a 100644
--- a/autogen/version.py
+++ b/autogen/version.py
@@ -2,4 +2,4 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-__version__ = "0.5.1"
+__version__ = "0.5.2b1"
diff --git a/setup.py b/setup.py
index fce7ff0afe..64a4e7cee0 100644
--- a/setup.py
+++ b/setup.py
@@ -79,6 +79,9 @@
 # NOTE: underscores in pip install, e.g. pip install ag2[graph_rag_falkor_db], will automatically
 # convert to hyphens. So, do not include underscores in the name of extras.
 
+# ** IMPORTANT: IF ADDING EXTRAS **
+# PLEASE add them in the setup_ag2.py and setup_autogen.py files
+
 extra_require = {
     "test": [
         "ipykernel",
diff --git a/setup_ag2.py b/setup_ag2.py
index 5427a26463..1a659ee875 100644
--- a/setup_ag2.py
+++ b/setup_ag2.py
@@ -24,6 +24,37 @@
     long_description=long_description,
     long_description_content_type="text/markdown",
     install_requires=["pyautogen==" + __version__],
+    extras_require={
+        "test": ["pyautogen[test]==" + __version__],
+        "blendsearch": ["pyautogen[blendsearch]==" + __version__],
+        "mathchat": ["pyautogen[mathchat]==" + __version__],
+        "retrievechat": ["pyautogen[retrievechat]==" + __version__],
+        "retrievechat-pgvector": ["pyautogen[retrievechat-pgvector]==" + __version__],
+        "retrievechat-mongodb": ["pyautogen[retrievechat-mongodb]==" + __version__],
+        "retrievechat-qdrant": ["pyautogen[retrievechat-qdrant]==" + __version__],
+        "graph-rag-falkor-db": ["pyautogen[graph-rag-falkor-db]==" + __version__],
+        "autobuild": ["pyautogen[autobuild]==" + __version__],
+        "captainagent": ["pyautogen[captainagent]==" + __version__],
+        "teachable": ["pyautogen[teachable]==" + __version__],
+        "lmm": ["pyautogen[lmm]==" + __version__],
+        "graph": ["pyautogen[graph]==" + __version__],
+        "gemini": ["pyautogen[gemini]==" + __version__],
+        "together": ["pyautogen[together]==" + __version__],
+        "websurfer": ["pyautogen[websurfer]==" + __version__],
+        "redis": ["pyautogen[redis]==" + __version__],
+        "cosmosdb": ["pyautogen[cosmosdb]==" + __version__],
+        "websockets": ["pyautogen[websockets]==" + __version__],
+        "jupyter-executor": ["pyautogen[jupyter-executor]==" + __version__],
+        "types": ["pyautogen[types]==" + __version__],
+        "long-context": ["pyautogen[long-context]==" + __version__],
+        "anthropic": ["pyautogen[anthropic]==" + __version__],
+        "cerebras": ["pyautogen[cerebras]==" + __version__],
+        "mistral": ["pyautogen[mistral]==" + __version__],
+        "groq": ["pyautogen[groq]==" + __version__],
+        "cohere": ["pyautogen[cohere]==" + __version__],
+        "ollama": ["pyautogen[ollama]==" + __version__],
+        "bedrock": ["pyautogen[bedrock]==" + __version__],
+    },
     url="https://github.com/ag2ai/ag2",
     author="Chi Wang & Qingyun Wu",
     author_email="support@ag2.ai",
diff --git a/setup_autogen.py b/setup_autogen.py
index 0b80410def..1e54c94369 100644
--- a/setup_autogen.py
+++ b/setup_autogen.py
@@ -24,6 +24,37 @@
     long_description=long_description,
     long_description_content_type="text/markdown",
     install_requires=["pyautogen==" + __version__],
+    extras_require={
+        "test": ["pyautogen[test]==" + __version__],
+        "blendsearch": ["pyautogen[blendsearch]==" + __version__],
+        "mathchat": ["pyautogen[mathchat]==" + __version__],
+        "retrievechat": ["pyautogen[retrievechat]==" + __version__],
+        "retrievechat-pgvector": ["pyautogen[retrievechat-pgvector]==" + __version__],
+        "retrievechat-mongodb": ["pyautogen[retrievechat-mongodb]==" + __version__],
+        "retrievechat-qdrant": ["pyautogen[retrievechat-qdrant]==" + __version__],
+        "graph-rag-falkor-db": ["pyautogen[graph-rag-falkor-db]==" + __version__],
+        "autobuild": ["pyautogen[autobuild]==" + __version__],
+        "captainagent": ["pyautogen[captainagent]==" + __version__],
+        "teachable": ["pyautogen[teachable]==" + __version__],
+        "lmm": ["pyautogen[lmm]==" + __version__],
+        "graph": ["pyautogen[graph]==" + __version__],
+        "gemini": ["pyautogen[gemini]==" + __version__],
+        "together": ["pyautogen[together]==" + __version__],
+        "websurfer": ["pyautogen[websurfer]==" + __version__],
+        "redis": ["pyautogen[redis]==" + __version__],
+        "cosmosdb": ["pyautogen[cosmosdb]==" + __version__],
+        "websockets": ["pyautogen[websockets]==" + __version__],
+        "jupyter-executor": ["pyautogen[jupyter-executor]==" + __version__],
+        "types": ["pyautogen[types]==" + __version__],
+        "long-context": ["pyautogen[long-context]==" + __version__],
+        "anthropic": ["pyautogen[anthropic]==" + __version__],
+        "cerebras": ["pyautogen[cerebras]==" + __version__],
+        "mistral": ["pyautogen[mistral]==" + __version__],
+        "groq": ["pyautogen[groq]==" + __version__],
+        "cohere": ["pyautogen[cohere]==" + __version__],
+        "ollama": ["pyautogen[ollama]==" + __version__],
+        "bedrock": ["pyautogen[bedrock]==" + __version__],
+    },
     url="https://github.com/ag2ai/ag2",
     author="Chi Wang & Qingyun Wu",
     author_email="support@ag2.ai",