diff --git a/buster/completers/base.py b/buster/completers/base.py index 7bd5c04..0f1fa80 100644 --- a/buster/completers/base.py +++ b/buster/completers/base.py @@ -1,9 +1,7 @@ import io import logging -import os import warnings from abc import ABC, abstractmethod -from dataclasses import dataclass, field from typing import Any, Iterator, Optional import openai @@ -16,23 +14,6 @@ logger = logging.getLogger(__name__) logging.basicConfig(level=logging.INFO) -# Check if an API key exists for promptlayer, if it does, use it -promptlayer_api_key = os.environ.get("PROMPTLAYER_API_KEY") -if promptlayer_api_key: - try: - import promptlayer - - logger.info("Enabling prompt layer...") - promptlayer.api_key = promptlayer_api_key - - # replace openai with the promptlayer wrapper - openai = promptlayer.openai - except Exception as e: - logger.exception("Something went wrong enabling promptlayer.") - -# Set openai credentials -openai.api_key = os.environ.get("OPENAI_API_KEY") - class Completion: def __init__( diff --git a/buster/completers/chatgpt.py b/buster/completers/chatgpt.py index 69eae48..ffa01cc 100644 --- a/buster/completers/chatgpt.py +++ b/buster/completers/chatgpt.py @@ -1,9 +1,28 @@ +import logging +import os from typing import Iterator import openai from buster.completers import Completer +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + +# Check if an API key exists for promptlayer, if it does, use it +promptlayer_api_key = os.environ.get("PROMPTLAYER_API_KEY") +if promptlayer_api_key: + try: + import promptlayer + + logger.info("Enabling prompt layer...") + promptlayer.api_key = promptlayer_api_key + + # replace openai with the promptlayer wrapper + openai = promptlayer.openai + except Exception as e: + logger.exception("Something went wrong enabling promptlayer.") + class ChatGPTCompleter(Completer): def complete(self, prompt: str, user_input, completion_kwargs=None) -> str | Iterator: