From b083d520c9d87ccb16766cabf4d354a3f9789a52 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 25 Feb 2024 23:00:53 -0300 Subject: [PATCH] Update embeddings and retrievers components --- .../embeddings/AmazonBedrockEmbeddings.py | 4 +- .../embeddings/VertexAIEmbeddings.py | 44 +++++++++++++++---- .../model_specs/AmazonBedrockSpecs.py | 4 +- .../components/models/HuggingFaceModel.py | 7 ++- .../components/retrievers/AmazonKendra.py | 6 ++- .../components/retrievers/MetalRetriever.py | 6 ++- 6 files changed, 55 insertions(+), 16 deletions(-) diff --git a/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py b/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py index 450e5ee3b..98deefcb8 100644 --- a/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py +++ b/src/backend/langflow/components/embeddings/AmazonBedrockEmbeddings.py @@ -1,7 +1,9 @@ from typing import Optional -from langchain.embeddings import BedrockEmbeddings from langchain.embeddings.base import Embeddings +from langchain_community.embeddings import BedrockEmbeddings + + from langflow import CustomComponent diff --git a/src/backend/langflow/components/embeddings/VertexAIEmbeddings.py b/src/backend/langflow/components/embeddings/VertexAIEmbeddings.py index 053fd6c7f..f7f4cf2ee 100644 --- a/src/backend/langflow/components/embeddings/VertexAIEmbeddings.py +++ b/src/backend/langflow/components/embeddings/VertexAIEmbeddings.py @@ -1,5 +1,5 @@ from langflow import CustomComponent -from langchain.embeddings import VertexAIEmbeddings +from langchain_community.embeddings import VertexAIEmbeddings from typing import Optional, List @@ -9,17 +9,45 @@ class VertexAIEmbeddingsComponent(CustomComponent): def build_config(self): return { - "credentials": {"display_name": "Credentials", "value": "", "file_types": [".json"], "field_type": "file"}, - "instance": {"display_name": "instance", "advanced": True, "field_type": "dict"}, - "location": {"display_name": "Location", "value": "us-central1", "advanced": True}, + "credentials": { + "display_name": "Credentials", + "value": "", + "file_types": [".json"], + "field_type": "file", + }, + "instance": { + "display_name": "instance", + "advanced": True, + "field_type": "dict", + }, + "location": { + "display_name": "Location", + "value": "us-central1", + "advanced": True, + }, "max_output_tokens": {"display_name": "Max Output Tokens", "value": 128}, - "max_retries": {"display_name": "Max Retries", "value": 6, "advanced": True}, - "model_name": {"display_name": "Model Name", "value": "textembedding-gecko"}, + "max_retries": { + "display_name": "Max Retries", + "value": 6, + "advanced": True, + }, + "model_name": { + "display_name": "Model Name", + "value": "textembedding-gecko", + }, "n": {"display_name": "N", "value": 1, "advanced": True}, "project": {"display_name": "Project", "advanced": True}, - "request_parallelism": {"display_name": "Request Parallelism", "value": 5, "advanced": True}, + "request_parallelism": { + "display_name": "Request Parallelism", + "value": 5, + "advanced": True, + }, "stop": {"display_name": "Stop", "advanced": True}, - "streaming": {"display_name": "Streaming", "value": False, "advanced": True}, + "streaming": { + "display_name": "Streaming", + "value": False, + "advanced": True, + }, "temperature": {"display_name": "Temperature", "value": 0.0}, "top_k": {"display_name": "Top K", "value": 40, "advanced": True}, "top_p": {"display_name": "Top P", "value": 0.95, "advanced": True}, diff --git a/src/backend/langflow/components/model_specs/AmazonBedrockSpecs.py b/src/backend/langflow/components/model_specs/AmazonBedrockSpecs.py index 5b931bb94..0c1e3b159 100644 --- a/src/backend/langflow/components/model_specs/AmazonBedrockSpecs.py +++ b/src/backend/langflow/components/model_specs/AmazonBedrockSpecs.py @@ -1,7 +1,9 @@ from typing import Optional from langchain.llms.base import BaseLLM -from langchain.llms.bedrock import Bedrock +from langchain_community.llms.bedrock import Bedrock + + from langflow import CustomComponent diff --git a/src/backend/langflow/components/models/HuggingFaceModel.py b/src/backend/langflow/components/models/HuggingFaceModel.py index c37cda530..4357ede61 100644 --- a/src/backend/langflow/components/models/HuggingFaceModel.py +++ b/src/backend/langflow/components/models/HuggingFaceModel.py @@ -1,7 +1,10 @@ from typing import Optional -from langflow import CustomComponent -from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint + from langchain_community.chat_models.huggingface import ChatHuggingFace +from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint + +from langflow import CustomComponent + from langflow.field_typing import Text diff --git a/src/backend/langflow/components/retrievers/AmazonKendra.py b/src/backend/langflow/components/retrievers/AmazonKendra.py index 827945a51..c0aedb1bc 100644 --- a/src/backend/langflow/components/retrievers/AmazonKendra.py +++ b/src/backend/langflow/components/retrievers/AmazonKendra.py @@ -1,7 +1,9 @@ from typing import Optional -from langflow import CustomComponent -from langchain.retrievers import AmazonKendraRetriever + from langchain.schema import BaseRetriever +from langchain_community.retrievers import AmazonKendraRetriever + +from langflow import CustomComponent class AmazonKendraRetrieverComponent(CustomComponent): diff --git a/src/backend/langflow/components/retrievers/MetalRetriever.py b/src/backend/langflow/components/retrievers/MetalRetriever.py index 88393f26d..1d37906d1 100644 --- a/src/backend/langflow/components/retrievers/MetalRetriever.py +++ b/src/backend/langflow/components/retrievers/MetalRetriever.py @@ -1,9 +1,11 @@ from typing import Optional -from langflow import CustomComponent -from langchain.retrievers import MetalRetriever + from langchain.schema import BaseRetriever +from langchain_community.retrievers import MetalRetriever from metal_sdk.metal import Metal # type: ignore +from langflow import CustomComponent + class MetalRetrieverComponent(CustomComponent): display_name: str = "Metal Retriever"