init project

tests/metagpt/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/29 16:01
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/actions/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 19:35
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/actions/mock.py

@@ -0,0 +1,363 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/18 23:51
+@Author  : alexanderwu
+@File    : mock.py
+"""
+
+PRD_SAMPLE = """产品/功能介绍：基于大语言模型的、私有知识库的搜索引擎
+
+目标：实现一个高效、准确、易用的搜索引擎，能够满足用户对私有知识库的搜索需求，提高工作效率和信息检索的准确性。
+
+用户和使用场景：该搜索引擎主要面向需要频繁使用私有知识库进行信息检索的用户，例如企业内部的知识管理者、研发人员和数据分析师等。用户需要通过输入关键词或短语，快速地获取与其相关的知识库内容。
+
+需求：
+1. 支持基于大语言模型的搜索算法，能够对用户输入的关键词或短语进行语义理解，提高搜索结果的准确性。
+2. 支持私有知识库的建立和维护，能够对知识库内容进行分类、标签和关键词的管理，方便用户进行信息检索。
+3. 提供简洁、直观的用户界面，支持多种搜索方式（如全文搜索、精确搜索、模糊搜索等），方便用户进行快速检索。
+4. 支持搜索结果的排序和过滤，能够根据相关度、时间等因素对搜索结果进行排序，方便用户找到最相关的信息。
+5. 支持多种数据格式的导入和导出，方便用户对知识库内容进行备份和分享。
+
+约束与限制：由于资源有限，需要在保证产品质量的前提下，控制开发成本和时间。同时，需要考虑用户的隐私保护和知识库内容的安全性。
+
+性能指标：
+1. 搜索响应时间：搜索引擎的搜索响应时间应该在毫秒级别，能够快速响应用户的搜索请求。
+2. 搜索准确率：搜索引擎应该能够准确地返回与用户搜索意图相关的知识库内容，提高搜索结果的准确率。
+3. 系统稳定性：搜索引擎应该具备良好的稳定性和可靠性，能够在高并发、大数据量等情况下保持正常运行。
+4. 用户体验：搜索引擎的用户界面应该简洁、直观、易用，让用户能够快速地找到所需的信息。
+"""
+
+DESIGN_LLM_KB_SEARCH_SAMPLE = """## 数据结构
+- 文档对象(Document Object)：表示知识库中的一篇文档，包含文档的标题、内容、标签等信息。
+- 知识库对象(Knowledge Base Object)：表示整个知识库，包含多篇文档对象，以及知识库的分类、标签等信息。
+
+## API接口
+- create_document(title, content, tags)：创建一篇新的文档，返回文档对象。
+- delete_document(document_id)：删除指定ID的文档。
+- update_document(document_id, title=None, content=None, tags=None)：更新指定ID的文档的标题、内容、标签等信息。
+- search_documents(query, mode='fulltext', limit=10, sort_by='relevance')：根据查询条件进行搜索，返回符合条件的文档列表。
+- create_knowledge_base(name, description=None)：创建一个新的知识库，返回知识库对象。
+- delete_knowledge_base(kb_id)：删除指定ID的知识库。
+- update_knowledge_base(kb_id, name=None, description=None)：更新指定ID的知识库的名称、描述等信息。
+
+## 调用流程（以dot语言描述）
+```dot
+digraph search_engine {
+  User -> UI [label="1. 输入查询关键词"];
+  UI -> API [label="2. 调用搜索API"];
+  API -> KnowledgeBase [label="3. 查询知识库"];
+  KnowledgeBase -> NLP [label="4. 进行自然语言处理"];
+  NLP -> API [label="5. 返回处理结果"];
+  API -> UI [label="6. 返回搜索结果"];
+  UI -> User [label="7. 显示搜索结果"];
+}
+```
+
+## 用户编写程序所需的全部、详尽的文件路径列表（以python字符串描述）
+- /api/main.py：主程序入口
+- /api/models/document.py：文档对象的定义
+- /api/models/knowledge_base.py：知识库对象的定义
+- /api/api/search_api.py：搜索API的实现
+- /api/api/knowledge_base_api.py：知识库API的实现
+- /api/nlp/nlp_engine.py：自然语言处理引擎的实现
+- /api/ui/search_ui.py：搜索界面的实现
+- /api/ui/knowledge_base_ui.py：知识库界面的实现
+- /api/utils/database.py：数据库连接和操作相关的工具函数
+- /api/utils/config.py：配置文件，包含数据库连接信息等配置项。
+"""
+
+
+WRITE_CODE_PROMPT_SAMPLE = """
+你是一个工程师。下面是背景信息与你的当前任务，请为任务撰写代码。
+撰写的代码应该符合PEP8，优雅，模块化，易于阅读与维护，代码本身应该有__main__入口来防止桩函数
+
+## 用户编写程序所需的全部、详尽的文件路径列表（只需要相对路径，并不需要前缀，组织形式应该符合PEP规范）
+
+- `main.py`: 主程序文件
+- `search_engine.py`: 搜索引擎实现文件
+- `knowledge_base.py`: 知识库管理文件
+- `user_interface.py`: 用户界面文件
+- `data_import.py`: 数据导入功能文件
+- `data_export.py`: 数据导出功能文件
+- `utils.py`: 工具函数文件
+
+## 数据结构
+
+- `KnowledgeBase`: 知识库类，用于管理私有知识库的内容、分类、标签和关键词。
+- `SearchEngine`: 搜索引擎类，基于大语言模型，用于对用户输入的关键词或短语进行语义理解，并提供准确的搜索结果。
+- `SearchResult`: 搜索结果类，包含与用户搜索意图相关的知识库内容的相关信息。
+- `UserInterface`: 用户界面类，提供简洁、直观的用户界面，支持多种搜索方式和搜索结果的排序和过滤。
+- `DataImporter`: 数据导入类，支持多种数据格式的导入功能，用于将外部数据导入到知识库中。
+- `DataExporter`: 数据导出类，支持多种数据格式的导出功能，用于将知识库内容进行备份和分享。
+
+## API接口
+
+- `KnowledgeBase`类接口:
+  - `add_entry(entry: str, category: str, tags: List[str], keywords: List[str]) -> bool`: 添加知识库条目。
+  - `delete_entry(entry_id: str) -> bool`: 删除知识库条目。
+  - `update_entry(entry_id: str, entry: str, category: str, tags: List[str], keywords: List[str]) -> bool`: 更新知识库条目。
+  - `search_entries(query: str) -> List[str]`: 根据查询词搜索知识库条目。
+
+- `SearchEngine`类接口:
+  - `search(query: str) -> SearchResult`: 根据用户查询词进行搜索，返回与查询意图相关的搜索结果。
+
+- `UserInterface`类接口:
+  - `display_search_results(results: List[SearchResult]) -> None`: 显示搜索结果。
+  - `filter_results(results: List[SearchResult], filters: Dict[str, Any]) -> List[SearchResult]`: 根据过滤条件对搜索结果进行过滤。
+  - `sort_results(results: List[SearchResult], key: str, reverse: bool = False) -> List[SearchResult]`: 根据指定的键对搜索结果进行排序。
+
+- `DataImporter`类接口:
+  - `import_data(file_path: str) -> bool`: 导入外部数据到知识库。
+
+- `DataExporter`类接口:
+  - `export_data(file_path: str) -> bool`: 导出知识库数据到外部文件。
+
+## 调用流程（以dot语言描述）
+
+```dot
+digraph call_flow {
+    rankdir=LR;
+
+    subgraph cluster_user_program {
+        label="User Program";
+        style=dotted;
+
+        main_py -> search_engine_py;
+        main_py -> knowledge_base_py;
+        main_py -> user_interface_py;
+        main_py -> data_import_py;
+        main_py -> data_export_py;
+
+        search_engine_py -> knowledge_base_py;
+        search_engine_py -> user_interface_py;
+
+        user_interface_py -> knowledge_base_py;
+        user_interface_py -> search_engine_py;
+
+        data_import_py -> knowledge_base_py;
+        data_import_py -> user_interface_py;
+
+        data_export_py -> knowledge_base_py;
+        data_export_py -> user_interface_py;
+    }
+
+    main_py [label="main.py"];
+    search_engine_py [label="search_engine.py"];
+    knowledge_base_py [label="knowledge_base.py"];
+    user_interface_py [label="user_interface.py"];
+    data_import_py [label="data_import.py"];
+    data_export_py [label="data_export.py"];
+}
+```
+
+这是一个简化的调用流程图，展示了各个模块之间的调用关系。用户程序的`main.py`文件通过调用其他模块实现搜索引擎的功能。`search_engine.py`模块与`knowledge_base.py`和`user_interface.py`模块进行交互，实现搜索算法和搜索结果的展示。`data_import.py`和`data_export.py`模块与`knowledge_base.py`和`user_interface.py`模块进行交互，实现数据导入和导出的功能。用户界面模块`user_interface.py`与其他模块进行交互，提供简洁、直观的用户界面，并支持搜索方式、排序和过滤等操作。
+
+## 当前任务
+
+"""
+
+TASKS = [
+    "添加数据API：接受用户输入的文档库，对文档库进行索引\n- 使用MeiliSearch连接并添加文档库",
+    "搜索API：接收用户输入的关键词，返回相关的搜索结果\n- 使用MeiliSearch连接并使用接口获得对应数据",
+    "多条件筛选API：接收用户选择的筛选条件，返回符合条件的搜索结果。\n- 使用MeiliSearch进行筛选并返回符合条件的搜索结果",
+    "智能推荐API：根据用户的搜索历史记录和搜索行为，推荐相关的搜索结果。"
+]
+
+TASKS_2 = [
+    "完成main.py的功能"
+]
+
+SEARCH_CODE_SAMPLE = """
+import requests
+
+
+class SearchAPI:
+    def __init__(self, elastic_search_url):
+        self.elastic_search_url = elastic_search_url
+
+    def search(self, keyword):
+        # 构建搜索请求的参数
+        params = {
+            'q': keyword,
+            'size': 10  # 返回结果数量
+        }
+
+        try:
+            # 发送搜索请求
+            response = requests.get(self.elastic_search_url, params=params)
+            if response.status_code == 200:
+                # 解析搜索结果
+                search_results = response.json()
+                formatted_results = self.format_results(search_results)
+                return formatted_results
+            else:
+                print('Error: Failed to retrieve search results.')
+        except requests.exceptions.RequestException as e:
+            print(f'Error: {e}')
+
+    def format_results(self, search_results):
+        formatted_results = []
+        hits = search_results.get('hits', {}).get('hits', [])
+        for hit in hits:
+            result = hit.get('_source', {})
+            title = result.get('title', '')
+            summary = result.get('summary', '')
+            url = result.get('url', '')
+            formatted_results.append({
+                'title': title,
+                'summary': summary,
+                'url': url
+            })
+        return formatted_results
+
+
+if __name__ == '__main__':
+    # 使用示例
+    elastic_search_url = 'http://localhost:9200/search'
+    search_api = SearchAPI(elastic_search_url)
+    keyword = input('Enter search keyword: ')
+    results = search_api.search(keyword)
+    if results:
+        for result in results:
+            print(result)
+    else:
+        print('No results found.')
+"""
+
+
+REFINED_CODE = '''
+import requests
+
+
+class SearchAPI:
+    def __init__(self, elastic_search_url):
+        """
+        初始化SearchAPI对象。
+
+        Args:
+            elastic_search_url (str): ElasticSearch的URL。
+        """
+        self.elastic_search_url = elastic_search_url
+
+    def search(self, keyword, size=10):
+        """
+        搜索关键词并返回相关的搜索结果。
+
+        Args:
+            keyword (str): 用户输入的搜索关键词。
+            size (int): 返回结果数量，默认为10。
+
+        Returns:
+            list: 包含搜索结果的列表，每个结果是一个字典，包含标题、摘要和URL等信息。如果没有搜索结果，返回一个空列表。
+        """
+        # 构建搜索请求的参数
+        params = {
+            'q': keyword,
+            'size': size
+        }
+
+        try:
+            # 发送搜索请求
+            response = requests.get(self.elastic_search_url, params=params)
+            response.raise_for_status()
+            # 解析搜索结果
+            search_results = response.json()
+            formatted_results = self.format_results(search_results)
+            return formatted_results
+        except requests.exceptions.RequestException as e:
+            print(f'Error: {e}')
+            return None
+
+    def format_results(self, search_results):
+        """
+        格式化搜索结果。
+
+        Args:
+            search_results (dict): ElasticSearch返回的搜索结果。
+
+        Returns:
+            list: 包含格式化搜索结果的列表，每个结果是一个字典，包含标题、摘要和URL等信息。如果搜索结果为空，返回None。
+        """
+        if not isinstance(search_results, dict):
+            return None
+
+        formatted_results = []
+        hits = search_results.get('hits', {}).get('hits', [])
+        for hit in hits:
+            result = hit.get('_source', {})
+            title = result.get('title', '')
+            summary = result.get('summary', '')
+            url = result.get('url', '')
+            formatted_results.append({
+                'title': title,
+                'summary': summary,
+                'url': url
+            })
+        return formatted_results if formatted_results else None
+
+
+if __name__ == '__main__':
+    # 使用示例
+    elastic_search_url = 'http://localhost:9200/search'
+    search_api = SearchAPI(elastic_search_url)
+    keyword = input('Enter search keyword: ')
+    results = search_api.search(keyword)
+    if results:
+        for result in results:
+            print(result)
+    else:
+        print('No results found.')
+'''
+
+MEILI_CODE = '''import meilisearch
+from typing import List
+
+
+class DataSource:
+    def __init__(self, name: str, url: str):
+        self.name = name
+        self.url = url
+
+
+class SearchEngine:
+    def __init__(self):
+        self.client = meilisearch.Client('http://localhost:7700')  # MeiliSearch服务器的URL
+
+    def add_documents(self, data_source: DataSource, documents: List[dict]):
+        index_name = f"{data_source.name}_index"
+        index = self.client.get_or_create_index(index_name)
+        index.add_documents(documents)
+
+
+# 示例用法
+if __name__ == '__main__':
+    search_engine = SearchEngine()
+
+    # 假设有一个名为"books"的数据源，包含要添加的文档库
+    books_data_source = DataSource(name='books', url='https://example.com/books')
+
+    # 假设有一个名为"documents"的文档库，包含要添加的文档
+    documents = [
+        {"id": 1, "title": "Book 1", "content": "This is the content of Book 1."},
+        {"id": 2, "title": "Book 2", "content": "This is the content of Book 2."},
+        # 其他文档...
+    ]
+
+    # 添加文档库到搜索引擎
+    search_engine.add_documents(books_data_source, documents)
+'''
+
+MEILI_ERROR = '''/usr/local/bin/python3.9 /Users/alexanderwu/git/metagpt/examples/search/meilisearch_index.py
+Traceback (most recent call last):
+  File "/Users/alexanderwu/git/metagpt/examples/search/meilisearch_index.py", line 44, in <module>
+    search_engine.add_documents(books_data_source, documents)
+  File "/Users/alexanderwu/git/metagpt/examples/search/meilisearch_index.py", line 25, in add_documents
+    index = self.client.get_or_create_index(index_name)
+AttributeError: 'Client' object has no attribute 'get_or_create_index'
+
+Process finished with exit code 1'''
+
+MEILI_CODE_REFINED = """
+"""
+

tests/metagpt/actions/test_action.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 14:43
+@Author  : alexanderwu
+@File    : test_action.py
+"""
+
+import pytest
+from metagpt.logs import logger
+from metagpt.actions import Action, WritePRD, WriteTest
+
+
+def test_action_repr():
+    actions = [Action(), WriteTest(), WritePRD()]
+    assert "WriteTest" in str(actions)

