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