from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session from pydantic import BaseModel from typing import List, Optional import logging import PyPDF2 import io import uuid from core.database import get_db from models import db_models from services.rag_service import rag_service from services.s3_service import s3_service from api.auth import get_current_user from core.config import settings from openai import OpenAI router = APIRouter(prefix="/api/rag", tags=["RAG Document Management"]) logger = logging.getLogger(__name__) # Request/Response Models class RAGIndexRequest(BaseModel): file_key: str # S3 key of source file to index class RAGIndexResponse(BaseModel): id: int filename: str azure_doc_id: str chunk_count: int message: str class RAGDocumentResponse(BaseModel): id: int filename: str azure_doc_id: str chunk_count: int source_id: Optional[int] created_at: str class Config: from_attributes = True def extract_text_from_pdf(file_content: bytes) -> str: """Extract text from PDF file.""" try: pdf_reader = PyPDF2.PdfReader(io.BytesIO(file_content)) text = "" for page in pdf_reader.pages: text += page.extract_text() + "\n" return text.strip() except Exception as e: logger.error(f"Error extracting PDF text: {e}") raise HTTPException(status_code=400, detail=f"Failed to extract text: {str(e)}") def chunk_text(text: str, chunk_size: int = 1000, overlap: int = 200) -> List[str]: """Split text into overlapping chunks.""" chunks = [] start = 0 while start < len(text): end = start + chunk_size chunks.append(text[start:end]) start += (chunk_size - overlap) return chunks @router.post("/index", response_model=RAGIndexResponse) async def index_document( request: RAGIndexRequest, current_user: db_models.User = Depends(get_current_user), db: Session = Depends(get_db)): """ Index a document for AI search (one-time operation). Downloads from S3, extracts text, generates embeddings, stores in Azure Search. """ try: # 1. Verify file ownership source = db.query(db_models.Source).filter( db_models.Source.s3_key == request.file_key, db_models.Source.user_id == current_user.id ).first() if not source: raise HTTPException(status_code=404, detail="File not found") # 2. Check if already indexed existing = db.query(db_models.RAGDocument).filter( db_models.RAGDocument.source_id == source.id, db_models.RAGDocument.user_id == current_user.id ).first() if existing: return RAGIndexResponse( id=existing.id, filename=existing.filename, azure_doc_id=existing.azure_doc_id, chunk_count=existing.chunk_count, message="Document already indexed" ) # 3. Download from S3 logger.info(f"Downloading {request.file_key}...") # Create temp local path import tempfile import os with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(source.filename)[1]) as tmp: temp_file = tmp.name s3_service.s3_client.download_file( settings.AWS_S3_BUCKET, request.file_key, temp_file ) # 4. Extract text try: with open(temp_file, "rb") as f: file_content = f.read() if source.filename.lower().endswith('.pdf'): text = extract_text_from_pdf(file_content) elif source.filename.lower().endswith('.txt'): text = file_content.decode('utf-8') else: raise HTTPException(status_code=400, detail="Only PDF and TXT supported") if len(text) < 10: raise HTTPException(status_code=400, detail="No readable text content found in file") # 5. Chunk text chunks = chunk_text(text) logger.info(f"Created {len(chunks)} chunks") # 6. Generate doc ID and index in Azure Search doc_id = str(uuid.uuid4()) chunk_count = rag_service.index_document( chunks=chunks, filename=source.filename, user_id=current_user.id, doc_id=doc_id ) # 7. Save to database rag_doc = db_models.RAGDocument( filename=source.filename, azure_doc_id=doc_id, chunk_count=chunk_count, user_id=current_user.id, source_id=source.id ) db.add(rag_doc) db.commit() db.refresh(rag_doc) logger.info(f"Successfully indexed {source.filename}") return RAGIndexResponse( id=rag_doc.id, filename=rag_doc.filename, azure_doc_id=rag_doc.azure_doc_id, chunk_count=rag_doc.chunk_count, message="Document indexed successfully for AI conversation" ) finally: if os.path.exists(temp_file): os.remove(temp_file) except HTTPException: raise except Exception as e: logger.error(f"Error indexing document: {e}", exc_info=True) raise HTTPException(status_code=500, detail=f"Indexing failed: {str(e)}") @router.get("/documents", response_model=List[RAGDocumentResponse]) async def list_indexed_documents( current_user: db_models.User = Depends(get_current_user), db: Session = Depends(get_db) ): """List all documents that have been processed and are ready for chatting.""" documents = db.query(db_models.RAGDocument).filter( db_models.RAGDocument.user_id == current_user.id ).order_by(db_models.RAGDocument.created_at.desc()).all() return [ RAGDocumentResponse( id=doc.id, filename=doc.filename, azure_doc_id=doc.azure_doc_id, chunk_count=doc.chunk_count, source_id=doc.source_id, created_at=doc.created_at.isoformat() ) for doc in documents ] @router.delete("/documents/{doc_id}") async def delete_indexed_document( doc_id: str, # Azure doc ID current_user: db_models.User = Depends(get_current_user), db: Session = Depends(get_db) ): """Remove a document from the AI search index.""" # Find document rag_doc = db.query(db_models.RAGDocument).filter( db_models.RAGDocument.azure_doc_id == doc_id, db_models.RAGDocument.user_id == current_user.id ).first() if not rag_doc: raise HTTPException(status_code=404, detail="Document index entry not found") try: # Delete from Azure Search rag_service.delete_document(doc_id) # Delete from database db.delete(rag_doc) db.commit() return {"message": "AI index for document deleted successfully"} except Exception as e: logger.error(f"Error deleting document index: {e}") raise HTTPException(status_code=500, detail=f"Deletion failed: {str(e)}") class RAGSummaryRequest(BaseModel): rag_doc_id: int @router.post("/summary") async def generate_document_summary( request: RAGSummaryRequest, current_user: db_models.User = Depends(get_current_user), db: Session = Depends(get_db) ): """ Generate an on-the-fly summary for an indexed document. No data is stored in the database. """ try: # 1. Verify existence and ownership rag_doc = db.query(db_models.RAGDocument).filter( db_models.RAGDocument.id == request.rag_doc_id, db_models.RAGDocument.user_id == current_user.id ).first() if not rag_doc: raise HTTPException(status_code=404, detail="Document not found") # 2. Fetch top chunks to build a summary # We search with a generic prompt to get a representative spread of content results = rag_service.search_document( query="Give me a general overview and executive summary of this document.", doc_id=rag_doc.azure_doc_id, user_id=current_user.id, top_k=8 # Fetch more context for a better summary ) if not results: return {"summary": "No content found to summarize."} context = "\n\n".join([r["content"] for r in results]) # 3. Generate summary using OpenAI openai_client = OpenAI(api_key=settings.OPENAI_API_KEY) response = openai_client.chat.completions.create( model="gpt-4o-mini", messages=[ { "role": "system", "content": "You are a professional document analyst. Provide a concise, high-level summary (3-5 sentences) of the document based on the provided context." }, {"role": "user", "content": f"Context from '{rag_doc.filename}':\n\n{context}"} ], temperature=0.5 ) return {"summary": response.choices[0].message.content} except Exception as e: logger.error(f"Summary generation failed: {e}") raise HTTPException(status_code=500, detail=f"Failed to generate summary: {str(e)}")