Skip to content

Commit

Permalink
Add aliases to script llm tool description
Browse files Browse the repository at this point in the history
  • Loading branch information
Shulyaka committed Jul 22, 2024
1 parent 5612e3a commit faba33e
Show file tree
Hide file tree
Showing 2 changed files with 32 additions and 7 deletions.
9 changes: 9 additions & 0 deletions homeassistant/helpers/llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -677,6 +677,15 @@ def on_homeassistant_close(event: Event) -> None:

self.parameters = vol.Schema(schema)

aliases = entity_entry.aliases
if aliases:
if self.description:
self.description = (
self.description + ". Aliases: " + str(list(aliases))
)
else:
self.description = "Aliases: " + str(list(aliases))

parameters_cache[entity_entry.unique_id] = (
self.description,
self.parameters,
Expand Down
30 changes: 23 additions & 7 deletions tests/helpers/test_llm.py
Original file line number Diff line number Diff line change
Expand Up @@ -412,7 +412,9 @@ async def test_assist_api_prompt(
)
hass.states.async_set(entry2.entity_id, "on", {"friendly_name": "Living Room"})

def create_entity(device: dr.DeviceEntry, write_state=True) -> None:
def create_entity(
device: dr.DeviceEntry, write_state=True, aliases: set[str] | None = None
) -> None:
"""Create an entity for a device and track entity_id."""
entity = entity_registry.async_get_or_create(
"light",
Expand All @@ -422,6 +424,8 @@ def create_entity(device: dr.DeviceEntry, write_state=True) -> None:
original_name=str(device.name or "Unnamed Device"),
suggested_object_id=str(device.name or "unnamed_device"),
)
if aliases:
entity_registry.async_update_entity(entity.entity_id, aliases=aliases)
if write_state:
entity.write_unavailable_state(hass)

Expand All @@ -433,7 +437,8 @@ def create_entity(device: dr.DeviceEntry, write_state=True) -> None:
manufacturer="Test Manufacturer",
model="Test Model",
suggested_area="Test Area",
)
),
aliases={"my test light"},
)
for i in range(3):
create_entity(
Expand Down Expand Up @@ -526,7 +531,7 @@ def create_entity(device: dr.DeviceEntry, write_state=True) -> None:
},
"light.test_device": {
"areas": "Test Area, Alternative name",
"names": "Test Device",
"names": "Test Device, my test light",
"state": "unavailable",
},
"light.test_device_2": {
Expand Down Expand Up @@ -633,6 +638,7 @@ def create_entity(device: dr.DeviceEntry, write_state=True) -> None:

async def test_script_tool(
hass: HomeAssistant,
entity_registry: er.EntityRegistry,
area_registry: ar.AreaRegistry,
floor_registry: fr.FloorRegistry,
) -> None:
Expand Down Expand Up @@ -676,6 +682,10 @@ async def test_script_tool(
)
async_expose_entity(hass, "conversation", "script.test_script", True)

entity_registry.async_update_entity(
"script.test_script", aliases={"my test script"}
)

area = area_registry.async_create("Living room")
floor = floor_registry.async_create("2")

Expand All @@ -688,7 +698,7 @@ async def test_script_tool(

tool = tools[0]
assert tool.name == "test_script"
assert tool.description == "This is a test script"
assert tool.description == "This is a test script. Aliases: ['my test script']"
schema = {
vol.Required("beer", description="Number of beers"): cv.string,
vol.Optional("wine"): selector.NumberSelector({"min": 0, "max": 3}),
Expand All @@ -701,7 +711,10 @@ async def test_script_tool(
assert tool.parameters.schema == schema

assert hass.data[llm.SCRIPT_PARAMETERS_CACHE] == {
"test_script": ("This is a test script", vol.Schema(schema))
"test_script": (
"This is a test script. Aliases: ['my test script']",
vol.Schema(schema),
)
}

tool_input = llm.ToolInput(
Expand Down Expand Up @@ -771,12 +784,15 @@ async def test_script_tool(

tool = tools[0]
assert tool.name == "test_script"
assert tool.description == "This is a new test script"
assert tool.description == "This is a new test script. Aliases: ['my test script']"
schema = {vol.Required("beer", description="Number of beers"): cv.string}
assert tool.parameters.schema == schema

assert hass.data[llm.SCRIPT_PARAMETERS_CACHE] == {
"test_script": ("This is a new test script", vol.Schema(schema))
"test_script": (
"This is a new test script. Aliases: ['my test script']",
vol.Schema(schema),
)
}


Expand Down

0 comments on commit faba33e

Please sign in to comment.