DeepCodeGeniusWeb-tools/site-packages/devchat/_cli/router.py

# pylint: disable=import-outside-toplevel
import json
import sys
from typing import List, Optional

from devchat.workflow.workflow import Workflow


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)

    parameters_data = config.dict(exclude_unset=True)
    if config_str:
        config_data = json.loads(config_str)
        parameters_data.update(config_data)
    return model, parameters_data

def _load_tool_functions(functions: Optional[str]):
    try:
        if functions:
            with open(functions, 'r', encoding="utf-8") as f_file:
                return json.load(f_file)
        return None
    except Exception:
        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:
        instruct_contents.extend(command_instructions)

    return instruct_contents


def before_prompt(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,
           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 = sys.stdin.read()

    if content == '':
        raise MissContentInPromptException()

    instruct_contents = _load_instruction_contents(content, instruct)
    context_contents = parse_files(context)
    tool_functions = _load_tool_functions(functions)

    model, parameters_data = _get_model_and_config(model, config_str)
    max_input_tokens = parameters_data.get("max_input_tokens", 4000)

    openai_config = OpenAIChatConfig(model=model, **parameters_data)
    chat = OpenAIChat(openai_config)
    chat_store = Store(repo_chat_dir, chat)

    assistant = Assistant(chat, chat_store, max_input_tokens, not not_store)
    assistant.make_prompt(
        request = content,
        instruct_contents = instruct_contents,
        context_contents = context_contents,
        functions = tool_functions,
        parent=parent,
        references=reference,
        function_name=function_name
    )

    return model, assistant, content

def llm_prompt(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,
           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
		)

        print(assistant.prompt.formatted_header())
        for response in assistant.iterate_response():
            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
		)

        print(assistant.prompt.formatted_header())
        command_result = run_command(
            model_name = model,
            history_messages = assistant.prompt.messages,
            input_text = content,
            parent_hash = parent,
            auto_fun = False)
        if command_result is not None:
            sys.exit(0)

        print("run command fail.")
        print(command_result)
    sys.exit(-1)


def llm_route(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,
           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
		)

        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,
            history_messages = assistant.prompt.messages,
            input_text = content,
            parent_hash = parent,
            auto_fun = auto)
        if command_result is not None:
            sys.exit(command_result[0])

        for response in assistant.iterate_response():
            print(response, end='', flush=True)