| 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 |
| } |
