PythonでちょっとChatGPT APIを使いたい時にちょうどいいclassを作った

　まだChatGPT APIを使っていろいろ試してるんですが、そろそろAPI呼び出し部分をちゃんとclassでまとめるかと思って作ったコードをChatGPTにコードレビューさせて以下の形になりました。

　なお、このClassは単発の会話用なので、会話のラリーをしたい場合は適切に会話履歴を保存するようにして下さい。（そのメソッドも作ろうとしたけど、面倒くさくなって辞めちゃった。）

ソースコード

secret.json

{
    "OPENAI_API_KEY": "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
}

chatgpt_text_gen.py

# ChatGPT APIを呼び出してテキストを生成するプログラム
import json
import openai
from typing import Optional

# ChatGPT APIを呼び出してテキストを生成するクラス
class ChatGPT:
    # 初期化時にはAPIキーをJSONファイルから読み込む
    def __init__(self, secret_file: str = "secret.json") -> None:
        # JSONファイルからAPIキーを読み込む
        try:
            with open(secret_file) as f:
                secrets = json.load(f)
                openai.api_key: str = secrets.get("OPENAI_API_KEY")
        except Exception as e:
            print(f"Error loading API key: {e}")
            return
        # 使用するモデルを指定
        self.model = "gpt-3.5-turbo"

    # プロンプトからテキストを生成
    def generate_text(
            self, 
            prompt: str, 
            max_tokens: Optional[int] = 100
            ) -> str:
        # messagesを作成
        messages = [{"role": "user", "content": prompt}]
        # OpenAI APIを呼び出してテキストを生成
        try:
            response = openai.ChatCompletion.create(
                model=self.model, 
                messages=messages, 
                max_tokens=max_tokens)
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error generating text: {e}")
            return ""
        
    # ファイルから固定部分のプロンプトを読み込み、追加のテキストと結合してテキストを生成
    def generate_from_file(
            self, 
            filepath: str, 
            additional_text: str, 
            max_tokens: Optional[int] = 100
            ) -> str:
        # ファイルからプロンプトを読み込む
        try:
            with open(filepath, 'r', encoding="utf-8") as f:
                file_prompt = f.read()
        except Exception as e:
            print(f"Error loading file: {e}")
            return ""
        # ファイルの内容と追加のテキストを結合
        full_prompt = file_prompt + additional_text
        # 結合したプロンプトを使ってテキストを生成
        return self.generate_text(full_prompt, max_tokens)

　私はAPIキーはsecret.jsonで保存する派なので、環境変数から読み込む派の人はその辺書き換えてください。（環境変数から読み込む方がセキュリティ的に良いです。）

使い方

　入力文章をそのままChatGPT APIに送る場合はgenerate_textメソッドを使う。

from chatgpt_text_gen import ChatGPT

chat_gpt = ChatGPT()
print(chat_gpt.generate_text("AIについて30文字以内で説明して下さい"))

AIは人工知能のことで、機械が自ら学習や判断を行い人間に代わってタスクを遂行する技術の総称です。

　ChatGPT APIの役割が決まっている場合は読み込むファイルに役割を書いてgenerate_from_fileメソッドを使う。

prompt.txt

次の文章を日本語訳してください：

from chatgpt_text_gen import ChatGPT

chat_gpt = ChatGPT()
additional_text = "Hello, how are you?"
print(chat_gpt.generate_from_file('prompt.txt', additional_text))

こんにちは、お元気ですか？

　近々このclassを使って作ったクソアプリの記事を書きます。