add ut for researcher

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,

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 ""

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")