diff --git a/metagpt/actions/research.py b/metagpt/actions/research.py index c47a77bdd..5057c3d3a 100644 --- a/metagpt/actions/research.py +++ b/metagpt/actions/research.py @@ -85,7 +85,7 @@ class CollectLinks(Action): llm: BaseGPTAPI = Field(default_factory=LLM) desc: str = "Collect links from a search engine." search_engine: SearchEngine = Field(default_factory=SearchEngine) - rank_func: Union[Callable[[list[str]], None], None] = None + rank_func: Optional[Callable[[list[str]], None]] = None async def run( self, @@ -180,18 +180,18 @@ class WebBrowseAndSummarize(Action): llm: BaseGPTAPI = Field(default_factory=LLM) desc: str = "Explore the web and provide summaries of articles and webpages." browse_func: Union[Callable[[list[str]], None], None] = None - web_browser_engine: WebBrowserEngine = Field( - default_factory=lambda: WebBrowserEngine( - engine=WebBrowserEngineType.CUSTOM if WebBrowseAndSummarize.browse_func else None, - run_func=WebBrowseAndSummarize.browse_func, - ) - ) + web_browser_engine: Optional[WebBrowserEngine] = None def __init__(self, **kwargs): super().__init__(**kwargs) if CONFIG.model_for_researcher_summary: self.llm.model = CONFIG.model_for_researcher_summary + self.web_browser_engine = WebBrowserEngine( + engine=WebBrowserEngineType.CUSTOM if self.browse_func else None, + run_func=self.browse_func, + ) + async def run( self, url: str, diff --git a/tests/metagpt/actions/test_research.py b/tests/metagpt/actions/test_research.py new file mode 100644 index 000000000..bc1982c5d --- /dev/null +++ b/tests/metagpt/actions/test_research.py @@ -0,0 +1,105 @@ +import pytest + +from metagpt.actions import research + + +@pytest.mark.asyncio +async def test_collect_links(mocker): + async def mock_llm_ask(self, prompt: str, system_msgs): + if "Please provide up to 2 necessary keywords" in prompt: + return '["metagpt", "llm"]' + + elif "Provide up to 4 queries related to your research topic" in prompt: + return ( + '["MetaGPT use cases", "The roadmap of MetaGPT", ' + '"The function of MetaGPT", "What llm MetaGPT support"]' + ) + elif "sort the remaining search results" in prompt: + return "[1,2]" + + mocker.patch("metagpt.provider.base_gpt_api.BaseGPTAPI.aask", mock_llm_ask) + resp = await research.CollectLinks().run("The application of MetaGPT") + for i in ["MetaGPT use cases", "The roadmap of MetaGPT", "The function of MetaGPT", "What llm MetaGPT support"]: + assert i in resp + + +@pytest.mark.asyncio +async def test_collect_links_with_rank_func(mocker): + rank_before = [] + rank_after = [] + url_per_query = 4 + + def rank_func(results): + results = results[:url_per_query] + rank_before.append(results) + results = results[::-1] + rank_after.append(results) + return results + + mocker.patch("metagpt.provider.base_gpt_api.BaseGPTAPI.aask", mock_collect_links_llm_ask) + resp = await research.CollectLinks(rank_func=rank_func).run("The application of MetaGPT") + for x, y, z in zip(rank_before, rank_after, resp.values()): + assert x[::-1] == y + assert [i["link"] for i in y] == z + + +@pytest.mark.asyncio +async def test_web_browse_and_summarize(mocker): + async def mock_llm_ask(*args, **kwargs): + return "metagpt" + + mocker.patch("metagpt.provider.base_gpt_api.BaseGPTAPI.aask", mock_llm_ask) + url = "https://github.com/geekan/MetaGPT" + url2 = "https://github.com/trending" + query = "What's new in metagpt" + resp = await research.WebBrowseAndSummarize().run(url, query=query) + + assert len(resp) == 1 + assert url in resp + assert resp[url] == "metagpt" + + resp = await research.WebBrowseAndSummarize().run(url, url2, query=query) + assert len(resp) == 2 + + async def mock_llm_ask(*args, **kwargs): + return "Not relevant." + + mocker.patch("metagpt.provider.base_gpt_api.BaseGPTAPI.aask", mock_llm_ask) + resp = await research.WebBrowseAndSummarize().run(url, query=query) + + assert len(resp) == 1 + assert url in resp + assert resp[url] is None + + +@pytest.mark.asyncio +async def test_conduct_research(mocker): + data = None + + async def mock_llm_ask(*args, **kwargs): + nonlocal data + data = f"# Research Report\n## Introduction\n{args} {kwargs}" + return data + + mocker.patch("metagpt.provider.base_gpt_api.BaseGPTAPI.aask", mock_llm_ask) + content = ( + "MetaGPT takes a one line requirement as input and " + "outputs user stories / competitive analysis / requirements / data structures / APIs / documents, etc." + ) + + resp = await research.ConductResearch().run("The application of MetaGPT", content) + assert resp == data + + +async def mock_collect_links_llm_ask(self, prompt: str, system_msgs): + if "Please provide up to 2 necessary keywords" in prompt: + return '["metagpt", "llm"]' + + elif "Provide up to 4 queries related to your research topic" in prompt: + return ( + '["MetaGPT use cases", "The roadmap of MetaGPT", ' '"The function of MetaGPT", "What llm MetaGPT support"]' + ) + elif "sort the remaining search results" in prompt: + return "[1,2]" + + return "" diff --git a/tests/metagpt/roles/test_researcher.py b/tests/metagpt/roles/test_researcher.py index dd130662d..83e90de66 100644 --- a/tests/metagpt/roles/test_researcher.py +++ b/tests/metagpt/roles/test_researcher.py @@ -32,3 +32,19 @@ async def test_researcher(mocker): researcher.RESEARCH_PATH = Path(dirname) await researcher.Researcher().run(topic) assert (researcher.RESEARCH_PATH / f"{topic}.md").read_text().startswith("# Research Report") + + +def test_write_report(mocker): + with TemporaryDirectory() as dirname: + for i, topic in enumerate( + [ + ("1./metagpt"), + ('2.:"metagpt'), + ("3.*?<>|metagpt"), + ("4. metagpt\n"), + ] + ): + researcher.RESEARCH_PATH = Path(dirname) + content = "# Research Report" + researcher.Researcher().write_report(topic, content) + assert (researcher.RESEARCH_PATH / f"{i+1}. metagpt.md").read_text().startswith("# Research Report")