Merge pull request #76 from devchat-ai/context-info

Show additional context info in summary
2024-03-15 16:56:57 +08:00 · 2024-03-15 16:56:57 +08:00 · f753fb7bb2
commit f753fb7bb2
parent ccc1d97c90 ef6554df5a
9 changed files with 163 additions and 138 deletions
--- a/unit_tests/assistants/directory_structure/relevant_file_finder.py
+++ b/unit_tests/assistants/directory_structure/relevant_file_finder.py
@ -16,7 +16,7 @@ from openai_util import create_chat_completion_content
 from tools.directory_viewer import ListViewer
 from tools.tiktoken_util import get_encoding

-MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-3.5-turbo"
+MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-4-turbo-preview"
 ENCODING = (
    get_encoding(DEFAULT_ENCODING)  # Use default encoding as an approximation
    if USE_USER_MODEL
--- a/unit_tests/assistants/rerank_files.py
+++ b/unit_tests/assistants/rerank_files.py
@ -8,7 +8,7 @@ from llm_conf import (
 )
 from openai_util import create_chat_completion_content

-MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-3.5-turbo"
+MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-4-turbo-preview"

 # ruff: noqa: E501

--- a/unit_tests/find_context.py
+++ b/unit_tests/find_context.py
@ -15,13 +15,14 @@ from tools.symbol_util import (

 sys.path.append(os.path.join(os.path.dirname(__file__), ".."))

-from libs.ide_services import IDEService, Location, SymbolNode
+from libs.ide_services import IDEService, Location, Range, SymbolNode


@dataclass
 class Context:
    file_path: str  # relative path to repo root
    content: str
+    range: Range

    def __hash__(self) -> int:
        return hash((self.file_path, self.content))
@ -39,7 +40,7 @@ def _extract_referenced_symbols_context(
    """
    referenced_symbols_context = defaultdict(list)
    func_content = func_to_test.func_content
-    referenced_symbols = []
+    referenced_symbols: List[SymbolNode] = []
    stack = [(s, 0) for s in symbols]

    while stack:
@ -60,7 +61,7 @@ def _extract_referenced_symbols_context(
    # Get the content of the symbols
    for s in referenced_symbols:
        content = get_symbol_content(s, file_content=func_to_test.file_content)
-        context = Context(file_path=func_to_test.file_path, content=content)
+        context = Context(file_path=func_to_test.file_path, content=content, range=s.range)
        referenced_symbols_context[s.name].append(context)
    return referenced_symbols_context

@ -106,7 +107,7 @@ def _find_children_symbols_type_def_context(
            for t, _ in targets:
                content = get_symbol_content(t, abspath=loc.abspath)
                relpath = os.path.relpath(loc.abspath, func_to_test.repo_root)
-                context = Context(file_path=relpath, content=content)
+                context = Context(file_path=relpath, content=content, range=t.range)
                type_defs[symbol_name].append(context)

    return type_defs
@ -168,7 +169,7 @@ def _extract_recommended_symbols_context(
            for t, _ in targets:
                content = get_symbol_content(t, abspath=loc.abspath)
                relpath = os.path.relpath(loc.abspath, func_to_test.repo_root)
-                context = Context(file_path=relpath, content=content)
+                context = Context(file_path=relpath, content=content, range=t.range)
                recommended_symbols[symbol_name].append(context)

    return recommended_symbols
--- a/unit_tests/locales/base.pot
+++ b/unit_tests/locales/base.pot
@ -8,7 +8,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-03-11 22:42+0800\n"
+"POT-Creation-Date: 2024-03-15 16:39+0800\n"
 "PO-Revision-Date: YEAR-MO-DA HO:MI+ZONE\n"
 "Last-Translator: FULL NAME <EMAIL@ADDRESS>\n"
 "Language-Team: LANGUAGE <LL@li.org>\n"
@ -17,77 +17,75 @@ msgstr ""
 "Content-Type: text/plain; charset=CHARSET\n"
 "Content-Transfer-Encoding: 8bit\n"

-#: main.py:79
+#: main.py:107
 msgid "Analyzing the function and current unit tests..."
 msgstr ""

-#: main.py:116
+#: main.py:146
 msgid "Select test cases to generate"
 msgstr ""

-#: main.py:121
+#: main.py:151
 msgid ""
 "You can add more test cases here\n"
 "(Multiple cases can be separated by line breaks)"
 msgstr ""

-#: main.py:127
+#: main.py:157
 msgid ""
 "Edit reference test file\n"
 "(Multiple files can be separated by line breaks)"
 msgstr ""

-#: main.py:134
+#: main.py:164
 msgid ""
 "Write your customized requirements(prompts) for tests here.\n"
 "(For example, what testing framework to use.)"
 msgstr ""

-#: main.py:153
+#: main.py:183
 msgid "No test case is selected. Quit generating tests."
 msgstr ""

-#: main.py:176
+#: main.py:220
 msgid "Will generate tests for the following cases."
 msgstr ""

-#: main.py:179
+#: main.py:223
 msgid ""
 "\n"
 "Test cases:"
 msgstr ""

-#: main.py:186
+#: main.py:230
 msgid ""
 "\n"
 "No valid reference file is provided. Will not use reference to generate "
 "tests."
 msgstr ""

-#: main.py:191
+#: main.py:235
 msgid ""
 "\n"
 "Will use the following reference files to generate tests."
 msgstr ""

-#: main.py:192
-msgid ""
-"\n"
-"Valid reference files:"
-msgstr ""
-
-#: main.py:197
-msgid ""
-"\n"
-"Invalid files:"
-msgstr ""
-
-#: main.py:201
+#: main.py:245
 msgid ""
 "\n"
 "Customized requirements(prompts):"
 msgstr ""

-#: main.py:330
+#: main.py:249
+msgid "No customized requirments."
+msgstr ""
+
+#: main.py:252
+msgid ""
+"\n"
+"Additional context:"
+msgstr ""
+
+#: main.py:348
 msgid "The function's size surpasses AI's context capacity."
 msgstr ""
--- a/unit_tests/locales/en/base.po
+++ b/unit_tests/locales/en/base.po
@ -7,7 +7,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-03-11 22:42+0800\n"
+"POT-Creation-Date: 2024-03-15 16:39+0800\n"
 "PO-Revision-Date: 2023-12-24 16:51+0800\n"
 "Last-Translator: kagami <mingjing@merico.dev>\n"
 "Language-Team: English\n"
@ -17,78 +17,76 @@ msgstr ""
 "Content-Transfer-Encoding: 8bit\n"
 "Plural-Forms: nplurals=2; plural=(n != 1);\n"

-#: main.py:79
+#: main.py:107
 msgid "Analyzing the function and current unit tests..."
 msgstr ""

-#: main.py:116
+#: main.py:146
 msgid "Select test cases to generate"
 msgstr ""

-#: main.py:121
+#: main.py:151
 msgid ""
 "You can add more test cases here\n"
 "(Multiple cases can be separated by line breaks)"
 msgstr ""

-#: main.py:127
+#: main.py:157
 msgid ""
 "Edit reference test file\n"
 "(Multiple files can be separated by line breaks)"
 msgstr ""

-#: main.py:134
+#: main.py:164
 msgid ""
 "Write your customized requirements(prompts) for tests here.\n"
 "(For example, what testing framework to use.)"
 msgstr ""

-#: main.py:153
+#: main.py:183
 msgid "No test case is selected. Quit generating tests."
 msgstr ""

-#: main.py:176
+#: main.py:220
 msgid "Will generate tests for the following cases."
 msgstr ""

-#: main.py:179
+#: main.py:223
 msgid ""
 "\n"
 "Test cases:"
 msgstr ""

-#: main.py:186
+#: main.py:230
 msgid ""
 "\n"
 "No valid reference file is provided. Will not use reference to generate "
 "tests."
 msgstr ""

-#: main.py:191
+#: main.py:235
 msgid ""
 "\n"
 "Will use the following reference files to generate tests."
 msgstr ""

-#: main.py:192
-msgid ""
-"\n"
-"Valid reference files:"
-msgstr ""
-
-#: main.py:197
-msgid ""
-"\n"
-"Invalid files:"
-msgstr ""
-
-#: main.py:201
+#: main.py:245
 msgid ""
 "\n"
 "Customized requirements(prompts):"
 msgstr ""

-#: main.py:330
+#: main.py:249
+msgid "No customized requirments."
+msgstr ""
+
+#: main.py:252
+msgid ""
+"\n"
+"Additional context:"
+msgstr ""
+
+#: main.py:348
 msgid "The function's size surpasses AI's context capacity."
 msgstr ""

--- a/unit_tests/locales/zh/LC_MESSAGES/base.mo
+++ b/unit_tests/locales/zh/LC_MESSAGES/base.mo
--- a/unit_tests/locales/zh/base.po
+++ b/unit_tests/locales/zh/base.po
@ -7,7 +7,7 @@ msgid ""
 msgstr ""
 "Project-Id-Version: PACKAGE VERSION\n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-03-11 22:42+0800\n"
+"POT-Creation-Date: 2024-03-15 16:39+0800\n"
 "PO-Revision-Date: 2023-12-24 16:51+0800\n"
 "Last-Translator: kagami <mingjing@merico.dev>\n"
 "Language-Team: Chinese\n"
@ -16,15 +16,15 @@ msgstr ""
 "Content-Type: text/plain; charset=UTF-8\n"
 "Content-Transfer-Encoding: 8bit\n"

-#: main.py:79
+#: main.py:107
 msgid "Analyzing the function and current unit tests..."
 msgstr "正在分析目标函数和项目中现有的单元测试……"

-#: main.py:116
+#: main.py:146
 msgid "Select test cases to generate"
 msgstr "选择测试用例以生成单元测试"

-#: main.py:121
+#: main.py:151
 msgid ""
 "You can add more test cases here\n"
 "(Multiple cases can be separated by line breaks)"
@ -32,7 +32,7 @@ msgstr ""
 "可在输入框中添加更多测试用例\n"
 "（多个测试用例用换行分隔）"

-#: main.py:127
+#: main.py:157
 msgid ""
 "Edit reference test file\n"
 "(Multiple files can be separated by line breaks)"
@ -40,7 +40,7 @@ msgstr ""
 "参考的测试文件\n"
 "（多个文件路径用换行分隔）"

-#: main.py:134
+#: main.py:164
 msgid ""
 "Write your customized requirements(prompts) for tests here.\n"
 "(For example, what testing framework to use.)"
@ -48,15 +48,15 @@ msgstr ""
 "可在以下输入框中填写自定义的测试需求（提示词）\n"
 "（例如指定特定的测试框架）"

-#: main.py:153
+#: main.py:183
 msgid "No test case is selected. Quit generating tests."
 msgstr "未选择测试用例，退出生成单元测试。"

-#: main.py:176
+#: main.py:220
 msgid "Will generate tests for the following cases."
 msgstr "将为以下用例生成测试。"

-#: main.py:179
+#: main.py:223
 msgid ""
 "\n"
 "Test cases:"
@ -64,7 +64,7 @@ msgstr ""
 "\n"
 "测试用例："

-#: main.py:186
+#: main.py:230
 msgid ""
 "\n"
 "No valid reference file is provided. Will not use reference to generate "
@ -73,7 +73,7 @@ msgstr ""
 "\n"
 "没有提供合法的文件，生成单元测试时将不使用参考。"

-#: main.py:191
+#: main.py:235
 msgid ""
 "\n"
 "Will use the following reference files to generate tests."
@ -81,23 +81,7 @@ msgstr ""
 "\n"
 "将参考以下文件生成单元测试。"

-#: main.py:192
-msgid ""
-"\n"
-"Valid reference files:"
-msgstr ""
-"\n"
-"参考文件："
-
-#: main.py:197
-msgid ""
-"\n"
-"Invalid files:"
-msgstr ""
-"\n"
-"不合法的文件："
-
-#: main.py:201
+#: main.py:245
 msgid ""
 "\n"
 "Customized requirements(prompts):"
@ -105,10 +89,36 @@ msgstr ""
 "\n"
 "自定义测试需求（提示词）："

-#: main.py:330
+#: main.py:249
+msgid "No customized requirments."
+msgstr "未添加自定义需求。"
+
+#: main.py:252
+msgid ""
+"\n"
+"Additional context:"
+msgstr ""
+"\n"
+"补充的上下文："
+
+#: main.py:348
 msgid "The function's size surpasses AI's context capacity."
 msgstr "由于当前函数过大，AI暂时无法处理。"

+#~ msgid ""
+#~ "\n"
+#~ "Valid reference files:"
+#~ msgstr ""
+#~ "\n"
+#~ "参考文件："
+
+#~ msgid ""
+#~ "\n"
+#~ "Invalid files:"
+#~ msgstr ""
+#~ "\n"
+#~ "不合法的文件："
+
 #~ msgid "proposed cases"
 #~ msgstr "建议的测试用例"

--- a/unit_tests/main.py
+++ b/unit_tests/main.py
@ -52,15 +52,43 @@ class UnitTestsWorkflow:
        contexts = set()
        for _, v in symbol_context.items():
            contexts.update(v)
+        contexts = list(contexts)

-        cases, files = self.step2_propose_cases_and_reference_files(list(contexts))
+        cases, files = self.step2_propose_cases_and_reference_files(contexts)

        res = self.step3_user_interaction(cases, files)
        cases = res[0]
        files = res[1]
        requirements = res[2]

-        self.step4_write_and_print_tests(cases, files, list(contexts), requirements)
+        self.step4_print_test_summary(cases, files, requirements, contexts)
+
+        self.step5_write_and_print_tests(cases, files, contexts, requirements)
+
+    def step1_find_symbol_context(self) -> Dict[str, List[Context]]:
+        symbol_context = find_symbol_context_by_static_analysis(
+            self.func_to_test, self.tui_lang.chat_language
+        )
+
+        known_context_for_llm: List[Context] = []
+        if self.func_to_test.container_content is not None:
+            known_context_for_llm.append(
+                Context(
+                    file_path=self.func_to_test.file_path,
+                    content=self.func_to_test.container_content,
+                )
+            )
+        known_context_for_llm += list(
+            {item for sublist in list(symbol_context.values()) for item in sublist}
+        )
+
+        recommended_context = find_symbol_context_of_llm_recommendation(
+            self.func_to_test, known_context_for_llm
+        )
+
+        symbol_context.update(recommended_context)
+
+        return symbol_context

    def step2_propose_cases_and_reference_files(
        self,
@ -99,7 +127,9 @@ class UnitTestsWorkflow:
        return test_cases, reference_files

    def step3_user_interaction(
-        self, test_cases: List[str], reference_files: List[str]
+        self,
+        test_cases: List[str],
+        reference_files: List[str],
    ) -> Tuple[List[str], List[str], str]:
        """
        Edit test cases and reference files by user.
@ -172,7 +202,21 @@ class UnitTestsWorkflow:
        )
        self.local_cache.set("user_requirements", requirements)

-        # Print summary
+        return cases, valid_files, requirements
+
+    # Tuple[List[str], List[str], str]:
+    def step4_print_test_summary(
+        self,
+        cases: List[str],
+        valid_files: List[str],
+        requirements: str,
+        contexts: List[Context],
+    ):
+        """
+        Print the summary message in Step
+        """
+        _i = get_translation(self.tui_lang)
+
        title = _i("Will generate tests for the following cases.")
        lines = []

@ -189,61 +233,35 @@ class UnitTestsWorkflow:
            )
        else:
            lines.append(_i("\nWill use the following reference files to generate tests."))
-            lines.append(_i("\nValid reference files:"))
+            # lines.append(_i("\nValid reference files:"))
            width = len(str(len(valid_files)))
            lines.extend([f"{(i+1):>{width}}. {f}" for i, f in enumerate(valid_files)])

-        if invalid_files:
-            lines.append(_i("\nInvalid files:"))
-            width = len(str(len(invalid_files)))
-            lines.extend([f"{(i+1):>{width}}. {f}" for i, f in enumerate(invalid_files)])
+        # if invalid_files:
+        #     lines.append(_i("\nInvalid files:"))
+        #     width = len(str(len(invalid_files)))
+        #     lines.extend([f"{(i+1):>{width}}. {f}" for i, f in enumerate(invalid_files)])

        lines.append(_i("\nCustomized requirements(prompts):"))
-        lines.append(requirements)
+        if requirements.strip():
+            lines.append(requirements)
+        else:
+            lines.append(_i("No customized requirments."))
+
+        if contexts:
+            lines.append(_i("\nAdditional context:"))
+            width = len(str(len(contexts)))
+            lines.extend(
+                [
+                    f"{(i+1):>{width}}. {c.file_path}:{c.range.start.line+1}-{c.range.end.line+1}"
+                    for i, c in enumerate(contexts)
+                ]
+            )

        with Step(title):
            print("\n".join(lines), flush=True)

-        return cases, valid_files, requirements
-
-    def step1_find_symbol_context(self) -> Dict[str, List[Context]]:
-        symbol_context = find_symbol_context_by_static_analysis(
-            self.func_to_test, self.tui_lang.chat_language
-        )
-
-        # with Step("Symbol context"):
-        #     for k, v in symbol_context.items():
-        #         print(f"\n- {k}: ")
-        #         for item in v:
-        #             print(f"{item.file_path}\n{item.content}")
-
-        known_context_for_llm: List[Context] = []
-        if self.func_to_test.container_content is not None:
-            known_context_for_llm.append(
-                Context(
-                    file_path=self.func_to_test.file_path,
-                    content=self.func_to_test.container_content,
-                )
-            )
-        known_context_for_llm += list(
-            {item for sublist in list(symbol_context.values()) for item in sublist}
-        )
-
-        recommended_context = find_symbol_context_of_llm_recommendation(
-            self.func_to_test, known_context_for_llm
-        )
-
-        # with Step("Recommended context"):
-        #     for k, v in recommended_context.items():
-        #         print(f"\n- {k}: ")
-        #         for item in v:
-        #             print(f"{item.file_path}\n{item.content}")
-
-        symbol_context.update(recommended_context)
-
-        return symbol_context
-
-    def step4_write_and_print_tests(
+    def step5_write_and_print_tests(
        self,
        cases: List[str],
        ref_files: List[str],
--- a/unit_tests/propose_test.py
+++ b/unit_tests/propose_test.py
@ -16,13 +16,13 @@ from openai_util import create_chat_completion_content
 from prompts import PROPOSE_TEST_PROMPT
 from tools.tiktoken_util import get_encoding

-MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-3.5-turbo"
+MODEL = USER_LLM_MODEL if USE_USER_MODEL else "gpt-4-turbo-preview"
 ENCODING = (
    get_encoding(DEFAULT_ENCODING)  # Use default encoding as an approximation
    if USE_USER_MODEL
    else get_encoding("cl100k_base")
 )
-TOKEN_BUDGET = int(CONTEXT_SIZE.get(MODEL, DEFAULT_CONTEXT_SIZE) * 0.9)
+TOKEN_BUDGET = int(CONTEXT_SIZE.get(MODEL, DEFAULT_CONTEXT_SIZE) * 0.95)


 def _mk_user_msg(