diff --git a/custom_components/average/sensor.py b/custom_components/average/sensor.py
index 9d1f6d7..61cf9a6 100644
--- a/custom_components/average/sensor.py
+++ b/custom_components/average/sensor.py
@@ -17,7 +17,6 @@
 import homeassistant.util.dt as dt_util
 import voluptuous as vol
 from _sha1 import sha1
-from homeassistant.components import history
 from homeassistant.components.climate import DOMAIN as CLIMATE_DOMAIN
 from homeassistant.components.group import expand_entity_ids
 from homeassistant.components.water_heater import DOMAIN as WATER_HEATER_DOMAIN
@@ -59,8 +58,12 @@
 )
 
 try:
+    # HA version >=2021.6
+    from homeassistant.components.recorder import history
     from homeassistant.components.recorder.models import LazyState
 except ImportError:
+    # HA version <=2021.5
+    from homeassistant.components import history
     from homeassistant.components.history import LazyState