From 789eb029fa34e7add1ab0342accf6a9100cd0d12 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Thu, 3 Jul 2025 21:26:20 +0000 Subject: [PATCH 1/5] Add AI task structured output --- homeassistant/components/ai_task/__init__.py | 34 +++- homeassistant/components/ai_task/const.py | 5 + .../components/ai_task/services.yaml | 9 ++ homeassistant/components/ai_task/strings.json | 4 + homeassistant/components/ai_task/task.py | 16 ++ homeassistant/helpers/service.py | 2 + tests/components/ai_task/conftest.py | 25 ++- tests/components/ai_task/test_entity.py | 39 +++++ tests/components/ai_task/test_init.py | 152 +++++++++++++++++- tests/components/ai_task/test_task.py | 24 ++- 10 files changed, 305 insertions(+), 5 deletions(-) diff --git a/homeassistant/components/ai_task/__init__.py b/homeassistant/components/ai_task/__init__.py index 692e5d410ae879..47a9207218eaf2 100644 --- a/homeassistant/components/ai_task/__init__.py +++ b/homeassistant/components/ai_task/__init__.py @@ -1,11 +1,12 @@ """Integration to offer AI tasks to Home Assistant.""" import logging +from typing import Any import voluptuous as vol from homeassistant.config_entries import ConfigEntry -from homeassistant.const import ATTR_ENTITY_ID +from homeassistant.const import ATTR_ENTITY_ID, CONF_DESCRIPTION, CONF_SELECTOR from homeassistant.core import ( HassJobType, HomeAssistant, @@ -14,12 +15,14 @@ SupportsResponse, callback, ) -from homeassistant.helpers import config_validation as cv, storage +from homeassistant.helpers import config_validation as cv, selector, storage from homeassistant.helpers.entity_component import EntityComponent from homeassistant.helpers.typing import UNDEFINED, ConfigType, UndefinedType from .const import ( ATTR_INSTRUCTIONS, + ATTR_REQUIRED, + ATTR_STRUCTURE, ATTR_TASK_NAME, DATA_COMPONENT, DATA_PREFERENCES, @@ -47,6 +50,29 @@ CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN) +STRUCTURE_FIELD_SCHEMA = vol.Schema( + { + vol.Optional(CONF_DESCRIPTION): str, + vol.Optional(ATTR_REQUIRED): bool, + vol.Required(CONF_SELECTOR): selector.validate_selector, + } +) + + +def _validate_structure(value: dict[str, Any]) -> vol.Schema: + """Validate the structure for the generate data task.""" + if not isinstance(value, dict): + raise vol.Invalid("Structure must be a dictionary") + fields = {} + for k, v in value.items(): + if not isinstance(v, dict): + raise vol.Invalid(f"Structure field '{k}' must be a dictionary") + field_class = vol.Required if v.get(ATTR_REQUIRED, False) else vol.Optional + fields[field_class(k, description=v.get(CONF_DESCRIPTION))] = selector.selector( + v[CONF_SELECTOR] + ) + return vol.Schema(fields) + async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: """Register the process service.""" @@ -64,6 +90,10 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: vol.Required(ATTR_TASK_NAME): cv.string, vol.Optional(ATTR_ENTITY_ID): cv.entity_id, vol.Required(ATTR_INSTRUCTIONS): cv.string, + vol.Optional(ATTR_STRUCTURE): vol.All( + vol.Schema({str: STRUCTURE_FIELD_SCHEMA}), + _validate_structure, + ), } ), supports_response=SupportsResponse.ONLY, diff --git a/homeassistant/components/ai_task/const.py b/homeassistant/components/ai_task/const.py index 8b612e90560d43..17dfcca4660042 100644 --- a/homeassistant/components/ai_task/const.py +++ b/homeassistant/components/ai_task/const.py @@ -21,6 +21,8 @@ ATTR_INSTRUCTIONS: Final = "instructions" ATTR_TASK_NAME: Final = "task_name" +ATTR_STRUCTURE: Final = "structure" +ATTR_REQUIRED: Final = "required" DEFAULT_SYSTEM_PROMPT = ( "You are a Home Assistant expert and help users with their tasks." @@ -32,3 +34,6 @@ class AITaskEntityFeature(IntFlag): GENERATE_DATA = 1 """Generate data based on instructions.""" + + GENERATE_STRUCTURED_DATA = 2 + """Generate structured data based on instructions.""" diff --git a/homeassistant/components/ai_task/services.yaml b/homeassistant/components/ai_task/services.yaml index a531ca599b16b7..eb26541a343825 100644 --- a/homeassistant/components/ai_task/services.yaml +++ b/homeassistant/components/ai_task/services.yaml @@ -17,3 +17,12 @@ generate_data: domain: ai_task supported_features: - ai_task.AITaskEntityFeature.GENERATE_DATA + structure: + advanced: true + required: false + example: '{ "name": { "selector": { "text": }, "description": "Name of the user", "required": "True" } } }, "age": { "selector": { "number": }, "description": "Age of the user" } }' + selector: + object: + filter: + supported_features: + - ai_task.AITaskEntityFeature.GENERATE_STRUCTURED_DATA diff --git a/homeassistant/components/ai_task/strings.json b/homeassistant/components/ai_task/strings.json index 877174de681fc6..92106c3baca210 100644 --- a/homeassistant/components/ai_task/strings.json +++ b/homeassistant/components/ai_task/strings.json @@ -15,6 +15,10 @@ "entity_id": { "name": "Entity ID", "description": "Entity ID to run the task on. If not provided, the preferred entity will be used." + }, + "structure": { + "name": "Structured output", + "description": "When set, the AI Task will output fields with this in structure. The structure is a dictionary where the keys are the field names and the values contain a 'description', a 'selector', and an optional 'required' field." } } } diff --git a/homeassistant/components/ai_task/task.py b/homeassistant/components/ai_task/task.py index 2e546897602356..b401091c6e4084 100644 --- a/homeassistant/components/ai_task/task.py +++ b/homeassistant/components/ai_task/task.py @@ -5,6 +5,8 @@ from dataclasses import dataclass from typing import Any +import voluptuous as vol + from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError @@ -17,6 +19,7 @@ async def async_generate_data( 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: @@ -34,10 +37,20 @@ async def async_generate_data( f"AI Task entity {entity_id} does not support generating data" ) + if structure is not None: + if ( + AITaskEntityFeature.GENERATE_STRUCTURED_DATA + not in entity.supported_features + ): + raise HomeAssistantError( + f"AI Task entity {entity_id} does not support generating structured data" + ) + return await entity.internal_async_generate_data( GenDataTask( name=task_name, instructions=instructions, + structure=structure, ) ) @@ -52,6 +65,9 @@ class GenDataTask: 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"" diff --git a/homeassistant/helpers/service.py b/homeassistant/helpers/service.py index 51d9c97ceebd12..c7d4a26c86e8f9 100644 --- a/homeassistant/helpers/service.py +++ b/homeassistant/helpers/service.py @@ -86,6 +86,7 @@ def _base_components() -> dict[str, ModuleType]: """Return a cached lookup of base components.""" from homeassistant.components import ( # noqa: PLC0415 + ai_task, alarm_control_panel, assist_satellite, calendar, @@ -107,6 +108,7 @@ def _base_components() -> dict[str, ModuleType]: ) return { + "ai_task": ai_task, "alarm_control_panel": alarm_control_panel, "assist_satellite": assist_satellite, "calendar": calendar, diff --git a/tests/components/ai_task/conftest.py b/tests/components/ai_task/conftest.py index 7efbd1ffcdb628..de039b2abd1f12 100644 --- a/tests/components/ai_task/conftest.py +++ b/tests/components/ai_task/conftest.py @@ -1,5 +1,7 @@ """Test helpers for AI Task integration.""" +import json + import pytest from homeassistant.components.ai_task import ( @@ -33,7 +35,9 @@ class MockAITaskEntity(AITaskEntity): """Mock AI Task entity for testing.""" _attr_name = "Test Task Entity" - _attr_supported_features = AITaskEntityFeature.GENERATE_DATA + _attr_supported_features = ( + AITaskEntityFeature.GENERATE_DATA | AITaskEntityFeature.GENERATE_STRUCTURED_DATA + ) def __init__(self) -> None: """Initialize the mock entity.""" @@ -42,6 +46,25 @@ def __init__(self) -> None: async def _async_generate_data( self, task: GenDataTask, chat_log: ChatLog + ) -> GenDataTaskResult: + """Mock handling of generate data task.""" + self.mock_generate_data_tasks.append(task) + if task.structure is not None: + data = {"name": "Tracy Chen", "age": 30} + data_chat_log = json.dumps(data) + else: + data = "Mock result" + data_chat_log = data + chat_log.async_add_assistant_content_without_tools( + AssistantContent(self.entity_id, data_chat_log) + ) + return GenDataTaskResult( + conversation_id=chat_log.conversation_id, + data=data, + ) + + async def _async_generate_structured_data( + self, task: GenDataTask, chat_log: ChatLog, structure: dict[str, dict] ) -> GenDataTaskResult: """Mock handling of generate data task.""" self.mock_generate_data_tasks.append(task) diff --git a/tests/components/ai_task/test_entity.py b/tests/components/ai_task/test_entity.py index 3ed1c3935883f5..08f1bb42836576 100644 --- a/tests/components/ai_task/test_entity.py +++ b/tests/components/ai_task/test_entity.py @@ -1,10 +1,12 @@ """Tests for the AI Task entity model.""" from freezegun import freeze_time +import voluptuous as vol from homeassistant.components.ai_task import async_generate_data from homeassistant.const import STATE_UNKNOWN from homeassistant.core import HomeAssistant +from homeassistant.helpers import selector from .conftest import TEST_ENTITY_ID, MockAITaskEntity @@ -37,3 +39,40 @@ async def test_state_generate_data( assert mock_ai_task_entity.mock_generate_data_tasks task = mock_ai_task_entity.mock_generate_data_tasks[0] assert task.instructions == "Test prompt" + + +async def test_generate_structured_data( + hass: HomeAssistant, + init_components: None, + mock_config_entry: MockConfigEntry, + mock_ai_task_entity: MockAITaskEntity, +) -> None: + """Test the entity can generate structured data.""" + result = await async_generate_data( + hass, + task_name="Test task", + entity_id=TEST_ENTITY_ID, + instructions="Please generate a profile for a new user", + structure=vol.Schema( + { + vol.Required("name"): selector.TextSelector(), + vol.Optional("age"): selector.NumberSelector( + config=selector.NumberSelectorConfig( + min=0, + max=120, + ) + ), + } + ), + ) + # Arbitrary data returned by the mock entity (not determined by above schema in test) + assert result.data == { + "name": "Tracy Chen", + "age": 30, + } + + assert mock_ai_task_entity.mock_generate_data_tasks + task = mock_ai_task_entity.mock_generate_data_tasks[0] + assert task.instructions == "Please generate a profile for a new user" + assert task.structure + assert isinstance(task.structure, vol.Schema) diff --git a/tests/components/ai_task/test_init.py b/tests/components/ai_task/test_init.py index fdfaaccd0a4f47..dcc7290c5ae587 100644 --- a/tests/components/ai_task/test_init.py +++ b/tests/components/ai_task/test_init.py @@ -1,13 +1,17 @@ """Test initialization of the AI Task component.""" +from typing import Any + from freezegun.api import FrozenDateTimeFactory import pytest +import voluptuous as vol from homeassistant.components.ai_task import AITaskPreferences from homeassistant.components.ai_task.const import DATA_PREFERENCES from homeassistant.core import HomeAssistant +from homeassistant.helpers import selector -from .conftest import TEST_ENTITY_ID +from .conftest import TEST_ENTITY_ID, MockAITaskEntity from tests.common import flush_store @@ -82,3 +86,149 @@ async def test_generate_data_service( ) assert result["data"] == "Mock result" + + +async def test_generate_data_service_structure( + hass: HomeAssistant, + init_components: None, + mock_ai_task_entity: MockAITaskEntity, +) -> None: + """Test the entity can generate structured data.""" + result = await hass.services.async_call( + "ai_task", + "generate_data", + { + "task_name": "Profile Generation", + "instructions": "Please generate a profile for a new user", + "entity_id": TEST_ENTITY_ID, + "structure": { + "name": { + "description": "First and last name of the user such as Alice Smith", + "required": True, + "selector": {"text": {}}, + }, + "age": { + "description": "Age of the user", + "selector": { + "number": { + "min": 0, + "max": 120, + } + }, + }, + }, + }, + blocking=True, + return_response=True, + ) + # Arbitrary data returned by the mock entity (not determined by above schema in test) + assert result["data"] == { + "name": "Tracy Chen", + "age": 30, + } + + assert mock_ai_task_entity.mock_generate_data_tasks + task = mock_ai_task_entity.mock_generate_data_tasks[0] + assert task.instructions == "Please generate a profile for a new user" + assert task.structure + assert isinstance(task.structure, vol.Schema) + schema = list(task.structure.schema.items()) + assert len(schema) == 2 + + name_key, name_value = schema[0] + assert name_key == "name" + assert isinstance(name_key, vol.Required) + assert name_key.description == "First and last name of the user such as Alice Smith" + assert isinstance(name_value, selector.TextSelector) + + age_key, age_value = schema[1] + assert age_key == "age" + assert isinstance(age_key, vol.Optional) + assert age_key.description == "Age of the user" + assert isinstance(age_value, selector.NumberSelector) + assert age_value.config["min"] == 0 + assert age_value.config["max"] == 120 + + +@pytest.mark.parametrize( + ("structure", "expected_exception", "expected_error"), + [ + ( + { + "name": { + "description": "First and last name of the user such as Alice Smith", + "selector": {"invalid-selector": {}}, + }, + }, + vol.Invalid, + r"Unknown selector type invalid-selector.*", + ), + ( + { + "name": { + "description": "First and last name of the user such as Alice Smith", + "selector": { + "text": { + "extra-config": False, + } + }, + }, + }, + vol.Invalid, + r"extra keys not allowed.*", + ), + ( + { + "name": { + "description": "First and last name of the user such as Alice Smith", + }, + }, + vol.Invalid, + r"required key not provided.*selector.*", + ), + (12345, vol.Invalid, r"expected a dictionary.*"), + ("name", vol.Invalid, r"expected a dictionary.*"), + (["name"], vol.Invalid, r"expected a dictionary.*"), + ( + { + "name": { + "description": "First and last name of the user such as Alice Smith", + "selector": {"text": {}}, + "extra-fields": "Some extra fields", + }, + }, + vol.Invalid, + r"extra keys not allowed .*", + ), + ], + ids=( + "invalid-selector", + "invalid-selector-config", + "missing-selector", + "structure-is-int-not-object", + "structure-is-str-not-object", + "structure-is-list-not-object", + "extra-fields", + ), +) +async def test_generate_data_service_invalid_structure( + hass: HomeAssistant, + init_components: None, + structure: Any, + expected_exception: Exception, + expected_error: str, +) -> None: + """Test the entity can generate structured data.""" + with pytest.raises(expected_exception, match=expected_error): + await hass.services.async_call( + "ai_task", + "generate_data", + { + "task_name": "Profile Generation", + "instructions": "Please generate a profile for a new user", + "entity_id": TEST_ENTITY_ID, + "structure": structure, + }, + blocking=True, + return_response=True, + ) diff --git a/tests/components/ai_task/test_task.py b/tests/components/ai_task/test_task.py index bed760c8a1d5af..a74c3f5064793a 100644 --- a/tests/components/ai_task/test_task.py +++ b/tests/components/ai_task/test_task.py @@ -3,13 +3,14 @@ from freezegun import freeze_time import pytest from syrupy.assertion import SnapshotAssertion +import voluptuous as vol from homeassistant.components.ai_task import AITaskEntityFeature, async_generate_data from homeassistant.components.conversation import async_get_chat_log from homeassistant.const import STATE_UNKNOWN from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError -from homeassistant.helpers import chat_session +from homeassistant.helpers import chat_session, selector from .conftest import TEST_ENTITY_ID, MockAITaskEntity @@ -127,3 +128,24 @@ async def test_run_data_task_updates_chat_log( async_get_chat_log(hass, session) as chat_log, ): assert chat_log.content == snapshot + + +async def test_run_task_structure_unsupported_feature( + hass: HomeAssistant, + init_components: None, + mock_ai_task_entity: MockAITaskEntity, +) -> None: + """Test running a task with an unknown entity.""" + + mock_ai_task_entity.supported_features = AITaskEntityFeature.GENERATE_DATA + with pytest.raises( + HomeAssistantError, + match="AI Task entity ai_task.test_task_entity does not support generating structured data", + ): + await async_generate_data( + hass, + task_name="Test Task", + instructions="Test prompt", + entity_id=TEST_ENTITY_ID, + structure=vol.Schema({vol.Required("name"): selector.TextSelector()}), + ) From afa30be64b304917bd19b9db85b87f4d34af2aa9 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Thu, 3 Jul 2025 21:32:10 +0000 Subject: [PATCH 2/5] Rename _validate_structure to _validate_schema --- homeassistant/components/ai_task/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/homeassistant/components/ai_task/__init__.py b/homeassistant/components/ai_task/__init__.py index 47a9207218eaf2..9bdb5d40d31c2b 100644 --- a/homeassistant/components/ai_task/__init__.py +++ b/homeassistant/components/ai_task/__init__.py @@ -59,8 +59,8 @@ ) -def _validate_structure(value: dict[str, Any]) -> vol.Schema: - """Validate the structure for the generate data task.""" +def _validate_schema(value: dict[str, Any]) -> vol.Schema: + """Validate the structure for the generate data task and convert to a vol Schema.""" if not isinstance(value, dict): raise vol.Invalid("Structure must be a dictionary") fields = {} @@ -92,7 +92,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: vol.Required(ATTR_INSTRUCTIONS): cv.string, vol.Optional(ATTR_STRUCTURE): vol.All( vol.Schema({str: STRUCTURE_FIELD_SCHEMA}), - _validate_structure, + _validate_schema, ), } ), From 7fb7bc8e50ee2b1174b0887742bdc06e9c6d47f7 Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Thu, 3 Jul 2025 14:57:52 -0700 Subject: [PATCH 3/5] Update conftest.py to revert conftest function --- tests/components/ai_task/conftest.py | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/tests/components/ai_task/conftest.py b/tests/components/ai_task/conftest.py index de039b2abd1f12..30f11b915552ef 100644 --- a/tests/components/ai_task/conftest.py +++ b/tests/components/ai_task/conftest.py @@ -63,19 +63,6 @@ async def _async_generate_data( data=data, ) - async def _async_generate_structured_data( - self, task: GenDataTask, chat_log: ChatLog, structure: dict[str, dict] - ) -> GenDataTaskResult: - """Mock handling of generate data task.""" - self.mock_generate_data_tasks.append(task) - chat_log.async_add_assistant_content_without_tools( - AssistantContent(self.entity_id, "Mock result") - ) - return GenDataTaskResult( - conversation_id=chat_log.conversation_id, - data="Mock result", - ) - @pytest.fixture def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry: From 5ba71a467530c11e9a740b992c7b6330e3b536fe Mon Sep 17 00:00:00 2001 From: Allen Porter Date: Fri, 4 Jul 2025 00:35:03 +0000 Subject: [PATCH 4/5] Forbid extra fields in the vol schema to ensure generated output is correct --- homeassistant/components/ai_task/__init__.py | 10 ++++------ tests/components/ai_task/test_init.py | 21 +++++++++++++++----- 2 files changed, 20 insertions(+), 11 deletions(-) diff --git a/homeassistant/components/ai_task/__init__.py b/homeassistant/components/ai_task/__init__.py index 9bdb5d40d31c2b..95c080cc4724e6 100644 --- a/homeassistant/components/ai_task/__init__.py +++ b/homeassistant/components/ai_task/__init__.py @@ -59,19 +59,17 @@ ) -def _validate_schema(value: dict[str, Any]) -> vol.Schema: - """Validate the structure for the generate data task and convert to a vol Schema.""" +def _validate_structure_fields(value: dict[str, Any]) -> vol.Schema: + """Validate the structure fields as a voluptuous Schema.""" if not isinstance(value, dict): raise vol.Invalid("Structure must be a dictionary") fields = {} for k, v in value.items(): - if not isinstance(v, dict): - raise vol.Invalid(f"Structure field '{k}' must be a dictionary") field_class = vol.Required if v.get(ATTR_REQUIRED, False) else vol.Optional fields[field_class(k, description=v.get(CONF_DESCRIPTION))] = selector.selector( v[CONF_SELECTOR] ) - return vol.Schema(fields) + return vol.Schema(fields, extra=vol.PREVENT_EXTRA) async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: @@ -92,7 +90,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: vol.Required(ATTR_INSTRUCTIONS): cv.string, vol.Optional(ATTR_STRUCTURE): vol.All( vol.Schema({str: STRUCTURE_FIELD_SCHEMA}), - _validate_schema, + _validate_structure_fields, ), } ), diff --git a/tests/components/ai_task/test_init.py b/tests/components/ai_task/test_init.py index dcc7290c5ae587..02980cd699d31b 100644 --- a/tests/components/ai_task/test_init.py +++ b/tests/components/ai_task/test_init.py @@ -88,12 +88,12 @@ async def test_generate_data_service( assert result["data"] == "Mock result" -async def test_generate_data_service_structure( +async def test_generate_data_service_structure_fields( hass: HomeAssistant, init_components: None, mock_ai_task_entity: MockAITaskEntity, ) -> None: - """Test the entity can generate structured data.""" + """Test the entity can generate structured data with a top level object schemea.""" result = await hass.services.async_call( "ai_task", "generate_data", @@ -186,9 +186,9 @@ async def test_generate_data_service_structure( vol.Invalid, r"required key not provided.*selector.*", ), - (12345, vol.Invalid, r"expected a dictionary.*"), - ("name", vol.Invalid, r"expected a dictionary.*"), - (["name"], vol.Invalid, r"expected a dictionary.*"), + (12345, vol.Invalid, r"xpected a dictionary.*"), + ("name", vol.Invalid, r"xpected a dictionary.*"), + (["name"], vol.Invalid, r"xpected a dictionary.*"), ( { "name": { @@ -200,6 +200,16 @@ async def test_generate_data_service_structure( vol.Invalid, r"extra keys not allowed .*", ), + ( + { + "name": { + "description": "First and last name of the user such as Alice Smith", + "selector": "invalid-schema", + }, + }, + vol.Invalid, + r"xpected a dictionary for dictionary.", + ), ], ids=( "invalid-selector", @@ -209,6 +219,7 @@ async def test_generate_data_service_structure( "structure-is-str-not-object", "structure-is-list-not-object", "extra-fields", + "invalid-selector-schema", ), ) async def test_generate_data_service_invalid_structure( From 9bd7ea78f099673a14a44736de9bd5cb77bcb093 Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Fri, 4 Jul 2025 13:22:48 +0200 Subject: [PATCH 5/5] Remove AITaskEntityFeature.GENERATE_STRUCTURED_DATA feature flag MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The structured data generation functionality is now available to all entities that support GENERATE_DATA. This simplifies the API by removing an unnecessary feature flag while maintaining all functionality. - Remove GENERATE_STRUCTURED_DATA from AITaskEntityFeature enum - Remove feature check in task.py - Update services.yaml to remove filter - Update tests to reflect the change 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude --- homeassistant/components/ai_task/const.py | 3 --- .../components/ai_task/services.yaml | 3 --- homeassistant/components/ai_task/task.py | 9 ------- tests/components/ai_task/conftest.py | 4 +--- tests/components/ai_task/test_task.py | 24 +------------------ 5 files changed, 2 insertions(+), 41 deletions(-) diff --git a/homeassistant/components/ai_task/const.py b/homeassistant/components/ai_task/const.py index 17dfcca4660042..fa8702ed69efa1 100644 --- a/homeassistant/components/ai_task/const.py +++ b/homeassistant/components/ai_task/const.py @@ -34,6 +34,3 @@ class AITaskEntityFeature(IntFlag): GENERATE_DATA = 1 """Generate data based on instructions.""" - - GENERATE_STRUCTURED_DATA = 2 - """Generate structured data based on instructions.""" diff --git a/homeassistant/components/ai_task/services.yaml b/homeassistant/components/ai_task/services.yaml index eb26541a343825..d55b0e60faceff 100644 --- a/homeassistant/components/ai_task/services.yaml +++ b/homeassistant/components/ai_task/services.yaml @@ -23,6 +23,3 @@ generate_data: example: '{ "name": { "selector": { "text": }, "description": "Name of the user", "required": "True" } } }, "age": { "selector": { "number": }, "description": "Age of the user" } }' selector: object: - filter: - supported_features: - - ai_task.AITaskEntityFeature.GENERATE_STRUCTURED_DATA diff --git a/homeassistant/components/ai_task/task.py b/homeassistant/components/ai_task/task.py index b401091c6e4084..b6defbfad31024 100644 --- a/homeassistant/components/ai_task/task.py +++ b/homeassistant/components/ai_task/task.py @@ -37,15 +37,6 @@ async def async_generate_data( f"AI Task entity {entity_id} does not support generating data" ) - if structure is not None: - if ( - AITaskEntityFeature.GENERATE_STRUCTURED_DATA - not in entity.supported_features - ): - raise HomeAssistantError( - f"AI Task entity {entity_id} does not support generating structured data" - ) - return await entity.internal_async_generate_data( GenDataTask( name=task_name, diff --git a/tests/components/ai_task/conftest.py b/tests/components/ai_task/conftest.py index 30f11b915552ef..e80e70ddaedf2b 100644 --- a/tests/components/ai_task/conftest.py +++ b/tests/components/ai_task/conftest.py @@ -35,9 +35,7 @@ class MockAITaskEntity(AITaskEntity): """Mock AI Task entity for testing.""" _attr_name = "Test Task Entity" - _attr_supported_features = ( - AITaskEntityFeature.GENERATE_DATA | AITaskEntityFeature.GENERATE_STRUCTURED_DATA - ) + _attr_supported_features = AITaskEntityFeature.GENERATE_DATA def __init__(self) -> None: """Initialize the mock entity.""" diff --git a/tests/components/ai_task/test_task.py b/tests/components/ai_task/test_task.py index a74c3f5064793a..bed760c8a1d5af 100644 --- a/tests/components/ai_task/test_task.py +++ b/tests/components/ai_task/test_task.py @@ -3,14 +3,13 @@ from freezegun import freeze_time import pytest from syrupy.assertion import SnapshotAssertion -import voluptuous as vol from homeassistant.components.ai_task import AITaskEntityFeature, async_generate_data from homeassistant.components.conversation import async_get_chat_log from homeassistant.const import STATE_UNKNOWN from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError -from homeassistant.helpers import chat_session, selector +from homeassistant.helpers import chat_session from .conftest import TEST_ENTITY_ID, MockAITaskEntity @@ -128,24 +127,3 @@ async def test_run_data_task_updates_chat_log( async_get_chat_log(hass, session) as chat_log, ): assert chat_log.content == snapshot - - -async def test_run_task_structure_unsupported_feature( - hass: HomeAssistant, - init_components: None, - mock_ai_task_entity: MockAITaskEntity, -) -> None: - """Test running a task with an unknown entity.""" - - mock_ai_task_entity.supported_features = AITaskEntityFeature.GENERATE_DATA - with pytest.raises( - HomeAssistantError, - match="AI Task entity ai_task.test_task_entity does not support generating structured data", - ): - await async_generate_data( - hass, - task_name="Test Task", - instructions="Test prompt", - entity_id=TEST_ENTITY_ID, - structure=vol.Schema({vol.Required("name"): selector.TextSelector()}), - ) pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy