From 7617f340f065b3d3d02421f0295328aec1a01d2a Mon Sep 17 00:00:00 2001
From: anovazzi1 <otavio2204@gmail.com>
Date: Fri, 16 Feb 2024 16:44:11 -0300
Subject: [PATCH] Add HuggingFaceEndpointsComponent to the project

---
 .../langflow/components/models/HuggingFace.py | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 src/backend/langflow/components/models/HuggingFace.py

diff --git a/src/backend/langflow/components/models/HuggingFace.py b/src/backend/langflow/components/models/HuggingFace.py
new file mode 100644
index 000000000..e49a60b08
--- /dev/null
+++ b/src/backend/langflow/components/models/HuggingFace.py
@@ -0,0 +1,50 @@
+from typing import Optional
+from langflow import CustomComponent
+from langchain.llms.huggingface_endpoint import HuggingFaceEndpoint
+from langchain_community.chat_models.huggingface import ChatHuggingFace
+from langflow.field_typing import Text
+
+
+class HuggingFaceEndpointsComponent(CustomComponent):
+    display_name: str = "Hugging Face Inference API models"
+    description: str = "LLM model from Hugging Face Inference API."
+
+    def build_config(self):
+        return {
+            "endpoint_url": {"display_name": "Endpoint URL", "password": True},
+            "task": {
+                "display_name": "Task",
+                "options": ["text2text-generation", "text-generation", "summarization"],
+            },
+            "huggingfacehub_api_token": {"display_name": "API token", "password": True},
+            "model_kwargs": {
+                "display_name": "Model Keyword Arguments",
+                "field_type": "code",
+            },
+            "code": {"show": False},
+            "inputs": {"display_name": "Input"},
+        }
+
+    def build(
+        self,
+        inputs: str,
+        endpoint_url: str,
+        task: str = "text2text-generation",
+        huggingfacehub_api_token: Optional[str] = None,
+        model_kwargs: Optional[dict] = None,
+    ) -> Text:
+        try:
+            llm = HuggingFaceEndpoint(
+                endpoint_url=endpoint_url,
+                task=task,
+                huggingfacehub_api_token=huggingfacehub_api_token,
+                model_kwargs=model_kwargs,
+            )
+        except Exception as e:
+            raise ValueError("Could not connect to HuggingFace Endpoints API.") from e
+        output = ChatHuggingFace(llm=llm)
+        message = output.invoke(inputs)
+        result = message.content if hasattr(message, "content") else message
+        self.status = result
+        return result
+