Spaces:
Runtime error
Runtime error
File size: 4,690 Bytes
f3f614f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 |
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
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')
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()
# 使用 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:
# 确保在发生异常时队列中的所有元素都被消费
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: # 确保消费完所有元素
break
yield response
if not isinstance(
response,
tuple) and response.state == AgentStatusCode.END:
break
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:
queue.close()
await queue.wait_closed()
inputs = request.inputs
agent = init_agent(lang=args.lang, model_format=args.model_format)
return EventSourceResponse(generate())
if __name__ == '__main__':
import uvicorn
uvicorn.run(app, host='0.0.0.0', port=8002, log_level='info')
|