api.py

from fastapi import FastAPI, File, UploadFile, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional, List
import os
import shutil
from rag_pipeline import RAGPipeline
from chunk_viewer import ChunkViewer
from config import Config

app = FastAPI(title="RAG System API", version="1.0.0")

# CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global instances
rag_pipeline = None
chunk_viewer = ChunkViewer()

# Request/Response models
class QueryRequest(BaseModel):
    question: str
    top_k: Optional[int] = 5

class QueryResponse(BaseModel):
    question: str
    answer: str
    sources: List[dict]
    context: Optional[str] = None

class StatusResponse(BaseModel):
    status: str
    message: str
    details: Optional[dict] = None

@app.on_event("startup")
async def startup_event():
    """Initialize RAG pipeline on startup"""
    global rag_pipeline
    try:
        Config.validate()
        rag_pipeline = RAGPipeline()
        # Try to load existing index
        rag_pipeline.load_existing_index()
    except Exception as e:
        print(f"Startup warning: {e}")

@app.get("/")
async def root():
    """Root endpoint"""
    return {
        "message": "RAG System API",
        "version": "1.0.0",
        "endpoints": {
            "POST /ingest": "Upload and ingest PDF",
            "POST /query": "Query the RAG system",
            "GET /chunks": "View all chunks",
            "GET /chunks/{chunk_id}": "View specific chunk",
            "GET /chunks/statistics": "Get chunk statistics",
            "GET /chunks/search": "Search chunks by text",
            "GET /embeddings": "View embeddings",
            "GET /status": "Get system status"
        }
    }

@app.post("/ingest", response_model=StatusResponse)
async def ingest_pdf(
    file: UploadFile = File(...),
    use_vision: bool = Query(True, description="Use GPT-4 Vision for image analysis")
):
    """Upload and ingest a PDF file"""
    global rag_pipeline
    
    if not file.filename.endswith('.pdf'):
        raise HTTPException(status_code=400, detail="Only PDF files are allowed")
    
    # Save uploaded file
    upload_dir = "uploads"
    os.makedirs(upload_dir, exist_ok=True)
    file_path = os.path.join(upload_dir, file.filename)
    
    try:
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(file.file, buffer)
        
        # Ingest PDF
        if rag_pipeline is None:
            rag_pipeline = RAGPipeline()
        
        rag_pipeline.ingest_pdf(file_path, use_vision=use_vision)
        
        return StatusResponse(
            status="success",
            message=f"PDF '{file.filename}' ingested successfully",
            details={
                "filename": file.filename,
                "vision_enabled": use_vision
            }
        )
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error ingesting PDF: {str(e)}")
    
    finally:
        # Clean up uploaded file
        if os.path.exists(file_path):
            os.remove(file_path)

@app.post("/query", response_model=QueryResponse)
async def query_rag(request: QueryRequest):
    """Query the RAG system"""
    global rag_pipeline
    
    if rag_pipeline is None or rag_pipeline.retriever is None:
        raise HTTPException(
            status_code=400,
            detail="No PDF has been ingested. Please upload a PDF first."
        )
    
    try:
        result = rag_pipeline.query(request.question, top_k=request.top_k)
        
        if 'error' in result:
            raise HTTPException(status_code=400, detail=result['error'])
        
        return QueryResponse(**result)
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error querying: {str(e)}")

@app.get("/chunks")
async def get_chunks(
    page_num: Optional[int] = Query(None, description="Filter by page number"),
    content_type: Optional[str] = Query(None, description="Filter by content type (text, vision_analysis)")
):
    """Get all chunks with optional filtering"""
    try:
        chunks = chunk_viewer.list_chunks(page_num=page_num, content_type=content_type)
        return {
            "total": len(chunks),
            "filters": {
                "page_num": page_num,
                "content_type": content_type
            },
            "chunks": chunks
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error retrieving chunks: {str(e)}")

@app.get("/chunks/{chunk_id}")
async def get_chunk(chunk_id: int):
    """Get a specific chunk by ID"""
    try:
        chunk = chunk_viewer.view_chunk(chunk_id)
        if chunk is None:
            raise HTTPException(status_code=404, detail=f"Chunk {chunk_id} not found")
        return chunk
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error retrieving chunk: {str(e)}")

@app.get("/chunks/statistics/summary")
async def get_chunk_statistics():
    """Get chunk statistics"""
    try:
        stats = chunk_viewer.get_chunk_statistics()
        if 'error' in stats:
            raise HTTPException(status_code=404, detail=stats['error'])
        return stats
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error getting statistics: {str(e)}")

@app.get("/chunks/search/text")
async def search_chunks(query: str = Query(..., description="Text to search for in chunks")):
    """Search chunks by text content"""
    try:
        results = chunk_viewer.search_chunks_by_text(query)
        return {
            "query": query,
            "total_results": len(results),
            "results": results
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error searching chunks: {str(e)}")

@app.get("/embeddings")
async def get_embeddings(limit: int = Query(10, description="Number of embeddings to return")):
    """Get embeddings data"""
    try:
        embeddings_data = chunk_viewer.load_embeddings()
        if 'error' in embeddings_data:
            raise HTTPException(status_code=404, detail=embeddings_data['error'])
        
        # Return limited results
        embeddings = embeddings_data.get('embeddings', [])[:limit]
        
        return {
            "total_embeddings": embeddings_data.get('total_embeddings'),
            "embedding_model": embeddings_data.get('embedding_model'),
            "embedding_dimension": embeddings_data.get('embedding_dimension'),
            "source_pdf": embeddings_data.get('source_pdf'),
            "showing": len(embeddings),
            "embeddings": embeddings
        }
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error retrieving embeddings: {str(e)}")

@app.get("/status", response_model=StatusResponse)
async def get_status():
    """Get system status"""
    global rag_pipeline
    
    status = {
        "rag_initialized": rag_pipeline is not None,
        "index_loaded": rag_pipeline is not None and rag_pipeline.retriever is not None,
        "config": {
            "embedding_model": Config.EMBEDDING_MODEL,
            "chat_model": Config.CHAT_MODEL,
            "chunk_size": Config.CHUNK_SIZE,
            "top_k": Config.TOP_K_RESULTS
        }
    }
    
    # Get chunk statistics if available
    try:
        stats = chunk_viewer.get_chunk_statistics()
        if 'error' not in stats:
            status['chunks_available'] = stats.get('total_chunks', 0)
    except:
        pass
    
    return StatusResponse(
        status="operational",
        message="System is running",
        details=status
    )

@app.delete("/reset")
async def reset_system():
    """Reset the system (clear current index)"""
    global rag_pipeline
    
    try:
        rag_pipeline = RAGPipeline()
        return StatusResponse(
            status="success",
            message="System reset successfully"
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error resetting system: {str(e)}")

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