117 lines
2.6 KiB
Python
117 lines
2.6 KiB
Python
"""
|
|
RAG 检索 API 接口
|
|
|
|
提供向量检索功能
|
|
"""
|
|
from typing import Optional
|
|
|
|
from fastapi import APIRouter, HTTPException, Query
|
|
from pydantic import BaseModel
|
|
|
|
from app.services.rag_service import rag_service
|
|
|
|
router = APIRouter(prefix="/rag", tags=["RAG检索"])
|
|
|
|
|
|
class SearchRequest(BaseModel):
|
|
query: str
|
|
top_k: int = 5
|
|
|
|
|
|
class SearchResult(BaseModel):
|
|
content: str
|
|
metadata: dict
|
|
score: float
|
|
doc_id: str
|
|
|
|
|
|
@router.post("/search")
|
|
async def search_rag(
|
|
request: SearchRequest
|
|
):
|
|
"""
|
|
RAG 语义检索
|
|
|
|
根据查询文本检索相关的文档片段或字段
|
|
|
|
Args:
|
|
request.query: 查询文本
|
|
request.top_k: 返回数量
|
|
|
|
Returns:
|
|
相关文档列表
|
|
"""
|
|
try:
|
|
results = rag_service.retrieve(
|
|
query=request.query,
|
|
top_k=request.top_k
|
|
)
|
|
|
|
return {
|
|
"success": True,
|
|
"results": results
|
|
}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"检索失败: {str(e)}")
|
|
|
|
|
|
@router.get("/status")
|
|
async def get_rag_status():
|
|
"""
|
|
获取 RAG 索引状态
|
|
|
|
Returns:
|
|
RAG 索引统计信息
|
|
"""
|
|
try:
|
|
count = rag_service.get_vector_count()
|
|
|
|
return {
|
|
"success": True,
|
|
"vector_count": count,
|
|
"collections": ["document_fields", "document_content"] # 预留
|
|
}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}")
|
|
|
|
|
|
@router.post("/rebuild")
|
|
async def rebuild_rag_index():
|
|
"""
|
|
重建 RAG 索引
|
|
|
|
从 MongoDB 中读取所有文档，重新构建向量索引
|
|
"""
|
|
from app.core.database import mongodb
|
|
|
|
try:
|
|
# 清空现有索引
|
|
rag_service.clear()
|
|
|
|
# 从 MongoDB 读取所有文档
|
|
cursor = mongodb.documents.find({})
|
|
count = 0
|
|
|
|
async for doc in cursor:
|
|
content = doc.get("content", "")
|
|
if content:
|
|
rag_service.index_document_content(
|
|
doc_id=str(doc["_id"]),
|
|
content=content[:5000],
|
|
metadata={
|
|
"filename": doc.get("metadata", {}).get("filename"),
|
|
"doc_type": doc.get("doc_type")
|
|
}
|
|
)
|
|
count += 1
|
|
|
|
return {
|
|
"success": True,
|
|
"message": f"已重建索引，共处理 {count} 个文档"
|
|
}
|
|
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"重建索引失败: {str(e)}")
|