oliverwang15/DAN_Chat

1import time2import json3from openai import OpenAI4import gradio as gr5 6class Backend:7    def __init__(self):8        self.return_instruction = """  Please only return in the following Json format:9        {{10        "Answer": "",11        "Reference Sentences": [""]12        }}"""13 14        self.chat_history = []15 16    def load_agent(self, openai_api_key, assistant_id):17        client = OpenAI(api_key=openai_api_key)18        assistant = client.beta.assistants.retrieve(assistant_id=assistant_id)19        return client, assistant20 21    def update_file(self, file_path):22        file = open(file_path, 'rb')23        file = self.client.files.create(file=file, purpose='assistants')24        return file25 26    def create_thread(self):27        thread = self.client.beta.threads.create()28        return thread29 30    def delate_thread(self, thread):31        self.client.beta.threads.delete(thread.id)32 33    def create_message(self, question, thread, file):34        message = self.client.beta.threads.messages.create(35            thread_id=thread.id,36            role="user",37            content= question + self.return_instruction,38            file_ids=[file.id]39        )40        return message41 42    def delate_message(self, message):43        self.client.beta.threads.messages.delete(message.id)44 45    def create_run(self, thread, assistant):46        run = self.client.beta.threads.runs.create(47            thread_id=thread.id,48            assistant_id=assistant.id,49            # instructions="""Please read PDF and answer the qusetions asked by users with professional knowledge."""50        )51        return run52    53    def delate_run(self, run):54        self.client.beta.threads.runs.delete(run.id)55 56    def get_massage(self, thread):57        messages = self.client.beta.threads.messages.list(58            thread_id=thread.id59        )60        return messages61 62    def phrase_massage(self, question, messages):63        mess = json.loads(messages.json())64        output = mess['data'][0]['content'][0]['text']['value']65        print(output)66        try:67            output = output.split("{")[1:]68            output = "{" + "".join(output)69            output = output.split("}")[:-1]70            output = "".join(output) + "}"   71            print(output)    72            output = eval(output)73            answer = output['Answer']74            reference = output['Reference Sentences']75        except:76            self.detete_message(message)77            answer = output78            reference = []79        reference = ' '.join(reference)80        reference = self.processing_html(reference)81        self.chat_history.append([question, answer])82        83        return self.chat_history, reference84 85    def phrase_massage_1(self, question, messages):86        mess = json.loads(messages.json())87        output = mess['data'][0]['content'][0]['text']['value']88        89        self.chat_history.append([question, output])90        91        return self.chat_history92 93    def processing_html(self, text):94        return f'<center><p> {text} </p></center>'95 96    def submit_passage(self, openai_key, assistant_id, file):97        # Create a new conversation98        self.client, self.assistant = self.load_agent(openai_key, assistant_id)99 100        # Update file101        self.file = self.update_file(file.name)102        103        # Create a new conversation104        self.thread = self.create_thread()105 106        gr.Info("Upload successful. Please can now chat with the assistant. Enjoy!")107 108    def submit_question(self, question):109        # print(question)110        # print(self.thread.id)111        # print(self.file.id)112        # Create a new message113        self.message = self.create_message(question, self.thread, self.file)114 115        # Create a new run116        run = self.create_run(self.thread, self.assistant)117 118        # Wait for the run to complete119        while True:120            run = self.client.beta.threads.runs.retrieve(thread_id=self.thread.id, run_id=run.id)121            if run.status not in ["queued", "in_progress"]:122                break123            time.sleep(1)124        125        # Get the answer126        messages = self.get_massage(self.thread)127        answer, reference = self.phrase_massage(question, messages)128 129        return answer, reference130 131    def submit_question_another(self, question):132        # Create a new message133        self.message = self.create_message(question, self.thread, self.file)134 135        # Create a new run136        run = self.create_run(self.thread, self.assistant)137 138        # Wait for the run to complete139        while True:140            run = self.client.beta.threads.runs.retrieve(thread_id=self.thread.id, run_id=run.id)141            if run.status not in ["queued", "in_progress"]:142                break143            time.sleep(1)144        145        # Get the answer146        messages = self.get_massage(self.thread)147        answer = self.phrase_massage_1(question, messages)148 149        return answer