phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

(test) stream chunk builder - azure prompt tokens

Browse source
This commit is contained in:
ishaan-jaff 2023-12-29 15:45:41 +05:30
parent b1077ebc38
commit 73f60b7315
2 changed files with 37 additions and 22 deletions
Download patch file Download diff file Expand all files Collapse all files

38
litellm/tests/test_stream_chunk_builder.py
View file

@ -135,4 +135,40 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message():
pytest.fail(f"An exception occurred - {str(e)}")
test_stream_chunk_builder_litellm_tool_call_regular_message()
# test_stream_chunk_builder_litellm_tool_call_regular_message()
def test_stream_chunk_builder_count_prompt_tokens():
# test the prompt tokens for streamed responses == prompt tokens for non-streamed
# test the model for streamed responses == model for non-streamed
try:
messages = [{"role": "user", "content": "Hey, how's it going?"}]
litellm.set_verbose = False
response = litellm.completion(
model="azure/chatgpt-v-2",
messages=messages,
stream=True,
max_tokens=1,
complete_response=True,
)
print(f"Stream Assembled response: {response}")
stream_prompt_tokens = response.usage.prompt_tokens
stream_model = response.model
response = litellm.completion(
model="azure/chatgpt-v-2", messages=messages, max_tokens=1
)
print(f"\nNon Stream Response: {response}")
non_stream_prompt_tokens = response.usage.prompt_tokens
non_stream_model = response.model
assert stream_prompt_tokens == non_stream_prompt_tokens
assert stream_model != non_stream_model
except Exception as e:
pytest.fail(f"An exception occurred - {str(e)}")
# test_stream_chunk_builder_count_prompt_tokens()

21
litellm/tests/test_utils.py
View file

@ -317,24 +317,3 @@ def test_token_counter():
# test_token_counter()
def test_token_counter_azure():
# test that prompt_tokens counted == prompt_tokens from Azure API
try:
messages = [{"role": "user", "content": "hi how are you what time is it"}]
tokens_counted = token_counter(model="gpt-35-turbo", messages=messages)
print("Tokens Counted=", tokens_counted)
response = litellm.completion(model="azure/chatgpt-v-2", messages=messages)
prompt_tokens = response.usage.prompt_tokens
print("Tokens from Azure API=", prompt_tokens)
assert prompt_tokens == tokens_counted
except:
pytest.fail(f"Error occurred: {e}")
# test_token_counter_azure()