itsprasun/pythonic-rag-fastapi-react

1import os2from typing import List3from fastapi import FastAPI, HTTPException, UploadFile, File4from fastapi.middleware.cors import CORSMiddleware5from pydantic import BaseModel6import asyncio7import tempfile8from aimakerspace.vectordatabase import VectorDatabase9from aimakerspace.openai_utils.chatmodel import ChatOpenAI10 11from app import (12    RetrievalAugmentedQAPipeline,13    process_file,14    system_role_prompt,15    user_role_prompt,16)17 18app = FastAPI()19 20# Update CORS middleware configuration21app.add_middleware(22    CORSMiddleware,23    allow_origins=[24        "http://localhost:3001",    # Development React server25        "http://localhost:7860",    # Production nginx server26        "http://localhost",         # Just in case27        "*",                        # Allow all origins in development28    ],29    allow_credentials=True,30    allow_methods=["*"],31    allow_headers=["*"],32    expose_headers=["*"],33)34 35class ChatResponse(BaseModel):36    response: str37    context: List[tuple]38 39class ChatRequest(BaseModel):40    query: str41 42@app.post("/api/upload", response_model=dict)43async def upload_file(file: UploadFile = File(...)):44    try:45        # Create a temporary file to store the upload46        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{file.filename.split('.')[-1]}") as temp_file:47            content = await file.read()48            temp_file.write(content)49            temp_file.flush()50 51            # Process the file using existing function52            texts = process_file(temp_file.name, file.filename)53 54            # Create vector database55            vector_db = VectorDatabase()56            vector_db = await vector_db.abuild_from_list(texts)57 58            # Create chat model59            chat_openai = ChatOpenAI()60 61            # Create pipeline62            pipeline = RetrievalAugmentedQAPipeline(63                vector_db_retriever=vector_db,64                llm=chat_openai65            )66 67            # Store the pipeline in memory (Note: this is not production-ready)68            if not hasattr(app, 'pipelines'):69                app.pipelines = {}70            pipeline_id = str(len(app.pipelines))71            app.pipelines[pipeline_id] = pipeline72 73            # Clean up temporary file74            os.unlink(temp_file.name)75 76            return {"pipeline_id": pipeline_id, "message": "File processed successfully"}77 78    except Exception as e:79        raise HTTPException(status_code=500, detail=str(e))80 81@app.post("/api/chat/{pipeline_id}", response_model=ChatResponse)82async def chat(pipeline_id: str, request: ChatRequest):83    try:84        if not hasattr(app, 'pipelines') or pipeline_id not in app.pipelines:85            raise HTTPException(status_code=404, detail="Pipeline not found. Please upload a file first.")86 87        pipeline = app.pipelines[pipeline_id]88        result = await pipeline.arun_pipeline(request.query)89 90        # Collect the streaming response91        response_text = ""92        async for chunk in result["response"]:93            response_text += chunk94 95        return ChatResponse(96            response=response_text,97            context=result["context"]98        )99 100    except Exception as e:101        raise HTTPException(status_code=500, detail=str(e))102 103if __name__ == "__main__":104    import uvicorn105    uvicorn.run(app, host="0.0.0.0", port=8000)