Allen Porter b3d9908cd9
Add AI task structured output (#148083)
Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>
2025-07-04 15:03:34 +02:00

83 lines
2.1 KiB
Python

"""AI tasks to be handled by agents."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
import voluptuous as vol
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from .const import DATA_COMPONENT, DATA_PREFERENCES, AITaskEntityFeature
async def async_generate_data(
hass: HomeAssistant,
*,
task_name: str,
entity_id: str | None = None,
instructions: str,
structure: vol.Schema | None = None,
) -> GenDataTaskResult:
"""Run a task in the AI Task integration."""
if entity_id is None:
entity_id = hass.data[DATA_PREFERENCES].gen_data_entity_id
if entity_id is None:
raise HomeAssistantError("No entity_id provided and no preferred entity set")
entity = hass.data[DATA_COMPONENT].get_entity(entity_id)
if entity is None:
raise HomeAssistantError(f"AI Task entity {entity_id} not found")
if AITaskEntityFeature.GENERATE_DATA not in entity.supported_features:
raise HomeAssistantError(
f"AI Task entity {entity_id} does not support generating data"
)
return await entity.internal_async_generate_data(
GenDataTask(
name=task_name,
instructions=instructions,
structure=structure,
)
)
@dataclass(slots=True)
class GenDataTask:
"""Gen data task to be processed."""
name: str
"""Name of the task."""
instructions: str
"""Instructions on what needs to be done."""
structure: vol.Schema | None = None
"""Optional structure for the data to be generated."""
def __str__(self) -> str:
"""Return task as a string."""
return f"<GenDataTask {self.name}: {id(self)}>"
@dataclass(slots=True)
class GenDataTaskResult:
"""Result of gen data task."""
conversation_id: str
"""Unique identifier for the conversation."""
data: Any
"""Data generated by the task."""
def as_dict(self) -> dict[str, Any]:
"""Return result as a dict."""
return {
"conversation_id": self.conversation_id,
"data": self.data,
}