mirror of
https://github.com/smkrv/ha-text-ai.git
synced 2026-07-29 11:23:57 +08:00
Release v2.0.0-alpha
This commit is contained in:
SMKRV
@@ -0,0 +1,377 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user