Files
ha-text-ai/custom_components/ha_text_ai/sensor.py
T
SMKRV 6e635f7e2f fix: v2.5.1 security/ML/reliability patch
Security:
- H1 DNS rebinding TOCTOU: validate_endpoint returns (endpoint, resolved_ips);
  create_pinned_session builds an aiohttp session with a custom resolver that
  returns only the pre-validated IPs, closing the re-resolve gap.
- H2/M3 Shared session cookie pollution: integration no longer uses HA shared
  clientsession; isolated session with DummyCookieJar prevents cross-integration
  cookie leaks.
- Cloud metadata / link-local block added to allow_local_network mode to
  prevent IMDS exfiltration on cloud VMs.
- L3 hard cap extended to history.async_get_history (not only service schema).

ML/LLM correctness:
- _is_openai_reasoning_model uses regex with gpt-5-chat* blacklist and handles
  OpenRouter-style openai/ prefix. Future o5/gpt-6 auto-recognized.
- reasoning_effort=minimal for gpt-5 family, low for o-series.
- Gemini 2.5 Pro gets thinking_budget=128 (Pro rejects 0, flash accepts 0).
- Anthropic extracts first type=text block instead of hardcoded content[0].
- /no_think dedup uses word-boundary regex instead of substring.
- DeepSeek-reasoner detected: skips /no_think; preserves reasoning_content.

Reliability:
- Exception chaining added to Gemini-block re-raises.
- Top-level imports for json/re/hashlib (no more lazy stdlib imports).
- allow_local_network logged at INFO, not WARNING on every setup.

Code style:
- PEP 604 type hints across all .py files (dict/list/| None).

No breaking changes for users. Internal API: validate_endpoint return type
changed from str to tuple[str, list[str]].
2026-04-17 01:58:16 +03:00

366 lines
13 KiB
Python

