Add latest changes from Google subentries

This commit is contained in:
Paulus Schoutsen
2025-06-24 01:59:51 +00:00
parent f88ea00be6
commit d8caa17266
7 changed files with 331 additions and 46 deletions

View File

@ -10,15 +10,14 @@ from homeassistant.config_entries import ConfigEntry, ConfigSubentry
from homeassistant.const import CONF_API_KEY, Platform
from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import config_validation as cv, entity_registry as er
from .const import (
CONF_CHAT_MODEL,
DEFAULT_CONVERSATION_NAME,
DOMAIN,
LOGGER,
RECOMMENDED_CHAT_MODEL,
from homeassistant.helpers import (
config_validation as cv,
device_registry as dr,
entity_registry as er,
)
from homeassistant.helpers.typing import ConfigType
from .const import CONF_CHAT_MODEL, DOMAIN, LOGGER, RECOMMENDED_CHAT_MODEL
PLATFORMS = (Platform.CONVERSATION,)
CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
@ -26,6 +25,12 @@ CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
type AnthropicConfigEntry = ConfigEntry[anthropic.AsyncClient]
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up Anthropic."""
await async_migrate_integration(hass)
return True
async def async_setup_entry(hass: HomeAssistant, entry: AnthropicConfigEntry) -> bool:
"""Set up Anthropic from a config entry."""
client = await hass.async_add_executor_job(
@ -58,41 +63,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: AnthropicConfigEntry) -> 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

@ -12,6 +12,7 @@ import voluptuous as vol
from homeassistant.config_entries import (
ConfigEntry,
ConfigEntryState,
ConfigFlow,
ConfigFlowResult,
ConfigSubentryFlow,
@ -82,6 +83,7 @@ class AnthropicConfigFlow(ConfigFlow, domain=DOMAIN):
errors = {}
if user_input is not None:
self._async_abort_entries_match(user_input)
try:
await validate_input(self.hass, user_input)
except anthropic.APITimeoutError:
@ -140,6 +142,10 @@ class ConversationSubentryFlowHandler(ConfigSubentryFlow):
self, user_input: dict[str, Any] | None = None
) -> SubentryFlowResult:
"""Set conversation options."""
# abort if entry is not loaded
if self._get_entry().state != ConfigEntryState.LOADED:
return self.async_abort(reason="entry_not_loaded")
errors: dict[str, str] = {}
if user_input is None:

View File

@ -52,7 +52,6 @@ from .const import (
CONF_PROMPT,
CONF_TEMPERATURE,
CONF_THINKING_BUDGET,
DEFAULT_CONVERSATION_NAME,
DOMAIN,
LOGGER,
MIN_THINKING_BUDGET,
@ -339,10 +338,11 @@ class AnthropicConversationEntity(
"""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)},
identifiers={(DOMAIN, subentry.subentry_id)},
name=subentry.title,
manufacturer="Anthropic",
model="Claude",
entry_type=dr.DeviceEntryType.SERVICE,

View File

