Release v1.0.3

custom_components/ha_text_ai/sensor.py

@@ -13,8 +13,18 @@ from homeassistant.core import HomeAssistant
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.helpers.typing import StateType
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
+from homeassistant.util import dt as dt_util
 
-from .const import DOMAIN, ATTR_QUESTION, ATTR_RESPONSE, ATTR_LAST_UPDATED
+from .const import (
+    DOMAIN,
+    ATTR_QUESTION,
+    ATTR_RESPONSE,
+    ATTR_LAST_UPDATED,
+    ATTR_MODEL,
+    ATTR_TEMPERATURE,
+    ATTR_MAX_TOKENS,
+    ATTR_TOTAL_RESPONSES,
+)
 from .coordinator import HATextAICoordinator
 
 _LOGGER = logging.getLogger(__name__)
@@ -46,42 +56,71 @@ class HATextAISensor(CoordinatorEntity, SensorEntity):
         self._config_entry = config_entry
         self._attr_unique_id = f"{config_entry.entry_id}"
         self._attr_name = "Last Response"
+        self._attr_suggested_display_precision = 0
 
     @property
     def state(self) -> StateType:
         """Return the state of the sensor."""
-        if not self.coordinator.data:
+        if not self.coordinator.data or not self.coordinator.last_update_success_time:
             return None
+
+        # Convert to local time
+        if isinstance(self.coordinator.last_update_success_time, datetime):
+            return dt_util.as_local(self.coordinator.last_update_success_time)
         return self.coordinator.last_update_success_time
 
     @property
     def extra_state_attributes(self) -> Optional[Dict[str, Any]]:
         """Return entity specific state attributes."""
         if not self.coordinator.data:
-            return None
+            return {
+                ATTR_TOTAL_RESPONSES: 0,
+                ATTR_MODEL: self.coordinator.model,
+                ATTR_TEMPERATURE: self.coordinator.temperature,
+                ATTR_MAX_TOKENS: self.coordinator.max_tokens,
+            }
 
         try:
-
             history = list(self.coordinator.data.items())
             if not history:
-                return None
+                return {
+                    ATTR_TOTAL_RESPONSES: 0,
+                    ATTR_MODEL: self.coordinator.model,
+                    ATTR_TEMPERATURE: self.coordinator.temperature,
+                    ATTR_MAX_TOKENS: self.coordinator.max_tokens,
+                }
 
             last_question, last_data = history[-1]
 
-
+            # Handle different response formats
             if isinstance(last_data, dict):
                 last_response = last_data.get("response", "")
+                last_updated = last_data.get("timestamp", self.coordinator.last_update_success_time)
             else:
                 last_response = str(last_data)
+                last_updated = self.coordinator.last_update_success_time
+
+            # Convert timestamp to local time if needed
+            if isinstance(last_updated, datetime):
+                last_updated = dt_util.as_local(last_updated)
 
             return {
                 ATTR_QUESTION: last_question,
                 ATTR_RESPONSE: last_response,
-                ATTR_LAST_UPDATED: self.coordinator.last_update_success_time,
+                ATTR_LAST_UPDATED: last_updated,
+                ATTR_TOTAL_RESPONSES: len(history),
+                ATTR_MODEL: self.coordinator.model,
+                ATTR_TEMPERATURE: self.coordinator.temperature,
+                ATTR_MAX_TOKENS: self.coordinator.max_tokens,
+            }
+        except Exception as err:
+            _LOGGER.error("Error getting attributes: %s", err, exc_info=True)
+            return {
+                ATTR_TOTAL_RESPONSES: 0,
+                ATTR_MODEL: self.coordinator.model,
+                ATTR_TEMPERATURE: self.coordinator.temperature,
+                ATTR_MAX_TOKENS: self.coordinator.max_tokens,
             }
-        except (IndexError, KeyError, AttributeError) as err:
-            _LOGGER.warning("Error getting attributes: %s", err)
-            return None
 
     @property
     def available(self) -> bool:
@@ -92,3 +131,8 @@ class HATextAISensor(CoordinatorEntity, SensorEntity):
     def should_poll(self) -> bool:
         """No need to poll. Coordinator notifies entity of updates."""
         return False
+
+    async def async_added_to_hass(self) -> None:
+        """When entity is added to hass."""
+        await super().async_added_to_hass()
+        self._handle_coordinator_update()