tests/metagpt/actions/test_debug_error.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:46
+@Author  : alexanderwu
+@File    : test_debug_error.py
+"""
+import pytest
+from metagpt.actions.debug_error import DebugError
+
+
+@pytest.mark.asyncio
+async def test_debug_error():
+    code = "def add(a, b):\n    return a - b"
+    error = "AssertionError: Expected add(1, 1) to equal 2 but got 0"
+    fixed_code = "def add(a, b):\n    return a + b"
+
+    debug_error = DebugError("debug_error")
+
+    result = await debug_error.run(code, error)
+
+    prompt = f"以下是一段Python代码:\n\n{code}\n\n执行时发生了以下错误:\n\n{error}\n\n请尝试修复这段代码中的错误。"
+    # mock_llm.ask.assert_called_once_with(prompt)
+    assert len(result) > 0

tests/metagpt/actions/test_design_api.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 19:26
+@Author  : alexanderwu
+@File    : test_design_api.py
+"""
+import pytest
+
+from metagpt.logs import logger
+
+from metagpt.actions.design_api import WriteDesign
+from metagpt.llm import LLM
+from metagpt.roles.architect import Architect
+
+
+@pytest.mark.asyncio
+async def test_design_api():
+    prd = "我们需要一个音乐播放器，它应该有播放、暂停、上一曲、下一曲等功能。"
+
+    design_api = WriteDesign("design_api")
+
+    result = await design_api.run(prd)
+    logger.info(result)
+    assert len(result) > 0
+
+
+@pytest.mark.asyncio
+async def test_design_api_calculator():
+    prd = """产品/功能介绍：基于大语言模型的、私有知识库的搜索引擎
+
+目标：实现一个高效、准确、易用的搜索引擎，能够满足用户对私有知识库的搜索需求，提高工作效率和信息检索的准确性。
+
+用户和使用场景：该搜索引擎主要面向需要频繁使用私有知识库进行信息检索的用户，例如企业内部的知识管理者、研发人员和数据分析师等。用户需要通过输入关键词或短语，快速地获取与其相关的知识库内容。
+
+需求：
+1. 支持基于大语言模型的搜索算法，能够对用户输入的关键词或短语进行语义理解，提高搜索结果的准确性。
+2. 支持私有知识库的建立和维护，能够对知识库内容进行分类、标签和关键词的管理，方便用户进行信息检索。
+3. 提供简洁、直观的用户界面，支持多种搜索方式（如全文搜索、精确搜索、模糊搜索等），方便用户进行快速检索。
+4. 支持搜索结果的排序和过滤，能够根据相关度、时间等因素对搜索结果进行排序，方便用户找到最相关的信息。
+5. 支持多种数据格式的导入和导出，方便用户对知识库内容进行备份和分享。
+
+约束与限制：由于资源有限，需要在保证产品质量的前提下，控制开发成本和时间。同时，需要考虑用户的隐私保护和知识库内容的安全性。
+
+性能指标：
+1. 搜索响应时间：搜索引擎的搜索响应时间应该在毫秒级别，能够快速响应用户的搜索请求。
+2. 搜索准确率：搜索引擎应该能够准确地返回与用户搜索意图相关的知识库内容，提高搜索结果的准确率。
+3. 系统稳定性：搜索引擎应该具备良好的稳定性和可靠性，能够在高并发、大数据量等情况下保持正常运行。
+4. 用户体验：搜索引擎的用户界面应该简洁、直观、易用，让用户能够快速地找到所需的信息。"""
+
+    design_api = WriteDesign("design_api")
+    result = await design_api.run(prd)
+    logger.info(result)
+
+    assert len(result) > 10

