phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(test_text_completion.py): fix test

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-08 23:03:17 +05:30
parent ff12e023ae
commit fd6ccfca7d
3 changed files with 14 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/main.py
View file

@ -3164,8 +3164,7 @@ def stream_chunk_builder_text_completion(chunks: list, messages: Optional[List]=
else: else:
completion_output = "" completion_output = ""
# # Update usage information if needed # # Update usage information if needed
print(f"INSIDE TEXT COMPLETION STREAM CHUNK BUILDER") _usage = litellm.Usage()
_usage = litellm.Usage
print(f"messages: {messages}") print(f"messages: {messages}")
_usage.prompt_tokens = token_counter( _usage.prompt_tokens = token_counter(
model=model, messages=messages, count_response_tokens=True model=model, messages=messages, count_response_tokens=True
@ -3180,6 +3179,7 @@ def stream_chunk_builder_text_completion(chunks: list, messages: Optional[List]=
_usage.prompt_tokens + _usage.completion_tokens _usage.prompt_tokens + _usage.completion_tokens
) )
response["usage"] = _usage response["usage"] = _usage
print(f"final usage: {_usage}")
return litellm.TextCompletionResponse(**response) return litellm.TextCompletionResponse(**response)
def stream_chunk_builder(chunks: list, messages: Optional[list] = None): def stream_chunk_builder(chunks: list, messages: Optional[list] = None):

5
litellm/tests/test_text_completion.py
View file

@ -2930,9 +2930,10 @@ test_async_text_completion_stream()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_async_text_completion_chat_model_stream(): async def test_async_text_completion_chat_model_stream():
try: try:
prompt="good morning"
response = await litellm.atext_completion( response = await litellm.atext_completion(
model="gpt-3.5-turbo", model="gpt-3.5-turbo",
prompt="good morning", prompt=prompt,
stream=True, stream=True,
max_tokens=10, max_tokens=10,
) )
@ -2948,7 +2949,7 @@ async def test_async_text_completion_chat_model_stream():
assert ( assert (
num_finish_reason == 1 num_finish_reason == 1
), f"expected only one finish reason. Got {num_finish_reason}" ), f"expected only one finish reason. Got {num_finish_reason}"
response_obj = litellm.stream_chunk_builder(chunks=chunks) response_obj = litellm.stream_chunk_builder(chunks=chunks, messages=[{"role": "user", "content": prompt}])
cost = litellm.completion_cost(completion_response=response_obj) cost = litellm.completion_cost(completion_response=response_obj)
assert cost > 0 assert cost > 0
except Exception as e: except Exception as e:

13
litellm/utils.py
View file

@ -2267,11 +2267,16 @@ def client(original_function):
and kwargs["complete_response"] == True and kwargs["complete_response"] == True
): ):
chunks = [] chunks = []
for idx, chunk in enumerate(result): async for chunk in result:
chunks.append(chunk) chunks.append(chunk)
return litellm.stream_chunk_builder( if call_type == CallTypes.acompletion.value:
chunks, messages=kwargs.get("messages", None) return litellm.stream_chunk_builder(
) chunks, messages=kwargs.get("messages")
)
elif call_type == CallTypes.atext_completion.value:
return litellm.stream_chunk_builder(
chunks, messages=[{"role": "user", "content": kwargs.get("prompt")}]
)
else: else:
return result return result