Add configurable keep_alive to Ollama integration, change default to 5m (#119341)

Co-authored-by: Franck Nijhof <git@frenck.dev>
This commit is contained in:
srwareham 2024-07-07 07:38:01 -07:00 committed by GitHub
parent 6937aed9fe
commit 76bdc4f5c4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 33 additions and 4 deletions

View File

@ -14,7 +14,14 @@ from homeassistant.core import HomeAssistant
from homeassistant.exceptions import ConfigEntryNotReady
from homeassistant.helpers import config_validation as cv
from .const import CONF_MAX_HISTORY, CONF_MODEL, CONF_PROMPT, DEFAULT_TIMEOUT, DOMAIN
from .const import (
CONF_KEEP_ALIVE,
CONF_MAX_HISTORY,
CONF_MODEL,
CONF_PROMPT,
DEFAULT_TIMEOUT,
DOMAIN,
)
_LOGGER = logging.getLogger(__name__)
@ -23,6 +30,7 @@ __all__ = [
"CONF_PROMPT",
"CONF_MODEL",
"CONF_MAX_HISTORY",
"CONF_KEEP_ALIVE",
"DOMAIN",
]

View File

@ -33,9 +33,11 @@ from homeassistant.helpers.selector import (
)
from .const import (
CONF_KEEP_ALIVE,
CONF_MAX_HISTORY,
CONF_MODEL,
CONF_PROMPT,
DEFAULT_KEEP_ALIVE,
DEFAULT_MAX_HISTORY,
DEFAULT_MODEL,
DEFAULT_PROMPT,
@ -235,6 +237,16 @@ def ollama_config_option_schema(options: MappingProxyType[str, Any]) -> dict:
min=0, max=sys.maxsize, step=1, mode=NumberSelectorMode.BOX
)
),
vol.Optional(
CONF_KEEP_ALIVE,
description={
"suggested_value": options.get(CONF_KEEP_ALIVE, DEFAULT_KEEP_ALIVE)
},
): NumberSelector(
NumberSelectorConfig(
min=-1, max=sys.maxsize, step=1, mode=NumberSelectorMode.BOX
)
),
}

View File

@ -72,6 +72,9 @@ An overview of the areas and the devices in this smart home:
Answer the user's questions using the information about this smart home.
Keep your answers brief and do not apologize."""
CONF_KEEP_ALIVE = "keep_alive"
DEFAULT_KEEP_ALIVE = -1 # seconds. -1 = indefinite, 0 = never
KEEP_ALIVE_FOREVER = -1
DEFAULT_TIMEOUT = 5.0 # seconds

View File

@ -26,13 +26,14 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.util import ulid
from .const import (
CONF_KEEP_ALIVE,
CONF_MAX_HISTORY,
CONF_MODEL,
CONF_PROMPT,
DEFAULT_KEEP_ALIVE,
DEFAULT_MAX_HISTORY,
DEFAULT_PROMPT,
DOMAIN,
KEEP_ALIVE_FOREVER,
MAX_HISTORY_SECONDS,
)
from .models import ExposedEntity, MessageHistory, MessageRole
@ -151,7 +152,8 @@ class OllamaConversationEntity(
# Make a copy of the messages because we mutate the list later
messages=list(message_history.messages),
stream=False,
keep_alive=KEEP_ALIVE_FOREVER,
# keep_alive requires specifying unit. In this case, seconds
keep_alive=f"{settings.get(CONF_KEEP_ALIVE, DEFAULT_KEEP_ALIVE)}s",
)
except (ollama.RequestError, ollama.ResponseError) as err:
_LOGGER.error("Unexpected error talking to Ollama server: %s", err)

View File

@ -25,7 +25,11 @@
"init": {
"data": {
"prompt": "Prompt template",
"max_history": "Max history messages"
"max_history": "Max history messages",
"keep_alive": "Keep alive"
},
"data_description": {
"keep_alive": "Duration in seconds for Ollama to keep model in memory. -1 = indefinite, 0 = never."
}
}
}