-
Notifications
You must be signed in to change notification settings - Fork 4.4k
/
test_token_counter.py
50 lines (34 loc) · 1.53 KB
/
test_token_counter.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
"""Embeddings."""
from llama_index.core.callbacks.schema import CBEventType
from llama_index.core.callbacks.token_counting import TokenCountingHandler
TEST_PAYLOAD = {"chunks": ["one"], "formatted_prompt": "two", "response": "three"}
TEST_ID = "my id"
def test_on_event_start() -> None:
"""Test event start."""
handler = TokenCountingHandler()
event_id = handler.on_event_start(
CBEventType.LLM, payload=TEST_PAYLOAD, event_id=TEST_ID
)
assert event_id == TEST_ID
event_id = handler.on_event_start(
CBEventType.EMBEDDING, payload=TEST_PAYLOAD, event_id=TEST_ID
)
assert event_id == TEST_ID
assert len(handler.llm_token_counts) == 0
assert len(handler.embedding_token_counts) == 0
def test_on_event_end() -> None:
"""Test event end."""
handler = TokenCountingHandler()
handler.on_event_end(CBEventType.LLM, payload=TEST_PAYLOAD, event_id=TEST_ID)
assert len(handler.llm_token_counts) == 1
assert len(handler.embedding_token_counts) == 0
handler.on_event_end(CBEventType.EMBEDDING, payload=TEST_PAYLOAD, event_id=TEST_ID)
assert len(handler.llm_token_counts) == 1
assert len(handler.embedding_token_counts) == 1
assert handler.embedding_token_counts[0].total_token_count == 1
assert handler.llm_token_counts[0].total_token_count == 2
# test actual counts
# LLM should be two (prompt plus response)
# Embedding should be one (single token chunk)
assert handler.total_llm_token_count == 2
assert handler.total_embedding_token_count == 1