Update app.py

app.py

@@ -1,136 +1,154 @@
-import asyncio
 import json
-import logging
-from copy import deepcopy
-from dataclasses import asdict
-from typing import Dict, List, Union
-
-import janus
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
+import os
+
+import gradio as gr
+import requests
 from lagent.schema import AgentStatusCode
-from pydantic import BaseModel
-from sse_starlette.sse import EventSourceResponse
-
-from mindsearch.agent import init_agent
-
-
-def parse_arguments():
-    import argparse
-    parser = argparse.ArgumentParser(description='MindSearch API')
-    parser.add_argument('--lang', default='cn', type=str, help='Language')
-    parser.add_argument('--model_format',
-                        default='internlm_server',
-                        type=str,
-                        help='Model format')
-    parser.add_argument('--search_engine',
-                       default='DuckDuckGoSearch',
-                       type=str,
-                       help='Search engine')
-    return parser.parse_args()
-
-
-args = parse_arguments()
-app = FastAPI(docs_url='/')
-
-app.add_middleware(CORSMiddleware,
-                   allow_origins=['*'],
-                   allow_credentials=True,
-                   allow_methods=['*'],
-                   allow_headers=['*'])
-
-
-class GenerationParams(BaseModel):
-    inputs: Union[str, List[Dict]]
-    agent_cfg: Dict = dict()
-
-
-@app.post('/solve')
-async def run(request: GenerationParams):
-
-    def convert_adjacency_to_tree(adjacency_input, root_name):
-
-        def build_tree(node_name):
-            node = {'name': node_name, 'children': []}
-            if node_name in adjacency_input:
-                for child in adjacency_input[node_name]:
-                    child_node = build_tree(child['name'])
-                    child_node['state'] = child['state']
-                    child_node['id'] = child['id']
-                    node['children'].append(child_node)
-            return node
-
-        return build_tree(root_name)
-
-    async def generate():
-        try:
-            queue = janus.Queue()
-            stop_event = asyncio.Event()
-
-            # Wrapping a sync generator as an async generator using run_in_executor
-            def sync_generator_wrapper():
-                try:
-                    for response in agent.stream_chat(inputs):
-                        queue.sync_q.put(response)
-                except Exception as e:
-                    logging.exception(
-                        f'Exception in sync_generator_wrapper: {e}')
-                finally:
-                    # Notify async_generator_wrapper that the data generation is complete.
-                    queue.sync_q.put(None)
-
-            async def async_generator_wrapper():
-                loop = asyncio.get_event_loop()
-                loop.run_in_executor(None, sync_generator_wrapper)
-                while True:
-                    response = await queue.async_q.get()
-                    if response is None:  # Ensure that all elements are consumed
-                        break
-                    yield response
-                    if not isinstance(
-                            response,
-                            tuple) and response.state == AgentStatusCode.END:
-                        break
-                stop_event.set()  # Inform sync_generator_wrapper to stop
-
-            async for response in async_generator_wrapper():
-                if isinstance(response, tuple):
-                    agent_return, node_name = response
-                else:
-                    agent_return = response
-                    node_name = None
-                origin_adj = deepcopy(agent_return.adjacency_list)
-                adjacency_list = convert_adjacency_to_tree(
-                    agent_return.adjacency_list, 'root')
-                assert adjacency_list[
-                    'name'] == 'root' and 'children' in adjacency_list
-                agent_return.adjacency_list = adjacency_list['children']
-                agent_return = asdict(agent_return)
-                agent_return['adj'] = origin_adj
-                response_json = json.dumps(dict(response=agent_return,
-                                                current_node=node_name),
-                                           ensure_ascii=False)
-                yield {'data': response_json}
-                # yield f'data: {response_json}\n\n'
-        except Exception as exc:
-            msg = 'An error occurred while generating the response.'
-            logging.exception(msg)
-            response_json = json.dumps(
-                dict(error=dict(msg=msg, details=str(exc))),
-                ensure_ascii=False)
-            yield {'data': response_json}
-            # yield f'data: {response_json}\n\n'
-        finally:
-            await stop_event.wait(
-            )  # Waiting for async_generator_wrapper to stop
-            queue.close()
-            await queue.wait_closed()
-
-    inputs = request.inputs
-    agent = init_agent(lang=args.lang, model_format=args.model_format,search_engine=args.search_engine)
-    return EventSourceResponse(generate())
-
-
-if __name__ == '__main__':
-    import uvicorn
-    uvicorn.run(app, host='0.0.0.0', port=8002, log_level='info')
+
+os.system("python -m mindsearch.app --lang cn --model_format internlm_silicon &")
+
+PLANNER_HISTORY = []
+SEARCHER_HISTORY = []
+
+
+def rst_mem(history_planner: list, history_searcher: list):
+    '''
+    Reset the chatbot memory.
+    '''
+    history_planner = []
+    history_searcher = []
+    if PLANNER_HISTORY:
+        PLANNER_HISTORY.clear()
+    return history_planner, history_searcher
+
+
+def format_response(gr_history, agent_return):
+    if agent_return['state'] in [
+            AgentStatusCode.STREAM_ING, AgentStatusCode.ANSWER_ING
+    ]:
+        gr_history[-1][1] = agent_return['response']
+    elif agent_return['state'] == AgentStatusCode.PLUGIN_START:
+        thought = gr_history[-1][1].split('```')[0]
+        if agent_return['response'].startswith('```'):
+            gr_history[-1][1] = thought + '\n' + agent_return['response']
+    elif agent_return['state'] == AgentStatusCode.PLUGIN_END:
+        thought = gr_history[-1][1].split('```')[0]
+        if isinstance(agent_return['response'], dict):
+            gr_history[-1][
+                1] = thought + '\n' + f'```json\n{json.dumps(agent_return["response"], ensure_ascii=False, indent=4)}\n```'  # noqa: E501
+    elif agent_return['state'] == AgentStatusCode.PLUGIN_RETURN:
+        assert agent_return['inner_steps'][-1]['role'] == 'environment'
+        item = agent_return['inner_steps'][-1]
+        gr_history.append([
+            None,
+            f"```json\n{json.dumps(item['content'], ensure_ascii=False, indent=4)}\n```"
+        ])
+        gr_history.append([None, ''])
+    return
+
+
+def predict(history_planner, history_searcher):
+
+    def streaming(raw_response):
+        for chunk in raw_response.iter_lines(chunk_size=8192,
+                                             decode_unicode=False,
+                                             delimiter=b'\n'):
+            if chunk:
+                decoded = chunk.decode('utf-8')
+                if decoded == '\r':
+                    continue
+                if decoded[:6] == 'data: ':
+                    decoded = decoded[6:]
+                elif decoded.startswith(': ping - '):
+                    continue
+                response = json.loads(decoded)
+                yield (response['response'], response['current_node'])
+
+    global PLANNER_HISTORY
+    PLANNER_HISTORY.append(dict(role='user', content=history_planner[-1][0]))
+    new_search_turn = True
+
+    url = 'http://localhost:8002/solve'
+    headers = {'Content-Type': 'application/json'}
+    data = {'inputs': PLANNER_HISTORY}
+    raw_response = requests.post(url,
+                                 headers=headers,
+                                 data=json.dumps(data),
+                                 timeout=20,
+                                 stream=True)
+
+    for resp in streaming(raw_response):
+        agent_return, node_name = resp
+        if node_name:
+            if node_name in ['root', 'response']:
+                continue
+            agent_return = agent_return['nodes'][node_name]['detail']
+            if new_search_turn:
+                history_searcher.append([agent_return['content'], ''])
+                new_search_turn = False
+            format_response(history_searcher, agent_return)
+            if agent_return['state'] == AgentStatusCode.END:
+                new_search_turn = True
+            yield history_planner, history_searcher
+        else:
+            new_search_turn = True
+            format_response(history_planner, agent_return)
+            if agent_return['state'] == AgentStatusCode.END:
+                PLANNER_HISTORY = agent_return['inner_steps']
+            yield history_planner, history_searcher
+    return history_planner, history_searcher
+
+
+with gr.Blocks() as demo:
+    gr.HTML("""<h1 align="center">MindSearch Gradio Demo</h1>""")
+    gr.HTML("""<p style="text-align: center; font-family: Arial, sans-serif;">MindSearch is an open-source AI Search Engine Framework with Perplexity.ai Pro performance. You can deploy your own Perplexity.ai-style search engine using either closed-source LLMs (GPT, Claude) or open-source LLMs (InternLM2.5-7b-chat).</p>""")
+    gr.HTML("""
+    <div style="text-align: center; font-size: 16px;">
+        <a href="https://github.com/InternLM/MindSearch" style="margin-right: 15px; text-decoration: none; color: #4A90E2;">🔗 GitHub</a>
+        <a href="https://arxiv.org/abs/2407.20183" style="margin-right: 15px; text-decoration: none; color: #4A90E2;">📄 Arxiv</a>
+        <a href="https://huggingface.co/papers/2407.20183" style="margin-right: 15px; text-decoration: none; color: #4A90E2;">📚 Hugging Face Papers</a>
+        <a href="https://huggingface.co/spaces/internlm/MindSearch" style="text-decoration: none; color: #4A90E2;">🤗 Hugging Face Demo</a>
+    </div>
+    """)
+    with gr.Row():
+        with gr.Column(scale=10):
+            with gr.Row():
+                with gr.Column():
+                    planner = gr.Chatbot(label='planner',
+                                         height=700,
+                                         show_label=True,
+                                         show_copy_button=True,
+                                         bubble_full_width=False,
+                                         render_markdown=True)
+                with gr.Column():
+                    searcher = gr.Chatbot(label='searcher',
+                                          height=700,
+                                          show_label=True,
+                                          show_copy_button=True,
+                                          bubble_full_width=False,
+                                          render_markdown=True)
+            with gr.Row():
+                user_input = gr.Textbox(show_label=False,
+                                        placeholder='帮我搜索一下 InternLM 开源体系',
+                                        lines=5,
+                                        container=False)
+            with gr.Row():
+                with gr.Column(scale=2):
+                    submitBtn = gr.Button('Submit')
+                with gr.Column(scale=1, min_width=20):
+                    emptyBtn = gr.Button('Clear History')
+
+    def user(query, history):
+        return '', history + [[query, '']]
+
+    submitBtn.click(user, [user_input, planner], [user_input, planner],
+                    queue=False).then(predict, [planner, searcher],
+                                      [planner, searcher])
+    emptyBtn.click(rst_mem, [planner, searcher], [planner, searcher],
+                   queue=False)
+
+demo.queue()
+demo.launch(server_name='0.0.0.0',
+            server_port=7860,
+            inbrowser=True,
+            share=True)