From b2c7e62147f57e80361652d4ee6c2892f41a1bcb Mon Sep 17 00:00:00 2001
From: kagami <mingjing@merico.dev>
Date: Fri, 10 May 2024 20:41:06 +0800
Subject: [PATCH] Revive click for cli

---
 devchat/_cli/__init__.py   |  4 ---
 devchat/_cli/click_main.py | 11 -------
 devchat/_cli/command.py    | 66 --------------------------------------
 devchat/_cli/log.py        | 30 +++++------------
 devchat/_cli/main.py       | 29 ++++++-----------
 devchat/_cli/prompt.py     | 24 +++++++-------
 devchat/_cli/route.py      | 20 ++++++------
 devchat/_cli/run.py        | 24 +++++++-------
 devchat/_cli/topic.py      | 11 +++----
 tests/test_cli_log.py      | 32 +++++++++---------
 tests/test_cli_prompt.py   | 30 ++++++++---------
 tests/test_cli_topic.py    |  8 ++---
 12 files changed, 92 insertions(+), 197 deletions(-)
 delete mode 100644 devchat/_cli/click_main.py
 delete mode 100755 devchat/_cli/command.py

diff --git a/devchat/_cli/__init__.py b/devchat/_cli/__init__.py
index 1cce4ad4..29b351dd 100644
--- a/devchat/_cli/__init__.py
+++ b/devchat/_cli/__init__.py
@@ -4,7 +4,6 @@
 from .run import run
 from .topic import topic
 from .route import route
-from .command import commands, command, Command
 
 script_dir = os.path.dirname(os.path.realpath(__file__))
 os.environ['TIKTOKEN_CACHE_DIR'] = os.path.join(script_dir, '..', 'tiktoken_cache')
@@ -15,7 +14,4 @@
     'run',
     'topic',
 	'route',
-    'commands',
-    'command',
-    'Command'
 ]
diff --git a/devchat/_cli/click_main.py b/devchat/_cli/click_main.py
deleted file mode 100644
index 6027b69d..00000000
--- a/devchat/_cli/click_main.py
+++ /dev/null
@@ -1,11 +0,0 @@
-import click
-from .main import main as main_from_argparse  # 导入修改后的main函数
-
-@click.command(context_settings={
-    "ignore_unknown_options": True,
-    "allow_extra_args": True
-})
-@click.pass_context
-def click_main(ctx):
-    """调用基于argparse的CLI程序，传递参数列表。"""
-    main_from_argparse(ctx.args)
diff --git a/devchat/_cli/command.py b/devchat/_cli/command.py
deleted file mode 100755
index ef61e870..00000000
--- a/devchat/_cli/command.py
+++ /dev/null
@@ -1,66 +0,0 @@
-commands = {}
-
-class Command:
-    def __init__(self, name, help_text):
-        self.parser = None
-        self.func = None
-        self.name = name
-        self.help = help_text
-        self.options = []
-        commands[name] = self
-
-    def add_option(self, option):
-        self.options.append(option)
-
-    @staticmethod
-    def option(*args, **kwargs):
-        def decorator(func):
-            if not hasattr(func, 'command_args'):
-                setattr(func, 'command_args', [])
-            func.command_args.append(("option", args, kwargs))
-            return func
-        return decorator
-
-    @staticmethod
-    def argument(*args, **kwargs):
-        def decorator(func):
-            if not hasattr(func, 'command_args'):
-                setattr(func, 'command_args', [])
-            func.command_args.append(("argument", args, kwargs))
-            return func
-        return decorator
-
-    def register(self, subparsers):
-        self.parser = subparsers.add_parser(self.name, help=self.help)
-        for option_type, args, kwargs in self.options:
-            is_flag = kwargs.pop('is_flag', None)
-            if is_flag:
-                kwargs['action'] = 'store_true'  # 如果是标志，则设置此动作
-            else:
-                nargs = kwargs.pop('multiple', None)
-                if nargs:
-                    kwargs['action'] = 'append'  # 表示至少需要一个参数，或'*'允许零个参数
-            required = kwargs.pop('required', None)
-            if required is not None:
-                kwargs['required'] = required
-
-            if option_type == "option":
-                self.parser.add_argument(*args, **kwargs)
-            elif option_type == "argument":
-                self.parser.add_argument(*args, **kwargs)
-        self.parser.set_defaults(func=self.func)
-
-# 命令装饰器工厂，每个命令通过这个工厂创建
-# pylint: disable=redefined-builtin
-def command(name, help=""):
-    def decorator(func):
-        cmd = Command(name, help_text=help)
-        cmd.func = func  # 将处理函数直接赋值给 Command 实例
-
-        # 注册命令参数
-        for option in getattr(func, 'command_args', []):
-            cmd.add_option(option)
-
-        commands[name] = cmd  # 将命令实例添加到全局命令字典中
-        return func
-    return decorator
diff --git a/devchat/_cli/log.py b/devchat/_cli/log.py
index c7059259..ac476311 100755
--- a/devchat/_cli/log.py
+++ b/devchat/_cli/log.py
@@ -6,7 +6,7 @@
 from typing import Optional, List, Dict
 from dataclasses import dataclass, field
 