tests/metagpt/actions/test_design_api_review.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 19:31
+@Author  : alexanderwu
+@File    : test_design_api_review.py
+"""
+import pytest
+
+from metagpt.actions.design_api_review import DesignReview
+
+
+@pytest.mark.asyncio
+async def test_design_api_review():
+    prd = "我们需要一个音乐播放器，它应该有播放、暂停、上一曲、下一曲等功能。"
+    api_design = """
+数据结构:
+1. Song: 包含歌曲信息，如标题、艺术家等。
+2. Playlist: 包含一系列歌曲。
+
+API列表:
+1. play(song: Song): 开始播放指定的歌曲。
+2. pause(): 暂停当前播放的歌曲。
+3. next(): 跳到播放列表的下一首歌曲。
+4. previous(): 跳到播放列表的上一首歌曲。
+"""
+    api_review = "API设计看起来非常合理，满足了PRD中的所有需求。"
+
+    design_api_review = DesignReview("design_api_review")
+
+    result = await design_api_review.run(prd, api_design)
+
+    prompt = f"以下是产品需求文档(PRD):\n\n{prd}\n\n以下是基于这个PRD设计的API列表:\n\n{api_design}\n\n请审查这个API设计是否满足PRD的需求，以及是否符合良好的设计实践。"
+    # mock_llm.ask.assert_called_once_with(prompt)
+    assert len(result) > 0

tests/metagpt/actions/test_project_management.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 19:12
+@Author  : alexanderwu
+@File    : test_project_management.py
+"""
+
+from metagpt.actions.project_management import WriteTasks, AssignTasks
+
+
+class TestCreateProjectPlan:
+    pass
+
+
+class TestAssignTasks:
+    pass

tests/metagpt/actions/test_run_code.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:46
+@Author  : alexanderwu
+@File    : test_run_code.py
+"""
+import pytest
+from metagpt.actions.run_code import RunCode
+
+
+@pytest.mark.asyncio
+async def test_run_code():
+    code = """
+def add(a, b):
+    return a + b
+result = add(1, 2)
+"""
+    run_code = RunCode("run_code")
+
+    result = await run_code.run(code)
+
+    assert result == 3
+
+
+@pytest.mark.asyncio
+async def test_run_code_with_error():
+    code = """
+def add(a, b):
+    return a + b
+result = add(1, '2')
+"""
+    run_code = RunCode("run_code")
+
+    result = await run_code.run(code)
+
+    assert "TypeError: unsupported operand type(s) for +" in result
+

tests/metagpt/actions/test_write_code.py

@@ -0,0 +1,33 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:45
+@Author  : alexanderwu
+@File    : test_write_code.py
+"""
+import pytest
+from metagpt.logs import logger
+from metagpt.actions.write_code import WriteCode
+from tests.metagpt.actions.mock import WRITE_CODE_PROMPT_SAMPLE, TASKS_2
+from metagpt.llm import LLM
+
+
+@pytest.mark.asyncio
+async def test_write_code():
+    api_design = "设计一个名为'add'的函数，该函数接受两个整数作为输入，并返回它们的和。"
+    write_code = WriteCode("write_code")
+
+    code = await write_code.run(api_design)
+    logger.info(code)
+
+    # 我们不能精确地预测生成的代码，但我们可以检查某些关键字
+    assert 'def add' in code
+    assert 'return' in code
+
+
+@pytest.mark.asyncio
+async def test_write_code_directly():
+    prompt = WRITE_CODE_PROMPT_SAMPLE + '\n' + TASKS_2[0]
+    llm = LLM()
+    rsp = await llm.aask(prompt)
+    logger.info(rsp)

tests/metagpt/actions/test_write_code_review.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:45
+@Author  : alexanderwu
+@File    : test_write_code_review.py
+"""
+import pytest
+from metagpt.logs import logger
+from metagpt.llm import LLM
+from metagpt.actions.write_code_review import WriteCodeReview
+from tests.metagpt.actions.mock import SEARCH_CODE_SAMPLE
+
+
+@pytest.mark.asyncio
+async def test_write_code_review():
+    code = """
+def add(a, b):
+    return a + b
+"""
+    write_code_review = WriteCodeReview("write_code_review")
+
+    review = await write_code_review.run(code)
+
+    # 我们不能精确地预测生成的代码评审，但我们可以检查返回的是否为字符串
+    assert isinstance(review, str)
+    assert len(review) > 0
+
+
+@pytest.mark.asyncio
+async def test_write_code_review_directly():
+    code = SEARCH_CODE_SAMPLE
+    write_code_review = WriteCodeReview("write_code_review")
+    review = await write_code_review.run(code)
+    logger.info(review)

tests/metagpt/actions/test_write_prd.py

@@ -0,0 +1,25 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:45
+@Author  : alexanderwu
+@File    : test_write_prd.py
+"""
+import pytest
+from metagpt.logs import logger
+from metagpt.actions import WritePRD, BossRequirement
+from metagpt.roles.product_manager import ProductManager
+from metagpt.schema import Message
+
+
+@pytest.mark.asyncio
+async def test_write_prd():
+    product_manager = ProductManager()
+    requirements = "开发一个基于大语言模型与私有知识库的搜索引擎，希望可以基于大语言模型进行搜索总结"
+    prd = await product_manager.handle(Message(content=requirements, cause_by=BossRequirement))
+    logger.info(requirements)
+    logger.info(prd)
+
+    # Assert the prd is not None or empty
+    assert prd is not None
+    assert prd != ""

tests/metagpt/actions/test_write_prd_review.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:45
+@Author  : alexanderwu
+@File    : test_write_prd_review.py
+"""
+import pytest
+from metagpt.actions.write_prd_review import WritePRDReview
+
+
+@pytest.mark.asyncio
+async def test_write_prd_review():
+    prd = """
+    Introduction: This is a new feature for our product.
+    Goals: The goal is to improve user engagement.
+    User Scenarios: The expected user group is millennials who like to use social media.
+    Requirements: The feature needs to be interactive and user-friendly.
+    Constraints: The feature needs to be implemented within 2 months.
+    Mockups: There will be a new button on the homepage that users can click to access the feature.
+    Metrics: We will measure the success of the feature by user engagement metrics.
+    Timeline: The feature should be ready for testing in 1.5 months.
+    """
+
+    write_prd_review = WritePRDReview("write_prd_review")
+
+    prd_review = await write_prd_review.run(prd)
+
+    # We cannot exactly predict the generated PRD review, but we can check if it is a string and if it is not empty
+    assert isinstance(prd_review, str)
+    assert len(prd_review) > 0

tests/metagpt/actions/test_write_test.py

@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 17:45
+@Author  : alexanderwu
+@File    : test_write_test.py
+"""
+import pytest
+from metagpt.logs import logger
+from metagpt.actions.write_test import WriteTest
+
+
+@pytest.mark.asyncio
+async def test_write_test():
+    code = """
+    def add(a, b):
+        return a + b
+    """
+
+    write_test = WriteTest("write_test")
+
+    test_cases = await write_test.run(code)
+
+    # We cannot exactly predict the generated test cases, but we can check if it is a string and if it is not empty
+    assert isinstance(test_cases, str)
+    assert len(test_cases) > 0

