Allow passing an LLM API to AI Task generate data (#151081)

Author: balloob

homeassistant/components/ai_task/entity.py

@@ -60,6 +60,10 @@ async def _async_get_ai_task_chat_log(
         task: GenDataTask | GenImageTask,
     ) -> AsyncGenerator[ChatLog]:
         """Context manager used to manage the ChatLog used during an AI Task."""
+        user_llm_hass_api: llm.API | None = None
+        if isinstance(task, GenDataTask):
+            user_llm_hass_api = task.llm_api
+
         # pylint: disable-next=contextmanager-generator-missing-cleanup
         with (
             async_get_chat_log(
@@ -77,6 +81,7 @@ async def _async_get_ai_task_chat_log(
                     device_id=None,
                 ),
                 user_llm_prompt=DEFAULT_SYSTEM_PROMPT,
+                user_llm_hass_api=user_llm_hass_api,
             )
 
             chat_log.async_add_user_content(

homeassistant/components/ai_task/task.py

@@ -16,6 +16,7 @@
 from homeassistant.components.http.auth import async_sign_path
 from homeassistant.core import HomeAssistant, ServiceResponse, callback
 from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers import llm
 from homeassistant.helpers.chat_session import ChatSession, async_get_chat_session
 from homeassistant.helpers.event import async_call_later
 from homeassistant.helpers.network import get_url
@@ -116,6 +117,7 @@ async def async_generate_data(
     instructions: str,
     structure: vol.Schema | None = None,
     attachments: list[dict] | None = None,
+    llm_api: llm.API | None = None,
 ) -> GenDataTaskResult:
     """Run a data generation task in the AI Task integration."""
     if entity_id is None:
@@ -151,6 +153,7 @@ async def async_generate_data(
                 instructions=instructions,
                 structure=structure,
                 attachments=resolved_attachments or None,
+                llm_api=llm_api,
             ),
         )
 
@@ -272,6 +275,9 @@ class GenDataTask:
     attachments: list[conversation.Attachment] | None = None
     """List of attachments to go along the instructions."""
 
+    llm_api: llm.API | None = None
+    """API to provide to the LLM."""
+
     def __str__(self) -> str:
         """Return task as a string."""
         return f"<GenDataTask {self.name}: {id(self)}>"

homeassistant/components/conversation/chat_log.py

@@ -507,14 +507,18 @@ async def async_update_llm_data(
     async def async_provide_llm_data(
         self,
         llm_context: llm.LLMContext,
-        user_llm_hass_api: str | list[str] | None = None,
+        user_llm_hass_api: str | list[str] | llm.API | None = None,
         user_llm_prompt: str | None = None,
         user_extra_system_prompt: str | None = None,
     ) -> None:
         """Set the LLM system prompt."""
         llm_api: llm.APIInstance | None = None
 
-        if user_llm_hass_api:
+        if user_llm_hass_api is None:
+            pass
+        elif isinstance(user_llm_hass_api, llm.API):
+            llm_api = await user_llm_hass_api.async_get_api_instance(llm_context)
+        else:
             try:
                 llm_api = await llm.async_get_api(
                     self.hass,

tests/components/ai_task/test_task.py

@@ -20,7 +20,7 @@
 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, llm
 from homeassistant.util import dt as dt_util
 
 from .conftest import TEST_ENTITY_ID, MockAITaskEntity
@@ -78,10 +78,12 @@ async def test_generate_data_preferred_entity(
     assert state is not None
     assert state.state == STATE_UNKNOWN
 
+    llm_api = llm.AssistAPI(hass)
     result = await async_generate_data(
         hass,
         task_name="Test Task",
         instructions="Test prompt",
+        llm_api=llm_api,
     )
     assert result.data == "Mock result"
     as_dict = result.as_dict()
@@ -91,6 +93,12 @@ async def test_generate_data_preferred_entity(
     assert state is not None
     assert state.state != STATE_UNKNOWN
 
+    with (
+        chat_session.async_get_chat_session(hass, result.conversation_id) as session,
+        async_get_chat_log(hass, session) as chat_log,
+    ):
+        assert chat_log.llm_api.api is llm_api
+
     mock_ai_task_entity.supported_features = AITaskEntityFeature(0)
     with pytest.raises(
         HomeAssistantError,