# 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))