diff --git a/homeassistant/components/ai_task/__init__.py b/homeassistant/components/ai_task/__init__.py index bc3b5db2cf27..c20b6cf1fb54 100644 --- a/homeassistant/components/ai_task/__init__.py +++ b/homeassistant/components/ai_task/__init__.py @@ -146,13 +146,15 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: async def async_service_generate_data(call: ServiceCall) -> ServiceResponse: """Run the data task service.""" - result = await async_generate_data(hass=call.hass, **call.data) + result = await async_generate_data( + hass=call.hass, context=call.context, **call.data + ) return result.as_dict() async def async_service_generate_image(call: ServiceCall) -> ServiceResponse: """Run the image task service.""" - return await async_generate_image(hass=call.hass, **call.data) + return await async_generate_image(hass=call.hass, context=call.context, **call.data) class AITaskPreferences: diff --git a/homeassistant/components/ai_task/entity.py b/homeassistant/components/ai_task/entity.py index 5d4f6a5c12fe..dc8f84817bb7 100644 --- a/homeassistant/components/ai_task/entity.py +++ b/homeassistant/components/ai_task/entity.py @@ -12,6 +12,7 @@ from homeassistant.components.conversation import ( async_get_chat_log, ) from homeassistant.const import STATE_UNAVAILABLE, STATE_UNKNOWN +from homeassistant.core import Context from homeassistant.helpers import llm from homeassistant.helpers.chat_session import ChatSession from homeassistant.helpers.restore_state import RestoreEntity @@ -61,6 +62,7 @@ class AITaskEntity(RestoreEntity): self, session: ChatSession, task: GenDataTask | GenImageTask, + context: Context | None, ) -> AsyncGenerator[ChatLog]: """Context manager used to manage the ChatLog used during an AI Task.""" user_llm_hass_api: llm.API | None = None @@ -78,7 +80,7 @@ class AITaskEntity(RestoreEntity): await chat_log.async_provide_llm_data( llm.LLMContext( platform=self.platform.domain, - context=None, + context=context, language=None, assistant=DOMAIN, device_id=None, @@ -98,11 +100,14 @@ class AITaskEntity(RestoreEntity): self, session: ChatSession, task: GenDataTask, + context: Context | None = None, ) -> GenDataTaskResult: """Run a gen data task.""" + if context is not None: + self.async_set_context(context) self.__last_activity = dt_util.utcnow().isoformat() self.async_write_ha_state() - async with self._async_get_ai_task_chat_log(session, task) as chat_log: + async with self._async_get_ai_task_chat_log(session, task, context) as chat_log: return await self._async_generate_data(task, chat_log) async def _async_generate_data( @@ -118,11 +123,14 @@ class AITaskEntity(RestoreEntity): self, session: ChatSession, task: GenImageTask, + context: Context | None = None, ) -> GenImageTaskResult: """Run a gen image task.""" + if context is not None: + self.async_set_context(context) self.__last_activity = dt_util.utcnow().isoformat() self.async_write_ha_state() - async with self._async_get_ai_task_chat_log(session, task) as chat_log: + async with self._async_get_ai_task_chat_log(session, task, context) as chat_log: return await self._async_generate_image(task, chat_log) async def _async_generate_image( diff --git a/homeassistant/components/ai_task/task.py b/homeassistant/components/ai_task/task.py index 245318dab133..a29d186537ba 100644 --- a/homeassistant/components/ai_task/task.py +++ b/homeassistant/components/ai_task/task.py @@ -12,7 +12,7 @@ import voluptuous as vol from homeassistant.components import camera, conversation, image, media_source from homeassistant.components.http.auth import async_sign_path -from homeassistant.core import HomeAssistant, ServiceResponse, callback +from homeassistant.core import Context, HomeAssistant, ServiceResponse, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import llm from homeassistant.helpers.chat_session import ChatSession, async_get_chat_session @@ -119,6 +119,7 @@ async def async_generate_data( structure: vol.Schema | None = None, attachments: list[dict] | None = None, llm_api: llm.API | None = None, + context: Context | None = None, ) -> GenDataTaskResult: """Run a data generation task in the AI Task integration.""" if entity_id is None: @@ -156,6 +157,7 @@ async def async_generate_data( attachments=resolved_attachments or None, llm_api=llm_api, ), + context, ) @@ -166,6 +168,7 @@ async def async_generate_image( entity_id: str | None = None, instructions: str, attachments: list[dict] | None = None, + context: Context | None = None, ) -> ServiceResponse: """Run an image generation task in the AI Task integration.""" if entity_id is None: @@ -201,6 +204,7 @@ async def async_generate_image( instructions=instructions, attachments=resolved_attachments or None, ), + context, ) service_result = task_result.as_dict() diff --git a/tests/components/ai_task/conftest.py b/tests/components/ai_task/conftest.py index ceffb7c055e8..88d231afc671 100644 --- a/tests/components/ai_task/conftest.py +++ b/tests/components/ai_task/conftest.py @@ -48,12 +48,14 @@ class MockAITaskEntity(AITaskEntity): super().__init__() self.mock_generate_data_tasks = [] self.mock_generate_image_tasks = [] + self.mock_chat_logs = [] 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) + self.mock_chat_logs.append(chat_log) if task.structure is not None: data = {"name": "Tracy Chen", "age": 30} data_chat_log = json.dumps(data) diff --git a/tests/components/ai_task/test_init.py b/tests/components/ai_task/test_init.py index 4286cb62991e..f1d43a0ff1b3 100644 --- a/tests/components/ai_task/test_init.py +++ b/tests/components/ai_task/test_init.py @@ -15,7 +15,7 @@ from homeassistant.components.ai_task.const import ( DATA_PREFERENCES, DOMAIN, ) -from homeassistant.core import HomeAssistant +from homeassistant.core import Context, HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import selector @@ -87,6 +87,7 @@ async def test_generate_data_service( mock_ai_task_entity: MockAITaskEntity, ) -> None: """Test the generate data service.""" + context = Context() preferences = hass.data[DATA_PREFERENCES] preferences.async_set_preferences(**set_preferences) @@ -108,9 +109,11 @@ async def test_generate_data_service( | msg_extra, blocking=True, return_response=True, + context=context, ) assert result["data"] == "Mock result" + assert hass.states.get(TEST_ENTITY_ID).context is context assert len(mock_ai_task_entity.mock_generate_data_tasks) == 1 task = mock_ai_task_entity.mock_generate_data_tasks[0] @@ -317,6 +320,7 @@ async def test_generate_image_service( mock_ai_task_entity: MockAITaskEntity, ) -> None: """Test the generate image service.""" + context = Context() preferences = hass.data[DATA_PREFERENCES] preferences.async_set_preferences(**set_preferences) @@ -335,9 +339,11 @@ async def test_generate_image_service( | msg_extra, blocking=True, return_response=True, + context=context, ) mock_upload_media.assert_called_once() + assert hass.states.get(TEST_ENTITY_ID).context is context assert "image_data" not in result assert ( result["media_source_id"] diff --git a/tests/components/ai_task/test_task.py b/tests/components/ai_task/test_task.py index 2a5add49a9bc..f49e8fd7e047 100644 --- a/tests/components/ai_task/test_task.py +++ b/tests/components/ai_task/test_task.py @@ -18,7 +18,7 @@ from homeassistant.components.camera import Image from homeassistant.components.conversation import async_get_chat_log from homeassistant.components.llm import AssistAPI from homeassistant.const import STATE_UNKNOWN -from homeassistant.core import HomeAssistant +from homeassistant.core import Context, HomeAssistant from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import chat_session from homeassistant.util import dt as dt_util @@ -78,14 +78,21 @@ async def test_generate_data_preferred_entity( assert state is not None assert state.state == STATE_UNKNOWN + context = Context() llm_api = AssistAPI(hass) result = await async_generate_data( hass, task_name="Test Task", instructions="Test prompt", llm_api=llm_api, + context=context, ) assert result.data == "Mock result" + + # The LLM API uses the context to check permissions when calling tools + chat_log = mock_ai_task_entity.mock_chat_logs[0] + assert chat_log.llm_api is not None + assert chat_log.llm_api.llm_context.context is context as_dict = result.as_dict() assert as_dict["conversation_id"] == result.conversation_id assert as_dict["data"] == "Mock result"