features/rag_chatbot/routes.py

from fastapi import APIRouter, Depends, HTTPException, UploadFile, File, Request
from fastapi.security import HTTPBearer
from pydantic import BaseModel, Field
from slowapi.util import get_remote_address
from slowapi import Limiter
from typing import Optional
from config import ACCESS_RATE, Config
from .controller import (
    handle_rag_query,
    handle_document_upload,
    handle_health_check,
    verify_token,
)

limiter = Limiter(key_func=get_remote_address)
router = APIRouter(prefix="/rag", tags=["RAG Chatbot"])
security = HTTPBearer()

class QueryInput(BaseModel):
    query: str = Field(..., min_length=1, max_length=1000, description="The question to ask")

class QueryResponse(BaseModel):
    answer: str
    source: str
    route: Optional[str] = None
    documents: Optional[list] = None
    error: Optional[str] = None

class UploadResponse(BaseModel):
    message: str
    filename: str
    text_length: int
    content_type: str

class HealthResponse(BaseModel):
    status: str
    components: Optional[dict] = None
    error: Optional[str] = None

@router.post("/question", response_model=QueryResponse)
@limiter.limit(ACCESS_RATE)
async def ask_question(
    request: Request, 
    data: QueryInput, 
    token: str = Depends(verify_token)
) -> QueryResponse:
    """
    Ask a question to the RAG chatbot.
    
    The chatbot can answer:
    - Company-related questions (based on uploaded documents)
    - Cybersecurity questions (from knowledge base)
    """
    response = await handle_rag_query(data.query)
    return QueryResponse(**response)

@router.post("/upload", response_model=UploadResponse)
@limiter.limit(ACCESS_RATE)
async def upload_document(
    request: Request, 
    file: UploadFile = File(..., description="Document file (PDF, DOCX, or TXT)"),
    token: str = Depends(verify_token)
) -> UploadResponse:
    """
    Upload a document to the company knowledge base.
    
    Supported formats:
    - PDF (.pdf)
    - Word documents (.docx)
    - Plain text (.txt)
    
    Maximum file size: 10MB
    """
    response = await handle_document_upload(file)
    return UploadResponse(**response)

@router.get("/health", response_model=HealthResponse)
@limiter.limit(ACCESS_RATE)
async def health_check(request: Request) -> HealthResponse:
    """
    Check the health status of the RAG system.
    
    Returns the status of all components:
    - ChromaDB connection
    - Vector store
    - AI chains
    """
    response = await handle_health_check()
    return HealthResponse(**response)

@router.get("/info")
@limiter.limit(ACCESS_RATE)
async def get_system_info(request: Request):
    """Get information about the RAG system capabilities."""
    return {
        "name": "RAG Chatbot",
        "version": "1.0.0",
        "description": "A specialized chatbot for cybersecurity and company-related questions",
        "capabilities": [
            "Company document Q&A (based on uploaded documents)",
            "Cybersecurity knowledge and best practices",
            "Document upload and processing (PDF, DOCX, TXT)"
        ],
        "supported_file_types": sorted(Config.RAG_SUPPORTED_CONTENT_TYPES),
        "max_file_size_mb": round(Config.RAG_MAX_FILE_SIZE / (1024 * 1024), 2),
        "max_query_length": Config.RAG_MAX_QUERY_LENGTH
    }