LLM to handle decimal attributes (#120257)

This commit is contained in:
Paulus Schoutsen 2024-06-23 17:09:57 -04:00 committed by GitHub
parent 143e8d09af
commit 19f97a3e53
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 2 deletions

View File

@ -4,6 +4,7 @@ from __future__ import annotations
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from dataclasses import dataclass from dataclasses import dataclass
from decimal import Decimal
from enum import Enum from enum import Enum
from functools import cache, partial from functools import cache, partial
from typing import Any from typing import Any
@ -461,7 +462,9 @@ def _get_exposed_entities(
info["areas"] = ", ".join(area_names) info["areas"] = ", ".join(area_names)
if attributes := { if attributes := {
attr_name: str(attr_value) if isinstance(attr_value, Enum) else attr_value attr_name: str(attr_value)
if isinstance(attr_value, (Enum, Decimal))
else attr_value
for attr_name, attr_value in state.attributes.items() for attr_name, attr_value in state.attributes.items()
if attr_name in interesting_attributes if attr_name in interesting_attributes
}: }:

View File

@ -1,5 +1,6 @@
"""Tests for the llm helpers.""" """Tests for the llm helpers."""
from decimal import Decimal
from unittest.mock import patch from unittest.mock import patch
import pytest import pytest
@ -402,7 +403,11 @@ async def test_assist_api_prompt(
suggested_object_id="living_room", suggested_object_id="living_room",
device_id=device.id, device_id=device.id,
) )
hass.states.async_set(entry1.entity_id, "on", {"friendly_name": "Kitchen"}) hass.states.async_set(
entry1.entity_id,
"on",
{"friendly_name": "Kitchen", "temperature": Decimal("0.9")},
)
hass.states.async_set(entry2.entity_id, "on", {"friendly_name": "Living Room"}) 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) -> None:
@ -510,6 +515,9 @@ async def test_assist_api_prompt(
entry1.entity_id: { entry1.entity_id: {
"names": "Kitchen", "names": "Kitchen",
"state": "on", "state": "on",
"attributes": {
"temperature": "0.9",
},
}, },
entry2.entity_id: { entry2.entity_id: {
"areas": "Test Area, Alternative name", "areas": "Test Area, Alternative name",