-from .command import command, Command
+import click
 
 @dataclass
 class PromptData:
@@ -19,27 +19,13 @@ class PromptData:
     response_tokens: int = 0
 
 
-@command('log', help='Process logs')
-@Command.option('--skip',
-                type=int,
-                default=0,
-                help='Skip number prompts before showing the prompt history.')
-@Command.option('-n',
-                '--max-count',
-                type=int,
-                default=1,
-                help='Limit the number of commits to output.')
-@Command.option('-t',
-                '--topic',
-                dest='topic_root',
-                default=None,
-                help='Hash of the root prompt of the topic to select prompts from.')
-@Command.option('--insert',
-                default=None,
-                help='JSON string of the prompt to insert into the log.')
-@Command.option('--delete',
-                default=None,
-                help='Hash of the leaf prompt to delete from the log.')
+@click.command(help='Process logs')
+@click.option('--skip', default=0, help='Skip number prompts before showing the prompt history.')
+@click.option('-n', '--max-count', default=1, help='Limit the number of commits to output.')
+@click.option('-t', '--topic', 'topic_root', default=None,
+              help='Hash of the root prompt of the topic to select prompts from.')
+@click.option('--insert', default=None, help='JSON string of the prompt to insert into the log.')
+@click.option('--delete', default=None, help='Hash of the leaf prompt to delete from the log.')
 def log(skip, max_count, topic_root, insert, delete):
     """
     Manage the prompt history.
diff --git a/devchat/_cli/main.py b/devchat/_cli/main.py
index 10bcc364..87b51674 100755
--- a/devchat/_cli/main.py
+++ b/devchat/_cli/main.py
@@ -1,34 +1,25 @@
 """
 This module contains the main function for the DevChat CLI.
 """
-import argparse
-import sys
+import click
+
 from devchat.utils import get_logger
-# pylint: disable=unused-import
 from devchat._cli import log
 from devchat._cli import prompt
 from devchat._cli import run
 from devchat._cli import topic
 from devchat._cli import route
-from devchat._cli import commands
 
 logger = get_logger(__name__)
 
 
-def main(argv=None):
-    if argv is None:
-        argv = sys.argv[1:]
-
-    parser = argparse.ArgumentParser(description="CLI tool")
-    subparsers = parser.add_subparsers(help='sub-command help')
-    for _1, cmd in commands.items():
-        cmd.register(subparsers)
+@click.group()
+def main():
+    """DevChat CLI: A command-line interface for DevChat."""
 
-    args = parser.parse_args(argv)
-    if hasattr(args, 'func'):
-        func_args = vars(args).copy()
-        del func_args['func']
 
-        args.func(**func_args)
-    else:
-        parser.print_help()
+main.add_command(prompt)
+main.add_command(log)
+main.add_command(run)
+main.add_command(topic)
+main.add_command(route)
diff --git a/devchat/_cli/prompt.py b/devchat/_cli/prompt.py
index 5829e60a..73a2aa17 100755
--- a/devchat/_cli/prompt.py
+++ b/devchat/_cli/prompt.py
@@ -2,26 +2,26 @@
 import sys
 from typing import List, Optional
 
-from .command import command, Command
+import click
 
 
-@command('prompt', help='Interact with the large language model (LLM).')
-@Command.argument('content')
-@Command.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.')
-@Command.option('-r', '--reference', multiple=True,
+@click.command(help='Interact with the large language model (LLM).')
+@click.argument('content', required=False)
+@click.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.')
+@click.option('-r', '--reference', multiple=True,
               help='Input one or more specific previous prompts to include in the current prompt.')
-@Command.option('-i', '--instruct', multiple=True,
+@click.option('-i', '--instruct', multiple=True,
               help='Add one or more files to the prompt as instructions.')
-@Command.option('-c', '--context', multiple=True,
+@click.option('-c', '--context', multiple=True,
               help='Add one or more files to the prompt as a context.')
-@Command.option('-m', '--model', help='Specify the model to use for the prompt.')
-@Command.option('--config', dest="config_str", required=False,
+@click.option('-m', '--model', help='Specify the model to use for the prompt.')
+@click.option('--config', 'config_str',
               help='Specify a JSON string to overwrite the default configuration for this prompt.')
-@Command.option('-f', '--functions',
+@click.option('-f', '--functions', type=click.Path(exists=True),
               help='Path to a JSON file with functions for the prompt.')
-@Command.option('-n', '--function-name',
+@click.option('-n', '--function-name',
               help='Specify the function name when the content is the output of a function.')
-@Command.option('-ns', '--not-store', is_flag=True, default=False, required=False,
+@click.option('-ns', '--not-store', is_flag=True, default=False, required=False,
               help='Do not save the conversation to the store.')
 def prompt(content: Optional[str], parent: Optional[str], reference: Optional[List[str]],
            instruct: Optional[List[str]], context: Optional[List[str]],
diff --git a/devchat/_cli/route.py b/devchat/_cli/route.py
index f2b596e1..106e0bcc 100755
--- a/devchat/_cli/route.py
+++ b/devchat/_cli/route.py
@@ -2,22 +2,22 @@
 import sys
 from typing import List, Optional
 
-from .command import command, Command
+import click
 
 
-@command('route', help='Route a prompt to the specified LLM')
-@Command.argument('content')
-@Command.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.')
-@Command.option('-r', '--reference', multiple=True,
+@click.command(help='Route a prompt to the specified LLM')
+@click.argument('content', required=False)
+@click.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.')
+@click.option('-r', '--reference', multiple=True,
               help='Input one or more specific previous prompts to include in the current prompt.')
-@Command.option('-i', '--instruct', multiple=True,
+@click.option('-i', '--instruct', multiple=True,
               help='Add one or more files to the prompt as instructions.')
-@Command.option('-c', '--context', multiple=True,
+@click.option('-c', '--context', multiple=True,
               help='Add one or more files to the prompt as a context.')
-@Command.option('-m', '--model', help='Specify the model to use for the prompt.')
-@Command.option('--config', dest='config_str',
+@click.option('-m', '--model', help='Specify the model to use for the prompt.')
+@click.option('--config', 'config_str',
               help='Specify a JSON string to overwrite the default configuration for this prompt.')
-@Command.option('-a', '--auto', is_flag=True, default=False, required=False,
+@click.option('-a', '--auto', is_flag=True, default=False, required=False,
               help='Answer question by function-calling.')
 def route(content: Optional[str], parent: Optional[str], reference: Optional[List[str]],
            instruct: Optional[List[str]], context: Optional[List[str]],
diff --git a/devchat/_cli/run.py b/devchat/_cli/run.py
index bf06491a..ab028413 100755
--- a/devchat/_cli/run.py
+++ b/devchat/_cli/run.py
@@ -1,26 +1,26 @@
 # pylint: disable=import-outside-toplevel
 from typing import List, Optional, Tuple
-from .command import command, Command
 
+import click
 
-@command('run',
+@click.command(
     help="The 'command' argument is the name of the command to run or get information about.")
-@Command.argument('command', nargs='?',  default='')
-@Command.option('--list', dest='list_flag', is_flag=True, default=False,
+@click.argument('command', required=False, default='')
+@click.option('--list', 'list_flag', is_flag=True, default=False,
               help='List all specified commands in JSON format.')
-@Command.option('--recursive', '-r', dest='recursive_flag', is_flag=True, default=True,
+@click.option('--recursive', '-r', 'recursive_flag', is_flag=True, default=True,
               help='List commands recursively.')
-@Command.option('--update-sys', dest='update_sys_flag', is_flag=True, default=False,
+@click.option('--update-sys', 'update_sys_flag', is_flag=True, default=False,
               help='Pull the `sys` command directory from the DevChat repository.')
-@Command.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.')
-@Command.option('--reference', multiple=True,
+@click.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.')
+@click.option('-r', '--reference', multiple=True,
               help='Input one or more specific previous prompts to include in the current prompt.')
-@Command.option('-i', '--instruct', multiple=True,
+@click.option('-i', '--instruct', multiple=True,
               help='Add one or more files to the prompt as instructions.')
-@Command.option('-c', '--context', multiple=True,
+@click.option('-c', '--context', multiple=True,
               help='Add one or more files to the prompt as a context.')
-@Command.option('-m', '--model', help='Specify the model to use for the prompt.')
-@Command.option('--config', dest='config_str',
+@click.option('-m', '--model', help='Specify the model to use for the prompt.')
+@click.option('--config', 'config_str',
               help='Specify a JSON string to overwrite the default configuration for this prompt.')
 # pylint: disable=redefined-outer-name
 def run(command: str, list_flag: bool, recursive_flag: bool, update_sys_flag: bool,
diff --git a/devchat/_cli/topic.py b/devchat/_cli/topic.py
index d89cd5d0..63c90e2c 100755
--- a/devchat/_cli/topic.py
+++ b/devchat/_cli/topic.py
@@ -1,12 +1,11 @@
 # pylint: disable=import-outside-toplevel
-from .command import command, Command
+import click
 
-
-@command('topic', help='Manage topics')
-@Command.option('--list', '-l', dest='list_topics', is_flag=True,
+@click.command(help='Manage topics')
+@click.option('--list', '-l', 'list_topics', is_flag=True,
               help='List topics in reverse chronological order.')
-@Command.option('--skip', default=0, help='Skip number of topics before showing the list.')
-@Command.option('-n', '--max-count', default=100, help='Limit the number of topics to output.')
+@click.option('--skip', default=0, help='Skip number of topics before showing the list.')
+@click.option('-n', '--max-count', default=100, help='Limit the number of topics to output.')
 def topic(list_topics: bool, skip: int, max_count: int):
     """
     Manage topics.
diff --git a/tests/test_cli_log.py b/tests/test_cli_log.py
index 1a201d4c..bb34fe6f 100755
--- a/tests/test_cli_log.py
+++ b/tests/test_cli_log.py
@@ -1,20 +1,20 @@
 import json
 from click.testing import CliRunner
 from devchat.utils import get_prompt_hash, get_content
-from devchat._cli.click_main import click_main
+from devchat._cli.main import main
 
 runner = CliRunner()
 
 
 def test_log_no_args(git_repo):  # pylint: disable=W0613
-    result = runner.invoke(click_main, ['log'])
+    result = runner.invoke(main, ['log'])
     assert result.exit_code == 0
     logs = json.loads(result.output)
     assert isinstance(logs, list)
 
 
 def test_log_with_skip_and_max_count(git_repo):  # pylint: disable=W0613
-    result = runner.invoke(click_main, ['log', '--skip', '1', '--max-count', '2'])
+    result = runner.invoke(main, ['log', '--skip', '1', '--max-count', '2'])
     assert result.exit_code == 0
     logs = json.loads(result.output)
     assert isinstance(logs, list)
@@ -52,7 +52,7 @@ def test_tokens_with_log(git_repo):  # pylint: disable=W0613
     # Group 1
     config_str = '--config={ "stream": false, "temperature": 0 }'
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '--model=gpt-3.5-turbo', config_str, request1]
     )
     assert result.exit_code == 0
