diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py
index da509328..3f7f42d8 100644
--- a/libs/vertexai/langchain_google_vertexai/chat_models.py
+++ b/libs/vertexai/langchain_google_vertexai/chat_models.py
@@ -4,7 +4,7 @@
import json
import logging
from dataclasses import dataclass, field
-from typing import Any, Dict, Iterator, List, Optional, Union, cast
+from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union, cast
import proto # type: ignore[import-untyped]
from google.cloud.aiplatform_v1beta1.types.content import Part as GapicPart
@@ -539,6 +539,48 @@ def _stream(
),
)
+ async def _astream(
+ self,
+ messages: List[BaseMessage],
+ stop: Optional[List[str]] = None,
+ run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+ **kwargs: Any,
+ ) -> AsyncIterator[ChatGenerationChunk]:
+ if not self._is_gemini_model:
+ raise NotImplementedError()
+ params = self._prepare_params(stop=stop, stream=True, **kwargs)
+ history_gemini = _parse_chat_history_gemini(
+ messages,
+ project=self.project,
+ convert_system_message_to_human=self.convert_system_message_to_human,
+ )
+ message = history_gemini.pop()
+ chat = self.client.start_chat(history=history_gemini)
+ raw_tools = params.pop("functions") if "functions" in params else None
+ tools = _format_tools_to_vertex_tool(raw_tools) if raw_tools else None
+ safety_settings = params.pop("safety_settings", None)
+ async for chunk in await chat.send_message_async(
+ message,
+ stream=True,
+ generation_config=params,
+ safety_settings=safety_settings,
+ tools=tools,
+ ):
+ message = _parse_response_candidate(chunk.candidates[0])
+ if run_manager:
+ await run_manager.on_llm_new_token(message.content)
+ yield ChatGenerationChunk(
+ message=AIMessageChunk(
+ content=message.content,
+ additional_kwargs=message.additional_kwargs,
+ ),
+ generation_info=get_generation_info(
+ chunk.candidates[0],
+ self._is_gemini_model,
+ usage_metadata=chunk.to_dict().get("usage_metadata"),
+ ),
+ )
+
def _start_chat(
self, history: _ChatHistory, **kwargs: Any
) -> Union[ChatSession, CodeChatSession]:
diff --git a/libs/vertexai/langchain_google_vertexai/llms.py b/libs/vertexai/langchain_google_vertexai/llms.py
index 8b6f1062..c236ab3b 100644
--- a/libs/vertexai/langchain_google_vertexai/llms.py
+++ b/libs/vertexai/langchain_google_vertexai/llms.py
@@ -1,7 +1,7 @@
from __future__ import annotations
from concurrent.futures import Executor
-from typing import Any, ClassVar, Dict, Iterator, List, Optional, Union
+from typing import Any, AsyncIterator, ClassVar, Dict, Iterator, List, Optional, Union
import vertexai # type: ignore[import-untyped]
from google.api_core.client_options import ClientOptions
@@ -95,6 +95,7 @@ async def _acompletion_with_retry(
llm: VertexAI,
prompt: str,
is_gemini: bool = False,
+ stream: bool = False,
run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> Any:
@@ -105,17 +106,22 @@ async def _acompletion_with_retry(
@retry_decorator
async def _acompletion_with_retry_inner(
- prompt: str, is_gemini: bool = False, **kwargs: Any
+ prompt: str, is_gemini: bool = False, stream: bool = False, **kwargs: Any
) -> Any:
if is_gemini:
return await llm.client.generate_content_async(
prompt,
generation_config=kwargs,
+ stream=stream,
safety_settings=kwargs.pop("safety_settings", None),
)
+ if stream:
+ raise ValueError("Async streaming is supported only for Gemini family!")
return await llm.client.predict_async(prompt, **kwargs)
- return await _acompletion_with_retry_inner(prompt, is_gemini, **kwargs)
+ return await _acompletion_with_retry_inner(
+ prompt, is_gemini, stream=stream, **kwargs
+ )
class _VertexAIBase(BaseModel):
@@ -453,6 +459,34 @@ def _stream(
verbose=self.verbose,
)
+ async def _astream(
+ self,
+ prompt: str,
+ stop: Optional[List[str]] = None,
+ run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+ **kwargs: Any,
+ ) -> AsyncIterator[GenerationChunk]:
+ params = self._prepare_params(stop=stop, stream=True, **kwargs)
+ if not self._is_gemini_model:
+ raise ValueError("Async streaming is supported only for Gemini family!")
+ async for chunk in await _acompletion_with_retry(
+ self,
+ prompt,
+ stream=True,
+ is_gemini=self._is_gemini_model,
+ run_manager=run_manager,
+ **params,
+ ):
+ usage_metadata = chunk.to_dict().get("usage_metadata")
+ chunk = self._candidate_to_generation(
+ chunk.candidates[0], stream=True, usage_metadata=usage_metadata
+ )
+ yield chunk
+ if run_manager:
+ await run_manager.on_llm_new_token(
+ chunk.text, chunk=chunk, verbose=self.verbose
+ )
+
class VertexAIModelGarden(_VertexAIBase, BaseLLM):
"""Large language models served from Vertex AI Model Garden."""
diff --git a/libs/vertexai/tests/integration_tests/test_chat_models.py b/libs/vertexai/tests/integration_tests/test_chat_models.py
index c3721af3..365c7a21 100644
--- a/libs/vertexai/tests/integration_tests/test_chat_models.py
+++ b/libs/vertexai/tests/integration_tests/test_chat_models.py
@@ -86,6 +86,14 @@ def test_vertexai_stream(model_name: str) -> None:
assert isinstance(chunk, AIMessageChunk)
+async def test_vertexai_astream() -> None:
+ model = ChatVertexAI(temperature=0, model_name="gemini-pro")
+ message = HumanMessage(content="Hello")
+
+ async for chunk in model.astream([message]):
+ assert isinstance(chunk, AIMessageChunk)
+
+
def test_vertexai_single_call_with_context() -> None:
model = ChatVertexAI()
raw_context = (
diff --git a/libs/vertexai/tests/integration_tests/test_image_utils.py b/libs/vertexai/tests/integration_tests/test_image_utils.py
index b297fa0d..87f75a81 100644
--- a/libs/vertexai/tests/integration_tests/test_image_utils.py
+++ b/libs/vertexai/tests/integration_tests/test_image_utils.py
@@ -1,9 +1,11 @@
+import pytest
from google.cloud import storage # type: ignore[attr-defined]
from google.cloud.exceptions import NotFound
from langchain_google_vertexai._image_utils import ImageBytesLoader
+@pytest.mark.skip("CI testing not set up")
def test_image_utils():
base64_image = (
"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABgAAAAYCAYAAADgdz34AAAA"
diff --git a/libs/vertexai/tests/integration_tests/test_llms.py b/libs/vertexai/tests/integration_tests/test_llms.py
index 29f29db2..6b252f42 100644
--- a/libs/vertexai/tests/integration_tests/test_llms.py
+++ b/libs/vertexai/tests/integration_tests/test_llms.py
@@ -113,6 +113,12 @@ async def test_vertex_consistency() -> None:
assert output.generations[0][0].text == async_output.generations[0][0].text
+async def test_astream() -> None:
+ llm = VertexAI(temperature=0, model_name="gemini-pro")
+ async for token in llm.astream("I'm Pickle Rick"):
+ assert isinstance(token, str)
+
+
@pytest.mark.skip("CI testing not set up")
@pytest.mark.parametrize(
"endpoint_os_variable_name,result_arg",