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.
    """
    router = APIRouter()

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

    # --- Chat Endpoint ---
    @router.post("/chat", response_model=schemas.ChatResponse, summary="Get AI-Generated Response", tags=["Chat"])
    async def chat_handler(
        request: schemas.ChatRequest,
        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 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}"
            )

    # --- Document Management Endpoints ---
    @router.post("/documents", response_model=schemas.DocumentResponse, summary="Add a New Document", tags=["Documents"])
    def add_document(
        doc: schemas.DocumentCreate,
        db: Session = Depends(get_db)
    ):
        """
        Adds a new document to the database and vector store.
        """
        try:
            doc_data = doc.model_dump()
            document_id = rag_service.add_document(db=db, doc_data=doc_data)
            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}")

    @router.get("/documents", response_model=schemas.DocumentListResponse, summary="List All Documents", tags=["Documents"])
    def get_documents(db: Session = Depends(get_db)):
        """
        Retrieves a list of all documents in the knowledge base.
        """
        try:
            documents_from_db = rag_service.get_all_documents(db=db)
            # **SIMPLIFICATION**: Just return the list of ORM objects.
            # FastAPI will use your Pydantic schema's new ORM mode to convert it.
            return {"documents": documents_from_db}
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"An error occurred: {e}")

    @router.delete("/documents/{document_id}", response_model=schemas.DocumentDeleteResponse, summary="Delete a Document", tags=["Documents"])
    def delete_document(document_id: int, db: Session = Depends(get_db)):
        """
        Deletes a document from the database and vector store by its ID.
        """
        try:
            # Note: You'll need to implement the `delete_document` method in your RAGService.
            # This method should return the ID of the deleted doc or raise an error if not found.
            deleted_id = rag_service.delete_document(db=db, document_id=document_id)
            if deleted_id is None:
                raise HTTPException(status_code=404, detail=f"Document with ID {document_id} not found.")
            
            return schemas.DocumentDeleteResponse(
                message="Document deleted successfully",
                document_id=deleted_id
            )
        except HTTPException:
            # Re-raise HTTPException to preserve the 404 status
            raise
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"An error occurred: {e}")

    return router