💬

対話的にChat Completions APIにテストコードを書いてもらう

2023/12/30に公開

導入

Fusic Advent Calendar 2023の11日目を担当します@seike460です!
信じられないくらい遅れました!ごめんなさい!

OpenAI流行ってますね。
テストコードはOpenAIに任せると精度が高いという話を聞いたので、
テストコードを書いてもらおうと思います。
opt-outを気にしなくて良いChat Completions APIに、
CLIで読み込んだファイルを投げてサクッとテストコード書いてもらいます。
そもそもこれじゃテストコード書きづらいよ!って状態の時もあると思うので、
テストコードがかけるようにリファクタリングも行います。

要素の解説

OpenAIへのアクセス

from openai import OpenAI

client = OpenAI()

上記のclientを利用しながら動かしていきます。
今回はCLIを前提としている為、API_KEYはCLI自体に仕込んでいきます。

API KEYに関しては以下を御覧ください。

https://platform.openai.com/docs/quickstart/step-2-setup-your-api-key?context=python

なお可能な限りToken数を増やした方ので、「GPT-4 Turbo」を利用しています。

CLIのオプションを設定

読み込ませるファイル、どの様なモードで動作させるかをオプションにて指定させたいのでargparseを入れます。

pip install argparse

以下のように利用します。

parser = argparse.ArgumentParser(description='テストコードをかくやつ')
parser.add_argument('-f', '--file', type=str, help='読み込むファイル名')
parser.add_argument('-t', '--Test', action='store_true', help='テストを記述する')
parser.add_argument('-r', '--Refactoring', action='store_true',
                        help='テストが書きやすいようにリファクタリングしてもらう')
args = parser.parse_args()

#以下変数にてオプションにアクセス
args.file
args.Test
args.Refactoring

ファイルを読み込ませる

ファイルを読み込ませる必要があるので、サクッとopen()を利用してファイルを読み込みます。
一応存在チェックは行ってます。

def read_file_if_exists(self, filepath):
if filepath and os.path.exists(filepath.strip()):
    with open(filepath, 'r') as file:
	self.contents = file.read()

ユーザーの入力を受け付ける

それぞれのモードに追加して、指示を出せるようにinput()を利用して
ユーザーからの入力を受け付けるようにします。

def get_user_input(self):
print("入力してください")
user_input = ""
while True:
    line = input()
    if line.strip() == '':
	break
    if user_input:
	user_input += "\n"
    user_input += line
return user_input

Chat Completions APIに投げる

以下を参考に実行します。
わりと待つこともあるので実行中だよーってことがわかるように実行中だよって投げてから
Responseを待ちます。

https://github.com/openai/openai-python?tab=readme-ov-file#usage

まずはプロンプトを生成します。
今回はすごく簡素にしてます。

def process_user_input(self, user_input, test_mode, refactoring_mode):
if self.contents.strip() != "":
    if test_mode:
	user_input = f"""
次のコードのテストを可能な限りカバレッジ率を高くするように書いてください。
{user_input}
----- Code Start -----
{self.contents}
----- Code End -----
"""
    elif refactoring_mode:
	user_input = f"""
次のコードのテスト記述が出来るようにリファクタリングしてください。
{user_input}
----- Code Start -----
{self.contents}
----- Code End -----
"""
    else:
	user_input = f"""
{user_input}
----- Code Start -----
{self.contents}
----- Code End -----
"""
def get_response_from_gpt(self, max_tokens=4096):
print("\n---------- Go gpt ----------\n")
response = self.client.chat.completions.create(
    model='gpt-4-1106-preview',
    messages=self.conversation_history,
    max_tokens=max_tokens,
    seed=-1
)
ai_response = response.choices[0].message.content
print(f"GPT: {ai_response}")
self.conversation_history.append(
    {"role": "assistant", "content": ai_response}
)

実行してみる

自分自身をリファクタリングさせてみます。

python chatbot.py -f chatbot.py -r

リファクタリングされたコード

import argparse
import os
from openai import OpenAI

# 定数はモジュールレベルの定数として定義
RED = "\033[31m"
BLUE = "\033[34m"
RESET = "\033[0m"


class ChatbotWrapper:
    def __init__(self, client, input_provider=input, display_func=print):
        self.client = client
        self.conversation_history = []
        self.input_provider = input_provider
        self.display_func = display_func

    def display(self, message, color=RESET):
        self.display_func(f"{color}{message}{RESET}")

    def read_file_contents(self, filepath):
        if os.path.exists(filepath.strip()):
            with open(filepath, 'r') as file:
                return file.read()
        return ""

    def acquire_user_input(self):
        self.display("入力してください")
        user_input = []
        try:
            while True:
                line = self.input_provider()
                if line.strip() == '':
                    break
                user_input.append(line)
        except EOFError:
            pass  # End of File (Ctrl+D) is a valid way to end the input loop
        return "\n".join(user_input)

    def format_user_input(self, user_input, contents, mode):
        template = "\n----- Code Start -----\n{}\n----- Code End -----\n"
        prompt = {
            "test": "次のコードのテストを可能な限りカバレッジ率を高くするように書いてください。",
            "refactoring": "次のコードのテスト記述が出来るようにリファクタリングしてください。",
            None: ""
        }[mode]
        return f"{prompt}\n{user_input}{template.format(contents)}" if contents.strip() else user_input

    def append_to_conversation_history(self, role, content):
        if role in ["user", "assistant"] and content:
            self.conversation_history.append({"role": role, "content": content})

    def generate_api_payload(self, max_tokens=4096):
        return {
            'model': 'gpt-4-1106-preview',
            'messages': self.conversation_history,
            'max_tokens': max_tokens,
            'seed': -1
        }

    def get_response_from_gpt(self, max_tokens=4096):
        self.display("\n---------- Go gpt ----------\n", BLUE)
        payload = self.generate_api_payload(max_tokens)
        response = self.client.chat.completions.create(**payload)
        ai_response = response.choices[0].message.content
        self.display(f"GPT: {ai_response}", BLUE)
        self.append_to_conversation_history("assistant", ai_response)

    def display_conversation_history(self):
        self.display("\n会話履歴:")
        for message in self.conversation_history:
            color = RED if message['role'] == "user" else BLUE
            self.display(f"{message['role'].upper()}: {message['content']}", color)

    def run_chat_session(self, contents, mode):
        while True:
            user_input = self.acquire_user_input()
            if user_input.strip():
                formatted_input = self.format_user_input(user_input, contents, mode)
                self.append_to_conversation_history("user", formatted_input)
                self.get_response_from_gpt()
            else:
                self.display("空の入力が検出されました。")
                break
        self.display_conversation_history()


