mirror of
https://github.com/home-assistant/core.git
synced 2025-05-19 21:39:19 +00:00
57ab567d08
* Update ollama to use the ChatLog/ChatSession APIs * Add documentation about history trimming. * Revert changes to chat_log.py * Explicitly check for SystemContent when converting system messages * Remove half of a comment
29 lines
641 B
Python
29 lines
641 B
Python
"""Models for Ollama integration."""
|
|
|
|
from dataclasses import dataclass
|
|
from enum import StrEnum
|
|
|
|
import ollama
|
|
|
|
|
|
class MessageRole(StrEnum):
|
|
"""Role of a chat message."""
|
|
|
|
SYSTEM = "system" # prompt
|
|
USER = "user"
|
|
ASSISTANT = "assistant"
|
|
TOOL = "tool"
|
|
|
|
|
|
@dataclass
|
|
class MessageHistory:
|
|
"""Chat message history."""
|
|
|
|
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.value for m in self.messages)
|