Newer
Older
cortex-hub / ai-hub / app / api_endpoints.py
from fastapi import APIRouter, HTTPException, Query, Depends
from pydantic import BaseModel, Field # Import Field here
from typing import Literal
from sqlalchemy.orm import Session
from app.core.rag_service import RAGService
from app.db_setup import get_db

# Pydantic Models for API requests
class ChatRequest(BaseModel):
    # Added min_length to ensure the prompt is not an empty string
    prompt: str = Field(..., min_length=1)

class DocumentCreate(BaseModel):
    title: str
    text: str
    source_url: str = None
    author: str = None
    user_id: str = "default_user"

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("/")
    def read_root():
        return {"status": "AI Model Hub is running!"}

    @router.post("/chat")
    async def chat_handler(
        request: ChatRequest,
        model: Literal["deepseek", "gemini"] = Query("deepseek", description="The AI model to use."),
        db: Session = Depends(get_db)
    ):
        try:
            response_text = await rag_service.chat_with_rag(
                db=db,
                prompt=request.prompt,
                model=model
            )
            return {"response": response_text, "model_used": model}
        except ValueError as e:
            raise HTTPException(status_code=400, detail=str(e))
        except Exception as e:
            raise HTTPException(status_code=500, detail=f"An error occurred with the {model} API: {e}")

    @router.post("/document")
    async def add_document(
        doc: DocumentCreate,
        db: Session = Depends(get_db)
    ):
        """
        Adds a new document to the database and its vector embedding to the FAISS index.
        """
        try:
            doc_data = doc.model_dump()
            document_id = rag_service.add_document(db=db, doc_data=doc_data)
            
            return {"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