From 2061f8ca4121d77ca6be9d1b0e7813a08f75ff56 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Fri, 1 Nov 2024 17:19:54 -0400 Subject: [PATCH] Improve firecrawl tool --- .../firecrawl_crawl_website_tool.py | 8 ++++++-- .../firecrawl_scrape_website_tool.py | 8 ++++++-- .../firecrawl_search_tool/firecrawl_search_tool.py | 10 +++++++--- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/crewai_tools/tools/firecrawl_crawl_website_tool/firecrawl_crawl_website_tool.py b/src/crewai_tools/tools/firecrawl_crawl_website_tool/firecrawl_crawl_website_tool.py index 80a8392e9..a535b6c63 100644 --- a/src/crewai_tools/tools/firecrawl_crawl_website_tool/firecrawl_crawl_website_tool.py +++ b/src/crewai_tools/tools/firecrawl_crawl_website_tool/firecrawl_crawl_website_tool.py @@ -1,9 +1,13 @@ -from typing import Any, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, Optional, Type from pydantic import BaseModel, Field from crewai_tools.tools.base_tool import BaseTool +# Type checking import +if TYPE_CHECKING: + from firecrawl import FirecrawlApp + class FirecrawlCrawlWebsiteToolSchema(BaseModel): url: str = Field(description="Website URL") @@ -20,7 +24,7 @@ class FirecrawlCrawlWebsiteTool(BaseTool): description: str = "Crawl webpages using Firecrawl and return the contents" args_schema: Type[BaseModel] = FirecrawlCrawlWebsiteToolSchema api_key: Optional[str] = None - firecrawl: Optional[Any] = None + firecrawl: Optional["FirecrawlApp"] = None def __init__(self, api_key: Optional[str] = None, **kwargs): super().__init__(**kwargs) diff --git a/src/crewai_tools/tools/firecrawl_scrape_website_tool/firecrawl_scrape_website_tool.py b/src/crewai_tools/tools/firecrawl_scrape_website_tool/firecrawl_scrape_website_tool.py index 0de3335c7..ee8e592ca 100644 --- a/src/crewai_tools/tools/firecrawl_scrape_website_tool/firecrawl_scrape_website_tool.py +++ b/src/crewai_tools/tools/firecrawl_scrape_website_tool/firecrawl_scrape_website_tool.py @@ -1,9 +1,13 @@ -from typing import Any, Dict, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, Optional, Type from pydantic import BaseModel, Field from crewai_tools.tools.base_tool import BaseTool +# Type checking import +if TYPE_CHECKING: + from firecrawl import FirecrawlApp + class FirecrawlScrapeWebsiteToolSchema(BaseModel): url: str = Field(description="Website URL") @@ -24,7 +28,7 @@ class FirecrawlScrapeWebsiteTool(BaseTool): description: str = "Scrape webpages url using Firecrawl and return the contents" args_schema: Type[BaseModel] = FirecrawlScrapeWebsiteToolSchema api_key: Optional[str] = None - firecrawl: Optional[Any] = None + firecrawl: Optional["FirecrawlApp"] = None # Updated to use TYPE_CHECKING def __init__(self, api_key: Optional[str] = None, **kwargs): super().__init__(**kwargs) diff --git a/src/crewai_tools/tools/firecrawl_search_tool/firecrawl_search_tool.py b/src/crewai_tools/tools/firecrawl_search_tool/firecrawl_search_tool.py index ad92e2661..13c3b82ee 100644 --- a/src/crewai_tools/tools/firecrawl_search_tool/firecrawl_search_tool.py +++ b/src/crewai_tools/tools/firecrawl_search_tool/firecrawl_search_tool.py @@ -1,9 +1,13 @@ -from typing import Any, Dict, List, Optional, Type +from typing import TYPE_CHECKING, Any, Dict, Optional, Type from pydantic import BaseModel, Field from crewai_tools.tools.base_tool import BaseTool +# Type checking import +if TYPE_CHECKING: + from firecrawl import FirecrawlApp + class FirecrawlSearchToolSchema(BaseModel): query: str = Field(description="Search query") @@ -20,7 +24,7 @@ class FirecrawlSearchTool(BaseTool): description: str = "Search webpages using Firecrawl and return the results" args_schema: Type[BaseModel] = FirecrawlSearchToolSchema api_key: Optional[str] = None - firecrawl: Optional[Any] = None + firecrawl: Optional["FirecrawlApp"] = None def __init__(self, api_key: Optional[str] = None, **kwargs): super().__init__(**kwargs) @@ -45,4 +49,4 @@ class FirecrawlSearchTool(BaseTool): result_options = {} options = {"pageOptions": page_options, "resultOptions": result_options} - return self.firecrawl.search(query, options) + return self.firecrawl.search(query, **options)