main
HuangHai 3 weeks ago
parent ccafd3c5cd
commit b2ee229e15

@ -45,7 +45,10 @@ async def lifespan(app: FastAPI):
warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host') warnings.filterwarnings('ignore', message='Unverified HTTPS request is being made to host')
yield yield
async def print_stream(stream):
async for chunk in stream:
if chunk:
print(chunk, end="", flush=True)
app = FastAPI(lifespan=lifespan) app = FastAPI(lifespan=lifespan)
# 挂载静态文件目录 # 挂载静态文件目录
@ -57,34 +60,28 @@ async def rag(request: fastapi.Request):
query = data.get('query', '') query = data.get('query', '')
working_dir = "./Test/Math" # 使用与T2_Query.py相同的目录 working_dir = "./Test/Math" # 使用与T2_Query.py相同的目录
async def generate_response_stream(): async def generate_response_stream(query: str):
rag = None
try: try:
logger.info(f"Initializing RAG with working_dir: {working_dir}") rag = await initialize_rag(working_dir="./Test/Math")
rag = await initialize_rag(working_dir) resp = await rag.aquery(query=query, stream=True)
logger.info("RAG initialized successfully")
resp = await rag.aquery(
query,
param=QueryParam(mode="hybrid", stream=True),
)
logger.info("Query execution started")
async for chunk in resp: async for chunk in resp:
if not chunk: # Skip empty chunks if not chunk:
continue continue
try:
yield f"data: {json.dumps({'reply': chunk}, ensure_ascii=False)}\n\n" # 正确的SSE格式应该是单层data:前缀
except Exception as e: yield f"data: {json.dumps({'reply': chunk})}\n\n"
logger.error(f"Error processing chunk: {e}")
yield f"data: {json.dumps({'error': str(e)})}\n\n" # 打印到控制台用于调试
print(chunk, end='', flush=True)
except Exception as e: except Exception as e:
yield f"data: {json.dumps({'error': str(e)})}\n\n" yield f"data: {json.dumps({'error': str(e)})}\n\n"
finally: finally:
if rag: if 'rag' in locals():
await rag.finalize_storages() await rag.finalize()
return EventSourceResponse(generate_response_stream()) return EventSourceResponse(generate_response_stream(query=query))
if __name__ == "__main__": if __name__ == "__main__":

Loading…
Cancel
Save