From 75dbb68dfcfe5790f3db48991b067324756a71ea Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Mon, 19 Aug 2024 14:05:59 -0300
Subject: [PATCH] feat: add ComponentTool to support converting Component to
 Tool (#3412)

* feat: Add ComponentTool to convert a Component to a Tool

* test(component): add unit test for ComponentTool with ChatInput input.

* feat: Add method to convert Component to ComponentTool.

* feat: Add unit test for ChatInput to Tool conversion.

* chore: add comment

* test: fix assertion

---------

Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com>
---
 .../langflow/base/tools/component_tool.py     | 41 ++++++++++++
 .../custom/custom_component/component.py      |  6 ++
 .../unit/base/tools/test_component_tool.py    | 62 +++++++++++++++++++
 .../component/test_component_to_tool.py       | 16 +++++
 4 files changed, 125 insertions(+)
 create mode 100644 src/backend/base/langflow/base/tools/component_tool.py
 create mode 100644 src/backend/tests/unit/base/tools/test_component_tool.py
 create mode 100644 src/backend/tests/unit/custom/component/test_component_to_tool.py

diff --git a/src/backend/base/langflow/base/tools/component_tool.py b/src/backend/base/langflow/base/tools/component_tool.py
new file mode 100644
index 000000000..3f26cde5e
--- /dev/null
+++ b/src/backend/base/langflow/base/tools/component_tool.py
@@ -0,0 +1,41 @@
+from typing import Any
+
+from langchain_core.tools import BaseTool, ToolException
+
+from langflow.custom.custom_component.component import Component
+
+
+class ComponentTool(BaseTool):
+    name: str
+    description: str
+    component: "Component"
+
+    def __init__(self, component: "Component") -> None:
+        """Initialize the tool."""
+        from langflow.io.schema import create_input_schema
+
+        name = component.name or component.__class__.__name__
+        description = component.description or ""
+        args_schema = create_input_schema(component.inputs)
+        super().__init__(name=name, description=description, args_schema=args_schema, component=component)
+        # self.component = component
+
+    @property
+    def args(self) -> dict:
+        schema = self.get_input_schema()
+        return schema.schema()["properties"]
+
+    def _run(
+        self,
+        *args: Any,
+        **kwargs: Any,
+    ) -> dict:
+        """Use the tool."""
+        try:
+            results, _ = self.component(**kwargs)
+            return results
+        except Exception as e:
+            raise ToolException(f"Error running {self.name}: {e}")
+
+
+ComponentTool.update_forward_refs()
diff --git a/src/backend/base/langflow/custom/custom_component/component.py b/src/backend/base/langflow/custom/custom_component/component.py
index 992e3724a..c78856d7f 100644
--- a/src/backend/base/langflow/custom/custom_component/component.py
+++ b/src/backend/base/langflow/custom/custom_component/component.py
@@ -653,3 +653,9 @@ class Component(CustomComponent):
 
     def _get_fallback_input(self, **kwargs):
         return Input(**kwargs)
+
+    def to_tool(self):
+        # TODO: This is a temporary solution to avoid circular imports
+        from langflow.base.tools.component_tool import ComponentTool
+
+        return ComponentTool(component=self)
diff --git a/src/backend/tests/unit/base/tools/test_component_tool.py b/src/backend/tests/unit/base/tools/test_component_tool.py
new file mode 100644
index 000000000..25c547099
--- /dev/null
+++ b/src/backend/tests/unit/base/tools/test_component_tool.py
@@ -0,0 +1,62 @@
+import pytest
+
+from langflow.base.tools.component_tool import ComponentTool
+from langflow.components.inputs.ChatInput import ChatInput
+
+
+@pytest.fixture
+def client():
+    pass
+
+
+def test_component_tool():
+    chat_input = ChatInput()
+    component_tool = ComponentTool(component=chat_input)
+    assert component_tool.name == "ChatInput"
+    assert component_tool.description == chat_input.description
+    assert component_tool.args == {
+        "input_value": {
+            "default": "",
+            "description": "Message to be passed as input.",
+            "title": "Input Value",
+            "type": "string",
+        },
+        "should_store_message": {
+            "default": True,
+            "description": "Store the message in the history.",
+            "title": "Should Store Message",
+            "type": "boolean",
+        },
+        "sender": {
+            "default": "User",
+            "description": "Type of sender.",
+            "enum": ["Machine", "User"],
+            "title": "Sender",
+            "type": "string",
+        },
+        "sender_name": {
+            "default": "User",
+            "description": "Name of the sender.",
+            "title": "Sender Name",
+            "type": "string",
+        },
+        "session_id": {
+            "default": "",
+            "description": "The session ID of the chat. If empty, the current session ID parameter will be used.",
+            "title": "Session Id",
+            "type": "string",
+        },
+        "files": {
+            "default": "",
+            "description": "Files to be sent with the message.",
+            "items": {"type": "string"},
+            "title": "Files",
+            "type": "array",
+        },
+    }
+    assert component_tool.component == chat_input
+
+    result = component_tool.invoke(input=dict(input_value="test"))
+    assert isinstance(result, dict)
+    assert hasattr(result["message"], "get_text")
+    assert result["message"].get_text() == "test"
diff --git a/src/backend/tests/unit/custom/component/test_component_to_tool.py b/src/backend/tests/unit/custom/component/test_component_to_tool.py
new file mode 100644
index 000000000..076240a4d
--- /dev/null
+++ b/src/backend/tests/unit/custom/component/test_component_to_tool.py
@@ -0,0 +1,16 @@
+import pytest
+
+from langflow.components.inputs.ChatInput import ChatInput
+
+
+@pytest.fixture
+def client():
+    pass
+
+
+def test_component_to_tool():
+    chat_input = ChatInput()
+    tool = chat_input.to_tool()
+    assert tool.name == "ChatInput"
+    assert tool.description == "Get chat inputs from the Playground."
+    assert tool.component._id == chat_input._id