Newer
Older
cortex-hub / ai-hub / app / app.py
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI
from dotenv import load_dotenv
from typing import List

# Import core application logic
from app.core.vector_store import FaissVectorStore
from app.core.retrievers import FaissDBRetriever, Retriever
from app.core.rag_service import RAGService

# Import the new files for database and API routes
from app.db_setup import create_db_tables
from app.api_endpoints import create_api_router

# Load environment variables from a .env file
load_dotenv()

# --- Application Factory Function ---
def create_app() -> FastAPI:
    """
    Factory function to create and configure the FastAPI application.
    This encapsulates all setup logic, making the main entry point clean.
    """
    # Initialize core services for RAG
    # CORRECTED: Now passing the required arguments to FaissVectorStore
    vector_store = FaissVectorStore(index_file_path="faiss_index.bin", dimension=768)
    retrievers: List[Retriever] = [
        FaissDBRetriever(vector_store=vector_store),
    ]
    rag_service = RAGService(vector_store=vector_store, retrievers=retrievers)

    @asynccontextmanager
    async def lifespan(app: FastAPI):
        """
        Initializes the database and vector store on startup and handles
        cleanup on shutdown.
        """
        print("Initializing application services...")
        create_db_tables()
        yield
        print("Shutting down application services...")
        vector_store.save_index()

    app = FastAPI(
        title="AI Model Hub Service",
        description="A extensible hub to route requests to various LLMs with RAG capabilities.",
        version="0.0.0",
        lifespan=lifespan
    )

    # Create and include the API router
    api_router = create_api_router(rag_service=rag_service)
    app.include_router(api_router)

    return app