diff --git a/homeassistant/components/google_generative_ai_conversation/__init__.py b/homeassistant/components/google_generative_ai_conversation/__init__.py index 7e9ca550275..98951cd9b2a 100644 --- a/homeassistant/components/google_generative_ai_conversation/__init__.py +++ b/homeassistant/components/google_generative_ai_conversation/__init__.py @@ -12,7 +12,7 @@ from google.genai.types import File, FileState from requests.exceptions import Timeout 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, @@ -26,13 +26,14 @@ from homeassistant.exceptions import ( ConfigEntryNotReady, HomeAssistantError, ) -from homeassistant.helpers import config_validation as cv +from homeassistant.helpers import config_validation as cv, entity_registry as er from homeassistant.helpers.issue_registry import IssueSeverity, async_create_issue from homeassistant.helpers.typing import ConfigType from .const import ( CONF_CHAT_MODEL, CONF_PROMPT, + DEFAULT_CONVERSATION_NAME, DOMAIN, FILE_POLLING_INTERVAL_SECONDS, LOGGER, @@ -209,3 +210,42 @@ async def async_unload_entry( return False return True + + +async def async_migrate_entry( + hass: HomeAssistant, entry: GoogleGenerativeAIConfigEntry +) -> bool: + """Migrate old entry.""" + if entry.version == 1: + # Migrate from version 1 to version 2 + # Move conversation-specific options to a subentry + subentry = ConfigSubentry( + data=entry.options, + subentry_type="conversation", + title=DEFAULT_CONVERSATION_NAME, + unique_id=None, + ) + hass.config_entries.async_add_subentry( + entry, + subentry, + ) + + # Migrate conversation entity to be linked to subentry + ent_reg = er.async_get(hass) + for entity_entry in er.async_entries_for_config_entry(ent_reg, entry.entry_id): + if entity_entry.domain == Platform.CONVERSATION: + ent_reg.async_update_entity( + entity_entry.entity_id, + config_subentry_id=subentry.subentry_id, + new_unique_id=subentry.subentry_id, + ) + break + + # Remove options from the main entry + hass.config_entries.async_update_entry( + entry, + options={}, + version=2, + ) + + return True diff --git a/homeassistant/components/google_generative_ai_conversation/config_flow.py b/homeassistant/components/google_generative_ai_conversation/config_flow.py index ae0f09b1037..55f0bd089a1 100644 --- a/homeassistant/components/google_generative_ai_conversation/config_flow.py +++ b/homeassistant/components/google_generative_ai_conversation/config_flow.py @@ -4,7 +4,6 @@ from __future__ import annotations from collections.abc import Mapping import logging -from types import MappingProxyType from typing import Any from google import genai @@ -17,10 +16,11 @@ from homeassistant.config_entries import ( ConfigEntry, ConfigFlow, ConfigFlowResult, - OptionsFlow, + ConfigSubentryFlow, + SubentryFlowResult, ) from homeassistant.const import 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.selector import ( NumberSelector, @@ -45,6 +45,7 @@ from .const import ( CONF_TOP_K, CONF_TOP_P, CONF_USE_GOOGLE_SEARCH_TOOL, + DEFAULT_CONVERSATION_NAME, DOMAIN, RECOMMENDED_CHAT_MODEL, RECOMMENDED_HARM_BLOCK_THRESHOLD, @@ -66,7 +67,7 @@ STEP_API_DATA_SCHEMA = vol.Schema( RECOMMENDED_OPTIONS = { CONF_RECOMMENDED: True, - CONF_LLM_HASS_API: llm.LLM_API_ASSIST, + CONF_LLM_HASS_API: [llm.LLM_API_ASSIST], CONF_PROMPT: llm.DEFAULT_INSTRUCTIONS_PROMPT, } @@ -90,7 +91,7 @@ async def validate_input(data: dict[str, Any]) -> None: class GoogleGenerativeAIConfigFlow(ConfigFlow, domain=DOMAIN): """Handle a config flow for Google Generative AI Conversation.""" - VERSION = 1 + VERSION = 2 async def async_step_api( self, user_input: dict[str, Any] | None = None @@ -117,7 +118,14 @@ class GoogleGenerativeAIConfigFlow(ConfigFlow, domain=DOMAIN): return self.async_create_entry( title="Google Generative AI", 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="api", @@ -156,58 +164,90 @@ class GoogleGenerativeAIConfigFlow(ConfigFlow, domain=DOMAIN): }, ) - @staticmethod - def async_get_options_flow( - config_entry: ConfigEntry, - ) -> OptionsFlow: - """Create the options flow.""" - return GoogleGenerativeAIOptionsFlow(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 GoogleGenerativeAIOptionsFlow(OptionsFlow): - """Google Generative AI config flow options handler.""" +class ConversationSubentryFlowHandler(ConfigSubentryFlow): + """Flow for managing conversation subentries.""" - def __init__(self, config_entry: ConfigEntry) -> None: - """Initialize options flow.""" - self.last_rendered_recommended = config_entry.options.get( - CONF_RECOMMENDED, False - ) - self._genai_client = config_entry.runtime_data + last_rendered_recommended = False + is_new: bool + start_data: dict[str, Any] - async def async_step_init( + @property + def _genai_client(self) -> genai.Client: + """Return the Google Generative AI client.""" + return self._get_entry().runtime_data + + async def async_step_user( self, user_input: dict[str, Any] | None = None - ) -> ConfigFlowResult: - """Manage the options.""" - options: dict[str, Any] | MappingProxyType[str, Any] = self.config_entry.options + ) -> SubentryFlowResult: + """Add a subentry.""" + self.is_new = True + self.start_data = RECOMMENDED_OPTIONS.copy() + return await self.async_step_set_options() + + async def async_step_reconfigure( + self, user_input: dict[str, Any] | None = None + ) -> SubentryFlowResult: + """Handle reconfiguration of a subentry.""" + self.is_new = False + self.start_data = self._get_reconfigure_subentry().data.copy() + return await self.async_step_set_options() + + async def async_step_set_options( + self, user_input: dict[str, Any] | None = None + ) -> SubentryFlowResult: + """Set conversation options.""" + options = self.start_data errors: dict[str, str] = {} if user_input is not None: if user_input[CONF_RECOMMENDED] == self.last_rendered_recommended: if not user_input.get(CONF_LLM_HASS_API): user_input.pop(CONF_LLM_HASS_API, None) + # Don't allow to save options that enable the Google Seearch tool with an Assist API if not ( user_input.get(CONF_LLM_HASS_API) and user_input.get(CONF_USE_GOOGLE_SEARCH_TOOL, False) is True ): - # Don't allow to save options that enable the Google Seearch tool with an Assist API - 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, + ) errors[CONF_USE_GOOGLE_SEARCH_TOOL] = "invalid_google_search_option" # Re-render the options again, now with the recommended options shown/hidden self.last_rendered_recommended = user_input[CONF_RECOMMENDED] options = user_input + else: + self.last_rendered_recommended = options.get(CONF_RECOMMENDED, False) schema = await google_generative_ai_config_option_schema( - self.hass, options, self._genai_client + self.hass, self.is_new, options, self._genai_client ) return self.async_show_form( - step_id="init", data_schema=vol.Schema(schema), errors=errors + step_id="set_options", data_schema=vol.Schema(schema), errors=errors ) async def google_generative_ai_config_option_schema( hass: HomeAssistant, + is_new: bool, options: Mapping[str, Any], genai_client: genai.Client, ) -> dict: @@ -224,23 +264,32 @@ async def google_generative_ai_config_option_schema( ): suggested_llm_apis = [suggested_llm_apis] - schema = { - vol.Optional( - CONF_PROMPT, - description={ - "suggested_value": options.get( - CONF_PROMPT, llm.DEFAULT_INSTRUCTIONS_PROMPT - ) - }, - ): TemplateSelector(), - vol.Optional( - CONF_LLM_HASS_API, - description={"suggested_value": suggested_llm_apis}, - ): SelectSelector(SelectSelectorConfig(options=hass_apis, multiple=True)), - vol.Required( - CONF_RECOMMENDED, default=options.get(CONF_RECOMMENDED, False) - ): bool, - } + if is_new: + schema: dict[vol.Required | vol.Optional, Any] = { + vol.Required(CONF_NAME, default=DEFAULT_CONVERSATION_NAME): str, + } + else: + schema = {} + + schema.update( + { + vol.Optional( + CONF_PROMPT, + description={ + "suggested_value": options.get( + CONF_PROMPT, llm.DEFAULT_INSTRUCTIONS_PROMPT + ) + }, + ): TemplateSelector(), + vol.Optional( + CONF_LLM_HASS_API, + description={"suggested_value": suggested_llm_apis}, + ): SelectSelector(SelectSelectorConfig(options=hass_apis, multiple=True)), + vol.Required( + CONF_RECOMMENDED, default=options.get(CONF_RECOMMENDED, False) + ): bool, + } + ) if options.get(CONF_RECOMMENDED): return schema diff --git a/homeassistant/components/google_generative_ai_conversation/const.py b/homeassistant/components/google_generative_ai_conversation/const.py index 831e7d8f508..6e523b2b5e8 100644 --- a/homeassistant/components/google_generative_ai_conversation/const.py +++ b/homeassistant/components/google_generative_ai_conversation/const.py @@ -6,6 +6,8 @@ DOMAIN = "google_generative_ai_conversation" LOGGER = logging.getLogger(__package__) CONF_PROMPT = "prompt" +DEFAULT_CONVERSATION_NAME = "Google Conversation" + ATTR_MODEL = "model" CONF_RECOMMENDED = "recommended" CONF_CHAT_MODEL = "chat_model" diff --git a/homeassistant/components/google_generative_ai_conversation/conversation.py b/homeassistant/components/google_generative_ai_conversation/conversation.py index 00199f5fe1f..d8eae3f6d0d 100644 --- a/homeassistant/components/google_generative_ai_conversation/conversation.py +++ b/homeassistant/components/google_generative_ai_conversation/conversation.py @@ -5,7 +5,7 @@ from __future__ import annotations from typing import Literal from homeassistant.components import assist_pipeline, conversation -from homeassistant.config_entries import ConfigEntry +from homeassistant.config_entries import ConfigEntry, ConfigSubentry from homeassistant.const import CONF_LLM_HASS_API, MATCH_ALL from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError @@ -22,8 +22,14 @@ async def async_setup_entry( async_add_entities: AddConfigEntryEntitiesCallback, ) -> None: """Set up conversation entities.""" - agent = GoogleGenerativeAIConversationEntity(config_entry) - async_add_entities([agent]) + for subentry in config_entry.subentries.values(): + if subentry.subentry_type != "conversation": + continue + + async_add_entities( + [GoogleGenerativeAIConversationEntity(config_entry, subentry)], + config_subentry_id=subentry.subentry_id, + ) class GoogleGenerativeAIConversationEntity( @@ -35,10 +41,10 @@ class GoogleGenerativeAIConversationEntity( _attr_supports_streaming = True - def __init__(self, entry: ConfigEntry) -> None: + def __init__(self, entry: ConfigEntry, subentry: ConfigSubentry) -> None: """Initialize the agent.""" - super().__init__(entry) - if self.entry.options.get(CONF_LLM_HASS_API): + super().__init__(entry, subentry) + if self.subentry.data.get(CONF_LLM_HASS_API): self._attr_supported_features = ( conversation.ConversationEntityFeature.CONTROL ) @@ -70,7 +76,7 @@ class GoogleGenerativeAIConversationEntity( chat_log: conversation.ChatLog, ) -> conversation.ConversationResult: """Call the API.""" - options = self.entry.options + options = self.subentry.data try: await chat_log.async_provide_llm_data( diff --git a/homeassistant/components/google_generative_ai_conversation/entity.py b/homeassistant/components/google_generative_ai_conversation/entity.py index 7eef3dbacff..718dd5a5632 100644 --- a/homeassistant/components/google_generative_ai_conversation/entity.py +++ b/homeassistant/components/google_generative_ai_conversation/entity.py @@ -24,7 +24,7 @@ from google.genai.types import ( from voluptuous_openapi import convert from homeassistant.components import conversation -from homeassistant.config_entries import ConfigEntry +from homeassistant.config_entries import ConfigEntry, ConfigSubentry from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers import device_registry as dr, llm from homeassistant.helpers.entity import Entity @@ -40,6 +40,7 @@ from .const import ( CONF_TOP_K, CONF_TOP_P, CONF_USE_GOOGLE_SEARCH_TOOL, + DEFAULT_CONVERSATION_NAME, DOMAIN, LOGGER, RECOMMENDED_CHAT_MODEL, @@ -301,14 +302,13 @@ async def _transform_stream( class GoogleGenerativeAILLMBaseEntity(Entity): """Google Generative AI base entity.""" - _attr_has_entity_name = True - _attr_name = None - - def __init__(self, entry: ConfigEntry) -> None: + def __init__(self, entry: ConfigEntry, subentry: ConfigSubentry) -> None: """Initialize the agent.""" self.entry = entry + self.subentry = subentry + self._attr_name = subentry.title or DEFAULT_CONVERSATION_NAME self._genai_client = entry.runtime_data - self._attr_unique_id = entry.entry_id + self._attr_unique_id = subentry.subentry_id self._attr_device_info = dr.DeviceInfo( identifiers={(DOMAIN, entry.entry_id)}, name=entry.title, @@ -322,7 +322,7 @@ class GoogleGenerativeAILLMBaseEntity(Entity): chat_log: conversation.ChatLog, ) -> None: """Generate an answer for the chat log.""" - options = self.entry.options + options = self.subentry.data tools: list[Tool | Callable[..., Any]] | None = None if chat_log.llm_api: diff --git a/homeassistant/components/google_generative_ai_conversation/strings.json b/homeassistant/components/google_generative_ai_conversation/strings.json index a57e2f78f53..227c7d6ef68 100644 --- a/homeassistant/components/google_generative_ai_conversation/strings.json +++ b/homeassistant/components/google_generative_ai_conversation/strings.json @@ -21,32 +21,41 @@ "reauth_successful": "[%key:common::config_flow::abort::reauth_successful%]" } }, - "options": { - "step": { - "init": { - "data": { - "recommended": "Recommended model settings", - "prompt": "Instructions", - "chat_model": "[%key:common::generic::model%]", - "temperature": "Temperature", - "top_p": "Top P", - "top_k": "Top K", - "max_tokens": "Maximum tokens to return in response", - "llm_hass_api": "[%key:common::config_flow::data::llm_hass_api%]", - "harassment_block_threshold": "Negative or harmful comments targeting identity and/or protected attributes", - "hate_block_threshold": "Content that is rude, disrespectful, or profane", - "sexual_block_threshold": "Contains references to sexual acts or other lewd content", - "dangerous_block_threshold": "Promotes, facilitates, or encourages harmful acts", - "enable_google_search_tool": "Enable Google Search tool" - }, - "data_description": { - "prompt": "Instruct how the LLM should respond. This can be a template.", - "enable_google_search_tool": "Only works if there is nothing selected in the \"Control Home Assistant\" setting. See docs for a workaround using it with \"Assist\"." + "config_subentries": { + "conversation": { + "initiate_flow": { + "user": "Add conversation agent", + "reconfigure": "Reconfigure conversation agent" + }, + "entry_type": "Conversation agent", + + "step": { + "set_options": { + "data": { + "name": "[%key:common::config_flow::data::name%]", + "recommended": "Recommended model settings", + "prompt": "Instructions", + "chat_model": "[%key:common::generic::model%]", + "temperature": "Temperature", + "top_p": "Top P", + "top_k": "Top K", + "max_tokens": "Maximum tokens to return in response", + "llm_hass_api": "[%key:common::config_flow::data::llm_hass_api%]", + "harassment_block_threshold": "Negative or harmful comments targeting identity and/or protected attributes", + "hate_block_threshold": "Content that is rude, disrespectful, or profane", + "sexual_block_threshold": "Contains references to sexual acts or other lewd content", + "dangerous_block_threshold": "Promotes, facilitates, or encourages harmful acts", + "enable_google_search_tool": "Enable Google Search tool" + }, + "data_description": { + "prompt": "Instruct how the LLM should respond. This can be a template.", + "enable_google_search_tool": "Only works if there is nothing selected in the \"Control Home Assistant\" setting. See docs for a workaround using it with \"Assist\"." + } } + }, + "error": { + "invalid_google_search_option": "Google Search can only be enabled if nothing is selected in the \"Control Home Assistant\" setting." } - }, - "error": { - "invalid_google_search_option": "Google Search can only be enabled if nothing is selected in the \"Control Home Assistant\" setting." } }, "services": { diff --git a/tests/components/google_generative_ai_conversation/conftest.py b/tests/components/google_generative_ai_conversation/conftest.py index f499f18bc15..36d99cd2764 100644 --- a/tests/components/google_generative_ai_conversation/conftest.py +++ b/tests/components/google_generative_ai_conversation/conftest.py @@ -5,8 +5,9 @@ from unittest.mock import AsyncMock, Mock, patch import pytest -from homeassistant.components.google_generative_ai_conversation.entity import ( +from homeassistant.components.google_generative_ai_conversation.const import ( CONF_USE_GOOGLE_SEARCH_TOOL, + DEFAULT_CONVERSATION_NAME, ) from homeassistant.config_entries import ConfigEntry from homeassistant.const import CONF_LLM_HASS_API @@ -26,6 +27,15 @@ def mock_config_entry(hass: HomeAssistant) -> MockConfigEntry: data={ "api_key": "bla", }, + version=2, + subentries_data=[ + { + "data": {}, + "subentry_type": "conversation", + "title": DEFAULT_CONVERSATION_NAME, + "unique_id": None, + } + ], ) entry.runtime_data = Mock() entry.add_to_hass(hass) @@ -38,8 +48,10 @@ async def mock_config_entry_with_assist( ) -> MockConfigEntry: """Mock a config entry with assist.""" with patch("google.genai.models.AsyncModels.get"): - hass.config_entries.async_update_entry( - mock_config_entry, options={CONF_LLM_HASS_API: llm.LLM_API_ASSIST} + hass.config_entries.async_update_subentry( + mock_config_entry, + next(iter(mock_config_entry.subentries.values())), + data={CONF_LLM_HASS_API: llm.LLM_API_ASSIST}, ) await hass.async_block_till_done() return mock_config_entry @@ -51,9 +63,10 @@ async def mock_config_entry_with_google_search( ) -> MockConfigEntry: """Mock a config entry with assist.""" with patch("google.genai.models.AsyncModels.get"): - hass.config_entries.async_update_entry( + hass.config_entries.async_update_subentry( mock_config_entry, - options={ + next(iter(mock_config_entry.subentries.values())), + data={ CONF_LLM_HASS_API: llm.LLM_API_ASSIST, CONF_USE_GOOGLE_SEARCH_TOOL: True, }, diff --git a/tests/components/google_generative_ai_conversation/test_config_flow.py b/tests/components/google_generative_ai_conversation/test_config_flow.py index 4234355cb5b..d429ace68da 100644 --- a/tests/components/google_generative_ai_conversation/test_config_flow.py +++ b/tests/components/google_generative_ai_conversation/test_config_flow.py @@ -30,7 +30,7 @@ from homeassistant.components.google_generative_ai_conversation.const import ( RECOMMENDED_TOP_P, RECOMMENDED_USE_GOOGLE_SEARCH_TOOL, ) -from homeassistant.const import CONF_LLM_HASS_API +from homeassistant.const import CONF_LLM_HASS_API, CONF_NAME from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import FlowResultType @@ -110,10 +110,58 @@ async def test_form(hass: HomeAssistant) -> None: assert result2["data"] == { "api_key": "bla", } - assert result2["options"] == RECOMMENDED_OPTIONS + assert result2["options"] == {} + assert result2["subentries"] == [ + { + "subentry_type": "conversation", + "data": RECOMMENDED_OPTIONS, + "title": "Google Conversation", + "unique_id": None, + } + ] assert len(mock_setup_entry.mock_calls) == 1 +async def test_creating_conversation_subentry( + hass: HomeAssistant, + mock_init_component: None, + mock_config_entry: MockConfigEntry, +) -> None: + """Test creating a conversation subentry.""" + mock_config_entry.add_to_hass(hass) + + with patch( + "google.genai.models.AsyncModels.list", + return_value=get_models_pager(), + ): + result = await hass.config_entries.subentries.async_init( + (mock_config_entry.entry_id, "conversation"), + context={"source": config_entries.SOURCE_USER}, + ) + + assert result["type"] is FlowResultType.FORM + assert result["step_id"] == "set_options" + assert not result["errors"] + + with patch( + "google.genai.models.AsyncModels.list", + return_value=get_models_pager(), + ): + result2 = await hass.config_entries.subentries.async_configure( + result["flow_id"], + {CONF_NAME: "Mock name", **RECOMMENDED_OPTIONS}, + ) + await hass.async_block_till_done() + + assert result2["type"] is FlowResultType.CREATE_ENTRY + assert result2["title"] == "Mock name" + + processed_options = RECOMMENDED_OPTIONS.copy() + processed_options[CONF_PROMPT] = processed_options[CONF_PROMPT].strip() + + assert result2["data"] == processed_options + + def will_options_be_rendered_again(current_options, new_options) -> bool: """Determine if options will be rendered again.""" return current_options.get(CONF_RECOMMENDED) != new_options.get(CONF_RECOMMENDED) @@ -283,7 +331,7 @@ def will_options_be_rendered_again(current_options, new_options) -> bool: ], ) @pytest.mark.usefixtures("mock_init_component") -async def test_options_switching( +async def test_subentry_options_switching( hass: HomeAssistant, mock_config_entry: MockConfigEntry, current_options, @@ -292,17 +340,18 @@ async def test_options_switching( errors, ) -> None: """Test the options form.""" + subentry = next(iter(mock_config_entry.subentries.values())) with patch("google.genai.models.AsyncModels.get"): - hass.config_entries.async_update_entry( - mock_config_entry, options=current_options + hass.config_entries.async_update_subentry( + mock_config_entry, subentry, data=current_options ) await hass.async_block_till_done() with patch( "google.genai.models.AsyncModels.list", return_value=get_models_pager(), ): - options_flow = await hass.config_entries.options.async_init( - mock_config_entry.entry_id + options_flow = await mock_config_entry.start_subentry_reconfigure_flow( + hass, subentry.subentry_type, subentry.subentry_id ) if will_options_be_rendered_again(current_options, new_options): retry_options = { @@ -313,7 +362,7 @@ async def test_options_switching( "google.genai.models.AsyncModels.list", return_value=get_models_pager(), ): - options_flow = await hass.config_entries.options.async_configure( + options_flow = await hass.config_entries.subentries.async_configure( options_flow["flow_id"], retry_options, ) @@ -321,14 +370,15 @@ async def test_options_switching( "google.genai.models.AsyncModels.list", return_value=get_models_pager(), ): - options = await hass.config_entries.options.async_configure( + options = await hass.config_entries.subentries.async_configure( options_flow["flow_id"], new_options, ) - await hass.async_block_till_done() + await hass.async_block_till_done() if errors is None: - assert options["type"] is FlowResultType.CREATE_ENTRY - assert options["data"] == expected_options + assert options["type"] is FlowResultType.ABORT + assert options["reason"] == "reconfigure_successful" + assert subentry.data == expected_options else: assert options["type"] is FlowResultType.FORM diff --git a/tests/components/google_generative_ai_conversation/test_conversation.py b/tests/components/google_generative_ai_conversation/test_conversation.py index 92aa6f08d42..e0a20b781bf 100644 --- a/tests/components/google_generative_ai_conversation/test_conversation.py +++ b/tests/components/google_generative_ai_conversation/test_conversation.py @@ -64,7 +64,7 @@ async def test_error_handling( "hello", None, Context(), - agent_id="conversation.google_generative_ai_conversation", + agent_id="conversation.google_conversation", ) assert result.response.response_type == intent.IntentResponseType.ERROR, result assert result.response.error_code == "unknown", result @@ -82,7 +82,7 @@ async def test_function_call( mock_send_message_stream: AsyncMock, ) -> None: """Test function calling.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -212,7 +212,7 @@ async def test_google_search_tool_is_sent( mock_send_message_stream: AsyncMock, ) -> None: """Test if the Google Search tool is sent to the model.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -278,7 +278,7 @@ async def test_blocked_response( mock_send_message_stream: AsyncMock, ) -> None: """Test blocked response.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -328,7 +328,7 @@ async def test_empty_response( ) -> None: """Test empty response.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -371,7 +371,7 @@ async def test_none_response( mock_send_message_stream: AsyncMock, ) -> None: """Test None response.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -403,10 +403,12 @@ async def test_converse_error( hass: HomeAssistant, mock_config_entry: MockConfigEntry ) -> None: """Test handling ChatLog raising ConverseError.""" + subentry = next(iter(mock_config_entry.subentries.values())) with patch("google.genai.models.AsyncModels.get"): - hass.config_entries.async_update_entry( + hass.config_entries.async_update_subentry( mock_config_entry, - options={**mock_config_entry.options, CONF_LLM_HASS_API: "invalid_llm_api"}, + next(iter(mock_config_entry.subentries.values())), + data={**subentry.data, CONF_LLM_HASS_API: "invalid_llm_api"}, ) await hass.async_block_till_done() @@ -415,7 +417,7 @@ async def test_converse_error( "hello", None, Context(), - agent_id="conversation.google_generative_ai_conversation", + agent_id="conversation.google_conversation", ) assert result.response.response_type == intent.IntentResponseType.ERROR, result @@ -593,7 +595,7 @@ async def test_empty_content_in_chat_history( mock_send_message_stream: AsyncMock, ) -> None: """Tests that in case of an empty entry in the chat history the google API will receive an injected space sign instead.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -648,7 +650,7 @@ async def test_history_always_user_first_turn( ) -> None: """Test that the user is always first in the chat history.""" - agent_id = "conversation.google_generative_ai_conversation" + agent_id = "conversation.google_conversation" context = Context() messages = [ @@ -674,7 +676,7 @@ async def test_history_always_user_first_turn( mock_chat_log.async_add_assistant_content_without_tools( conversation.AssistantContent( - agent_id="conversation.google_generative_ai_conversation", + agent_id="conversation.google_conversation", content="Garage door left open, do you want to close it?", ) ) diff --git a/tests/components/google_generative_ai_conversation/test_init.py b/tests/components/google_generative_ai_conversation/test_init.py index 6cc0bdd5f44..306c567d8e5 100644 --- a/tests/components/google_generative_ai_conversation/test_init.py +++ b/tests/components/google_generative_ai_conversation/test_init.py @@ -7,9 +7,14 @@ import pytest from requests.exceptions import Timeout from syrupy.assertion import SnapshotAssertion +from homeassistant.components.google_generative_ai_conversation import ( + async_migrate_entry, +) +from homeassistant.components.google_generative_ai_conversation.const import DOMAIN from homeassistant.config_entries import ConfigEntryState from homeassistant.core import HomeAssistant from homeassistant.exceptions import HomeAssistantError +from homeassistant.helpers import device_registry as dr, entity_registry as er from . import API_ERROR_500, CLIENT_ERROR_API_KEY_INVALID @@ -387,3 +392,65 @@ async def test_load_entry_with_unloaded_entries( "text": stubbed_generated_content, } assert [tuple(mock_call) for mock_call in mock_generate.mock_calls] == snapshot + + +async def test_migration_from_v1_to_v2( + hass: HomeAssistant, + device_registry: dr.DeviceRegistry, + entity_registry: er.EntityRegistry, +) -> None: + """Test migration from version 1 to version 2.""" + # Create a v1 config entry with conversation options and an entity + OPTIONS = { + "recommended": True, + "llm_hass_api": ["assist"], + "prompt": "You are a helpful assistant", + "chat_model": "models/gemini-2.0-flash", + } + mock_config_entry = MockConfigEntry( + domain=DOMAIN, + data={"api_key": "1234"}, + options=OPTIONS, + version=1, + title="Google Generative AI", + ) + mock_config_entry.add_to_hass(hass) + + device = device_registry.async_get_or_create( + config_entry_id=mock_config_entry.entry_id, + identifiers={(DOMAIN, mock_config_entry.entry_id)}, + name=mock_config_entry.title, + manufacturer="Google", + model="Generative AI", + entry_type=dr.DeviceEntryType.SERVICE, + ) + entity = entity_registry.async_get_or_create( + "conversation", + DOMAIN, + "mock_config_entry.entry_id", + config_entry=mock_config_entry, + device_id=device.id, + suggested_object_id="google_generative_ai_conversation", + ) + + # Run migration + result = await async_migrate_entry(hass, mock_config_entry) + + assert result is True + assert mock_config_entry.version == 2 + assert mock_config_entry.data == {"api_key": "1234"} + assert mock_config_entry.options == {} + + assert len(mock_config_entry.subentries) == 1 + + subentry = next(iter(mock_config_entry.subentries.values())) + assert subentry.unique_id is None + assert subentry.title == "Google Conversation" + assert subentry.subentry_type == "conversation" + assert subentry.data == OPTIONS + + migrated_entity = entity_registry.async_get(entity.entity_id) + assert migrated_entity is not None + assert migrated_entity.config_entry_id == mock_config_entry.entry_id + assert migrated_entity.config_subentry_id == subentry.subentry_id + assert migrated_entity.device_id == device.id