4b2073ca59
* Add LLM tools support for Ollama * fix tests * coverage * Separate call for tool parameters * Fix example * hint on parameters schema if LLM forgot to request it * Switch to native tool call functionality * Fix tests * Fix tools list * update strings and default model * Ignore mypy error until fixed upstream * Ignore mypy error until fixed upstream * Add missing prompt part * Update default model
31 lines
707 B
Python
31 lines
707 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."""
|
|
|
|
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.value for m in self.messages)
|