Move default conversation agent to manager (home-assistant#152479)

Author: arturpragacz

homeassistant/components/conversation/__init__.py

@@ -50,14 +50,13 @@
     ATTR_LANGUAGE,
     ATTR_TEXT,
     DATA_COMPONENT,
-    DATA_DEFAULT_ENTITY,
     DOMAIN,
     HOME_ASSISTANT_AGENT,
     SERVICE_PROCESS,
     SERVICE_RELOAD,
     ConversationEntityFeature,
 )
-from .default_agent import DefaultAgent, async_setup_default_agent
+from .default_agent import async_setup_default_agent
 from .entity import ConversationEntity
 from .http import async_setup as async_setup_conversation_http
 from .models import AbstractConversationAgent, ConversationInput, ConversationResult
@@ -142,7 +141,7 @@ def async_unset_agent(
     hass: HomeAssistant,
     config_entry: ConfigEntry,
 ) -> None:
-    """Set the agent to handle the conversations."""
+    """Unset the agent to handle the conversations."""
     get_agent_manager(hass).async_unset_agent(config_entry.entry_id)
 
 
@@ -241,10 +240,10 @@ async def async_handle_sentence_triggers(
 
     Returns None if no match occurred.
     """
-    default_agent = async_get_agent(hass)
-    assert isinstance(default_agent, DefaultAgent)
+    agent = get_agent_manager(hass).default_agent
+    assert agent is not None
 
-    return await default_agent.async_handle_sentence_triggers(user_input)
+    return await agent.async_handle_sentence_triggers(user_input)
 
 
 async def async_handle_intents(
@@ -257,12 +256,10 @@ async def async_handle_intents(
 
     Returns None if no match occurred.
     """
-    default_agent = async_get_agent(hass)
-    assert isinstance(default_agent, DefaultAgent)
+    agent = get_agent_manager(hass).default_agent
+    assert agent is not None
 
-    return await default_agent.async_handle_intents(
-        user_input, intent_filter=intent_filter
-    )
+    return await agent.async_handle_intents(user_input, intent_filter=intent_filter)
 
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
@@ -298,9 +295,9 @@ async def handle_process(service: ServiceCall) -> ServiceResponse:
 
     async def handle_reload(service: ServiceCall) -> None:
         """Reload intents."""
-        await hass.data[DATA_DEFAULT_ENTITY].async_reload(
-            language=service.data.get(ATTR_LANGUAGE)
-        )
+        agent = get_agent_manager(hass).default_agent
+        if agent is not None:
+            await agent.async_reload(language=service.data.get(ATTR_LANGUAGE))
 
     hass.services.async_register(
         DOMAIN,

homeassistant/components/conversation/agent_manager.py

@@ -4,15 +4,15 @@
 
 import dataclasses
 import logging
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import voluptuous as vol
 
 from homeassistant.core import Context, HomeAssistant, async_get_hass, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import config_validation as cv, intent, singleton
 
-from .const import DATA_COMPONENT, DATA_DEFAULT_ENTITY, HOME_ASSISTANT_AGENT
+from .const import DATA_COMPONENT, HOME_ASSISTANT_AGENT
 from .entity import ConversationEntity
 from .models import (
     AbstractConversationAgent,
@@ -28,6 +28,9 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+if TYPE_CHECKING:
+    from .default_agent import DefaultAgent
+
 
 @singleton.singleton("conversation_agent")
 @callback
@@ -49,8 +52,10 @@ def async_get_agent(
     hass: HomeAssistant, agent_id: str | None = None
 ) -> AbstractConversationAgent | ConversationEntity | None:
     """Get specified agent."""
+    manager = get_agent_manager(hass)
+
     if agent_id is None or agent_id == HOME_ASSISTANT_AGENT:
-        return hass.data[DATA_DEFAULT_ENTITY]
+        return manager.default_agent
 
     if "." in agent_id:
         return hass.data[DATA_COMPONENT].get_entity(agent_id)
@@ -134,6 +139,7 @@ def __init__(self, hass: HomeAssistant) -> None:
         """Initialize the conversation agents."""
         self.hass = hass
         self._agents: dict[str, AbstractConversationAgent] = {}
+        self.default_agent: DefaultAgent | None = None
 
     @callback
     def async_get_agent(self, agent_id: str) -> AbstractConversationAgent | None:
@@ -182,3 +188,7 @@ def async_set_agent(self, agent_id: str, agent: AbstractConversationAgent) -> No
     def async_unset_agent(self, agent_id: str) -> None:
         """Unset the agent."""
         self._agents.pop(agent_id, None)
+
+    async def async_setup_default_agent(self, agent: DefaultAgent) -> None:
+        """Set up the default agent."""
+        self.default_agent = agent

homeassistant/components/conversation/const.py

@@ -10,11 +10,9 @@
 if TYPE_CHECKING:
     from homeassistant.helpers.entity_component import EntityComponent
 
-    from .default_agent import DefaultAgent
     from .entity import ConversationEntity
 
 DOMAIN = "conversation"
-DEFAULT_EXPOSED_ATTRIBUTES = {"device_class"}
 HOME_ASSISTANT_AGENT = "conversation.home_assistant"
 
 ATTR_TEXT = "text"
@@ -26,7 +24,6 @@
 SERVICE_RELOAD = "reload"
 
 DATA_COMPONENT: HassKey[EntityComponent[ConversationEntity]] = HassKey(DOMAIN)
-DATA_DEFAULT_ENTITY: HassKey[DefaultAgent] = HassKey(f"{DOMAIN}_default_entity")
 
 
 class ConversationEntityFeature(IntFlag):

homeassistant/components/conversation/default_agent.py

@@ -68,13 +68,9 @@
 from homeassistant.util import language as language_util
 from homeassistant.util.json import JsonObjectType, json_loads_object
 
+from .agent_manager import get_agent_manager
 from .chat_log import AssistantContent, ChatLog
-from .const import (
-    DATA_DEFAULT_ENTITY,
-    DEFAULT_EXPOSED_ATTRIBUTES,
-    DOMAIN,
-    ConversationEntityFeature,
-)
+from .const import DOMAIN, ConversationEntityFeature
 from .entity import ConversationEntity
 from .models import ConversationInput, ConversationResult
 from .trace import ConversationTraceEventType, async_conversation_trace_append
@@ -83,6 +79,8 @@
 _DEFAULT_ERROR_TEXT = "Sorry, I couldn't understand that"
 _ENTITY_REGISTRY_UPDATE_FIELDS = ["aliases", "name", "original_name"]
 
+_DEFAULT_EXPOSED_ATTRIBUTES = {"device_class"}
+
 REGEX_TYPE = type(re.compile(""))
 TRIGGER_CALLBACK_TYPE = Callable[
     [ConversationInput, RecognizeResult], Awaitable[str | None]
@@ -209,9 +207,9 @@ async def async_setup_default_agent(
     config_intents: dict[str, Any],
 ) -> None:
     """Set up entity registry listener for the default agent."""
-    entity = DefaultAgent(hass, config_intents)
-    await entity_component.async_add_entities([entity])
-    hass.data[DATA_DEFAULT_ENTITY] = entity
+    agent = DefaultAgent(hass, config_intents)
+    await entity_component.async_add_entities([agent])
+    await get_agent_manager(hass).async_setup_default_agent(agent)
 
     @core.callback
     def async_entity_state_listener(
@@ -846,7 +844,7 @@ def _get_entity_name_tuples(
             context = {"domain": state.domain}
             if state.attributes:
                 # Include some attributes
-                for attr in DEFAULT_EXPOSED_ATTRIBUTES:
+                for attr in _DEFAULT_EXPOSED_ATTRIBUTES:
                     if attr not in state.attributes:
                         continue
                     context[attr] = state.attributes[attr]

homeassistant/components/conversation/http.py

@@ -25,7 +25,7 @@
     async_get_agent,
     get_agent_manager,
 )
-from .const import DATA_COMPONENT, DATA_DEFAULT_ENTITY
+from .const import DATA_COMPONENT
 from .default_agent import (
     METADATA_CUSTOM_FILE,
     METADATA_CUSTOM_SENTENCE,
@@ -169,7 +169,8 @@ async def websocket_list_sentences(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
 ) -> None:
     """List custom registered sentences."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = get_agent_manager(hass).default_agent
+    assert agent is not None
 
     sentences = []
     for trigger_data in agent.trigger_sentences:
@@ -191,7 +192,8 @@ async def websocket_hass_agent_debug(
     hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
 ) -> None:
     """Return intents that would be matched by the default agent for a list of sentences."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = get_agent_manager(hass).default_agent
+    assert agent is not None
 
     # Return results for each sentence in the same order as the input.
     result_dicts: list[dict[str, Any] | None] = []

homeassistant/components/conversation/trigger.py

@@ -20,7 +20,8 @@
 from homeassistant.helpers.trigger import TriggerActionType, TriggerInfo
 from homeassistant.helpers.typing import UNDEFINED, ConfigType
 
-from .const import DATA_DEFAULT_ENTITY, DOMAIN
+from .agent_manager import get_agent_manager
+from .const import DOMAIN
 from .models import ConversationInput
 
 
@@ -123,4 +124,6 @@ async def call_action(
         # two trigger copies for who will provide a response.
         return None
 
-    return hass.data[DATA_DEFAULT_ENTITY].register_trigger(sentences, call_action)
+    agent = get_agent_manager(hass).default_agent
+    assert agent is not None
+    return agent.register_trigger(sentences, call_action)

tests/components/assist_pipeline/conftest.py

@@ -298,7 +298,7 @@ async def async_setup_entry_wake_word_platform(
     assert await async_setup_component(hass, "conversation", {"conversation": {}})
 
     # Disable fuzzy matching by default for tests
-    agent = hass.data[conversation.DATA_DEFAULT_ENTITY]
+    agent = conversation.async_get_agent(hass)
     agent.fuzzy_matching = False
 
     config_entry = MockConfigEntry(domain="test")

tests/components/conversation/conftest.py

@@ -6,6 +6,7 @@
 import pytest
 
 from homeassistant.components import conversation
+from homeassistant.components.conversation import async_get_agent, default_agent
 from homeassistant.components.shopping_list import intent as sl_intent
 from homeassistant.const import MATCH_ALL
 from homeassistant.core import Context, HomeAssistant
@@ -77,5 +78,6 @@ async def init_components(hass: HomeAssistant):
     assert await async_setup_component(hass, "conversation", {conversation.DOMAIN: {}})
 
     # Disable fuzzy matching by default for tests
-    agent = hass.data[conversation.DATA_DEFAULT_ENTITY]
+    agent = async_get_agent(hass)
+    assert isinstance(agent, default_agent.DefaultAgent)
     agent.fuzzy_matching = False

tests/components/conversation/test_default_agent.py

@@ -12,8 +12,7 @@
 import yaml
 
 from homeassistant.components import conversation, cover, media_player, weather
-from homeassistant.components.conversation import default_agent
-from homeassistant.components.conversation.const import DATA_DEFAULT_ENTITY
+from homeassistant.components.conversation import async_get_agent, default_agent
 from homeassistant.components.conversation.default_agent import METADATA_CUSTOM_SENTENCE
 from homeassistant.components.conversation.models import ConversationInput
 from homeassistant.components.cover import SERVICE_OPEN_COVER
@@ -87,7 +86,8 @@ async def init_components(hass: HomeAssistant) -> None:
     assert await async_setup_component(hass, "intent", {})
 
     # Disable fuzzy matching by default for tests
-    agent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = async_get_agent(hass)
+    assert isinstance(agent, default_agent.DefaultAgent)
     agent.fuzzy_matching = False
 
 
@@ -215,7 +215,7 @@ async def test_exposed_areas(
 @pytest.mark.usefixtures("init_components")
 async def test_conversation_agent(hass: HomeAssistant) -> None:
     """Test DefaultAgent."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = async_get_agent(hass)
     with patch(
         "homeassistant.components.conversation.default_agent.get_languages",
         return_value=["dwarvish", "elvish", "entish"],
@@ -415,8 +415,7 @@ async def test_trigger_sentences(hass: HomeAssistant) -> None:
     trigger_sentences = ["It's party time", "It is time to party"]
     trigger_response = "Cowabunga!"
 
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     callback = AsyncMock(return_value=trigger_response)
     unregister = agent.register_trigger(trigger_sentences, callback)
@@ -462,8 +461,7 @@ async def test_trigger_sentence_response_translation(
     """Test translation of default response 'done'."""
     hass.config.language = language
 
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     translations = {
         "en": {"component.conversation.conversation.agent.done": "English done"},
@@ -2525,8 +2523,7 @@ async def test_non_default_response(hass: HomeAssistant, init_components) -> Non
     hass.states.async_set("cover.front_door", "closed")
     calls = async_mock_service(hass, "cover", SERVICE_OPEN_COVER)
 
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     result = await agent.async_process(
         ConversationInput(
@@ -2872,8 +2869,7 @@ async def test_query_same_name_different_areas(
 @pytest.mark.usefixtures("init_components")
 async def test_intent_cache_exposed(hass: HomeAssistant) -> None:
     """Test that intent recognition results are cached for exposed entities."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     entity_id = "light.test_light"
     hass.states.async_set(entity_id, "off")
@@ -2912,8 +2908,7 @@ async def test_intent_cache_exposed(hass: HomeAssistant) -> None:
 @pytest.mark.usefixtures("init_components")
 async def test_intent_cache_all_entities(hass: HomeAssistant) -> None:
     """Test that intent recognition results are cached for all entities."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     entity_id = "light.test_light"
     hass.states.async_set(entity_id, "off")
@@ -2952,8 +2947,7 @@ async def test_intent_cache_all_entities(hass: HomeAssistant) -> None:
 @pytest.mark.usefixtures("init_components")
 async def test_intent_cache_fuzzy(hass: HomeAssistant) -> None:
     """Test that intent recognition results are cached for fuzzy matches."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     # There is no entity named test light
     user_input = ConversationInput(
@@ -2982,8 +2976,7 @@ async def test_intent_cache_fuzzy(hass: HomeAssistant) -> None:
 @pytest.mark.usefixtures("init_components")
 async def test_entities_filtered_by_input(hass: HomeAssistant) -> None:
     """Test that entities are filtered by the input text before intent matching."""
-    agent = hass.data[DATA_DEFAULT_ENTITY]
-    assert isinstance(agent, default_agent.DefaultAgent)
+    agent = async_get_agent(hass)
 
     # Only the switch is exposed
     hass.states.async_set("light.test_light", "off")
@@ -3165,7 +3158,7 @@ async def test_handle_intents_with_response_errors(
     assert await async_setup_component(hass, "climate", {})
     area_registry.async_create("living room")
 
-    agent: default_agent.DefaultAgent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = async_get_agent(hass)
 
     user_input = ConversationInput(
         text="What is the temperature in the living room?",
@@ -3203,7 +3196,7 @@ async def test_handle_intents_filters_results(
     assert await async_setup_component(hass, "climate", {})
     area_registry.async_create("living room")
 
-    agent: default_agent.DefaultAgent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = async_get_agent(hass)
 
     user_input = ConversationInput(
         text="What is the temperature in the living room?",
@@ -3363,7 +3356,7 @@ async def test_fuzzy_matching(
     assert await async_setup_component(hass, "intent", {})
     await light_intent.async_setup_intents(hass)
 
-    agent = hass.data[DATA_DEFAULT_ENTITY]
+    agent = async_get_agent(hass)
     agent.fuzzy_matching = fuzzy_matching
 
     area_office = area_registry.async_get_or_create("office_id")