added tokens() function to Message and Chat

This commit is contained in:
juk0de 2023-09-01 08:57:54 +02:00 committed by Oleksandr Kozachuk
parent 815a21893c
commit 6737fa98c7
2 changed files with 19 additions and 0 deletions

View File

@ -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:

View File

@ -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