75 lines
1.8 KiB
Python
75 lines
1.8 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.
|
|
"""
|
|
|
|
ID: str
|
|
name: str
|
|
config: AIConfig
|
|
|
|
def request(self,
|
|
question: Message,
|
|
chat: Chat,
|
|
num_answers: int = 1,
|
|
otags: Optional[set[Tag]] = None) -> AIResponse:
|
|
"""
|
|
Make an AI request. Parameters:
|
|
* question: the question to ask
|
|
* chat: the chat history to be added as context
|
|
* num_answers: nr. of requested answers (corresponds
|
|
to the nr. of messages in the 'AIResponse')
|
|
* otags: the output tags, i. e. the tags that all
|
|
returned messages should contain
|
|
"""
|
|
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
|
|
|
|
def print(self) -> None:
|
|
"""
|
|
Print some info about the current AI, like system message.
|
|
"""
|
|
pass
|