From 728a86aab05da99677ed1d0401b622aac5af1366 Mon Sep 17 00:00:00 2001
From: Cristhian Zanforlin Lousa <cristhian.lousa@gmail.com>
Date: Tue, 3 Dec 2024 15:21:26 -0300
Subject: [PATCH] refactor: Move RetrieverToolComponent to langchain_utilities
 and add legacy support (#5003)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* ✨ (langchain_utilities): add RetrieverToolComponent to langchain_utilities package
♻️ (tools): remove RetrieverToolComponent from tools package to avoid duplication and improve organization

* style: fix import block un-sorted

* 🔧 (twoEdges.spec.ts): remove unnecessary click on "input-inspection-retriever" element to improve test clarity and efficiency

---------

Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com>
---
 .../langchain_utilities/__init__.py           |  2 ++
 .../langchain_utilities/retriever.py          | 32 +++++++++++++++++++
 .../tests/extended/features/twoEdges.spec.ts  |  1 -
 .../regression/generalBugs-shard-12.spec.ts   | 12 ++++++-
 4 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 src/backend/base/langflow/components/langchain_utilities/retriever.py

diff --git a/src/backend/base/langflow/components/langchain_utilities/__init__.py b/src/backend/base/langflow/components/langchain_utilities/__init__.py
index 819cc4af5..6154ae880 100644
--- a/src/backend/base/langflow/components/langchain_utilities/__init__.py
+++ b/src/backend/base/langflow/components/langchain_utilities/__init__.py
@@ -15,6 +15,7 @@ from .openai_tools import OpenAIToolsAgentComponent
 from .openapi import OpenAPIAgentComponent
 from .recursive_character import RecursiveCharacterTextSplitterComponent
 from .retrieval_qa import RetrievalQAComponent
+from .retriever import RetrieverToolComponent
 from .runnable_executor import RunnableExecComponent
 from .self_query import SelfQueryRetrieverComponent
 from .spider import SpiderTool
@@ -56,4 +57,5 @@ __all__ = [
     "VectorStoreInfoComponent",
     "VectorStoreRouterAgentComponent",
     "XMLAgentComponent",
+    "RetrieverToolComponent",
 ]
diff --git a/src/backend/base/langflow/components/langchain_utilities/retriever.py b/src/backend/base/langflow/components/langchain_utilities/retriever.py
new file mode 100644
index 000000000..c016ec0df
--- /dev/null
+++ b/src/backend/base/langflow/components/langchain_utilities/retriever.py
@@ -0,0 +1,32 @@
+from langchain_core.tools import create_retriever_tool
+
+from langflow.custom import CustomComponent
+from langflow.field_typing import BaseRetriever, Tool
+
+
+class RetrieverToolComponent(CustomComponent):
+    display_name = "RetrieverTool"
+    description = "Tool for interacting with retriever"
+    name = "RetrieverTool"
+    legacy = True
+    icon = "LangChain"
+
+    def build_config(self):
+        return {
+            "retriever": {
+                "display_name": "Retriever",
+                "info": "Retriever to interact with",
+                "type": BaseRetriever,
+                "input_types": ["Retriever"],
+            },
+            "name": {"display_name": "Name", "info": "Name of the tool"},
+            "description": {"display_name": "Description", "info": "Description of the tool"},
+        }
+
+    def build(self, retriever: BaseRetriever, name: str, description: str, **kwargs) -> Tool:
+        _ = kwargs
+        return create_retriever_tool(
+            retriever=retriever,
+            name=name,
+            description=description,
+        )
diff --git a/src/frontend/tests/extended/features/twoEdges.spec.ts b/src/frontend/tests/extended/features/twoEdges.spec.ts
index 495bbf866..b366a2e58 100644
--- a/src/frontend/tests/extended/features/twoEdges.spec.ts
+++ b/src/frontend/tests/extended/features/twoEdges.spec.ts
@@ -49,7 +49,6 @@ test(
     await page.getByTestId("zoom_out").click();
     await page.getByTestId("zoom_out").click();
 
-    await page.getByTestId("input-inspection-retriever").first().click();
     await page.getByText("Retriever", { exact: true }).first().isHidden();
     await page.getByTestId("icon-ChevronDown").last().isVisible();
     await page.getByTestId("icon-ChevronDown").last().click();
diff --git a/src/frontend/tests/extended/regression/generalBugs-shard-12.spec.ts b/src/frontend/tests/extended/regression/generalBugs-shard-12.spec.ts
index 68c2e45a4..ccba7f1ef 100644
--- a/src/frontend/tests/extended/regression/generalBugs-shard-12.spec.ts
+++ b/src/frontend/tests/extended/regression/generalBugs-shard-12.spec.ts
@@ -36,7 +36,17 @@ test(
 
     await page.waitForTimeout(1000);
 
-    let modelElement = await page.getByTestId("toolsRetrieverTool");
+    await page.getByTestId("sidebar-options-trigger").click();
+    await page
+      .getByTestId("sidebar-legacy-switch")
+      .isVisible({ timeout: 5000 });
+    await page.getByTestId("sidebar-legacy-switch").click();
+    await expect(page.getByTestId("sidebar-legacy-switch")).toBeChecked();
+    await page.getByTestId("sidebar-options-trigger").click();
+
+    let modelElement = await page.getByTestId(
+      "langchain_utilitiesRetrieverTool",
+    );
     let targetElement = await page.locator('//*[@id="react-flow-id"]');
     await modelElement.dragTo(targetElement);