能够正确调用rag,rag执行完成后,无法返回内容

This commit is contained in:
yangdx
2025-01-15 18:47:01 +08:00
parent 4e5517a602
commit 828af49d6b

View File

@@ -659,38 +659,55 @@ def create_app(args):
cleaned_query, mode = parse_query_mode(query) cleaned_query, mode = parse_query_mode(query)
# 调用RAG进行查询 # 调用RAG进行查询
response = await rag.aquery(
cleaned_query,
param=QueryParam(
mode=mode,
stream=request.stream,
)
)
if request.stream: if request.stream:
response = await rag.aquery(
cleaned_query,
param=QueryParam(
mode=mode,
stream=True,
only_need_context=False
),
)
async def stream_generator(): async def stream_generator():
async for chunk in response: try:
yield OllamaChatResponse( async for chunk in response:
model=LIGHTRAG_MODEL, yield {
created_at=LIGHTRAG_CREATED_AT, "model": LIGHTRAG_MODEL,
message=OllamaMessage( "created_at": LIGHTRAG_CREATED_AT,
role="assistant", "message": {
content=chunk "role": "assistant",
), "content": chunk
done=False },
) "done": False
# 发送一个空的完成消息 }
yield OllamaChatResponse( yield {
model=LIGHTRAG_MODEL, "model": LIGHTRAG_MODEL,
created_at=LIGHTRAG_CREATED_AT, "created_at": LIGHTRAG_CREATED_AT,
message=OllamaMessage( "message": {
role="assistant", "role": "assistant",
content="" "content": ""
), },
done=True "done": True
) }
return stream_generator() except Exception as e:
logging.error(f"Error in stream_generator: {str(e)}")
raise
from fastapi.responses import StreamingResponse
import json
return StreamingResponse(
(f"data: {json.dumps(chunk)}\n\n" async for chunk in stream_generator()),
media_type="text/event-stream"
)
else: else:
response = await rag.aquery(
cleaned_query,
param=QueryParam(
mode=mode,
stream=False,
only_need_context=False
),
)
return OllamaChatResponse( return OllamaChatResponse(
model=LIGHTRAG_MODEL, model=LIGHTRAG_MODEL,
created_at=LIGHTRAG_CREATED_AT, created_at=LIGHTRAG_CREATED_AT,