llmriddles/llms/chatgpt.py

import os
from functools import lru_cache

import openai

from .base import register_llm


@lru_cache()
def _setup_openai():
    current_path = os.path.dirname(os.path.realpath(__file__))
    parent_dir = os.path.dirname(current_path)
    if 'OPENAI_KEY' in os.environ:
        openai.api_key = os.environ['OPENAI_KEY']
    else:
        openai.api_key_path = f'{parent_dir}/.key'
    return openai.api_key_path


def ask_chatgpt(message: str):
    _setup_openai()

    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {"role": "user", "content": message}
        ],
    )
    return response["choices"][0]["message"]["content"].strip()


register_llm('chatgpt', ask_chatgpt)