"""The HA Text AI coordinator.""" from __future__ import annotations import logging from datetime import datetime, timedelta from typing import Any, Dict, List, Optional from homeassistant.core import HomeAssistant from homeassistant.helpers.update_coordinator import DataUpdateCoordinator from homeassistant.util import dt as dt_util from homeassistant.exceptions import HomeAssistantError from .const import ( DOMAIN, STATE_READY, STATE_PROCESSING, STATE_ERROR, STATE_RATE_LIMITED, STATE_MAINTENANCE, DEFAULT_MAX_TOKENS, DEFAULT_TEMPERATURE, DEFAULT_MAX_HISTORY, DEFAULT_CONTEXT_MESSAGES, ) _LOGGER = logging.getLogger(__name__) class HATextAICoordinator(DataUpdateCoordinator): def __init__( self, hass: HomeAssistant, client: Any, model: str, update_interval: int, instance_name: str, max_tokens: int = DEFAULT_MAX_TOKENS, temperature: float = DEFAULT_TEMPERATURE, max_history_size: int = DEFAULT_MAX_HISTORY, context_messages: int = DEFAULT_CONTEXT_MESSAGES, is_anthropic: bool = False, ) -> None: """Initialize coordinator.""" self.instance_name = instance_name self.hass = hass self.client = client self.model = model self.temperature = temperature self.max_tokens = max_tokens self.max_history_size = max_history_size self.is_anthropic = is_anthropic # Initialize with default state self._initial_state = { "state": STATE_READY, "metrics": { "total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0, "successful_requests": 0, "failed_requests": 0, "total_errors": 0, "average_latency": 0, "max_latency": 0, "min_latency": float('inf'), }, "last_response": { "timestamp": dt_util.utcnow().isoformat(), "question": "", "response": "", "model": model, "instance": instance_name, "error": None }, "is_processing": False, "is_rate_limited": False, "is_maintenance": False, "endpoint_status": "ready", "uptime": 0, "system_prompt": None, "history_size": 0, "conversation_history": [], } update_interval_td = timedelta(seconds=update_interval) super().__init__( hass, _LOGGER, name=instance_name, update_interval=update_interval_td, ) # Register instance self.hass.data.setdefault(DOMAIN, {}) self.hass.data[DOMAIN][instance_name] = self self.context_messages = context_messages self._system_prompt = None self._conversation_history = [] self._performance_metrics = self._initial_state["metrics"].copy() self._is_processing = False self._is_rate_limited = False self._is_maintenance = False self.endpoint_status = "ready" self.last_response = self._initial_state["last_response"].copy() self._start_time = dt_util.utcnow() _LOGGER.info(f"Initialized HA Text AI coordinator with instance: {instance_name}") async def _async_update_data(self) -> Dict[str, Any]: """Update data via library.""" try: current_state = self._get_current_state() _LOGGER.debug(f"Updating data for {self.instance_name}, current state: {current_state}") data = { "state": current_state, "metrics": self._performance_metrics, "last_response": self.last_response, "is_processing": self._is_processing, "is_rate_limited": self._is_rate_limited, "is_maintenance": self._is_maintenance, "endpoint_status": self.endpoint_status, "uptime": (dt_util.utcnow() - self._start_time).total_seconds(), "system_prompt": self._system_prompt, "history_size": len(self._conversation_history), "conversation_history": self._conversation_history, } # Validate data if not isinstance(data, dict): raise ValueError("Invalid data format") _LOGGER.debug(f"Updated data for {self.instance_name}: {data}") return data except Exception as err: _LOGGER.error(f"Error updating data for {self.instance_name}: {err}") return self._initial_state async def async_update_ha_state(self) -> None: """Update Home Assistant state.""" try: _LOGGER.debug(f"Requesting state update for {self.instance_name}") await self.async_request_refresh() # Force update of all entities for entity_id in self.hass.states.async_entity_ids(): if entity_id.startswith(f"sensor.ha_text_ai_{self.instance_name}"): self.hass.states.async_set(entity_id, self._get_current_state()) except Exception as err: _LOGGER.error(f"Error updating HA state for {self.instance_name}: {err}") def _get_current_state(self) -> str: """Get current state based on internal flags.""" if self._is_processing: return STATE_PROCESSING elif self._is_rate_limited: return STATE_RATE_LIMITED elif self._is_maintenance: return STATE_MAINTENANCE elif self.last_response.get("error"): return STATE_ERROR return STATE_READY def _calculate_context_tokens(self, messages: List[Dict[str, str]]) -> int: try: if self.is_anthropic and hasattr(self.client, 'count_tokens'): return sum(self.client.count_tokens(msg['content']) for msg in messages) return sum(len(msg['content']) // 4 for msg in messages) except Exception as e: _LOGGER.warning(f"Error calculating context tokens: {e}") return 0 async def async_ask_question( self, question: str, model: Optional[str] = None, temperature: Optional[float] = None, max_tokens: Optional[int] = None, system_prompt: Optional[str] = None, context_messages: Optional[int] = None, ) -> dict: """Process a question with optional parameters.""" return await self.async_process_question( question, model, temperature, max_tokens, system_prompt, context_messages ) async def async_process_question( self, question: str, model: Optional[str] = None, temperature: Optional[float] = None, max_tokens: Optional[int] = None, system_prompt: Optional[str] = None, context_messages: Optional[int] = None, ) -> dict: temp_context_messages = context_messages or self.context_messages if not question: raise ValueError("Question cannot be empty") _LOGGER.debug(f"Processing question for instance {self.instance_name}") try: self._is_processing = True await self.async_update_ha_state() temp_model = model or self.model temp_temperature = temperature or self.temperature temp_max_tokens = max_tokens or self.max_tokens temp_system_prompt = system_prompt or self._system_prompt start_time = dt_util.utcnow() messages = [] if temp_system_prompt: if self.is_anthropic: system_content = f"\n\nHuman: {temp_system_prompt}\n\nAssistant: I understand and will follow these instructions." messages.append({"role": "user", "content": system_content}) else: messages.append({"role": "system", "content": temp_system_prompt}) # Add conversation history context_history = self._conversation_history[-temp_context_messages:] for entry in context_history: messages.append({"role": "user", "content": entry["question"]}) messages.append({"role": "assistant", "content": entry["response"]}) messages.append({"role": "user", "content": question}) kwargs = { "model": temp_model, "temperature": temp_temperature, "max_tokens": temp_max_tokens, "messages": messages, } response = await self.async_process_message(question, **kwargs) # Update metrics end_time = dt_util.utcnow() latency = (end_time - start_time).total_seconds() self._update_metrics(latency, response) # Update history self._update_history(question, response) return response except Exception as err: self._handle_error(err) raise HomeAssistantError(f"Failed to process question: {err}") finally: self._is_processing = False await self.async_update_ha_state() async def async_process_message(self, question: str, **kwargs) -> dict: """Process message using the AI client.""" try: if self.is_anthropic: response = await self._process_anthropic_message(question, **kwargs) else: response = await self._process_openai_message(question, **kwargs) self.last_response = { "timestamp": dt_util.utcnow().isoformat(), "question": question, "response": response["content"], "model": kwargs.get("model", self.model), "instance": self.instance_name, "error": None } return response except Exception as err: self._handle_error(err) raise async def _process_anthropic_message(self, question: str, **kwargs) -> dict: """Process message using Anthropic API.""" response = await self.client.messages.create( model=kwargs["model"], max_tokens=kwargs["max_tokens"], messages=kwargs["messages"], temperature=kwargs["temperature"], ) return { "content": response.content[0].text, "tokens": { "prompt": response.usage.input_tokens, "completion": response.usage.output_tokens, "total": response.usage.input_tokens + response.usage.output_tokens } } async def _process_openai_message(self, question: str, **kwargs) -> dict: """Process message using OpenAI API.""" try: response = await self.client.create( model=kwargs["model"], messages=kwargs["messages"], temperature=kwargs["temperature"], max_tokens=kwargs["max_tokens"], ) return { "content": response["choices"][0]["message"]["content"], "tokens": { "prompt": response["usage"]["prompt_tokens"], "completion": response["usage"]["completion_tokens"], "total": response["usage"]["total_tokens"] } } except Exception as e: _LOGGER.error(f"Error in OpenAI API call: {str(e)}") raise def _update_metrics(self, latency: float, response: dict) -> None: """Update performance metrics.""" metrics = self._performance_metrics tokens = response.get("tokens", {}) metrics["total_tokens"] += tokens.get("total", 0) metrics["prompt_tokens"] += tokens.get("prompt", 0) metrics["completion_tokens"] += tokens.get("completion", 0) metrics["successful_requests"] += 1 metrics["average_latency"] = ( (metrics["average_latency"] * (metrics["successful_requests"] - 1) + latency) / metrics["successful_requests"] ) metrics["max_latency"] = max(metrics["max_latency"], latency) metrics["min_latency"] = min(metrics["min_latency"], latency) def _update_history(self, question: str, response: dict) -> None: """Update conversation history.""" self._conversation_history.append({ "timestamp": dt_util.utcnow().isoformat(), "question": question, "response": response["content"] }) while len(self._conversation_history) > self.max_history_size: self._conversation_history.pop(0) def _handle_error(self, error: Exception) -> None: """Handle error and update metrics.""" self._performance_metrics["total_errors"] += 1 self._performance_metrics["failed_requests"] += 1 self.last_response = { "timestamp": dt_util.utcnow().isoformat(), "question": "", "response": "", "model": self.model, "instance": self.instance_name, "error": str(error) } async def async_clear_history(self) -> None: """Clear conversation history.""" self._conversation_history = [] await self.async_update_ha_state() async def async_get_history(self) -> List[Dict[str, str]]: """Get conversation history.""" return self._conversation_history async def async_set_system_prompt(self, prompt: str) -> None: """Set system prompt.""" self._system_prompt = prompt await self.async_update_ha_state()