Add AI Task support in Ollama (home-assistant#148226)

Co-authored-by: Paulus Schoutsen <[email protected]>

Author: allenporter

homeassistant/components/ollama/__init__.py

@@ -28,6 +28,7 @@
     CONF_NUM_CTX,
     CONF_PROMPT,
     CONF_THINK,
+    DEFAULT_AI_TASK_NAME,
     DEFAULT_NAME,
     DEFAULT_TIMEOUT,
     DOMAIN,
@@ -47,7 +48,7 @@
 ]
 
 CONFIG_SCHEMA = cv.config_entry_only_config_schema(DOMAIN)
-PLATFORMS = (Platform.CONVERSATION,)
+PLATFORMS = (Platform.AI_TASK, Platform.CONVERSATION)
 
 type OllamaConfigEntry = ConfigEntry[ollama.AsyncClient]
 
@@ -118,6 +119,7 @@ async def async_migrate_integration(hass: HomeAssistant) -> None:
         parent_entry = api_keys_entries[entry.data[CONF_URL]]
 
         hass.config_entries.async_add_subentry(parent_entry, subentry)
+
         conversation_entity = entity_registry.async_get_entity_id(
             "conversation",
             DOMAIN,
@@ -208,6 +210,31 @@ async def async_migrate_entry(hass: HomeAssistant, entry: OllamaConfigEntry) ->
             minor_version=1,
         )
 
+    if entry.version == 3 and entry.minor_version == 1:
+        # Add AI Task subentry with default options. We can only create a new
+        # subentry if we can find an existing model in the entry. The model
+        # was removed in the previous migration step, so we need to
+        # check the subentries for an existing model.
+        existing_model = next(
+            iter(
+                model
+                for subentry in entry.subentries.values()
+                if (model := subentry.data.get(CONF_MODEL)) is not None
+            ),
+            None,
+        )
+        if existing_model:
+            hass.config_entries.async_add_subentry(
+                entry,
+                ConfigSubentry(
+                    data=MappingProxyType({CONF_MODEL: existing_model}),
+                    subentry_type="ai_task_data",
+                    title=DEFAULT_AI_TASK_NAME,
+                    unique_id=None,
+                ),
+            )
+        hass.config_entries.async_update_entry(entry, minor_version=2)
+
     _LOGGER.debug(
         "Migration to version %s:%s successful", entry.version, entry.minor_version
     )

homeassistant/components/ollama/ai_task.py

@@ -0,0 +1,77 @@
+"""AI Task integration for Ollama."""
+
+from __future__ import annotations
+
+from json import JSONDecodeError
+import logging
+
+from homeassistant.components import ai_task, conversation
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
+from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
+from homeassistant.util.json import json_loads
+
+from .entity import OllamaBaseLLMEntity
+
+_LOGGER = logging.getLogger(__name__)
+
+
+async def async_setup_entry(
+    hass: HomeAssistant,
+    config_entry: ConfigEntry,
+    async_add_entities: AddConfigEntryEntitiesCallback,
+) -> None:
+    """Set up AI Task entities."""
+    for subentry in config_entry.subentries.values():
+        if subentry.subentry_type != "ai_task_data":
+            continue
+
+        async_add_entities(
+            [OllamaTaskEntity(config_entry, subentry)],
+            config_subentry_id=subentry.subentry_id,
+        )
+
+
+class OllamaTaskEntity(
+    ai_task.AITaskEntity,
+    OllamaBaseLLMEntity,
+):
+    """Ollama AI Task entity."""
+
+    _attr_supported_features = ai_task.AITaskEntityFeature.GENERATE_DATA
+
+    async def _async_generate_data(
+        self,
+        task: ai_task.GenDataTask,
+        chat_log: conversation.ChatLog,
+    ) -> ai_task.GenDataTaskResult:
+        """Handle a generate data task."""
+        await self._async_handle_chat_log(chat_log, task.structure)
+
+        if not isinstance(chat_log.content[-1], conversation.AssistantContent):
+            raise HomeAssistantError(
+                "Last content in chat log is not an AssistantContent"
+            )
+
+        text = chat_log.content[-1].content or ""
+
+        if not task.structure:
+            return ai_task.GenDataTaskResult(
+                conversation_id=chat_log.conversation_id,
+                data=text,
+            )
+        try:
+            data = json_loads(text)
+        except JSONDecodeError as err:
+            _LOGGER.error(
+                "Failed to parse JSON response: %s. Response: %s",
+                err,
+                text,
+            )
+            raise HomeAssistantError("Error with Ollama structured response") from err
+
+        return ai_task.GenDataTaskResult(
+            conversation_id=chat_log.conversation_id,
+            data=data,
+        )

homeassistant/components/ollama/config_flow.py

