Add latest changes from Google subentries

This commit is contained in:
Paulus Schoutsen
2025-06-24 02:54:53 +00:00
parent 566c5c6313
commit a604cefb8e
6 changed files with 320 additions and 38 deletions

View File

@ -34,6 +34,7 @@ from homeassistant.exceptions import (
)
from homeassistant.helpers import (
config_validation as cv,
device_registry as dr,
entity_registry as er,
selector,
)
@ -48,7 +49,6 @@ from .const import (
CONF_REASONING_EFFORT,
CONF_TEMPERATURE,
CONF_TOP_P,
DEFAULT_CONVERSATION_NAME,
DOMAIN,
LOGGER,
RECOMMENDED_CHAT_MODEL,
@ -78,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."""
@ -290,41 +291,66 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
async def async_migrate_entry(hass: HomeAssistant, entry: OpenAIConfigEntry) -> bool:
"""Migrate old entry."""
if entry.version == 1:
# Migrate from version 1 to version 2
# Move conversation-specific options to a subentry
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=DEFAULT_CONVERSATION_NAME,
title=entry.title,
unique_id=None,
)
hass.config_entries.async_add_subentry(
entry,
subentry,
)
if entry.data[CONF_API_KEY] not in api_keys_entries:
use_existing = True
api_keys_entries[entry.data[CONF_API_KEY]] = entry
# Migrate conversation entity to be linked to subentry
ent_reg = er.async_get(hass)
conversation_entity = ent_reg.async_get_entity_id(
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:
ent_reg.async_update_entity(
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,
)
# Remove options from the main entry
hass.config_entries.async_update_entry(
entry,
options={},
version=2,
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,
)
return True
if not use_existing:
await hass.config_entries.async_remove(entry.entry_id)
else:
hass.config_entries.async_update_entry(
entry,
options={},
version=2,
)

View File

@ -13,6 +13,7 @@ from voluptuous_openapi import convert
from homeassistant.components.zone import ENTITY_ID_HOME
from homeassistant.config_entries import (
ConfigEntry,
ConfigEntryState,
ConfigFlow,
ConfigFlowResult,
ConfigSubentryFlow,
@ -110,6 +111,7 @@ class OpenAIConfigFlow(ConfigFlow, domain=DOMAIN):
errors: dict[str, str] = {}
self._async_abort_entries_match(user_input)
try:
await validate_input(self.hass, user_input)
except openai.APIConnectionError:
@ -175,6 +177,9 @@ class ConversationSubentryFlowHandler(ConfigSubentryFlow):
self, user_input: dict[str, Any] | None = None
) -> 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] = [

View File

@ -56,7 +56,6 @@ from .const import (
CONF_WEB_SEARCH_REGION,
CONF_WEB_SEARCH_TIMEZONE,
CONF_WEB_SEARCH_USER_LOCATION,
DEFAULT_CONVERSATION_NAME,
DOMAIN,
LOGGER,
RECOMMENDED_CHAT_MODEL,
@ -242,13 +241,13 @@ class OpenAIConversationEntity(
"""Initialize the agent."""
self.entry = entry
self.subentry = subentry
self._attr_name = subentry.title or DEFAULT_CONVERSATION_NAME
self._attr_name = subentry.title
self._attr_unique_id = subentry.subentry_id
self._attr_device_info = dr.DeviceInfo(
identifiers={(DOMAIN, entry.entry_id)},
name=entry.title,
identifiers={(DOMAIN, subentry.subentry_id)},
name=subentry.title,
manufacturer="OpenAI",
model="ChatGPT",
model=entry.data.get(CONF_CHAT_MODEL, RECOMMENDED_CHAT_MODEL),
entry_type=dr.DeviceEntryType.SERVICE,
)
if self.subentry.data.get(CONF_LLM_HASS_API):

View File

@ -11,6 +11,9 @@
"cannot_connect": "[%key:common::config_flow::error::cannot_connect%]",
"invalid_auth": "[%key:common::config_flow::error::invalid_auth%]",
"unknown": "[%key:common::config_flow::error::unknown%]"
},
"abort": {
"already_configured": "[%key:common::config_flow::abort::already_configured_service%]"
}
},
"config_subentries": {
@ -59,7 +62,8 @@
}
},
"abort": {
"reconfigure_successful": "[%key:common::config_flow::abort::reconfigure_successful%]"
"reconfigure_successful": "[%key:common::config_flow::abort::reconfigure_successful%]",
"entry_not_loaded": "Cannot add things while the configuration is disabled."
},
"error": {
"model_not_supported": "This model is not supported, please select a different model"

View File

@ -30,7 +30,7 @@ from homeassistant.components.openai_conversation.const import (
RECOMMENDED_MAX_TOKENS,
RECOMMENDED_TOP_P,
)
from homeassistant.const import CONF_LLM_HASS_API
from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType
@ -85,6 +85,33 @@ async def test_form(hass: HomeAssistant) -> None:
assert len(mock_setup_entry.mock_calls) == 1
async def test_duplicate_entry(hass: HomeAssistant) -> None:
"""Test we abort on duplicate config entry."""
MockConfigEntry(
domain=DOMAIN,
data={CONF_API_KEY: "bla"},
).add_to_hass(hass)
result = await hass.config_entries.flow.async_init(
DOMAIN, context={"source": config_entries.SOURCE_USER}
)
assert result["type"] is FlowResultType.FORM
assert not result["errors"]
with patch(
"homeassistant.components.openai_conversation.config_flow.openai.resources.models.AsyncModels.list",
):
result = await hass.config_entries.flow.async_configure(
result["flow_id"],
{
CONF_API_KEY: "bla",
},
)
assert result["type"] is FlowResultType.ABORT
assert result["reason"] == "already_configured"
async def test_creating_conversation_subentry(
hass: HomeAssistant,
mock_init_component: None,
@ -117,13 +144,33 @@ async def test_creating_conversation_subentry(
assert result2["data"] == processed_options
async def test_creating_conversation_subentry_not_loaded(
hass: HomeAssistant,
mock_init_component,
mock_config_entry: MockConfigEntry,
) -> None:
"""Test creating a conversation subentry when entry is not loaded."""
await hass.config_entries.async_unload(mock_config_entry.entry_id)
with patch(
"homeassistant.components.openai_conversation.config_flow.openai.resources.models.AsyncModels.list",
return_value=[],
):
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.ABORT
assert result["reason"] == "entry_not_loaded"
async def test_subentry_recommended(
hass: HomeAssistant, mock_config_entry, mock_init_component
) -> None:
"""Test the subentry flow with recommended settings."""
subentry = next(iter(mock_config_entry.subentries.values()))
subentry_flow = await mock_config_entry.start_subentry_reconfigure_flow(
hass, subentry.subentry_type, subentry.subentry_id
hass, subentry.subentry_id
)
options = await hass.config_entries.subentries.async_configure(
subentry_flow["flow_id"],
@ -144,7 +191,7 @@ async def test_subentry_unsupported_model(
"""Test the subentry form giving error about models not supported."""
subentry = next(iter(mock_config_entry.subentries.values()))
subentry_flow = await mock_config_entry.start_subentry_reconfigure_flow(
hass, subentry.subentry_type, subentry.subentry_id
hass, subentry.subentry_id
)
assert subentry_flow["type"] == FlowResultType.FORM
assert subentry_flow["step_id"] == "init"
@ -551,8 +598,9 @@ async def test_subentry_switching(
hass.config_entries.async_update_subentry(
mock_config_entry, subentry, data=current_options
)
await hass.async_block_till_done()
subentry_flow = await mock_config_entry.start_subentry_reconfigure_flow(
hass, subentry.subentry_type, subentry.subentry_id
hass, subentry.subentry_id
)
assert subentry_flow["step_id"] == "init"
@ -589,7 +637,7 @@ async def test_subentry_web_search_user_location(
"""Test fetching user location."""
subentry = next(iter(mock_config_entry.subentries.values()))
subentry_flow = await mock_config_entry.start_subentry_reconfigure_flow(
hass, subentry.subentry_type, subentry.subentry_id
hass, subentry.subentry_id
)
assert subentry_flow["type"] == FlowResultType.FORM
assert subentry_flow["step_id"] == "init"

View File

@ -15,10 +15,7 @@ from openai.types.responses import Response, ResponseOutputMessage, ResponseOutp
import pytest
from homeassistant.components.openai_conversation import CONF_FILENAMES
from homeassistant.components.openai_conversation.const import (
DEFAULT_CONVERSATION_NAME,
DOMAIN,
)
from homeassistant.components.openai_conversation.const import DOMAIN
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
from homeassistant.helpers import device_registry as dr, entity_registry as er
@ -588,6 +585,7 @@ async def test_migration_from_v1_to_v2(
return_value=True,
):
await hass.config_entries.async_setup(mock_config_entry.entry_id)
await hass.async_block_till_done()
assert mock_config_entry.version == 2
assert mock_config_entry.data == {"api_key": "1234"}
@ -597,7 +595,7 @@ async def test_migration_from_v1_to_v2(
subentry = next(iter(mock_config_entry.subentries.values()))
assert subentry.unique_id is None
assert subentry.title == DEFAULT_CONVERSATION_NAME
assert subentry.title == "ChatGPT"
assert subentry.subentry_type == "conversation"
assert subentry.data == OPTIONS
@ -605,4 +603,206 @@ async def test_migration_from_v1_to_v2(
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
assert migrated_entity.unique_id == subentry.subentry_id
# Check device migration
assert not device_registry.async_get_device(
identifiers={(DOMAIN, mock_config_entry.entry_id)}
)
assert (
migrated_device := device_registry.async_get_device(
identifiers={(DOMAIN, subentry.subentry_id)}
)
)
assert migrated_device.identifiers == {(DOMAIN, subentry.subentry_id)}
assert migrated_device.id == device.id
async def test_migration_from_v1_to_v2_with_multiple_keys(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test migration from version 1 to version 2 with different API keys."""
# Create two v1 config entries with different API keys
options = {
"recommended": True,
"llm_hass_api": ["assist"],
"prompt": "You are a helpful assistant",
"chat_model": "gpt-4o-mini",
}
mock_config_entry = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "1234"},
options=options,
version=1,
title="ChatGPT 1",
)
mock_config_entry.add_to_hass(hass)
mock_config_entry_2 = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "12345"},
options=options,
version=1,
title="ChatGPT 2",
)
mock_config_entry_2.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="OpenAI",
model="ChatGPT 1",
entry_type=dr.DeviceEntryType.SERVICE,
)
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="chatgpt_1",
)
device_2 = device_registry.async_get_or_create(
config_entry_id=mock_config_entry_2.entry_id,
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
name=mock_config_entry_2.title,
manufacturer="OpenAI",
model="ChatGPT 2",
entry_type=dr.DeviceEntryType.SERVICE,
)
entity_registry.async_get_or_create(
"conversation",
DOMAIN,
mock_config_entry_2.entry_id,
config_entry=mock_config_entry_2,
device_id=device_2.id,
suggested_object_id="chatgpt_2",
)
# Run migration
with patch(
"homeassistant.components.openai_conversation.async_setup_entry",
return_value=True,
):
await hass.config_entries.async_setup(mock_config_entry.entry_id)
await hass.async_block_till_done()
await hass.async_block_till_done()
entries = hass.config_entries.async_entries(DOMAIN)
assert len(entries) == 2
for idx, entry in enumerate(entries):
assert entry.version == 2
assert not entry.options
assert len(entry.subentries) == 1
subentry = list(entry.subentries.values())[0]
assert subentry.subentry_type == "conversation"
assert subentry.data == options
assert subentry.title == f"ChatGPT {idx + 1}"
dev = device_registry.async_get_device(
identifiers={(DOMAIN, list(entry.subentries.values())[0].subentry_id)}
)
assert dev is not None
async def test_migration_from_v1_to_v2_with_same_keys(
hass: HomeAssistant,
device_registry: dr.DeviceRegistry,
entity_registry: er.EntityRegistry,
) -> None:
"""Test migration from version 1 to version 2 with same API keys consolidates entries."""
# Create two v1 config entries with the same API key
options = {
"recommended": True,
"llm_hass_api": ["assist"],
"prompt": "You are a helpful assistant",
"chat_model": "gpt-4o-mini",
}
mock_config_entry = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "1234"},
options=options,
version=1,
title="ChatGPT",
)
mock_config_entry.add_to_hass(hass)
mock_config_entry_2 = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "1234"}, # Same API key
options=options,
version=1,
title="ChatGPT 2",
)
mock_config_entry_2.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="OpenAI",
model="ChatGPT",
entry_type=dr.DeviceEntryType.SERVICE,
)
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="chatgpt",
)
device_2 = device_registry.async_get_or_create(
config_entry_id=mock_config_entry_2.entry_id,
identifiers={(DOMAIN, mock_config_entry_2.entry_id)},
name=mock_config_entry_2.title,
manufacturer="OpenAI",
model="ChatGPT",
entry_type=dr.DeviceEntryType.SERVICE,
)
entity_registry.async_get_or_create(
"conversation",
DOMAIN,
mock_config_entry_2.entry_id,
config_entry=mock_config_entry_2,
device_id=device_2.id,
suggested_object_id="chatgpt_2",
)
# Run migration
with patch(
"homeassistant.components.openai_conversation.async_setup_entry",
return_value=True,
):
await hass.config_entries.async_setup(mock_config_entry.entry_id)
await hass.async_block_till_done()
await hass.async_block_till_done()
# Should have only one entry left (consolidated)
entries = hass.config_entries.async_entries(DOMAIN)
assert len(entries) == 1
entry = entries[0]
assert entry.version == 2
assert not entry.options
assert len(entry.subentries) == 2 # Two subentries from the two original entries
# Check both subentries exist with correct data
subentries = list(entry.subentries.values())
titles = [sub.title for sub in subentries]
assert "ChatGPT" in titles
assert "ChatGPT 2" in titles
for subentry in subentries:
assert subentry.subentry_type == "conversation"
assert subentry.data == options
# Check devices were migrated correctly
dev = device_registry.async_get_device(
identifiers={(DOMAIN, subentry.subentry_id)}
)
assert dev is not None