Add bearer token as optional setting to Ollama

homeassistant/components/ollama/__init__.py

@@ -10,7 +10,7 @@
 import ollama
 
 from homeassistant.config_entries import ConfigEntry, ConfigSubentry
-from homeassistant.const import CONF_URL, Platform
+from homeassistant.const import CONF_API_KEY, CONF_URL, Platform
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 from homeassistant.helpers import (
@@ -62,7 +62,12 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
 async def async_setup_entry(hass: HomeAssistant, entry: OllamaConfigEntry) -> bool:
     """Set up Ollama from a config entry."""
     settings = {**entry.data, **entry.options}
-    client = ollama.AsyncClient(host=settings[CONF_URL], verify=get_default_context())
+    api_key = settings.get(CONF_API_KEY) or ""
+    client = ollama.AsyncClient(
+        host=settings[CONF_URL],
+        headers={"Authorization": f"Bearer {api_key}"} if api_key else None,
+        verify=get_default_context(),
+    )
     try:
         async with asyncio.timeout(DEFAULT_TIMEOUT):
             await client.list()

homeassistant/components/ollama/config_flow.py

@@ -20,7 +20,7 @@
     ConfigSubentryFlow,
     SubentryFlowResult,
 )
-from homeassistant.const import CONF_LLM_HASS_API, CONF_NAME, CONF_URL
+from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API, CONF_NAME, CONF_URL
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import config_validation as cv, llm
 from homeassistant.helpers.selector import (
@@ -68,7 +68,10 @@
         vol.Required(CONF_URL): TextSelector(
             TextSelectorConfig(type=TextSelectorType.URL)
         ),
-    }
+        vol.Optional(CONF_API_KEY, default=""): TextSelector(
+            TextSelectorConfig(type=TextSelectorType.PASSWORD)
+        ),
+    },
 )
 
 
@@ -93,6 +96,7 @@ async def async_step_user(
 
         errors = {}
         url = user_input[CONF_URL]
+        api_key = user_input.get(CONF_API_KEY) or ""
 
         self._async_abort_entries_match({CONF_URL: url})
 
@@ -109,7 +113,12 @@ async def async_step_user(
             )
 
         try:
-            client = ollama.AsyncClient(host=url, verify=get_default_context())
+            client = ollama.AsyncClient(
+                host=url,
+                headers={"Authorization": f"Bearer {api_key}"} if api_key else None,
+                verify=get_default_context(),
+            )
+
             async with asyncio.timeout(DEFAULT_TIMEOUT):
                 await client.list()
         except TimeoutError, httpx.ConnectError:
@@ -129,7 +138,7 @@ async def async_step_user(
 
         return self.async_create_entry(
             title=url,
-            data={CONF_URL: url},
+            data={CONF_URL: url, CONF_API_KEY: api_key},
         )
 
     @classmethod

homeassistant/components/ollama/strings.json

@@ -11,6 +11,7 @@
     "step": {
       "user": {
         "data": {
+          "api_key": "[%key:common::config_flow::data::api_key%]",
           "url": "[%key:common::config_flow::data::url%]"
         }
       }

tests/components/ollama/test_config_flow.py

@@ -1,14 +1,15 @@
 """Test the Ollama config flow."""
 
 import asyncio
-from unittest.mock import patch
+from unittest.mock import ANY, AsyncMock, patch
 
 from httpx import ConnectError
 import pytest
 
 from homeassistant import config_entries
 from homeassistant.components import ollama
-from homeassistant.const import CONF_LLM_HASS_API, CONF_NAME
+from homeassistant.components.ollama.const import DOMAIN
+from homeassistant.const import CONF_API_KEY, CONF_LLM_HASS_API, CONF_NAME, CONF_URL
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 
@@ -50,6 +51,7 @@ async def test_form(hass: HomeAssistant) -> None:
     assert result2["type"] is FlowResultType.CREATE_ENTRY
     assert result2["data"] == {
         ollama.CONF_URL: "http://localhost:11434",
+        ollama.CONF_API_KEY: "",  # Default API key should be empty string
     }
     # No subentries created by default
     assert len(result2.get("subentries", [])) == 0
@@ -557,3 +559,50 @@ async def test_ai_task_subentry_not_loaded(
 
     assert result.get("type") is FlowResultType.ABORT
     assert result.get("reason") == "entry_not_loaded"
+
+
+@pytest.mark.parametrize(
+    ("user_input", "expected_headers"),
+    [
+        (
+            {CONF_URL: "http://localhost:11434", CONF_API_KEY: "my-secret-token"},
+            {"Authorization": "Bearer my-secret-token"},
+        ),
+        (
+            {CONF_URL: "http://localhost:11434", CONF_API_KEY: ""},
+            None,
+        ),
+        (
+            {CONF_URL: "http://localhost:11434"},
+            None,
+        ),
+    ],
+)
+async def test_user_step_async_client_headers(
+    hass: HomeAssistant,
+    user_input: dict[str, str],
+    expected_headers: dict[str, str] | None,
+) -> None:
+    """Test Authorization header passed to AsyncClient with/without api_key."""
+    with patch(
+        "homeassistant.components.ollama.config_flow.ollama.AsyncClient",
+    ) as mock_async_client:
+        mock_async_client.return_value.list = AsyncMock(return_value={"models": []})
+
+        result = await hass.config_entries.flow.async_init(
+            DOMAIN, context={"source": "user"}
+        )
+        assert result["type"] is FlowResultType.FORM
+
+        result = await hass.config_entries.flow.async_configure(
+            result["flow_id"],
+            user_input=user_input,
+        )
+        await hass.async_block_till_done()
+
+    assert result["type"] is FlowResultType.CREATE_ENTRY
+    mock_async_client.assert_called_with(
+        host="http://localhost:11434",
+        headers=expected_headers,
+        verify=ANY,
+    )