QualityFlow：LLMが協調してプログラムを高品質化する新たなワークフロー

!pip install openai

import os
import openai
import re
import traceback

# ★★★ ここにあなたのAPIキーを設定してください ★★★
openai.api_key = "sk-XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"

CODE_GENERATOR_PROMPT = """\
You are a helpful coding assistant. 
Given the following problem statement and test cases, generate a Python function or code to solve the problem.

Problem Statement:
{problem_statement}

Test Cases (example inputs/outputs):
{test_cases}

Constraints:
- Return only the code (no extra explanations).
"""

TEST_DESIGNER_PROMPT = """\
You are a test-case designer. We want more test cases for the following problem:

Problem Statement:
{problem_statement}

Current Code (for reference):
{current_code}

Current Test Cases:
{current_test_cases}

Please generate up to 10 new test cases in Python's unittest-like format (just the function calls and expected values).
"""

SELF_DEBUGGER_PROMPT = """\
You are a self-debugging agent. 
We have a problem statement and some test results (with errors or mismatches).
Analyze the error messages and the current code, then propose a corrected code.

Problem Statement:
{problem_statement}

Current Code:
{current_code}

Test Failures (error outputs or mismatch details):
{test_failures}

Now provide a revised version of the code that fixes these issues.
"""

QUALITY_CHECKER_PROMPT = """\
You are a code quality checker. You will perform an imagined execution of the following code against each test case and reason step-by-step about the output.
Only answer "PASS" if you are absolutely certain the code would produce the correct outputs for all test cases. Otherwise, answer "FAIL".

Code:
{current_code}

Test Cases (input -> expected):
{test_cases}

Your step-by-step reasoning:
"""

TEST_QUALITY_CHECKER_PROMPT = """\
You are a test case quality checker. 
We have newly generated test cases for the following problem:

Problem Statement:
{problem_statement}

Test Cases:
{test_cases}

Analyze if these test cases are valid, relevant, and correctly framed to test the solution. 
List the invalid or irrelevant test cases if any, and explain briefly.
"""

PROBLEM_CLARIFIER_PROMPT = """\
You are a clarifier agent. We have repeatedly failed to solve the following problem. 
Review the original problem statement, and consider alternative interpretations or missing details.

Problem Statement:
{problem_statement}

History of attempts or known issues:
{history}

Provide a revised or clarified problem statement that may resolve the ambiguities.
"""

# OpenAI APIを呼び出す簡単なラッパ関数
def call_openai_api(prompt, model="gpt-3.5-turbo", temperature=0.7):
    completion = openai.ChatCompletion.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
    )
    return completion.choices[0].message.content.strip()

class CodeGenerator:
    def __init__(self, api_caller):
        self.api_caller = api_caller

    def generate_code(self, problem_statement: str, test_cases: str) -> str:
        prompt = CODE_GENERATOR_PROMPT.format(
            problem_statement=problem_statement,
            test_cases=test_cases
        )
        response = self.api_caller(prompt)
        return response

class TestDesigner:
    def __init__(self, api_caller):
        self.api_caller = api_caller

    def generate_test_cases(self, problem_statement: str, current_code: str, current_test_cases: str) -> str:
        prompt = TEST_DESIGNER_PROMPT.format(
            problem_statement=problem_statement,
            current_code=current_code,
            current_test_cases=current_test_cases
        )
        response = self.api_caller(prompt)
        return response

class SelfDebugger:
    def __init__(self, api_caller):
        self.api_caller = api_caller

    def debug_code(self, problem_statement: str, current_code: str, test_failures: str) -> str:
        prompt = SELF_DEBUGGER_PROMPT.format(
            problem_statement=problem_statement,
            current_code=current_code,
            test_failures=test_failures
        )
        response = self.api_caller(prompt)
        return response

class QualityChecker:
    def __init__(self, api_caller):
        self.api_caller = api_caller

    def check_code_quality(self, current_code: str, test_cases: str) -> bool:
        prompt = QUALITY_CHECKER_PROMPT.format(
            current_code=current_code,
            test_cases=test_cases
        )
        response = self.api_caller(prompt)
        # 非常に単純化した判定
        if "PASS" in response:
            return True
        return False

class TestQualityChecker:
    def __init__(self, api_caller):
        self.api_caller = api_caller

    def check_test_quality(self, problem_statement: str, test_cases: str):
        prompt = TEST_QUALITY_CHECKER_PROMPT.format(
            problem_statement=problem_statement,
            test_cases=test_cases
        )
        response = self.api_caller(prompt)
        # 簡易的に「Invalid Test Case: ～」を検索
        invalid_cases = re.findall(r"Invalid Test Case: (.+)", response)
        return invalid_cases

class ProblemClarifier:
    def __init__(self, api_caller):
        self.api_caller = api_caller

    def clarify_problem(self, problem_statement: str, history: str) -> str:
        prompt = PROBLEM_CLARIFIER_PROMPT.format(
            problem_statement=problem_statement,
            history=history
        )
        response = self.api_caller(prompt)
        return response

