AtHeartEngineering/ai_automation_suggester
1
0
Fork 0
mirror of https://github.com/AtHeartEngineering/ai_automation_suggester.git synced 2026-01-09 22:27:55 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add Model sensor

Browse Source
This commit is contained in:
RmG152
2025-05-09 09:43:03 +02:00
parent 83e4f1e428
commit dbd38f9a2b
1 changed files with 95 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
custom_components/ai_automation_suggester/sensor.py
View File

@@ -20,6 +20,16 @@ from .const import (
DEFAULT_MAX_INPUT_TOKENS,
CONF_MAX_OUTPUT_TOKENS,
DEFAULT_MAX_OUTPUT_TOKENS,
CONF_OPENAI_MODEL,
CONF_ANTHROPIC_MODEL,
CONF_GOOGLE_MODEL,
CONF_GROQ_MODEL,
CONF_LOCALAI_MODEL,
CONF_OLLAMA_MODEL,
CONF_CUSTOM_OPENAI_MODEL,
CONF_MISTRAL_MODEL,
CONF_PERPLEXITY_MODEL,
DEFAULT_MODELS,
)
_LOGGER = logging.getLogger(__name__)
@@ -51,6 +61,13 @@ MAX_OUTPUT_TOKENS_SENSOR = SensorEntityDescription(
entity_category=EntityCategory.DIAGNOSTIC,
)
MODEL_SENSOR = SensorEntityDescription(
key="model",
name="AI Model",
icon="mdi:brain",
entity_category=EntityCategory.DIAGNOSTIC,
)
async def async_setup_entry(hass, entry, async_add_entities):
"""Set up AI Automation Suggester sensors."""
@@ -61,6 +78,7 @@ async def async_setup_entry(hass, entry, async_add_entities):
AIProviderStatusSensor(coordinator, entry, STATUS_SENSOR),
InputTokensSensor(coordinator, entry, MAX_INPUT_TOKENS_SENSOR),
OutputTokensSensor(coordinator, entry, MAX_OUTPUT_TOKENS_SENSOR),
AIModelSensor(coordinator, entry, MODEL_SENSOR),
]
async_add_entities(entities, True)
_LOGGER.debug("Sensor platform setup complete")
@@ -83,7 +101,7 @@ class AISuggestionsSensor(CoordinatorEntity, SensorEntity):
self._attr_unique_id = f"{entry.entry_id}_{description.key}"
self._attr_device_info = {
"identifiers": {(DOMAIN, entry.entry_id)},
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, 'unknown')})",
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, "unknown")})",
"manufacturer": "Community",
"model": entry.data.get(CONF_PROVIDER, "unknown"),
"sw_version": entry.version,
@@ -154,7 +172,7 @@ class AIProviderStatusSensor(CoordinatorEntity, SensorEntity):
self._attr_unique_id = f"{entry.entry_id}_{description.key}"
self._attr_device_info = {
"identifiers": {(DOMAIN, entry.entry_id)},
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, 'unknown')})",
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, "unknown")})",
"manufacturer": "Community",
"model": entry.data.get(CONF_PROVIDER, "unknown"),
"sw_version": entry.version,
@@ -207,7 +225,7 @@ class InputTokensSensor(CoordinatorEntity, SensorEntity):
self._attr_unique_id = f"{entry.entry_id}_{description.key}"
self._attr_device_info = {
"identifiers": {(DOMAIN, entry.entry_id)},
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, 'unknown')})",
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, "unknown")})",
"manufacturer": "Community",
"model": entry.data.get(CONF_PROVIDER, "unknown"),
"sw_version": entry.version,
@@ -246,7 +264,7 @@ class OutputTokensSensor(CoordinatorEntity, SensorEntity):
self._attr_unique_id = f"{entry.entry_id}_{description.key}"
self._attr_device_info = {
"identifiers": {(DOMAIN, entry.entry_id)},
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, 'unknown')})",
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, "unknown")})",
"manufacturer": "Community",
"model": entry.data.get(CONF_PROVIDER, "unknown"),
"sw_version": entry.version,
@@ -270,4 +288,76 @@ class OutputTokensSensor(CoordinatorEntity, SensorEntity):
CONF_MAX_OUTPUT_TOKENS,
self._entry.data.get(CONF_MAX_OUTPUT_TOKENS, DEFAULT_MAX_OUTPUT_TOKENS)
)
self.async_write_ha_state()
self.async_write_ha_state()
# ─────────────────────────────────────────────────────────────
# Model Sensor
# ─────────────────────────────────────────────────────────────
class AIModelSensor(CoordinatorEntity, SensorEntity):
"""Shows the currently configured AI model."""
def __init__(
self,
coordinator: DataUpdateCoordinator,
entry,
description: SensorEntityDescription,
) -> None:
super().__init__(coordinator)
self.entity_description = description
self._attr_unique_id = f"{entry.entry_id}_{description.key}"
self._attr_device_info = {
"identifiers": {(DOMAIN, entry.entry_id)},
"name": f"AI Automation Suggester ({entry.data.get(CONF_PROVIDER, "unknown")})",
"manufacturer": "Community",
"model": entry.data.get(CONF_PROVIDER, "unknown"),
"sw_version": entry.version,
}
self._entry = entry
provider = entry.data.get(CONF_PROVIDER, "unknown")
# Determinamos qué clave de modelo usar según el proveedor
model_key = {
"OpenAI": CONF_OPENAI_MODEL,
"Anthropic": CONF_ANTHROPIC_MODEL,
"Google": CONF_GOOGLE_MODEL,
"Groq": CONF_GROQ_MODEL,
"LocalAI": CONF_LOCALAI_MODEL,
"Ollama": CONF_OLLAMA_MODEL,
"Custom OpenAI": CONF_CUSTOM_OPENAI_MODEL,
"Mistral AI": CONF_MISTRAL_MODEL,
"Perplexity AI": CONF_PERPLEXITY_MODEL,
}.get(provider)
# Obtenemos el modelo configurado o el valor por defecto
self._attr_native_value = entry.options.get(
model_key,
entry.data.get(model_key, DEFAULT_MODELS.get(provider, "unknown"))
) if model_key else "unknown"
@property
def name(self) -> str:
provider = self._entry.data.get(CONF_PROVIDER, "unknown")
return f"AI Model ({provider})"
@callback
def _handle_coordinator_update(self) -> None:
"""Update when options change."""
provider = self._entry.data.get(CONF_PROVIDER, "unknown")
model_key = {
"OpenAI": CONF_OPENAI_MODEL,
"Anthropic": CONF_ANTHROPIC_MODEL,
"Google": CONF_GOOGLE_MODEL,
"Groq": CONF_GROQ_MODEL,
"LocalAI": CONF_LOCALAI_MODEL,
"Ollama": CONF_OLLAMA_MODEL,
"Custom OpenAI": CONF_CUSTOM_OPENAI_MODEL,
"Mistral AI": CONF_MISTRAL_MODEL,
"Perplexity AI": CONF_PERPLEXITY_MODEL,
}.get(provider)
self._attr_native_value = self._entry.options.get(
model_key,
self._entry.data.get(model_key, DEFAULT_MODELS.get(provider, "unknown"))
) if model_key else "unknown"
self.async_write_ha_state()