diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 03652c84..3e9f61ae 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -3,7 +3,7 @@ name: tests
 on:
   push:
     branches:
-      - master
+      - main
   pull_request:
     types:
       - opened
@@ -24,7 +24,7 @@ jobs:
           python-version: "3.10"
       - name: Install dependencies
         run: |
-          pip install .[lint]
+          pip install . --group lint
       - name: Lint Python code with ruff
         run: |
           ruff check .
@@ -45,7 +45,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          pip install ".[dev, memory]"
+          pip install ".[memory]" --group dev
       - name: Run tests
         run: |
           pytest tests/
@@ -66,7 +66,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .
-          pip install integrations/langchain[dev]
+          pip install integrations/langchain --group dev
       - name: Run tests
         run: |
           pytest integrations/langchain/tests/unit_tests
@@ -87,7 +87,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .[memory]
-          pip install "integrations/langchain[dev, memory]"
+          pip install "integrations/langchain[memory]" --group dev
       - name: Run tests
         run: |
           pytest tests/databricks_ai_bridge/test_lakebase.py
@@ -131,7 +131,7 @@ jobs:
           cp -r older-version/integrations/langchain integrations/
       - name: Install langchain dependency
         run: |
-          pip install integrations/langchain[dev]
+          pip install integrations/langchain --group dev
       - name: Run tests
         run: |
           # Only testing initialization since functionality can change
@@ -156,7 +156,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .
-          pip install integrations/openai[dev]
+          pip install integrations/openai --group dev
       - name: Run tests
         run: |
           pytest integrations/openai/tests/unit_tests
@@ -198,7 +198,8 @@ jobs:
           cp -r older-version/integrations/openai integrations/
       - name: Install openai dependency
         run: |
-          pip install integrations/openai[dev]
+          pip install .
+          pip install integrations/openai --group dev
       - name: Run tests
         run: |
           # Only testing initialization since functionality can change
@@ -220,7 +221,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .
-          pip install integrations/llamaindex[dev]
+          pip install integrations/llamaindex --group dev
       - name: Run tests
         run: |
           pytest integrations/llamaindex/tests/unit_tests
@@ -241,7 +242,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .
-          pip install -e databricks_mcp[dev]
+          pip install -e databricks_mcp --group dev
       - name: Run tests
         run: |
           pytest databricks_mcp/tests/unit_tests
@@ -262,7 +263,7 @@ jobs:
       - name: Install dependencies
         run: |
           pip install .
-          pip install integrations/dspy[dev]
+          pip install integrations/dspy --group dev
       - name: Run tests
         run: |
           pytest integrations/dspy/tests/unit_tests
diff --git a/databricks_mcp/pyproject.toml b/databricks_mcp/pyproject.toml
index 487696a5..0a381968 100644
--- a/databricks_mcp/pyproject.toml
+++ b/databricks_mcp/pyproject.toml
@@ -15,7 +15,7 @@ dependencies = [
     "mlflow>=3.1"
 ]
 
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
   "pytest",
   "pytest-asyncio",
diff --git a/integrations/dspy/pyproject.toml b/integrations/dspy/pyproject.toml
index 922b63d7..70729688 100644
--- a/integrations/dspy/pyproject.toml
+++ b/integrations/dspy/pyproject.toml
@@ -14,7 +14,7 @@ dependencies = [
     "mlflow>=3.0.0",
 ]
 
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
   "pytest",
   "ruff",
diff --git a/integrations/langchain/pyproject.toml b/integrations/langchain/pyproject.toml
index 83ad9000..e25cf4bd 100644
--- a/integrations/langchain/pyproject.toml
+++ b/integrations/langchain/pyproject.toml
@@ -20,6 +20,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+memory = [
+  "langgraph-checkpoint-postgres>=2.0.0",
+  "databricks-ai-bridge[memory]>=0.10.0",
+]
+
+[dependency-groups]
 dev = [
   "pytest",
   "typing_extensions",
@@ -33,10 +39,6 @@ integration = [
   "anyio",
 ]
 
-memory = [
-  "langgraph-checkpoint-postgres>=2.0.0",
-  "databricks-ai-bridge[memory]>=0.10.0",
-]
 
 [build-system]
 requires = ["hatchling"]
diff --git a/integrations/llamaindex/pyproject.toml b/integrations/llamaindex/pyproject.toml
index 84460ea9..46e1c7c4 100644
--- a/integrations/llamaindex/pyproject.toml
+++ b/integrations/llamaindex/pyproject.toml
@@ -15,7 +15,7 @@ dependencies = [
     "unitycatalog-llamaindex[databricks]>=0.2.0",
 ]
 
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
     "pytest",
     "typing_extensions",
@@ -69,4 +69,4 @@ docstring-code-format = true
 docstring-code-line-length = 88
 
 [tool.ruff.lint.pydocstyle]
-convention = "google"
\ No newline at end of file
+convention = "google"
diff --git a/integrations/openai/pyproject.toml b/integrations/openai/pyproject.toml
index c33cd559..9a3b2ae1 100644
--- a/integrations/openai/pyproject.toml
+++ b/integrations/openai/pyproject.toml
@@ -20,7 +20,7 @@ dependencies = [
     "openai-agents>=0.5.0"
 ]
 
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
     "pytest",
     "pytest-asyncio",
diff --git a/pyproject.toml b/pyproject.toml
index ba3e1fae..52a149c8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,19 +21,12 @@ dependencies = [
 [project.license]
 file = "LICENSE.txt"
 
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
-[tool.hatch.build]
-include = [
-  "src/*"
+[project.optional-dependencies]
+memory = [
+  "psycopg[binary,pool]>=3.1",
 ]
 
-[tool.hatch.build.targets.wheel]
-packages = ["src/databricks_ai_bridge"]
-
-[project.optional-dependencies]
+[dependency-groups]
 dev = [
   "hatch",
   "pytest",
@@ -56,10 +49,18 @@ lint = [
   "ruff==0.12.10",
 ]
 
-memory = [
-  "psycopg[binary,pool]>=3.1",
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build]
+include = [
+  "src/*"
 ]
 
+[tool.hatch.build.targets.wheel]
+packages = ["src/databricks_ai_bridge"]
+
 [tool.ruff]
 line-length = 100
 target-version = "py39"