""" FastAPI backend for the arXivCSRAG application. """ import os from typing import List, Optional from pathlib import Path from datetime import datetime from pydantic import BaseModel from fastapi import FastAPI, UploadFile, File, Form, HTTPException from fastapi.staticfiles import StaticFiles from fastapi.middleware.cors import CORSMiddleware from utils.setup_logger import setup_logger from src.config import TEMP_DIR, ROOT_DIR from src.fetcher.arxiv_fetcher import ArxivFetcher from src.data_extraction.extractor import extract_from_pdf, separate_content_types from src.processors.text_processor import TextProcessor from src.processors.table_processor import TableProcessor from src.processors.image_processor import ImageProcessor from src.storage.vectorstore import VectorStore from src.rag.pipeline import RAGPipeline # Configure logging logger = setup_logger(__name__) # Initialize the FastAPI app app = FastAPI( title = 'arXivCSRAG API', description = 'API for the arXivCSRAG Multimodal RAG Application', version = '1.0.0', ) # CORS configuration app.add_middleware( CORSMiddleware, allow_origins = ['*'], allow_credentials = True, allow_methods = ['*'], allow_headers = ['*'], ) # Models class APIKeys(BaseModel): gemini_api_key : str huggingface_token: str class SearchQuery(BaseModel): subject_tags: Optional[List[str]] = None start_date : Optional[str] = None end_date : Optional[str] = None max_results : int = 10 query : str class PaperID(BaseModel): arxiv_id: str class ChatMessage(BaseModel): message: str # Initialize components arxiv_fetcher = ArxivFetcher() text_processor = TextProcessor() table_processor = TableProcessor() image_processor = ImageProcessor() vector_store = VectorStore() rag_pipeline = RAGPipeline(vector_store.retriever) # API endpoints @app.post('/api/configure') async def configure_api_keys(api_keys: APIKeys): """Configure API keys for the application.""" try: # Set environment variables os.environ['GOOGLE_API_KEY'] = api_keys.gemini_api_key os.environ['HF_TOKEN'] = api_keys.huggingface_token logger.info('API keys configured successfully') return {'status' : 'success', 'message': 'API keys configured successfully'} except Exception as e: logger.error(f"Error configuring API keys: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/fetch-papers') async def fetch_papers(search_query: SearchQuery): """Fetch papers from arXiv based on search query and filters.""" try: papers = arxiv_fetcher.fetch_papers( subject_tags = search_query.subject_tags, start_date = search_query.start_date, end_date = search_query.end_date, max_results = search_query.max_results, query = search_query.query ) logger.info(f"Fetched {len(papers)} papers") return {'status': 'success', 'papers': papers} except Exception as e: logger.error(f"Error fetching papers: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/paper-metadata') async def get_paper_metadata(paper_id: PaperID): """Get metadata for a specific paper.""" try: search = arxiv_fetcher.fetch_papers(f"id:{paper_id.arxiv_id}", max_results=1) if not search: raise HTTPException(status_code=404, detail='Paper not found') return {'status': 'success', 'metadata': search[0]} except Exception as e: logger.error(f"Error getting paper metadata: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/download-paper') async def download_paper(paper_id: PaperID): """Download a paper's PDF from arXiv.""" try: pdf_path = arxiv_fetcher.download_paper(paper_id.arxiv_id) if not pdf_path: raise HTTPException(status_code=404, detail="Failed to download paper") logger.info(f"Downloaded paper {paper_id.arxiv_id} to {pdf_path}") return {'status': 'success', 'file_path': str(pdf_path)} except Exception as e: logger.error(f"Error downloading paper: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/upload-paper') async def upload_paper(file: UploadFile = File(...)): """Upload a paper's PDF file.""" try: # Create a unique filename timestamp = datetime.now().strftime('%Y%m%d%H%M%S') filename = f"uploaded_{timestamp}.pdf" filepath = TEMP_DIR / filename # Save the uploaded file with open(filepath, 'wb') as f: f.write(await file.read()) logger.info(f"Uploaded paper saved at {filepath}") return {'status': 'success', 'file_path': str(filepath)} except Exception as e: logger.error(f"Error uploading paper: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/process-paper') async def process_paper(file_path: str = Form(...)): """Process a paper for RAG.""" try: # # Reset the vector store # vector_store.reset() # # Set the new retriever for the RAG pipeline # rag_pipeline.retriever = vector_store.retriever # Process the paper pdf_path = Path(file_path) logger.info(f"Processing paper at {pdf_path}") if not pdf_path.exists(): raise HTTPException(status_code=404, detail='PDF file not found') # Extract content from PDF logger.info(f"Extracting content from {pdf_path}") chunks = extract_from_pdf(pdf_path) # Separate content types logger.info(f"Separating content types from {len(chunks)} chunks") content = separate_content_types(chunks) # Process and summarize content logger.info(f"Processing {len(content['texts'])} text content") text_summaries = text_processor.process(content['texts']) logger.info(f"Processing {len(content['tables'])} table content") table_summaries = table_processor.process(content['tables']) logger.info(f"Processing {len(content['images'])} image content") image_summaries = image_processor.process(content['images']) # Add to vector store logger.info("Adding processed content to vector store") vector_store.add_contents( content['texts'] , text_summaries, content['tables'], table_summaries, content['images'], image_summaries ) logger.info(f"Processed paper {pdf_path.name} successfully") return { 'status': 'success', 'stats' : { 'texts' : len(content['texts']), 'tables': len(content['tables']), 'images': len(content['images']) } } except Exception as e: logger.error(f"Error processing paper: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/chat') async def chat_with_paper(message: ChatMessage): """ Chat with a processed paper. Returns: - status: success or error - response: The generated text response - citations: Dictionary containing three keys: - texts: List of text excerpts used as citations - images: List of base64-encoded image strings - tables: List of HTML-formatted table strings """ try: rag_pipeline.retriever = vector_store.retriever # Query the RAG pipeline logger.info(f"Chatting with paper: {message.message}") response = rag_pipeline.query(message.message) # Get the retrieved documents retrieved_docs = vector_store.retrieve(message.message) parsed_docs = rag_pipeline.parse_docs(retrieved_docs) return { 'status' : 'success', 'response' : response['response'], 'citations': parsed_docs } except Exception as e: logger.error(f"Error chatting with paper: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/fetch-citations') async def fetch_citations(message: ChatMessage): """ Fetch citations for a specific query without generating a response. This is useful for retrieving only the source documents that would be used to answer a query without generating the complete answer. Returns: - status: success or error - citations: Dictionary containing three keys: - texts: List of text excerpts used as citations - images: List of base64-encoded image strings - tables: List of HTML-formatted table strings """ try: # Get the retrieved documents retrieved_docs = vector_store.retrieve(message.message) parsed_docs = rag_pipeline.parse_docs(retrieved_docs) return { 'status' : 'success', 'citations': parsed_docs } except Exception as e: logger.error(f"Error fetching citations: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.post('/api/reset-chat') async def reset_chat(): """Reset the chat and vector store.""" try: logger.info("Resetting chat and vector store") vector_store.reset() rag_pipeline.retriever = vector_store.retriever rag_pipeline.reset_memory() return {'status': 'success', 'message': 'Chat reset successfully'} except Exception as e: logger.error(f"Error resetting chat: {e}") raise HTTPException(status_code=500, detail=str(e)) # Serve static files app.mount('/static', StaticFiles(directory=ROOT_DIR / 'static', html=False), name='static') app.mount('/data' , StaticFiles(directory=ROOT_DIR / 'static/data') , name='data') app.mount('/' , StaticFiles(directory=ROOT_DIR / 'static', html=True) , name='root')