Module `ai_shell.code_generate`

Generate code for AI Shell using the docstrings as source data.

Currently, targets: - CLI interface - Toolkit - JsonSchema dict

Expand source code

"""
Generate code for AI Shell using the docstrings as source data.

Currently, targets:
- CLI interface
- Toolkit
- JsonSchema dict
"""

from ai_shell.code_generate.generate_cli import generate_the_cli
from ai_shell.code_generate.generate_schema import generate_the_schema
from ai_shell.code_generate.generate_toolkit import generate_the_toolkit

__all__ = ["generate_the_schema", "generate_the_toolkit", "generate_the_cli"]

Sub-modules

ai_shell.code_generate.generate_cli: Code generate reverse client
ai_shell.code_generate.generate_schema: Generates JsonSchema in form of a python file.
ai_shell.code_generate.generate_toolkit: Code generate reverse client
ai_shell.code_generate.method_to_jsonschema: Utility for generating jsonschema from class methods.

Functions

def generate_the_cli(target_file: str) ‑> None

Main entry point.

Args

target_file : str: The target file path for the generated code.

Expand source code

def generate_the_cli(target_file: str) -> None:
    """Main entry point.
    Args:
        target_file (str): The target file path for the generated code.
    """
    tools = ""
    for _ns, json_schema in schemas._SCHEMAS.items():
        for tool, _ in json_schema.items():
            tools += f'            "{tool}" : self.{tool},\n'

    meta: dict[str, dict[str, str]] = {}
    meta["headtail"] = {
        "module": "ai_shell.head_tail_tool",
        "class": "HeadTailTool",
    }
    meta["cat"] = {
        "module": "ai_shell.cat_tool",
        "class": "CatTool",
    }
    meta["cut"] = {
        "module": "ai_shell.cut_tool",
        "class": "CutTool",
    }
    meta["find"] = {
        "module": "ai_shell.find_tool",
        "class": "FindTool",
    }
    meta["git"] = {
        "module": "ai_shell.git_tool",
        "class": "GitTool",
    }
    meta["patch"] = {
        "module": "ai_shell.patch_tool",
        "class": "PatchTool",
    }
    meta["ls"] = {
        "module": "ai_shell.ls_tool",
        "class": "LsTool",
    }
    meta["grep"] = {
        "module": "ai_shell.grep_tool",
        "class": "GrepTool",
    }
    meta["token_counter"] = {
        "module": "ai_shell.token_tool",
        "class": "TokenCounterTool",
    }
    meta["todo"] = {
        "module": "ai_shell.todo_tool",
        "class": "TodoTool",
    }
    meta["insert"] = {
        "module": "ai_shell.insert_tool",
        "class": "InsertTool",
    }
    meta["replace"] = {
        "module": "ai_shell.replace_tool",
        "class": "ReplaceTool",
    }
    meta["rewrite"] = {
        "module": "ai_shell.rewrite_tool",
        "class": "RewriteTool",
    }
    meta["answer_collector"] = {
        "module": "ai_shell.answer_tool",
        "class": "AnswerCollectorTool",
    }
    meta["pytest"] = {
        "module": "ai_shell.pytest_tool",
        "class": "PytestTool",
    }

    header = """\"\"\"
Generated code, do not edit.
\"\"\"
import argparse
from ai_shell.utils.console_utils import pretty_console
from ai_shell.utils.config_manager import Config
from ai_shell.__about__ import __version__, __description__

CONFIG = Config()
# pylint: disable=unused-argument
"""

    for _key, value in meta.items():
        header += f"\nfrom {value['module']} import {value['class']}"
    header += "\n\n"

    middle = ""
    for ns, data in meta.items():
        middle += "\n"
        for method, method_data in schemas._SCHEMAS[ns].items():
            middle += "\n"
            middle += f"def {method}_command(args):\n"
            middle += f'    """Invoke {method}"""\n'
            middle += f"    tool = {data['class']}('.', CONFIG)\n"
            middle += f"    pretty_console(tool.{method}("
            for arg_name, _arg_details in cast(dict[str, Any], method_data["properties"]).items():
                if arg_name != "mime_type":
                    # don't know how to handle mime types in CLI yet.
                    middle += f"\n        {arg_name} = args.{arg_name},"
            middle += "\n    ))"

    argparse_part = """\n\ndef run():
        \"\"\"Create the main parser\"\"\"
        program = 'ais'
        parser = argparse.ArgumentParser(
        prog=program,
        allow_abbrev=False,
        description=__description__,
        epilog=f\"\"\"
    Examples:

        ais cat hello.py
\"\"\",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    parser.add_argument(
        "-V",
        "--version",
        action="version",
        version=f"%(prog)s {__version__}",
        help="Show program's version number and exit.",
    )
        
        )
        subparsers = parser.add_subparsers(dest='subcommand', help='sub-command help')
    """
    for ns, _data in meta.items():
        for method, method_data in schemas._SCHEMAS[ns].items():
            escaped_method_description = (
                cast(str, method_data.get("description", "")).replace("'", "\\'").replace("\n", "\\n")
            )
            argparse_part += f'    # Create a parser for the "{method}" command\n'
            argparse_part += f"""    {method}_parser = subparsers.add_parser('{method}', help=\"\"\"{escaped_method_description}.\"\"\")\n"""
            for arg_name, arg_details in cast(dict[str, Any], method_data["properties"]).items():
                dashed_arg_name = arg_name.replace("_", "-")
                escaped_description = arg_details.get("description", "").replace("'", "\\'")

                argparse_part += f"\n    {method}_parser.add_argument('--{dashed_arg_name}', "
                bool_part = ""
                if arg_details["type"] == "boolean":
                    bool_part = "action='store_true', "

                if method_data.get("default"):
                    argparse_part += (
                        f"    dest='{arg_name}', {bool_part} default={method_data['default']},\n"
                        f"    help='{escaped_description}, defaults to {method_data['default']}')\n"
                    )
                else:
                    argparse_part += f'    dest=\'{arg_name}\', {bool_part} help="""{escaped_description}""")\n'

                # if list
                # cat_parser.add_argument('file_paths', nargs='+', type=str, help='Paths to the files to be concatenated')

            argparse_part += f"    {method}_parser.set_defaults(func={method}_command)\n\n"

    argparse_part += """        # Parse the arguments
        args = parser.parse_args()
    """

    footer = """
    # Execute the appropriate command
    if hasattr(args, 'func'):
        args.func(args)
    else:
        parser.print_help()

if __name__ == '__main__':
    run()
"""

    # Generate method code
    with open(target_file, "w", encoding="utf-8") as code:
        code.write(header)
        code.write(middle)
        code.write(argparse_part)
        code.write(footer)

