109 lines
4.0 KiB
Python
109 lines
4.0 KiB
Python
import asyncio
|
|
import json
|
|
from http.client import HTTPException
|
|
from typing import List
|
|
|
|
from fastapi import UploadFile, File, Form
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
from starlette.requests import Request
|
|
from starlette.responses import StreamingResponse
|
|
|
|
from app.utils.db_utils import AsyncSessionDep
|
|
from app.utils.knowledge_utils import knowledge_service
|
|
from app.utils.llm_utils import create_llm
|
|
from app.utils.milvus_utils import milvus_service, KnowledgeQueryParam
|
|
|
|
|
|
def add_knowledge_route(app):
|
|
@app.post("/knowledge/search")
|
|
async def knowledge_search(param: KnowledgeQueryParam):
|
|
return await milvus_service.async_search(param)
|
|
|
|
@app.post("/knowledge/recall/stream")
|
|
async def knowledge_search(body: dict):
|
|
|
|
param = KnowledgeQueryParam(**body.get('input'))
|
|
|
|
async def generator_function():
|
|
search_response = await milvus_service.async_search(param)
|
|
|
|
chain = create_llm() | StrOutputParser()
|
|
|
|
# 先把检索结果返回前端
|
|
yield f'data: {json.dumps({"type": "retrieve", "data": json.dumps(search_response.model_dump())}, ensure_ascii=False)}\n\n'
|
|
|
|
external_prompt = body.get('input', {}).get('prompt', '')
|
|
|
|
# print("external_prompt", external_prompt)
|
|
|
|
async for chunk in chain.astream([
|
|
SystemMessage(content=f"""你需要根据如下内容来回答用户的问题：{search_response.answer}，如果内容为'在提供的资料中未找到相关信息'，你需要根据你自己的知识来回答用户问题。""" + external_prompt),
|
|
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")
|
|
|
|
# @app.post("/knowledge/embed_text")
|
|
# async def knowledge_embed_text(text_list: List[str]):
|
|
# id_list = await next_id(len(text_list))
|
|
# document_list = [
|
|
# Document(
|
|
# text=text,
|
|
# id=id_list[index],
|
|
# metadata={"kb_id": "cde"}
|
|
# )
|
|
# for index, text in enumerate(text_list)
|
|
# ]
|
|
# await milvus_service.async_create_index_from_documents(document_list)
|
|
# search_result = await milvus_service.async_search("hello")
|
|
# return {
|
|
# "result": "嵌入成功:",
|
|
# "origin_documents": [document.to_dict() for document in document_list],
|
|
# "search_documents": search_result,
|
|
# }
|
|
#
|
|
@app.post("/knowledge/delete")
|
|
async def knowledge_embed_text(body: dict):
|
|
await milvus_service.async_delete(body.get('id'))
|
|
return {
|
|
"result": "删除成功",
|
|
}
|
|
|
|
@app.post("/knowledge/upload_files")
|
|
async def knowledge_search(
|
|
session: AsyncSessionDep,
|
|
request: Request,
|
|
files: List[UploadFile] = File(...),
|
|
kb_code: str = Form(..., description="所属知识库的编码"),
|
|
):
|
|
"""
|
|
批量上传文档并嵌入到向量数据库（同步响应版本）
|
|
"""
|
|
if not files:
|
|
raise HTTPException(status_code=400, detail="No files provided")
|
|
|
|
# 限制文件数量
|
|
if len(files) > 20:
|
|
raise HTTPException(status_code=400, detail="Maximum 20 files allowed per upload")
|
|
|
|
# 先将所有文件直接保存到本地
|
|
task_list = [asyncio.create_task(knowledge_service.save_file_with_new_session(file=file)) for file in files]
|
|
task_result_list = await asyncio.gather(*task_list)
|
|
file_dict_list = [item["result"] for item in task_result_list]
|
|
|
|
# 插入对应的文档对象记录
|
|
doc_cls_list = await knowledge_service.save_knowledge_doc_list(
|
|
session=session,
|
|
file_dict_list=file_dict_list,
|
|
kb_code=kb_code,
|
|
user=request.state.user,
|
|
)
|
|
|
|
# 异步处理嵌入文档，不再等待
|
|
[asyncio.create_task(knowledge_service.process_doc_cls(doc_cls=doc_cls)) for doc_cls in doc_cls_list]
|
|
|
|
# 直接返回文档对象
|
|
return {"message": f"正在处理 {len(doc_cls_list)} 个文档，请刷新列表查看文档状态。", "result": [item.model_dump() for item in doc_cls_list]}
|