Nithi1509/MicroServices

1"""2RAG Microservice for HuggingFace Space3Handles document processing and vector search4"""5import os6import tempfile7from typing import Optional, List8from fastapi import FastAPI, HTTPException, UploadFile, File, Header,Request9from fastapi.middleware.cors import CORSMiddleware10from pydantic import BaseModel11import httpx12from langchain.text_splitter import RecursiveCharacterTextSplitter13from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader, UnstructuredPowerPointLoader14from langchain_community.vectorstores import FAISS15from langchain_community.embeddings import HuggingFaceEmbeddings16 17app = FastAPI(title="Pack AI RAG Service", version="1.0")18 19# CORS20app.add_middleware(21    CORSMiddleware,22    allow_origins=["*"],23    allow_credentials=True,24    allow_methods=["*"],25    allow_headers=["*"],26)27 28# In-memory storage for vector stores29vector_stores = {}30 31# Initialize embeddings model (lightweight)32embeddings = HuggingFaceEmbeddings(33    model_name="sentence-transformers/all-MiniLM-L6-v2",34    model_kwargs={'device': 'cpu'}35)36 37# API Key for authentication38API_KEY = os.getenv("RAG_API_KEY", "your-secret-key-here")39 40 41# Pydantic models42class ProcessDocumentRequest(BaseModel):43    document_url: str44    document_id: str45    file_type: str = "application/pdf"46 47 48class QueryRequest(BaseModel):49    document_id: str50    query: str51    num_results: int = 452 53 54class ProcessResponse(BaseModel):55    success: bool56    document_id: str57    num_chunks: int = 058    message: str = ""59    error: str = ""60 61 62class QueryResponse(BaseModel):63    success: bool64    chunks: List[str] = []65    document_id: str66    error: str = ""67 68 69# Helper functions70def verify_api_key(authorization: str = Header(None)):71    """Verify API key"""72    if not authorization:73        raise HTTPException(status_code=401, detail="Missing authorization header")74    75    if authorization != f"Bearer {API_KEY}":76        raise HTTPException(status_code=403, detail="Invalid API key")77 78 79async def download_file(url: str, suffix: str) -> str:80    """Download file from URL to temp file"""81    async with httpx.AsyncClient() as client:82        response = await client.get(url, follow_redirects=True)83        response.raise_for_status()84        85        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp_file:86            tmp_file.write(response.content)87            return tmp_file.name88 89 90def load_document(file_path: str, file_type: str):91    """Load document based on file type"""92    if 'pdf' in file_type.lower():93        loader = PyPDFLoader(file_path)94    elif 'word' in file_type.lower() or 'docx' in file_type.lower():95        loader = Docx2txtLoader(file_path)96    elif 'powerpoint' in file_type.lower() or 'pptx' in file_type.lower():97        loader = UnstructuredPowerPointLoader(file_path)98    else:99        raise ValueError(f"Unsupported file type: {file_type}")100    101    return loader.load()102 103 104# Endpoints105@app.get("/")106async def root():107    return {108        "service": "Pack AI RAG Service",109        "version": "1.0",110        "status": "running",111        "documents_loaded": len(vector_stores)112    }113 114 115@app.get("/health")116async def health():117    return {118        "status": "healthy",119        "documents_in_memory": len(vector_stores)120    }121 122 123@app.post("/process-document", response_model=ProcessResponse)124async def process_document(125    request: ProcessDocumentRequest,126    authorization: str = Header(None)127):128    """129    Process a document from URL and create vector store130    """131    verify_api_key(authorization)132    133    try:134        # Determine file suffix135        suffix = '.pdf'136        if 'word' in request.file_type.lower() or 'docx' in request.file_type.lower():137            suffix = '.docx'138        elif 'powerpoint' in request.file_type.lower() or 'pptx' in request.file_type.lower():139            suffix = '.pptx'140        141        # Download file142        file_path = await download_file(request.document_url, suffix)143        144        try:145            # Load document146            documents = load_document(file_path, request.file_type)147            148            # Split into chunks149            text_splitter = RecursiveCharacterTextSplitter(150                chunk_size=1000,151                chunk_overlap=200,152                length_function=len153            )154            chunks = text_splitter.split_documents(documents)155            156            # Create vector store157            vector_store = FAISS.from_documents(chunks, embeddings)158            159            # Store in memory160            vector_stores[request.document_id] = vector_store161            162            return ProcessResponse(163                success=True,164                document_id=request.document_id,165                num_chunks=len(chunks),166                message=f"Document processed successfully with {len(chunks)} chunks"167            )168            169        finally:170            # Clean up temp file171            try:172                os.unlink(file_path)173            except:174                pass175                176    except Exception as e:177        return ProcessResponse(178            success=False,179            document_id=request.document_id,180            error=str(e)181        )182 183 184@app.post("/query", response_model=QueryResponse)185async def query_document(186    request: QueryRequest,187    authorization: str = Header(None)188):189    """190    Query a processed document191    """192    verify_api_key(authorization)193    194    try:195        # Check if document exists196        if request.document_id not in vector_stores:197            return QueryResponse(198                success=False,199                document_id=request.document_id,200                error="Document not found. Please process it first."201            )202        203        # Get vector store204        vector_store = vector_stores[request.document_id]205        206        # Search for relevant chunks207        docs = vector_store.similarity_search(request.query, k=request.num_results)208        209        # Extract text from documents210        chunks = [doc.page_content for doc in docs]211        212        return QueryResponse(213            success=True,214            chunks=chunks,215            document_id=request.document_id216        )217        218    except Exception as e:219        return QueryResponse(220            success=False,221            document_id=request.document_id,222            error=str(e)223        )224 225 226@app.delete("/document/{document_id}")227async def delete_document(228    document_id: str,229    authorization: str = Header(None)230):231    """232    Delete a document from memory233    """234    verify_api_key(authorization)235    236    if document_id in vector_stores:237        del vector_stores[document_id]238        return {239            "success": True,240            "message": f"Document {document_id} deleted successfully"241        }242    else:243        raise HTTPException(status_code=404, detail="Document not found")244 245 246@app.get("/documents")247async def list_documents(authorization: str = Header(None)):248    """249    List all documents in memory250    """251    verify_api_key(authorization)252    253    return {254        "success": True,255        "documents": list(vector_stores.keys()),256        "count": len(vector_stores)257    }258 259@app.api_route("/health", methods=["GET", "HEAD"])260def health_check(request: Request):261    return {262        "status": "ok-Perfect"263    }264 265    266if __name__ == "__main__":267    import uvicorn268    uvicorn.run(app, host="0.0.0.0", port=7860)269