feat(tavily): enhance TavilyExtractorTool and TavilySearchTool with additional parameters and improved error handling

Author: prdai

crewai_tools/tools/tavily_extractor_tool/tavily_extractor_tool.py

@@ -23,18 +23,6 @@ class TavilyExtractorToolSchema(BaseModel):
         ...,
         description="The URL(s) to extract data from. Can be a single URL or a list of URLs.",
     )
-    include_images: Optional[bool] = Field(
-        default=False,
-        description="Whether to include images in the extraction.",
-    )
-    extract_depth: Literal["basic", "advanced"] = Field(
-        default="basic",
-        description="The depth of extraction. 'basic' for basic extraction, 'advanced' for advanced extraction.",
-    )
-    timeout: int = Field(
-        default=60,
-        description="The timeout for the extraction request in seconds.",
-    )
 
 
 class TavilyExtractorTool(BaseTool):
@@ -49,26 +37,41 @@ class TavilyExtractorTool(BaseTool):
         args_schema: The schema for the tool's arguments.
         api_key: The Tavily API key.
         proxies: Optional proxies for the API requests.
+        include_images: Whether to include images in the extraction.
+        extract_depth: The depth of extraction.
+        timeout: The timeout for the extraction request in seconds.
     """
 
-    model_config = {}
-    client: TavilyClient = None
-    async_client: AsyncTavilyClient = None
+    model_config = {"arbitrary_types_allowed": True}
+    client: Optional[TavilyClient] = None
+    async_client: Optional[AsyncTavilyClient] = None
     name: str = "TavilyExtractorTool"
     description: str = (
         "Extracts content from one or more web pages using the Tavily API. Returns structured data."
     )
     args_schema: Type[BaseModel] = TavilyExtractorToolSchema
     api_key: Optional[str] = Field(
-        default=os.getenv("TAVILY_API_KEY"),
+        default_factory=lambda: os.getenv("TAVILY_API_KEY"),
         description="The Tavily API key. If not provided, it will be loaded from the environment variable TAVILY_API_KEY.",
     )
     proxies: Optional[dict[str, str]] = Field(
         default=None,
         description="Optional proxies to use for the Tavily API requests.",
     )
+    include_images: bool = Field(
+        default=False,
+        description="Whether to include images in the extraction.",
+    )
+    extract_depth: Literal["basic", "advanced"] = Field(
+        default="basic",
+        description="The depth of extraction. 'basic' for basic extraction, 'advanced' for advanced extraction.",
+    )
+    timeout: int = Field(
+        default=60,
+        description="The timeout for the extraction request in seconds.",
+    )
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any):
         """
         Initializes the TavilyExtractorTool.
 
@@ -82,75 +85,84 @@ def __init__(self, **kwargs):
                 api_key=self.api_key, proxies=self.proxies
             )
         else:
-            import click
+            try:
+                import click
+                import subprocess
+            except ImportError:
+                raise ImportError(
+                    "The 'tavily-python' package is required. 'click' and 'subprocess' are also needed to assist with installation if the package is missing. "
+                    "Please install 'tavily-python' manually (e.g., 'pip install tavily-python') and ensure 'click' and 'subprocess' are available."
+                )
 
             if click.confirm(
-                "The 'tavily-python' package is required to use the TavilyExtractorTool. "
-                "Would you like to install it?"
+                "You are missing the 'tavily-python' package, which is required for TavilyExtractorTool. Would you like to install it?"
             ):
-                import subprocess
-
-                subprocess.run(["uv", "add", "tavily-python"], check=True)
+                try:
+                    subprocess.run(["pip", "install", "tavily-python"], check=True)
+                    raise ImportError(
+                        "'tavily-python' has been installed. Please restart your Python application to use the TavilyExtractorTool."
+                    )
+                except subprocess.CalledProcessError as e:
+                    raise ImportError(
+                        f"Attempted to install 'tavily-python' but failed: {e}. "
+                        f"Please install it manually to use the TavilyExtractorTool."
+                    )
             else:
                 raise ImportError(
                     "The 'tavily-python' package is required to use the TavilyExtractorTool. "
-                    "Please install it with: uv add tavily-python"
+                    "Please install it with: pip install tavily-python"
                 )
 
     def _run(
         self,
         urls: Union[List[str], str],
-        include_images: bool = False,
-        extract_depth: Literal["basic", "advanced"] = "basic",
-        timeout: int = 60,
     ) -> str:
         """
         Synchronously extracts content from the given URL(s).
 
         Args:
             urls: The URL(s) to extract data from.
-            include_images: Whether to include images in the extraction.
-            extract_depth: The depth of extraction ('basic' or 'advanced').
-            timeout: The timeout for the request in seconds.
 
         Returns:
             A JSON string containing the extracted data.
         """
+        if not self.client:
+            raise ValueError(
+                "Tavily client is not initialized. Ensure 'tavily-python' is installed and API key is set."
+            )
+
         return json.dumps(
             self.client.extract(
                 urls=urls,
-                extract_depth=extract_depth,
-                include_images=include_images,
-                timeout=timeout,
+                extract_depth=self.extract_depth,
+                include_images=self.include_images,
+                timeout=self.timeout,
             ),
             indent=2,
         )
 
     async def _arun(
         self,
         urls: Union[List[str], str],
-        include_images: bool = False,
-        extract_depth: Literal["basic", "advanced"] = "basic",
-        timeout: int = 60,
     ) -> str:
         """
         Asynchronously extracts content from the given URL(s).
 
         Args:
             urls: The URL(s) to extract data from.
-            include_images: Whether to include images in the extraction.
-            extract_depth: The depth of extraction ('basic' or 'advanced').
-            timeout: The timeout for the request in seconds.
 
         Returns:
             A JSON string containing the extracted data.
         """
-        return json.dumps(
-            self.async_client.extract(
-                urls=urls,
-                extract_depth=extract_depth,
-                include_images=include_images,
-                timeout=timeout,
-            ),
-            indent=2,
+        if not self.async_client:
+            raise ValueError(
+                "Tavily async client is not initialized. Ensure 'tavily-python' is installed and API key is set."
+            )
+
+        results = await self.async_client.extract(
+            urls=urls,
+            extract_depth=self.extract_depth,
+            include_images=self.include_images,
+            timeout=self.timeout,
         )
+        return json.dumps(results, indent=2)