From c40188f6b63e6130d83a64abbadbc212499db50a Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Tue, 11 Jun 2024 20:40:52 -0300 Subject: [PATCH] refactor: Update VertexAIEmbeddingsComponent to use Embeddings field type Update the VertexAIEmbeddingsComponent to use the `Embeddings` field type instead of importing it from `langchain_google_vertexai`. This ensures compatibility with the latest version of the langflow library and improves code readability. --- .../components/embeddings/VertexAIEmbeddings.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/backend/base/langflow/components/embeddings/VertexAIEmbeddings.py b/src/backend/base/langflow/components/embeddings/VertexAIEmbeddings.py index c0d249326..1a3eedc7c 100644 --- a/src/backend/base/langflow/components/embeddings/VertexAIEmbeddings.py +++ b/src/backend/base/langflow/components/embeddings/VertexAIEmbeddings.py @@ -1,8 +1,7 @@ from typing import List, Optional -from langchain_google_vertexai import VertexAIEmbeddings - from langflow.custom import CustomComponent +from langflow.field_typing import Embeddings class VertexAIEmbeddingsComponent(CustomComponent): @@ -71,7 +70,14 @@ class VertexAIEmbeddingsComponent(CustomComponent): temperature: float = 0.0, top_k: int = 40, top_p: float = 0.95, - ) -> VertexAIEmbeddings: + ) -> Embeddings: + try: + from langchain_google_vertexai import VertexAIEmbeddings + except ImportError: + raise ImportError( + "Please install the langchain-google-vertexai package to use the VertexAIEmbeddings component." + ) + return VertexAIEmbeddings( instance=instance, credentials=credentials,