def generate_the_schema(target_file: str) ‑> None

Main entrypoint.

Args

target_file : str: The target file path for the generated code.

Expand source code

def generate_the_schema(target_file: str) -> None:
    """Main entrypoint.

    Args:
        target_file (str): The target file path for the generated code.
    """
    # Setup logging for my_app
    # We will only setup a console handler
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    ch = logging.StreamHandler()
    ch.setFormatter(logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s"))
    logger.addHandler(ch)

    schemas = {}
    schemas["cat"] = convert_to_json_schema(ai_shell.CatTool)
    schemas["cut"] = convert_to_json_schema(ai_shell.CutTool)
    schemas["ed"] = convert_to_json_schema(ai_shell.EdTool)
    schemas["sed"] = convert_to_json_schema(ai_shell.SedTool)
    schemas["insert"] = convert_to_json_schema(ai_shell.InsertTool)
    schemas["replace"] = convert_to_json_schema(ai_shell.ReplaceTool)
    schemas["headtail"] = convert_to_json_schema(ai_shell.HeadTailTool)
    schemas["edlin"] = convert_to_json_schema(ai_shell.EdlinTool)
    schemas["find"] = convert_to_json_schema(ai_shell.FindTool)
    schemas["git"] = convert_to_json_schema(ai_shell.GitTool)
    schemas["patch"] = convert_to_json_schema(ai_shell.PatchTool)
    schemas["ls"] = convert_to_json_schema(ai_shell.LsTool)
    schemas["grep"] = convert_to_json_schema(ai_shell.GrepTool)
    schemas["pycat"] = convert_to_json_schema(ai_shell.PyCatTool)
    schemas["token_counter"] = convert_to_json_schema(ai_shell.TokenCounterTool)
    schemas["todo"] = convert_to_json_schema(ai_shell.TodoTool)
    schemas["answer_collector"] = convert_to_json_schema(ai_shell.AnswerCollectorTool)
    schemas["rewrite"] = convert_to_json_schema(ai_shell.RewriteTool)
    schemas["pytest"] = convert_to_json_schema(ai_shell.PytestTool)

    with open(target_file, "w", encoding="utf-8") as source:
        source.write('"""jsonschema for functions"""')
        source.write("\n\n")
        source.write(f"_SCHEMAS = {pformat(schemas, width=500, indent=2)}")

def generate_the_toolkit(target_file: str) ‑> None

Main entry point

Args

target_file : str: Target file to write the generated code.

Expand source code

def generate_the_toolkit(target_file: str) -> None:
    """Main entry point
    Args:
        target_file (str): Target file to write the generated code.
    """
    tools = ""
    for _ns, json_schema in schemas._SCHEMAS.items():
        for tool, _ in json_schema.items():
            tools += f'            "{tool}" : self.{tool},\n'

    header = f"""\"\"\"
Generate code, do not edit.
\"\"\"
from ai_shell.openai_support import ToolKitBase
from typing import Any, cast, Callable, Optional

from ai_shell.utils.config_manager import Config
from ai_shell.cat_tool import CatTool
from ai_shell.find_tool import FindTool
from ai_shell.git_tool import GitTool
from ai_shell.grep_tool import GrepTool
from ai_shell.ls_tool import LsTool
from ai_shell.token_tool import TokenCounterTool
from ai_shell.edlin_tool import EdlinTool
from ai_shell.pycat_tool import PyCatTool
from ai_shell.ed_tool import EdTool
from ai_shell.cut_tool import CutTool
from ai_shell.head_tail_tool import HeadTailTool
from ai_shell.patch_tool import PatchTool
from ai_shell.sed_tool import SedTool
from ai_shell.insert_tool import InsertTool
from ai_shell.replace_tool import ReplaceTool
from ai_shell.todo_tool import TodoTool
from ai_shell.answer_tool import AnswerCollectorTool
from ai_shell.rewrite_tool import RewriteTool
from ai_shell.pytest_tool import PytestTool

# pylint: disable=unused-argument

class ToolKit(ToolKitBase):
    \"\"\"AI Shell Toolkit\"\"\"\n\n
    def __init__(self, root_folder: str, token_model: str, global_max_lines: int, permitted_tools: list[str], config:Config) -> None:
        super().__init__(root_folder, token_model, global_max_lines, permitted_tools, config)
        self._lookup: dict[str, Callable[[dict[str, Any]], Any]] = {{
            {tools}
        }}
        # Stateful tool support. Useless assignment to make mypy happy
        self.tool_answer_collector = AnswerCollectorTool(self.root_folder, self.config)
"""

    prologue = {}
    prologue["cat"] = "tool = CatTool(self.root_folder, self.config)"
    prologue["headtail"] = "tool = HeadTailTool(self.root_folder, self.config)"
    prologue["cut"] = "tool = CutTool(self.root_folder, self.config)"
    prologue["pycat"] = "tool = PyCatTool(self.root_folder, self.config)"
    prologue["ed"] = "tool = EdTool(self.root_folder, self.config)"
    prologue["sed"] = "tool = SedTool(self.root_folder, self.config)"
    prologue["insert"] = "tool = InsertTool(self.root_folder, self.config)"
    prologue["replace"] = "tool = ReplaceTool(self.root_folder, self.config)"
    prologue["edlin"] = "tool = EdlinTool(self.root_folder, self.config)"
    prologue["find"] = "tool = FindTool(self.root_folder, self.config)"
    prologue["git"] = "tool = GitTool(self.root_folder, self.config)"
    prologue["patch"] = "tool = PatchTool(self.root_folder, self.config)"
    prologue["ls"] = "tool = LsTool(self.root_folder, self.config)"
    prologue["grep"] = "tool = GrepTool(self.root_folder, self.config)"
    prologue["token_counter"] = "tool = TokenCounterTool(self.root_folder, self.config)"  # nosec
    prologue["todo"] = "tool = TodoTool(self.root_folder, self.config)"
    prologue["answer_collector"] = "self.tool_answer_collector = AnswerCollectorTool(self.root_folder, self.config)"
    prologue["rewrite"] = "tool = RewriteTool(self.root_folder, self.config)"
    prologue["pytest"] = "tool = PytestTool(self.root_folder, self.config)"

    # Generate method code
    with open(target_file, "w", encoding="utf-8") as code:
        code.write(header)
        for ns, json_schema in schemas._SCHEMAS.items():
            pro = prologue[ns]
            method_code = generate_method_code(json_schema, pro, ns)
            code.write(method_code)