diff --git a/custom_components/ha_text_ai/sensor.py b/custom_components/ha_text_ai/sensor.py
index 655f1e9..3ac6a87 100644
--- a/custom_components/ha_text_ai/sensor.py
+++ b/custom_components/ha_text_ai/sensor.py
@@ -1,7 +1,7 @@
 """Sensor platform for HA text AI."""
 from typing import Any, Callable, Dict, Optional
 
-from homeassistant.components.sensor import SensorEntity
+from homeassistant.components.sensor import SensorEntity, SensorStateClass
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
@@ -56,4 +56,3 @@ class HATextAISensor(CoordinatorEntity, SensorEntity):
             ATTR_RESPONSE: last_response,
             ATTR_LAST_UPDATED: self.coordinator.last_update_success_time,
         }
-        
diff --git a/ha-text-ai.zip b/ha-text-ai.zip
index 49ebdb9..ab37191 100644
Binary files a/ha-text-ai.zip and b/ha-text-ai.zip differ