mirror of
https://github.com/smkrv/ha-text-ai.git
synced 2026-07-21 22:54:00 +08:00
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]].
366 lines
13 KiB
Python
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()
|