AI task generate_text -> generate_data (#147370)

homeassistant/components/ai_task/task.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
+from typing import Any
 
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
@@ -10,16 +11,16 @@ from homeassistant.exceptions import HomeAssistantError
 from .const import DATA_COMPONENT, DATA_PREFERENCES, AITaskEntityFeature
 
 
-async def async_generate_text(
+async def async_generate_data(
     hass: HomeAssistant,
     *,
     task_name: str,
     entity_id: str | None = None,
     instructions: str,
-) -> GenTextTaskResult:
+) -> GenDataTaskResult:
     """Run a task in the AI Task integration."""
     if entity_id is None:
-        entity_id = hass.data[DATA_PREFERENCES].gen_text_entity_id
+        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")
@@ -28,13 +29,13 @@ async def async_generate_text(
     if entity is None:
         raise HomeAssistantError(f"AI Task entity {entity_id} not found")
 
-    if AITaskEntityFeature.GENERATE_TEXT not in entity.supported_features:
+    if AITaskEntityFeature.GENERATE_DATA not in entity.supported_features:
         raise HomeAssistantError(
-            f"AI Task entity {entity_id} does not support generating text"
+            f"AI Task entity {entity_id} does not support generating data"
         )
 
-    return await entity.internal_async_generate_text(
-        GenTextTask(
+    return await entity.internal_async_generate_data(
+        GenDataTask(
             name=task_name,
             instructions=instructions,
         )
@@ -42,8 +43,8 @@ async def async_generate_text(
 
 
 @dataclass(slots=True)
-class GenTextTask:
-    """Gen text task to be processed."""
+class GenDataTask:
+    """Gen data task to be processed."""
 
     name: str
     """Name of the task."""
@@ -53,22 +54,22 @@ class GenTextTask:
 
     def __str__(self) -> str:
         """Return task as a string."""
-        return f"<GenTextTask {self.name}: {id(self)}>"
+        return f"<GenDataTask {self.name}: {id(self)}>"
 
 
 @dataclass(slots=True)
-class GenTextTaskResult:
-    """Result of gen text task."""
+class GenDataTaskResult:
+    """Result of gen data task."""
 
     conversation_id: str
     """Unique identifier for the conversation."""
 
-    text: str
-    """Generated text."""
+    data: Any
+    """Data generated by the task."""
 
-    def as_dict(self) -> dict[str, str]:
+    def as_dict(self) -> dict[str, Any]:
         """Return result as a dict."""
         return {
             "conversation_id": self.conversation_id,
-            "text": self.text,
+            "data": self.data,
         }