tests/metagpt/document_store/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/27 20:19
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/document_store/test_chromadb_store.py

@@ -0,0 +1,30 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/6/6 00:41
+@Author  : alexanderwu
+@File    : test_chromadb_store.py
+"""
+import pytest
+from sentence_transformers import SentenceTransformer
+
+from metagpt.document_store.chromadb_store import ChromaStore
+
+
+# @pytest.mark.skip()
+def test_chroma_store():
+    """FIXME：chroma使用感觉很诡异，一用Python就挂，测试用例里也是"""
+    # 创建 ChromaStore 实例，使用 'sample_collection' 集合
+    document_store = ChromaStore('sample_collection_1')
+
+    # 使用 write 方法添加多个文档
+    document_store.write(["This is document1", "This is document2"],
+                [{"source": "google-docs"}, {"source": "notion"}],
+                ["doc1", "doc2"])
+
+    # 使用 add 方法添加一个文档
+    document_store.add("This is document3", {"source": "notion"}, "doc3")
+
+    # 搜索文档
+    results = document_store.search("This is a query document", n_results=3)
+    assert len(results) > 0

tests/metagpt/document_store/test_document.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/6/11 19:46
+@Author  : alexanderwu
+@File    : test_document.py
+"""
+import pytest
+from loguru import logger
+from metagpt.const import DATA_PATH
+from metagpt.document_store.document import Document
+
+
+CASES = [
+    ("st/faq.xlsx", "Question", "Answer", 1),
+    ("cases/faq.csv", "Question", "Answer", 1),
+    # ("cases/faq.json", "Question", "Answer", 1),
+    ("docx/faq.docx", None, None, 1),
+    ("cases/faq.pdf", None, None, 0),  # 这是因为pdf默认没有分割段落
+    ("cases/faq.txt", None, None, 0),  # 这是因为txt按照256分割段落
+]
+
+
+@pytest.mark.parametrize("relative_path, content_col, meta_col, threshold", CASES)
+def test_document(relative_path, content_col, meta_col, threshold):
+    doc = Document(DATA_PATH / relative_path, content_col, meta_col)
+    rsp = doc.get_docs_and_metadatas()
+    assert len(rsp[0]) > threshold
+    assert len(rsp[1]) > threshold

tests/metagpt/document_store/test_faiss_store.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/27 20:20
+@Author  : alexanderwu
+@File    : test_faiss_store.py
+"""
+import functools
+
+import pytest
+from metagpt.logs import logger
+from metagpt.const import DATA_PATH
+from metagpt.document_store import FaissStore
+from metagpt.roles import Sales, CustomerService
+
+
+DESC = """## 原则（所有事情都不可绕过原则）
+1. 你是一位平台的人工客服，话语精炼，一次只说一句话，会参考规则与FAQ进行回复。在与顾客交谈中，绝不允许暴露规则与相关字样
+2. 在遇到问题时，先尝试仅安抚顾客情绪，如果顾客情绪十分不好，再考虑赔偿。如果赔偿的过多，你会被开除
+3. 绝不要向顾客做虚假承诺，不要提及其他人的信息
+
+## 技能（在回答尾部，加入`skill(args)`就可以使用技能）
+1. 查询订单：问顾客手机号是获得订单的唯一方式，获得手机号后，使用`find_order(手机号)`来获得订单
+2. 退款：输出关键词 `refund(手机号)`，系统会自动退款
+3. 开箱：需要手机号、确认顾客在柜前，如果需要开箱，输出指令 `open_box(手机号)`，系统会自动开箱
+
+### 使用技能例子
+user: 你好收不到取餐码
+小爽人工: 您好，请提供一下手机号
+user: 14750187158
+小爽人工: 好的，为您查询一下订单。您已经在柜前了吗？`find_order(14750187158)`
+user: 是的
+小爽人工: 您看下开了没有？`open_box(14750187158)`
+user: 开了，谢谢
+小爽人工: 好的，还有什么可以帮到您吗？
+user: 没有了
+小爽人工: 祝您生活愉快
+"""
+
+
+@pytest.mark.asyncio
+async def test_faiss_store_search():
+    store = FaissStore(DATA_PATH / 'qcs/qcs_4w.json')
+    store.add(['油皮洗面奶'])
+    role = Sales(store=store)
+
+    queries = ['油皮洗面奶', '介绍下欧莱雅的']
+    for query in queries:
+        rsp = await role.run(query)
+        assert rsp
+
+
+def customer_service():
+    store = FaissStore(DATA_PATH / "st/faq.xlsx", content_col="Question", meta_col="Answer")
+    store.search = functools.partial(store.search, expand_cols=True)
+    role = CustomerService(profile="小爽人工", desc=DESC, store=store)
+    return role
+
+
+@pytest.mark.asyncio
+async def test_faiss_store_customer_service():
+    allq = [
+        # ["我的餐怎么两小时都没到", "退货吧"],
+        ["你好收不到取餐码，麻烦帮我开箱", "14750187158", ]
+    ]
+    role = customer_service()
+    for queries in allq:
+        for query in queries:
+            rsp = await role.run(query)
+            assert rsp
+
+
+def test_faiss_store_no_file():
+    with pytest.raises(FileNotFoundError):
+        FaissStore(DATA_PATH / 'wtf.json')

tests/metagpt/document_store/test_milvus_store.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/6/11 21:08
+@Author  : alexanderwu
+@File    : test_milvus_store.py
+"""
+import random
+import numpy as np
+from metagpt.logs import logger
+from metagpt.document_store.milvus_store import MilvusStore, MilvusConnection
+
+
+book_columns = {'idx': int, 'name': str, 'desc': str, 'emb': np.ndarray, 'price': float}
+book_data = [
+    [i for i in range(10)],
+    [f"book-{i}" for i in range(10)],
+    [f"book-desc-{i}" for i in range(10000, 10010)],
+    [[random.random() for _ in range(2)] for _ in range(10)],
+    [random.random() for _ in range(10)],
+]
+
+
+def test_milvus_store():
+    milvus_connection = MilvusConnection(alias="default", host="192.168.50.161", port="30530")
+    milvus_store = MilvusStore(milvus_connection)
+    milvus_store.drop('Book')
+    milvus_store.create_collection('Book', book_columns)
+    milvus_store.add(book_data)
+    milvus_store.build_index('emb')
+    milvus_store.load_collection()
+
+    results = milvus_store.search([[1.0, 1.0]], field='emb')
+    logger.info(results)
+    assert results

tests/metagpt/gpt_provider/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/6 17:32
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/gpt_provider/test_azure_gpt_api.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/16 10:12
+@Author  : alexanderwu
+@File    : test_azure_gpt_api.py
+"""
+
+from metagpt.provider import AzureGPTAPI
+
+
+def test_azure_gpt_api():
+    api = AzureGPTAPI()
+    rsp = api.ask('hello')
+    assert len(rsp) > 0

tests/metagpt/gpt_provider/test_base_gpt_api.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/7 17:40
+@Author  : alexanderwu
+@File    : test_base_gpt_api.py
+"""
+
+from metagpt.schema import Message
+
+
+def test_message():
+    message = Message(role='user', content='wtf')
+    assert 'role' in message.to_dict()
+    assert 'user' in str(message)

