diff --git a/chatmastermind/chat.py b/chatmastermind/chat.py index 6b85ba5..366988b 100644 --- a/chatmastermind/chat.py +++ b/chatmastermind/chat.py @@ -129,6 +129,13 @@ class Chat: tags |= m.filter_tags(prefix, contain) return tags + def tokens(self) -> int: + """ + Returns the nr. of AI language tokens used by all messages in this chat. + If unknown, 0 is returned. + """ + return sum(m.tokens() for m in self.messages) + def print(self, dump: bool = False, source_code_only: bool = False, with_tags: bool = False, with_file: bool = False, paged: bool = True) -> None: diff --git a/chatmastermind/message.py b/chatmastermind/message.py index 3eca26e..675ab3a 100644 --- a/chatmastermind/message.py +++ b/chatmastermind/message.py @@ -132,6 +132,7 @@ class Question(str): """ A single question with a defined header. """ + tokens: int = 0 # tokens used by this question txt_header: ClassVar[str] = '=== QUESTION ===' yaml_key: ClassVar[str] = 'question' @@ -165,6 +166,7 @@ class Answer(str): """ A single answer with a defined header. """ + tokens: int = 0 # tokens used by this answer txt_header: ClassVar[str] = '=== ANSWER ===' yaml_key: ClassVar[str] = 'answer' @@ -502,3 +504,13 @@ class Message(): def as_dict(self) -> dict[str, Any]: return asdict(self) + + def tokens(self) -> int: + """ + Returns the nr. of AI language tokens used by this message. + If unknown, 0 is returned. + """ + if self.answer: + return self.question.tokens + self.answer.tokens + else: + return self.question.tokens