diff --git a/providers/imap/src/airflow/providers/imap/hooks/imap.py b/providers/imap/src/airflow/providers/imap/hooks/imap.py
index f14b720137f78..413d557b64f03 100644
--- a/providers/imap/src/airflow/providers/imap/hooks/imap.py
+++ b/providers/imap/src/airflow/providers/imap/hooks/imap.py
@@ -32,11 +32,7 @@
 from typing import TYPE_CHECKING, Any
 
 from airflow.exceptions import AirflowException
-
-try:
-    from airflow.sdk import BaseHook
-except ImportError:
-    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+from airflow.providers.imap.version_compat import BaseHook
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
diff --git a/providers/imap/src/airflow/providers/imap/sensors/imap_attachment.py b/providers/imap/src/airflow/providers/imap/sensors/imap_attachment.py
index 7f6bbb55d9e81..cee07c6ce2379 100644
--- a/providers/imap/src/airflow/providers/imap/sensors/imap_attachment.py
+++ b/providers/imap/src/airflow/providers/imap/sensors/imap_attachment.py
@@ -23,12 +23,7 @@
 from typing import TYPE_CHECKING
 
 from airflow.providers.imap.hooks.imap import ImapHook
-from airflow.providers.imap.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseSensorOperator
-else:
-    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
+from airflow.providers.imap.version_compat import BaseSensorOperator
 
 if TYPE_CHECKING:
     try:
diff --git a/providers/imap/src/airflow/providers/imap/version_compat.py b/providers/imap/src/airflow/providers/imap/version_compat.py
index 48d122b669696..2002927fcf990 100644
--- a/providers/imap/src/airflow/providers/imap/version_compat.py
+++ b/providers/imap/src/airflow/providers/imap/version_compat.py
@@ -33,3 +33,16 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]:
 
 
 AIRFLOW_V_3_0_PLUS = get_base_airflow_version_tuple() >= (3, 0, 0)
+AIRFLOW_V_3_1_PLUS: bool = get_base_airflow_version_tuple() >= (3, 1, 0)
+
+if AIRFLOW_V_3_1_PLUS:
+    from airflow.sdk import BaseHook
+else:
+    from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
+
+if AIRFLOW_V_3_0_PLUS:
+    from airflow.sdk import BaseSensorOperator
+else:
+    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
+
+__all__ = ["AIRFLOW_V_3_0_PLUS", "AIRFLOW_V_3_1_PLUS", "BaseHook", "BaseSensorOperator"]