Newer
Older
cortex-hub / ai-hub / app / api / routes.py
from fastapi import APIRouter, HTTPException, Depends
from sqlalchemy.orm import Session
from app.core.services import RAGService 
from app.api.dependencies import get_db
from app.api import schemas

def create_api_router(rag_service: RAGService) -> APIRouter:
    """
    Creates and returns an APIRouter with all the application's endpoints.
    
    This function takes the RAGService instance as an argument, so it can be
    injected from the main application factory.
    """
    router = APIRouter()

    @router.get("/", summary="Check Service Status")
    def read_root():
        return {"status": "AI Model Hub is running!"}

    # Use the schemas for request body validation and to define the response model
    @router.post("/chat", response_model=schemas.ChatResponse, summary="Get AI-Generated Response")
    async def chat_handler(
        request: schemas.ChatRequest, # <-- Use the imported schema for the request body
        db: Session = Depends(get_db)
    ):
        """
        Handles a chat request using the prompt and model from the request body.
        """
        try:
            response_text = await rag_service.chat_with_rag(
                db=db,
                prompt=request.prompt,
                model=request.model
            )
            # Return an instance of the response schema for automatic serialization
            return schemas.ChatResponse(answer=response_text, model_used=request.model)
        except Exception as e:
            raise HTTPException(
                status_code=500, 
                detail=f"An unexpected error occurred with the {request.model} API: {e}"
            )

    # Use the schemas for the /document endpoint as well
    @router.post("/document", response_model=schemas.DocumentResponse, summary="Add a New Document")
    def add_document(
        doc: schemas.DocumentCreate, # <-- Use the imported schema for the request body
        db: Session = Depends(get_db)
    ):
        """
        Adds a new document to the database and vector store.
        """
        try:
            # The 'doc' object is already a validated Pydantic model
            doc_data = doc.model_dump()
            document_id = rag_service.add_document(db=db, doc_data=doc_data)
            
            # Return an instance of the response schema
            return schemas.DocumentResponse(
                message=f"Document '{doc.title}' added successfully with ID {document_id}"
            )
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"An error occurred: {e}")

    return router