vertexai[patch]: use default astream for gemini models (#526)

langchain-ai · Oct 1, 2024 · 39fd967 · 39fd967
2 parents c03162e + e362a16
commit 39fd967
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py
@@ -1593,8 +1593,13 @@ async def _astream(
         run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
         **kwargs: Any,
     ) -> AsyncIterator[ChatGenerationChunk]:
+        # TODO: Update to properly support async streaming from gemini.
         if not self._is_gemini_model:
-            raise NotImplementedError()
+            async for chunk in super()._astream(
+                messages, stop=stop, run_manager=run_manager, **kwargs
+            ):
+                yield chunk
+            return
         request = self._prepare_request_gemini(messages=messages, stop=stop, **kwargs)
 
         response_iter = _acompletion_with_retry(