from fastapi import FastAPI
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field
from test_model import mark_pii_context
class MarkRequest(BaseModel):
text: str = Field(..., min_length=1, description="Input text to tag PII")
class MarkResponse(BaseModel):
sentence: str
original_text: str
entities: list[dict]
context: str
app = FastAPI(title="PII Marking API", version="1.0.0")
@app.get("/", response_class=HTMLResponse)
def test_page() -> str:
return """
PII Marking Test
PII Marking Quick Test
sentence
-
original_text
-
"""
@app.get("/health")
def health() -> dict:
return {"status": "ok"}
@app.post("/mark", response_model=MarkResponse)
def mark_text(payload: MarkRequest) -> MarkResponse:
try:
result = mark_pii_context(payload.text)
return MarkResponse(
sentence=result["sentence"],
original_text=result["original_text"],
entities=result["entities"],
context=result["context"],
)
except Exception as e:
from fastapi import HTTPException
raise HTTPException(status_code=500, detail=str(e))