fix linter

langchain-ai · Apr 11, 2024 · cf84c17 · cf84c17
1 parent d2bd4fa
commit cf84c17
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git 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
@@ -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
@@ -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"