72fed878b4
* Add ollama conversation agent * Change iot class * Much better default template * Slight adjustment to prompt * Make casing consistent * Switch to ollama Python fork * Add prompt to tests * Rename to "ollama" * Download models in config flow * Update homeassistant/components/ollama/config_flow.py --------- Co-authored-by: Paulus Schoutsen <balloob@gmail.com>
47 lines
1 KiB
Python
47 lines
1 KiB
Python
"""Models for Ollama integration."""
|
|
|
|
from dataclasses import dataclass
|
|
from enum import StrEnum
|
|
from functools import cached_property
|
|
|
|
import ollama
|
|
|
|
from homeassistant.core import State
|
|
|
|
|
|
class MessageRole(StrEnum):
|
|
"""Role of a chat message."""
|
|
|
|
SYSTEM = "system" # prompt
|
|
USER = "user"
|
|
|
|
|
|
@dataclass
|
|
class MessageHistory:
|
|
"""Chat message history."""
|
|
|
|
timestamp: float
|
|
"""Timestamp of last use in seconds."""
|
|
|
|
messages: list[ollama.Message]
|
|
"""List of message history, including system prompt and assistant responses."""
|
|
|
|
@property
|
|
def num_user_messages(self) -> int:
|
|
"""Return a count of user messages."""
|
|
return sum(m["role"] == MessageRole.USER for m in self.messages)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ExposedEntity:
|
|
"""Relevant information about an exposed entity."""
|
|
|
|
entity_id: str
|
|
state: State
|
|
names: list[str]
|
|
area_names: list[str]
|
|
|
|
@cached_property
|
|
def domain(self) -> str:
|
|
"""Get domain from entity id."""
|
|
return self.entity_id.split(".", maxsplit=1)[0]
|