llmbb/LLMBB-Agent

1import json2import os3from pathlib import Path4import gradio as gr5import jsonlines6from agent.actions import RetrievalQA7from agent.llm import ChatAsOAI8from agent.memory import Memory9from utils import service, cache_file, max_ref_token10 11llm = ChatAsOAI(model="gpt-3.5-turbo")12mem = Memory(llm=llm, stream=False)13 14with open('css/main.css', 'r') as f:15    css = f.read()16with open('js/main.js', 'r') as f:17    js = f.read()18 19 20def add_text(history, text):21    history = history + [(text, None)]22    return history, gr.update(value='', interactive=False)23 24 25def rm_text(history):26    if not history:27        gr.Warning('No input content!')28    elif not history[-1][1]:29        return history, gr.update(value='', interactive=False)30    else:31        history = history[:-1] + [(history[-1][0], None)]32        return history, gr.update(value='', interactive=False)33 34 35def add_file(history, file):36    history = history + [((file.name,), None)]37    return history38 39 40def initialize(request: gr.Request):41    # print(request.kwargs)42    access_token = request.query_params["access_token"]43    url = request.query_params["url"]44    is_valid = False45    if access_token:46        account_info = json.loads(service.get(access_token, "info.json", False))47        if account_info and account_info["enabled"]:48            is_valid = True49    if not is_valid:50        gr.Info("The token is not valid, Please reset!")51        return52    return access_token, url53 54 55def bot(history, access_token, page_url):56    if not history:57        yield history58    else:59        now_page = None60        _ref = ''61        if not service.exists(access_token, page_url):62            gr.Info("Please add this page to LLMBB's Reading List first!")63        else:64            now_page = json.loads(service.get(access_token, page_url))65            if not now_page:66                gr.Info(67                    "This page has not yet been added to the LLMBB's reading list!"68                )69            elif not now_page['raw']:70                gr.Info('Please reopen later, LLMBB is analyzing this page...')71            else:72                _ref_list = mem.get(73                    history[-1][0], [now_page],74                    max_token=max_ref_token)75                if _ref_list:76                    _ref = '\n'.join(77                        json.dumps(x, ensure_ascii=False) for x in _ref_list)78                else:79                    _ref = ''80 81        # TODO: considering history for retrieval qa82        agent = RetrievalQA(stream=True, llm=llm)83        history[-1][1] = ''84        response = agent.run(user_request=history[-1][0], ref_doc=_ref)85 86        for chunk in response:87            if chunk is not None:88                history[-1][1] += chunk89            yield history90 91        # save history92        if now_page:93            now_page['session'] = history94            service.upsert(access_token, page_url, json.dumps(now_page, ensure_ascii=False))95 96 97def load_history_session(history, access_token, page_url):98    now_page = None99    if not service.exists(access_token, page_url):100        gr.Info("Please add this page to LLMBB's Reading List first!")101        return []102    now_page = json.loads(service.get(access_token, page_url))103    if not now_page:104        gr.Info("Please add this page to LLMBB's Reading List first!")105        return []106    if not now_page['raw']:107        gr.Info('Please wait, LLMBB is analyzing this page...')108        return []109    return now_page['session']110 111 112def clear_session(access_token, page_url):113    if not service.exists(access_token, page_url):114        return None115    now_page = json.loads(service.get(access_token, page_url))116    if not now_page:117        return None118    now_page['session'] = []119    service.upsert(access_token, page_url, json.dumps(now_page, ensure_ascii=False))120    return None121 122 123with gr.Blocks(css=css, theme='soft') as demo:124    access_token = gr.State("")125    page_url = gr.State("")126    chatbot = gr.Chatbot([], elem_id='chatbot', height=480, avatar_images=(None, 'img/logo.png'))127    with gr.Row():128        with gr.Column(scale=7):129            txt = gr.Textbox(show_label=False,130                             placeholder='Chat with LLMBB...',131                             container=False)132        # with gr.Column(scale=0.06, min_width=0):133        #     smt_bt = gr.Button('⏎')134        with gr.Column(scale=1, min_width=0):135            clr_bt = gr.Button('🧹', elem_classes='bt_small_font')136        with gr.Column(scale=1, min_width=0):137            stop_bt = gr.Button('🚫', elem_classes='bt_small_font')138        with gr.Column(scale=1, min_width=0):139            re_bt = gr.Button('🔁', elem_classes='bt_small_font')140 141    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt],142                         queue=False).then(bot, [chatbot, access_token, page_url], chatbot)143    txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)144 145    # txt_msg_bt = smt_bt.click(add_text, [chatbot, txt], [chatbot, txt],146    #                           queue=False).then(bot, chatbot, chatbot)147    # txt_msg_bt.then(lambda: gr.update(interactive=True),148    #                 None, [txt],149    #                 queue=False)150 151    clr_bt.click(clear_session, [access_token, page_url], chatbot, queue=False)152    re_txt_msg = re_bt.click(rm_text, [chatbot], [chatbot, txt],153                             queue=False).then(bot, [chatbot, access_token, page_url], chatbot)154    re_txt_msg.then(lambda: gr.update(interactive=True),155                    None, [txt],156                    queue=False)157 158    stop_bt.click(None, None, None, cancels=[txt_msg, re_txt_msg], queue=False)159 160    demo.load(initialize, [], [access_token, page_url]).then(load_history_session, [chatbot, access_token, page_url], chatbot)161    demo.queue()162 163# demo.queue().launch(server_name=server_config.server.server_host, server_port=server_config.server.app_in_browser_port)164