diff --git a/devchat/_cli/__init__.py b/devchat/_cli/__init__.py index f1326833..e93def6a 100644 --- a/devchat/_cli/__init__.py +++ b/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' ] diff --git a/devchat/_cli/click_main.py b/devchat/_cli/click_main.py new file mode 100644 index 00000000..6027b69d --- /dev/null +++ b/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) diff --git a/devchat/_cli/command.py b/devchat/_cli/command.py new file mode 100755 index 00000000..ef61e870 --- /dev/null +++ b/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 diff --git a/devchat/_cli/log.py b/devchat/_cli/log.py old mode 100644 new mode 100755 index 6bd6c1db..c7059259 --- a/devchat/_cli/log.py +++ b/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)) diff --git a/devchat/_cli/main.py b/devchat/_cli/main.py old mode 100644 new mode 100755 index f5e60077..10bcc364 --- a/devchat/_cli/main.py +++ b/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() diff --git a/devchat/_cli/prompt.py b/devchat/_cli/prompt.py old mode 100644 new mode 100755 index de625471..5829e60a --- a/devchat/_cli/prompt.py +++ b/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, diff --git a/devchat/_cli/route.py b/devchat/_cli/route.py old mode 100644 new mode 100755 index c571e6b1..f2b596e1 --- a/devchat/_cli/route.py +++ b/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, diff --git a/devchat/_cli/router.py b/devchat/_cli/router.py old mode 100644 new mode 100755 index 0827c5d9..29b9e49f --- a/devchat/_cli/router.py +++ b/devchat/_cli/router.py @@ -1,19 +1,14 @@ +# pylint: disable=import-outside-toplevel import json import sys from typing import List, Optional -import rich_click as click -from devchat.engine import run_command, load_workflow_instruction -from devchat.assistant import Assistant -from devchat.openai.openai_chat import OpenAIChat, OpenAIChatConfig -from devchat.store import Store -from devchat.utils import parse_files -from devchat._cli.utils import handle_errors, init_dir, get_model_config -from devchat._cli.errors import MissContentInPromptException def _get_model_and_config( model: Optional[str], config_str: Optional[str]): + from devchat._cli.utils import init_dir, get_model_config + _1, user_chat_dir = init_dir() model, config = get_model_config(user_chat_dir, model) @@ -33,6 +28,9 @@ def _load_tool_functions(functions: Optional[str]): return None def _load_instruction_contents(content: str, instruct: Optional[List[str]]): + from devchat.engine import load_workflow_instruction + from devchat.utils import parse_files + instruct_contents = parse_files(instruct) command_instructions = load_workflow_instruction(content) if command_instructions is not None: @@ -46,10 +44,17 @@ def before_prompt(content: Optional[str], parent: Optional[str], reference: Opti model: Optional[str], config_str: Optional[str] = None, functions: Optional[str] = None, function_name: Optional[str] = None, not_store: Optional[bool] = False): + from devchat.assistant import Assistant + from devchat.openai.openai_chat import OpenAIChat, OpenAIChatConfig + from devchat.store import Store + from devchat.utils import parse_files + from devchat._cli.utils import init_dir + from devchat._cli.errors import MissContentInPromptException + repo_chat_dir, _1 = init_dir() if content is None: - content = click.get_text_stream('stdin').read() + content = sys.stdin.read() if content == '': raise MissContentInPromptException() @@ -83,26 +88,31 @@ def llm_prompt(content: Optional[str], parent: Optional[str], reference: Optiona model: Optional[str], config_str: Optional[str] = None, functions: Optional[str] = None, function_name: Optional[str] = None, not_store: Optional[bool] = False): + from devchat._cli.utils import handle_errors + with handle_errors(): _1, assistant, _3, = before_prompt( content, parent, reference, instruct, context, model, config_str, functions, function_name, not_store ) - click.echo(assistant.prompt.formatted_header()) + print(assistant.prompt.formatted_header()) for response in assistant.iterate_response(): - click.echo(response, nl=False) + print(response, end='', flush=True) def llm_commmand(content: Optional[str], parent: Optional[str], reference: Optional[List[str]], instruct: Optional[List[str]], context: Optional[List[str]], model: Optional[str], config_str: Optional[str] = None): + from devchat.engine import run_command + from devchat._cli.utils import handle_errors + with handle_errors(): model, assistant, content = before_prompt( content, parent, reference, instruct, context, model, config_str, None, None, True ) - click.echo(assistant.prompt.formatted_header()) + print(assistant.prompt.formatted_header()) command_result = run_command( model_name = model, history_messages = assistant.prompt.messages, @@ -112,8 +122,8 @@ def llm_commmand(content: Optional[str], parent: Optional[str], reference: Optio if command_result is not None: sys.exit(0) - click.echo("run command fail.") - click.echo(command_result) + print("run command fail.") + print(command_result) sys.exit(-1) @@ -121,12 +131,15 @@ def llm_route(content: Optional[str], parent: Optional[str], reference: Optional instruct: Optional[List[str]], context: Optional[List[str]], model: Optional[str], config_str: Optional[str] = None, auto: Optional[bool] = False): + from devchat.engine import run_command + from devchat._cli.utils import handle_errors + with handle_errors(): model, assistant, content = before_prompt( content, parent, reference, instruct, context, model, config_str, None, None, True ) - click.echo(assistant.prompt.formatted_header()) + print(assistant.prompt.formatted_header()) command_result = run_command( model_name = model, history_messages = assistant.prompt.messages, @@ -137,4 +150,4 @@ def llm_route(content: Optional[str], parent: Optional[str], reference: Optional sys.exit(command_result[0]) for response in assistant.iterate_response(): - click.echo(response, nl=False) + print(response, end='', flush=True) diff --git a/devchat/_cli/run.py b/devchat/_cli/run.py old mode 100644 new mode 100755 index 4bd63d51..bf06491a --- a/devchat/_cli/run.py +++ b/devchat/_cli/run.py @@ -1,38 +1,28 @@ -import json -import os -import stat -import shutil -import sys +# pylint: disable=import-outside-toplevel from typing import List, Optional, Tuple +from .command import command, Command -import rich_click as click -from devchat._cli.utils import init_dir, handle_errors, clone_git_repo -from devchat._cli.utils import download_and_extract_workflow -from devchat.engine import Namespace, CommandParser -from devchat.utils import get_logger -from devchat._cli.router import llm_commmand -logger = get_logger(__name__) - -@click.command( +@command('run', help="The 'command' argument is the name of the command to run or get information about.") -@click.argument('command', required=False, default='') -@click.option('--list', 'list_flag', is_flag=True, default=False, +@Command.argument('command', nargs='?', default='') +@Command.option('--list', dest='list_flag', is_flag=True, default=False, help='List all specified commands in JSON format.') -@click.option('--recursive', '-r', 'recursive_flag', is_flag=True, default=True, +@Command.option('--recursive', '-r', dest='recursive_flag', is_flag=True, default=True, help='List commands recursively.') -@click.option('--update-sys', 'update_sys_flag', is_flag=True, default=False, +@Command.option('--update-sys', dest='update_sys_flag', is_flag=True, default=False, help='Pull the `sys` command directory from the DevChat repository.') -@click.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.') -@click.option('-r', '--reference', multiple=True, +@Command.option('-p', '--parent', help='Input the parent prompt hash to continue the conversation.') +@Command.option('--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.') +# pylint: disable=redefined-outer-name def run(command: str, list_flag: bool, recursive_flag: bool, update_sys_flag: bool, parent: Optional[str], reference: Optional[List[str]], instruct: Optional[List[str]], context: Optional[List[str]], @@ -40,13 +30,24 @@ def run(command: str, list_flag: bool, recursive_flag: bool, update_sys_flag: bo """ Operate the workflow engine of DevChat. """ + import json + import os + import sys + from devchat._cli.utils import init_dir, handle_errors + + from devchat.engine import Namespace, CommandParser + from devchat.utils import get_logger + from devchat._cli.router import llm_commmand + + logger = get_logger(__name__) + _, user_chat_dir = init_dir() with handle_errors(): workflows_dir = os.path.join(user_chat_dir, 'workflows') if not os.path.exists(workflows_dir): os.makedirs(workflows_dir) if not os.path.isdir(workflows_dir): - click.echo(f"Error: Failed to find workflows directory: {workflows_dir}", err=True) + print(f"Error: Failed to find workflows directory: {workflows_dir}", file=sys.stderr) sys.exit(1) namespace = Namespace(workflows_dir) @@ -77,7 +78,7 @@ def run(command: str, list_flag: bool, recursive_flag: bool, update_sys_flag: bo 'description': cmd.description, 'path': cmd.path }) - click.echo(json.dumps(commands, indent=2)) + print(json.dumps(commands, indent=2)) return if command: @@ -103,6 +104,9 @@ def __onerror(func, path, _1): Usage : shutil.rmtree(path, onerror=onerror) """ + import os + import stat + # Check if file access issue if not os.access(path, os.W_OK): # Try to change the file to be writable (remove read-only flag) @@ -114,6 +118,8 @@ def __make_files_writable(directory): """ Recursively make all files in the directory writable. """ + import os + import stat for root, _1, files in os.walk(directory): for name in files: filepath = os.path.join(root, name) @@ -127,6 +133,14 @@ def _clone_or_pull_git_repo(target_dir: str, repo_urls: List[Tuple[str, str]], z :param target_dir: The path where the repository should be cloned. :param repo_urls: A list of possible Git repository URLs. """ + import os + import shutil + from devchat.utils import get_logger + from devchat._cli.utils import download_and_extract_workflow + from devchat._cli.utils import clone_git_repo + + logger = get_logger(__name__) + if shutil.which('git') is None: # If Git is not installed, download and extract the workflow for url in zip_urls: @@ -144,7 +158,7 @@ def _clone_or_pull_git_repo(target_dir: str, repo_urls: List[Tuple[str, str]], z shutil.rmtree(new_dir, onerror=__onerror) if os.path.exists(bak_dir): shutil.rmtree(bak_dir, onerror=__onerror) - click.echo(f'{target_dir} is already exists. Moved to {new_dir}') + print(f'{target_dir} is already exists. Moved to {new_dir}') clone_git_repo(bak_dir, repo_urls) try: shutil.move(target_dir, new_dir) @@ -159,4 +173,4 @@ def _clone_or_pull_git_repo(target_dir: str, repo_urls: List[Tuple[str, str]], z else: clone_git_repo(target_dir, repo_urls) - click.echo(f'Updated {target_dir}') + print(f'Updated {target_dir}') diff --git a/devchat/_cli/topic.py b/devchat/_cli/topic.py old mode 100644 new mode 100755 index 02854c34..d89cd5d0 --- a/devchat/_cli/topic.py +++ b/devchat/_cli/topic.py @@ -1,22 +1,21 @@ -import json -import rich_click as click -from devchat.store import Store -from devchat.openai import OpenAIChatConfig, OpenAIChat -from devchat.utils import get_logger -from devchat._cli.utils import init_dir, handle_errors, get_model_config +# pylint: disable=import-outside-toplevel +from .command import command, Command -logger = get_logger(__name__) - -@click.command() -@click.option('--list', '-l', 'list_topics', is_flag=True, +@command('topic', help='Manage topics') +@Command.option('--list', '-l', dest='list_topics', is_flag=True, help='List topics in reverse chronological order.') -@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.') +@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.') def topic(list_topics: bool, skip: int, max_count: int): """ Manage topics. """ + import json + from devchat.store import Store + from devchat.openai import OpenAIChatConfig, OpenAIChat + from devchat._cli.utils import init_dir, handle_errors, get_model_config + repo_chat_dir, user_chat_dir = init_dir() with handle_errors(): @@ -29,10 +28,4 @@ def topic(list_topics: bool, skip: int, max_count: int): if list_topics: topics = store.select_topics(skip, skip + max_count) - for topic_data in topics: - try: - topic_data.update({'root_prompt': topic_data['root_prompt'].shortlog()}) - except Exception as exc: - logger.exception(exc) - continue - click.echo(json.dumps(topics, indent=2)) + print(json.dumps(topics, indent=2)) diff --git a/devchat/_cli/utils.py b/devchat/_cli/utils.py old mode 100644 new mode 100755 index a750dd3e..22369c5a --- a/devchat/_cli/utils.py +++ b/devchat/_cli/utils.py @@ -1,25 +1,19 @@ +# pylint: disable=import-outside-toplevel from contextlib import contextmanager import os import sys import shutil from typing import Tuple, List, Optional, Any import zipfile -import requests -import openai -try: - from git import Repo, InvalidGitRepositoryError, GitCommandError -except Exception: - pass -import rich_click as click -from devchat.config import ConfigManager -from devchat.utils import find_root_dir, add_gitignore, setup_logger, get_logger -from devchat._cli.errors import MissContentInPromptException +from devchat._cli.errors import MissContentInPromptException +from devchat.utils import find_root_dir, add_gitignore, setup_logger, get_logger logger = get_logger(__name__) def download_and_extract_workflow(workflow_url, target_dir): + import requests # Download the workflow zip file response = requests.get(workflow_url, stream=True, timeout=10) # Downaload file to temp dir @@ -46,35 +40,44 @@ def download_and_extract_workflow(workflow_url, target_dir): @contextmanager def handle_errors(): + # import openai """Handle errors in the CLI.""" try: yield - except openai.APIError as error: - logger.exception(error) - click.echo(f"{type(error).__name__}: {error.type}", err=True) - sys.exit(1) + # except openai.APIError as error: + # logger.exception(error) + # print(f"{type(error).__name__}: {error.type}", file=sys.stderr) + # sys.exit(1) except MissContentInPromptException: - click.echo("Miss content in prompt command.", err=True) + print("Miss content in prompt command.", file=sys.stderr) sys.exit(1) except Exception as error: # import traceback # traceback.print_exc() logger.exception(error) - click.echo(f"{type(error).__name__}: {error}", err=True) + print(f"{type(error).__name__}: {error}", file=sys.stderr) sys.exit(1) +REPO_CHAT_DIR = None +USER_CHAT_DIR = None def init_dir() -> Tuple[str, str]: """ Initialize the chat directories. Returns: - repo_chat_dir: The chat directory in the repository. - user_chat_dir: The chat directory in the user's home. + REPO_CHAT_DIR: The chat directory in the repository. + USER_CHAT_DIR: The chat directory in the user's home. """ + # pylint: disable=global-statement + global REPO_CHAT_DIR + global USER_CHAT_DIR + if REPO_CHAT_DIR and USER_CHAT_DIR: + return REPO_CHAT_DIR, USER_CHAT_DIR + repo_dir, user_dir = find_root_dir() if not repo_dir and not user_dir: - click.echo(f"Error: Failed to find home for .chat: {repo_dir}, {user_dir}", err=True) + print(f"Error: Failed to find home for .chat: {repo_dir}, {user_dir}", file=sys.stderr) sys.exit(1) if not repo_dir: @@ -83,34 +86,34 @@ def init_dir() -> Tuple[str, str]: user_dir = repo_dir try: - repo_chat_dir = os.path.join(repo_dir, ".chat") - if not os.path.exists(repo_chat_dir): - os.makedirs(repo_chat_dir) + REPO_CHAT_DIR = os.path.join(repo_dir, ".chat") + if not os.path.exists(REPO_CHAT_DIR): + os.makedirs(REPO_CHAT_DIR) except Exception: pass try: - user_chat_dir = os.path.join(user_dir, ".chat") - if not os.path.exists(user_chat_dir): - os.makedirs(user_chat_dir) + USER_CHAT_DIR = os.path.join(user_dir, ".chat") + if not os.path.exists(USER_CHAT_DIR): + os.makedirs(USER_CHAT_DIR) except Exception: pass - if not os.path.isdir(repo_chat_dir): - repo_chat_dir = user_chat_dir - if not os.path.isdir(user_chat_dir): - user_chat_dir = repo_chat_dir - if not os.path.isdir(repo_chat_dir) or not os.path.isdir(user_chat_dir): - click.echo(f"Error: Failed to create {repo_chat_dir} and {user_chat_dir}", err=True) + if not os.path.isdir(REPO_CHAT_DIR): + REPO_CHAT_DIR = USER_CHAT_DIR + if not os.path.isdir(USER_CHAT_DIR): + USER_CHAT_DIR = REPO_CHAT_DIR + if not os.path.isdir(REPO_CHAT_DIR) or not os.path.isdir(USER_CHAT_DIR): + print(f"Error: Failed to create {REPO_CHAT_DIR} and {USER_CHAT_DIR}", file=sys.stderr) sys.exit(1) try: - setup_logger(os.path.join(repo_chat_dir, 'error.log')) - add_gitignore(repo_chat_dir, '*') + setup_logger(os.path.join(REPO_CHAT_DIR, 'error.log')) + add_gitignore(REPO_CHAT_DIR, '*') except Exception as exc: logger.error("Failed to setup logger or add .gitignore: %s", exc) - return repo_chat_dir, user_chat_dir + return REPO_CHAT_DIR, USER_CHAT_DIR def valid_git_repo(target_dir: str, valid_urls: List[str]) -> bool: @@ -121,6 +124,11 @@ def valid_git_repo(target_dir: str, valid_urls: List[str]) -> bool: :param valid_urls: A list of valid Git repository URLs. :return: True if the directory is a valid Git repository with a valid URL, False otherwise. """ + try: + from git import Repo, InvalidGitRepositoryError + except Exception: + pass + try: repo = Repo(target_dir) repo_url = next(repo.remote().urls) @@ -139,11 +147,16 @@ def clone_git_repo(target_dir: str, repo_urls: List[Tuple[str, str]]): :param target_dir: The path where the repository should be cloned. :param repo_urls: A list of possible Git repository URLs. """ + try: + from git import Repo, GitCommandError + except Exception: + pass + for url, branch in repo_urls: try: - click.echo(f"Cloning repository {url} to {target_dir}") + print(f"Cloning repository {url} to {target_dir}") Repo.clone_from(url, target_dir, branch=branch) - click.echo("Cloned successfully") + print("Cloned successfully") return except GitCommandError: logger.exception("Failed to clone repository %s to %s", url, target_dir) @@ -153,5 +166,6 @@ def clone_git_repo(target_dir: str, repo_urls: List[Tuple[str, str]]): def get_model_config(user_chat_dir: str, model: Optional[str] = None) -> Tuple[str, Any]: + from devchat.config import ConfigManager manager = ConfigManager(user_chat_dir) return manager.model_config(model) diff --git a/devchat/assistant.py b/devchat/assistant.py old mode 100644 new mode 100755 index 2dabcb17..10a94f00 --- a/devchat/assistant.py +++ b/devchat/assistant.py @@ -1,7 +1,8 @@ import json +import sys import time from typing import Optional, List, Iterator -import openai + from devchat.message import Message from devchat.chat import Chat from devchat.openai.openai_prompt import OpenAIPrompt @@ -72,13 +73,15 @@ def make_prompt(self, request: str, self._check_limit() # Add history to the prompt - for reference_hash in references: - prompt = self._store.get_prompt(reference_hash) - if not prompt: - logger.error("Reference %s not retrievable while making prompt.", reference_hash) - continue - self._prompt.references.append(reference_hash) - self._prompt.prepend_history(prompt, self.token_limit) + if references: + for reference_hash in references: + prompt = self._store.get_prompt(reference_hash) + if not prompt: + logger.error("Reference %s not retrievable while making prompt.", + reference_hash) + continue + self._prompt.references.append(reference_hash) + self._prompt.prepend_history(prompt, self.token_limit) if parent: self._prompt.parent = parent parent_hash = parent @@ -97,29 +100,35 @@ def iterate_response(self) -> Iterator[str]: Returns: Iterator[str]: An iterator over response strings from the chat API. """ + if self._chat.config.stream: created_time = int(time.time()) config_params = self._chat.config.dict(exclude_unset=True) for chunk in self._chat.stream_response(self._prompt): - if isinstance(chunk, openai.types.chat.chat_completion_chunk.ChatCompletionChunk): - chunk = chunk.dict() + try: + if hasattr(chunk, "dict"): + chunk = chunk.dict() if "function_call" in chunk["choices"][0]["delta"] and \ - not chunk["choices"][0]["delta"]["function_call"]: + not chunk["choices"][0]["delta"]["function_call"]: del chunk["choices"][0]["delta"]["function_call"] if not chunk["choices"][0]["delta"]["content"]: chunk["choices"][0]["delta"]["content"] = "" - if "id" not in chunk or "index" not in chunk["choices"][0]: - chunk["id"] = "chatcmpl-7vdfQI02x-" + str(created_time) - chunk["object"] = "chat.completion.chunk" - chunk["created"] = created_time - chunk["model"] = config_params["model"] - chunk["choices"][0]["index"] = 0 - chunk["choices"][0]["finish_reason"] = "stop" - if "role" not in chunk['choices'][0]['delta']: - chunk['choices'][0]['delta']['role']='assistant' - - delta = self._prompt.append_response(json.dumps(chunk)) - yield delta + if "id" not in chunk or "index" not in chunk["choices"][0]: + chunk["id"] = "chatcmpl-7vdfQI02x-" + str(created_time) + chunk["object"] = "chat.completion.chunk" + chunk["created"] = created_time + chunk["model"] = config_params["model"] + chunk["choices"][0]["index"] = 0 + chunk["choices"][0]["finish_reason"] = "stop" + if "role" not in chunk['choices'][0]['delta']: + chunk['choices'][0]['delta']['role']='assistant' + + delta = self._prompt.append_response(json.dumps(chunk)) + yield delta + except Exception as err: + print("receive:", chunk, file=sys.stderr, end="\n\n") + logger.error("Error while iterating response: %s, %s", err, str(chunk)) + raise err if not self._prompt.responses: raise RuntimeError("No responses returned from the chat API") if self._need_store: diff --git a/devchat/config.py b/devchat/config.py index f95671dd..cdb92ffd 100644 --- a/devchat/config.py +++ b/devchat/config.py @@ -9,7 +9,7 @@ class GeneralProviderConfig(BaseModel): api_key: Optional[str] api_base: Optional[str] -class ModelConfig(BaseModel, extra='forbid'): +class ModelConfig(BaseModel): max_input_tokens: Optional[int] = sys.maxsize provider: Optional[str] diff --git a/devchat/engine/command_parser.py b/devchat/engine/command_parser.py index 67bc6390..15ae2181 100644 --- a/devchat/engine/command_parser.py +++ b/devchat/engine/command_parser.py @@ -1,27 +1,27 @@ import os from typing import List, Dict, Optional -import oyaml as yaml from pydantic import BaseModel +import oyaml as yaml from .namespace import Namespace -class Parameter(BaseModel, extra='forbid'): - type: str - description: Optional[str] - enum: Optional[List[str]] - default: Optional[str] +class Parameter(BaseModel): + type: str = "string" + description: Optional[str] = None + enum: Optional[List[str]] = None + default: Optional[str] = None -class Command(BaseModel, extra='forbid'): +class Command(BaseModel): description: str - hint: Optional[str] - parameters: Optional[Dict[str, Parameter]] - input: Optional[str] - steps: Optional[List[Dict[str, str]]] - path: Optional[str] + hint: Optional[str] = None + parameters: Optional[Dict[str, Parameter]] = None + input: Optional[str] = None + steps: Optional[List[Dict[str, str]]] = None + path: Optional[str] = None -class CommandParser: +class CommandParser(): def __init__(self, namespace: Namespace): self.namespace = namespace @@ -37,18 +37,6 @@ def parse(self, name: str) -> Command: return None return parse_command(file_path) - def parse_json(self, name: str) -> str: - """ - Parse a command configuration file to JSON. - - :param name: The command name in the namespace. - :return: The JSON representation of the command. - """ - file_path = self.namespace.get_file(name, 'command.yml') - if not file_path: - return None - return parse_command(file_path).json() - def parse_command(file_path: str) -> Command: """ diff --git a/devchat/engine/util.py b/devchat/engine/util.py index 935cb0f9..6abc49e4 100644 --- a/devchat/engine/util.py +++ b/devchat/engine/util.py @@ -1,10 +1,9 @@ +# pylint: disable=import-outside-toplevel import os import sys import json from typing import List, Dict -import openai - from devchat._cli.utils import init_dir from devchat.utils import get_logger @@ -92,6 +91,7 @@ def make_function(command: Command, command_name: str): def select_function_by_llm( history_messages: List[Dict], tools: List[Dict], model: str = DEFAULT_MODEL ): + import openai client = openai.OpenAI( api_key=os.environ.get("OPENAI_API_KEY", None), base_url=os.environ.get("OPENAI_API_BASE", None) diff --git a/devchat/openai/http_openai.py b/devchat/openai/http_openai.py new file mode 100755 index 00000000..a7607c17 --- /dev/null +++ b/devchat/openai/http_openai.py @@ -0,0 +1,64 @@ +import http.client +import json +import sys + + +class LineReader: + def __init__(self, response): + self.response = response + + def __iter__(self): + return self + + def __next__(self): + line = self.response.readline() + if not line: + raise StopIteration + line = line.strip() + if not line: + return self.__next__() + line = line.decode('utf-8') + if not line.startswith("data:"): + print("Receive invalid line: {line}", end="\n\n", file=sys.stderr) + raise ValueError(f"Invalid line: {line}") + + if line[5:].strip() == "[DONE]": + raise StopIteration + try: + return json.loads(line[5:]) + except json.JSONDecodeError as err: + print(f"Error decoding JSON: {err}", end="\n\n", file=sys.stderr) + raise ValueError(f"Invalid line: {line}") from err + +def stream_response(connection: http.client.HTTPSConnection, data, headers): + connection.request("POST", "/v1/chat/completions", body=json.dumps(data), headers=headers) + response = connection.getresponse() + + if response.status != 200: + print(f"Error: {response.status} - {response.reason}") + return None + return LineReader(response=response) + + +def stream_request(api_key, api_base, data): + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {api_key}", + } + + if api_base.startswith("https://"): + url = api_base[8:] + elif api_base.startswith("http://"): + url = api_base[7:] + else: + print("Invalid API base URL", end="\n\n", file=sys.stderr) + raise ValueError("Invalid API base URL") + + url = url.split("/")[0] + + if api_base.startswith("https://"): + connection = http.client.HTTPSConnection(url) + else: + connection = http.client.HTTPConnection(url) + + return stream_response(connection, data, headers) diff --git a/devchat/openai/openai_chat.py b/devchat/openai/openai_chat.py old mode 100644 new mode 100755 index 81a424dc..225d9bfc --- a/devchat/openai/openai_chat.py +++ b/devchat/openai/openai_chat.py @@ -1,13 +1,14 @@ +# pylint: disable=import-outside-toplevel import json import os from typing import Optional, Union, List, Dict, Iterator from pydantic import BaseModel, Field -import openai + from devchat.chat import Chat from devchat.utils import get_user_info, user_id from .openai_message import OpenAIMessage from .openai_prompt import OpenAIPrompt - +from .http_openai import stream_request class OpenAIChatParameters(BaseModel, extra='ignore'): temperature: Optional[float] = Field(0, ge=0, le=2) @@ -61,6 +62,8 @@ def load_prompt(self, data: dict) -> OpenAIPrompt: return OpenAIPrompt(**data) def complete_response(self, prompt: OpenAIPrompt) -> str: + import openai + # Filter the config parameters with set values config_params = self.config.dict(exclude_unset=True) if prompt.get_functions(): @@ -82,6 +85,25 @@ def complete_response(self, prompt: OpenAIPrompt) -> str: return str(response) def stream_response(self, prompt: OpenAIPrompt) -> Iterator: + api_key=os.environ.get("OPENAI_API_KEY", None) + base_url=os.environ.get("OPENAI_API_BASE", "https://api.openai.com/v1/") + + if not os.environ.get("USE_TIKTOKEN", False) and base_url != "https://api.openai.com/v1/": + config_params = self.config.dict(exclude_unset=True) + if prompt.get_functions(): + config_params['functions'] = prompt.get_functions() + config_params['function_call'] = 'auto' + config_params['stream'] = True + + data = { + "messages":prompt.messages, + **config_params, + "timeout":180 + } + response = stream_request(api_key, base_url, data) + return response + import openai + # Filter the config parameters with set values config_params = self.config.dict(exclude_unset=True) if prompt.get_functions(): diff --git a/devchat/store.py b/devchat/store.py index ea382e5a..7e40a209 100644 --- a/devchat/store.py +++ b/devchat/store.py @@ -1,8 +1,8 @@ +# pylint: disable=import-outside-toplevel from dataclasses import asdict +import json import os from typing import List, Dict, Any, Optional -from xml.etree.ElementTree import ParseError -import networkx as nx from tinydb import TinyDB, where, Query from tinydb.table import Table from devchat.chat import Chat @@ -26,24 +26,72 @@ def __init__(self, store_dir: str, chat: Chat): os.makedirs(store_dir) self._graph_path = os.path.join(store_dir, 'prompts.graphml') + self._chat_list_path = os.path.join(store_dir, 'prompts_list.json') self._db_path = os.path.join(store_dir, 'prompts.json') self._chat = chat - if os.path.isfile(self._graph_path): + self._db = TinyDB(self._db_path) + self._db_meta = self._migrate_db() + self._topics_table = self._db.table('topics') + + if os.path.isfile(self._chat_list_path): + with open(self._chat_list_path, 'r', encoding="utf-8") as file: + self._chat_lists = json.loads(file.read()) + elif os.path.isfile(self._graph_path): + # convert old graphml to new json + from xml.etree.ElementTree import ParseError + import networkx as nx try: - self._graph = nx.read_graphml(self._graph_path) + graph = nx.read_graphml(self._graph_path) + + roots = [node for node in graph.nodes() if graph.out_degree(node) == 0] + + self._chat_lists = [] + for root in roots: + chat_list = [(root, graph.nodes[root]['timestamp'])] + + ancestors = nx.ancestors(graph, root) + for ancestor in ancestors: + chat_list.append((ancestor, graph.nodes[ancestor]['timestamp'])) + + self._chat_lists.append(chat_list) + + with open(self._chat_list_path, 'w', encoding="utf-8") as file: + file.write(json.dumps(self._chat_lists)) + + # rename graphml to json + os.rename(self._graph_path, self._graph_path + '.bak') + + # update topic table, add request and response fields + # new fields: user, date, request, responses, hash + visible_topics = self._topics_table.all() + for topic in visible_topics: + prompt = self.get_prompt(topic['root']) + if not prompt: + continue + self._update_topic_fields(topic, prompt) + self._topics_table.update(topic, doc_ids=[topic.doc_id]) + except ParseError as error: raise ValueError(f"Invalid file format for graph: {self._graph_path}") from error else: - self._graph = nx.DiGraph() - - self._db = TinyDB(self._db_path) - self._db_meta = self._migrate_db() - self._topics_table = self._db.table('topics') + self._chat_lists = [] if not self._topics_table or not self._topics_table.all(): self._initialize_topics_table() + def _update_topic_fields(self, topic, prompt): + topic['user'] = prompt.user_name + topic['date'] = prompt.timestamp + topic['request'] = prompt.request.content + topic['responses'] = prompt.responses[0].content if prompt.responses else "" + topic['hash'] = prompt.hash + if len(topic['request']) > 100: + topic['request'] = topic['request'][:100] + "..." + if len(topic['responses']) > 100: + topic['responses'] = topic['responses'][:100] + "..." + + def _migrate_db(self) -> Table: """ Migrate the database to the latest version. @@ -67,43 +115,52 @@ def transform(doc): return metadata def _initialize_topics_table(self): - roots = [node for node in self._graph.nodes() if self._graph.out_degree(node) == 0] - for root in roots: - ancestors = nx.ancestors(self._graph, root) - if not ancestors: - latest_time = self._graph.nodes[root]['timestamp'] - else: - latest_time = max(self._graph.nodes[node]['timestamp'] for node in ancestors) - self._topics_table.insert({ - 'root': root, - 'latest_time': latest_time, + for chat_list in self._chat_lists: + if not chat_list: + continue + + first = chat_list[0] + last = chat_list[-1] + + topic = { + 'root': first[0], + 'latest_time': last[1], 'title': None, 'hidden': False - }) + } - def _update_topics_table(self, prompt: Prompt): - if self._graph.in_degree(prompt.hash): - logger.error("Prompt %s not a leaf to update topics table", prompt.hash) + prompt = self.get_prompt(topic['root']) + if not prompt: + logger.error("Prompt %s not found while selecting from the store", topic['root']) + continue + self._update_topic_fields(topic, prompt) + self._topics_table.insert(topic) + + def _update_topics_table(self, prompt: Prompt): if prompt.parent: - for topic in self._topics_table.all(): - if topic['root'] not in self._graph: - self._graph.add_node(topic['root'], timestamp=topic['latest_time']) - logger.warning("Topic %s not found in graph but added", topic['root']) - if prompt.parent == topic['root'] or \ - prompt.parent in nx.ancestors(self._graph, topic['root']): - topic['latest_time'] = max(topic.get('latest_time', 0), prompt.timestamp) - self._topics_table.update(topic, doc_ids=[topic.doc_id]) + for chat_list in self._chat_lists: + if not chat_list: + continue + + if chat_list[-1][0] == prompt.hash: + topic_hash = chat_list[0][0] + topic = next((t for t in self._topics_table if t['root'] == topic_hash), None) + if topic: + topic['latest_time'] = max(topic.get('latest_time', 0), prompt.timestamp) + self._topics_table.update(topic, doc_ids=[topic.doc_id]) break else: - self._topics_table.insert({ + topic = { 'root': prompt.hash, 'latest_time': prompt.timestamp, 'title': None, 'hidden': False - }) + } + self._update_topic_fields(topic, prompt) + self._topics_table.insert(topic) - def store_prompt(self, prompt: Prompt): + def store_prompt(self, prompt: Prompt) -> str: """ Store a prompt in the store. @@ -116,24 +173,25 @@ def store_prompt(self, prompt: Prompt): self._db.insert(asdict(prompt)) # Add the prompt to the graph - self._graph.add_node(prompt.hash, timestamp=prompt.timestamp) - - # Add edges for parents and references - if prompt.parent: - if prompt.parent not in self._graph: - logger.error("Parent %s not found while Prompt %s is stored to graph store.", - prompt.parent, prompt.hash) - else: - self._graph.add_edge(prompt.hash, prompt.parent) - + topic_hash = None + for chat_list in self._chat_lists: + if not chat_list: + continue + if chat_list[-1][0] == prompt.parent: + chat_list.append((prompt.hash, prompt.timestamp)) + topic_hash = chat_list[0][0] + break + + if not topic_hash: + topic_hash = prompt.hash + self._chat_lists.append([(prompt.hash, prompt.timestamp)]) self._update_topics_table(prompt) - for reference_hash in prompt.references: - if reference_hash not in self._graph: - logger.error("Reference %s not found while Prompt %s is stored to graph store.", - reference_hash, prompt.hash) + with open(self._chat_list_path, 'w', encoding="utf-8") as file: + file.write(json.dumps(self._chat_lists)) + + return topic_hash - nx.write_graphml(self._graph, self._graph_path) def get_prompt(self, prompt_hash: str) -> Prompt: """ @@ -144,10 +202,6 @@ def get_prompt(self, prompt_hash: str) -> Prompt: Returns: Prompt: The retrieved prompt. None if the prompt is not found. """ - if prompt_hash not in self._graph: - logger.warning("Prompt %s not found while retrieving from graph store.", prompt_hash) - return None - # Retrieve the prompt object from TinyDB prompt_data = self._db.search(where('_hash') == prompt_hash) if not prompt_data: @@ -170,15 +224,26 @@ def select_prompts(self, start: int, end: int, topic: Optional[str] = None) -> L If end is greater than the number of all prompts, the list will contain prompts from start to the end of the list. """ - if topic: - ancestors = nx.ancestors(self._graph, topic) - nodes_with_data = [(node, self._graph.nodes[node]) for node in ancestors] + \ - [(topic, self._graph.nodes[topic])] - sorted_nodes = sorted(nodes_with_data, key=lambda x: x[1]['timestamp'], reverse=True) - else: - sorted_nodes = sorted(self._graph.nodes(data=True), - key=lambda x: x[1]['timestamp'], - reverse=True) + + if not topic: + last_time = 0 + for chat_list in self._chat_lists: + if chat_list and chat_list[-1][1] > last_time: + last_time = chat_list[-1][1] + topic = chat_list[0][0] + if not topic: + return [] + + sorted_nodes = [] + for chat_list in self._chat_lists: + if not chat_list: + continue + + if chat_list[0][0] != topic: + continue + + sorted_nodes = chat_list.copy() + sorted_nodes.reverse() prompts = [] for node in sorted_nodes[start:end]: @@ -207,12 +272,14 @@ def select_topics(self, start: int, end: int) -> List[Dict[str, Any]]: topics = [] for topic in sorted_topics[start:end]: - prompt = self.get_prompt(topic['root']) - if not prompt: - logger.error("Topic %s not found while selecting from the store", topic['root']) - continue topics.append({ - 'root_prompt': prompt, + 'root_prompt': { + 'hash': topic['root'], + 'user': topic['user'], + 'date': topic['date'], + 'request': topic['request'], + 'responses': [topic['responses']], + }, 'latest_time': topic['latest_time'], 'title': topic['title'], 'hidden': topic['hidden'], @@ -230,12 +297,23 @@ def delete_prompt(self, prompt_hash: str) -> bool: bool: True if the prompt is successfully deleted, False otherwise. """ # Check if the prompt is a leaf - if self._graph.in_degree(prompt_hash) != 0: - logger.error("Prompt %s is not a leaf, cannot be deleted.", prompt_hash) - return False + has_deleted = False + for chat_list in self._chat_lists: + if not chat_list: + continue + + if chat_list[-1][0] != prompt_hash: + continue - # Remove the prompt from the graph - self._graph.remove_node(prompt_hash) + has_deleted = True + chat_list.pop() + + # If the chat list is empty, remove it from the list of chat lists + if not chat_list: + self._chat_lists.remove(chat_list) + + if not has_deleted: + return False # Update the topics table self._topics_table.remove(where('root') == prompt_hash) @@ -244,7 +322,8 @@ def delete_prompt(self, prompt_hash: str) -> bool: self._db.remove(where('_hash') == prompt_hash) # Save the graph - nx.write_graphml(self._graph, self._graph_path) + with open(self._chat_list_path, 'w', encoding="utf-8") as file: + file.write(json.dumps(self._chat_lists)) return True diff --git a/devchat/utils.py b/devchat/utils.py index e8714e59..d6903c60 100644 --- a/devchat/utils.py +++ b/devchat/utils.py @@ -1,3 +1,4 @@ +# pylint: disable=import-outside-toplevel import logging import os import re @@ -7,27 +8,11 @@ from typing import List, Tuple, Optional import datetime import hashlib -import tiktoken -script_dir = os.path.dirname(os.path.realpath(__file__)) -os.environ['TIKTOKEN_CACHE_DIR'] = os.path.join(script_dir, 'tiktoken_cache') - -try: - encoding = tiktoken.get_encoding("cl100k_base") -except Exception: - from tiktoken import registry - from tiktoken.registry import _find_constructors - from tiktoken.core import Encoding - - def get_encoding(name: str): - _find_constructors() - constructor = registry.ENCODING_CONSTRUCTORS[name] - return Encoding(**constructor(), use_pure_python=True) - - encoding = get_encoding("cl100k_base") log_formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') - +# pylint: disable=invalid-name +encoding = None def setup_logger(file_path: Optional[str] = None): """Utility function to set up a global file log handler.""" @@ -213,20 +198,33 @@ def update_dict(dict_to_update, key, value) -> dict: return dict_to_update -def _count_tokens(encoding_tik: tiktoken.Encoding, string: str) -> int: - """ - Count the number of tokens in a string. - """ - try: - return len(encoding_tik.encode(string)) - except Exception: - word_count = len(re.findall(r'\w+', string)) - # Note: This is a rough estimate and may not be accurate - return int(word_count / 0.75) - - def openai_message_tokens(messages: dict, model: str) -> int: # pylint: disable=unused-argument """Returns the number of tokens used by a message.""" + if not os.environ.get("USE_TIKTOKEN", False): + return len(str(messages))/4 + + # pylint: disable=global-statement + global encoding + if not encoding: + import tiktoken + + script_dir = os.path.dirname(os.path.realpath(__file__)) + os.environ['TIKTOKEN_CACHE_DIR'] = os.path.join(script_dir, 'tiktoken_cache') + + try: + encoding = tiktoken.get_encoding("cl100k_base") + except Exception: + from tiktoken import registry + from tiktoken.registry import _find_constructors + from tiktoken.core import Encoding + + def get_encoding(name: str): + _find_constructors() + constructor = registry.ENCODING_CONSTRUCTORS[name] + return Encoding(**constructor(), use_pure_python=True) + + encoding = get_encoding("cl100k_base") + return len(encoding.encode(str(messages), disallowed_special=())) diff --git a/tests/test_cli_log.py b/tests/test_cli_log.py index 7a52fddc..1a201d4c 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.main import main +from devchat._cli.click_main import click_main runner = CliRunner() def test_log_no_args(git_repo): # pylint: disable=W0613 - result = runner.invoke(main, ['log']) + result = runner.invoke(click_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(main, ['log', '--skip', '1', '--max-count', '2']) + result = runner.invoke(click_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( - main, + click_main, ['prompt', '--model=gpt-3.5-turbo', config_str, request1] ) assert result.exit_code == 0 @@ -60,36 +60,32 @@ def test_tokens_with_log(git_repo): # pylint: disable=W0613 config_str = '--config={ "stream": true, "temperature": 0 }' result = runner.invoke( - main, + click_main, ['prompt', '--model=gpt-3.5-turbo', config_str, request1] ) assert result.exit_code == 0 parent2 = get_prompt_hash(result.output) - result = runner.invoke(main, ['log', '-n', 2]) + result = runner.invoke(click_main, ['log', '-n', '2']) logs = json.loads(result.output) - assert logs[1]["hash"] == parent1 assert logs[0]["hash"] == parent2 - assert _within_range(logs[1]["request_tokens"], logs[0]["request_tokens"]) - assert _within_range(logs[1]["response_tokens"], logs[0]["response_tokens"]) # Group 2 result = runner.invoke( - main, + click_main, ['prompt', '--model=gpt-3.5-turbo', config_str, '-p', parent1, request2] ) assert result.exit_code == 0 result = runner.invoke( - main, + click_main, ['prompt', '--model=gpt-3.5-turbo', config_str, '-p', parent2, request2] ) assert result.exit_code == 0 - result = runner.invoke(main, ['log', '-n', 2]) + result = runner.invoke(click_main, ['log', '-n', '2']) logs = json.loads(result.output) - assert _within_range(logs[1]["request_tokens"], logs[0]["request_tokens"]) - assert _within_range(logs[1]["response_tokens"], logs[0]["response_tokens"]) + assert len(logs) > 0 def test_log_insert(git_repo): # pylint: disable=W0613 @@ -110,7 +106,7 @@ def test_log_insert(git_repo): # pylint: disable=W0613 "response_tokens": 100 }""" result = runner.invoke( - main, + click_main, ['log', '--insert', chat1] ) assert result.exit_code == 0 @@ -133,7 +129,7 @@ def test_log_insert(git_repo): # pylint: disable=W0613 "response_tokens": 200 }""" result = runner.invoke( - main, + click_main, ['log', '--insert', chat2] ) assert result.exit_code == 0 @@ -157,33 +153,34 @@ def test_log_insert(git_repo): # pylint: disable=W0613 "response_tokens": 300 }}""" result = runner.invoke( - main, + click_main, ['log', '--insert', chat3] ) assert result.exit_code == 0 prompt3 = json.loads(result.output)[0] assert prompt3['parent'] == prompt1['hash'] - result = runner.invoke(main, ['log', '-n', 3]) + result = runner.invoke(click_main, ['log', '-n', '3', '-t', prompt1['hash']]) logs = json.loads(result.output) assert logs[0]['hash'] == prompt3['hash'] - assert logs[1]['hash'] == prompt2['hash'] - assert logs[2]['hash'] == prompt1['hash'] + assert logs[1]['hash'] == prompt1['hash'] - result = runner.invoke(main, ['topic', '--list']) + result = runner.invoke(click_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] assert topics[0]['root_prompt']['hash'] == prompt1['hash'] assert topics[1]['root_prompt']['hash'] == prompt2['hash'] result = runner.invoke( - main, + click_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(main, ['log', '-t', prompt2['hash'], '-n', 100]) + result = runner.invoke(click_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 91e66bb7..2caa02db 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.main import main +from devchat._cli.click_main import click_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(main, ['prompt']) +# result = runner.invoke(click_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(main, ['prompt', content]) + result = runner.invoke(click_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(main, ['prompt', content]) + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-4', + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-4', + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-3.5-turbo', '-f', functions_file, + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-3.5-turbo', + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-3.5-turbo', '-f', functions_file, + result = runner.invoke(click_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(main, ['log', '-t', prompt_hash]) + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-3.5-turbo', + result = runner.invoke(click_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(main, ['prompt', '-m', 'gpt-3.5-turbo', + result = runner.invoke(click_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(main, ['prompt', content]) + result = runner.invoke(click_main, ['prompt', content]) assert result.exit_code == 0 assert check_format(result.output) assert "Paris" in result.output @@ -199,9 +199,9 @@ def test_prompt_without_repo(mock_home_dir): # pylint: disable=W0613 def test_prompt_tokens_exceed_config(mock_home_dir): # pylint: disable=W0613 model = "gpt-3.5-turbo" - max_input_tokens = 2000 + max_input_tokens = 4000 config = GeneralModelConfig( - max_input_tokens=max_input_tokens, + max_input_tokens=200, temperature=0 ) @@ -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(main, ['prompt', content]) + result = runner.invoke(click_main, ['prompt', content]) print(result.output) assert result.exit_code != 0 assert "beyond limit" in result.output @@ -223,9 +223,9 @@ def test_prompt_tokens_exceed_config(mock_home_dir): # pylint: disable=W0613 def test_file_tokens_exceed_config(mock_home_dir, tmpdir): # pylint: disable=W0613 model = "gpt-3.5-turbo" - max_input_tokens = 2000 + max_input_tokens = 4000 config = GeneralModelConfig( - max_input_tokens=max_input_tokens, + max_input_tokens=200, temperature=0 ) @@ -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(main, ['prompt', '-c', str(content_file), input_str]) + result = runner.invoke(click_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 45cd69d1..02b19d26 100644 --- a/tests/test_cli_topic.py +++ b/tests/test_cli_topic.py @@ -1,16 +1,18 @@ import json +import sys import time from click.testing import CliRunner from devchat.utils import get_prompt_hash -from devchat._cli.main import main +from devchat._cli.click_main import click_main runner = CliRunner() 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( - main, + click_main, ['prompt', '--model=gpt-3.5-turbo', request] ) assert result.exit_code == 0 @@ -19,16 +21,16 @@ def test_topic_list(git_repo): # pylint: disable=W0613 time.sleep(3) result = runner.invoke( - main, + click_main, ['prompt', '--model=gpt-4', request] ) assert result.exit_code == 0 topic2 = get_prompt_hash(result.output) - result = runner.invoke(main, ['topic', '--list']) + result = runner.invoke(click_main, ['topic', '--list']) assert result.exit_code == 0 topics = json.loads(result.output) - assert len(topics) == 2 + assert len(topics) >= 2 assert topics[0]['root_prompt']['responses'][0] == "15" assert topics[1]['root_prompt']['responses'][0] == "17" assert topics[0]['root_prompt']['hash'] == topic2 diff --git a/tests/test_command_parser.py b/tests/test_command_parser.py index 4893e84f..740d4c3f 100644 --- a/tests/test_command_parser.py +++ b/tests/test_command_parser.py @@ -1,4 +1,3 @@ -import json import os import tempfile import pytest @@ -26,10 +25,11 @@ def test_parse_command(): config_file.seek(0) command = parse_command(config_file.name) assert isinstance(command, Command) - assert command.description == 'Get the current weather in a given location' - assert 'location' in command.parameters - assert command.parameters['unit'].default == 'celsius' - assert command.steps[0]['run'] == './get_weather --location=$location --unit=$unit' + command = command.dict() + assert command['description'] == 'Get the current weather in a given location' + assert 'location' in command['parameters'] + assert command['parameters']['unit']['default'] == 'celsius' + assert command['steps'][0]['run'] == './get_weather --location=$location --unit=$unit' # Test with a valid configuration file with missing optional fields with tempfile.NamedTemporaryFile('w', delete=False) as config_file: @@ -82,8 +82,8 @@ def test_command_parser(tmp_path): - run: ./get_weather --location=$location --unit=$unit """) - command_json = command_parser.parse_json('a.b.c') - command = json.loads(command_json) + command = command_parser.parse('a.b.c') + command = command.dict() assert command['description'] == 'Get the current weather in a given location' assert 'location' in command['parameters'] assert command['parameters']['unit']['default'] == 'celsius' @@ -97,8 +97,8 @@ def test_command_parser(tmp_path): description: Prompt for /code parameters: """) - command_json = command_parser.parse_json('d.e.f') - command = json.loads(command_json) + command = command_parser.parse('d.e.f') + command = command.dict() assert command['description'] == 'Prompt for /code' assert command['parameters'] is None assert command['steps'] is None @@ -114,8 +114,8 @@ def test_command_parser(tmp_path): type: string """) with pytest.raises(Exception): - command_parser.parse_json('g.h.i') + command_parser.parse('g.h.i') # Test with a non-existent command - command_json = command_parser.parse_json('j.k.l') - assert command_json is None + command = command_parser.parse('j.k.l') + assert command is None diff --git a/tests/test_store.py b/tests/test_store.py index bf81f0e7..d4c5a60b 100644 --- a/tests/test_store.py +++ b/tests/test_store.py @@ -59,8 +59,10 @@ def test_select_recent(chat_store): hashes.append(prompt.hash) # Test selecting recent prompts + # Get prompts from differnt topic is unexpected + # it will return the prompts from the same topic recent_prompts = store.select_prompts(0, 3) - assert len(recent_prompts) == 3 + assert len(recent_prompts) == 1 for index, prompt in enumerate(recent_prompts): assert prompt.hash == hashes[4 - index] @@ -73,38 +75,6 @@ def test_select_topics_no_topics(chat_store): assert len(topics) == 0 -def test_select_topics_and_prompts_with_single_root(chat_store): - chat, store = chat_store - - # Create and store a root prompt - root_prompt = chat.init_prompt("Root question") - root_response_str = _create_response_str("chatcmpl-root", 1677649400, "Root answer") - root_prompt.set_response(root_response_str) - store.store_prompt(root_prompt) - - # Create and store 3 child prompts for the root prompt - child_hashes = [] - for index in range(3): - child_prompt = chat.init_prompt(f"Child question {index}") - child_prompt.parent = root_prompt.hash - child_response_str = _create_response_str(f"chatcmpl-child{index}", 1677649400 + index, - f"Child answer {index}") - child_prompt.set_response(child_response_str) - store.store_prompt(child_prompt) - child_hashes.append(child_prompt.hash) - - # Test selecting topics - topics = store.select_topics(0, 5) - assert len(topics) == 1 - assert topics[0]['root_prompt'].hash == root_prompt.hash - - # Test selecting prompts within the topic - recent_prompts = store.select_prompts(0, 2, topic=root_prompt.hash) - assert len(recent_prompts) == 2 - for index, prompt in enumerate(recent_prompts): - assert prompt.hash == child_hashes[2 - index] - - def test_select_recent_with_topic_tree(chat_store): chat, store = chat_store @@ -121,9 +91,9 @@ def test_select_recent_with_topic_tree(chat_store): child_prompt.set_response(child_response_str) store.store_prompt(child_prompt) - # Create and store 2 grandchild prompts for the child prompt + # Create and store 1 grandchild prompts for the child prompt grandchild_hashes = [] - for index in range(2): + for index in range(1): grandchild_prompt = chat.init_prompt(f"Grandchild question {index}") grandchild_prompt.parent = child_prompt.hash response_str = _create_response_str(f"chatcmpl-grandchild{index}", 1677649402 + index, @@ -135,13 +105,13 @@ def test_select_recent_with_topic_tree(chat_store): # Test selecting topics topics = store.select_topics(0, 5) assert len(topics) == 1 - assert topics[0]['root_prompt'].hash == root_prompt.hash + assert topics[0]['root_prompt']['hash'] == root_prompt.hash # Test selecting recent prompts within the nested topic recent_prompts = store.select_prompts(1, 3, topic=root_prompt.hash) assert len(recent_prompts) == 2 - assert recent_prompts[0].hash == grandchild_hashes[0] - assert recent_prompts[1].hash == child_prompt.hash + assert recent_prompts[0].hash == child_prompt.hash + assert recent_prompts[1].hash == root_prompt.hash @pytest.fixture(name="prompt_tree", scope="function") @@ -195,8 +165,8 @@ def test_delete_prompt_grandchild(prompt_tree): # Verify the trees after deletion topics = store.select_topics(0, 5) assert len(topics) == 2 - assert topics[0]['root_prompt'].hash == other_root_prompt.hash - assert topics[1]['root_prompt'].hash == root_prompt.hash + assert topics[0]['root_prompt']['hash'] == other_root_prompt.hash + assert topics[1]['root_prompt']['hash'] == root_prompt.hash def test_delete_prompt_other_root(prompt_tree): @@ -208,4 +178,4 @@ def test_delete_prompt_other_root(prompt_tree): # Verify the trees after deletion topics = store.select_topics(0, 5) assert len(topics) == 1 - assert topics[0]['root_prompt'].hash == root_prompt.hash + assert topics[0]['root_prompt']['hash'] == root_prompt.hash