2025-07-06 19:33:41 +02:00

126 lines
3.5 KiB
Python

"""AI tasks to be handled by agents."""
from __future__ import annotations
from dataclasses import dataclass, fields
from typing import Any
import voluptuous as vol
from homeassistant.components import media_source
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError
from .const import DATA_COMPONENT, DATA_PREFERENCES, AITaskEntityFeature
@dataclass(slots=True)
class PlayMediaWithId(media_source.PlayMedia):
"""Play media with a media content ID."""
media_content_id: str
"""Media source ID to play."""
def __str__(self) -> str:
"""Return media source ID as a string."""
return f"<PlayMediaWithId {self.media_content_id}>"
async def async_generate_data(
hass: HomeAssistant,
*,
task_name: str,
entity_id: str | None = None,
instructions: str,
structure: vol.Schema | None = None,
attachments: list[dict] | 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"
)
# Resolve attachments
resolved_attachments: list[PlayMediaWithId] | None = None
if attachments:
if AITaskEntityFeature.SUPPORT_ATTACHMENTS not in entity.supported_features:
raise HomeAssistantError(
f"AI Task entity {entity_id} does not support attachments"
)
resolved_attachments = []
for attachment in attachments:
media = await media_source.async_resolve_media(
hass, attachment["media_content_id"], None
)
resolved_attachments.append(
PlayMediaWithId(
**{
field.name: getattr(media, field.name)
for field in fields(media)
},
media_content_id=attachment["media_content_id"],
)
)
return await entity.internal_async_generate_data(
GenDataTask(
name=task_name,
instructions=instructions,
structure=structure,
attachments=resolved_attachments,
)
)
@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."""
attachments: list[PlayMediaWithId] | None = None
"""List of attachments to go along the instructions."""
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,
}