576 lines
17 KiB
Python
576 lines
17 KiB
Python
import hashlib
|
||
from typing import List, Any, Dict
|
||
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, BackgroundTasks, Query
|
||
from sqlalchemy.orm import Session
|
||
from langchain_chroma import Chroma
|
||
from sqlalchemy import text
|
||
import logging
|
||
from datetime import datetime, timedelta
|
||
from pydantic import BaseModel
|
||
from sqlalchemy.orm import selectinload
|
||
import time
|
||
import asyncio
|
||
|
||
from app.db.session import get_db
|
||
from app.models.user import User
|
||
from app.core.security import get_current_user
|
||
from app.models.knowledge import KnowledgeBase, Document, ProcessingTask, DocumentChunk, DocumentUpload
|
||
from app.schemas.knowledge import (
|
||
KnowledgeBaseCreate,
|
||
KnowledgeBaseResponse,
|
||
KnowledgeBaseUpdate,
|
||
DocumentResponse,
|
||
PreviewRequest
|
||
)
|
||
from app.services.document_processor import process_document_background, upload_document, preview_document, PreviewResult
|
||
from app.core.config import settings
|
||
from app.core.minio import get_minio_client
|
||
from minio.error import MinioException
|
||
from app.services.vector_store import VectorStoreFactory
|
||
from app.services.embedding.embedding_factory import EmbeddingsFactory
|
||
|
||
router = APIRouter()
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
class TestRetrievalRequest(BaseModel):
|
||
query: str
|
||
kb_id: int
|
||
top_k: int
|
||
|
||
@router.post("", response_model=KnowledgeBaseResponse)
|
||
def create_knowledge_base(
|
||
*,
|
||
db: Session = Depends(get_db),
|
||
kb_in: KnowledgeBaseCreate,
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Any:
|
||
"""
|
||
Create new knowledge base.
|
||
"""
|
||
kb = KnowledgeBase(
|
||
name=kb_in.name,
|
||
description=kb_in.description,
|
||
user_id=current_user.id
|
||
)
|
||
db.add(kb)
|
||
db.commit()
|
||
db.refresh(kb)
|
||
logger.info(f"Knowledge base created: {kb.name} for user {current_user.id}")
|
||
return kb
|
||
|
||
@router.get("", response_model=List[KnowledgeBaseResponse])
|
||
def get_knowledge_bases(
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user),
|
||
skip: int = 0,
|
||
limit: int = 100
|
||
) -> Any:
|
||
"""
|
||
Retrieve knowledge bases.
|
||
"""
|
||
knowledge_bases = (
|
||
db.query(KnowledgeBase)
|
||
.filter(KnowledgeBase.user_id == current_user.id)
|
||
.offset(skip)
|
||
.limit(limit)
|
||
.all()
|
||
)
|
||
return knowledge_bases
|
||
|
||
@router.get("/{kb_id}", response_model=KnowledgeBaseResponse)
|
||
def get_knowledge_base(
|
||
*,
|
||
db: Session = Depends(get_db),
|
||
kb_id: int,
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Any:
|
||
"""
|
||
Get knowledge base by ID.
|
||
"""
|
||
from sqlalchemy.orm import joinedload
|
||
|
||
kb = (
|
||
db.query(KnowledgeBase)
|
||
.options(
|
||
joinedload(KnowledgeBase.documents)
|
||
.joinedload(Document.processing_tasks)
|
||
)
|
||
.filter(
|
||
KnowledgeBase.id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
)
|
||
.first()
|
||
)
|
||
|
||
if not kb:
|
||
raise HTTPException(status_code=404, detail="Knowledge base not found")
|
||
|
||
return kb
|
||
|
||
@router.put("/{kb_id}", response_model=KnowledgeBaseResponse)
|
||
def update_knowledge_base(
|
||
*,
|
||
db: Session = Depends(get_db),
|
||
kb_id: int,
|
||
kb_in: KnowledgeBaseUpdate,
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Any:
|
||
"""
|
||
Update knowledge base.
|
||
"""
|
||
kb = db.query(KnowledgeBase).filter(
|
||
KnowledgeBase.id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
|
||
if not kb:
|
||
raise HTTPException(status_code=404, detail="Knowledge base not found")
|
||
|
||
for field, value in kb_in.dict(exclude_unset=True).items():
|
||
setattr(kb, field, value)
|
||
|
||
db.add(kb)
|
||
db.commit()
|
||
db.refresh(kb)
|
||
logger.info(f"Knowledge base updated: {kb.name} for user {current_user.id}")
|
||
return kb
|
||
|
||
@router.delete("/{kb_id}")
|
||
async def delete_knowledge_base(
|
||
*,
|
||
db: Session = Depends(get_db),
|
||
kb_id: int,
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Any:
|
||
"""
|
||
Delete knowledge base and all associated resources.
|
||
"""
|
||
logger = logging.getLogger(__name__)
|
||
|
||
kb = (
|
||
db.query(KnowledgeBase)
|
||
.filter(
|
||
KnowledgeBase.id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
)
|
||
.first()
|
||
)
|
||
if not kb:
|
||
raise HTTPException(status_code=404, detail="Knowledge base not found")
|
||
|
||
try:
|
||
# Get all document file paths before deletion
|
||
document_paths = [doc.file_path for doc in kb.documents]
|
||
|
||
# Initialize services
|
||
minio_client = get_minio_client()
|
||
embeddings = EmbeddingsFactory.create()
|
||
|
||
vector_store = VectorStoreFactory.create(
|
||
store_type=settings.VECTOR_STORE_TYPE,
|
||
collection_name=f"kb_{kb_id}",
|
||
embedding_function=embeddings,
|
||
)
|
||
|
||
# Clean up external resources first
|
||
cleanup_errors = []
|
||
|
||
# 1. Clean up MinIO files
|
||
try:
|
||
# Delete all objects with prefix kb_{kb_id}/
|
||
objects = minio_client.list_objects(settings.MINIO_BUCKET_NAME, prefix=f"kb_{kb_id}/")
|
||
for obj in objects:
|
||
minio_client.remove_object(settings.MINIO_BUCKET_NAME, obj.object_name)
|
||
logger.info(f"Cleaned up MinIO files for knowledge base {kb_id}")
|
||
except MinioException as e:
|
||
cleanup_errors.append(f"Failed to clean up MinIO files: {str(e)}")
|
||
logger.error(f"MinIO cleanup error for kb {kb_id}: {str(e)}")
|
||
|
||
# 2. Clean up vector store
|
||
try:
|
||
vector_store._store.delete_collection(f"kb_{kb_id}")
|
||
logger.info(f"Cleaned up vector store for knowledge base {kb_id}")
|
||
except Exception as e:
|
||
cleanup_errors.append(f"Failed to clean up vector store: {str(e)}")
|
||
logger.error(f"Vector store cleanup error for kb {kb_id}: {str(e)}")
|
||
|
||
# Finally, delete database records in a single transaction
|
||
db.delete(kb)
|
||
db.commit()
|
||
|
||
# Report any cleanup errors in the response
|
||
if cleanup_errors:
|
||
return {
|
||
"message": "Knowledge base deleted with cleanup warnings",
|
||
"warnings": cleanup_errors
|
||
}
|
||
|
||
return {"message": "Knowledge base and all associated resources deleted successfully"}
|
||
except Exception as e:
|
||
db.rollback()
|
||
logger.error(f"Failed to delete knowledge base {kb_id}: {str(e)}")
|
||
raise HTTPException(status_code=500, detail=f"Failed to delete knowledge base: {str(e)}")
|
||
|
||
# Batch upload documents
|
||
@router.post("/{kb_id}/documents/upload")
|
||
async def upload_kb_documents(
|
||
kb_id: int,
|
||
files: List[UploadFile],
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user)
|
||
):
|
||
"""
|
||
Upload multiple documents to MinIO.
|
||
"""
|
||
kb = db.query(KnowledgeBase).filter(
|
||
KnowledgeBase.id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
if not kb:
|
||
raise HTTPException(status_code=404, detail="Knowledge base not found")
|
||
|
||
results = []
|
||
for file in files:
|
||
# 1. 计算文件 hash
|
||
file_content = await file.read()
|
||
file_hash = hashlib.sha256(file_content).hexdigest()
|
||
|
||
# 2. 检查是否存在完全相同的文件（名称和hash都相同）
|
||
existing_document = db.query(Document).filter(
|
||
Document.file_name == file.filename,
|
||
Document.file_hash == file_hash,
|
||
Document.knowledge_base_id == kb_id
|
||
).first()
|
||
|
||
if existing_document:
|
||
# 完全相同的文件，直接返回
|
||
results.append({
|
||
"document_id": existing_document.id,
|
||
"file_name": existing_document.file_name,
|
||
"status": "exists",
|
||
"message": "文件已存在且已处理完成",
|
||
"skip_processing": True
|
||
})
|
||
continue
|
||
|
||
# 3. 上传到临时目录
|
||
temp_path = f"kb_{kb_id}/temp/{file.filename}"
|
||
await file.seek(0)
|
||
try:
|
||
minio_client = get_minio_client()
|
||
file_size = len(file_content) # 使用之前读取的文件内容长度
|
||
minio_client.put_object(
|
||
bucket_name=settings.MINIO_BUCKET_NAME,
|
||
object_name=temp_path,
|
||
data=file.file,
|
||
length=file_size, # 指定文件大小
|
||
content_type=file.content_type
|
||
)
|
||
except MinioException as e:
|
||
logger.error(f"Failed to upload file to MinIO: {str(e)}")
|
||
raise HTTPException(status_code=500, detail="Failed to upload file")
|
||
|
||
# 4. 创建上传记录
|
||
upload = DocumentUpload(
|
||
knowledge_base_id=kb_id,
|
||
file_name=file.filename,
|
||
file_hash=file_hash,
|
||
file_size=len(file_content),
|
||
content_type=file.content_type,
|
||
temp_path=temp_path
|
||
)
|
||
db.add(upload)
|
||
db.commit()
|
||
db.refresh(upload)
|
||
|
||
results.append({
|
||
"upload_id": upload.id,
|
||
"file_name": file.filename,
|
||
"temp_path": temp_path,
|
||
"status": "pending",
|
||
"skip_processing": False
|
||
})
|
||
|
||
return results
|
||
|
||
@router.post("/{kb_id}/documents/preview")
|
||
async def preview_kb_documents(
|
||
kb_id: int,
|
||
preview_request: PreviewRequest,
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Dict[int, PreviewResult]:
|
||
"""
|
||
Preview multiple documents' chunks.
|
||
"""
|
||
results = {}
|
||
for doc_id in preview_request.document_ids:
|
||
document = db.query(Document).join(KnowledgeBase).filter(
|
||
Document.id == doc_id,
|
||
Document.knowledge_base_id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
|
||
if document:
|
||
file_path = document.file_path
|
||
else:
|
||
upload = db.query(DocumentUpload).join(KnowledgeBase).filter(
|
||
DocumentUpload.id == doc_id,
|
||
DocumentUpload.knowledge_base_id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
|
||
if not upload:
|
||
raise HTTPException(status_code=404, detail=f"Document {doc_id} not found")
|
||
|
||
file_path = upload.temp_path
|
||
|
||
preview = await preview_document(
|
||
file_path,
|
||
chunk_size=preview_request.chunk_size,
|
||
chunk_overlap=preview_request.chunk_overlap
|
||
)
|
||
results[doc_id] = preview
|
||
|
||
return results
|
||
|
||
@router.post("/{kb_id}/documents/process")
|
||
async def process_kb_documents(
|
||
kb_id: int,
|
||
upload_results: List[dict],
|
||
background_tasks: BackgroundTasks,
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user)
|
||
):
|
||
"""
|
||
Process multiple documents asynchronously.
|
||
"""
|
||
start_time = time.time()
|
||
|
||
kb = db.query(KnowledgeBase).filter(
|
||
KnowledgeBase.id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
|
||
if not kb:
|
||
raise HTTPException(status_code=404, detail="Knowledge base not found")
|
||
|
||
task_info = []
|
||
upload_ids = []
|
||
|
||
for result in upload_results:
|
||
if result.get("skip_processing"):
|
||
continue
|
||
upload_ids.append(result["upload_id"])
|
||
|
||
if not upload_ids:
|
||
return {"tasks": []}
|
||
|
||
uploads = db.query(DocumentUpload).filter(DocumentUpload.id.in_(upload_ids)).all()
|
||
uploads_dict = {upload.id: upload for upload in uploads}
|
||
|
||
all_tasks = []
|
||
for upload_id in upload_ids:
|
||
upload = uploads_dict.get(upload_id)
|
||
if not upload:
|
||
continue
|
||
|
||
task = ProcessingTask(
|
||
document_upload_id=upload_id,
|
||
knowledge_base_id=kb_id,
|
||
status="pending"
|
||
)
|
||
all_tasks.append(task)
|
||
|
||
db.add_all(all_tasks)
|
||
db.commit()
|
||
|
||
for task in all_tasks:
|
||
db.refresh(task)
|
||
|
||
task_data = []
|
||
for i, upload_id in enumerate(upload_ids):
|
||
if i < len(all_tasks):
|
||
task = all_tasks[i]
|
||
upload = uploads_dict.get(upload_id)
|
||
|
||
task_info.append({
|
||
"upload_id": upload_id,
|
||
"task_id": task.id
|
||
})
|
||
|
||
if upload:
|
||
task_data.append({
|
||
"task_id": task.id,
|
||
"upload_id": upload_id,
|
||
"temp_path": upload.temp_path,
|
||
"file_name": upload.file_name
|
||
})
|
||
|
||
background_tasks.add_task(
|
||
add_processing_tasks_to_queue,
|
||
task_data,
|
||
kb_id
|
||
)
|
||
|
||
return {"tasks": task_info}
|
||
|
||
async def add_processing_tasks_to_queue(task_data, kb_id):
|
||
"""Helper function to add document processing tasks to the queue without blocking the main response."""
|
||
for data in task_data:
|
||
asyncio.create_task(
|
||
process_document_background(
|
||
data["temp_path"],
|
||
data["file_name"],
|
||
kb_id,
|
||
data["task_id"],
|
||
None
|
||
)
|
||
)
|
||
logger.info(f"Added {len(task_data)} document processing tasks to queue")
|
||
|
||
@router.post("/cleanup")
|
||
async def cleanup_temp_files(
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user)
|
||
):
|
||
"""
|
||
Clean up expired temporary files.
|
||
"""
|
||
expired_time = datetime.utcnow() - timedelta(hours=24)
|
||
expired_uploads = db.query(DocumentUpload).filter(
|
||
DocumentUpload.created_at < expired_time
|
||
).all()
|
||
|
||
minio_client = get_minio_client()
|
||
for upload in expired_uploads:
|
||
try:
|
||
minio_client.remove_object(
|
||
bucket_name=settings.MINIO_BUCKET_NAME,
|
||
object_name=upload.temp_path
|
||
)
|
||
except MinioException as e:
|
||
logger.error(f"Failed to delete temp file {upload.temp_path}: {str(e)}")
|
||
|
||
db.delete(upload)
|
||
|
||
db.commit()
|
||
|
||
return {"message": f"Cleaned up {len(expired_uploads)} expired uploads"}
|
||
|
||
@router.get("/{kb_id}/documents/tasks")
|
||
async def get_processing_tasks(
|
||
kb_id: int,
|
||
task_ids: str = Query(..., description="Comma-separated list of task IDs to check status for"),
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user)
|
||
):
|
||
"""
|
||
Get status of multiple processing tasks.
|
||
"""
|
||
task_id_list = [int(id.strip()) for id in task_ids.split(",")]
|
||
|
||
kb = db.query(KnowledgeBase).filter(
|
||
KnowledgeBase.id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
|
||
if not kb:
|
||
raise HTTPException(status_code=404, detail="Knowledge base not found")
|
||
|
||
tasks = (
|
||
db.query(ProcessingTask)
|
||
.options(
|
||
selectinload(ProcessingTask.document_upload)
|
||
)
|
||
.filter(
|
||
ProcessingTask.id.in_(task_id_list),
|
||
ProcessingTask.knowledge_base_id == kb_id
|
||
)
|
||
.all()
|
||
)
|
||
|
||
return {
|
||
task.id: {
|
||
"document_id": task.document_id,
|
||
"status": task.status,
|
||
"error_message": task.error_message,
|
||
"upload_id": task.document_upload_id,
|
||
"file_name": task.document_upload.file_name if task.document_upload else None
|
||
}
|
||
for task in tasks
|
||
}
|
||
|
||
@router.get("/{kb_id}/documents/{doc_id}", response_model=DocumentResponse)
|
||
async def get_document(
|
||
*,
|
||
db: Session = Depends(get_db),
|
||
kb_id: int,
|
||
doc_id: int,
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Any:
|
||
"""
|
||
Get document details by ID.
|
||
"""
|
||
document = (
|
||
db.query(Document)
|
||
.join(KnowledgeBase)
|
||
.filter(
|
||
Document.id == doc_id,
|
||
Document.knowledge_base_id == kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
)
|
||
.first()
|
||
)
|
||
|
||
if not document:
|
||
raise HTTPException(status_code=404, detail="Document not found")
|
||
|
||
return document
|
||
|
||
@router.post("/test-retrieval")
|
||
async def test_retrieval(
|
||
request: TestRetrievalRequest,
|
||
background_tasks: BackgroundTasks,
|
||
db: Session = Depends(get_db),
|
||
current_user: User = Depends(get_current_user)
|
||
) -> Any:
|
||
"""
|
||
Test retrieval quality for a given query against a knowledge base.
|
||
"""
|
||
try:
|
||
kb = db.query(KnowledgeBase).filter(
|
||
KnowledgeBase.id == request.kb_id,
|
||
KnowledgeBase.user_id == current_user.id
|
||
).first()
|
||
|
||
if not kb:
|
||
raise HTTPException(
|
||
status_code=404,
|
||
detail=f"Knowledge base {request.kb_id} not found",
|
||
)
|
||
|
||
embeddings = EmbeddingsFactory.create()
|
||
|
||
vector_store = VectorStoreFactory.create(
|
||
store_type=settings.VECTOR_STORE_TYPE,
|
||
collection_name=f"kb_{request.kb_id}",
|
||
embedding_function=embeddings,
|
||
)
|
||
|
||
results = vector_store.similarity_search_with_score(request.query, k=request.top_k)
|
||
|
||
response = []
|
||
for doc, score in results:
|
||
response.append({
|
||
"content": doc.page_content,
|
||
"metadata": doc.metadata,
|
||
"score": float(score)
|
||
})
|
||
|
||
return {"results": response}
|
||
|
||
except Exception as e:
|
||
raise HTTPException(status_code=500, detail=str(e))
|