diff --git a/libs/langchain/langchain/chat_models/vertexai.py b/libs/langchain/langchain/chat_models/vertexai.py index dbe9919db2b..d6010e27fc7 100644 --- a/libs/langchain/langchain/chat_models/vertexai.py +++ b/libs/langchain/langchain/chat_models/vertexai.py @@ -132,16 +132,14 @@ class ChatVertexAI(_VertexAICommon, BaseChatModel): """Validate that the python package exists in environment.""" cls._try_init_vertexai(values) try: - if is_codey_model(values["model_name"]): - from vertexai.preview.language_models import CodeChatModel - - values["client"] = CodeChatModel.from_pretrained(values["model_name"]) - else: - from vertexai.preview.language_models import ChatModel - - values["client"] = ChatModel.from_pretrained(values["model_name"]) + from vertexai.language_models import ChatModel, CodeChatModel except ImportError: raise_vertex_import_error() + if is_codey_model(values["model_name"]): + model_cls = CodeChatModel + else: + model_cls = ChatModel + values["client"] = model_cls.from_pretrained(values["model_name"]) return values def _generate( diff --git a/libs/langchain/langchain/embeddings/vertexai.py b/libs/langchain/langchain/embeddings/vertexai.py index 54eb164a34a..26eba1565ef 100644 --- a/libs/langchain/langchain/embeddings/vertexai.py +++ b/libs/langchain/langchain/embeddings/vertexai.py @@ -17,7 +17,7 @@ class VertexAIEmbeddings(_VertexAICommon, Embeddings): """Validates that the python package exists in environment.""" cls._try_init_vertexai(values) try: - from vertexai.preview.language_models import TextEmbeddingModel + from vertexai.language_models import TextEmbeddingModel except ImportError: raise_vertex_import_error() values["client"] = TextEmbeddingModel.from_pretrained(values["model_name"]) diff --git a/libs/langchain/langchain/llms/vertexai.py b/libs/langchain/langchain/llms/vertexai.py index 994cfdf6bcf..177c64d5de2 100644 --- a/libs/langchain/langchain/llms/vertexai.py +++ b/libs/langchain/langchain/llms/vertexai.py @@ -195,19 +195,19 @@ class _VertexAICommon(_VertexAIBase): @property def _default_params(self) -> Dict[str, Any]: - if self.is_codey_model: - return { - "temperature": self.temperature, - "max_output_tokens": self.max_output_tokens, - } - else: - return { - "temperature": self.temperature, - "max_output_tokens": self.max_output_tokens, - "top_k": self.top_k, - "top_p": self.top_p, - "candidate_count": self.n, - } + params = { + "temperature": self.temperature, + "max_output_tokens": self.max_output_tokens, + "candidate_count": self.n, + } + if not self.is_codey_model: + params.update( + { + "top_k": self.top_k, + "top_p": self.top_p, + } + ) + return params @classmethod def _try_init_vertexai(cls, values: Dict) -> None: @@ -250,24 +250,33 @@ class VertexAI(_VertexAICommon, BaseLLM): tuned_model_name = values.get("tuned_model_name") model_name = values["model_name"] try: - if not is_codey_model(model_name): - from vertexai.preview.language_models import TextGenerationModel + from vertexai.language_models import ( + CodeGenerationModel, + TextGenerationModel, + ) + from vertexai.preview.language_models import ( + CodeGenerationModel as PreviewCodeGenerationModel, + ) + from vertexai.preview.language_models import ( + TextGenerationModel as PreviewTextGenerationModel, + ) - if tuned_model_name: - values["client"] = TextGenerationModel.get_tuned_model( - tuned_model_name - ) - else: - values["client"] = TextGenerationModel.from_pretrained(model_name) + if is_codey_model(model_name): + model_cls = CodeGenerationModel + preview_model_cls = PreviewCodeGenerationModel else: - from vertexai.preview.language_models import CodeGenerationModel + model_cls = TextGenerationModel + preview_model_cls = PreviewTextGenerationModel + + if tuned_model_name: + values["client"] = model_cls.get_tuned_model(tuned_model_name) + values["client_preview"] = preview_model_cls.get_tuned_model( + tuned_model_name + ) + else: + values["client"] = model_cls.from_pretrained(model_name) + values["client_preview"] = preview_model_cls.from_pretrained(model_name) - if tuned_model_name: - values["client"] = CodeGenerationModel.get_tuned_model( - tuned_model_name - ) - else: - values["client"] = CodeGenerationModel.from_pretrained(model_name) except ImportError: raise_vertex_import_error() @@ -287,12 +296,9 @@ class VertexAI(_VertexAICommon, BaseLLM): The integer number of tokens in the text. """ try: - result = self.client.count_tokens([text]) + result = self.client_preview.count_tokens([text]) except AttributeError: - raise NotImplementedError( - "Your google-cloud-aiplatform version didn't implement count_tokens." - "Please, install it with pip install google-cloud-aiplatform>=1.35.0" - ) + raise_vertex_import_error() return result.total_tokens @@ -319,12 +325,7 @@ class VertexAI(_VertexAICommon, BaseLLM): res = completion_with_retry( self, prompt, run_manager=run_manager, **params ) - if self.is_codey_model: - generations.append([_response_to_generation(res)]) - else: - generations.append( - [_response_to_generation(r) for r in res.candidates] - ) + generations.append([_response_to_generation(r) for r in res.candidates]) return LLMResult(generations=generations) async def _agenerate( diff --git a/libs/langchain/langchain/utilities/vertexai.py b/libs/langchain/langchain/utilities/vertexai.py index 334ab57f2ed..a60d5eea39a 100644 --- a/libs/langchain/langchain/utilities/vertexai.py +++ b/libs/langchain/langchain/utilities/vertexai.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from google.auth.credentials import Credentials -def raise_vertex_import_error(minimum_expected_version: str = "1.35.0") -> None: +def raise_vertex_import_error(minimum_expected_version: str = "1.36.0") -> None: """Raise ImportError related to Vertex SDK being not available. Args: @@ -16,7 +16,7 @@ def raise_vertex_import_error(minimum_expected_version: str = "1.35.0") -> None: ImportError: an ImportError that mentions a required version of the SDK. """ raise ImportError( - "Could not import VertexAI. Please, install it with " + "Please, install or upgrade the google-cloud-aiplatform library: " f"pip install google-cloud-aiplatform>={minimum_expected_version}" ) diff --git a/libs/langchain/tests/integration_tests/llms/test_vertexai.py b/libs/langchain/tests/integration_tests/llms/test_vertexai.py index bcc8ec273a9..1408d71ffa7 100644 --- a/libs/langchain/tests/integration_tests/llms/test_vertexai.py +++ b/libs/langchain/tests/integration_tests/llms/test_vertexai.py @@ -1,8 +1,6 @@ """Test Vertex AI API wrapper. -In order to run this test, you need to install VertexAI SDK (that is is the private -preview) and be whitelisted to list the models themselves: -In order to run this test, you need to install VertexAI SDK -pip install google-cloud-aiplatform>=1.35.0 +In order to run this test, you need to install VertexAI SDK: +pip install google-cloud-aiplatform>=1.36.0 Your end-user credentials would be used to make the calls (make sure you've run `gcloud auth login` first). @@ -150,10 +148,10 @@ async def test_model_garden_agenerate( assert len(output.generations) == 2 -def test_vertex_call_trigger_count_tokens() -> None: +def test_vertex_call_count_tokens() -> None: llm = VertexAI() - output = llm.get_num_tokens("Hi") - assert output == 2 + output = llm.get_num_tokens("How are you?") + assert output == 4 @pytest.mark.requires("google.cloud.aiplatform")