Migrate OpenAI to config subentries (home-assistant#147282)

* Migrate OpenAI to config subentries

* Add latest changes from Google subentries

* Update homeassistant/components/openai_conversation/__init__.py

Co-authored-by: Joost Lekkerkerker <[email protected]>

---------

Co-authored-by: Joost Lekkerkerker <[email protected]>

Author: balloob

homeassistant/components/openai_conversation/__init__.py

@@ -19,7 +19,7 @@
 )
 import voluptuous as vol
 
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigSubentry
 from homeassistant.const import CONF_API_KEY, Platform
 from homeassistant.core import (
     HomeAssistant,
@@ -32,7 +32,12 @@
     HomeAssistantError,
     ServiceValidationError,
 )
-from homeassistant.helpers import config_validation as cv, selector
+from homeassistant.helpers import (
+    config_validation as cv,
+    device_registry as dr,
+    entity_registry as er,
+    selector,
+)
 from homeassistant.helpers.httpx_client import get_async_client
 from homeassistant.helpers.typing import ConfigType
 
@@ -73,6 +78,7 @@ def encode_file(file_path: str) -> tuple[str, str]:
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     """Set up OpenAI Conversation."""
+    await async_migrate_integration(hass)
 
     async def render_image(call: ServiceCall) -> ServiceResponse:
         """Render an image with dall-e."""
@@ -118,7 +124,21 @@ async def send_prompt(call: ServiceCall) -> ServiceResponse:
                 translation_placeholders={"config_entry": entry_id},
             )
 
-        model: str = entry.options.get(CONF_CHAT_MODEL, RECOMMENDED_CHAT_MODEL)
+        # Get first conversation subentry for options
+        conversation_subentry = next(
+            (
+                sub
+                for sub in entry.subentries.values()
+                if sub.subentry_type == "conversation"
+            ),
+            None,
+        )
+        if not conversation_subentry:
+            raise ServiceValidationError("No conversation configuration found")
+
+        model: str = conversation_subentry.data.get(
+            CONF_CHAT_MODEL, RECOMMENDED_CHAT_MODEL
+        )
         client: openai.AsyncClient = entry.runtime_data
 
         content: ResponseInputMessageContentListParam = [
@@ -169,11 +189,11 @@ def append_files_to_content() -> None:
             model_args = {
                 "model": model,
                 "input": messages,
-                "max_output_tokens": entry.options.get(
+                "max_output_tokens": conversation_subentry.data.get(
                     CONF_MAX_TOKENS, RECOMMENDED_MAX_TOKENS
                 ),
-                "top_p": entry.options.get(CONF_TOP_P, RECOMMENDED_TOP_P),
-                "temperature": entry.options.get(
+                "top_p": conversation_subentry.data.get(CONF_TOP_P, RECOMMENDED_TOP_P),
+                "temperature": conversation_subentry.data.get(
                     CONF_TEMPERATURE, RECOMMENDED_TEMPERATURE
                 ),
                 "user": call.context.user_id,
@@ -182,7 +202,7 @@ def append_files_to_content() -> None:
 
             if model.startswith("o"):
                 model_args["reasoning"] = {
-                    "effort": entry.options.get(
+                    "effort": conversation_subentry.data.get(
                         CONF_REASONING_EFFORT, RECOMMENDED_REASONING_EFFORT
                     )
                 }
@@ -269,3 +289,68 @@ async def async_setup_entry(hass: HomeAssistant, entry: OpenAIConfigEntry) -> bo
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload OpenAI."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
+
+
+async def async_migrate_integration(hass: HomeAssistant) -> None:
+    """Migrate integration entry structure."""
+
+    entries = hass.config_entries.async_entries(DOMAIN)
+    if not any(entry.version == 1 for entry in entries):
+        return
+
+    api_keys_entries: dict[str, ConfigEntry] = {}
+    entity_registry = er.async_get(hass)
+    device_registry = dr.async_get(hass)
+
+    for entry in entries:
+        use_existing = False
+        subentry = ConfigSubentry(
+            data=entry.options,
+            subentry_type="conversation",
+            title=entry.title,
+            unique_id=None,
+        )
+        if entry.data[CONF_API_KEY] not in api_keys_entries:
+            use_existing = True
+            api_keys_entries[entry.data[CONF_API_KEY]] = entry
+
+        parent_entry = api_keys_entries[entry.data[CONF_API_KEY]]
+
+        hass.config_entries.async_add_subentry(parent_entry, subentry)
+        conversation_entity = entity_registry.async_get_entity_id(
+            "conversation",
+            DOMAIN,
+            entry.entry_id,
+        )
+        if conversation_entity is not None:
+            entity_registry.async_update_entity(
+                conversation_entity,
+                config_entry_id=parent_entry.entry_id,
+                config_subentry_id=subentry.subentry_id,
+                new_unique_id=subentry.subentry_id,
+            )
+
+        device = device_registry.async_get_device(
+            identifiers={(DOMAIN, entry.entry_id)}
+        )
+        if device is not None:
+            device_registry.async_update_device(
+                device.id,
+                new_identifiers={(DOMAIN, subentry.subentry_id)},
+                add_config_subentry_id=subentry.subentry_id,
+                add_config_entry_id=parent_entry.entry_id,
+            )
+            if parent_entry.entry_id != entry.entry_id:
+                device_registry.async_update_device(
+                    device.id,
+                    remove_config_entry_id=entry.entry_id,
+                )
+
+        if not use_existing:
+            await hass.config_entries.async_remove(entry.entry_id)
+        else:
+            hass.config_entries.async_update_entry(
+                entry,
+                options={},
+                version=2,
+            )

homeassistant/components/openai_conversation/config_flow.py

@@ -13,17 +13,20 @@
 from homeassistant.components.zone import ENTITY_ID_HOME
 from homeassistant.config_entries import (
     ConfigEntry,
+    ConfigEntryState,
     ConfigFlow,
     ConfigFlowResult,
-    OptionsFlow,
+    ConfigSubentryFlow,
+    SubentryFlowResult,
 )
 from homeassistant.const import (
     ATTR_LATITUDE,
     ATTR_LONGITUDE,
     CONF_API_KEY,
     CONF_LLM_HASS_API,
+    CONF_NAME,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import llm
 from homeassistant.helpers.httpx_client import get_async_client
 from homeassistant.helpers.selector import (
@@ -52,6 +55,7 @@
     CONF_WEB_SEARCH_REGION,
     CONF_WEB_SEARCH_TIMEZONE,
     CONF_WEB_SEARCH_USER_LOCATION,
+    DEFAULT_CONVERSATION_NAME,
     DOMAIN,
     RECOMMENDED_CHAT_MODEL,
     RECOMMENDED_MAX_TOKENS,
@@ -94,7 +98,7 @@ async def validate_input(hass: HomeAssistant, data: dict[str, Any]) -> None:
 class OpenAIConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for OpenAI Conversation."""
 
-    VERSION = 1
+    VERSION = 2
 
     async def async_step_user(
         self, user_input: dict[str, Any] | None = None
@@ -107,6 +111,7 @@ async def async_step_user(
 
         errors: dict[str, str] = {}
 
+        self._async_abort_entries_match(user_input)
         try:
             await validate_input(self.hass, user_input)
         except openai.APIConnectionError:
@@ -120,32 +125,61 @@ async def async_step_user(
             return self.async_create_entry(
                 title="ChatGPT",
                 data=user_input,
-                options=RECOMMENDED_OPTIONS,
+                subentries=[
+                    {
+                        "subentry_type": "conversation",
+                        "data": RECOMMENDED_OPTIONS,
+                        "title": DEFAULT_CONVERSATION_NAME,
+                        "unique_id": None,
+                    }
+                ],
             )
 
         return self.async_show_form(
             step_id="user", data_schema=STEP_USER_DATA_SCHEMA, errors=errors
         )
 
-    @staticmethod
-    def async_get_options_flow(
-        config_entry: ConfigEntry,
-    ) -> OptionsFlow:
-        """Create the options flow."""
-        return OpenAIOptionsFlow(config_entry)
+    @classmethod
+    @callback
+    def async_get_supported_subentry_types(
+        cls, config_entry: ConfigEntry
+    ) -> dict[str, type[ConfigSubentryFlow]]:
+        """Return subentries supported by this integration."""
+        return {"conversation": ConversationSubentryFlowHandler}
+
 
+class ConversationSubentryFlowHandler(ConfigSubentryFlow):
+    """Flow for managing conversation subentries."""
 
-class OpenAIOptionsFlow(OptionsFlow):
-    """OpenAI config flow options handler."""
+    last_rendered_recommended = False
+    options: dict[str, Any]
+
+    @property
+    def _is_new(self) -> bool:
+        """Return if this is a new subentry."""
+        return self.source == "user"
+
+    async def async_step_user(
+        self, user_input: dict[str, Any] | None = None
+    ) -> SubentryFlowResult:
+        """Add a subentry."""
+        self.options = RECOMMENDED_OPTIONS.copy()
+        return await self.async_step_init()
 
-    def __init__(self, config_entry: ConfigEntry) -> None:
-        """Initialize options flow."""
-        self.options = config_entry.options.copy()
+    async def async_step_reconfigure(
+        self, user_input: dict[str, Any] | None = None
+    ) -> SubentryFlowResult:
+        """Handle reconfiguration of a subentry."""
+        self.options = self._get_reconfigure_subentry().data.copy()
+        return await self.async_step_init()
 
     async def async_step_init(
         self, user_input: dict[str, Any] | None = None
-    ) -> ConfigFlowResult:
+    ) -> SubentryFlowResult:
         """Manage initial options."""
+        # abort if entry is not loaded
+        if self._get_entry().state != ConfigEntryState.LOADED:
+            return self.async_abort(reason="entry_not_loaded")
         options = self.options
 
         hass_apis: list[SelectOptionDict] = [
@@ -160,25 +194,47 @@ async def async_step_init(
         ):
             options[CONF_LLM_HASS_API] = [suggested_llm_apis]
 
-        step_schema: VolDictType = {
-            vol.Optional(
-                CONF_PROMPT,
-                description={"suggested_value": llm.DEFAULT_INSTRUCTIONS_PROMPT},
-            ): TemplateSelector(),
-            vol.Optional(CONF_LLM_HASS_API): SelectSelector(
-                SelectSelectorConfig(options=hass_apis, multiple=True)
-            ),
-            vol.Required(
-                CONF_RECOMMENDED, default=options.get(CONF_RECOMMENDED, False)
-            ): bool,
-        }
+        step_schema: VolDictType = {}
+
+        if self._is_new:
+            step_schema[vol.Required(CONF_NAME, default=DEFAULT_CONVERSATION_NAME)] = (
+                str
+            )
+
+        step_schema.update(
+            {
+                vol.Optional(
+                    CONF_PROMPT,
+                    description={
+                        "suggested_value": options.get(
+                            CONF_PROMPT, llm.DEFAULT_INSTRUCTIONS_PROMPT
+                        )
+                    },
+                ): TemplateSelector(),
+                vol.Optional(CONF_LLM_HASS_API): SelectSelector(
+                    SelectSelectorConfig(options=hass_apis, multiple=True)
+                ),
+                vol.Required(
+                    CONF_RECOMMENDED, default=options.get(CONF_RECOMMENDED, False)
+                ): bool,
+            }
+        )
 
         if user_input is not None:
             if not user_input.get(CONF_LLM_HASS_API):
                 user_input.pop(CONF_LLM_HASS_API, None)
 
             if user_input[CONF_RECOMMENDED]:
-                return self.async_create_entry(title="", data=user_input)
+                if self._is_new:
+                    return self.async_create_entry(
+                        title=user_input.pop(CONF_NAME),
+                        data=user_input,
+                    )
+                return self.async_update_and_abort(
+                    self._get_entry(),
+                    self._get_reconfigure_subentry(),
+                    data=user_input,
+                )
 
             options.update(user_input)
             if CONF_LLM_HASS_API in options and CONF_LLM_HASS_API not in user_input:
@@ -194,7 +250,7 @@ async def async_step_init(
 
     async def async_step_advanced(
         self, user_input: dict[str, Any] | None = None
-    ) -> ConfigFlowResult:
+    ) -> SubentryFlowResult:
         """Manage advanced options."""
         options = self.options
         errors: dict[str, str] = {}
@@ -236,7 +292,7 @@ async def async_step_advanced(
 
     async def async_step_model(
         self, user_input: dict[str, Any] | None = None
-    ) -> ConfigFlowResult:
+    ) -> SubentryFlowResult:
         """Manage model-specific options."""
         options = self.options
         errors: dict[str, str] = {}
@@ -303,7 +359,16 @@ async def async_step_model(
             }
 
         if not step_schema:
-            return self.async_create_entry(title="", data=options)
+            if self._is_new:
+                return self.async_create_entry(
+                    title=options.pop(CONF_NAME, DEFAULT_CONVERSATION_NAME),
+                    data=options,
+                )
+            return self.async_update_and_abort(
+                self._get_entry(),
+                self._get_reconfigure_subentry(),
+                data=options,
+            )
 
         if user_input is not None:
             if user_input.get(CONF_WEB_SEARCH):
@@ -316,7 +381,16 @@ async def async_step_model(
                     options.pop(CONF_WEB_SEARCH_TIMEZONE, None)
 
             options.update(user_input)
-            return self.async_create_entry(title="", data=options)
+            if self._is_new:
+                return self.async_create_entry(
+                    title=options.pop(CONF_NAME, DEFAULT_CONVERSATION_NAME),
+                    data=options,
+                )
+            return self.async_update_and_abort(
+                self._get_entry(),
+                self._get_reconfigure_subentry(),
+                data=options,
+            )
 
         return self.async_show_form(
             step_id="model",
@@ -332,7 +406,7 @@ async def _get_location_data(self) -> dict[str, str]:
         zone_home = self.hass.states.get(ENTITY_ID_HOME)
         if zone_home is not None:
             client = openai.AsyncOpenAI(
-                api_key=self.config_entry.data[CONF_API_KEY],
+                api_key=self._get_entry().data[CONF_API_KEY],
                 http_client=get_async_client(self.hass),
             )
             location_schema = vol.Schema(