Skip to content

Commit

Permalink
partners/mistralai: Fix KeyError in Vertex AI stream (#28624)
Browse files Browse the repository at this point in the history
- **Description:** Streaming response from Mistral model using Vertex AI
raises KeyError when trying to access `choices` key, that the last chunk
doesn't have. The fix is to access the key safely using `get()`.
  - **Issue:** #27886
  - **Dependencies:**
  - **Twitter handle:**
  • Loading branch information
nikitajoyn authored Dec 9, 2024
1 parent bdb4cf7 commit 9fcd203
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions libs/partners/mistralai/langchain_mistralai/chat_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -595,7 +595,7 @@ def _stream(
for chunk in self.completion_with_retry(
messages=message_dicts, run_manager=run_manager, **params
):
if len(chunk["choices"]) == 0:
if len(chunk.get("choices", [])) == 0:
continue
new_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)
# make future chunks same type as first chunk
Expand All @@ -621,7 +621,7 @@ async def _astream(
async for chunk in await acompletion_with_retry(
self, messages=message_dicts, run_manager=run_manager, **params
):
if len(chunk["choices"]) == 0:
if len(chunk.get("choices", [])) == 0:
continue
new_chunk = _convert_chunk_to_message_chunk(chunk, default_chunk_class)
# make future chunks same type as first chunk
Expand Down

0 comments on commit 9fcd203

Please sign in to comment.