tests/metagpt/management/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/6/6 12:38
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/management/test_skill_manager.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/6/6 12:38
+@Author  : alexanderwu
+@File    : test_skill_manager.py
+"""
+from metagpt.actions import WritePRD, WriteTest
+from metagpt.logs import logger
+from metagpt.management.skill_manager import SkillManager
+
+
+def test_skill_manager():
+    manager = SkillManager()
+    logger.info(manager._store)
+
+    write_prd = WritePRD("WritePRD")
+    write_prd.desc = "基于老板或其他人的需求进行PRD的撰写，包括用户故事、需求分解等"
+    write_test = WriteTest("WriteTest")
+    write_test.desc = "进行测试用例的撰写"
+    manager.add_skill(write_prd)
+    manager.add_skill(write_test)
+
+    skill = manager.get_skill("WriteTest")
+    logger.info(skill)
+
+    rsp = manager.retrieve_skill("写PRD")
+    logger.info(rsp)
+    assert rsp[0] == "WritePRD"
+
+    rsp = manager.retrieve_skill("写测试用例")
+    logger.info(rsp)
+    assert rsp[0] == 'WriteTest'
+
+    rsp = manager.retrieve_skill_scored("写PRD")
+    logger.info(rsp)

tests/metagpt/roles/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/12 10:14
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/roles/mock.py

@@ -0,0 +1,261 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/12 13:05
+@Author  : alexanderwu
+@File    : mock.py
+"""
+from metagpt.actions import WritePRD, BossRequirement, WriteDesign, WriteTasks
+from metagpt.schema import Message
+
+BOSS_REQUIREMENT = """开发一个基于大语言模型与私有知识库的搜索引擎，希望可以基于大语言模型进行搜索总结"""
+
+DETAIL_REQUIREMENT = """需求：开发一个基于LLM（大语言模型）与私有知识库的搜索引擎，希望有几点能力
+1. 用户可以在私有知识库进行搜索，再根据大语言模型进行总结，输出的结果包括了总结
+2. 私有知识库可以实时更新，底层基于 ElasticSearch
+3. 私有知识库支持pdf、word、txt等各种文件格式上传，上传后可以在服务端解析为文本，存储ES
+
+资源：
+1. 大语言模型已经有前置的抽象、部署，可以通过 `from metagpt.llm import LLM`，再使用`LLM().ask(prompt)`直接调用
+2. Elastic已有[部署](http://192.168.50.82:9200/)，代码可以直接使用这个部署"""
+
+
+PRD = '''## 原始需求
+```python
+"""
+我们希望开发一个基于大语言模型与私有知识库的搜索引擎。该搜索引擎应当能根据用户输入的查询进行智能搜索，并基于大语言模型对搜索结果进行总结，以便用户能够快速获取他们所需要的信息。该搜索引擎应当能够处理大规模的数据，同时保持搜索结果的准确性和相关性。我们希望这个产品能够降低用户在查找、筛选和理解信息时的工作负担，提高他们的工作效率。
+"""
+```
+
+## 产品目标
+```python
+[
+    "提供高准确性、高相关性的搜索结果，满足用户的查询需求",
+    "基于大语言模型对搜索结果进行智能总结，帮助用户快速获取所需信息",
+    "处理大规模数据，保证搜索的速度和效率，提高用户的工作效率"
+]
+```
+
+## 用户故事
+```python
+[
+    "假设用户是一名研究员，他正在为一项关于全球气候变化的报告做研究。他输入了'全球气候变化的最新研究'，我们的搜索引擎快速返回了相关的文章、报告、数据集等。并且基于大语言模型对这些信息进行了智能总结，研究员可以快速了解到最新的研究趋势和发现。",
+    "用户是一名学生，正在为即将到来的历史考试复习。他输入了'二战的主要战役'，搜索引擎返回了相关的资料，大语言模型总结出主要战役的时间、地点、结果等关键信息，帮助学生快速记忆。",
+    "用户是一名企业家，他正在寻找关于最新的市场趋势信息。他输入了'2023年人工智能市场趋势'，搜索引擎返回了各种报告、新闻和分析文章。大语言模型对这些信息进行了总结，用户能够快速了解到市场的最新动态和趋势。"
+]
+```
+
+## 竞品分析
+```python
+[
+    "Google Search：Google搜索是市场上最主要的搜索引擎，它能够提供海量的搜索结果。但Google搜索并不提供搜索结果的总结功能，用户需要自己去阅读和理解搜索结果。",
+    "Microsoft Bing：Bing搜索也能提供丰富的搜索结果，同样没有提供搜索结果的总结功能。",
+    "Wolfram Alpha：Wolfram Alpha是一个基于知识库的计算型搜索引擎，能够针对某些特定类型的查询提供直接的答案和总结，但它的知识库覆盖范围有限，无法处理大规模的数据。"
+]
+```
+
+## 开发需求池
+```python
+[
+    ("开发基于大语言模型的智能总结功能", 5),
+    ("开发搜索引擎核心算法，包括索引构建、查询处理、结果排序等", 7),
+    ("设计和实现用户界面，包括查询输入、搜索结果展示、总结结果展示等", 3),
+    ("构建和维护私有知识库，包括数据采集、清洗、更新等", 7),
+    ("优化搜索引擎性能，包括搜索速度、准确性、相关性等", 6),
+    ("开发用户反馈机制，包括反馈界面、反馈处理等", 2),
+    ("开发安全防护机制，防止恶意查询和攻击", 3),
+    ("集成大语言模型，包括模型选择、优化、更新等", 5),
+    ("进行大规模的测试，包括功能测试、性能测试、压力测试等", 5),
+    ("开发数据监控和日志系统，用于监控搜索引擎的运行状态和性能", 4)
+]
+```
+'''
+
+SYSTEM_DESIGN = '''## Python package name
+```python
+"smart_search_engine"
+```
+
+## Task list:
+```python
+[
+    "smart_search_engine/__init__.py",
+    "smart_search_engine/main.py",
+    "smart_search_engine/search.py",
+    "smart_search_engine/index.py",
+    "smart_search_engine/ranking.py",
+    "smart_search_engine/summary.py",
+    "smart_search_engine/knowledge_base.py",
+    "smart_search_engine/interface.py",
+    "smart_search_engine/user_feedback.py",
+    "smart_search_engine/security.py",
+    "smart_search_engine/testing.py",
+    "smart_search_engine/monitoring.py"
+]
+```
+
+## Data structures and interface definitions
+```mermaid
+classDiagram
+    class Main {
+        -SearchEngine search_engine
+        +main() str
+    }
+    class SearchEngine {
+        -Index index
+        -Ranking ranking
+        -Summary summary
+        +search(query: str) str
+    }
+    class Index {
+        -KnowledgeBase knowledge_base
+        +create_index(data: dict)
+        +query_index(query: str) list
+    }
+    class Ranking {
+        +rank_results(results: list) list
+    }
+    class Summary {
+        +summarize_results(results: list) str
+    }
+    class KnowledgeBase {
+        +update(data: dict)
+        +fetch_data(query: str) dict
+    }
+    Main --> SearchEngine
+    SearchEngine --> Index
+    SearchEngine --> Ranking
+    SearchEngine --> Summary
+    Index --> KnowledgeBase
+```
+
+## Program call flow
+```mermaid
+sequenceDiagram
+    participant M as Main
+    participant SE as SearchEngine
+    participant I as Index
+    participant R as Ranking
+    participant S as Summary
+    participant KB as KnowledgeBase
+    M->>SE: search(query)
+    SE->>I: query_index(query)
+    I->>KB: fetch_data(query)
+    KB-->>I: return data
+    I-->>SE: return results
+    SE->>R: rank_results(results)
+    R-->>SE: return ranked_results
+    SE->>S: summarize_results(ranked_results)
+    S-->>SE: return summary
+    SE-->>M: return summary
+```
+'''
+
+
+TASKS = '''## Logic Analysis
+
+在这个项目中，所有的模块都依赖于“SearchEngine”类，这是主入口，其他的模块（Index、Ranking和Summary）都通过它交互。另外，"Index"类又依赖于"KnowledgeBase"类，因为它需要从知识库中获取数据。
+
+- "main.py"包含"Main"类，是程序的入口点，它调用"SearchEngine"进行搜索操作，所以在其他任何模块之前，"SearchEngine"必须首先被定义。
+- "search.py"定义了"SearchEngine"类，它依赖于"Index"、"Ranking"和"Summary"，因此，这些模块需要在"search.py"之前定义。
+- "index.py"定义了"Index"类，它从"knowledge_base.py"获取数据来创建索引，所以"knowledge_base.py"需要在"index.py"之前定义。
+- "ranking.py"和"summary.py"相对独立，只需确保在"search.py"之前定义。
+- "knowledge_base.py"是独立的模块，可以优先开发。
+- "interface.py"、"user_feedback.py"、"security.py"、"testing.py"和"monitoring.py"看起来像是功能辅助模块，可以在主要功能模块开发完成后并行开发。
+
+## Task list
+
+```python
+task_list = [
+    "smart_search_engine/knowledge_base.py",
+    "smart_search_engine/index.py",
+    "smart_search_engine/ranking.py",
+    "smart_search_engine/summary.py",
+    "smart_search_engine/search.py",
+    "smart_search_engine/main.py",
+    "smart_search_engine/interface.py",
+    "smart_search_engine/user_feedback.py",
+    "smart_search_engine/security.py",
+    "smart_search_engine/testing.py",
+    "smart_search_engine/monitoring.py",
+]
+```
+这个任务列表首先定义了最基础的模块，然后是依赖这些模块的模块，最后是辅助模块。可以根据团队的能力和资源，同时开发多个任务，只要满足依赖关系。例如，在开发"search.py"之前，可以同时开发"knowledge_base.py"、"index.py"、"ranking.py"和"summary.py"。
+'''
+
+
+TASKS_TOMATO_CLOCK = '''## Required Python third-party packages: Provided in requirements.txt format
+```python
+Flask==2.1.1
+Jinja2==3.1.0
+Bootstrap==5.3.0-alpha1
+```
+
+## Logic Analysis: Provided as a Python str, analyze the dependencies between the files, which work should be done first
+```python
+"""
+1. Start by setting up the Flask app, config.py, and requirements.txt to create the basic structure of the web application.
+2. Create the timer functionality using JavaScript and the Web Audio API in the timer.js file.
+3. Develop the frontend templates (index.html and settings.html) using Jinja2 and integrate the timer functionality.
+4. Add the necessary static files (main.css, main.js, and notification.mp3) for styling and interactivity.
+5. Implement the ProgressBar class in main.js and integrate it with the Timer class in timer.js.
+6. Write tests for the application in test_app.py.
+"""
+```
+
+## Task list: Provided as Python list[str], each str is a file, the more at the beginning, the more it is a prerequisite dependency, should be done first
+```python
+task_list = [
+    'app.py',
+    'config.py',
+    'requirements.txt',
+    'static/js/timer.js',
+    'templates/index.html',
+    'templates/settings.html',
+    'static/css/main.css',
+    'static/js/main.js',
+    'static/audio/notification.mp3',
+    'static/js/progressbar.js',
+    'tests/test_app.py'
+]
+```
+'''
+
+
+
+TASK = """smart_search_engine/knowledge_base.py"""
+
+
+STRS_FOR_PARSING = [
+"""
+## 1
+```python
+a
+```
+""",
+"""
+##2
+```python
+"a"
+```
+""",
+"""
+##  3
+```python
+a = "a"
+```
+""",
+"""
+## 4
+```python
+a =  'a'
+```
+"""
+]
+
+
+class MockMessages:
+    req = Message(role="Boss", content=BOSS_REQUIREMENT, cause_by=BossRequirement)
+    prd = Message(role="Product Manager", content=PRD, cause_by=WritePRD)
+    system_design = Message(role="Architect", content=SYSTEM_DESIGN, cause_by=WriteDesign)
+    tasks = Message(role="Project Manager", content=TASKS, cause_by=WriteTasks)

