Repository: geekan/MetaGPT
File: data/inference/run_api.py (Python)

Code before:
from metagpt.logs import logger
from metagpt.roles.di.data_interpreter import DataInterpreter
from metagpt.utils.recovery_util import save_history

dotenv.load_dotenv()

# Replace with your own
MAX_TOKEN = 128000


@retry(wait=wait_random_exponential(min=30, max=600), stop=stop_after_attempt(5))
async def call_chat(inputs, interpreter):
    """
    Calls the openai API to generate completions for the given inputs.

    Args:
    inputs (str): The inputs to generate completions for.
    interpreter (DataInterpreter): The data interpreter to use for execution.
    """
    system_messages = inputs.split("\n", 1)[0]
    user_message = inputs.split("\n", 1)[1]
    try:
        await interpreter.run([system_messages, user_message])
        code = interpreter.execute_code.nb.cells[-1].source
        return code
    except Exception as e:
        logger.error(e)
        logger.error(f"Inputs: {inputs}")
        traceback.print_exc()
        raise e


def gpt_tokenize(string: str, encoding) -> int:
    """Returns the number of tokens in a text string."""
    num_tokens = len(encoding.encode(string))
    return num_tokens


async def openai_inference(
    test_dataset,
    model_name_or_path,
    output_file,
    existing_ids,
):
    """
    Runs inference on a dataset using the openai API.

    Args:
    test_dataset (datasets.Dataset): The dataset to run inference on.
    model_name_or_path (str): The name or path of the model to use.
    output_file (str): The path to the output file.

Review comment (suggestion):
1. I think most of this function should be written inside the interpreter
2. The two lines of logger.error should be reduced to one line

Code after:
from metagpt.utils import count_string_tokens
from metagpt.utils.recovery_util import save_history

# Replace with your own
MAX_TOKEN = 128000


@retry(wait=wait_random_exponential(min=30, max=600), stop=stop_after_attempt(5))
async def call_chat(inputs, interpreter):
    """
    Calls the openai API to generate completions for the given inputs.

    Args:
    inputs (str): The inputs to generate completions for.
    interpreter (DataInterpreter): The data interpreter to use for execution.
    """
    requirement = "Please rewrite the code to address the issues existing in the repository and generate the correct code. Then, use the `git diff` command to output the patch based on the correct code."
    system_messages = inputs.split("\n", 1)[0]
    user_message = inputs.split("\n", 1)[1]
    try:
        await interpreter.run([requirement, system_messages, user_message])
        return interpreter.get_last_cell_source
    except Exception as e:
        logger.error(f"Error: {e}\nInputs: {inputs}")
        traceback.print_exc()
        raise e


async def openai_inference(
    test_dataset,
    model_name_or_path,
    output_file,
    existing_ids,
    use_reflection,
):
    """
    Runs inference on a dataset using the openai API.

    Args:
    test_dataset (datasets.Dataset): The dataset to run inference on.
    model_name_or_path (str): The name or path of the model to use.
    output_file (str): The path to the output file.
    existing_ids (set): A set of ids that have already been processed.
    """
    test_dataset = test_dataset.filter(
        lambda x: count_string_tokens(x["text"], model_name_or_path) <= MAX_TOKEN,
        desc="Filtering",
        load_from_cache_file=False,
    )
    basic_args = {
        "model_name_or_path": model_name_or_path,