Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

refactor(component): Convert Tavily Search to standard component pattern #5430

Open
wants to merge 5 commits into
base: main
Choose a base branch
from
+65 −130
Open

refactor(component): Convert Tavily Search to standard component pattern #5430

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
7bf2bd4
refactor(components): Convert Tavily Search to standard component pat…
raphaelchristi Dec 24, 2024
d035895
[autofix.ci] apply automated fixes
autofix-ci[bot] Dec 24, 2024
cac8d97
fix: improve error handling in tavily component
raphaelchristi Dec 24, 2024
28d3b1f
fix: Fix linting error in __init__.py
raphaelchristi Dec 24, 2024
24201c0
Merge branch 'main' into refactor/tavily-standard-component
raphaelchristi Dec 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/backend/base/langflow/components/tools/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
from .search_api import SearchAPIComponent
from .searxng import SearXNGToolComponent
from .serp_api import SerpAPIComponent
from .tavily_search import TavilySearchToolComponent
from .tavily_search import TavilyComponent
from .wikidata_api import WikidataAPIComponent
from .wikipedia_api import WikipediaAPIComponent
from .wolfram_alpha_api import WolframAlphaAPIComponent
Expand All @@ -41,7 +41,7 @@
"SearXNGToolComponent",
"SearchAPIComponent",
"SerpAPIComponent",
"TavilySearchToolComponent",
"TavilyComponent",
"WikidataAPIComponent",
"WikipediaAPIComponent",
"WolframAlphaAPIComponent",
Expand Down
191 changes: 63 additions & 128 deletions src/backend/base/langflow/components/tools/tavily_search.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,45 +1,19 @@
from enum import Enum

import httpx
from langchain.tools import StructuredTool
from langchain_core.tools import ToolException
from loguru import logger
from pydantic import BaseModel, Field

from langflow.base.langchain_utilities.model import LCToolComponent
from langflow.field_typing import Tool
from langflow.inputs import BoolInput, DropdownInput, IntInput, MessageTextInput, SecretStrInput
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 TavilySearchDepth(Enum):
BASIC = "basic"
ADVANCED = "advanced"


class TavilySearchTopic(Enum):
GENERAL = "general"
NEWS = "news"


class TavilySearchSchema(BaseModel):
query: str = Field(..., description="The search query you want to execute with Tavily.")
search_depth: TavilySearchDepth = Field(TavilySearchDepth.BASIC, description="The depth of the search.")
topic: TavilySearchTopic = Field(TavilySearchTopic.GENERAL, description="The category of the search.")
max_results: int = Field(5, description="The maximum number of search results to return.")
include_images: bool = Field(default=False, description="Include a list of query-related images in the response.")
include_answer: bool = Field(default=False, description="Include a short answer to original query.")


class TavilySearchToolComponent(LCToolComponent):
class TavilyComponent(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. It can be used independently or as an agent tool.

Note: Check 'Advanced' for all options.
"""
aimed at efficient, quick, and persistent search results."""
icon = "TavilyIcon"
name = "TavilyAISearch"
name = "TavilyAI"
documentation = "https://docs.tavily.com/"

inputs = [
Expand All @@ -53,21 +27,22 @@ class TavilySearchToolComponent(LCToolComponent):
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=list(TavilySearchDepth),
value=TavilySearchDepth.ADVANCED,
options=["basic", "advanced"],
value="advanced",
advanced=True,
),
DropdownInput(
name="topic",
display_name="Search Topic",
info="The category of the search.",
options=list(TavilySearchTopic),
value=TavilySearchTopic.GENERAL,
options=["general", "news"],
value="general",
advanced=True,
),
IntInput(
Expand All @@ -93,63 +68,12 @@ class TavilySearchToolComponent(LCToolComponent):
),
]

def run_model(self) -> list[Data]:
# Convert string values to enum instances with validation
try:
search_depth_enum = (
self.search_depth
if isinstance(self.search_depth, TavilySearchDepth)
else TavilySearchDepth(str(self.search_depth).lower())
)
except ValueError as e:
error_message = f"Invalid search depth value: {e!s}"
self.status = error_message
return [Data(data={"error": error_message})]

try:
topic_enum = (
self.topic if isinstance(self.topic, TavilySearchTopic) else TavilySearchTopic(str(self.topic).lower())
)
except ValueError as e:
error_message = f"Invalid topic value: {e!s}"
self.status = error_message
return [Data(data={"error": error_message})]

return self._tavily_search(
self.query,
search_depth=search_depth_enum,
topic=topic_enum,
max_results=self.max_results,
include_images=self.include_images,
include_answer=self.include_answer,
)

def build_tool(self) -> Tool:
return StructuredTool.from_function(
name="tavily_search",
description="Perform a web search using the Tavily API.",
func=self._tavily_search,
args_schema=TavilySearchSchema,
)

def _tavily_search(
self,
query: str,
*,
search_depth: TavilySearchDepth = TavilySearchDepth.BASIC,
topic: TavilySearchTopic = TavilySearchTopic.GENERAL,
max_results: int = 5,
include_images: bool = False,
include_answer: bool = False,
) -> list[Data]:
# Validate enum values
if not isinstance(search_depth, TavilySearchDepth):
msg = f"Invalid search_depth value: {search_depth}"
raise TypeError(msg)
if not isinstance(topic, TavilySearchTopic):
msg = f"Invalid topic value: {topic}"
raise TypeError(msg)
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 = {
Expand All @@ -158,12 +82,12 @@ def _tavily_search(
}
payload = {
"api_key": self.api_key,
"query": query,
"search_depth": search_depth.value,
"topic": topic.value,
"max_results": max_results,
"include_images": include_images,
"include_answer": include_answer,
"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:
Expand All @@ -172,34 +96,45 @@ def _tavily_search(
response.raise_for_status()
search_results = response.json()

data_results = [
Data(
data={
"title": result.get("title"),
"url": result.get("url"),
"content": result.get("content"),
"score": result.get("score"),
}
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"),
},
)
)
for result in search_results.get("results", [])
]

if include_answer and search_results.get("answer"):
data_results.insert(0, Data(data={"answer": search_results["answer"]}))

if include_images and search_results.get("images"):
data_results.append(Data(data={"images": search_results["images"]}))

self.status = data_results # type: ignore[assignment]

except httpx.HTTPStatusError as e:
error_message = f"HTTP error: {e.response.status_code} - {e.response.text}"
logger.debug(error_message)
self.status = error_message
raise ToolException(error_message) from e
except Exception as e:
error_message = f"Unexpected error: {e}"
logger.opt(exception=True).debug("Error running Tavily Search")
self.status = error_message
raise ToolException(error_message) from e
return data_results
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)
Loading