-
Notifications
You must be signed in to change notification settings - Fork 0
/
api.py
37 lines (28 loc) · 885 Bytes
/
api.py
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
from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from config import get_config
config = get_config()
OPENAI_API_BASE = config["bot"]["openai_api_base"]
OPENAI_API_KEY = config["bot"]["openai_api_key"]
"""
Handles API calls to the OPENAI API
"""
class InvalidApiResponseError(RuntimeError):
pass
def chat_completion(history):
chat = ChatOpenAI(
openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY
)
messages = history.get_messages()
response = chat.generate(messages=[messages])
try:
return response.generations[0][0].text
except:
raise InvalidApiResponseError
def prompt_completion(prompt):
print(prompt)
llm = OpenAI(openai_api_base=OPENAI_API_BASE, openai_api_key=OPENAI_API_KEY)
try:
return llm.invoke(prompt)
except:
raise InvalidApiResponseError