@@ -46,6 +46,8 @@
     CONF_NUM_CTX,
     CONF_PROMPT,
     CONF_THINK,
+    DEFAULT_AI_TASK_NAME,
+    DEFAULT_CONVERSATION_NAME,
     DEFAULT_KEEP_ALIVE,
     DEFAULT_MAX_HISTORY,
     DEFAULT_MODEL,
@@ -74,7 +76,7 @@ class OllamaConfigFlow(ConfigFlow, domain=DOMAIN):
     """Handle a config flow for Ollama."""
 
     VERSION = 3
-    MINOR_VERSION = 1
+    MINOR_VERSION = 2
 
     def __init__(self) -> None:
         """Initialize config flow."""
@@ -136,11 +138,14 @@ def async_get_supported_subentry_types(
         cls, config_entry: ConfigEntry
     ) -> dict[str, type[ConfigSubentryFlow]]:
         """Return subentries supported by this integration."""
-        return {"conversation": ConversationSubentryFlowHandler}
+        return {
+            "conversation": OllamaSubentryFlowHandler,
+            "ai_task_data": OllamaSubentryFlowHandler,
+        }
 
 
-class ConversationSubentryFlowHandler(ConfigSubentryFlow):
-    """Flow for managing conversation subentries."""
+class OllamaSubentryFlowHandler(ConfigSubentryFlow):
+    """Flow for managing Ollama subentries."""
 
     def __init__(self) -> None:
         """Initialize the subentry flow."""
@@ -201,7 +206,11 @@ async def async_step_set_options(
                 step_id="set_options",
                 data_schema=vol.Schema(
                     ollama_config_option_schema(
-                        self.hass, self._is_new, options, models_to_list
+                        self.hass,
+                        self._is_new,
+                        self._subentry_type,
+                        options,
+                        models_to_list,
                     )
                 ),
             )
