# AUTOGENERATED! DO NOT EDIT! File to edit: 00_wrapper.ipynb (unless otherwise specified).

__all__ = ['wrapper']

# Cell
import openai, pandas, numpy, datetime, json, os

# Cell
class wrapper:
    """
    Wrapper holds the OpenAI API
    """

    autosave=False

    def __init__(self, outdir="queries"):
        """
        build the wrapper, inheriting above
        """
        self.outdir = outdir
        os.makedirs(outdir, exist_ok=True)

    def add_key(self, key:str):
        """
        sets the openai api key
        """
        openai.api_key = key

    def toggleAutoSave(self, status:bool):
        """
        set self.autosave to true or false
        """
        self.autosave=status

    def query(self, verbose=False, **kwargs):
        """
        send kwargs to the API
        """

        r=""

        if verbose:
            print(kwargs)

        r = openai.Completion.create(**kwargs)

        if self.autosave:
            with open("{}/{}".format(self.outdir, datetime.datetime.now().strftime("%Y%m%d%H%m%S")), "w") as fh:
                json.dump({"promt":kwargs["prompt"], "response":r}, fh, indent=4)

        return r
