Allow storing AI Task generate image preferred entity (#26959)

Author: balloob

src/data/ai_task.ts

@@ -1,8 +1,14 @@
 import type { HomeAssistant } from "../types";
 import type { Selector } from "./selector";
 
+export const enum AITaskEntityFeature {
+  GENERATE_DATA = 1,
+  SUPPORT_ATTACHMENTS = 2,
+  GENERATE_IMAGE = 4,
+}
 export interface AITaskPreferences {
   gen_data_entity_id: string | null;
+  gen_image_entity_id: string | null;
 }
 
 export interface GenDataTask {

src/panels/config/core/ai-task-pref.ts

@@ -1,5 +1,6 @@
 import { mdiHelpCircle, mdiStarFourPoints } from "@mdi/js";
 import { css, html, LitElement } from "lit";
+import type { HassEntity } from "home-assistant-js-websocket";
 import { customElement, property, state } from "lit/decorators";
 import { isComponentLoaded } from "../../../common/config/is_component_loaded";
 import type { HaProgressButton } from "../../../components/buttons/ha-progress-button";
@@ -8,13 +9,23 @@ import type { HaEntityPicker } from "../../../components/entity/ha-entity-picker
 import "../../../components/ha-card";
 import "../../../components/ha-settings-row";
 import {
+  AITaskEntityFeature,
   fetchAITaskPreferences,
   saveAITaskPreferences,
   type AITaskPreferences,
 } from "../../../data/ai_task";
 import type { HomeAssistant } from "../../../types";
 import { brandsUrl } from "../../../util/brands-url";
 import { documentationUrl } from "../../../util/documentation-url";
+import { computeDomain } from "../../../common/entity/compute_domain";
+import { supportsFeature } from "../../../common/entity/supports-feature";
+
+const filterGenData = (entity: HassEntity) =>
+  computeDomain(entity.entity_id) === "ai_task" &&
+  supportsFeature(entity, AITaskEntityFeature.GENERATE_DATA);
+const filterGenImage = (entity: HassEntity) =>
+  computeDomain(entity.entity_id) === "ai_task" &&
+  supportsFeature(entity, AITaskEntityFeature.GENERATE_IMAGE);
 
 @customElement("ai-task-pref")
 export class AITaskPref extends LitElement {
@@ -26,6 +37,8 @@ export class AITaskPref extends LitElement {
 
   private _gen_data_entity_id?: string | null;
 
+  private _gen_image_entity_id?: string | null;
+
   protected firstUpdated(changedProps) {
     super.firstUpdated(changedProps);
     if (!this.hass || !isComponentLoaded(this.hass, "ai_task")) {
@@ -90,7 +103,27 @@ export class AITaskPref extends LitElement {
               isComponentLoaded(this.hass, "ai_task")}
               .value=${this._gen_data_entity_id ||
               this._prefs?.gen_data_entity_id}
-              .includeDomains=${["ai_task"]}
+              .entityFilter=${filterGenData}
+              @value-changed=${this._handlePrefChange}
+            ></ha-entity-picker>
+          </ha-settings-row>
+          <ha-settings-row .narrow=${this.narrow}>
+            <span slot="heading">
+              ${this.hass!.localize("ui.panel.config.ai_task.gen_image_header")}
+            </span>
+            <span slot="description">
+              ${this.hass!.localize(
+                "ui.panel.config.ai_task.gen_image_description"
+              )}
+            </span>
+            <ha-entity-picker
+              data-name="gen_image_entity_id"
+              .hass=${this.hass}
+              .disabled=${this._prefs === undefined &&
+              isComponentLoaded(this.hass, "ai_task")}
+              .value=${this._gen_image_entity_id ||
+              this._prefs?.gen_image_entity_id}
+              .entityFilter=${filterGenImage}
               @value-changed=${this._handlePrefChange}
             ></ha-entity-picker>
           </ha-settings-row>
@@ -121,6 +154,7 @@ export class AITaskPref extends LitElement {
     const oldPrefs = this._prefs;
     const update: Partial<AITaskPreferences> = {
       gen_data_entity_id: this._gen_data_entity_id,
+      gen_image_entity_id: this._gen_image_entity_id,
     };
     this._prefs = { ...this._prefs!, ...update };
     try {

src/translations/en.json

@@ -2312,7 +2312,9 @@
           "header": "AI suggestions",
           "description": "Home Assistant can use generative AI to help you with tasks. Look for the button with the {button} icon throughout Home Assistant to get suggestions. Select an AI task entity to use this feature.",
           "gen_data_header": "Data generation tasks",
-          "gen_data_description": "Suggest automation names."
+          "gen_data_description": "Suggest automation names.",
+          "gen_image_header": "Image generation tasks",
+          "gen_image_description": "Generate images."
         },
         "category": {
           "caption": "Categories",