Add promptlayer (#21)

buster/chatbot.py

@@ -1,9 +1,11 @@
 import logging
+import os
 from dataclasses import dataclass, field
 
 import numpy as np
 import openai
 import pandas as pd
+import promptlayer
 from openai.embeddings_utils import cosine_similarity, get_embedding
 
 from buster.docparser import EMBEDDING_MODEL, read_documents
@@ -11,6 +13,16 @@ from buster.docparser import EMBEDDING_MODEL, read_documents
 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:
+    logger.info("Enabling prompt layer...")
+    promptlayer.api_key = promptlayer_api_key
+
+    # replace openai with the promptlayer wrapper
+    openai = promptlayer.openai
+    openai.api_key = os.environ.get("OPENAI_API_KEY")
+
 
 @dataclass
 class ChatbotConfig:

requirements.txt

@@ -5,6 +5,7 @@ pandas
 tabulate
 tenacity
 tiktoken
+promptlayer
 openai
 
 # all openai[embeddings] deps, their list breaks our CI, see: https://github.com/openai/openai-python/issues/210