diff --git a/backend/app/agent/factory/developer.py b/backend/app/agent/factory/developer.py index 96359e055..43250e964 100644 --- a/backend/app/agent/factory/developer.py +++ b/backend/app/agent/factory/developer.py @@ -20,11 +20,13 @@ from app.agent.agent_model import agent_model from app.agent.listen_chat_agent import logger from app.agent.prompt import DEVELOPER_SYS_PROMPT +from app.agent.toolkit.craw4ai_toolkit import Crawl4AIToolkit from app.agent.toolkit.human_toolkit import HumanToolkit # TODO: Remove NoteTakingToolkit and use TerminalToolkit instead from app.agent.toolkit.note_taking_toolkit import NoteTakingToolkit from app.agent.toolkit.screenshot_toolkit import ScreenshotToolkit +from app.agent.toolkit.search_toolkit import SearchToolkit from app.agent.toolkit.terminal_toolkit import TerminalToolkit from app.agent.toolkit.web_deploy_toolkit import WebDeployToolkit from app.agent.utils import NOW_STR @@ -70,6 +72,15 @@ async def developer_agent(options: Chat): ) terminal_toolkit = message_integration.register_toolkits(terminal_toolkit) + search_tools = SearchToolkit.get_can_use_tools(options.project_id) + if search_tools: + search_tools = message_integration.register_functions(search_tools) + else: + search_tools = [] + + crawl_toolkit = Crawl4AIToolkit(options.project_id) + crawl_toolkit = message_integration.register_toolkits(crawl_toolkit) + tools = [ *HumanToolkit.get_can_use_tools( options.project_id, Agents.developer_agent @@ -78,6 +89,8 @@ async def developer_agent(options: Chat): *web_deploy_toolkit.get_tools(), *terminal_toolkit.get_tools(), *screenshot_toolkit.get_tools(), + *search_tools, + *crawl_toolkit.get_tools(), ] system_message = DEVELOPER_SYS_PROMPT.format( platform_system=platform.system(), @@ -99,5 +112,7 @@ async def developer_agent(options: Chat): TerminalToolkit.toolkit_name(), NoteTakingToolkit.toolkit_name(), WebDeployToolkit.toolkit_name(), + SearchToolkit.toolkit_name(), + Crawl4AIToolkit.toolkit_name(), ], ) diff --git a/backend/app/agent/factory/document.py b/backend/app/agent/factory/document.py index 0d3b2c897..cb487fdc4 100644 --- a/backend/app/agent/factory/document.py +++ b/backend/app/agent/factory/document.py @@ -19,6 +19,7 @@ from app.agent.agent_model import agent_model from app.agent.listen_chat_agent import logger from app.agent.prompt import DOCUMENT_SYS_PROMPT +from app.agent.toolkit.craw4ai_toolkit import Crawl4AIToolkit from app.agent.toolkit.excel_toolkit import ExcelToolkit from app.agent.toolkit.file_write_toolkit import FileToolkit from app.agent.toolkit.google_drive_mcp_toolkit import GoogleDriveMCPToolkit @@ -28,6 +29,7 @@ # TODO: Remove NoteTakingToolkit and use TerminalToolkit instead from app.agent.toolkit.note_taking_toolkit import NoteTakingToolkit from app.agent.toolkit.pptx_toolkit import PPTXToolkit +from app.agent.toolkit.search_toolkit import SearchToolkit from app.agent.toolkit.terminal_toolkit import TerminalToolkit from app.agent.utils import NOW_STR from app.model.chat import Chat @@ -82,6 +84,15 @@ async def document_agent(options: Chat): options.project_id, options.get_bun_env() ) + search_tools = SearchToolkit.get_can_use_tools(options.project_id) + if search_tools: + search_tools = message_integration.register_functions(search_tools) + else: + search_tools = [] + + crawl_toolkit = Crawl4AIToolkit(options.project_id) + crawl_toolkit = message_integration.register_toolkits(crawl_toolkit) + tools = [ *file_write_toolkit.get_tools(), *pptx_toolkit.get_tools(), @@ -93,6 +104,8 @@ async def document_agent(options: Chat): *note_toolkit.get_tools(), *terminal_toolkit.get_tools(), *google_drive_tools, + *search_tools, + *crawl_toolkit.get_tools(), ] system_message = DOCUMENT_SYS_PROMPT.format( platform_system=platform.system(), @@ -118,5 +131,7 @@ async def document_agent(options: Chat): NoteTakingToolkit.toolkit_name(), TerminalToolkit.toolkit_name(), GoogleDriveMCPToolkit.toolkit_name(), + SearchToolkit.toolkit_name(), + Crawl4AIToolkit.toolkit_name(), ], ) diff --git a/backend/app/agent/factory/multi_modal.py b/backend/app/agent/factory/multi_modal.py index fea7b1c29..374431b65 100644 --- a/backend/app/agent/factory/multi_modal.py +++ b/backend/app/agent/factory/multi_modal.py @@ -22,6 +22,7 @@ from app.agent.listen_chat_agent import logger from app.agent.prompt import MULTI_MODAL_SYS_PROMPT from app.agent.toolkit.audio_analysis_toolkit import AudioAnalysisToolkit +from app.agent.toolkit.craw4ai_toolkit import Crawl4AIToolkit from app.agent.toolkit.human_toolkit import HumanToolkit from app.agent.toolkit.image_analysis_toolkit import ImageAnalysisToolkit @@ -75,6 +76,16 @@ def multi_modal_agent(options: Chat): working_directory=working_directory, ) note_toolkit = message_integration.register_toolkits(note_toolkit) + + search_tools = SearchToolkit.get_can_use_tools(options.project_id) + if search_tools: + search_tools = message_integration.register_functions(search_tools) + else: + search_tools = [] + + crawl_toolkit = Crawl4AIToolkit(options.project_id) + crawl_toolkit = message_integration.register_toolkits(crawl_toolkit) + tools = [ *video_download_toolkit.get_tools(), *image_analysis_toolkit.get_tools(), @@ -83,6 +94,8 @@ def multi_modal_agent(options: Chat): ), *terminal_toolkit.get_tools(), *note_toolkit.get_tools(), + *search_tools, + *crawl_toolkit.get_tools(), ] if options.is_cloud(): # TODO: check llm has this model @@ -147,5 +160,6 @@ def multi_modal_agent(options: Chat): TerminalToolkit.toolkit_name(), NoteTakingToolkit.toolkit_name(), SearchToolkit.toolkit_name(), + Crawl4AIToolkit.toolkit_name(), ], ) diff --git a/backend/app/agent/factory/social_media.py b/backend/app/agent/factory/social_media.py index 649613bb6..ca9b5bce7 100644 --- a/backend/app/agent/factory/social_media.py +++ b/backend/app/agent/factory/social_media.py @@ -16,6 +16,7 @@ from app.agent.agent_model import agent_model from app.agent.listen_chat_agent import logger from app.agent.prompt import SOCIAL_MEDIA_SYS_PROMPT +from app.agent.toolkit.craw4ai_toolkit import Crawl4AIToolkit from app.agent.toolkit.google_calendar_toolkit import GoogleCalendarToolkit from app.agent.toolkit.google_gmail_mcp_toolkit import GoogleGmailMCPToolkit from app.agent.toolkit.human_toolkit import HumanToolkit @@ -25,6 +26,7 @@ from app.agent.toolkit.note_taking_toolkit import NoteTakingToolkit from app.agent.toolkit.notion_mcp_toolkit import NotionMCPToolkit from app.agent.toolkit.reddit_toolkit import RedditToolkit +from app.agent.toolkit.search_toolkit import SearchToolkit from app.agent.toolkit.terminal_toolkit import TerminalToolkit from app.agent.toolkit.twitter_toolkit import TwitterToolkit from app.agent.toolkit.whatsapp_toolkit import WhatsAppToolkit @@ -70,6 +72,8 @@ async def social_media_agent(options: Chat): Agents.social_media_agent, working_directory=working_directory, ).get_tools(), + *SearchToolkit.get_can_use_tools(options.project_id), + *Crawl4AIToolkit(options.project_id).get_tools(), # *DiscordToolkit(options.project_id).get_tools(), # *GoogleSuiteToolkit(options.project_id).get_tools(), ] @@ -94,5 +98,7 @@ async def social_media_agent(options: Chat): HumanToolkit.toolkit_name(), TerminalToolkit.toolkit_name(), NoteTakingToolkit.toolkit_name(), + SearchToolkit.toolkit_name(), + Crawl4AIToolkit.toolkit_name(), ], ) diff --git a/backend/tests/app/agent/factory/test_developer.py b/backend/tests/app/agent/factory/test_developer.py index 0a4e55ce7..420930d88 100644 --- a/backend/tests/app/agent/factory/test_developer.py +++ b/backend/tests/app/agent/factory/test_developer.py @@ -42,6 +42,8 @@ async def test_developer_agent_creation(sample_chat_data): patch(f"{_mod}.WebDeployToolkit") as mock_web_toolkit, patch(f"{_mod}.ScreenshotToolkit") as mock_screenshot_toolkit, patch(f"{_mod}.TerminalToolkit") as mock_terminal_toolkit, + patch(f"{_mod}.SearchToolkit") as mock_search_toolkit, + patch(f"{_mod}.Crawl4AIToolkit") as mock_crawl_toolkit, patch(f"{_mod}.ToolkitMessageIntegration"), ): # Mock all toolkit instances @@ -50,6 +52,8 @@ async def test_developer_agent_creation(sample_chat_data): mock_web_toolkit.return_value.get_tools.return_value = [] mock_screenshot_toolkit.return_value.get_tools.return_value = [] mock_terminal_toolkit.return_value.get_tools.return_value = [] + mock_search_toolkit.get_can_use_tools.return_value = [] + mock_crawl_toolkit.return_value.get_tools.return_value = [] mock_agent = MagicMock() mock_agent_model.return_value = mock_agent @@ -88,6 +92,8 @@ async def test_developer_agent_with_multiple_toolkits(sample_chat_data): patch(f"{_mod}.WebDeployToolkit") as mock_web_toolkit, patch(f"{_mod}.ScreenshotToolkit") as mock_screenshot_toolkit, patch(f"{_mod}.TerminalToolkit") as mock_terminal_toolkit, + patch(f"{_mod}.SearchToolkit") as mock_search_toolkit, + patch(f"{_mod}.Crawl4AIToolkit") as mock_crawl_toolkit, patch(f"{_mod}.ToolkitMessageIntegration"), ): # Mock all toolkit instances @@ -96,6 +102,8 @@ async def test_developer_agent_with_multiple_toolkits(sample_chat_data): mock_web_toolkit.return_value.get_tools.return_value = [] mock_screenshot_toolkit.return_value.get_tools.return_value = [] mock_terminal_toolkit.return_value.get_tools.return_value = [] + mock_search_toolkit.get_can_use_tools.return_value = [] + mock_crawl_toolkit.return_value.get_tools.return_value = [] mock_agent = MagicMock() mock_agent_model.return_value = mock_agent diff --git a/backend/tests/app/agent/factory/test_document.py b/backend/tests/app/agent/factory/test_document.py index 034a1696c..21ea994ba 100644 --- a/backend/tests/app/agent/factory/test_document.py +++ b/backend/tests/app/agent/factory/test_document.py @@ -45,6 +45,8 @@ async def test_document_agent_creation(sample_chat_data): patch(f"{_mod}.NoteTakingToolkit") as mock_note_toolkit, patch(f"{_mod}.TerminalToolkit") as mock_terminal_toolkit, patch(f"{_mod}.GoogleDriveMCPToolkit") as mock_gdrive_toolkit, + patch(f"{_mod}.SearchToolkit") as mock_search_toolkit, + patch(f"{_mod}.Crawl4AIToolkit") as mock_crawl_toolkit, patch(f"{_mod}.ToolkitMessageIntegration"), ): # Mock all toolkit instances @@ -56,6 +58,8 @@ async def test_document_agent_creation(sample_chat_data): mock_note_toolkit.return_value.get_tools.return_value = [] mock_terminal_toolkit.return_value.get_tools.return_value = [] mock_gdrive_toolkit.get_can_use_tools = AsyncMock(return_value=[]) + mock_search_toolkit.get_can_use_tools.return_value = [] + mock_crawl_toolkit.return_value.get_tools.return_value = [] mock_agent = MagicMock() mock_agent_model.return_value = mock_agent diff --git a/backend/tests/app/agent/factory/test_multi_modal.py b/backend/tests/app/agent/factory/test_multi_modal.py index 308def0ab..887e3ab14 100644 --- a/backend/tests/app/agent/factory/test_multi_modal.py +++ b/backend/tests/app/agent/factory/test_multi_modal.py @@ -44,6 +44,7 @@ def test_multi_modal_agent_creation(sample_chat_data): patch(f"{_mod}.TerminalToolkit") as mock_terminal_toolkit, patch(f"{_mod}.NoteTakingToolkit") as mock_note_toolkit, patch(f"{_mod}.SearchToolkit") as mock_search_toolkit, + patch(f"{_mod}.Crawl4AIToolkit") as mock_crawl_toolkit, patch(f"{_mod}.ToolkitMessageIntegration"), ): # Mock all toolkit instances @@ -54,7 +55,8 @@ def test_multi_modal_agent_creation(sample_chat_data): mock_audio_toolkit.return_value.get_tools.return_value = [] mock_terminal_toolkit.return_value.get_tools.return_value = [] mock_note_toolkit.return_value.get_tools.return_value = [] - mock_search_toolkit.return_value.get_tools.return_value = [] + mock_search_toolkit.get_can_use_tools.return_value = [] + mock_crawl_toolkit.return_value.get_tools.return_value = [] mock_agent = MagicMock() mock_agent_model.return_value = mock_agent diff --git a/backend/tests/app/agent/factory/test_social_media.py b/backend/tests/app/agent/factory/test_social_media.py index 56b3ee426..30f7c4f17 100644 --- a/backend/tests/app/agent/factory/test_social_media.py +++ b/backend/tests/app/agent/factory/test_social_media.py @@ -47,6 +47,8 @@ async def test_social_media_agent_creation(sample_chat_data): patch(f"{mod}.HumanToolkit") as mock_human_toolkit, patch(f"{mod}.TerminalToolkit") as mock_terminal_toolkit, patch(f"{mod}.NoteTakingToolkit") as mock_note_toolkit, + patch(f"{mod}.SearchToolkit") as mock_search_toolkit, + patch(f"{mod}.Crawl4AIToolkit") as mock_crawl_toolkit, ): # Mock all toolkit instances mock_whatsapp_toolkit.get_can_use_tools.return_value = [] @@ -59,6 +61,8 @@ async def test_social_media_agent_creation(sample_chat_data): mock_human_toolkit.get_can_use_tools.return_value = [] mock_terminal_toolkit.return_value.get_tools.return_value = [] mock_note_toolkit.return_value.get_tools.return_value = [] + mock_search_toolkit.get_can_use_tools.return_value = [] + mock_crawl_toolkit.return_value.get_tools.return_value = [] mock_agent = MagicMock() mock_agent_model.return_value = mock_agent