wehub-resource-sync
142 lines
5.0 KiB
Python
142 lines
5.0 KiB
Python
# !/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
# @author: CS_木成河
|
|
# @time: 2024/11/20 16:32
|
|
# @blog: https://blog.csdn.net/weixin_47936614
|
|
|
|
import os
|
|
import shutil
|
|
|
|
from fastapi import FastAPI, Request, File, UploadFile
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.responses import HTMLResponse
|
|
from fastapi.templating import Jinja2Templates
|
|
from langchain.chains.retrieval_qa.base import RetrievalQA
|
|
from langchain_core.prompts import PromptTemplate
|
|
from langchain_openai import ChatOpenAI
|
|
from pymilvus import Collection, connections
|
|
from pymilvus.orm import utility
|
|
|
|
from milvus_vector import vector_store, config
|
|
from file_process import RagFileProcessor
|
|
from protocol.prompts import prompt_template
|
|
from protocol.mode import ChatRequest
|
|
|
|
# 初始化 FastAPI 应用
|
|
app = FastAPI(title="Knowledge_QA_RAG API", description="API for data process and retrieval using Milvus and LangChain.")
|
|
|
|
# 挂载静态文件目录
|
|
app.mount("/static", StaticFiles(directory="static"), name="static")
|
|
|
|
# 配置模板目录
|
|
templates = Jinja2Templates(directory="templates")
|
|
|
|
|
|
# 渲染主页
|
|
@app.get("/", response_class=HTMLResponse)
|
|
async def read_home(request: Request):
|
|
return templates.TemplateResponse("qa.html", {"request": request})
|
|
|
|
|
|
@app.post("/rag/chat/")
|
|
async def chat(request: ChatRequest):
|
|
"""
|
|
根据用户问题，从向量库检索并返回回答。
|
|
"""
|
|
print(f"Q: {request.question}")
|
|
try:
|
|
# 初始化 OpenAI Chat 模型
|
|
llm = ChatOpenAI(model=config.llm_model_name, api_key=config.api_key, base_url=config.base_url)
|
|
# 定义 Prompt 模板
|
|
qa_prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
|
|
|
|
# 定义搜索参数
|
|
search_kwargs = {"score_threshold": 0.3, "k": 5}
|
|
retriever = vector_store.as_retriever(search_type="similarity_score_threshold", search_kwargs=search_kwargs)
|
|
qa_chain = RetrievalQA.from_chain_type(
|
|
llm=llm,
|
|
chain_type="stuff",
|
|
retriever=retriever,
|
|
chain_type_kwargs={"prompt": qa_prompt},
|
|
return_source_documents=True
|
|
)
|
|
result = qa_chain.invoke({"query": request.question})
|
|
answer_result = result.get("result", "")
|
|
print(f"A: {answer_result}")
|
|
|
|
source = {"source_documents": [{"content": doc.page_content, "metadata": doc.metadata} for doc in
|
|
result.get("source_documents", [])]}
|
|
print(f"source: {source}")
|
|
|
|
return {"answer": answer_result}
|
|
except Exception as e:
|
|
return {"status": "error", "message": str(e)}
|
|
|
|
|
|
@app.post("/rag/clear/")
|
|
async def clear_knowledge(collection_name: str = config.milvus_collection_name,
|
|
host: str = config.milvus_host,
|
|
port: int = config.milvus_port):
|
|
"""
|
|
清空 Milvus 知识库集合，并删除指定目录中的文件
|
|
"""
|
|
folder = "./upload_files"
|
|
try:
|
|
connections.connect("default", host=host, port=port)
|
|
if utility.has_collection(collection_name):
|
|
collection = Collection(name=collection_name)
|
|
collection.drop()
|
|
print(f"Collection '{collection_name}' 成功删除.")
|
|
else:
|
|
print(f"Collection '{collection_name}' 不存在.")
|
|
connections.disconnect("default")
|
|
|
|
for filename in os.listdir(folder):
|
|
file_path = os.path.join(folder, filename)
|
|
if os.path.isfile(file_path) or os.path.islink(file_path):
|
|
os.unlink(file_path)
|
|
elif os.path.isdir(file_path):
|
|
shutil.rmtree(file_path)
|
|
|
|
return {"message": f"知识库清空: Collection '{collection_name}' 删除, 文件夹 '{folder}' 清空."}
|
|
except Exception as e:
|
|
return {"error": f"知识库清空失败. 原因: {str(e)}"}
|
|
|
|
|
|
@app.post("/rag/create/")
|
|
async def create_knowledge(file: UploadFile = File(...)):
|
|
"""
|
|
上传文件到指定目录后，处理文件内容并添加到向量库。
|
|
"""
|
|
folder = './upload_files' # 文件存储目录
|
|
os.makedirs(folder, exist_ok=True) # 确保目录存在
|
|
file_path = os.path.join(folder, file.filename)
|
|
|
|
try:
|
|
# 保存文件
|
|
with open(file_path, "wb") as f:
|
|
content = await file.read()
|
|
f.write(content)
|
|
|
|
folder: str = './upload_files'
|
|
file_path = os.path.join(folder, file.filename)
|
|
|
|
# 初始化文件处理器
|
|
file_processor = RagFileProcessor(chunk_size=64)
|
|
# 处理文件内容并插入到向量库
|
|
text_datas = file_processor.get_data(file_path=file_path)
|
|
|
|
# 连接到 Milvus
|
|
vector_store.add_texts(**text_datas)
|
|
return {
|
|
"status": "success",
|
|
"message": f"文件 '{file.filename}' 上传成功并添加至向量数据库.",
|
|
}
|
|
except Exception as e:
|
|
return {
|
|
"status": "error",
|
|
"message": f"文件 '{file.filename}'处理失败. 原因: {str(e)}",
|
|
}
|
|
|
|
|