tests/metagpt/roles/test_architect.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/20 14:37
+@Author  : alexanderwu
+@File    : test_architect.py
+"""
+import pytest
+
+from metagpt.actions import BossRequirement
+from metagpt.logs import logger
+from metagpt.roles import Architect
+from metagpt.schema import Message
+from tests.metagpt.roles.mock import PRD, DETAIL_REQUIREMENT, BOSS_REQUIREMENT, MockMessages
+
+
+@pytest.mark.asyncio
+async def test_architect():
+    role = Architect()
+    role.recv(MockMessages.req)
+    rsp = await role.handle(MockMessages.prd)
+    logger.info(rsp)
+    assert len(rsp.content) > 0

tests/metagpt/roles/test_engineer.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/12 10:14
+@Author  : alexanderwu
+@File    : test_engineer.py
+"""
+import re
+import ast
+import pytest
+from metagpt.logs import logger
+from metagpt.utils.common import CodeParser
+from metagpt.roles.engineer import Engineer
+from metagpt.schema import Message
+from tests.metagpt.roles.mock import SYSTEM_DESIGN, TASKS, PRD, MockMessages, STRS_FOR_PARSING, \
+    TASKS_TOMATO_CLOCK
+
+
+@pytest.mark.asyncio
+async def test_engineer():
+    engineer = Engineer()
+
+    engineer.recv(MockMessages.req)
+    engineer.recv(MockMessages.prd)
+    engineer.recv(MockMessages.system_design)
+    rsp = await engineer.handle(MockMessages.tasks)
+
+    logger.info(rsp)
+    assert "all done." == rsp.content
+
+
+def test_parse_str():
+    for idx, i in enumerate(STRS_FOR_PARSING):
+        text = CodeParser.parse_str(f"{idx+1}", i)
+        # logger.info(text)
+        assert text == 'a'
+
+
+def test_parse_blocks():
+    tasks = CodeParser.parse_blocks(TASKS)
+    logger.info(tasks.keys())
+    assert 'Task list' in tasks.keys()
+
+
+target_list = [
+    "smart_search_engine/knowledge_base.py",
+    "smart_search_engine/index.py",
+    "smart_search_engine/ranking.py",
+    "smart_search_engine/summary.py",
+    "smart_search_engine/search.py",
+    "smart_search_engine/main.py",
+    "smart_search_engine/interface.py",
+    "smart_search_engine/user_feedback.py",
+    "smart_search_engine/security.py",
+    "smart_search_engine/testing.py",
+    "smart_search_engine/monitoring.py",
+]
+
+
+def test_parse_file_list():
+    tasks = CodeParser.parse_file_list("任务列表", TASKS)
+    logger.info(tasks)
+    assert isinstance(tasks, list)
+    assert target_list == tasks
+
+
+target_code = """task_list = [
+    "smart_search_engine/knowledge_base.py",
+    "smart_search_engine/index.py",
+    "smart_search_engine/ranking.py",
+    "smart_search_engine/summary.py",
+    "smart_search_engine/search.py",
+    "smart_search_engine/main.py",
+    "smart_search_engine/interface.py",
+    "smart_search_engine/user_feedback.py",
+    "smart_search_engine/security.py",
+    "smart_search_engine/testing.py",
+    "smart_search_engine/monitoring.py",
+]
+"""
+
+
+def test_parse_code():
+    code = CodeParser.parse_code("任务列表", TASKS, lang="python")
+    logger.info(code)
+    assert isinstance(code, str)
+    assert target_code == code
+
+
+def test_parse_file_list():
+    file_list = CodeParser.parse_file_list("Task list", TASKS_TOMATO_CLOCK, lang="python")
+    logger.info(file_list)

tests/metagpt/roles/test_product_manager.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/16 14:50
+@Author  : alexanderwu
+@File    : test_product_manager.py
+"""
+import pytest
+from metagpt.logs import logger
+
+from metagpt.actions import BossRequirement
+from metagpt.roles import ProductManager
+from metagpt.schema import Message
+from tests.metagpt.roles.mock import DETAIL_REQUIREMENT, BOSS_REQUIREMENT, MockMessages
+
+
+@pytest.mark.asyncio
+async def test_product_manager():
+    product_manager = ProductManager()
+    rsp = await product_manager.handle(MockMessages.req)
+    logger.info(rsp)
+    assert len(rsp.content) > 0
+    assert "产品目标" in rsp.content

tests/metagpt/roles/test_project_manager.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/12 10:23
+@Author  : alexanderwu
+@File    : test_project_manager.py
+"""
+import pytest
+from metagpt.logs import logger
+from metagpt.roles import ProjectManager
+from metagpt.schema import Message
+from tests.metagpt.roles.mock import SYSTEM_DESIGN, MockMessages
+
+
+@pytest.mark.asyncio
+async def test_project_manager():
+    project_manager = ProjectManager()
+    rsp = await project_manager.handle(MockMessages.system_design)
+    logger.info(rsp)

tests/metagpt/roles/test_qa_engineer.py

@@ -0,0 +1,8 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/12 12:01
+@Author  : alexanderwu
+@File    : test_qa_engineer.py
+"""
+

tests/metagpt/test_action.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 14:44
+@Author  : alexanderwu
+@File    : test_action.py
+"""

