Spaces:
Runtime error
Runtime error
File size: 3,080 Bytes
c8a1687 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import logging
import os
from abc import ABC, abstractmethod
import openai
import promptlayer
from buster.formatter.base import Response
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")
class Completer(ABC):
def __init__(self, cfg):
self.cfg = cfg
@abstractmethod
def complete(self, prompt) -> str:
...
def generate_response(self, user_input, documents) -> Response:
# Call the API to generate a response
prompt = self.prepare_prompt(user_input, documents)
name = self.cfg["name"]
logger.info(f"querying model {name}...")
logger.info(f"{prompt=}")
try:
completion_kwargs = self.cfg["completion_kwargs"]
completion = self.complete(prompt=prompt, **completion_kwargs)
except Exception as e:
# log the error and return a generic response instead.
logger.exception("Error connecting to OpenAI API. See traceback:")
return Response("", True, "We're having trouble connecting to OpenAI right now... Try again soon!")
return Response(completion)
class GPT3Completer(Completer):
def prepare_prompt(
self,
user_input: str,
documents: str,
) -> str:
"""
Prepare the prompt with prompt engineering.
"""
text_before_docs = self.cfg["text_before_documents"]
text_before_prompt = self.cfg["text_before_prompt"]
return text_before_docs + documents + text_before_prompt + user_input
def complete(self, prompt, **completion_kwargs):
response = openai.Completion.create(prompt=prompt, **completion_kwargs)
return response["choices"][0]["text"]
class ChatGPTCompleter(Completer):
def prepare_prompt(
self,
user_input: str,
documents: str,
) -> list:
"""
Prepare the prompt with prompt engineering.
"""
text_before_docs = self.cfg["text_before_documents"]
text_before_prompt = self.cfg["text_before_prompt"]
prompt = [
{"role": "system", "content": text_before_docs + documents + text_before_prompt},
{"role": "user", "content": user_input},
]
return prompt
def complete(self, prompt, **completion_kwargs) -> str:
response = openai.ChatCompletion.create(
messages=prompt,
**completion_kwargs,
)
return response["choices"][0]["message"]["content"]
def get_completer(completer_cfg):
name = completer_cfg["name"]
completers = {
"GPT3": GPT3Completer,
"ChatGPT": ChatGPTCompleter,
}
return completers[name](completer_cfg)
|