@@ -60,30 +60,30 @@ def test_tokens_with_log(git_repo):  # pylint: disable=W0613
 
     config_str = '--config={ "stream": true, "temperature": 0 }'
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '--model=gpt-3.5-turbo', config_str, request1]
     )
     assert result.exit_code == 0
     parent2 = get_prompt_hash(result.output)
 
-    result = runner.invoke(click_main, ['log', '-n', '2'])
+    result = runner.invoke(main, ['log', '-n', '2'])
     logs = json.loads(result.output)
     assert logs[0]["hash"] == parent2
 
     # Group 2
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '--model=gpt-3.5-turbo', config_str, '-p', parent1, request2]
     )
     assert result.exit_code == 0
 
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '--model=gpt-3.5-turbo', config_str, '-p', parent2, request2]
     )
     assert result.exit_code == 0
 
-    result = runner.invoke(click_main, ['log', '-n', '2'])
+    result = runner.invoke(main, ['log', '-n', '2'])
     logs = json.loads(result.output)
     assert len(logs) > 0
 
@@ -106,7 +106,7 @@ def test_log_insert(git_repo):  # pylint: disable=W0613
         "response_tokens": 100
     }"""
     result = runner.invoke(
-        click_main,
+        main,
         ['log', '--insert', chat1]
     )
     assert result.exit_code == 0
@@ -129,7 +129,7 @@ def test_log_insert(git_repo):  # pylint: disable=W0613
         "response_tokens": 200
     }"""
     result = runner.invoke(
-        click_main,
+        main,
         ['log', '--insert', chat2]
     )
     assert result.exit_code == 0