"""
Sensor platform for HA Text AI.
@license: MIT (https://opensource.org/licenses/MIT)
@author: SMKRV
@github: https://github.com/smkrv/ha-text-ai
@source: https://github.com/smkrv/ha-text-ai
"""
from __future__ import annotations
import logging
import math
from typing import Any
from homeassistant.components.sensor import (
SensorEntity,
SensorEntityDescription,
)
from homeassistant.config_entries import ConfigEntry
from homeassistant.core import HomeAssistant
from homeassistant.helpers.device_registry import DeviceEntryType, DeviceInfo
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 homeassistant.util import slugify
from .const import (
DOMAIN,
CONF_MODEL,
CONF_API_PROVIDER,
ATTR_TOTAL_RESPONSES,
ATTR_TOTAL_ERRORS,
ATTR_AVG_RESPONSE_TIME,
ATTR_LAST_REQUEST_TIME,
ATTR_LAST_ERROR,
ATTR_IS_PROCESSING,
ATTR_IS_RATE_LIMITED,
ATTR_IS_MAINTENANCE,
ATTR_API_VERSION,
ATTR_ENDPOINT_STATUS,
ATTR_PERFORMANCE_METRICS,
ATTR_HISTORY_SIZE,
ATTR_UPTIME,
ATTR_API_PROVIDER,
ATTR_MODEL,
ATTR_SYSTEM_PROMPT,
ATTR_RESPONSE,
ATTR_QUESTION,
METRIC_TOTAL_TOKENS,
METRIC_PROMPT_TOKENS,
METRIC_COMPLETION_TOKENS,
METRIC_SUCCESSFUL_REQUESTS,
METRIC_FAILED_REQUESTS,
METRIC_AVERAGE_LATENCY,
METRIC_MAX_LATENCY,
METRIC_MIN_LATENCY,
STATE_READY,
STATE_PROCESSING,
STATE_ERROR,
STATE_INITIALIZING,
STATE_MAINTENANCE,
STATE_RATE_LIMITED,
STATE_DISCONNECTED,
ENTITY_ICON,
ENTITY_ICON_ERROR,
ENTITY_ICON_PROCESSING,
DEFAULT_NAME_PREFIX,
CONF_MAX_HISTORY_SIZE,
VERSION,
)
from .coordinator import HATextAICoordinator
from .utils import safe_log_data
_LOGGER = logging.getLogger(__name__)
# HA Recorder limit is 16384 bytes for state_attributes.
# Budget per field to stay well within the limit.
_ATTR_TEXT_LIMIT = 2048
_ATTR_PROMPT_LIMIT = 512
async def async_setup_entry(
hass: HomeAssistant,
entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up the HA Text AI sensor."""
_LOGGER.debug("Starting sensor setup for entry: %s", entry.entry_id)
try:
coordinator = hass.data[DOMAIN][entry.entry_id]
_LOGGER.debug("Found coordinator for entry %s", entry.entry_id)
instance_name = coordinator.instance_name
_LOGGER.debug("Setting up sensor with instance: %s", instance_name)
sensor = HATextAISensor(coordinator, entry)
_LOGGER.debug("Created sensor instance: %s", sensor.entity_id)
async_add_entities([sensor], True)
_LOGGER.debug("Added sensor entity: %s", sensor.entity_id)
except Exception as err:
_LOGGER.exception("Error setting up sensor: %s", err)
raise
class HATextAISensor(CoordinatorEntity, SensorEntity):
"""HA Text AI Sensor."""
coordinator: HATextAICoordinator
def __init__(
self,
coordinator: HATextAICoordinator,
config_entry: ConfigEntry,
) -> None:
"""Initialize the sensor."""
_LOGGER.debug("Initializing sensor with config entry: %s", safe_log_data(dict(config_entry.data)))
super().__init__(coordinator)
self._config_entry = config_entry
self._instance_name = coordinator.instance_name
self._normalized_name = coordinator.normalized_name
_LOGGER.debug("Instance name: %s", self._instance_name)
_LOGGER.debug("Normalized name: %s", self._normalized_name)
self._conversation_history = []
self._system_prompt = None
self._attr_has_entity_name = True
self._attr_name = self._instance_name
self.entity_id = f"sensor.ha_text_ai_{self._normalized_name}"
self._attr_unique_id = config_entry.entry_id
_LOGGER.debug("Created sensor with entity_id: %s", self.entity_id)
_LOGGER.debug("Sensor name: %s", self._attr_name)
_LOGGER.debug("Unique ID: %s", self._attr_unique_id)
self.entity_description = SensorEntityDescription(
key=f"ha_text_ai_{self._normalized_name.lower()}",
entity_registry_enabled_default=True,
)
self._current_state = STATE_INITIALIZING
self._error_count = 0
self._last_error = None
self._last_update = None
self._is_processing = False
self._last_response = {}
self._metrics = {}
model = config_entry.data.get(CONF_MODEL, "Unknown")
api_provider = config_entry.data.get(CONF_API_PROVIDER, "Unknown")
self._attr_device_info = DeviceInfo(
identifiers={(DOMAIN, self._attr_unique_id)},
name=self._attr_name,
manufacturer="Community",
model=f"{model} ({api_provider} provider)",
sw_version=VERSION,
entry_type=DeviceEntryType.SERVICE,
)
_LOGGER.debug(
"Initialized sensor: %s for instance: %s",
self.entity_id, self._instance_name,
)
@property
def available(self) -> bool:
"""Return if entity is available."""
return (
self.coordinator.last_update_success
and self.coordinator.data is not None
and self._current_state != STATE_DISCONNECTED
)
def _sanitize_value(self, value: Any) -> Any:
"""Sanitize values for JSON serialization."""
if isinstance(value, float):
if math.isinf(value) or math.isnan(value):
return None
return value
def _sanitize_attributes(self, attributes: dict[str, Any]) -> dict[str, Any]:
"""Sanitize all attributes for JSON serialization."""
sanitized = {
key: self._sanitize_value(value)
for key, value in attributes.items()
if value is not None
}
# Log metrics for debugging
metrics_keys = [
METRIC_TOTAL_TOKENS,
METRIC_PROMPT_TOKENS,
METRIC_COMPLETION_TOKENS,
METRIC_SUCCESSFUL_REQUESTS,
METRIC_FAILED_REQUESTS,
METRIC_AVERAGE_LATENCY,
METRIC_MAX_LATENCY,
METRIC_MIN_LATENCY,
]
metrics_values = {k: sanitized.get(k) for k in metrics_keys if k in sanitized}
_LOGGER.debug("Metrics for %s: %s", self.entity_id, metrics_values)
return sanitized
@property
def native_value(self) -> StateType:
"""Return the native value of the sensor."""
if not self.coordinator.last_update_success or not self.coordinator.data:
self._current_state = STATE_DISCONNECTED
return self._current_state
status = self.coordinator.data.get("state", STATE_READY)
self._current_state = status
return status
@property
def icon(self) -> str:
"""Return the icon based on the current state."""
if self._current_state == STATE_ERROR:
return ENTITY_ICON_ERROR
elif self._current_state == STATE_PROCESSING:
return ENTITY_ICON_PROCESSING
return ENTITY_ICON
@property
def extra_state_attributes(self) -> dict[str, Any]:
"""Return entity specific state attributes."""
if not self.coordinator.data:
return {}
try:
data = self.coordinator.data
metrics = data.get("metrics", {})
# Base attributes
attributes = {
ATTR_MODEL: self._config_entry.data.get(CONF_MODEL, "Unknown"),
ATTR_API_PROVIDER: self._config_entry.data.get(CONF_API_PROVIDER, "Unknown"),
ATTR_TOTAL_ERRORS: metrics.get("total_errors", 0),
"instance_name": self._instance_name,
"normalized_name": self._normalized_name,
ATTR_SYSTEM_PROMPT: (data.get("system_prompt", "")[:_ATTR_PROMPT_LIMIT]
if data.get("system_prompt") else None),
ATTR_IS_PROCESSING: data.get("is_processing", False),
ATTR_IS_RATE_LIMITED: data.get("is_rate_limited", False),
ATTR_IS_MAINTENANCE: data.get("is_maintenance", False),
ATTR_ENDPOINT_STATUS: data.get("endpoint_status", "unknown"),
ATTR_UPTIME: round(data.get("uptime", 0), 2),
ATTR_HISTORY_SIZE: data.get("history_size", 0),
}
# Conversation history preview (compact: last 3, truncated to 256 chars).
# Full history is available via ha_text_ai.get_history service.
conversation_history = data.get("conversation_history", [])
if conversation_history:
preview = conversation_history[-3:]
attributes["conversation_history"] = [
{
"timestamp": entry["timestamp"],
"question": entry["question"][:256],
"response": entry["response"][:256],
}
for entry in preview
]
# Metrics
if isinstance(metrics, dict):
attributes.update({
METRIC_TOTAL_TOKENS: metrics.get("total_tokens", 0),
METRIC_PROMPT_TOKENS: metrics.get("prompt_tokens", 0),
METRIC_COMPLETION_TOKENS: metrics.get("completion_tokens", 0),
METRIC_SUCCESSFUL_REQUESTS: metrics.get("successful_requests", 0),
METRIC_FAILED_REQUESTS: metrics.get("failed_requests", 0),
METRIC_AVERAGE_LATENCY: round(metrics.get("average_latency", 0), 2),
METRIC_MAX_LATENCY: round(metrics.get("max_latency", 0), 2),
METRIC_MIN_LATENCY: metrics.get("min_latency", 0) or None,
})
# Last response handling
last_response = data.get("last_response", {})
if isinstance(last_response, dict):
attributes.update({
ATTR_RESPONSE: last_response.get("response", "")[:_ATTR_TEXT_LIMIT],
ATTR_QUESTION: last_response.get("question", "")[:_ATTR_TEXT_LIMIT],
"last_model": last_response.get("model", ""),
"last_timestamp": last_response.get("timestamp", ""),
"last_error": (last_response.get("error", "")[:_ATTR_TEXT_LIMIT]
if last_response.get("error") else None),
})
return self._sanitize_attributes(attributes)
except Exception as err:
_LOGGER.error("Error preparing attributes: %s", err, exc_info=True)
return {}
async def async_added_to_hass(self) -> None:
"""When entity is added to hass."""
await super().async_added_to_hass()
self._handle_coordinator_update()
_LOGGER.debug("Entity %s added to Home Assistant", self.entity_id)
def _handle_coordinator_update(self) -> None:
"""Handle updated data from the coordinator."""
try:
data = self.coordinator.data
if not self.coordinator.last_update_success or not data:
self._current_state = STATE_DISCONNECTED
_LOGGER.warning("No data available for %s", self.entity_id)
self.async_write_ha_state()
return
self._is_processing = data.get("is_processing", False)
# Update metrics
metrics = data.get("metrics", {})
if isinstance(metrics, dict):
self._metrics.update(metrics)
_LOGGER.debug("Updated metrics for %s: %s", self.entity_id, self._metrics)
# Update conversation history and system prompt
self._conversation_history = data.get("conversation_history", [])
self._system_prompt = data.get("system_prompt")
# Update state based on conditions
if self._is_processing:
self._current_state = STATE_PROCESSING
elif data.get("is_rate_limited"):
self._current_state = STATE_RATE_LIMITED
elif data.get("is_maintenance"):
self._current_state = STATE_MAINTENANCE
elif data.get("error"):
self._current_state = STATE_ERROR
self._last_error = data["error"]
self._error_count += 1
else:
self._current_state = data.get("state", STATE_READY)
# Update last update timestamp
self._last_update = dt_util.utcnow()
_LOGGER.debug(
"Updated %s state to: %s (available: %s)",
self.entity_id, self._current_state, self.available,
)
except Exception as err:
self._current_state = STATE_ERROR
self._last_error = str(err)
self._error_count += 1
_LOGGER.error(
"Error handling update for %s: %s",
self.entity_id,
err,
exc_info=True,
)
self.async_write_ha_state()