Reset to current local state

This commit is contained in:
SMKRV
2024-11-18 11:35:23 +03:00
parent 6c610ade76
commit b6f6b1e09e
3 changed files with 126 additions and 69 deletions
+126 -69
View File
@@ -8,6 +8,7 @@ from homeassistant.config_entries import ConfigEntry
from homeassistant.const import CONF_API_KEY, Platform
from homeassistant.core import HomeAssistant, ServiceCall
import homeassistant.helpers.config_validation as cv
from homeassistant.exceptions import HomeAssistantError
from .const import (
DOMAIN,
@@ -26,74 +27,96 @@ from .coordinator import HATextAICoordinator
_LOGGER = logging.getLogger(__name__)
async def async_setup(hass: HomeAssistant, config: dict) -> bool:
async def async_setup(hass: HomeAssistant, config: dict[str, Any]) -> bool:
"""Set up the HA text AI component from configuration.yaml."""
hass.data.setdefault(DOMAIN, {})
return True
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up HA text AI from a config entry."""
coordinator = HATextAICoordinator(
hass,
api_key=entry.data[CONF_API_KEY],
endpoint=entry.data.get(CONF_API_ENDPOINT),
model=entry.data.get(CONF_MODEL),
temperature=entry.data.get(CONF_TEMPERATURE),
max_tokens=entry.data.get(CONF_MAX_TOKENS),
request_interval=entry.data.get(CONF_REQUEST_INTERVAL),
)
await coordinator.async_config_entry_first_refresh()
hass.data.setdefault(DOMAIN, {})
hass.data[DOMAIN][entry.entry_id] = coordinator
await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
async def async_ask_question(call: ServiceCall) -> None:
"""Handle the ask_question service call."""
question = call.data["question"]
model = call.data.get("model", coordinator.model)
temperature = call.data.get("temperature", coordinator.temperature)
max_tokens = call.data.get("max_tokens", coordinator.max_tokens)
# Temporarily update parameters if they were overridden
original_model = coordinator.model
original_temperature = coordinator.temperature
original_max_tokens = coordinator.max_tokens
"""Handle the ask_question service call.
Args:
call: Service call containing question and optional parameters.
"""
try:
coordinator.model = model
coordinator.temperature = temperature
coordinator.max_tokens = max_tokens
await coordinator.async_ask_question(question)
finally:
# Restore original parameters
coordinator.model = original_model
coordinator.temperature = original_temperature
coordinator.max_tokens = original_max_tokens
# Get the coordinator from the first config entry
if not hass.data[DOMAIN]:
raise HomeAssistantError("No AI Text integration configured")
coordinator = next(iter(hass.data[DOMAIN].values()))
question = call.data["question"]
model = call.data.get("model", coordinator.model)
temperature = call.data.get("temperature", coordinator.temperature)
max_tokens = call.data.get("max_tokens", coordinator.max_tokens)
# Temporarily update parameters if they were overridden
original_model = coordinator.model
original_temperature = coordinator.temperature
original_max_tokens = coordinator.max_tokens
try:
coordinator.model = model
coordinator.temperature = temperature
coordinator.max_tokens = max_tokens
await coordinator.async_ask_question(question)
finally:
# Restore original parameters
coordinator.model = original_model
coordinator.temperature = original_temperature
coordinator.max_tokens = original_max_tokens
except Exception as ex:
_LOGGER.error("Error asking question: %s", str(ex))
raise HomeAssistantError(f"Failed to ask question: {str(ex)}")
async def async_clear_history(call: ServiceCall) -> None:
"""Handle the clear_history service call."""
coordinator._responses.clear()
await coordinator.async_refresh()
try:
if not hass.data[DOMAIN]:
raise HomeAssistantError("No AI Text integration configured")
async def async_get_history(call: ServiceCall) -> None:
"""Handle the get_history service call."""
limit = call.data.get("limit", 10)
history = list(coordinator._responses.items())[-limit:]
return {
"history": [
{"question": q, "response": r} for q, r in history
]
}
coordinator = next(iter(hass.data[DOMAIN].values()))
coordinator._responses.clear()
await coordinator.async_refresh()
except Exception as ex:
_LOGGER.error("Error clearing history: %s", str(ex))
raise HomeAssistantError(f"Failed to clear history: {str(ex)}")
async def async_get_history(call: ServiceCall) -> dict[str, list]:
"""Handle the get_history service call.
Returns:
Dictionary containing chat history.
"""
try:
if not hass.data[DOMAIN]:
raise HomeAssistantError("No AI Text integration configured")
coordinator = next(iter(hass.data[DOMAIN].values()))
limit = call.data.get("limit", 10)
history = list(coordinator._responses.items())[-limit:]
return {
"history": [
{"question": q, "response": r} for q, r in history
]
}
except Exception as ex:
_LOGGER.error("Error getting history: %s", str(ex))
raise HomeAssistantError(f"Failed to get history: {str(ex)}")
async def async_set_system_prompt(call: ServiceCall) -> None:
"""Handle the set_system_prompt service call."""
prompt = call.data["prompt"]
coordinator.system_prompt = prompt
try:
if not hass.data[DOMAIN]:
raise HomeAssistantError("No AI Text integration configured")
# Register all services
coordinator = next(iter(hass.data[DOMAIN].values()))
prompt = call.data["prompt"]
coordinator.system_prompt = prompt
except Exception as ex:
_LOGGER.error("Error setting system prompt: %s", str(ex))
raise HomeAssistantError(f"Failed to set system prompt: {str(ex)}")
# Register services
hass.services.async_register(
DOMAIN,
SERVICE_ASK_QUESTION,
@@ -101,8 +124,12 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
schema=vol.Schema({
vol.Required("question"): cv.string,
vol.Optional("model"): cv.string,
vol.Optional("temperature"): vol.Coerce(float),
vol.Optional("max_tokens"): vol.Coerce(int),
vol.Optional("temperature"): vol.All(
vol.Coerce(float), vol.Range(min=0, max=2)
),
vol.Optional("max_tokens"): vol.All(
vol.Coerce(int), vol.Range(min=1, max=4096)
),
})
)
@@ -118,7 +145,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
SERVICE_GET_HISTORY,
async_get_history,
schema=vol.Schema({
vol.Optional("limit"): vol.Coerce(int),
vol.Optional("limit", default=10): vol.All(
vol.Coerce(int), vol.Range(min=1)
),
})
)
@@ -133,18 +162,46 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return True
async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up HA text AI from a config entry."""
try:
coordinator = HATextAICoordinator(
hass,
api_key=entry.data[CONF_API_KEY],
endpoint=entry.data.get(CONF_API_ENDPOINT),
model=entry.data.get(CONF_MODEL),
temperature=entry.data.get(CONF_TEMPERATURE),
max_tokens=entry.data.get(CONF_MAX_TOKENS),
request_interval=entry.data.get(CONF_REQUEST_INTERVAL),
)
await coordinator.async_config_entry_first_refresh()
hass.data[DOMAIN][entry.entry_id] = coordinator
return await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
except Exception as ex:
_LOGGER.error("Error setting up entry: %s", str(ex))
raise ConfigEntryNotReady from ex
async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload a config entry."""
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
if unload_ok:
hass.data[DOMAIN].pop(entry.entry_id)
# Unregister services
for service in [
SERVICE_ASK_QUESTION,
SERVICE_CLEAR_HISTORY,
SERVICE_GET_HISTORY,
SERVICE_SET_SYSTEM_PROMPT
]:
hass.services.async_remove(DOMAIN, service)
try:
unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
if unload_ok:
hass.data[DOMAIN].pop(entry.entry_id)
return unload_ok
# Only remove services if this is the last entry
if not hass.data[DOMAIN]:
for service in [
SERVICE_ASK_QUESTION,
SERVICE_CLEAR_HISTORY,
SERVICE_GET_HISTORY,
SERVICE_SET_SYSTEM_PROMPT
]:
hass.services.async_remove(DOMAIN, service)
return unload_ok
except Exception as ex:
_LOGGER.error("Error unloading entry: %s", str(ex))
return False
BIN
View File
Binary file not shown.
BIN
View File
Binary file not shown.