From 6c55822861306a0ec48837291ea308053aa6897b Mon Sep 17 00:00:00 2001
From: Black Roland <mail@roland.black>
Date: Thu, 12 Dec 2024 00:37:11 +0300
Subject: [PATCH] Use async SDK by default (AsyncYCloudML)

Throught the testing I didn't notice any performance issues. It seems AsyncYCloudML is good to use it as the default option.
---
 custom_components/yandexgpt_conversation/__init__.py |  6 ++----
 .../yandexgpt_conversation/config_flow.py            |  4 ----
 custom_components/yandexgpt_conversation/const.py    |  1 -
 .../yandexgpt_conversation/conversation.py           | 12 +++---------
 4 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/custom_components/yandexgpt_conversation/__init__.py b/custom_components/yandexgpt_conversation/__init__.py
index 7a1a778..6657093 100644
--- a/custom_components/yandexgpt_conversation/__init__.py
+++ b/custom_components/yandexgpt_conversation/__init__.py
@@ -97,10 +97,8 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up YandexGPT from a config entry."""
     settings = {**entry.data, **entry.options}
 
-    sdk_conf = {"folder_id": settings[CONF_FOLDER_ID], "auth": settings[CONF_API_KEY]}
-    entry.runtime_data = sdk_factory(
-        YCloudML(**sdk_conf),
-        AsyncYCloudML(**sdk_conf),
+    entry.runtime_data = AsyncYCloudML(
+        folder_id=settings[CONF_FOLDER_ID], auth=settings[CONF_API_KEY]
     )
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
diff --git a/custom_components/yandexgpt_conversation/config_flow.py b/custom_components/yandexgpt_conversation/config_flow.py
index 73317a7..eb29b1f 100644
--- a/custom_components/yandexgpt_conversation/config_flow.py
+++ b/custom_components/yandexgpt_conversation/config_flow.py
@@ -30,7 +30,6 @@
 )
 
 from .const import (
-    CONF_ASYNC_MODE,
     CONF_CHAT_MODEL,
     CONF_FOLDER_ID,
     CONF_MAX_TOKENS,
@@ -189,9 +188,6 @@ def yandexgpt_config_option_schema(
             ): SelectSelector(
                 SelectSelectorConfig(mode=SelectSelectorMode.DROPDOWN, options=models)
             ),
-            vol.Required(
-                CONF_ASYNC_MODE, default=options.get(CONF_ASYNC_MODE, False)
-            ): bool,
             vol.Optional(
                 CONF_TEMPERATURE,
                 description={"suggested_value": options.get(CONF_TEMPERATURE)},
diff --git a/custom_components/yandexgpt_conversation/const.py b/custom_components/yandexgpt_conversation/const.py
index cc879b3..32de3fd 100644
--- a/custom_components/yandexgpt_conversation/const.py
+++ b/custom_components/yandexgpt_conversation/const.py
@@ -15,7 +15,6 @@
 CONF_MAX_TOKENS = "max_tokens"
 CONF_TEMPERATURE = "temperature"
 CONF_CHAT_MODEL = "chat_model"
-CONF_ASYNC_MODE = "async_mode"
 
 RECOMMENDED_CHAT_MODEL = "yandexgpt-lite/latest"
 RECOMMENDED_MAX_TOKENS = 1024
diff --git a/custom_components/yandexgpt_conversation/conversation.py b/custom_components/yandexgpt_conversation/conversation.py
index 714f1cf..f1331f2 100644
--- a/custom_components/yandexgpt_conversation/conversation.py
+++ b/custom_components/yandexgpt_conversation/conversation.py
@@ -18,12 +18,11 @@
 from homeassistant.helpers import intent, llm, template
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import ulid
-from yandex_cloud_ml_sdk import AsyncYCloudML, YCloudML
+from yandex_cloud_ml_sdk import AsyncYCloudML
 from yandex_cloud_ml_sdk._models.completions.message import TextMessage
 
 from .const import (
     BASE_PROMPT_RU,
-    CONF_ASYNC_MODE,
     CONF_CHAT_MODEL,
     CONF_MAX_TOKENS,
     CONF_PROMPT,
@@ -187,8 +186,7 @@ async def async_process(
             {"messages": messages},
         )
 
-        is_async = options.get(CONF_ASYNC_MODE, False)
-        client: YCloudML | AsyncYCloudML = self.entry.runtime_data(is_async)
+        client: AsyncYCloudML = self.entry.runtime_data
         # TODO: Allow selecting model version
         model_name_ver = zip(
             ["model_name", "model_version"],
@@ -201,11 +199,7 @@ async def async_process(
 
         try:
             model = client.models.completions(**dict(model_name_ver))
-            model.configure(**model_conf)
-            if options.get(CONF_ASYNC_MODE, False):
-                result = await model.run(messages)
-            else:
-                result = await self.hass.async_add_executor_job(model.run, messages)
+            result = await model.configure(**model_conf).run(messages)
         except Exception as err:
             LOGGER.exception(err)