From 6deeacfb05349da2d8b5a02a055de90af0d9d130 Mon Sep 17 00:00:00 2001
From: Jeremy Pinto <jerpint@gmail.com>
Date: Wed, 27 Sep 2023 14:46:13 -0400
Subject: [PATCH] fix import for promptlayer

---
 buster/completers/base.py    | 19 -------------------
 buster/completers/chatgpt.py | 19 +++++++++++++++++++
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/buster/completers/base.py b/buster/completers/base.py
index 7bd5c048..0f1fa806 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 69eae487..ffa01ccf 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: