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

(test) completion tokens counting + azure stream

Browse source
This commit is contained in:
ishaan-jaff 2024-01-03 12:06:35 +05:30
parent 96cb6f3b10
commit 2bea0c742e
1 changed files with 6 additions and 6 deletions
Download patch file Download diff file Expand all files Collapse all files

12
litellm/tests/test_stream_chunk_builder.py
View file

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