tiger94me/ChatPaper

1"""2A simple wrapper for the official ChatGPT API3"""4import json5import os6import threading7import time8import requests9import tiktoken10from typing import Generator11from queue import PriorityQueue as PQ12import json13import os14import time15ENCODER = tiktoken.get_encoding("gpt2")16class chatPaper:17    """18    Official ChatGPT API19    """20    def __init__(21        self,22        api_keys: list,23        proxy = None,24        api_proxy = None,25        max_tokens: int = 4000,26        temperature: float = 0.5,27        top_p: float = 1.0,28        model_name: str = "gpt-3.5-turbo",29        reply_count: int = 1,30        system_prompt = "You are ChatPaper, A paper reading bot",31        lastAPICallTime = time.time()-100,32        apiTimeInterval = 20,33    ) -> None:34        self.model_name = model_name35        self.system_prompt = system_prompt36        self.apiTimeInterval = apiTimeInterval37        self.session = requests.Session()38        self.api_keys = PQ()39        for key in api_keys:40            self.api_keys.put((lastAPICallTime,key))41        self.proxy = proxy42        if self.proxy:43            proxies = {44                "http": self.proxy,45                "https": self.proxy,46            }47            self.session.proxies = proxies48        self.max_tokens = max_tokens49        self.temperature = temperature50        self.top_p = top_p51        self.reply_count = reply_count52        self.decrease_step = 25053        self.conversation = {}54        if self.token_str(self.system_prompt) > self.max_tokens:55            raise Exception("System prompt is too long")56        self.lock = threading.Lock()57        58    def get_api_key(self):59        with self.lock:60            apiKey = self.api_keys.get()61            delay = self._calculate_delay(apiKey)62            time.sleep(delay)63            self.api_keys.put((time.time(), apiKey[1]))64            return apiKey[1]65 66    def _calculate_delay(self, apiKey):67        elapsed_time = time.time() - apiKey[0]68        if elapsed_time < self.apiTimeInterval:69            return self.apiTimeInterval - elapsed_time70        else:71            return 072 73    def add_to_conversation(self, message: str, role: str, convo_id: str = "default"):74        if(convo_id not in self.conversation):75            self.reset(convo_id)76        self.conversation[convo_id].append({"role": role, "content": message})77 78    def __truncate_conversation(self, convo_id: str = "default"):79        """80        Truncate the conversation81        """82        last_dialog = self.conversation[convo_id][-1]83        query = str(last_dialog['content'])84        if(len(ENCODER.encode(str(query)))>self.max_tokens):85            query = query[:int(1.5*self.max_tokens)]86        while(len(ENCODER.encode(str(query)))>self.max_tokens):87            query = query[:self.decrease_step]88        self.conversation[convo_id] = self.conversation[convo_id][:-1]89        full_conversation = "\n".join([str(x["content"]) for x in self.conversation[convo_id]],)90        if len(ENCODER.encode(full_conversation)) > self.max_tokens:91            self.conversation_summary(convo_id=convo_id)92        full_conversation = ""93        for x in self.conversation[convo_id]:94            full_conversation = str(x["content"]) + "\n" + full_conversation95        while True:96            if (len(ENCODER.encode(full_conversation+query)) > self.max_tokens):97                query = query[:self.decrease_step]98            else:99                break100        last_dialog['content'] = str(query)101        self.conversation[convo_id].append(last_dialog)102 103    def ask_stream(104        self,105        prompt: str,106        role: str = "user",107        convo_id: str = "default",108        **kwargs,109    ) -> Generator:110        if convo_id not in self.conversation:111            self.reset(convo_id=convo_id)112        self.add_to_conversation(prompt, "user", convo_id=convo_id)113        self.__truncate_conversation(convo_id=convo_id)114        apiKey = self.get_api_key()115        response = self.session.post(116            "https://api.openai.com/v1/chat/completions",117            headers={"Authorization": f"Bearer {kwargs.get('api_key', apiKey)}"},118            json={119                "model": self.model_name,120                "messages": self.conversation[convo_id],121                "stream": True,122                # kwargs123                "temperature": kwargs.get("temperature", self.temperature),124                "top_p": kwargs.get("top_p", self.top_p),125                "n": kwargs.get("n", self.reply_count),126                "user": role,127            },128            stream=True,129        )130        if response.status_code != 200:131            raise Exception(132                f"Error: {response.status_code} {response.reason} {response.text}",133            )134        for line in response.iter_lines():135            if not line:136                continue137            # Remove "data: "138            line = line.decode("utf-8")[6:]139            if line == "[DONE]":140                break141            resp: dict = json.loads(line)142            choices = resp.get("choices")143            if not choices:144                continue145            delta = choices[0].get("delta")146            if not delta:147                continue148            if "content" in delta:149                content = delta["content"]150                yield content151    def ask(self, prompt: str, role: str = "user", convo_id: str = "default", **kwargs):152        """153        Non-streaming ask154        """155        response = self.ask_stream(156            prompt=prompt,157            role=role,158            convo_id=convo_id,159            **kwargs,160        )161        full_response: str = "".join(response)162        self.add_to_conversation(full_response, role, convo_id=convo_id)163        usage_token = self.token_str(prompt)164        com_token = self.token_str(full_response)165        total_token = self.token_cost(convo_id=convo_id)166        return full_response, usage_token, com_token, total_token167 168    def check_api_available(self):169        response = self.session.post(170            "https://api.openai.com/v1/chat/completions",171            headers={"Authorization": f"Bearer {self.get_api_key()}"},172            json={173                "model": self.model_name,174                "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "print A"}],175                "stream": True,176                # kwargs177                "temperature": self.temperature,178                "top_p": self.top_p,179                "n": self.reply_count,180                "user": "user",181            },182            stream=True,183        )184        if response.status_code == 200:185            return True186        else:187            return False188    def reset(self, convo_id: str = "default", system_prompt = None):189        """190        Reset the conversation191        """192        self.conversation[convo_id] = [193            {"role": "system", "content": str(system_prompt or self.system_prompt)},194        ]195    def conversation_summary(self, convo_id: str = "default"):196        input = ""197        role = ""198        for conv in self.conversation[convo_id]:199            if (conv["role"]=='user'):200                role = 'User'201            else:202                role = 'ChatGpt'203            input+=role+' : '+conv['content']+'\n'204        prompt = "Your goal is to summarize the provided conversation in English. Your summary should be concise and focus on the key information to facilitate better dialogue for the large language model.Ensure that you include all necessary details and relevant information while still reducing the length of the conversation as much as possible. Your summary should be clear and easily understandable for the ChatGpt model providing a comprehensive and concise summary of the conversation."205        if(self.token_str(str(input)+prompt)>self.max_tokens):206            input = input[self.token_str(str(input))-self.max_tokens:]207        while self.token_str(str(input)+prompt)>self.max_tokens:208            input = input[self.decrease_step:]209        prompt = prompt.replace("{conversation}", input)210        self.reset(convo_id='conversationSummary')211        response = self.ask(prompt,convo_id='conversationSummary')212        while self.token_str(str(response))>self.max_tokens:213            response = response[:-self.decrease_step]214        self.reset(convo_id='conversationSummary',system_prompt='Summariaze our diaglog')215        self.conversation[convo_id] = [216            {"role": "system", "content": self.system_prompt},217            {"role": "user", "content": "Summariaze our diaglog"},218            {"role": 'assistant', "content": response},219        ]220        return self.conversation[convo_id]221    def token_cost(self,convo_id: str = "default"):222        return len(ENCODER.encode("\n".join([x["content"] for x in self.conversation[convo_id]])))223    def token_str(self,content:str):224        return len(ENCODER.encode(content))225def main():226    return227