@ -12,6 +12,9 @@
"timeout_connect": "[%key:common::config_flow::error::timeout_connect%]",
"authentication_error": "[%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": {
@ -41,7 +44,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": {
"thinking_budget_too_large": "Maximum tokens must be greater than the thinking budget."

View File

@ -267,11 +267,6 @@ async def async_migrate_integration(hass: HomeAssistant) -> None:
add_config_entry_id=parent_entry.entry_id,
)
if parent_entry.entry_id != entry.entry_id:
device_registry.async_update_device(
device.id,
add_config_subentry_id=subentry.subentry_id,
add_config_entry_id=parent_entry.entry_id,
)
device_registry.async_update_device(
device.id,
remove_config_entry_id=entry.entry_id,

View File

@ -1,6 +1,6 @@
"""Test the Anthropic config flow."""
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, Mock, patch
from anthropic import (
APIConnectionError,
@ -28,7 +28,7 @@ from homeassistant.components.anthropic.const import (
RECOMMENDED_MAX_TOKENS,
RECOMMENDED_THINKING_BUDGET,
)
from homeassistant.const import CONF_LLM_HASS_API, CONF_NAME
from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API, CONF_NAME
from homeassistant.core import HomeAssistant
from homeassistant.data_entry_flow import FlowResultType
@ -84,6 +84,34 @@ 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(
"anthropic.resources.models.AsyncModels.retrieve",
return_value=Mock(display_name="Claude 3.5 Sonnet"),
):
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_config_entry, mock_init_component
) -> None:
@ -112,13 +140,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(
"anthropic.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_options_thinking_budget_more_than_max(
hass: HomeAssistant, mock_config_entry, mock_init_component
) -> None:
"""Test error about thinking budget being more than max tokens."""
subentry = next(iter(mock_config_entry.subentries.values()))
options_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(
options_flow["flow_id"],
@ -128,6 +176,7 @@ async def test_subentry_options_thinking_budget_more_than_max(
"chat_model": "claude-3-7-sonnet-latest",
"temperature": 1,
"thinking_budget": 16384,
"recommended": False,
},
)
await hass.async_block_till_done()
@ -285,7 +334,7 @@ async def test_subentry_options_switching(
await hass.async_block_till_done()
options_flow = await mock_config_entry.start_subentry_reconfigure_flow(
hass, subentry.subentry_type, subentry.subentry_id
hass, subentry.subentry_id
)
if current_options.get(CONF_RECOMMENDED) != new_options.get(CONF_RECOMMENDED):
options_flow = await hass.config_entries.subentries.async_configure(

View File

@ -11,7 +11,7 @@ from anthropic import (
from httpx import URL, Request, Response
import pytest
from homeassistant.components.anthropic.const import DEFAULT_CONVERSATION_NAME, DOMAIN
from homeassistant.components.anthropic.const import DOMAIN
from homeassistant.core import HomeAssistant
from homeassistant.helpers import device_registry as dr, entity_registry as er
from homeassistant.setup import async_setup_component
@ -91,8 +91,8 @@ async def test_migration_from_v1_to_v2(
config_entry_id=mock_config_entry.entry_id,
identifiers={(DOMAIN, mock_config_entry.entry_id)},
name=mock_config_entry.title,
manufacturer="Claude",
model="Anthropic",
manufacturer="Anthropic",
model="Claude",
entry_type=dr.DeviceEntryType.SERVICE,
)
entity = entity_registry.async_get_or_create(
@ -110,6 +110,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"}
@ -119,7 +120,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 == "Claude"
assert subentry.subentry_type == "conversation"
assert subentry.data == OPTIONS
@ -127,4 +128,204 @@ 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": "claude-3-haiku-20240307",
}
mock_config_entry = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "1234"},
options=options,
version=1,
title="Claude 1",
)
mock_config_entry.add_to_hass(hass)
mock_config_entry_2 = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "12345"},
options=options,
version=1,
title="Claude 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="Anthropic",
model="Claude 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="claude_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="Anthropic",
model="Claude 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="claude_2",
)
# Run migration
with patch(
"homeassistant.components.anthropic.async_setup_entry",
return_value=True,
):
await hass.config_entries.async_setup(mock_config_entry.entry_id)
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"Claude {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": "claude-3-haiku-20240307",
}
mock_config_entry = MockConfigEntry(
domain=DOMAIN,
data={"api_key": "1234"},
options=options,
version=1,
title="Claude",
)
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="Claude 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="Anthropic",
model="Claude",
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="claude",
)
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="Anthropic",
model="Claude",
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="claude_2",
)
# Run migration
with patch(
"homeassistant.components.anthropic.async_setup_entry",
return_value=True,
):
await hass.config_entries.async_setup(mock_config_entry.entry_id)
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 "Claude" in titles
assert "Claude 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