From e362a164cf7ee32aa3f374bd33c6bbc603f10345 Mon Sep 17 00:00:00 2001
From: Bagatur <baskaryan@gmail.com>
Date: Tue, 1 Oct 2024 11:45:41 -0700
Subject: [PATCH] vertexai[patch]: use default astream for gemini models

---
 libs/vertexai/langchain_google_vertexai/chat_models.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py
index 1dc1398b..b4a5dcac 100644
--- 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(