tests/metagpt/test_environment.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/12 00:47
+@Author  : alexanderwu
+@File    : test_environment.py
+"""
+
+import pytest
+
+from metagpt.logs import logger
+from metagpt.manager import Manager
+from metagpt.environment import Environment
+from metagpt.roles import ProductManager, Architect, Role
+from metagpt.schema import Message
+from metagpt.actions import BossRequirement
+
+
+@pytest.fixture
+def env():
+    return Environment()
+
+
+def test_add_role(env: Environment):
+    role = ProductManager("Alice", "product manager", "create a new product", "limited resources")
+    env.add_role(role)
+    assert env.get_role(role.profile) == role
+
+
+def test_get_roles(env: Environment):
+    role1 = Role("Alice", "product manager", "create a new product", "limited resources")
+    role2 = Role("Bob", "engineer", "develop the new product", "short deadline")
+    env.add_role(role1)
+    env.add_role(role2)
+    roles = env.get_roles()
+    assert roles == {role1.profile: role1, role2.profile: role2}
+
+
+def test_set_manager(env: Environment):
+    manager = Manager()
+    env.set_manager(manager)
+    assert env.manager == manager
+
+
+@pytest.mark.asyncio
+async def test_publish_and_process_message(env: Environment):
+    product_manager = ProductManager("Alice", "Product Manager", "做AI Native产品", "资源有限")
+    architect = Architect("Bob", "Architect", "设计一个可用、高效、较低成本的系统，包括数据结构与接口", "资源有限，需要节省成本")
+
+    env.add_roles([product_manager, architect])
+    env.set_manager(Manager())
+    env.publish_message(Message(role="BOSS", content="需要一个基于LLM做总结的搜索引擎", cause_by=BossRequirement))
+
+    await env.run(k=2)
+    logger.info(f"{env.history=}")
+    assert len(env.history) > 10

tests/metagpt/test_gpt.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/29 19:47
+@Author  : alexanderwu
+@File    : test_gpt.py
+"""
+
+import pytest
+from metagpt.logs import logger
+
+
+@pytest.mark.usefixtures("llm_api")
+class TestGPT:
+    def test_llm_api_ask(self, llm_api):
+        answer = llm_api.ask('hello chatgpt')
+        assert len(answer) > 0
+
+    # def test_gptapi_ask_batch(self, llm_api):
+    #     answer = llm_api.ask_batch(['请扮演一个Google Python专家工程师，如果理解，回复明白', '写一个hello world'])
+    #     assert len(answer) > 0
+
+    def test_llm_api_ask_code(self, llm_api):
+        answer = llm_api.ask_code(['请扮演一个Google Python专家工程师，如果理解，回复明白', '写一个hello world'])
+        assert len(answer) > 0
+
+    @pytest.mark.asyncio
+    async def test_llm_api_aask(self, llm_api):
+        answer = await llm_api.aask('hello chatgpt')
+        assert len(answer) > 0
+
+    @pytest.mark.asyncio
+    async def test_llm_api_aask_code(self, llm_api):
+        answer = await llm_api.aask_code(['请扮演一个Google Python专家工程师，如果理解，回复明白', '写一个hello world'])
+        assert len(answer) > 0
+
+    @pytest.mark.asyncio
+    async def test_llm_api_costs(self, llm_api):
+        answer = await llm_api.aask('hello chatgpt')
+        costs = llm_api.get_costs()
+        logger.info(costs)
+        assert costs.total_cost > 0

tests/metagpt/test_llm.py

@@ -0,0 +1,34 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 14:45
+@Author  : alexanderwu
+@File    : test_llm.py
+"""
+
+import pytest
+from metagpt.llm import LLM
+
+
+@pytest.fixture()
+def llm():
+    return LLM()
+
+
+@pytest.mark.asyncio
+async def test_llm_aask(llm):
+    assert len(await llm.aask('hello world')) > 0
+
+
+@pytest.mark.asyncio
+async def test_llm_aask_batch(llm):
+    assert len(await llm.aask_batch(['hi', 'write python hello world.'])) > 0
+
+
+@pytest.mark.asyncio
+async def test_llm_aask(llm):
+
+    hello_msg = [{'role': 'user', 'content': 'hello'}]
+    assert len(await llm.acompletion(hello_msg)) > 0
+    assert len(await llm.acompletion_batch([hello_msg])) > 0
+    assert len(await llm.acompletion_batch_text([hello_msg])) > 0

tests/metagpt/test_manager.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 14:45
+@Author  : alexanderwu
+@File    : test_manager.py
+"""

tests/metagpt/test_message.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/16 10:57
+@Author  : alexanderwu
+@File    : test_message.py
+"""
+import pytest
+
+from metagpt.schema import Message, UserMessage, SystemMessage, AIMessage, RawMessage
+
+
+def test_message():
+    msg = Message(role='User', content='WTF')
+    assert msg.to_dict()['role'] == 'User'
+    assert 'User' in str(msg)
+
+
+def test_all_messages():
+    test_content = 'test_message'
+    msgs = [
+        UserMessage(test_content),
+        SystemMessage(test_content),
+        AIMessage(test_content),
+        Message(test_content, role='QA')
+    ]
+    for msg in msgs:
+        assert msg.content == test_content
+
+
+def test_raw_message():
+    msg = RawMessage(role='user', content='raw')
+    assert msg['role'] == 'user'
+    assert msg['content'] == 'raw'
+    with pytest.raises(KeyError):
+        assert msg['1'] == 1, "KeyError: '1'"

tests/metagpt/test_parser.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/26 20:54
+@Author  : alexanderwu
+@File    : test_parser.py
+"""
+from langchain.schema import AgentAction, AgentFinish, OutputParserException
+from metagpt.parsers import BasicParser
+
+def test_basic_parser():
+    parser = BasicParser()
+    action_sample = "I need to calculate the 0.23 power of Elon Musk's current age.\nAction: Calculator\nAction Input: 49 raised to the 0.23 power"
+    final_answer_sample = "I now know the answer to the question.\nFinal Answer: 2.447626228522259"
+
+    rsp = parser.parse(action_sample)
+    assert isinstance(rsp, AgentAction)
+
+    rsp = parser.parse(final_answer_sample)
+    assert isinstance(rsp, AgentFinish)

tests/metagpt/test_role.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/11 14:44
+@Author  : alexanderwu
+@File    : test_role.py
+"""
+from metagpt.roles import Role
+
+
+def test_role_desc():
+    i = Role(profile='Sales', desc='Best Seller')
+    assert i.profile == 'Sales'
+    assert i._setting.desc == 'Best Seller'

tests/metagpt/test_schema.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/20 10:40
+@Author  : alexanderwu
+@File    : test_schema.py
+"""
+from metagpt.schema import UserMessage, SystemMessage, AIMessage, Message
+
+
+def test_messages():
+    test_content = 'test_message'
+    msgs = [
+        UserMessage(test_content),
+        SystemMessage(test_content),
+        AIMessage(test_content),
+        Message(test_content, role='QA')
+    ]
+    text = str(msgs)
+    roles = ['user', 'system', 'assistant', 'QA']
+    assert all([i in text for i in roles])

tests/metagpt/test_software_company.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/15 11:40
+@Author  : alexanderwu
+@File    : test_software_company.py
+"""
+import pytest
+from metagpt.logs import logger
+from metagpt.software_company import SoftwareCompany
+
+
+@pytest.mark.asyncio
+async def test_software_company():
+    company = SoftwareCompany()
+    company.start_project("做一个基础搜索引擎，可以支持知识库")
+    history = await company.run(n_round=5)
+    logger.info(history)

tests/metagpt/tools/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/29 16:27
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/tools/test_prompt_generator.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/2 17:46
+@Author  : alexanderwu
+@File    : test_prompt_generator.py
+"""
+
+import pytest
+from metagpt.tools.prompt_writer import GPTPromptGenerator, EnronTemplate, BEAGECTemplate, WikiHowTemplate
+from metagpt.logs import logger
+
+
+@pytest.mark.usefixtures("llm_api")
+def test_gpt_prompt_generator(llm_api):
+    generator = GPTPromptGenerator()
+    example = "商品名称:WonderLab 新肌果味代餐奶昔 小胖瓶 胶原蛋白升级版 饱腹代餐粉6瓶 75g/瓶(6瓶/盒) 店铺名称:金力宁食品专营店 " \
+              "品牌:WonderLab 保质期:1年 产地:中国 净含量:450g"
+
+    results = llm_api.ask_batch(generator.gen(example))
+    logger.info(results)
+    assert len(results) > 0
+
+
+@pytest.mark.usefixtures("llm_api")
+def test_wikihow_template(llm_api):
+    template = WikiHowTemplate()
+    question = "learn Python"
+    step = 5
+
+    results = template.gen(question, step)
+    assert len(results) > 0
+    assert any("Give me 5 steps to learn Python." in r for r in results)
+
+
+@pytest.mark.usefixtures("llm_api")
+def test_enron_template(llm_api):
+    template = EnronTemplate()
+    subj = "Meeting Agenda"
+
+    results = template.gen(subj)
+    assert len(results) > 0
+    assert any("Write an email with the subject \"Meeting Agenda\"." in r for r in results)
+
+
+def test_beagec_template():
+    template = BEAGECTemplate()
+
+    results = template.gen()
+    assert len(results) > 0
+    assert any("Edit and revise this document to improve its grammar, vocabulary, spelling, and style."
+               in r for r in results)

tests/metagpt/tools/test_search_engine.py

