Spaces:
Sleeping
Sleeping
File size: 3,296 Bytes
29fbb51 4d6298c 29fbb51 4d6298c 29fbb51 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | 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
} |