app.py

import asyncio

from fastapi import FastAPI
from starlette.exceptions import HTTPException
from pydantic import BaseModel

import last_layer

app = FastAPI()
class Request(BaseModel):
    text: str


@app.post("/scan-prompt/")
async def scan_prompt(chunk: Request) -> last_layer.RiskModel:
    try:
        result = await asyncio.to_thread(last_layer.scan_prompt, chunk.text)
        return result
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"An error occurred: {str(e)}")


@app.post("/scan-llm/")
async def scan_llm(chunk: Request):
    return {"message": "None"}