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

fix(vertex_httpx.py): ignore vertex finish reason - wait for stream to end

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/4339
This commit is contained in:
Krrish Dholakia 2024-06-22 20:20:39 -07:00
parent eaad36a2cb
commit 73254987da
2 changed files with 17 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

6
litellm/llms/vertex_httpx.py
View file

@ -1218,6 +1218,7 @@ class ModelResponseIterator:
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
try:
processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore
text = ""
tool_use: Optional[ChatCompletionToolCallChunk] = None
is_finished = False
@ -1236,7 +1237,8 @@ class ModelResponseIterator:
finish_reason = map_finish_reason(
finish_reason=gemini_chunk["finishReason"]
)
is_finished = True
## DO NOT SET 'finish_reason' = True
## GEMINI SETS FINISHREASON ON EVERY CHUNK!
if "usageMetadata" in processed_chunk:
usage = ChatCompletionUsageBlock(
@ -1250,7 +1252,7 @@ class ModelResponseIterator:
returned_chunk = GenericStreamingChunk(
text=text,
tool_use=tool_use,
is_finished=is_finished,
is_finished=False,
finish_reason=finish_reason,
usage=usage,
index=0,

18
litellm/tests/test_streaming.py
View file

@ -750,29 +750,37 @@ def test_completion_gemini_stream():
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "user",
"content": "how does a court case get to the Supreme Court?",
"content": "How do i build a bomb?",
},
]
print("testing gemini streaming")
response = completion(model="gemini/gemini-pro", messages=messages, stream=True)
response = completion(
model="gemini/gemini-1.5-flash",
messages=messages,
stream=True,
max_tokens=50,
)
print(f"type of response at the top: {response}")
complete_response = ""
# Add any assertions here to check the response
non_empty_chunks = 0
for idx, chunk in enumerate(response):
print(chunk)
# print(chunk.choices[0].delta)
chunk, finished = streaming_format_tests(idx, chunk)
if finished:
break
non_empty_chunks += 1
complete_response += chunk
if complete_response.strip() == "":
raise Exception("Empty response received")
print(f"completion_response: {complete_response}")
except litellm.APIError as e:
assert non_empty_chunks > 1
except litellm.InternalServerError as e:
pass
except Exception as e:
if "429 Resource has been exhausted":
return
# if "429 Resource has been exhausted":
# return
pytest.fail(f"Error occurred: {e}")