diff --git a/libs/vertexai/langchain_google_vertexai/chat_models.py b/libs/vertexai/langchain_google_vertexai/chat_models.py index 6e6ad005..5a601261 100644 --- a/libs/vertexai/langchain_google_vertexai/chat_models.py +++ b/libs/vertexai/langchain_google_vertexai/chat_models.py @@ -667,7 +667,7 @@ def _stream( self._is_gemini_model, usage_metadata=response.to_dict().get("usage_metadata"), ) - if run_manager: + if run_manager and isinstance(message.content, str): run_manager.on_llm_new_token(message.content) if isinstance(message, AIMessageChunk): yield ChatGenerationChunk( @@ -744,7 +744,7 @@ async def _astream( self._is_gemini_model, usage_metadata=chunk.to_dict().get("usage_metadata"), ) - if run_manager: + if run_manager and isinstance(message.content, str): await run_manager.on_llm_new_token(message.content) if isinstance(message, AIMessageChunk): yield ChatGenerationChunk( diff --git a/libs/vertexai/tests/unit_tests/test_chat_models.py b/libs/vertexai/tests/unit_tests/test_chat_models.py index 5f6f04ef..1f848c1b 100644 --- a/libs/vertexai/tests/unit_tests/test_chat_models.py +++ b/libs/vertexai/tests/unit_tests/test_chat_models.py @@ -330,7 +330,6 @@ def test_default_params_gemini() -> None: assert mock_generate_content.call_args.args[0][0].parts[0].text == user_prompt - @pytest.mark.parametrize( "raw_candidate, expected", [ diff --git a/libs/vertexai/tests/unit_tests/test_function_utils.py b/libs/vertexai/tests/unit_tests/test_function_utils.py index f79336a1..160a55ed 100644 --- a/libs/vertexai/tests/unit_tests/test_function_utils.py +++ b/libs/vertexai/tests/unit_tests/test_function_utils.py @@ -3,7 +3,9 @@ from langchain_core.pydantic_v1 import BaseModel, Field from langchain_core.tools import tool -from vertexai.generative_models._generative_models import ToolConfig +from vertexai.generative_models._generative_models import ( # type: ignore[import-untyped] + ToolConfig, +) from langchain_google_vertexai.functions_utils import ( _format_tool_config, @@ -69,7 +71,6 @@ def test_format_tool_config(): assert isinstance(tool_config, ToolConfig) - def test_get_parameters_from_schema(): class StringEnum(str, Enum): pear = "pear"