IoTVirtualAssistant/tools/ha.py

import os
from homeassistant_api import Client
from langchain.tools import BaseTool
from homeassistant_api.errors import EndpointNotFoundError
ha_client = Client(
    os.environ.get('HA_URL', 'http://localhost:8123'),
    os.environ.get('HA_API_KEY', '')
)


class HALightControl(BaseTool):
    name = "Control the lights in the room"
    description  = (
        "use this tool when you need to turn on or off the light in the room. "
        "given the light entity name and an action like turnon or turnoff, this tool will turn the lights on or off. "
        "To use the tool you must provide exactly two of the following parameters "
        "['entity_name', 'action']"
    )
    return_direct = False

    def _run(
        self,
        action: str,
        entity_name: str,
    ):
        light = ha_client.get_domain("light")
        if entity_name:
            entity = "light.{entity_name}".format(entity_name=entity_name)
            if action == "turnon":
                light.turn_on(entity_id=entity)
                return "{entity_name} turned on".format(entity_name=entity_name)
            elif action == "turnoff":
                light.turn_off(entity_id=entity)
                return "{entity_name} turned off".format(entity_name=entity_name)
        return ""
    
    def _arun(self, query: str):
        raise NotImplementedError("This tool does not support async")


class HASensorReading(BaseTool):
    name = "Home Assistant Sensor Reading Tool"
    description  = (
        "use this tool when you need to read data from a sensor."
        "To use the tool you must provide exactly the following parameter "
        "'entity_name'"
    )
    return_direct = False

    def _get_entity_name(self, entity_name: str):
        # is entity_name starts with sensor. then return it as is
        if entity_name.startswith("sensor."):
            return entity_name
        return f"sensor.{entity_name}"
    
    def _run(
        self,
        entity_name: str,
    ):
        entity = self._get_entity_name(entity_name)
        try:
            sensor = ha_client.get_entity(entity_id=entity)
            state = sensor.get_state()
        except EndpointNotFoundError as e:
            return "No sensor found with name {entity_name}".format(entity_name=entity_name)
        except Exception as e:
            return "An error occurred while trying to get the sensor {entity_name}".format(entity_name=entity_name)

        return state
        
    
    def _arun(self, query: str):
        raise NotImplementedError("This tool does not support async")
    

class HATempHumReading(HASensorReading):
    name = "Room temperature and humidity readings"
    description  = (
        "use this tool when you need to get room temperature and humidity readings. "
        "To use the tool you must provide exactly the following parameter "
        "'entity_name'"
    )
    

class HAGeolocation(HASensorReading):
    name = "Home Assistant Geolocation Tool"
    description  = (
        "use this tool when you need to get geolocation of an entity, person, object. "
        "To use the tool you must provide exactly the following parameter "
        "'entity_name'"
    )

    def _get_entity_name(self, entity_name: str):
        return "sensor.{entity_name}_geocoded_location".format(entity_name=entity_name)