From 88111c3b344bd81ce319d08299baa4ca54fa860b Mon Sep 17 00:00:00 2001 From: Raphael Valdetaro <79842132+raphaelchristi@users.noreply.github.com> Date: Fri, 17 Jan 2025 18:13:39 -0300 Subject: [PATCH] refactor(component): Convert Tavily Search to standard component pattern (#5430) * refactor(components): Convert Tavily Search to standard component pattern * [autofix.ci] apply automated fixes * fix: improve error handling in tavily component * fix: Fix linting error in __init__.py * fix: rename TavilyComponent to TavilySearchToolComponent maintaining backward compatibility * fix: rename component class to avoid conflict with legacy version --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Edwin Jose --- .../langflow/components/tools/__init__.py | 2 + .../base/langflow/components/tools/tavily.py | 138 ++++++++++++++++++ .../components/tools/tavily_search.py | 3 +- 3 files changed, 142 insertions(+), 1 deletion(-) create mode 100644 src/backend/base/langflow/components/tools/tavily.py diff --git a/src/backend/base/langflow/components/tools/__init__.py b/src/backend/base/langflow/components/tools/__init__.py index 55c99c36708a..90aee6c3d219 100644 --- a/src/backend/base/langflow/components/tools/__init__.py +++ b/src/backend/base/langflow/components/tools/__init__.py @@ -16,6 +16,7 @@ from .search_api import SearchAPIComponent from .searxng import SearXNGToolComponent from .serp_api import SerpAPIComponent +from .tavily import TavilySearchComponent from .tavily_search import TavilySearchToolComponent from .wikidata_api import WikidataAPIComponent from .wikipedia_api import WikipediaAPIComponent @@ -45,6 +46,7 @@ "SearXNGToolComponent", "SearchAPIComponent", "SerpAPIComponent", + "TavilySearchComponent", "TavilySearchToolComponent", "WikidataAPIComponent", "WikipediaAPIComponent", diff --git a/src/backend/base/langflow/components/tools/tavily.py b/src/backend/base/langflow/components/tools/tavily.py new file mode 100644 index 000000000000..7a64c52b248e --- /dev/null +++ b/src/backend/base/langflow/components/tools/tavily.py @@ -0,0 +1,138 @@ +import httpx +from loguru import logger + +from langflow.custom import Component +from langflow.helpers.data import data_to_text +from langflow.io import BoolInput, DropdownInput, IntInput, MessageTextInput, Output, SecretStrInput +from langflow.schema import Data +from langflow.schema.message import Message + + +class TavilySearchComponent(Component): + display_name = "Tavily AI Search" + description = """**Tavily AI** is a search engine optimized for LLMs and RAG, \ + aimed at efficient, quick, and persistent search results.""" + icon = "TavilyIcon" + + inputs = [ + SecretStrInput( + name="api_key", + display_name="Tavily API Key", + required=True, + info="Your Tavily API Key.", + ), + MessageTextInput( + name="query", + display_name="Search Query", + info="The search query you want to execute with Tavily.", + tool_mode=True, + ), + DropdownInput( + name="search_depth", + display_name="Search Depth", + info="The depth of the search.", + options=["basic", "advanced"], + value="advanced", + advanced=True, + ), + DropdownInput( + name="topic", + display_name="Search Topic", + info="The category of the search.", + options=["general", "news"], + value="general", + advanced=True, + ), + IntInput( + name="max_results", + display_name="Max Results", + info="The maximum number of search results to return.", + value=5, + advanced=True, + ), + BoolInput( + name="include_images", + display_name="Include Images", + info="Include a list of query-related images in the response.", + value=True, + advanced=True, + ), + BoolInput( + name="include_answer", + display_name="Include Answer", + info="Include a short answer to original query.", + value=True, + advanced=True, + ), + ] + + outputs = [ + Output(display_name="Data", name="data", method="fetch_content"), + Output(display_name="Text", name="text", method="fetch_content_text"), + ] + + def fetch_content(self) -> list[Data]: + try: + url = "https://api.tavily.com/search" + headers = { + "content-type": "application/json", + "accept": "application/json", + } + payload = { + "api_key": self.api_key, + "query": self.query, + "search_depth": self.search_depth, + "topic": self.topic, + "max_results": self.max_results, + "include_images": self.include_images, + "include_answer": self.include_answer, + } + + with httpx.Client() as client: + response = client.post(url, json=payload, headers=headers) + + response.raise_for_status() + search_results = response.json() + + data_results = [] + + if self.include_answer and search_results.get("answer"): + data_results.append(Data(text=search_results["answer"])) + + for result in search_results.get("results", []): + content = result.get("content", "") + data_results.append( + Data( + text=content, + data={ + "title": result.get("title"), + "url": result.get("url"), + "content": content, + "score": result.get("score"), + }, + ) + ) + + if self.include_images and search_results.get("images"): + data_results.append(Data(text="Images found", data={"images": search_results["images"]})) + except httpx.HTTPStatusError as exc: + error_message = f"HTTP error occurred: {exc.response.status_code} - {exc.response.text}" + logger.error(error_message) + return [Data(text=error_message, data={"error": error_message})] + except httpx.RequestError as exc: + error_message = f"Request error occurred: {exc}" + logger.error(error_message) + return [Data(text=error_message, data={"error": error_message})] + except ValueError as exc: + error_message = f"Invalid response format: {exc}" + logger.error(error_message) + return [Data(text=error_message, data={"error": error_message})] + else: + self.status = data_results + return data_results + + def fetch_content_text(self) -> Message: + data = self.fetch_content() + result_string = data_to_text("{text}", data) + self.status = result_string + return Message(text=result_string) diff --git a/src/backend/base/langflow/components/tools/tavily_search.py b/src/backend/base/langflow/components/tools/tavily_search.py index 21496dd30c7c..da755a0bbfd8 100644 --- a/src/backend/base/langflow/components/tools/tavily_search.py +++ b/src/backend/base/langflow/components/tools/tavily_search.py @@ -32,7 +32,7 @@ class TavilySearchSchema(BaseModel): class TavilySearchToolComponent(LCToolComponent): - display_name = "Tavily AI Search" + display_name = "Tavily AI Search [DEPRECATED]" description = """**Tavily AI** is a search engine optimized for LLMs and RAG, \ aimed at efficient, quick, and persistent search results. It can be used independently or as an agent tool. @@ -41,6 +41,7 @@ class TavilySearchToolComponent(LCToolComponent): icon = "TavilyIcon" name = "TavilyAISearch" documentation = "https://docs.tavily.com/" + legacy = True inputs = [ SecretStrInput(