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 <edwin.jose@datastax.com>
This commit is contained in:
Raphael Valdetaro 2025-01-17 18:13:39 -03:00 • committed by GitHub
commit 88111c3b34
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 142 additions and 1 deletions

View file

@ -16,6 +16,7 @@ from .python_repl import PythonREPLToolComponent
from .search_api import SearchAPIComponent from .search_api import SearchAPIComponent
from .searxng import SearXNGToolComponent from .searxng import SearXNGToolComponent
from .serp_api import SerpAPIComponent from .serp_api import SerpAPIComponent
from .tavily import TavilySearchComponent
from .tavily_search import TavilySearchToolComponent from .tavily_search import TavilySearchToolComponent
from .wikidata_api import WikidataAPIComponent from .wikidata_api import WikidataAPIComponent
from .wikipedia_api import WikipediaAPIComponent from .wikipedia_api import WikipediaAPIComponent
@ -45,6 +46,7 @@ __all__ = [
"SearXNGToolComponent", "SearXNGToolComponent",
"SearchAPIComponent", "SearchAPIComponent",
"SerpAPIComponent", "SerpAPIComponent",
"TavilySearchComponent",
"TavilySearchToolComponent", "TavilySearchToolComponent",
"WikidataAPIComponent", "WikidataAPIComponent",
"WikipediaAPIComponent", "WikipediaAPIComponent",

View file

@ -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)

View file

@ -32,7 +32,7 @@ class TavilySearchSchema(BaseModel):
class TavilySearchToolComponent(LCToolComponent): 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, \ 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. aimed at efficient, quick, and persistent search results. It can be used independently or as an agent tool.
@ -41,6 +41,7 @@ Note: Check 'Advanced' for all options.
icon = "TavilyIcon" icon = "TavilyIcon"
name = "TavilyAISearch" name = "TavilyAISearch"
documentation = "https://docs.tavily.com/" documentation = "https://docs.tavily.com/"
legacy = True
inputs = [ inputs = [
SecretStrInput( SecretStrInput(