mirror of
https://github.com/home-assistant/core.git
synced 2025-07-07 05:17:07 +00:00
126 lines
3.5 KiB
Python
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,
|
|
}
|