108 lines
3.7 KiB
Python
108 lines
3.7 KiB
Python
import asyncio
|
||
import json
|
||
import logging
|
||
|
||
import fastapi
|
||
import uvicorn
|
||
from fastapi import FastAPI
|
||
from lightrag import LightRAG
|
||
from lightrag.kg.shared_storage import initialize_pipeline_status
|
||
from raganything import RAGAnything
|
||
from sse_starlette import EventSourceResponse
|
||
from starlette.staticfiles import StaticFiles
|
||
|
||
from Util.RagUtil import create_llm_model_func, create_vision_model_func, create_embedding_func
|
||
|
||
# 在程序开始时添加以下配置
|
||
logging.basicConfig(
|
||
level=logging.INFO, # 设置日志级别为INFO
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
|
||
# 或者如果你想更详细地控制日志输出
|
||
logger = logging.getLogger('lightrag')
|
||
logger.setLevel(logging.INFO)
|
||
handler = logging.StreamHandler()
|
||
handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
|
||
logger.addHandler(handler)
|
||
|
||
|
||
async def lifespan(app: FastAPI):
|
||
yield
|
||
|
||
|
||
async def print_stream(stream):
|
||
async for chunk in stream:
|
||
if chunk:
|
||
print(chunk, end="", flush=True)
|
||
|
||
|
||
app = FastAPI(lifespan=lifespan)
|
||
|
||
# 挂载静态文件目录
|
||
app.mount("/static", StaticFiles(directory="Static"), name="static")
|
||
|
||
|
||
@app.post("/api/rag")
|
||
async def rag(request: fastapi.Request):
|
||
data = await request.json()
|
||
topic = data.get("topic") # Chinese, Math
|
||
# 拼接路径
|
||
WORKING_PATH= "./Topic/" + topic
|
||
# 查询的问题
|
||
query = data.get("query")
|
||
# 关闭参考资料
|
||
user_prompt="\n 1、不要输出参考资料 或者 References ！"
|
||
user_prompt = user_prompt + "\n 2、如果问题与提供的知识库内容不符，则明确告诉未在知识库范围内提到！"
|
||
|
||
async def generate_response_stream(query: str):
|
||
try:
|
||
# 初始化RAG组件
|
||
llm_model_func = create_llm_model_func()
|
||
vision_model_func = create_vision_model_func(llm_model_func)
|
||
embedding_func = create_embedding_func()
|
||
|
||
lightrag_instance = LightRAG(
|
||
working_dir=WORKING_PATH,
|
||
llm_model_func=llm_model_func,
|
||
embedding_func=embedding_func
|
||
)
|
||
|
||
await lightrag_instance.initialize_storages()
|
||
await initialize_pipeline_status()
|
||
|
||
# 创建RAG实例并保存到app.state
|
||
app.state.rag = RAGAnything(
|
||
lightrag=lightrag_instance,
|
||
vision_model_func=vision_model_func,
|
||
)
|
||
# 直接使用app.state中已初始化的rag实例
|
||
# 修改为直接获取完整响应
|
||
resp = await app.state.rag.aquery(
|
||
query=query,
|
||
mode="hybrid",
|
||
stream=False # 改为False获取完整响应
|
||
)
|
||
|
||
# 将完整响应作为单个块返回
|
||
#yield f"data: {json.dumps({'reply': resp})}\n\n"
|
||
#print(resp, end='', flush=True)
|
||
# 添加逐字输出效果
|
||
for i in range(0, len(resp), 5): # 每次输出5个字符
|
||
chunk = resp[i:i + 5]
|
||
yield f"data: {json.dumps({'reply': chunk})}\n\n"
|
||
await asyncio.sleep(0.1) # 控制输出速度
|
||
|
||
except Exception as e:
|
||
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
||
logger.error(f"处理查询时出错: {query}. 错误: {str(e)}")
|
||
finally:
|
||
# 清理资源
|
||
await app.state.rag.lightrag.finalize_storages()
|
||
|
||
return EventSourceResponse(generate_response_stream(query=query))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
uvicorn.run(app, host="0.0.0.0", port=8000)
|