chat_completion.py

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)