diff --git a/devchat/__init__.py b/devchat/__init__.py index e69de29b..35426eed 100644 --- a/devchat/__init__.py +++ b/devchat/__init__.py @@ -0,0 +1,4 @@ +import os + +script_dir = os.path.dirname(os.path.realpath(__file__)) +os.environ["TIKTOKEN_CACHE_DIR"] = os.path.join(script_dir, "tiktoken_cache") \ No newline at end of file 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..f1672a65 100755 --- a/devchat/_cli/main.py +++ b/devchat/_cli/main.py @@ -1,34 +1,29 @@ """ 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 + +from devchat.workflow.cli import workflow logger = get_logger(__name__) -def main(argv=None): - if argv is None: - argv = sys.argv[1:] +@click.group() +def main(): + """DevChat CLI: A command-line interface for DevChat.""" - parser = argparse.ArgumentParser(description="CLI tool") - subparsers = parser.add_subparsers(help='sub-command help') - for _1, cmd in commands.items(): - cmd.register(subparsers) - args = parser.parse_args(argv) - if hasattr(args, 'func'): - func_args = vars(args).copy() - del func_args['func'] +main.add_command(prompt) +main.add_command(log) +main.add_command(run) +main.add_command(topic) +main.add_command(route) - args.func(**func_args) - else: - parser.print_help() +main.add_command(workflow) 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/router.py b/devchat/_cli/router.py index 29b9e49f..03cf8733 100755 --- a/devchat/_cli/router.py +++ b/devchat/_cli/router.py @@ -3,6 +3,8 @@ import sys from typing import List, Optional +from devchat.workflow.workflow import Workflow + def _get_model_and_config( model: Optional[str], @@ -139,6 +141,28 @@ def llm_route(content: Optional[str], parent: Optional[str], reference: Optional content, parent, reference, instruct, context, model, config_str, None, None, True ) + name, user_input = Workflow.parse_trigger(content) + workflow = Workflow.load(name) if name else None + if workflow: + print(assistant.prompt.formatted_header()) + + return_code = 0 + if workflow.should_show_help(user_input): + doc = workflow.get_help_doc(user_input) + print(doc) + + else: + # run the workflow + workflow.setup( + model_name=model, + user_input=user_input, + history_messages=assistant.prompt.messages, + parent_hash=parent, + ) + return_code = workflow.run_steps() + + sys.exit(return_code) + print(assistant.prompt.formatted_header()) command_result = run_command( model_name = model, 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/devchat/workflow/README.md b/devchat/workflow/README.md new file mode 100644 index 00000000..0db96b37 --- /dev/null +++ b/devchat/workflow/README.md @@ -0,0 +1,5 @@ +# Workflow Engine + +The Workflow Engine allows use to create, manage, and run workflows in DevChat. + +This is the refactored and enhanced version of the engine/ module and some commands in _cli/ module. diff --git a/devchat/workflow/__init__.py b/devchat/workflow/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/devchat/workflow/cli.py b/devchat/workflow/cli.py new file mode 100644 index 00000000..90038958 --- /dev/null +++ b/devchat/workflow/cli.py @@ -0,0 +1,20 @@ +import click +from devchat.workflow.command.update import update +from devchat.workflow.command.list import list_cmd +from devchat.workflow.command.env import env +from devchat.workflow.command.config import config_cmd + + +@click.group(help="CLI for devchat workflow engine.") +def workflow(): + pass + + +workflow.add_command(update) +workflow.add_command(list_cmd) +workflow.add_command(env) +workflow.add_command(config_cmd) + + +if __name__ == "__main__": + workflow() diff --git a/devchat/workflow/command/__init__.py b/devchat/workflow/command/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/devchat/workflow/command/config.py b/devchat/workflow/command/config.py new file mode 100644 index 00000000..4f4f50e0 --- /dev/null +++ b/devchat/workflow/command/config.py @@ -0,0 +1,26 @@ +import json + +from pathlib import Path +from typing import NamedTuple, List, Set, Tuple +import click +import oyaml as yaml + +from devchat.workflow.path import WORKFLOWS_BASE, WORKFLOWS_CONFIG_FILENAME + + +@click.command(help="Workflow configuration.", name="config") +@click.option("--json", "in_json", is_flag=True, help="Output in json format.") +def config_cmd(in_json: bool): + + config_path = Path(WORKFLOWS_BASE) / WORKFLOWS_CONFIG_FILENAME + config_content = {} + if config_path.exists(): + with open(config_path, "r", encoding="utf-8") as f: + config_content = yaml.safe_load(f.read()) + + if not in_json: + click.echo(config_content) + + else: + json_format = json.dumps(config_content) + click.echo(json_format) diff --git a/devchat/workflow/command/env.py b/devchat/workflow/command/env.py new file mode 100644 index 00000000..e1ce6c30 --- /dev/null +++ b/devchat/workflow/command/env.py @@ -0,0 +1,92 @@ +# pylint: disable=invalid-name + +""" +Commands for managing the python environment of workflows. +""" + +import sys +from pathlib import Path +from typing import Optional, List +import click +from devchat.workflow.env_manager import PyEnvManager, MAMBA_PY_ENVS + + +def _get_all_env_names() -> List[str]: + """ + Get all the python env names of workflows. + """ + # devchat reserved envs + excludes = ["devchat", "devchat-ask", "devchat-commands"] + + envs_path = Path(MAMBA_PY_ENVS) + envs = [ + env.name + for env in envs_path.iterdir() + if env.is_dir() and env.name not in excludes + ] + return envs + + +@click.command(help="List all the python envs of workflows.", name="list") +def list_envs(): + envs = _get_all_env_names() + click.echo(f"Found {len(envs)} python envs of workflows:") + click.echo("\n".join(envs)) + + +@click.command(help="Remove a specific workflow python env.") +@click.option( + "--env-name", + "-n", + help="The name of the python env to remove.", + required=False, + type=str, +) +@click.option( + "--all", "all_flag", help="Remove all the python envs of workflows.", is_flag=True +) +def remove(env_name: Optional[str] = None, all_flag: bool = False): + if not env_name and not all_flag: + click.echo("Please provide the name of the python env to remove.") + sys.exit(1) + + if env_name: + manager = PyEnvManager() + remove_ok = manager.remove(env_name) + + if remove_ok: + click.echo(f"Removed python env: {env_name}") + sys.exit(0) + + else: + click.echo(f"Failed to remove python env: {env_name}") + sys.exit(1) + + if all_flag: + envs = _get_all_env_names() + manager = PyEnvManager() + ok = [] + failed = [] + for name in envs: + remove_ok = manager.remove(name) + if remove_ok: + ok.append(name) + else: + failed.append(name) + + click.echo(f"Removed {len(ok)} python envs of workflows:") + click.echo("\n".join(ok)) + if failed: + click.echo(f"Failed to remove {len(failed)} python envs of workflows:") + click.echo("\n".join(failed)) + + sys.exit(0) + + +@click.group(help="Manage the python environment of workflows.") +def env(): + pass + + +env.add_command(list_envs) +env.add_command(remove) diff --git a/devchat/workflow/command/list.py b/devchat/workflow/command/list.py new file mode 100644 index 00000000..42a1c642 --- /dev/null +++ b/devchat/workflow/command/list.py @@ -0,0 +1,93 @@ +import json +from pathlib import Path +from typing import NamedTuple, List, Set, Tuple, Optional, Dict +from dataclasses import dataclass, asdict, field + +import click +import oyaml as yaml + +from devchat.workflow.namespace import get_prioritized_namespace_path +from devchat.workflow.path import COMMAND_FILENAMES + + +@dataclass +class WorkflowMeta: + name: str + namespace: str + active: bool + command_conf: Dict = field( + default_factory=dict + ) # content of command.yml, excluding "steps" field + + def __str__(self): + return f"{'*' if self.active else ' '} {self.name} ({self.namespace})" + + +def iter_namespace( + ns_path: str, existing_names: Set[str] +) -> Tuple[List[WorkflowMeta], Set[str]]: + """ + Get all workflows under the namespace path. + + Args: + ns_path: the namespace path + existing_names: the existing workflow names to check if the workflow is the first priority + + Returns: + List[WorkflowMeta]: the workflows + Set[str]: the updated existing workflow names + """ + root = Path(ns_path) + interest_files = set(COMMAND_FILENAMES) + result = [] + unique_names = set(existing_names) + for f in root.rglob("*"): + if f.is_file() and f.name in interest_files: + rel_path = f.relative_to(root) + parts = rel_path.parts + workflow_name = ".".join(parts[:-1]) + is_first = workflow_name not in unique_names + unique_names.add(workflow_name) + + # load the config content from f + with open(f, "r", encoding="utf-8") as fi: + yaml_content = fi.read() + command_conf = yaml.safe_load(yaml_content) + # pop the "steps" field + command_conf.pop("steps", None) + + workflow = WorkflowMeta( + name=workflow_name, + namespace=root.name, + active=is_first, + command_conf=command_conf, + ) + result.append(workflow) + + return result, unique_names + + +@click.command(help="List all local workflows.", name="list") +@click.option("--json", "in_json", is_flag=True, help="Output in json format.") +def list_cmd(in_json: bool): + namespace_paths = get_prioritized_namespace_path() + + workflows: List[WorkflowMeta] = [] + visited_names = set() + for ns_path in namespace_paths: + ws, visited_names = iter_namespace(ns_path, visited_names) + workflows.extend(ws) + + if not in_json: + # print basic info + active_count = len([w for w in workflows if w.active]) + total_count = len(workflows) + click.echo(f"workflows (active/total): {active_count}/{total_count}") + for w in workflows: + click.echo(w) + + else: + # convert workflows to json + data = [asdict(w) for w in workflows] + json_format = json.dumps(data) + click.echo(json_format) diff --git a/devchat/workflow/command/update.py b/devchat/workflow/command/update.py new file mode 100644 index 00000000..6711cd28 --- /dev/null +++ b/devchat/workflow/command/update.py @@ -0,0 +1,316 @@ +# pylint: disable=invalid-name + +import os +import shutil +import tempfile +import zipfile +from typing import List, Optional, Tuple +from pathlib import Path +from datetime import datetime + +import click +import requests + +from devchat.workflow.path import ( + CHAT_DIR, + WORKFLOWS_BASE, + WORKFLOWS_BASE_NAME, + CUSTOM_BASE, +) +from devchat.utils import get_logger + +HAS_GIT = False +try: + from git import Repo, InvalidGitRepositoryError, GitCommandError +except ImportError: + pass +else: + HAS_GIT = True + + +REPO_NAME = "workflows" +DEFAULT_BRANCH = "scripts" +REPO_URLS = [ + # url, branch + ("git@github.com:devchat-ai/workflows.git", DEFAULT_BRANCH), + ("https://github.com/devchat-ai/workflows.git", DEFAULT_BRANCH), + ("https://gitlab.com/devchat-ai/workflows.git", DEFAULT_BRANCH), +] +ZIP_URLS = [ + "https://codeload.github.com/devchat-ai/workflows/zip/refs/heads/scripts", + "https://gitlab.com/devchat-ai/workflows/-/archive/scripts/workflows-scripts.zip", +] + + +logger = get_logger(__name__) + + +def _backup(workflow_base: Path, n: int = 5) -> Optional[Path]: + """ + Backup the current workflow base dir to zip with timestamp under .backup. + + Args: + n: the number of backups to keep, default 3 + + Returns: + Path: the backup zip path + """ + + if not workflow_base.exists(): + return + + backup_dir = workflow_base.parent / ".backup" + backup_dir.mkdir(exist_ok=True) + + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + backup_zip = backup_dir / f"{WORKFLOWS_BASE_NAME}_{timestamp}" + shutil.make_archive(backup_zip, "zip", workflow_base) + click.echo(f"Backup {workflow_base} to {backup_zip}") + + # keep the last n backups + backups = sorted(backup_dir.glob(f"{WORKFLOWS_BASE_NAME}_*"), reverse=True) + for backup in backups[n:]: + backup.unlink() + click.echo(f"Remove old backup {backup}") + + return backup_zip + + +def _download_zip_to_dir(candidate_urls: List[str], dst_dir: Path) -> bool: + """ + Download the zip file with the first successful url + in the candidate_urls to the target_dir. + + Args: + candidate_urls: the list of candidate urls + dst_dir: the dst dir of the extracted zip file, should not exist + + Returns: + bool: True if success else False + """ + assert not dst_dir.exists() + + download_ok = False + with tempfile.TemporaryDirectory() as tmp_dir: + tmp_dir_path = Path(tmp_dir) + + for url in candidate_urls: + try: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + zip_path = tmp_dir_path / f"{WORKFLOWS_BASE_NAME}_{timestamp}.zip" + click.echo(f"Downloading workflow from {url} to {zip_path}") + + # Download the workflow zip file + response = requests.get(url, stream=True, timeout=10) + with open(zip_path, "wb") as zip_f: + for chunk in response.iter_content(chunk_size=8192): + if chunk: + zip_f.write(chunk) + + # Extract the zip file + with zipfile.ZipFile(zip_path, "r") as zip_f: + zip_f.extractall(tmp_dir_path) + + # move the extracted dir to target dir + extracted_dir = tmp_dir_path / f"{REPO_NAME}-{DEFAULT_BRANCH}" + shutil.move(extracted_dir, dst_dir) + click.echo(f"Extracted to {dst_dir}") + + download_ok = True + break + + except Exception as e: + click.echo(f"Failed to download from {url}: {e}") + + return download_ok + + +def update_by_zip(workflow_base: Path): + click.echo("Updating by zip file...") + parent = workflow_base.parent + + if not workflow_base.exists(): + # No previous workflows, download to the workflow_base directly + download_ok = _download_zip_to_dir(ZIP_URLS, workflow_base) + if not download_ok: + click.echo("Failed to download from all zip urls. Please Try again later.") + return + + click.echo(f"Updated {workflow_base} successfully by zip.") + + else: + # Has previous workflows, download as tmp_new + tmp_new = parent / f"{WORKFLOWS_BASE_NAME}_new" + if tmp_new.exists(): + shutil.rmtree(tmp_new) + # TODO: handle error? + # shutil.rmtree(tmp_new, onerror=__onerror) + + download_ok = _download_zip_to_dir(ZIP_URLS, tmp_new) + + if not download_ok: + click.echo("Failed to download from all zip urls. Skip update.") + return + + # backup the current workflows + backup_zip = _backup(workflow_base) + + # rename the new dir to the workflow_base + shutil.rmtree(workflow_base) + shutil.move(tmp_new, workflow_base) + + click.echo(f"Updated {workflow_base} by zip. (backup: {backup_zip})") + + +def _clone_repo_to_dir(candidates: List[Tuple[str, str]], dst_dir: Path) -> bool: + """ + Clone the git repo with the first successful url + in the candidates to the dst_dir. + + Args: + candidates: the list of candidate git url and branch pairs. + dst_dir: the dst dir of the cloned repo, should not exist + + Returns: + bool: True if success else False + """ + + assert not dst_dir.exists() + + clone_ok = False + for url, branch in candidates: + try: + repo = Repo.clone_from(url, to_path=dst_dir, branch=branch) + click.echo(f"Cloned from {url}|{branch} to {dst_dir}") + clone_ok = True + break + except GitCommandError as e: + click.echo(f"Failed to clone from {url}|{branch}: {e}") + + return clone_ok + + +def update_by_git(workflow_base: Path): + click.echo("Updating by git...") + parent = workflow_base.parent + + if not workflow_base.exists(): + # No previous workflows, clone to the workflow_base directly + clone_ok = _clone_repo_to_dir(REPO_URLS, workflow_base) + if not clone_ok: + click.echo("Failed to clone from all git urls. Please Try again later.") + return + + click.echo(f"Updated {workflow_base} by git.") + + else: + # Has previous workflows + repo = None + try: + # check if the workflow base dir is a valid git repo + repo = Repo(workflow_base) + head_name = repo.head.reference.name + except InvalidGitRepositoryError: + pass + except Exception: + repo = None + + if repo is None: + # current workflow base dir is not a valid git repo + # try to clone the new repo to tmp_new + tmp_new = parent / f"{WORKFLOWS_BASE_NAME}_new" + if tmp_new.exists(): + shutil.rmtree(tmp_new) + + clone_ok = _clone_repo_to_dir(REPO_URLS, tmp_new) + if not clone_ok: + click.echo("Failed to clone from all git urls. Skip update.") + return + + # backup the current workflows + backup_zip = _backup(workflow_base) + + # rename the new dir to the workflow_base + shutil.rmtree(workflow_base) + shutil.move(tmp_new, workflow_base) + click.echo(f"Updated {workflow_base} by git. (backup: {backup_zip})") + return + + # current workflow base dir is a valid git repo + if head_name != DEFAULT_BRANCH: + click.echo( + f"Current workflow branch is not the default one[{DEFAULT_BRANCH}]: " + f"<{head_name}>. Skip update." + ) + return + + try: + repo.git.fetch("origin") + except GitCommandError as e: + click.echo(f"Failed to fetch from origin. Skip update. {e}") + return + + local_main_hash = repo.head.commit.hexsha + remote_main_hash = repo.commit(f"origin/{DEFAULT_BRANCH}").hexsha + + if local_main_hash == remote_main_hash: + click.echo( + f"Local branch is up-to-date with remote {DEFAULT_BRANCH}. Skip update." + ) + return + + try: + # backup the current workflows + backup_zip = _backup(workflow_base) + + # discard the local changes and force update to the latest main + repo.git.reset("--hard", "HEAD") + repo.git.clean("-df") + repo.git.fetch("origin") + repo.git.reset("--hard", f"origin/{DEFAULT_BRANCH}") + click.echo( + f"Updated {workflow_base} from <{local_main_hash[:8]}> to" + f" <{remote_main_hash[:8]}>. (backup: {backup_zip})" + ) + except GitCommandError as e: + click.echo(f"Failed to update to the latest main: {e}. Skip update.") + return + return + + +def copy_workflows_usr(): + """ + Copy workflows/usr to scripts/custom/usr for engine migration. + """ + old_usr_dir = os.path.join(CHAT_DIR, "workflows", "usr") + new_usr_dir = os.path.join(CUSTOM_BASE, "usr") + + old_exists = os.path.exists(old_usr_dir) + new_exists = os.path.exists(new_usr_dir) + + if old_exists and not new_exists: + shutil.copytree(old_usr_dir, new_usr_dir) + click.echo(f"Copied {old_usr_dir} to {new_usr_dir} successfully.") + else: + click.echo( + f"Skip copying usr dir. old exists: {old_exists}, new exists: {new_exists}." + ) + + +@click.command(help="Update the workflow_base dir.") +@click.option( + "-f", "--force", is_flag=True, help="Force update the workflows to the latest main." +) +def update(force: bool): + click.echo(f"Updating wf repo... force: {force}") + click.echo(f"WORKFLOWS_BASE: {WORKFLOWS_BASE}") + + base_path = Path(WORKFLOWS_BASE) + # # backup(base_path) + + if HAS_GIT: + update_by_git(base_path) + else: + update_by_zip(base_path) + + copy_workflows_usr() diff --git a/devchat/workflow/command/update_flowchart.md b/devchat/workflow/command/update_flowchart.md new file mode 100644 index 00000000..5950ed67 --- /dev/null +++ b/devchat/workflow/command/update_flowchart.md @@ -0,0 +1,38 @@ + +```mermaid +flowchart TD + A([Start]) --> B{is git installed?} + B -->|No| C(Update by downloading zip file) + C --> E{is workflow_base/ dir exists?} + E -->|No| F(Try to download latest zip \nand extract it to workflow_base/ ) + F --> H{downloaded and extracted successfully?} + H -->|No| I([Failed to update. \nWorkflow is unavailable.\nShould try again later.]) + H -->|Yes| J([Update completed]) + E -->|Yes| G(Try to download latest zip \nand extract it to workflow_base_new/) + G --> K{downloaded and extracted successfully?} + K -->|No| L([Skip update this time.]) + K -->|Yes| M(1. Archive workflow_base/ to the backup dir\n2. Rename workflow_base_new/ to workflow_base/) + M --> N([Update completed]) + + + B -->|Yes| D(Update by git command) + D --> D1{is workflow_base/ dir exists?} + D1 -->|No| D2(Try to clone the repo to workflow_base/) + D2 --> D3{cloned successfully?} + D3 -->|No| D4([Failed to update. \nWorkflow is unavailable.\nShould try again later.]) + D3 -->|Yes| D5([Update completed]) + + D1 -->|Yes| D6{is workflow_base/ a git repo?} + D6 -->|No| D7(Try to clone the repo to workflow_base_new/) + D7 --> D8{cloned successfully?} + D8 -->|No| D9([Skip update this time.]) + D8 -->|Yes| D10(1. Archive workflow_base/ to the backup dir\n2. Rename workflow_base_new/ to workflow_base/) + D10 --> D11([Update completed]) + + D6 -->|Yes| D12{is the repo on main branch?} + D12 -->|No| D13([Skip update because it is on dev branch.]) + D12 -->|Yes| D14(1. Archive current workflow_base/ to the backup dir\n2. Try to pull the latest main) + D14 --> D15{pulled successfully?} + D15 -->|No| D16([Skip update this time.]) + D15 -->|Yes| D17([Update completed]) +``` diff --git a/devchat/workflow/env_manager.py b/devchat/workflow/env_manager.py new file mode 100644 index 00000000..7bfe8c1e --- /dev/null +++ b/devchat/workflow/env_manager.py @@ -0,0 +1,205 @@ +# pylint: disable=invalid-name +import os +import sys +import subprocess +from typing import Optional, Dict + +from .envs import MAMBA_BIN_PATH +from .path import MAMBA_PY_ENVS, MAMBA_ROOT +from .user_setting import USER_SETTINGS +from .schema import ExternalPyConf + + +# CONDA_FORGE = [ +# "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/", +# "conda-forge", +# ] +CONDA_FORGE_TUNA = "https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/" +PYPI_TUNA = "https://pypi.tuna.tsinghua.edu.cn/simple" + + +def _get_external_envs() -> Dict[str, ExternalPyConf]: + """ + Get the external python environments info from the user settings. + """ + external_pythons: Dict[str, ExternalPyConf] = {} + for conf in USER_SETTINGS.external_workflow_python: + external_pythons[conf.env_name] = conf + + return external_pythons + +EXTERNAL_ENVS = _get_external_envs() + +class PyEnvManager: + mamba_bin = MAMBA_BIN_PATH + mamba_root = MAMBA_ROOT + + def __init__(self): + pass + + @staticmethod + def get_py_version(py: str) -> Optional[str]: + """ + Get the version of the python executable. + """ + py_version_cmd = [py, "--version"] + with subprocess.Popen( + py_version_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) as proc: + proc.wait() + + if proc.returncode != 0: + return None + + out = proc.stdout.read().decode("utf-8") + return out.split()[1] + + def install(self, env_name: str, requirements_file: str) -> bool: + """ + Install requirements into the python environment. + + Args: + requirements: the absolute path to the requirements file. + """ + py = self.get_py(env_name) + if not py: + # TODO: raise error? + return False + + if not os.path.exists(requirements_file): + # TODO: raise error? + return False + + cmd = [ + py, + "-m", + "pip", + "install", + "-r", + requirements_file, + "-i", + PYPI_TUNA, + ] + env = os.environ.copy() + env.pop("PYTHONPATH") + with subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env + ) as proc: + proc.wait() + + if proc.returncode != 0: + print( + f"Failed to install requirements: {requirements_file}", flush=True + ) + return False + + return True + + def ensure(self, env_name: str, py_version: str) -> Optional[str]: + """ + Ensure the python environment exists with the given name and version. + + return the python executable path. + """ + py = self.get_py(env_name) + should_remove = False + + if py: + # check the version of the python executable + current_version = self.get_py_version(py) + + if current_version == py_version: + return py + + should_remove = True + + print("\n```Step\n# Setting up workflow environment", flush=True) + print(f"\nenv_name: {env_name}") + print(f"python: {py_version}", flush=True) + + if should_remove: + self.remove(env_name) + + # create the environment + create_ok = self.create(env_name, py_version) + print("\n```", flush=True) + + if not create_ok: + return None + return self.get_py(env_name) + + def create(self, env_name: str, py_version: str) -> bool: + """ + Create a new python environment using mamba. + """ + is_exist = os.path.exists(os.path.join(MAMBA_PY_ENVS, env_name)) + if is_exist: + return True + + # create the environment + cmd = [ + self.mamba_bin, + "create", + "-n", + env_name, + "-c", + CONDA_FORGE_TUNA, + "-r", + self.mamba_root, + f"python={py_version}", + "-y", + ] + with subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) as proc: + proc.wait() + + if proc.returncode != 0: + return False + + return True + + def remove(self, env_name: str) -> bool: + """ + Remove the python environment. + """ + is_exist = os.path.exists(os.path.join(MAMBA_PY_ENVS, env_name)) + if not is_exist: + return True + + # remove the environment + cmd = [ + self.mamba_bin, + "env", + "remove", + "-n", + env_name, + "-r", + self.mamba_root, + "-y", + ] + with subprocess.Popen( + cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) as proc: + proc.wait() + + if proc.returncode != 0: + return False + + return True + + def get_py(self, env_name: str) -> Optional[str]: + """ + Get the python executable path of the given environment. + """ + env_path = None + if sys.platform == "win32": + env_path = os.path.join(MAMBA_PY_ENVS, env_name, "python.exe") + # env_path = os.path.join(MAMBA_PY_ENVS, env_name, "Scripts", "python.exe") + else: + env_path = os.path.join(MAMBA_PY_ENVS, env_name, "bin", "python") + + if env_path and os.path.exists(env_path): + return env_path + + return None diff --git a/devchat/workflow/envs.py b/devchat/workflow/envs.py new file mode 100644 index 00000000..f50b6997 --- /dev/null +++ b/devchat/workflow/envs.py @@ -0,0 +1,10 @@ +""" +Explicitly define the environment variables used in the workflow engine. +""" +import os + +PYTHON_PATH = os.environ.get("PYTHONPATH", "") +DEVCHAT_PYTHON_PATH = os.environ.get("DEVCHAT_PYTHONPATH", PYTHON_PATH) + +# the path to the mamba binary +MAMBA_BIN_PATH = os.environ.get("MAMBA_BIN_PATH", "") diff --git a/devchat/workflow/namespace.py b/devchat/workflow/namespace.py new file mode 100644 index 00000000..fa1c62e3 --- /dev/null +++ b/devchat/workflow/namespace.py @@ -0,0 +1,74 @@ +""" +Namespace management for workflows +""" + +import os +from typing import List +from pydantic import BaseModel, Extra, ValidationError +import oyaml as yaml +from devchat.utils import get_logger + +from .path import ( + CUSTOM_BASE, + MERICO_WORKFLOWS, + COMMUNITY_WORKFLOWS, + CUSTOM_CONFIG_FILE, +) + +logger = get_logger(__name__) + + +class CustomConfig(BaseModel): + namespaces: List[str] = [] # active namespaces ordered by priority + + class Config: + extra = Extra.ignore + + +def _load_custom_config() -> CustomConfig: + """ + Load the custom config file. + """ + config = CustomConfig() + + if not os.path.exists(CUSTOM_CONFIG_FILE): + return config + + with open(CUSTOM_CONFIG_FILE, "r", encoding="utf-8") as f: + content = f.read() + yaml_content = yaml.safe_load(content) + try: + if yaml_content: + config = CustomConfig.parse_obj(yaml_content) + except ValidationError as e: + logger.warning(f"Invalid custom config file: {e}") + + return config + + +def get_prioritized_namespace_path() -> List[str]: + """ + Get the prioritized namespaces. + + priority: custom > merico > community + """ + config = _load_custom_config() + + namespaces = config.namespaces + + namespace_paths = [os.path.join(CUSTOM_BASE, ns) for ns in namespaces] + + namespace_paths.append(MERICO_WORKFLOWS) + namespace_paths.append(COMMUNITY_WORKFLOWS) + + return namespace_paths + + +def main(): + paths = get_prioritized_namespace_path() + for p in paths: + print(p) + + +if __name__ == "__main__": + main() diff --git a/devchat/workflow/path.py b/devchat/workflow/path.py new file mode 100644 index 00000000..49dc6a0e --- /dev/null +++ b/devchat/workflow/path.py @@ -0,0 +1,35 @@ +import os + +# ------------------------------- +# devchat basic paths +# ------------------------------- +USE_DIR = os.path.expanduser("~") +CHAT_DIR = os.path.join(USE_DIR, ".chat") + + +# ------------------------------- +# workflow scripts paths +# ------------------------------- +WORKFLOWS_BASE_NAME = "scripts" +WORKFLOWS_BASE = os.path.join(CHAT_DIR, WORKFLOWS_BASE_NAME) # TODO: a temporary name +WORKFLOWS_CONFIG_FILENAME = "config.yml" # Under WORKFLOWS_BASE + +MERICO_WORKFLOWS = os.path.join(WORKFLOWS_BASE, "merico") +COMMUNITY_WORKFLOWS = os.path.join(WORKFLOWS_BASE, "community") + +COMMAND_FILENAMES = ["command.yml", "command.yaml"] + + +# ------------------------------- +# config & settings files paths +# ------------------------------- +CUSTOM_BASE = os.path.join(WORKFLOWS_BASE, "custom") +CUSTOM_CONFIG_FILE = os.path.join(CUSTOM_BASE, "config.yml") +USER_SETTINGS_FILENAME = "user_settings.yml" # Under WORKFLOWS_BASE + + +# ------------------------------- +# Python environments paths +# ------------------------------- +MAMBA_ROOT = os.path.join(CHAT_DIR, "mamba") +MAMBA_PY_ENVS = os.path.join(MAMBA_ROOT, "envs") diff --git a/devchat/workflow/schema.py b/devchat/workflow/schema.py new file mode 100644 index 00000000..4b27a699 --- /dev/null +++ b/devchat/workflow/schema.py @@ -0,0 +1,60 @@ +import re +from typing import Optional, List, Dict, Union + +from pydantic import BaseModel, validator, Extra, ValidationError + + +class WorkflowPyConf(BaseModel): + version: str # python version + dependencies: str # absolute path to the requirements file + env_name: Optional[str] # python env name, will use the workflow name if not set + + @validator("version") + def validate_version(cls, value): # pylint: disable=no-self-argument + pattern = r"^\d+\.\d+(\.\d+)?$" + if not re.match(pattern, value): + raise ValidationError( + f"Invalid version format: {value}. Expected format is x.y or x.y.z" + ) + return value + + +class ExternalPyConf(BaseModel): + env_name: str # the env_name of workflow python to act as + python_bin: str # the python executable path + + +class UserSettings(BaseModel): + external_workflow_python: List[ExternalPyConf] = [] + + class Config: + extra = Extra.ignore + + +class WorkflowConfig(BaseModel): + description: str + root_path: str # the root path of the workflow + steps: List[Dict] + input_required: bool = False # True for required + hint: Optional[str] = None + workflow_python: Optional[WorkflowPyConf] = None + help: Optional[Union[str, Dict[str, str]]] = None + + @validator("input_required", pre=True) + def to_boolean(cls, value): # pylint: disable=no-self-argument + return value.lower() == "required" + + class Config: + extra = Extra.ignore + + +class RuntimeParameter(BaseModel): + model_name: str + devchat_python: str + workflow_python: str = "" + user_input: Optional[str] + history_messages: Optional[Dict] + parent_hash: Optional[str] + + class Config: + extra = Extra.allow diff --git a/devchat/workflow/step.py b/devchat/workflow/step.py new file mode 100644 index 00000000..4317a448 --- /dev/null +++ b/devchat/workflow/step.py @@ -0,0 +1,178 @@ +# pylint: disable=invalid-name + +import os +import sys +import threading +import subprocess +import shlex +import json +from typing import Dict, Tuple, List +from enum import Enum +from .schema import WorkflowConfig, RuntimeParameter +from .path import WORKFLOWS_BASE + + +class BuiltInVars(str, Enum): + """ + Built-in variables within the workflow step command. + """ + + devchat_python = "$devchat_python" + command_path = "$command_path" + user_input = "$input" + workflow_python = "$workflow_python" + + +class BuiltInEnvs(str, Enum): + """ + Built-in environment variables for the step subprocess. + """ + + llm_model = "LLM_MODEL" + parent_hash = "PARENT_HASH" + context_contents = "CONTEXT_CONTENTS" + + +class WorkflowStep: + def __init__(self, **kwargs): + """ + Initialize a workflow step with the given configuration. + """ + self._kwargs = kwargs + + @property + def command_raw(self) -> str: + """ + The raw command string from the config. + """ + return self._kwargs.get("run", "") + + def _setup_env( + self, wf_config: WorkflowConfig, rt_param: RuntimeParameter + ) -> Dict[str, str]: + """ + Setup the environment variables for the subprocess. + """ + command_raw = self.command_raw + + env = os.environ.copy() + + # set PYTHONPATH for the subprocess + python_path = env.get("PYTHONPATH", "") + devchat_python_path = env.get("DEVCHAT_PYTHONPATH", python_path) + new_paths = [WORKFLOWS_BASE] + if (BuiltInVars.devchat_python in command_raw) and devchat_python_path: + # only add devchat pythonpath when it's used in the command + new_paths.append(devchat_python_path) + + paths = [os.path.normpath(p) for p in new_paths] + paths = [p.replace("\\", "\\\\") for p in paths] + joined = os.pathsep.join(paths) + + env["PYTHONPATH"] = joined + env[BuiltInEnvs.llm_model] = rt_param.model_name or "" + env[BuiltInEnvs.parent_hash] = rt_param.parent_hash or "" + env[BuiltInEnvs.context_contents] = "" + if rt_param.history_messages: + # convert dict to json string + env[BuiltInEnvs.context_contents] = json.dumps(rt_param.history_messages) + + return env + + def _validate_and_interpolate( + self, wf_config: WorkflowConfig, rt_param: RuntimeParameter + ) -> List[str]: + """ + Validate the step configuration and interpolate variables in the command. + + Return the command parts as a list of strings. + """ + command_raw = self.command_raw + parts = shlex.split(command_raw) + + # if the command_raw use $workflow_python, + # it must be set in workflow config + if BuiltInVars.workflow_python in command_raw: + if not rt_param.workflow_python: + raise ValueError( + "The command uses $workflow_python, " + "but the workflow_python is not set yet." + ) + + args = [] + for p in parts: + arg = p + + if p.startswith(BuiltInVars.workflow_python): + if not rt_param.workflow_python: + raise ValueError( + "The command uses $workflow_python, " + "but the workflow_python is not set yet." + ) + arg = arg.replace(BuiltInVars.workflow_python, rt_param.workflow_python) + + if p.startswith(BuiltInVars.devchat_python): + arg = arg.replace(BuiltInVars.devchat_python, rt_param.devchat_python) + + if p.startswith(BuiltInVars.command_path): + # NOTE: 在文档中说明 command.yml 中表示路径采用 POSIX 标准 + # 即，使用 / 分隔路径，而非 \ (Windows) + path_parts = p.split("/") + # replace "$command_path" with the root path in path_parts + arg = os.path.join(wf_config.root_path, *path_parts[1:]) + + if BuiltInVars.user_input in p: + arg = arg.replace(BuiltInVars.user_input, rt_param.user_input) + + args.append(arg) + + return args + + + def run( + self, wf_config: WorkflowConfig, rt_param: RuntimeParameter + ) -> Tuple[int, str, str]: + """ + Run the step in a subprocess. + + Returns the return code, stdout, and stderr. + """ + # setup the environment variables + env = self._setup_env(wf_config, rt_param) + + command_args = self._validate_and_interpolate(wf_config, rt_param) + + def _pipe_reader(pipe, data, out_file): + """ + Read from the pipe, then write and save the data. + """ + while pipe: + pipe_data = pipe.read(1) + if pipe_data == "": + break + data["data"] += pipe_data + print(pipe_data, end="", file=out_file, flush=True) + + with subprocess.Popen( + command_args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + text=True, + ) as proc: + stdout_data, stderr_data = {"data": ""}, {"data": ""} + stdout_thread = threading.Thread( + target=_pipe_reader, args=(proc.stdout, stdout_data, sys.stdout) + ) + stderr_thread = threading.Thread( + target=_pipe_reader, args=(proc.stderr, stderr_data, sys.stderr) + ) + stdout_thread.start() + stderr_thread.start() + stdout_thread.join() + stderr_thread.join() + + proc.wait() + return_code = proc.returncode + return return_code, stdout_data["data"], stderr_data["data"] + diff --git a/devchat/workflow/user_setting.py b/devchat/workflow/user_setting.py new file mode 100644 index 00000000..fcacd5a9 --- /dev/null +++ b/devchat/workflow/user_setting.py @@ -0,0 +1,24 @@ +from pathlib import Path +import oyaml as yaml +from .schema import UserSettings +from .path import WORKFLOWS_BASE, USER_SETTINGS_FILENAME + + +def _load_user_settings() -> UserSettings: + """ + Load the user settings from the settings.yml file. + """ + settings_path = Path(WORKFLOWS_BASE) / USER_SETTINGS_FILENAME + if not settings_path.exists(): + return UserSettings() + + with open(settings_path, "r", encoding="utf-8") as f: + content = yaml.safe_load(f.read()) + + if content: + return UserSettings.parse_obj(content) + + return UserSettings() + +USER_SETTINGS = _load_user_settings() + diff --git a/devchat/workflow/workflow.py b/devchat/workflow/workflow.py new file mode 100644 index 00000000..ba327bd3 --- /dev/null +++ b/devchat/workflow/workflow.py @@ -0,0 +1,215 @@ +# pylint: disable=invalid-name + +import os +import sys +from typing import Optional, Tuple, List, Dict +import oyaml as yaml +from .step import WorkflowStep +from .schema import WorkflowConfig, RuntimeParameter +from .path import COMMAND_FILENAMES +from .namespace import get_prioritized_namespace_path + +from .env_manager import PyEnvManager, EXTERNAL_ENVS + + +class Workflow: + TRIGGER_PREFIX = "/" + HELP_FLAG_PREFIX = "--help" + + def __init__(self, config: WorkflowConfig): + self._config = config + + self._runtime_param = None + + @property + def config(self): + return self._config + + @property + def runtime_param(self): + return self._runtime_param + + @staticmethod + def parse_trigger(user_input: str) -> Tuple[Optional[str], Optional[str]]: + """ + Check if the user input should trigger a workflow. + Return a tuple of (workflow_name, the input without workflow trigger). + + User input is considered a workflow trigger if it starts with the Workflow.PREFIX. + The workflow name is the first word after the prefix. + """ + striped = user_input.strip() + if not striped: + return None, user_input + if striped[0] != Workflow.TRIGGER_PREFIX: + return None, user_input + + workflow_name = striped.split()[0][1:] + + # remove the trigger prefix and the workflow name + actual_input = user_input.replace( + f"{Workflow.TRIGGER_PREFIX}{workflow_name}", "", 1 + ) + return workflow_name, actual_input + + @staticmethod + def load(workflow_name: str) -> Optional["Workflow"]: + """ + Load a workflow from the command.yml by name. + A workflow name is the relative path of command.yml + to the /workflows dir joined by "." + e.g + - "unit_tests": means the command file of the workflow is unit_tests/command.yml + - "commit.en": means the command file is commit/en/command.yml + - "pr.review.zh": means the command file is pr/review/zh/command.yml + """ + path_parts = workflow_name.split(".") + if len(path_parts) < 1: + return None + # path_parts.append(COMMAND_FILENAME) + rel_path = os.path.join(*path_parts) + + found = False + workflow_dir = "" + prioritized_dirs = get_prioritized_namespace_path() + for wf_dir in prioritized_dirs: + for fn in COMMAND_FILENAMES: + yaml_file = os.path.join(wf_dir, rel_path, fn) + if os.path.exists(yaml_file): + workflow_dir = wf_dir + found = True + break + if found: + break + if not found: + return None + + # Load and override yaml conf in top-down order + config_dict = {} + for i in range(len(path_parts)): + cur_path = os.path.join(workflow_dir, *path_parts[: i + 1]) + for fn in COMMAND_FILENAMES: + cur_yaml = os.path.join(cur_path, fn) + + if os.path.exists(cur_yaml): + with open(cur_yaml, "r", encoding="utf-8") as file: + yaml_content = file.read() + cur_conf = yaml.safe_load(yaml_content) + cur_conf["root_path"] = cur_path + + # convert relative path to absolute path for dependencies file + if cur_conf.get("workflow_python", {}).get("dependencies"): + rel_dep = cur_conf["workflow_python"]["dependencies"] + abs_dep = os.path.join(cur_path, rel_dep) + cur_conf["workflow_python"]["dependencies"] = abs_dep + + config_dict.update(cur_conf) + + config = WorkflowConfig.parse_obj(config_dict) + + if config.workflow_python and config.workflow_python.env_name is None: + # use the workflow name as the env name if not set + config.workflow_python.env_name = workflow_name + + return Workflow(config) + + def setup( + self, + model_name: Optional[str], + user_input: Optional[str], + history_messages: Optional[List[Dict]], + parent_hash: Optional[str], + ): + """ + Setup the workflow with the runtime parameters and env variables. + """ + workflow_py = "" + if self.config.workflow_python: + pyconf = self.config.workflow_python + if pyconf.env_name in EXTERNAL_ENVS: + # Use the external python set in the user settings + workflow_py = EXTERNAL_ENVS[pyconf.env_name].python_bin + print( + "\n```Step\n# Using exteranl Python from user settings\n", + flush=True, + ) + print(f"env_name: {pyconf.env_name}") + print(f"python_bin: {workflow_py}") + print( + "\nThis Python environment's version and dependencies should be " + "ensured by the user to meet the requirements.", + ) + print("\n```", flush=True) + + else: + # TODO: 有没有更好的时机判断方法？既保证运行时一定安装了依赖、又不用每次都检查？ + # TODO: 只在插件(IDE)启动后workflow第一次使用时ensure环境和依赖？ + # Create workflow python env + manager = PyEnvManager() + workflow_py = manager.ensure(pyconf.env_name, pyconf.version) + + r_file = pyconf.dependencies + _ = manager.install(pyconf.env_name, r_file) + + runtime_param = { + # from user interaction + "model_name": model_name, + "user_input": user_input, + "history_messages": history_messages, + "parent_hash": parent_hash, + # from user setting or system + "devchat_python": sys.executable, + "workflow_python": workflow_py, + } + + self._runtime_param = RuntimeParameter.parse_obj(runtime_param) + + def run_steps(self) -> int: + """ + Run the steps of the workflow. + """ + steps = self.config.steps + + for s in steps: + step = WorkflowStep(**s) + result = step.run(self.config, self.runtime_param) + return_code = result[0] + if return_code != 0: + # stop the workflow if any step fails + return return_code + print("\n\n") + + return 0 + + def get_help_doc(self, user_input: str) -> str: + """ + Get the help doc content of the workflow. + """ + help_info = self.config.help + help_file = None + + if isinstance(help_info, str): + # return the only help doc + help_file = help_info + + elif isinstance(help_info, dict): + first = next(iter(help_info)) + default_file = help_info.get(first) + print(f"default_file: {default_file}") + + # get language code from user input + code = user_input.strip().removeprefix(Workflow.HELP_FLAG_PREFIX) + code = code.removeprefix(".").strip() + help_file = help_info.get(code, default_file) + + if not help_file: + return "" + + help_path = os.path.join(self.config.root_path, help_file) + if os.path.exists(help_path): + with open(help_path, "r", encoding="utf-8") as file: + return file.read() + return "" + + def should_show_help(self, user_input) -> bool: + return user_input.strip().startswith(Workflow.HELP_FLAG_PREFIX) 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