Spaces:
Runtime error
Runtime error
File size: 2,219 Bytes
6ba5e8d 5759e22 6ba5e8d a797463 6ba5e8d |
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 |
import linecache
import re
from typing import Dict, List, Optional
import openai
class ChatCompletion:
def __init__(self, model: str = 'gpt-3.5-turbo',
api_key: Optional[str] = None, api_key_path: str = './openai_api_key'):
if api_key is None:
openai.api_key = api_key
api_key = linecache.getline(api_key_path, 2).strip('\n')
if len(api_key) == 0:
raise EnvironmentError
openai.api_key = api_key
self.model = model
self.system_messages = []
self.user_messages = []
def chat(self, msg: str, setting: Optional[str] = None, model: Optional[str] = None) -> str:
if self._context_length() > 2048:
self.reset()
if setting is not None:
if setting not in self.system_messages:
self.system_messages.append(setting)
if not self.user_messages or msg != self.user_messages[-1]:
self.user_messages.append(msg)
return self._run(model)
def retry(self, model: Optional[str] = None) -> str:
return self._run(model)
def reset(self):
self.system_messages.clear()
self.user_messages.clear()
def _make_message(self) -> List[Dict]:
sys_messages = [{'role': 'system', 'content': msg} for msg in self.system_messages]
user_messages = [{'role': 'user', 'content': msg} for msg in self.user_messages]
return sys_messages + user_messages
def _context_length(self) -> int:
return len(''.join(self.system_messages)) + len(''.join(self.user_messages))
def _run(self, model: Optional[str] = None) -> str:
if model is None:
model = self.model
try:
response = openai.ChatCompletion.create(model=model, messages=self._make_message())
ans = response['choices'][0]['message']['content']
ans = re.sub(r'^\n+', '', ans)
except openai.error.OpenAIError as e:
ans = e
except Exception as e:
print(e)
return ans
def __call__(self, msg: str, setting: Optional[str] = None, model: Optional[str] = None) -> str:
return self.chat(msg, setting, model)
|