Cleanup OpenAI and Ollama conversation entities (#116714)

This commit is contained in:
Allen Porter 2024-05-03 07:59:08 -07:00 committed by GitHub
parent 1eed5442e2
commit 9b4099950c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 13 deletions

View File

@ -45,7 +45,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up conversation entities."""
agent = OllamaConversationEntity(hass, config_entry)
agent = OllamaConversationEntity(config_entry)
async_add_entities([agent])
@ -56,9 +56,8 @@ class OllamaConversationEntity(
_attr_has_entity_name = True
def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
def __init__(self, entry: ConfigEntry) -> None:
"""Initialize the agent."""
self.hass = hass
self.entry = entry
# conversation id -> message history
@ -223,21 +222,21 @@ class OllamaConversationEntity(
]
for state in exposed_states:
entity = entity_registry.async_get(state.entity_id)
entity_entry = entity_registry.async_get(state.entity_id)
names = [state.name]
area_names = []
if entity is not None:
if entity_entry is not None:
# Add aliases
names.extend(entity.aliases)
if entity.area_id and (
area := area_registry.async_get_area(entity.area_id)
names.extend(entity_entry.aliases)
if entity_entry.area_id and (
area := area_registry.async_get_area(entity_entry.area_id)
):
# Entity is in area
area_names.append(area.name)
area_names.extend(area.aliases)
elif entity.device_id and (
device := device_registry.async_get(entity.device_id)
elif entity_entry.device_id and (
device := device_registry.async_get(entity_entry.device_id)
):
# Check device area
if device.area_id and (

View File

@ -35,7 +35,7 @@ async def async_setup_entry(
async_add_entities: AddEntitiesCallback,
) -> None:
"""Set up conversation entities."""
agent = OpenAIConversationEntity(hass, config_entry)
agent = OpenAIConversationEntity(config_entry)
async_add_entities([agent])
@ -46,9 +46,8 @@ class OpenAIConversationEntity(
_attr_has_entity_name = True
def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
def __init__(self, entry: ConfigEntry) -> None:
"""Initialize the agent."""
self.hass = hass
self.entry = entry
self.history: dict[str, list[dict]] = {}
self._attr_name = entry.title