api.py

# File: api.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from data_ingestion import run_ingestion_pipeline, fetch_arxiv_metadata, persist_to_dataset, ArxivQueryRequest
from retrieval import rag_query
import logging
import traceback

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

app = FastAPI()

class QueryRequest(BaseModel):
    query: str

@app.post("/ingest")
async def ingest(request: ArxivQueryRequest):
    try:
        logging.info(f"Received ingestion request with query: {request.query} and max_results: {request.max_results}")
        result = run_ingestion_pipeline(request.query, request.max_results)
        logging.info(f"Ingestion result: {result}")
        return {"message": result}
    except Exception as e:
        error_message = f"Ingestion error for query: {request.query} - {str(e)}\n{traceback.format_exc()}"
        logging.error(error_message)
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/query")
async def query(request: QueryRequest):
    try:
        result = rag_query(request.query)
        return {"result": result}
    except Exception as e:
        logging.error(f"Query error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/fetch_and_persist_arxiv")
async def fetch_and_persist_arxiv(request: ArxivQueryRequest):
    try:
        metadata = fetch_arxiv_metadata(request.query, request.max_results)
        if not metadata:
            return {"message": "No results found or error occurred during fetching."}
        persist_to_dataset(metadata)
        return {"message": f"Successfully fetched and persisted {len(metadata)} papers."}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)