diff --git a/chatmastermind/ai.py b/chatmastermind/ai.py new file mode 100644 index 0000000..4a8b914 --- /dev/null +++ b/chatmastermind/ai.py @@ -0,0 +1,63 @@ +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