diff --git a/cosmos/operators/base.py b/cosmos/operators/base.py
index 552172f746..04b5955649 100644
--- a/cosmos/operators/base.py
+++ b/cosmos/operators/base.py
@@ -9,7 +9,12 @@
 
 import yaml
 from airflow.utils.context import Context, context_merge
-from airflow.utils.operator_helpers import context_to_airflow_vars
+
+try:
+    from airflow.utils.operator_helpers import context_to_airflow_vars
+except ImportError:  # pragma: no cover
+    from airflow.sdk.execution_time.context import context_to_airflow_vars  # type: ignore
+
 from airflow.utils.strings import to_boolean
 
 from cosmos.dbt.executable import get_system_dbt
diff --git a/scripts/airflow3/README.md b/scripts/airflow3/README.md
index 02c37f4380..dfe8f174ba 100644
--- a/scripts/airflow3/README.md
+++ b/scripts/airflow3/README.md
@@ -49,7 +49,7 @@ sh scripts/airflow3/setup.sh
 Activate the virtual env created in previous step and run airflow
 
 ```commandline
-source "$(pwd)/scripts/airflow3/venv/bin/activate"
+source "$(pwd)/scripts/airflow3/venv-af3/bin/activate"
 
 airflow standalone
 ```
@@ -67,7 +67,7 @@ Once Airflow is running, you can also run tests.
 ```commandline
 source scripts/airflow3/env.sh
 
-source "$(pwd)/scripts/airflow3/venv/bin/activate"
+source "$(pwd)/scripts/airflow3/venv-af3/bin/activate"
 
 sh scripts/airflow3/tests.sh
 ```
diff --git a/scripts/airflow3/env.sh b/scripts/airflow3/env.sh
index 5f1ecf5791..0631049c82 100644
--- a/scripts/airflow3/env.sh
+++ b/scripts/airflow3/env.sh
@@ -1,8 +1,9 @@
 #!/bin/bash
 
 set -e
-set -x
 
+PYTHONPATH="$PWD"
+export PYTHONPATH
 AIRFLOW_HOME="$PWD/scripts/airflow3"
 export AIRFLOW_HOME
 export AIRFLOW__LOGGING__BASE_LOG_FOLDER="$AIRFLOW_HOME/logs"
diff --git a/scripts/airflow3/requirements.txt b/scripts/airflow3/requirements.txt
index 7ce5ec81b1..decc3f9ed8 100644
--- a/scripts/airflow3/requirements.txt
+++ b/scripts/airflow3/requirements.txt
@@ -1,9 +1,7 @@
---find-links https://dist.apache.org/repos/dist/dev/airflow/3.0.0b2/
-apache-airflow==3.0.0b2
-apache-airflow-task-sdk==1.0.0b2
-apache-airflow-providers-standard==1.0.0b2
-apache-airflow-providers-fab==2.0.0b2
-apache-airflow-providers-celery==3.11.0b2
-apache-airflow-providers-cncf-kubernetes==10.4.0b2
+apache-airflow
+apache-airflow-task-sdk
+apache-airflow-providers-standard
+apache-airflow-providers-fab
 psycopg2
 asyncpg
+dbt-postgres
diff --git a/scripts/airflow3/setup.sh b/scripts/airflow3/setup.sh
index 315d186598..79cc7f2eb8 100644
--- a/scripts/airflow3/setup.sh
+++ b/scripts/airflow3/setup.sh
@@ -4,26 +4,17 @@
 set -e
 
 # Create a UV virtual environment named 'env' (you can change this as needed)
-echo "Creating UV virtual environment at $(pwd)/tools"
-uv venv venv --directory "$(pwd)/scripts/airflow3"
+echo "Creating virtual environment at $(pwd)/tools"
+python3 -m venv "$(pwd)/scripts/airflow3/venv-af3"
 
 # Activate the virtual environment
 echo "Activating virtual environment..."
-source "$(pwd)/scripts/airflow3/venv/bin/activate"
+source "$(pwd)/scripts/airflow3/venv-af3/bin/activate"
 
 # Install dependencies in the virtual environment
 echo "Installing dependencies..."
-uv pip install -r "$(pwd)/scripts/airflow3/requirements.txt"
+pip3 install --pre -r "$(pwd)/scripts/airflow3/requirements.txt"
 
-# Install Cosmos
-uv pip install build
-rm -rf "$(pwd)/dist/"
-python3 -m build
-for wheel in "$(pwd)"/dist/*.whl; do
-    uv pip install "$wheel"
-done
-
-uv pip install dbt-core
-uv pip install dbt-postgres
+pip3 install ".[test]"
 
 echo "UV virtual environment setup and dependencies installed successfully!"