xhy/Quantum-based-PaperRead

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 time15import base6416import uuid17def random_machine_id():18    uuid_bytes = uuid.uuid4().bytes19    return base64.b64encode(uuid_bytes).decode('utf-8')20 21ENCODER = tiktoken.get_encoding("gpt2")22class chatPaper:23    """24    Official ChatGPT API25    """26    def __init__(27        self,28        api_keys: list,29        proxy = None,30        api_proxy = None,31        max_tokens: int = 4000,32        temperature: float = 0.5,33        top_p: float = 1.0,34        model_name: str = "gpt-3.5-turbo",35        reply_count: int = 1,36        system_prompt = "You are ChatPaper, A paper reading bot",37        lastAPICallTime = time.time()-100,38        apiTimeInterval = 20,39    ) -> None:40        self.model_name = "gpt-3.5-turbo"41        self.system_prompt = system_prompt42        self.apiTimeInterval = apiTimeInterval43        self.session = requests.Session()44        self.api_keys = PQ()45        for key in api_keys:46            self.api_keys.put((lastAPICallTime,key))47        self.proxy = proxy48        if self.proxy:49            proxies = {50                "http": self.proxy,51                "https": self.proxy,52            }53            self.session.proxies = proxies54        self.max_tokens = max_tokens55        self.temperature = temperature56        self.top_p = top_p57        self.reply_count = reply_count58        self.decrease_step = 25059        self.conversation = {}60        if self.token_str(self.system_prompt) > self.max_tokens:61            raise Exception("System prompt is too long")62        self.lock = threading.Lock()63        64    def get_api_key(self):65        with self.lock:66            apiKey = self.api_keys.get()67            delay = self._calculate_delay(apiKey)68            time.sleep(delay)69            self.api_keys.put((time.time(), apiKey[1]))70            return apiKey[1]71 72    def _calculate_delay(self, apiKey):73        elapsed_time = time.time() - apiKey[0]74        if elapsed_time < self.apiTimeInterval:75            return self.apiTimeInterval - elapsed_time76        else:77            return 078 79    def add_to_conversation(self, message: str, role: str, convo_id: str = "default"):80        if(convo_id not in self.conversation):81            self.reset(convo_id)82        self.conversation[convo_id].append({"role": role, "content": message})83 84    def __truncate_conversation(self, convo_id: str = "default"):85        """86        Truncate the conversation87        """88        last_dialog = self.conversation[convo_id][-1]89        query = str(last_dialog['content'])90        if(len(ENCODER.encode(str(query)))>self.max_tokens):91            query = query[:int(1.5*self.max_tokens)]92        while(len(ENCODER.encode(str(query)))>self.max_tokens):93            query = query[:self.decrease_step]94        self.conversation[convo_id] = self.conversation[convo_id][:-1]95        full_conversation = "\n".join([str(x["content"]) for x in self.conversation[convo_id]],)96        if len(ENCODER.encode(full_conversation)) > self.max_tokens:97            self.conversation_summary(convo_id=convo_id)98        full_conversation = ""99        for x in self.conversation[convo_id]:100            full_conversation = str(x["content"]) + "\n" + full_conversation101        while True:102            if (len(ENCODER.encode(full_conversation+query)) > self.max_tokens):103                query = query[:self.decrease_step]104            else:105                break106        last_dialog['content'] = str(query)107        self.conversation[convo_id].append(last_dialog)108 109 110    def ask_stream(111        self,112        prompt: str,113        role: str = "user",114        convo_id: str = "default",115        **kwargs,116    ) -> Generator:117        if convo_id not in self.conversation:118            self.reset(convo_id=convo_id)119        self.add_to_conversation(prompt, "user", convo_id=convo_id)120        self.__truncate_conversation(convo_id=convo_id)121        apiKey = self.get_api_key()122        response = self.session.post(123            "https://api.vekun.com/v1/chat/completions",124            headers={125                'machineId': random_machine_id(),126                'key': None127            },128            json={129                "model": self.model_name,130                "messages": self.conversation[convo_id],131                "stream": True,132                # kwargs133                "temperature": kwargs.get("temperature", self.temperature),134                "top_p": kwargs.get("top_p", self.top_p),135                "n": kwargs.get("n", self.reply_count),136                "user": role,137            },138            stream=True,139        )140        if response.status_code != 200:141            raise Exception(142                f"Error: {response.status_code} {response.reason} {response.text}",143            )144        for line in response.iter_lines():145            if not line:146                continue147            # Remove "data: "148            line = line.decode("utf-8")[6:]149            if line == "[DONE]":150                break151            resp: dict = json.loads(line)152            choices = resp.get("choices")153            if not choices:154                continue155            delta = choices[0].get("delta")156            if not delta:157                continue158            if "content" in delta:159                content = delta["content"]160                yield content161    def ask(self, prompt: str, role: str = "user", convo_id: str = "default", **kwargs):162        """163        Non-streaming ask164        """165        response = self.ask_stream(166            prompt=prompt,167            role=role,168            convo_id=convo_id,169            **kwargs,170        )171        full_response: str = "".join(response)172        self.add_to_conversation(full_response, role, convo_id=convo_id)173        usage_token = self.token_str(prompt)174        com_token = self.token_str(full_response)175        total_token = self.token_cost(convo_id=convo_id)176        return full_response, usage_token, com_token, total_token177 178    def check_api_available(self):179        response = self.session.post(180            "https://api.vekun.com/v1/chat/completions",181            headers={182                'machineId': random_machine_id(),183                'key': None184            },185            json={186                "model": self.model_name,187                "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "print A"}],188                "stream": True,189                # kwargs190                "temperature": self.temperature,191                "top_p": self.top_p,192                "n": self.reply_count,193                "user": "user",194            },195            stream=True,196        )197        if response.status_code == 200:198            return True199        else:200            return False201    def reset(self, convo_id: str = "default", system_prompt = None):202        """203        Reset the conversation204        """205        self.conversation[convo_id] = [206            {"role": "system", "content": str(system_prompt or self.system_prompt)},207        ]208    def conversation_summary(self, convo_id: str = "default"):209        input = ""210        role = ""211        for conv in self.conversation[convo_id]:212            if (conv["role"]=='user'):213                role = 'User'214            else:215                role = 'ChatGpt'216            input+=role+' : '+conv['content']+'\n'217        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."218        if(self.token_str(str(input)+prompt)>self.max_tokens):219            input = input[self.token_str(str(input))-self.max_tokens:]220        while self.token_str(str(input)+prompt)>self.max_tokens:221            input = input[self.decrease_step:]222        prompt = prompt.replace("{conversation}", input)223        self.reset(convo_id='conversationSummary')224        response = self.ask(prompt,convo_id='conversationSummary')225        while self.token_str(str(response))>self.max_tokens:226            response = response[:-self.decrease_step]227        self.reset(convo_id='conversationSummary',system_prompt='Summariaze our diaglog')228        self.conversation[convo_id] = [229            {"role": "system", "content": self.system_prompt},230            {"role": "user", "content": "Summariaze our diaglog"},231            {"role": 'assistant', "content": response},232        ]233        return self.conversation[convo_id]234    def token_cost(self,convo_id: str = "default"):235        return len(ENCODER.encode("\n".join([x["content"] for x in self.conversation[convo_id]])))236    def token_str(self,content:str):237        return len(ENCODER.encode(content))238def main():239    return240