Spaces:
Running
Running
File size: 2,961 Bytes
75e2b6c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 |
# app/routers/admin.py
from fastapi import APIRouter, Depends, HTTPException, UploadFile, File
from typing import List
import os
from app.database.database_query import DatabaseQuery
from app.services.vector_database_search import VectorDatabaseSearch
from app.middleware.auth import get_current_user
from pydantic import BaseModel
router = APIRouter()
vector_db = VectorDatabaseSearch()
query = DatabaseQuery()
TEMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'temp')
os.makedirs(TEMP_DIR, exist_ok=True)
class SearchQuery(BaseModel):
query: str
k: int = 5
@router.get('/books')
async def get_books(username: str = Depends(get_current_user)):
try:
book_info = vector_db.get_book_info()
return {
'status': 'success',
'data': book_info
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@router.post('/books', status_code=201)
async def add_books(files: List[UploadFile] = File(...), username: str = Depends(get_current_user)):
try:
pdf_paths = []
for file in files:
if file.filename.endswith('.pdf'):
safe_filename = os.path.basename(file.filename)
temp_path = os.path.join(TEMP_DIR, safe_filename)
with open(temp_path, "wb") as buffer:
content = await file.read()
buffer.write(content)
pdf_paths.append(temp_path)
if not pdf_paths:
raise HTTPException(status_code=400, detail="No valid PDF files provided")
success_count = 0
for pdf_path in pdf_paths:
if vector_db.add_pdf(pdf_path):
success_count += 1
# Clean up temporary files
for path in pdf_paths:
try:
if os.path.exists(path):
os.remove(path)
except Exception:
pass
return {
'status': 'success',
'message': f'Successfully added {success_count} of {len(pdf_paths)} books'
}
except Exception as e:
# Clean up temporary files in case of error
for path in pdf_paths:
try:
if os.path.exists(path):
os.remove(path)
except:
pass
if isinstance(e, HTTPException):
raise e
raise HTTPException(status_code=500, detail=str(e))
@router.post('/search')
async def search_books(search_data: SearchQuery, username: str = Depends(get_current_user)):
try:
query_text = search_data.query
k = search_data.k
results = vector_db.search(
query=query_text,
top_k=k
)
return {
'status': 'success',
'data': results
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e)) |