def main():
    parser = argparse.ArgumentParser(description='ChatBot Wrapper')
    parser.add_argument('-f', '--file', type=str, help='File to load')
    parser.add_argument('-t', '--test', action='store_const', const='test', dest='mode',
                        help='Describe tests to increase coverage')
    parser.add_argument('-r', '--refactor', action='store_const', const='refactoring', dest='mode',
                        help='Refactor to make it easier to write tests')
    args = parser.parse_args()

    client = OpenAI()
    chatbot = ChatbotWrapper(client)

    contents = chatbot.read_file_contents(args.file) if args.file else ""
    chatbot.run_chat_session(contents, args.mode)


if __name__ == "__main__":
    main()

続いてこのコードに対するテストコードを生成してみます

テストコード生成

python chatbot.py -f chatbot.py -t

生成されたテストコードは以下

import unittest
from unittest.mock import patch, Mock, mock_open
from chatbot import ChatbotWrapper, RED, BLUE, RESET  # chatbot.pyとして保存されていると仮定
import os

class TestChatbotWrapper(unittest.TestCase):

    def setUp(self):
        self.mock_client = Mock()

        # OpenAI APIのモック
        self.mock_response = Mock()
        self.mock_response.choices = [{"message": {"content": "Mocked Response"}}]
        self.mock_client.chat.completions.create.return_value = self.mock_response

        # os.path.existsを常にTrueを返すモックにする
        self.patcher = patch('os.path.exists', return_value=True)
        self.addCleanup(self.patcher.stop)
        self.patcher.start()

    def test_read_file_if_exists(self):
        # OpenAIクライアントをモックしてチャットボットを初期化
        chatbot_wrapper = ChatbotWrapper(self.mock_client)
        fake_filepath = "/path/to/fake_file.txt"

        # ファイルの内容をモック
        with patch('builtins.open', mock_open(read_data='fake file content')) as mocked_file:
            chatbot_wrapper.read_file_if_exists(fake_filepath)
            mocked_file.assert_called_with(fake_filepath, 'r')

        # ファイルの内容がセットされたことを確認
        self.assertEqual(chatbot_wrapper.contents, 'fake file content')

    def test_process_user_input_with_file(self):
        chatbot_wrapper = ChatbotWrapper(self.mock_client)
        chatbot_wrapper.contents = 'Test content'
        user_input = 'Sample user input'

        # テストモードがTrueのときの処理
        processed_input = chatbot_wrapper.process_user_input(user_input, test_mode=True, refactoring_mode=False)
        self.assertIn('次のコードのテストを可能な限りカバレッジ率を高くするように書いてください。', processed_input)

        # リファクタリングモードがTrueのときの処理
        processed_input = chatbot_wrapper.process_user_input(user_input, test_mode=False, refactoring_mode=True)
        self.assertIn('次のコードのテスト記述が出来るようにリファクタリングしてください。', processed_input)

        # どちらのモードでもないときの処理
        processed_input = chatbot_wrapper.process_user_input(user_input, test_mode=False, refactoring_mode=False)
        self.assertIn(user_input, processed_input)

    # 必要に応じて他のテストケースを追加...

if __name__ == '__main__':
    unittest.main()

coverage率を測ってみました。

pip install coverage
coverage report
Name              Stmts   Miss  Cover
-------------------------------------
chatbot.py           68     38    44%
chatbot_test.py      32      1    97%
-------------------------------------
TOTAL               100     39    61%

もうちょっと粘る

何度かリファクタリングとテストコードを書き直してもらった結果が以下

(例)リファクタリング

python chatbot.py -f chatbot.py -r
入力してください
テストコードを記述しやすいように可能な限りメソッド分割を行ってください。

(例)テストコード

python chatbot.py -f chatbot.py -t
入力してください
可能な限りカバレッジ率を100%に近づける事
  • 結果
Name              Stmts   Miss  Cover
-------------------------------------
chatbot.py           74     21    72%
chatbot_test.py      50      1    98%
-------------------------------------
TOTAL               124     22    82%

少し上がりましたね。
テストコードを書かせるという意味では、最初の取っ掛かりとして良いのかと思いました。
ただし、このコードはだめかもなと思う修正や、そのまま動かない修正もあったりしたので、
その部分はよしなに調整する必要はありました。

趣味の範囲での利用や、プロジェクトの取っ掛かりとして利用する分には使えるかもなと思いましたが、
やはり全てを任せるというのはハードルがありましたので、あくまで壁打ち役として利用するのが良さそうです。
それでも便利だと思いましたので、今後も利用方法を検討していきたいです。

Fusic 技術ブログ

Discussion