64 lines
1.5 KiB
Python
64 lines
1.5 KiB
Python
from dataclasses import dataclass
|
|
from typing import Protocol, Optional, Union
|
|
from .configuration import AIConfig
|
|
from .tags import Tag
|
|
from .message import Message
|
|
from .chat import Chat
|
|
|
|
|
|
class AIError(Exception):
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class Tokens:
|
|
prompt: int = 0
|
|
completion: int = 0
|
|
total: int = 0
|
|
|
|
|
|
@dataclass
|
|
class AIResponse:
|
|
"""
|
|
The response to an AI request. Consists of one or more messages
|
|
(each containing the question and a single answer) and the nr.
|
|
of used tokens.
|
|
"""
|
|
messages: list[Message]
|
|
tokens: Optional[Tokens] = None
|
|
|
|
|
|
class AI(Protocol):
|
|
"""
|
|
The base class for AI clients.
|
|
"""
|
|
|
|
name: str
|
|
config: AIConfig
|
|
|
|
def request(self,
|
|
question: Message,
|
|
context: Chat,
|
|
num_answers: int = 1,
|
|
otags: Optional[set[Tag]] = None) -> AIResponse:
|
|
"""
|
|
Make an AI request, asking the given question with the given
|
|
context (i. e. chat history). The nr. of requested answers
|
|
corresponds to the nr. of messages in the 'AIResponse'.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def models(self) -> list[str]:
|
|
"""
|
|
Return all models supported by this AI.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def tokens(self, data: Union[Message, Chat]) -> int:
|
|
"""
|
|
Computes the nr. of AI language tokens for the given message
|
|
or chat. Note that the computation may not be 100% accurate
|
|
and is not implemented for all AIs.
|
|
"""
|
|
raise NotImplementedError
|