CoolFace
Apppublic

ruby2210/rag-chatbot

sourceHugging Faceupdated 9mo agoView on Hugging Face
0likes
main.py70 linesDownload Raw Back to api
1"""2Main FastAPI application for the RAG Chatbot.3"""4from fastapi import FastAPI5from fastapi.middleware.cors import CORSMiddleware6from .chat_endpoints import router as chat_router7from .ingestion_endpoints import router as ingestion_router8from ..utils.database import init_db9from ..utils.vector_db import vector_db10from ..utils.logging import get_logger11 12 13logger = get_logger(__name__)14 15# Create the FastAPI app16app = FastAPI(17    title="RAG Chatbot API",18    description="API for the Retrieval-Augmented Generation Chatbot for AI-Spec-Driven Book",19    version="1.0.0"20)21 22# Add CORS middleware23app.add_middleware(24    CORSMiddleware,25    allow_origins=["*"],  # In production, configure this properly26    allow_credentials=True,27    allow_methods=["*"],28    allow_headers=["*"],29)30 31# Include the routers32app.include_router(chat_router, prefix="/api", tags=["chat"])33app.include_router(ingestion_router, prefix="/api", tags=["ingestion"])34 35# Initialize the database on startup36@app.on_event("startup")37async def startup_event():38    """39    Initialize database tables and vector database collection on startup.40    """41    try:42        # Initialize database tables43        init_db()44        logger.info("Database initialized successfully")45 46        # Initialize vector database collection with Cohere embedding size (1024)47        success = vector_db.create_collection(vector_size=1024)48        if success:49            logger.info("Vector database collection created/verified successfully")50        else:51            logger.error("Failed to create vector database collection")52    except Exception as e:53        logger.error(f"Error during startup: {e}")54        raise55 56 57@app.get("/")58async def root():59    """60    Root endpoint for health check.61    """62    return {"message": "RAG Chatbot API is running"}63 64 65@app.get("/health")66async def health_check():67    """68    Health check endpoint.69    """70    return {"status": "healthy", "service": "RAG Chatbot API"}