class QualityFlow:
    def __init__(
        self,
        code_generator,
        test_designer,
        self_debugger,
        quality_checker,
        test_quality_checker,
        problem_clarifier,
        max_attempts=3,
        max_clarify=1
    ):
        self.code_generator = code_generator
        self.test_designer = test_designer
        self.self_debugger = self_debugger
        self.quality_checker = quality_checker
        self.test_quality_checker = test_quality_checker
        self.problem_clarifier = problem_clarifier

        self.max_attempts = max_attempts
        self.max_clarify = max_clarify

    def run(self, problem_statement, initial_test_cases):
        """
        QualityFlowを実行する主なメソッド。
        :param problem_statement: 文字列で問題文
        :param initial_test_cases: 文字列またはフォーマット化されたテストケース
        :return: 最終的に得られたプログラム
        """

        # (1) Code Generator: 初期コード生成
        current_code = self.code_generator.generate_code(problem_statement, initial_test_cases)

        attempt_count = 0
        clarify_count = 0
        history_log = []

        while True:
            # (2) Quality Checker: 想像実行（Imaged Execution）による正誤判定
            is_passing = self.quality_checker.check_code_quality(current_code, initial_test_cases)

            if is_passing:
                print("[INFO] Code passed the quality check. Returning result.")
                return current_code  # 合格したので終了

            # 不合格の場合
            attempt_count += 1
            history_log.append(f"Attempt {attempt_count} failed.")
            print(f"[WARN] Quality check failed at attempt {attempt_count}.")

            if attempt_count > self.max_attempts:
                # ある回数以上失敗したら問題クラリファイアを呼ぶ
                clarify_count += 1
                if clarify_count > self.max_clarify:
                    print("[ERROR] Reached max clarify limit. Returning the latest code as fallback.")
                    return current_code  # 妥協策: 最後のコードを返して終了
                else:
                    print("[WARN] Invoking ProblemClarifier...")
                    # 問題文を修正し、最初から作り直す
                    revised_problem = self.problem_clarifier.clarify_problem(problem_statement, "\n".join(history_log))
                    history_log.append("Problem was clarified.")
                    # 改訂された問題文で再度最初からコード生成
                    current_code = self.code_generator.generate_code(revised_problem, initial_test_cases)
                    attempt_count = 0
                    problem_statement = revised_problem
                    continue

            # (3) テストデザイナーで新しいテストケースを生成
            new_test_cases = self.test_designer.generate_test_cases(problem_statement, current_code, initial_test_cases)
            invalid_cases = self.test_quality_checker.check_test_quality(problem_statement, new_test_cases)

            # ※ここでは単純に生成されたテストをそのまま追加
            combined_test_cases = f"{initial_test_cases}\n# Additional test cases:\n{new_test_cases}"

            # (4) Self Debugger: 実行してエラー収集 → コード修正
            test_failures_info = self._run_tests_and_collect_failures(current_code, combined_test_cases)
            current_code = self.self_debugger.debug_code(problem_statement, current_code, test_failures_info)

    def _run_tests_and_collect_failures(self, code_str, test_str):
        """
        簡易実装: exec()でコードを実行してテストする。安全上注意。
        test_strには "function_call -> expected" の形式を想定。
        """
        test_failures = []

        try:
            local_env = {}
            exec(code_str, local_env)  # 関数定義などを取り込む
        except Exception as e:
            return f"[ERROR] Runtime error in code:\n{traceback.format_exc()}"

        lines = test_str.strip().splitlines()
        for line in lines:
            line = line.strip()
            if "->" in line:
                parts = line.split("->")
                func_call = parts[0].strip()
                expected = parts[1].strip()
                try:
                    result = eval(func_call, local_env)
                    # 厳密な比較は用途次第。ここでは文字列比較。
                    if str(result) != expected:
                        test_failures.append(f"Test failed: {func_call}, expected={expected}, got={result}")
                except Exception as e:
                    test_failures.append(f"Runtime error: {func_call}, error={str(e)}")

        if not test_failures:
            return "All tests passed."
        else:
            return "\n".join(test_failures)

# エージェントのインスタンスを作成
code_generator = CodeGenerator(call_openai_api)
test_designer = TestDesigner(call_openai_api)
self_debugger = SelfDebugger(call_openai_api)
quality_checker = QualityChecker(call_openai_api)
test_quality_checker = TestQualityChecker(call_openai_api)
problem_clarifier = ProblemClarifier(call_openai_api)

# QualityFlowを組み立て
flow = QualityFlow(
    code_generator,
    test_designer,
    self_debugger,
    quality_checker,
    test_quality_checker,
    problem_clarifier,
    max_attempts=3,
    max_clarify=1
)

# 問題文（例：足し算関数）
problem_statement = """\
Write a Python function `add_two_numbers(a, b)` that returns the sum of a and b.
If a or b is not an integer, try to convert them to integer before summing.
"""

# 初期テストケース（例）
initial_test_cases = """\
add_two_numbers(1, 2) -> 3
add_two_numbers(10, 5) -> 15
"""

# 実行
final_code = flow.run(problem_statement, initial_test_cases)

print("=== Final Code Output ===")
print(final_code)