Merge

Co-authored-by: Wendelin <12148533+wendevlin@users.noreply.github.com>

src/panels/config/automation/automation-save-dialog/dialog-automation-save.ts

@@ -1,5 +1,4 @@
 import { mdiClose, mdiPlus } from "@mdi/js";
-import { dump } from "js-yaml";
 import type { CSSResultGroup } from "lit";
 import { css, html, LitElement, nothing } from "lit";
 import { customElement, property, state } from "lit/decorators";
@@ -19,13 +18,8 @@ import "../../../../components/ha-textarea";
 import "../../../../components/ha-textfield";
 import "../../category/ha-category-picker";
 
-import { computeStateDomain } from "../../../../common/entity/compute_state_domain";
 import { supportsMarkdownHelper } from "../../../../common/translations/markdown_support";
-import { subscribeOne } from "../../../../common/util/subscribe-one";
 import type { GenDataTaskResult } from "../../../../data/ai_task";
-import { fetchCategoryRegistry } from "../../../../data/category_registry";
-import { subscribeEntityRegistry } from "../../../../data/entity/entity_registry";
-import { subscribeLabelRegistry } from "../../../../data/label/label_registry";
 import type { HassDialog } from "../../../../dialogs/make-dialog-manager";
 import { haStyle, haStyleDialog } from "../../../../resources/styles";
 import type { HomeAssistant } from "../../../../types";
@@ -33,6 +27,11 @@ import type {
   EntityRegistryUpdate,
   SaveDialogParams,
 } from "./show-dialog-automation-save";