@@ -153,19 +153,19 @@ def test_log_insert(git_repo):  # pylint: disable=W0613
         "response_tokens": 300
     }}"""
     result = runner.invoke(
-        click_main,
+        main,
         ['log', '--insert', chat3]
     )
     assert result.exit_code == 0
     prompt3 = json.loads(result.output)[0]
     assert prompt3['parent'] == prompt1['hash']
 
-    result = runner.invoke(click_main, ['log', '-n', '3', '-t', prompt1['hash']])
+    result = runner.invoke(main, ['log', '-n', '3', '-t', prompt1['hash']])
     logs = json.loads(result.output)
     assert logs[0]['hash'] == prompt3['hash']
     assert logs[1]['hash'] == prompt1['hash']
 
-    result = runner.invoke(click_main, ['topic', '--list'])
+    result = runner.invoke(main, ['topic', '--list'])
     topics = json.loads(result.output)
     prompt_hashes = [prompt1['hash'], prompt2['hash']]
     topics = [topic for topic in topics if topic['root_prompt']['hash'] in prompt_hashes]
@@ -173,14 +173,14 @@ def test_log_insert(git_repo):  # pylint: disable=W0613
     assert topics[1]['root_prompt']['hash'] == prompt2['hash']
 
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '-p', prompt2['hash'], 'Again, just reply the topic number.']
     )
     assert result.exit_code == 0
     content = get_content(result.output)
     assert content.strip() == "Topic 2" or content.strip() == "2"
 
-    result = runner.invoke(click_main, ['log', '-t', prompt2['hash'], '-n', '100'])
+    result = runner.invoke(main, ['log', '-t', prompt2['hash'], '-n', '100'])
     assert result.exit_code == 0
     logs = json.loads(result.output)
     assert len(logs) == 2
diff --git a/tests/test_cli_prompt.py b/tests/test_cli_prompt.py
index 2caa02db..2759c60b 100644
--- a/tests/test_cli_prompt.py
+++ b/tests/test_cli_prompt.py
@@ -3,7 +3,7 @@
 import pytest
 from click.testing import CliRunner
 from devchat.config import ConfigManager, GeneralModelConfig
-from devchat._cli.click_main import click_main
+from devchat._cli.main import main
 from devchat.utils import openai_response_tokens
 from devchat.utils import check_format, get_content, get_prompt_hash
 
@@ -11,13 +11,13 @@
 
 
 # def test_prompt_no_args(git_repo):  # pylint: disable=W0613
-#     result = runner.invoke(click_main, ['prompt'])
+#     result = runner.invoke(main, ['prompt'])
 #     assert result.exit_code == 0
 
 
 def test_prompt_with_content(git_repo):  # pylint: disable=W0613
     content = "What is the capital of France?"
-    result = runner.invoke(click_main, ['prompt', content])
+    result = runner.invoke(main, ['prompt', content])
     print(result.output)
     assert result.exit_code == 0
     assert check_format(result.output)
@@ -45,7 +45,7 @@ def test_prompt_with_temp_config_file(mock_home_dir):
         config_file.write(config_data)
 
     content = "What is the capital of Spain?"
-    result = runner.invoke(click_main, ['prompt', content])
+    result = runner.invoke(main, ['prompt', content])
     print(result.output)
     assert result.exit_code == 0
     assert check_format(result.output)
@@ -95,7 +95,7 @@ def fixture_functions_file(tmpdir):
 
 
 def test_prompt_with_instruct(git_repo, temp_files):  # pylint: disable=W0613
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-4',
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-4',
                                   '-i', temp_files[0], '-i', temp_files[1],
                                   "It is really scorching."])
     assert result.exit_code == 0
@@ -103,7 +103,7 @@ def test_prompt_with_instruct(git_repo, temp_files):  # pylint: disable=W0613
 
 
 def test_prompt_with_instruct_and_context(git_repo, temp_files):  # pylint: disable=W0613
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-4',
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-4',
                                   '-i', temp_files[0], '-i', temp_files[2],
                                   '--context', temp_files[3],
                                   "It is really scorching."])
@@ -113,7 +113,7 @@ def test_prompt_with_instruct_and_context(git_repo, temp_files):  # pylint: disa
 
 def test_prompt_with_functions(git_repo, functions_file):  # pylint: disable=W0613
     # call with -f option
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-3.5-turbo', '-f', functions_file,
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-3.5-turbo', '-f', functions_file,
                                   "What is the weather like in Boston?"])
     if result.exit_code:
         print(result.output)
@@ -124,7 +124,7 @@ def test_prompt_with_functions(git_repo, functions_file):  # pylint: disable=W06
     assert '"name": "get_current_weather"' in content
 
     # compare with no -f options
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-3.5-turbo',
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-3.5-turbo',
                                   'What is the weather like in Boston?'])
 
     content = get_content(result.output)
@@ -135,13 +135,13 @@ def test_prompt_with_functions(git_repo, functions_file):  # pylint: disable=W06
 
 def test_prompt_log_with_functions(git_repo, functions_file):  # pylint: disable=W0613
     # call with -f option
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-3.5-turbo', '-f', functions_file,
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-3.5-turbo', '-f', functions_file,
                                   'What is the weather like in Boston?'])
     if result.exit_code:
         print(result.output)
     assert result.exit_code == 0
     prompt_hash = get_prompt_hash(result.output)
-    result = runner.invoke(click_main, ['log', '-t', prompt_hash])
+    result = runner.invoke(main, ['log', '-t', prompt_hash])
 
     result_json = json.loads(result.output)
     assert result.exit_code == 0
@@ -167,7 +167,7 @@ def test_prompt_log_compatibility():
 
 # test prompt with function replay
 def test_prompt_with_function_replay(git_repo, functions_file):  # pylint: disable=W0613
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-3.5-turbo',
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-3.5-turbo',
                                   '-f', functions_file,
                                   '-n', 'get_current_weather',
                                   '{"temperature": "22", "unit": "celsius", "weather": "Sunny"}'])
@@ -178,7 +178,7 @@ def test_prompt_with_function_replay(git_repo, functions_file):  # pylint: disab
     assert 'sunny' in content or 'Sunny' in content
 
     prompt_hash = get_prompt_hash(result.output)
-    result = runner.invoke(click_main, ['prompt', '-m', 'gpt-3.5-turbo',
+    result = runner.invoke(main, ['prompt', '-m', 'gpt-3.5-turbo',
                                   '-p', prompt_hash,
                                   'what is the function tool name?'])
 
@@ -191,7 +191,7 @@ def test_prompt_with_function_replay(git_repo, functions_file):  # pylint: disab
 
 def test_prompt_without_repo(mock_home_dir):  # pylint: disable=W0613
     content = "What is the capital of France?"
-    result = runner.invoke(click_main, ['prompt', content])
+    result = runner.invoke(main, ['prompt', content])
     assert result.exit_code == 0
     assert check_format(result.output)
     assert "Paris" in result.output
@@ -215,7 +215,7 @@ def test_prompt_tokens_exceed_config(mock_home_dir):  # pylint: disable=W0613
     content = ""
     while openai_response_tokens({"content": content}, model) < max_input_tokens:
         content += "This is a test. Ignore what I say.\n"
-    result = runner.invoke(click_main, ['prompt', content])
+    result = runner.invoke(main, ['prompt', content])
     print(result.output)
     assert result.exit_code != 0
     assert "beyond limit" in result.output
@@ -243,6 +243,6 @@ def test_file_tokens_exceed_config(mock_home_dir, tmpdir):  # pylint: disable=W0
     content_file.write(content)
 
     input_str = "This is a test. Ignore what I say."
-    result = runner.invoke(click_main, ['prompt', '-c', str(content_file), input_str])
+    result = runner.invoke(main, ['prompt', '-c', str(content_file), input_str])
     assert result.exit_code != 0
     assert "beyond limit" in result.output
diff --git a/tests/test_cli_topic.py b/tests/test_cli_topic.py
index 02b19d26..595d6237 100644
--- a/tests/test_cli_topic.py
+++ b/tests/test_cli_topic.py
@@ -3,7 +3,7 @@
 import time
 from click.testing import CliRunner
 from devchat.utils import get_prompt_hash
-from devchat._cli.click_main import click_main
+from devchat._cli.main import main
 
 runner = CliRunner()
 
@@ -12,7 +12,7 @@ def test_topic_list(git_repo):  # pylint: disable=W0613
     request = "Complete the sequence 1, 1, 3, 5, 9, ( ). Reply the number only."
     sys.argv = ['prompt', '--model=gpt-3.5-turbo', request]
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '--model=gpt-3.5-turbo', request]
     )
     assert result.exit_code == 0
@@ -21,13 +21,13 @@ def test_topic_list(git_repo):  # pylint: disable=W0613
     time.sleep(3)
 
     result = runner.invoke(
-        click_main,
+        main,
         ['prompt', '--model=gpt-4', request]
     )
     assert result.exit_code == 0
     topic2 = get_prompt_hash(result.output)
 
-    result = runner.invoke(click_main, ['topic', '--list'])
+    result = runner.invoke(main, ['topic', '--list'])
     assert result.exit_code == 0
     topics = json.loads(result.output)
     assert len(topics) >= 2