@@ -300,13 +309,19 @@ async def async_step_finish(
 def ollama_config_option_schema(
     hass: HomeAssistant,
     is_new: bool,
+    subentry_type: str,
     options: Mapping[str, Any],
     models_to_list: list[SelectOptionDict],
 ) -> dict:
     """Ollama options schema."""
     if is_new:
+        if subentry_type == "ai_task_data":
+            default_name = DEFAULT_AI_TASK_NAME
+        else:
+            default_name = DEFAULT_CONVERSATION_NAME
+
         schema: dict = {
-            vol.Required(CONF_NAME, default="Ollama Conversation"): str,
+            vol.Required(CONF_NAME, default=default_name): str,
         }
     else:
         schema = {}
@@ -319,29 +334,38 @@ def ollama_config_option_schema(
             ): SelectSelector(
                 SelectSelectorConfig(options=models_to_list, custom_value=True)
             ),
-            vol.Optional(
-                CONF_PROMPT,
-                description={
-                    "suggested_value": options.get(
-                        CONF_PROMPT, llm.DEFAULT_INSTRUCTIONS_PROMPT
-                    )
-                },
-            ): TemplateSelector(),
-            vol.Optional(
-                CONF_LLM_HASS_API,
-                description={"suggested_value": options.get(CONF_LLM_HASS_API)},
-            ): SelectSelector(
-                SelectSelectorConfig(
-                    options=[
-                        SelectOptionDict(
-                            label=api.name,
-                            value=api.id,
+        }
+    )
+    if subentry_type == "conversation":
+        schema.update(
+            {
+                vol.Optional(
+                    CONF_PROMPT,
+                    description={
+                        "suggested_value": options.get(
+                            CONF_PROMPT, llm.DEFAULT_INSTRUCTIONS_PROMPT
                         )
-                        for api in llm.async_get_apis(hass)
-                    ],
-                    multiple=True,
-                )
-            ),
+                    },
+                ): TemplateSelector(),
+                vol.Optional(
+                    CONF_LLM_HASS_API,
+                    description={"suggested_value": options.get(CONF_LLM_HASS_API)},
+                ): SelectSelector(
+                    SelectSelectorConfig(
+                        options=[
+                            SelectOptionDict(
+                                label=api.name,
+                                value=api.id,
+                            )
+                            for api in llm.async_get_apis(hass)
+                        ],
+                        multiple=True,
+                    )
+                ),
+            }
+        )
+    schema.update(
+        {
             vol.Optional(
                 CONF_NUM_CTX,
                 description={

homeassistant/components/ollama/const.py

@@ -159,3 +159,10 @@
     "zephyr",
 ]
 DEFAULT_MODEL = "llama3.2:latest"
+
+DEFAULT_CONVERSATION_NAME = "Ollama Conversation"
+DEFAULT_AI_TASK_NAME = "Ollama AI Task"
+
+RECOMMENDED_CONVERSATION_OPTIONS = {
+    CONF_MAX_HISTORY: DEFAULT_MAX_HISTORY,
+}

homeassistant/components/ollama/entity.py

@@ -8,6 +8,7 @@
 from typing import Any
 
 import ollama
+import voluptuous as vol
 from voluptuous_openapi import convert
 
 from homeassistant.components import conversation
@@ -180,6 +181,7 @@ def __init__(self, entry: OllamaConfigEntry, subentry: ConfigSubentry) -> None:
     async def _async_handle_chat_log(
         self,
         chat_log: conversation.ChatLog,
+        structure: vol.Schema | None = None,
     ) -> None:
         """Generate an answer for the chat log."""
         settings = {**self.entry.data, **self.subentry.data}
@@ -200,6 +202,17 @@ async def _async_handle_chat_log(
         max_messages = int(settings.get(CONF_MAX_HISTORY, DEFAULT_MAX_HISTORY))
         self._trim_history(message_history, max_messages)
 
+        output_format: dict[str, Any] | None = None
+        if structure:
+            output_format = convert(
+                structure,
+                custom_serializer=(
+                    chat_log.llm_api.custom_serializer
+                    if chat_log.llm_api
+                    else llm.selector_serializer
+                ),
+            )
+
         # Get response
         # To prevent infinite loops, we limit the number of iterations
         for _iteration in range(MAX_TOOL_ITERATIONS):
@@ -214,6 +227,7 @@ async def _async_handle_chat_log(
                     keep_alive=f"{settings.get(CONF_KEEP_ALIVE, DEFAULT_KEEP_ALIVE)}s",
                     options={CONF_NUM_CTX: settings.get(CONF_NUM_CTX, DEFAULT_NUM_CTX)},
                     think=settings.get(CONF_THINK),
+                    format=output_format,
                 )
             except (ollama.RequestError, ollama.ResponseError) as err:
                 _LOGGER.error("Unexpected error talking to Ollama server: %s", err)

homeassistant/components/ollama/strings.json

@@ -55,6 +55,44 @@
       "progress": {
         "download": "Please wait while the model is downloaded, which may take a very long time. Check your Ollama server logs for more details."
       }
+    },
+    "ai_task_data": {
+      "initiate_flow": {
+        "user": "Add Generate data with AI service",
+        "reconfigure": "Reconfigure Generate data with AI service"
+      },
+      "entry_type": "Generate data with AI service",
+      "step": {
+        "set_options": {
+          "data": {
+            "model": "[%key:component::ollama::config_subentries::conversation::step::set_options::data::model%]",
+            "name": "[%key:common::config_flow::data::name%]",
+            "prompt": "[%key:component::ollama::config_subentries::conversation::step::set_options::data::prompt%]",
+            "max_history": "[%key:component::ollama::config_subentries::conversation::step::set_options::data::max_history%]",
+            "num_ctx": "[%key:component::ollama::config_subentries::conversation::step::set_options::data::num_ctx%]",
+            "keep_alive": "[%key:component::ollama::config_subentries::conversation::step::set_options::data::keep_alive%]",
+            "think": "[%key:component::ollama::config_subentries::conversation::step::set_options::data::think%]"
+          },
+          "data_description": {
+            "prompt": "[%key:component::ollama::config_subentries::conversation::step::set_options::data_description::prompt%]",
+            "keep_alive": "[%key:component::ollama::config_subentries::conversation::step::set_options::data_description::keep_alive%]",
+            "num_ctx": "[%key:component::ollama::config_subentries::conversation::step::set_options::data_description::num_ctx%]",
+            "think": "[%key:component::ollama::config_subentries::conversation::step::set_options::data_description::think%]"
+          }
+        },
+        "download": {
+          "title": "[%key:component::ollama::config_subentries::conversation::step::download::title%]"
+        }
+      },
+      "abort": {
+        "reconfigure_successful": "[%key:common::config_flow::abort::reconfigure_successful%]",
+        "entry_not_loaded": "[%key:component::ollama::config_subentries::conversation::abort::entry_not_loaded%]",
+        "download_failed": "[%key:component::ollama::config_subentries::conversation::abort::download_failed%]",
+        "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]"
+      },
+      "progress": {
+        "download": "[%key:component::ollama::config_subentries::conversation::progress::download%]"
+      }
     }
   }
 }

tests/components/ollama/__init__.py

@@ -12,3 +12,8 @@
     ollama.CONF_MAX_HISTORY: 2,
     ollama.CONF_MODEL: "test_model:latest",
 }
+
+TEST_AI_TASK_OPTIONS = {
+    ollama.CONF_MAX_HISTORY: 2,
+    ollama.CONF_MODEL: "test_model:latest",
+}