add scrape function

custom_components/extended_openai_conversation/__init__.py

@@ -61,6 +61,7 @@
     ScriptFunctionExecutor,
     TemplateFunctionExecutor,
     RestFunctionExecutor,
+    ScrapeFunctionExecutor,
     convert_to_template,
 )
 
@@ -76,6 +77,7 @@
     "script": ScriptFunctionExecutor(),
     "template": TemplateFunctionExecutor(),
     "rest": RestFunctionExecutor(),
+    "scrape": ScrapeFunctionExecutor(),
 }
 
 # hass.data key for agent.
@@ -240,7 +242,7 @@ def get_functions(self):
             if result:
                 for setting in result:
                     for function in setting["function"].values():
-                        convert_to_template(function)
+                        convert_to_template(function, hass=self.hass)
             return result
         except:
             _LOGGER.error("Failed to load functions")

custom_components/extended_openai_conversation/helpers.py

@@ -3,9 +3,10 @@
 import os
 import yaml
 import time
+from bs4 import BeautifulSoup
+from typing import Any
 
-
-from homeassistant.components import automation, rest
+from homeassistant.components import automation, rest, scrape
 from homeassistant.components.automation.config import _async_validate_config_item
 from homeassistant.const import (
     SERVICE_RELOAD,
@@ -15,11 +16,13 @@
     CONF_VALUE_TEMPLATE,
     CONF_RESOURCE,
     CONF_RESOURCE_TEMPLATE,
+    CONF_NAME,
+    CONF_ATTRIBUTE,
 )
 from homeassistant.config import AUTOMATION_CONFIG_PATH
 from homeassistant.components import conversation
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import template
+from homeassistant.helpers.template import Template
 from homeassistant.helpers.script import (
     Script,
     SCRIPT_MODE_SINGLE,
@@ -44,25 +47,54 @@
 
 
 def convert_to_template(
-    settings, template_keys=["data", "event_data", "target", "service"]
+    settings,
+    template_keys=["data", "event_data", "target", "service"],
+    hass: HomeAssistant | None = None,
 ):
-    _convert_to_template(settings, template_keys, [])
+    _convert_to_template(settings, template_keys, hass, [])
 
 
-def _convert_to_template(settings, template_keys, parents: list[str]):
+def _convert_to_template(settings, template_keys, hass, parents: list[str]):
     if isinstance(settings, dict):
         for key, value in settings.items():
             if isinstance(value, str) and (
                 key in template_keys or set(parents).intersection(template_keys)
             ):
-                settings[key] = template.Template(value)
+                settings[key] = Template(value, hass)
             if isinstance(value, dict):
                 parents.append(key)
-                _convert_to_template(value, template_keys, parents)
+                _convert_to_template(value, template_keys, hass, parents)
+                parents.pop()
+            if isinstance(value, list):
+                parents.append(key)
+                for item in value:
+                    _convert_to_template(item, template_keys, hass, parents)
                 parents.pop()
     if isinstance(settings, list):
         for setting in settings:
-            _convert_to_template(setting, template_keys, parents)
+            _convert_to_template(setting, template_keys, hass, parents)
+
+
+def _get_rest_data(hass, rest_config, arguments):
+    rest_config.setdefault(CONF_METHOD, rest.const.DEFAULT_METHOD)
+    rest_config.setdefault(CONF_VERIFY_SSL, rest.const.DEFAULT_VERIFY_SSL)
+    rest_config.setdefault(CONF_TIMEOUT, rest.data.DEFAULT_TIMEOUT)
+    rest_config.setdefault(rest.const.CONF_ENCODING, rest.const.DEFAULT_ENCODING)
+
+    convert_to_template(
+        rest_config,
+        template_keys=[CONF_VALUE_TEMPLATE, CONF_RESOURCE_TEMPLATE],
+        hass=hass,
+    )
+
+    resource_template: Template | None = rest_config.get(CONF_RESOURCE_TEMPLATE)
+    if resource_template is not None:
+        rest_config.pop(CONF_RESOURCE_TEMPLATE)
+        rest_config[CONF_RESOURCE] = resource_template.async_render(
+            arguments, parse_result=False
+        )
+
+    return rest.create_rest_data_from_config(hass, rest_config)
 
 
 class FunctionExecutor(ABC):
@@ -229,9 +261,7 @@ async def execute(
         user_input: conversation.ConversationInput,
         exposed_entities,
     ) -> str:
-        return template.Template(
-            function["function"]["value_template"], hass
-        ).async_render(
+        return Template(function["function"]["value_template"], hass).async_render(
             arguments,
             parse_result=False,
         )
@@ -250,34 +280,99 @@ async def execute(
         exposed_entities,
     ) -> str:
         config = function["function"]
-        config.setdefault(CONF_METHOD, rest.const.DEFAULT_METHOD)
-        config.setdefault(CONF_VERIFY_SSL, rest.const.DEFAULT_VERIFY_SSL)
-        config.setdefault(CONF_TIMEOUT, rest.data.DEFAULT_TIMEOUT)
-        config.setdefault(rest.const.CONF_ENCODING, rest.const.DEFAULT_ENCODING)
+        rest_data = _get_rest_data(hass, config, arguments)
 
-        convert_to_template(
-            config, template_keys=[CONF_VALUE_TEMPLATE, CONF_RESOURCE_TEMPLATE]
+        await rest_data.async_update()
+        value = rest_data.data_without_xml()
+        value_template = config.get(CONF_VALUE_TEMPLATE)
+
+        if value is not None and value_template is not None:
+            value = value_template.async_render_with_possible_json_value(
+                value, None, arguments
+            )
+
+        return value
+
+
+class ScrapeFunctionExecutor(FunctionExecutor):
+    def __init__(self) -> None:
+        """initialize Scrape function"""
+
+    async def execute(
+        self,
+        hass: HomeAssistant,
+        function,
+        arguments,
+        user_input: conversation.ConversationInput,
+        exposed_entities,
+    ) -> str:
+        config = function["function"]
+        rest_data = _get_rest_data(hass, config, arguments)
+        coordinator = scrape.coordinator.ScrapeCoordinator(
+            hass,
+            rest_data,
+            scrape.const.DEFAULT_SCAN_INTERVAL,
         )
+        await coordinator.async_config_entry_first_refresh()
+
+        new_arguments = dict(arguments)
+
+        for sensor_config in config["sensor"]:
+            name: str = sensor_config.get(CONF_NAME)
+            value = self._async_update_from_rest_data(
+                coordinator.data, sensor_config, arguments
+            )
+            new_arguments["value"] = value
+            if name:
+                new_arguments[name] = value
+
+        result = new_arguments["value"]
         value_template = config.get(CONF_VALUE_TEMPLATE)
+
         if value_template is not None:
-            value_template.hass = hass
-
-        resource_template: template.Template | None = config.get(CONF_RESOURCE_TEMPLATE)
-        if resource_template is not None:
-            config.pop(CONF_RESOURCE_TEMPLATE)
-            resource_template.hass = hass
-            config[CONF_RESOURCE] = resource_template.async_render(
-                arguments, parse_result=False
+            result = value_template.async_render_with_possible_json_value(
+                result, None, new_arguments
             )
 
-        rest_data = rest.create_rest_data_from_config(hass, config)
+        return result
 
-        await rest_data.async_update()
-        value = rest_data.data_without_xml()
+    def _async_update_from_rest_data(
+        self,
+        data: BeautifulSoup,
+        sensor_config: dict[str, Any],
+        arguments: dict[str, Any],
+    ) -> None:
+        """Update state from the rest data."""
+        value = self._extract_value(data, sensor_config)
+        value_template = sensor_config.get(CONF_VALUE_TEMPLATE)
 
-        if value is not None and value_template is not None:
+        if value_template is not None:
             value = value_template.async_render_with_possible_json_value(
                 value, None, arguments
             )
 
         return value
+
+    def _extract_value(self, data: BeautifulSoup, sensor_config: dict[str, Any]) -> Any:
+        """Parse the html extraction in the executor."""
+        value: str | list[str] | None
+        select = sensor_config[scrape.const.CONF_SELECT]
+        index = sensor_config.get(scrape.const.CONF_INDEX, 0)
+        attr = sensor_config.get(CONF_ATTRIBUTE)
+        try:
+            if attr is not None:
+                value = data.select(select)[index][attr]
+            else:
+                tag = data.select(select)[index]
+                if tag.name in ("style", "script", "template"):
+                    value = tag.string
+                else:
+                    value = tag.text
+        except IndexError:
+            _LOGGER.warning("Index '%s' not found", index)
+            value = None
+        except KeyError:
+            _LOGGER.warning("Attribute '%s' not found", attr)
+            value = None
+        _LOGGER.debug("Parsed value: %s", value)
+        return value