+import {
+  generateMetadataSuggestionTask,
+  processMetadataSuggestion,
+  type MetadataSuggestionResult,
+} from "../../common/suggest-metadata-ai";
 
 @customElement("ha-dialog-automation-save")
 class DialogAutomationSave extends LitElement implements HassDialog {
@@ -335,180 +334,49 @@ class DialogAutomationSave extends LitElement implements HassDialog {
     this.closeDialog();
   }
 
-  private _getSuggestData() {
-    return Promise.all([
-      subscribeOne(this.hass.connection, subscribeLabelRegistry).then((labs) =>
-        Object.fromEntries(labs.map((lab) => [lab.label_id, lab.name]))
-      ),
-      subscribeOne(this.hass.connection, subscribeEntityRegistry).then((ents) =>
-        Object.fromEntries(ents.map((ent) => [ent.entity_id, ent]))
-      ),
-      fetchCategoryRegistry(this.hass.connection, "automation").then((cats) =>
-        Object.fromEntries(cats.map((cat) => [cat.category_id, cat.name]))
-      ),
-    ]);
-  }
-
-  private _generateTask = async (): Promise<SuggestWithAIGenerateTask> => {
-    const [labels, entities, categories] = await this._getSuggestData();
-    const inspirations: string[] = [];
-
-    const domain = this._params.domain;
-
-    for (const entity of Object.values(this.hass.states)) {
-      const entityEntry = entities[entity.entity_id];
-      if (
-        computeStateDomain(entity) !== domain ||
-        entity.attributes.restored ||
-        !entity.attributes.friendly_name ||
-        !entityEntry
-      ) {
-        continue;
-      }
-
-      let inspiration = `- ${entity.attributes.friendly_name}`;
-
-      const category = categories[entityEntry.categories.automation];
-      if (category) {
-        inspiration += ` (category: ${category})`;
-      }
-
-      if (entityEntry.labels.length) {
-        inspiration += ` (labels: ${entityEntry.labels
-          .map((label) => labels[label])
-          .join(", ")})`;
-      }
-
-      inspirations.push(inspiration);
-    }
-
-    const term = this._params.domain === "script" ? "script" : "automation";
-
-    return {
-      type: "data",
-      task: {
-        task_name: `frontend__${term}__save`,
-        instructions: `Suggest in language "${this.hass.language}" a name, description, category and labels for the following Home Assistant ${term}.
-
-The name should be relevant to the ${term}'s purpose.
-${
-  inspirations.length
-    ? `The name should be in same style and sentence capitalization as existing ${term}s.
-Suggest a category and labels if relevant to the ${term}'s purpose.
-Only suggest category and labels that are already used by existing ${term}s.`
-    : `The name should be short, descriptive, sentence case, and written in the language ${this.hass.language}.`
-}
-If the ${term} contains 5+ steps, include a short description.
-
-For inspiration, here are existing ${term}s:
-${inspirations.join("\n")}
-
-The ${term} configuration is as follows:
-
-${dump(this._params.config)}
-`,
-        structure: {
-          name: {
-            description: "The name of the automation",
-            required: true,
-            selector: {
-              text: {},
-            },
-          },
-          description: {
-            description: "A short description of the automation",
-            required: false,
-            selector: {
-              text: {},
-            },
-          },
-          labels: {
-            description: "Labels for the automation",
-            required: false,
-            selector: {
-              text: {
-                multiple: true,
-              },
-            },
-          },
-          category: {
-            description: "The category of the automation",
-            required: false,
-            selector: {
-              select: {
-                options: Object.entries(categories).map(([id, name]) => ({
-                  value: id,
-                  label: name,
-                })),
-              },
-            },
-          },
-        },
-      },
-    };
-  };
+  private _generateTask = async (): Promise<SuggestWithAIGenerateTask> =>
+    generateMetadataSuggestionTask(this.hass, {
+      domain: this._params.domain,
+      config: this._params.config,
+      includeDescription: true,
+    });
 
   private async _handleSuggestion(
-    event: CustomEvent<
-      GenDataTaskResult<{
-        name: string;
-        description?: string;
-        category?: string;
-        labels?: string[];
-      }>
-    >
+    event: CustomEvent<GenDataTaskResult<MetadataSuggestionResult>>
   ) {
     const result = event.detail;
-    const [labels, _entities, categories] = await this._getSuggestData();
+    const processed = await processMetadataSuggestion(
+      this.hass,
+      this._params.domain,
+      result
+    );
 
-    this._newName = result.data.name;
-    if (result.data.description) {
-      this._newDescription = result.data.description;
+    this._newName = processed.name;
+
+    if (processed.description) {
+      this._newDescription = processed.description;
       if (!this._visibleOptionals.includes("description")) {
         this._visibleOptionals = [...this._visibleOptionals, "description"];
       }
     }
-    if (result.data.category) {
-      // We get back category name, convert it to ID
-      const categoryId = Object.entries(categories).find(
-        ([, name]) => name === result.data.category
-      )?.[0];
-      if (categoryId) {
-        this._entryUpdates = {
-          ...this._entryUpdates,
-          category: categoryId,
-        };
-        if (!this._visibleOptionals.includes("category")) {
-          this._visibleOptionals = [...this._visibleOptionals, "category"];
-        }
+
+    if (processed.categoryId) {
+      this._entryUpdates = {
+        ...this._entryUpdates,
+        category: processed.categoryId,
+      };
+      if (!this._visibleOptionals.includes("category")) {
+        this._visibleOptionals = [...this._visibleOptionals, "category"];
       }
     }
-    if (result.data.labels?.length) {
-      // We get back label names, convert them to IDs
-      const newLabels: Record<string, undefined | string> = Object.fromEntries(
-        result.data.labels.map((name) => [name, undefined])
-      );
-      let toFind = result.data.labels.length;
-      for (const [labelId, labelName] of Object.entries(labels)) {
-        if (labelName in newLabels && newLabels[labelName] === undefined) {
-          newLabels[labelName] = labelId;
-          toFind--;
-          if (toFind === 0) {
-            break;
-          }
-        }
-      }
-      const foundLabels = Object.values(newLabels).filter(
-        (labelId) => labelId !== undefined
-      );
-      if (foundLabels.length) {
-        this._entryUpdates = {
-          ...this._entryUpdates,
-          labels: foundLabels,
-        };
-        if (!this._visibleOptionals.includes("labels")) {
-          this._visibleOptionals = [...this._visibleOptionals, "labels"];
-        }
+
+    if (processed.labelIds?.length) {
+      this._entryUpdates = {
+        ...this._entryUpdates,
+        labels: processed.labelIds,
+      };
+      if (!this._visibleOptionals.includes("labels")) {
+        this._visibleOptionals = [...this._visibleOptionals, "labels"];
       }
     }
   }

src/panels/config/common/suggest-metadata-ai.ts

@@ -0,0 +1,263 @@
+import { dump } from "js-yaml";
+import { computeStateDomain } from "../../../common/entity/compute_state_domain";
+import { subscribeOne } from "../../../common/util/subscribe-one";
+import type { AITaskStructure, GenDataTaskResult } from "../../../data/ai_task";
+import { fetchCategoryRegistry } from "../../../data/category_registry";
+import {
+  subscribeEntityRegistry,
+  type EntityRegistryEntry,
+} from "../../../data/entity/entity_registry";
+import { subscribeLabelRegistry } from "../../../data/label/label_registry";
+import type { HomeAssistant } from "../../../types";
+import type { SuggestWithAIGenerateTask } from "../../../components/ha-suggest-with-ai-button";
+
+export interface MetadataSuggestionResult {
+  name: string;
+  description?: string;
+  category?: string;
+  labels?: string[];
+}
+
+export interface ProcessedMetadataSuggestionResult {
+  name: string;
+  description?: string;
+  categoryId?: string;
+  labelIds?: string[];
+}
+
+export interface MetadataSuggestionConfig {
+  /** The domain to suggest metadata for (automation, script) */
+  domain: "automation" | "script";
+  /** The configuration to suggest metadata for */
+  config: any;
+  /** Whether to include description field in the suggestion */
+  includeDescription?: boolean;
+  /** Whether to include icon field in the suggestion (scripts only) */
+  includeIcon?: boolean;
+}
+
+type Categories = Record<string, string>;
+type Entities = Record<string, EntityRegistryEntry>;
+type Labels = Record<string, string>;
+
+const fetchCategories = (
+  connection: HomeAssistant["connection"],
+  domain: MetadataSuggestionConfig["domain"]
+): Promise<Categories> =>
+  fetchCategoryRegistry(connection, domain).then((cats) =>
+    Object.fromEntries(cats.map((cat) => [cat.category_id, cat.name]))
+  );
+
+const fetchEntities = (
+  connection: HomeAssistant["connection"]
+): Promise<Entities> =>
+  subscribeOne(connection, subscribeEntityRegistry).then((ents) =>
+    Object.fromEntries(ents.map((ent) => [ent.entity_id, ent]))
+  );
+
+const fetchLabels = (
+  connection: HomeAssistant["connection"]
+): Promise<Labels> =>
+  subscribeOne(connection, subscribeLabelRegistry).then((labs) =>
+    Object.fromEntries(labs.map((lab) => [lab.label_id, lab.name]))
+  );
+
+function buildMetadataInspirations(
+  states: HomeAssistant["states"],
+  entities: Record<string, EntityRegistryEntry>,
+  categories: Categories,
+  labels: Labels,
+  domain: MetadataSuggestionConfig["domain"]
+): string[] {
+  const inspirations: string[] = [];
+
+  for (const entity of Object.values(states)) {
+    const entityEntry = entities[entity.entity_id];
+    if (
+      !entityEntry ||
+      computeStateDomain(entity) !== domain ||
+      entity.attributes.restored ||
+      !entity.attributes.friendly_name
+    ) {
+      continue;
+    }
+
+    let inspiration = `- ${entity.attributes.friendly_name}`;
+
+    // Get the category for this domain
+    const category = categories[entityEntry.categories[domain]];
+    if (category) {
+      inspiration += ` (category: ${category})`;
+    }
+
+    if (entityEntry.labels.length) {
+      inspiration += ` (labels: ${entityEntry.labels
+        .map((label) => labels[label])
+        .join(", ")})`;
+    }
+
+    inspirations.push(inspiration);
+  }
+
+  return inspirations;
+}
+
+export async function generateMetadataSuggestionTask(
+  connection: HomeAssistant["connection"],
+  states: HomeAssistant["states"],
+  language: HomeAssistant["language"],
+  suggestionConfig: MetadataSuggestionConfig
+): Promise<SuggestWithAIGenerateTask> {
+  const { domain, config, includeDescription } = suggestionConfig;
+
+  let categories: Categories = {};
+  let entities: Entities = {};
+  let labels: Labels = {};
+  try {
+    [categories, entities, labels] = await Promise.all([
+      fetchCategories(connection, domain),
+      fetchEntities(connection),
+      fetchLabels(connection),
+    ]);
+  } catch (error) {
+    // eslint-disable-next-line no-console
+    console.error("Error getting suggest metadata:", error);
+  }
+
+  const inspirations = buildMetadataInspirations(
+    states,
+    entities,
+    categories,
+    labels,
+    domain
+  );
+
+  const structure: AITaskStructure = {
+    name: {
+      description: `The name of the ${domain}`,
+      required: true,
+      selector: {
+        text: {},
+      },
+    },
+    ...(includeDescription && {
+      description: {
+        description: `A short description of the ${domain}`,
+        required: false,
+        selector: {
+          text: {},
+        },
+      },
+    }),
+    labels: {
+      description: `Labels for the ${domain}`,
+      required: false,
+      selector: {
+        text: {
+          multiple: true,
+        },
+      },
+    },
+    category: {
+      description: `The category of the ${domain}`,
+      required: false,
+      selector: {
+        select: {
+          options: Object.entries(categories).map(([id, name]) => ({
+            value: id,
+            label: name,
+          })),
+        },
+      },
+    },
+  };
+
+  return {
+    type: "data",
+    task: {
+      task_name: `frontend__${domain}__save`,
+      instructions: `Suggest in language "${language}" a name${includeDescription ? ", description" : ""}, category and labels for the following Home Assistant ${domain}.
+
+The name should be relevant to the ${domain}'s purpose.
+${
+  inspirations.length
+    ? `The name should be in same style and sentence capitalization as existing ${domain}s.
+Suggest a category and labels if relevant to the ${domain}'s purpose.
+Only suggest category and labels that are already used by existing ${domain}s.`
+    : `The name should be short, descriptive, sentence case, and written in the language ${language}.`
+}${
+        includeDescription
+          ? `
+If the ${domain} contains 5+ steps, include a short description.`
+          : ""
+      }
+
+For inspiration, here are existing ${domain}s:
+${inspirations.join("\n")}
+
+The ${domain} configuration is as follows:
+
+${dump(config)}
+`,
+      structure,
+    },
+  };
+}
+
+export async function processMetadataSuggestion(
+  connection: HomeAssistant["connection"],
+  domain: MetadataSuggestionConfig["domain"],
+  result: GenDataTaskResult<MetadataSuggestionResult>
+): Promise<ProcessedMetadataSuggestionResult> {
+  let categories: Categories = {};
+  let labels: Labels = {};
+  try {
+    [categories, labels] = await Promise.all([
+      fetchCategories(connection, domain),
+      fetchLabels(connection),
+    ]);
+  } catch (error) {
+    // eslint-disable-next-line no-console
+    console.error("Error getting suggest metadata:", error);
+  }
+
+  const processed: ProcessedMetadataSuggestionResult = {
+    name: result.data.name,
+    description: result.data.description ?? undefined,
+  };
+
+  // Convert category name to ID
+  if (result.data.category) {
+    const categoryId = Object.entries(categories).find(
+      ([, name]) => name === result.data.category
+    )?.[0];
+    if (categoryId) {
+      processed.categoryId = categoryId;
+    }
+  }
+
+  // Convert label names to IDs
+  if (result.data.labels?.length) {
+    const newLabels: Record<string, undefined | string> = Object.fromEntries(
+      result.data.labels.map((name) => [name, undefined])
+    );
+    let toFind = result.data.labels.length;
+    for (const [labelId, labelName] of Object.entries(labels)) {
+      if (labelName in newLabels && newLabels[labelName] === undefined) {
+        newLabels[labelName] = labelId;
+        toFind--;
+        if (toFind === 0) {
+          break;
+        }
+      }
+    }
+    const foundLabels = Object.values(newLabels).filter(
+      (labelId) => labelId !== undefined
+    );
+    if (foundLabels.length) {
+      processed.labelIds = foundLabels;
+    }
+  }
+
+  return processed;
+}