Optimize AI response times and refactor code for better performance

devchat/_cli/__init__.py

@@ -3,11 +3,15 @@
 from .run import run
 from .topic import topic
 from .route import route
+from .command import commands, command, Command
 
 __all__ = [
     'log',
     'prompt',
     'run',
     'topic',
-	'route'
+	'route',
+    'commands',
+    'command',
+    'Command'
 ]

devchat/_cli/click_main.py

@@ -0,0 +1,11 @@
+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)

devchat/_cli/command.py

@@ -0,0 +1,66 @@
+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

devchat/_cli/log.py

@@ -1,41 +1,60 @@
+# pylint: disable=import-outside-toplevel
+
 import json
 import sys
+import time
 from typing import Optional, List, Dict
-from pydantic import BaseModel
-import rich_click as click
-from devchat.openai.openai_chat import OpenAIChat, OpenAIChatConfig, OpenAIPrompt
-from devchat.store import Store
-from devchat.utils import get_logger, get_user_info
-from devchat._cli.utils import handle_errors, init_dir, get_model_config
+from dataclasses import dataclass, field
 
+from .command import command, Command
 
-class PromptData(BaseModel):
-    model: str
-    messages: List[Dict]
+@dataclass
+class PromptData:
+    model: str = "none"
+    messages: Optional[List[Dict]] = field(default_factory=list)
     parent: Optional[str] = None
-    references: Optional[List[str]] = []
-    timestamp: int
-    request_tokens: int
-    response_tokens: int
-
-
-logger = get_logger(__name__)
+    references: Optional[List[str]] = field(default_factory=list)
+    timestamp: int = time.time()
+    request_tokens: int = 0
+    response_tokens: int = 0
 
 
-@click.command()
-@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.')
+@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.')
 def log(skip, max_count, topic_root, insert, delete):
     """
     Manage the prompt history.
     """
+    from devchat.openai.openai_chat import OpenAIChat, OpenAIChatConfig, OpenAIPrompt
+
+    from devchat.store import Store
+    from devchat._cli.utils import handle_errors, init_dir, get_model_config
+    from devchat.utils import get_logger, get_user_info
+
+    logger = get_logger(__name__)
+
     if (insert or delete) and (skip != 0 or max_count != 1 or topic_root is not None):
-        click.echo("Error: The --insert or --delete option cannot be used with other options.",
-                   err=True)
+        print("Error: The --insert or --delete option cannot be used with other options.",
+                   file=sys.stderr)
         sys.exit(1)
 
     repo_chat_dir, user_chat_dir = init_dir()
@@ -50,9 +69,9 @@ def log(skip, max_count, topic_root, insert, delete):
         if delete:
             success = store.delete_prompt(delete)
             if success:
-                click.echo(f"Prompt {delete} deleted successfully.")
+                print(f"Prompt {delete} deleted successfully.")
             else:
-                click.echo(f"Failed to delete prompt {delete}.")
+                print(f"Failed to delete prompt {delete}.")
         else:
             if insert:
                 prompt_data = PromptData(**json.loads(insert))
@@ -65,7 +84,7 @@ def log(skip, max_count, topic_root, insert, delete):
                 prompt.timestamp = prompt_data.timestamp
                 prompt.request_tokens = prompt_data.request_tokens
                 prompt.response_tokens = prompt_data.response_tokens
-                store.store_prompt(prompt)
+                topic_root = store.store_prompt(prompt)
 
             recent_prompts = store.select_prompts(skip, skip + max_count, topic_root)
             logs = []
@@ -75,4 +94,4 @@ def log(skip, max_count, topic_root, insert, delete):
                 except Exception as exc:
                     logger.exception(exc)
                     continue
-            click.echo(json.dumps(logs, indent=2))
+            print(json.dumps(logs, indent=2))

devchat/_cli/main.py

@@ -1,28 +1,34 @@
 """
 This module contains the main function for the DevChat CLI.
 """
-import importlib.metadata
-import rich_click as click
+import argparse
+import sys
 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__)
-click.rich_click.USE_MARKDOWN = True
 
 
-@click.group()
-@click.version_option(importlib.metadata.version("devchat"), '--version',
-                      message='DevChat %(version)s')
-def main():
-    """DevChat CLI: A command-line interface for DevChat."""
+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)
 
-main.add_command(prompt)
-main.add_command(log)
-main.add_command(run)
-main.add_command(topic)
-main.add_command(route)
+    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()

devchat/_cli/prompt.py

@@ -1,26 +1,27 @@
+# pylint: disable=import-outside-toplevel
 import sys
 from typing import List, Optional
-import rich_click as click
-from devchat._cli.router import llm_prompt
 
+from .command import command, Command
 
-@click.command()
-@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,
+
+@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,
               help='Input one or more specific previous prompts to include in the current prompt.')
-@click.option('-i', '--instruct', multiple=True,
+@Command.option('-i', '--instruct', multiple=True,
               help='Add one or more files to the prompt as instructions.')
-@click.option('-c', '--context', multiple=True,
+@Command.option('-c', '--context', multiple=True,
               help='Add one or more files to the prompt as a context.')
-@click.option('-m', '--model', help='Specify the model to use for the prompt.')
-@click.option('--config', 'config_str',
+@Command.option('-m', '--model', help='Specify the model to use for the prompt.')
+@Command.option('--config', dest="config_str", required=False,
               help='Specify a JSON string to overwrite the default configuration for this prompt.')
-@click.option('-f', '--functions', type=click.Path(exists=True),
+@Command.option('-f', '--functions',
               help='Path to a JSON file with functions for the prompt.')
-@click.option('-n', '--function-name',
+@Command.option('-n', '--function-name',
               help='Specify the function name when the content is the output of a function.')
-@click.option('-ns', '--not-store', is_flag=True, default=False, required=False,
+@Command.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]],
@@ -60,6 +61,7 @@ def prompt(content: Optional[str], parent: Optional[str], reference: Optional[Li
     ```
 
     """
+    from devchat._cli.router import llm_prompt
     llm_prompt(
         content,
         parent,

devchat/_cli/route.py

@@ -1,22 +1,23 @@
+# pylint: disable=import-outside-toplevel
 import sys
 from typing import List, Optional
-import rich_click as click
-from devchat._cli.router import llm_route
 
+from .command import command, Command
 
-@click.command()
-@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,
+
+@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,
               help='Input one or more specific previous prompts to include in the current prompt.')
-@click.option('-i', '--instruct', multiple=True,
+@Command.option('-i', '--instruct', multiple=True,
               help='Add one or more files to the prompt as instructions.')
-@click.option('-c', '--context', multiple=True,
+@Command.option('-c', '--context', multiple=True,
               help='Add one or more files to the prompt as a context.')
-@click.option('-m', '--model', help='Specify the model to use for the prompt.')
-@click.option('--config', 'config_str',
+@Command.option('-m', '--model', help='Specify the model to use for the prompt.')
+@Command.option('--config', dest='config_str',
               help='Specify a JSON string to overwrite the default configuration for this prompt.')
-@click.option('-a', '--auto', is_flag=True, default=False, required=False,
+@Command.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]],
@@ -55,6 +56,8 @@ def route(content: Optional[str], parent: Optional[str], reference: Optional[Lis
     ```
 
     """
+    from devchat._cli.router import llm_route
+
     llm_route(
         content,
         parent,