Added support for max_tokens attribute in ChatVertexAI (#714)

langchain-ai · Jan 28, 2025 · c3d4061 · c3d4061
1 parent 106872f
commit c3d4061
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/libs/vertexai/langchain_google_vertexai/_base.py b/libs/vertexai/langchain_google_vertexai/_base.py
@@ -225,6 +225,10 @@ def _is_gemini_model(self) -> bool:
     def _llm_type(self) -> str:
         return "vertexai"
 
+    @property
+    def max_tokens(self) -> int | None:
+        return self.max_output_tokens
+
     @property
     def _identifying_params(self) -> Dict[str, Any]:
         """Gets the identifying parameters."""

diff --git a/libs/vertexai/tests/unit_tests/test_chat_models.py b/libs/vertexai/tests/unit_tests/test_chat_models.py
@@ -119,10 +119,13 @@ def test_tuned_model_name() -> None:
         model_name="gemini-pro",
         project="test-project",
         tuned_model_name="projects/123/locations/europe-west4/endpoints/456",
+        max_tokens=500,
     )
     assert llm.model_name == "gemini-pro"
     assert llm.tuned_model_name == "projects/123/locations/europe-west4/endpoints/456"
     assert llm.full_model_name == "projects/123/locations/europe-west4/endpoints/456"
+    assert llm.max_output_tokens == 500
+    assert llm.max_tokens == 500
 
 
 def test_parse_examples_correct() -> None: