diff --git a/app/controller/add_knowledge_route.py b/app/controller/add_knowledge_route.py index e2c6c32..1c15cd8 100644 --- a/app/controller/add_knowledge_route.py +++ b/app/controller/add_knowledge_route.py @@ -4,10 +4,8 @@ from http.client import HTTPException from typing import List from fastapi import UploadFile, File, Form -from langchain_core.messages import HumanMessage +from langchain_core.messages import HumanMessage, SystemMessage from langchain_core.output_parsers import StrOutputParser -from langchain_core.prompts import ChatPromptTemplate -from langgraph.prebuilt import create_react_agent from starlette.requests import Request from starlette.responses import StreamingResponse @@ -26,20 +24,19 @@ def add_knowledge_route(app): async def knowledge_search(body: dict): param = KnowledgeQueryParam(**body.get('input')) - search_response = await milvus_service.async_search(param) - - chain = ChatPromptTemplate.from_template( - f"""你需要根据如下内容来回答用户的问题:{search_response.answer}""" - ) | create_llm() | StrOutputParser() async def generator_function(): + search_response = await milvus_service.async_search(param) + + chain = create_llm() | StrOutputParser() + # 先把检索结果返回前端 yield f'data: {json.dumps({"type": "retrieve", "data": search_response.model_dump()}, ensure_ascii=False)}\n\n' - async for chunk in chain.astream( - {"messages": [HumanMessage(content=param.question)]}, - stream_mode="messages" - ): + async for chunk in chain.astream([ + SystemMessage(content=f"""你需要根据如下内容来回答用户的问题:{search_response.answer}"""), + HumanMessage(content=param.question) + ]): yield f'data: {json.dumps({"type": "messages", "data": chunk}, ensure_ascii=False)}\n\n' return StreamingResponse(generator_function(), media_type="text/event-stream")