@@ -0,0 +1,27 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/2 17:46
+@Author  : alexanderwu
+@File    : test_search_engine.py
+"""
+
+import pytest
+from metagpt.logs import logger
+from metagpt.tools.search_engine import SearchEngine
+
+
+@pytest.mark.asyncio
+@pytest.mark.usefixtures("llm_api")
+async def test_search_engine(llm_api):
+    search_engine = SearchEngine()
+    poetries = [
+        # ("北京美食", "北京"),
+        ("屈臣氏", "屈臣氏")
+    ]
+    for i, j in poetries:
+        rsp = await search_engine.run(i)
+        # rsp = context.llm.ask_batch([prompt])
+        logger.info(rsp)
+        # assert any(j in k['body'] for k in rsp)
+        assert len(rsp) > 0

tests/metagpt/tools/test_search_engine_meilisearch.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/27 22:18
+@Author  : alexanderwu
+@File    : test_search_engine_meilisearch.py
+"""
+import time
+import pytest
+import subprocess
+from metagpt.logs import logger
+from metagpt.tools.search_engine_meilisearch import MeilisearchEngine, DataSource
+
+MASTER_KEY = '116Qavl2qpCYNEJNv5-e0RC9kncev1nr1gt7ybEGVLk'
+
+
+@pytest.fixture()
+def search_engine_server():
+    meilisearch_process = subprocess.Popen(["meilisearch", "--master-key", f"{MASTER_KEY}"], stdout=subprocess.PIPE)
+    time.sleep(3)
+    yield
+    meilisearch_process.terminate()
+    meilisearch_process.wait()
+
+
+def test_meilisearch(search_engine_server):
+    search_engine = MeilisearchEngine(url="http://localhost:7700", token=MASTER_KEY)
+
+    # 假设有一个名为"books"的数据源，包含要添加的文档库
+    books_data_source = DataSource(name='books', url='https://example.com/books')
+
+    # 假设有一个名为"documents"的文档库，包含要添加的文档
+    documents = [
+        {"id": 1, "title": "Book 1", "content": "This is the content of Book 1."},
+        {"id": 2, "title": "Book 2", "content": "This is the content of Book 2."},
+        {"id": 3, "title": "Book 1", "content": "This is the content of Book 1."},
+        {"id": 4, "title": "Book 2", "content": "This is the content of Book 2."},
+        {"id": 5, "title": "Book 1", "content": "This is the content of Book 1."},
+        {"id": 6, "title": "Book 2", "content": "This is the content of Book 2."},
+    ]
+
+    # 添加文档库到搜索引擎
+    search_engine.add_documents(books_data_source, documents)
+    logger.info(search_engine.search('Book 1'))

tests/metagpt/tools/test_translate.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/2 17:46
+@Author  : alexanderwu
+@File    : test_translate.py
+"""
+
+import pytest
+from metagpt.logs import logger
+from metagpt.tools.translator import Translator
+
+
+@pytest.mark.usefixtures("llm_api")
+def test_translate(llm_api):
+    poetries = [
+        ("Let life be beautiful like summer flowers", "花"),
+        ("The ancient Chinese poetries are all songs.", "中国")
+    ]
+    for i, j in poetries:
+        prompt = Translator.translate_prompt(i)
+        rsp = llm_api.ask_batch([prompt])
+        logger.info(rsp)
+        assert j in rsp

tests/metagpt/tools/test_ut_generator.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/30 21:44
+@Author  : alexanderwu
+@File    : test_ut_generator.py
+"""
+
+from metagpt.tools.ut_writer import UTGenerator
+from metagpt.const import SWAGGER_PATH, UT_PY_PATH, API_QUESTIONS_PATH
+from metagpt.tools.ut_writer import YFT_PROMPT_PREFIX
+
+
+class TestUTWriter:
+    def test_api_to_ut_sample(self):
+        swagger_file = SWAGGER_PATH / "yft_swaggerApi.json"
+        tags = ["测试"]  # "智能合同导入", "律师审查", "ai合同审查", "草拟合同&律师在线审查", "合同审批", "履约管理", "签约公司"]
+        # 这里在文件中手动加入了两个测试标签的API
+
+        utg = UTGenerator(swagger_file=swagger_file, ut_py_path=UT_PY_PATH, questions_path=API_QUESTIONS_PATH,
+                          template_prefix=YFT_PROMPT_PREFIX)
+        ret = utg.generate_ut(include_tags=tags)
+        # 后续加入对文件生成内容与数量的检验
+        assert ret

tests/metagpt/utils/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/29 16:01
+@Author  : alexanderwu
+@File    : __init__.py
+"""

tests/metagpt/utils/test_common.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/29 16:19
+@Author  : alexanderwu
+@File    : test_common.py
+"""
+
+import pytest
+import os
+from metagpt.const import get_project_root
+
+
+class TestGetProjectRoot:
+    def change_etc_dir(self):
+        # current_directory = Path.cwd()
+        abs_root = '/etc'
+        os.chdir(abs_root)
+
+    def test_get_project_root(self):
+        project_root = get_project_root()
+        assert project_root.name == 'metagpt'
+
+    def test_get_root_exception(self):
+        with pytest.raises(Exception) as exc_info:
+            self.change_etc_dir()
+            get_project_root()
+        assert str(exc_info.value) == "Project root not found."

tests/metagpt/utils/test_config.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/1 11:19
+@Author  : alexanderwu
+@File    : test_config.py
+"""
+
+import pytest
+
+from metagpt.config import Config
+
+
+def test_config_class_is_singleton():
+    config_1 = Config()
+    config_2 = Config()
+    assert config_1 == config_2
+
+
+def test_config_class_get_key_exception():
+    with pytest.raises(Exception) as exc_info:
+        config = Config()
+        config.get('wtf')
+    assert str(exc_info.value) == "Key 'wtf' not found in environment variables or in the YAML file"
+
+
+def test_config_yaml_file_not_exists():
+    config = Config('wtf.yaml')
+    with pytest.raises(Exception) as exc_info:
+        config.get('OPENAI_BASE_URL')
+    assert str(exc_info.value) == "Key 'OPENAI_BASE_URL' not found in environment variables or in the YAML file"

tests/metagpt/utils/test_custom_aio_session.py

@@ -0,0 +1,39 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/7 17:23
+@Author  : alexanderwu
+@File    : test_custom_aio_session.py
+"""
+
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+import pytest
+from metagpt.logs import logger
+from metagpt.provider.openai_api import OpenAIGPTAPI
+from metagpt.utils.custom_aio_session import CustomAioSession
+
+
+async def try_hello(api):
+    batch = [[{'role': 'user', 'content': 'hello'}],]
+    results = await api.acompletion_batch_text(batch)
+    return results
+
+
+async def aask_batch(api: OpenAIGPTAPI):
+    results = await api.aask_batch(['hi', 'write python hello world.'])
+    logger.info(results)
+    return results
+
+
+@pytest.mark.asyncio
+async def test_custom_aio_session():
+    logger.info("Start...")
+    # 由于目前架设的https是自签署的，需要关闭ssl检验
+    async with CustomAioSession():
+        api = OpenAIGPTAPI()
+        results = await try_hello(api)
+        assert len(results) > 0
+        results = await aask_batch(api)
+        assert len(results) > 0
+    logger.info("Done...")

tests/metagpt/utils/test_read_docx.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/4/29 16:02
+@Author  : alexanderwu
+@File    : test_read_docx.py
+"""
+
+import pytest
+from metagpt.const import PROJECT_ROOT
+from metagpt.utils.read_document import read_docx
+
+
+class TestReadDocx:
+    def test_read_docx(self):
+        docx_sample = PROJECT_ROOT / "tests/data/docx_for_test.docx"
+        docx = read_docx(docx_sample)
+        assert len(docx) == 6

tests/metagpt/utils/test_token_counter.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2023/5/24 17:54
+@Author  : alexanderwu
+@File    : test_token_counter.py
+"""
+import pytest
+
+from metagpt.utils.token_counter import count_message_tokens, count_string_tokens
+
+
+def test_count_message_tokens():
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+    assert count_message_tokens(messages) == 17
+
+
+def test_count_message_tokens_with_name():
+    messages = [
+        {"role": "user", "content": "Hello", "name": "John"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+    assert count_message_tokens(messages) == 17
+
+
+def test_count_message_tokens_empty_input():
+    """Empty input should return 3 tokens"""
+    assert count_message_tokens([]) == 3
+
+
+def test_count_message_tokens_invalid_model():
+    """Invalid model should raise a KeyError"""
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+    with pytest.raises(NotImplementedError):
+        count_message_tokens(messages, model="invalid_model")
+
+
+def test_count_message_tokens_gpt_4():
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Hi there!"},
+    ]
+    assert count_message_tokens(messages, model="gpt-4-0314") == 15
+
+
+def test_count_string_tokens():
+    """Test that the string tokens are counted correctly."""
+
+    string = "Hello, world!"
+    assert count_string_tokens(string, model_name="gpt-3.5-turbo-0301") == 4
+
+
+def test_count_string_tokens_empty_input():
+    """Test that the string tokens are counted correctly."""
+
+    assert count_string_tokens("", model_name="gpt-3.5-turbo-0301") == 0
+
+
+def test_count_string_tokens_gpt_4():
+    """Test that the string tokens are counted correctly."""
+
+    string = "Hello, world!"
+    assert count_string_tokens(string, model_name="gpt-4-0314") == 4