sense-backend / api.py
SHAIK ADAM SHAFI
improved the pii explanation feature
371e762
Raw
History Blame Contribute Delete
88.3 kB
import os
from fastapi import FastAPI, File, UploadFile, HTTPException, Form, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from fastapi.encoders import jsonable_encoder
from pydantic import BaseModel, Field, EmailStr
from typing import Optional, List, Dict, Any
import pandas as pd
import io
import json
import time
import asyncio
from backend import RegexClassifier
from email_service import send_welcome_email
import evaluator_api
import video_router
import video_job_queue
# db.schema_inspector provides unified SQLAlchemy-based schema reflection (TICKET-4)
from db.schema_inspector import SchemaInspector
_schema_inspector = SchemaInspector() # module-level singleton — stateless, safe to share
# Initialize FastAPI app
app = FastAPI(
title="Segmento Sense API",
description="AI-powered PII Detection and Data Classification Platform",
version="1.0.0"
)
# CORS Configuration
# Keep all known origins: production domains + local dev ports
app.add_middleware(
CORSMiddleware,
allow_origins=[
# ── Production ──────────────────────────────────────
"https://segmento.in",
"https://www.segmento.in",
"https://segmento-sense.vercel.app",
# ── Local development ────────────────────────────────
"http://localhost:3000",
"http://localhost:3001",
"http://localhost:3002",
],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Initialize the classifier
classifier = RegexClassifier()
# Wire evaluator API — must happen after classifier is ready
evaluator_api.setup(classifier)
app.include_router(evaluator_api.router)
# Wire video API — async job queue (startup worker launched below)
video_job_queue.setup(classifier)
app.include_router(video_router.router)
@app.on_event("startup")
async def _start_video_worker():
"""Launch the background asyncio video processing worker."""
await video_job_queue.startup()
# Maximum file size (1GB)
MAX_FILE_SIZE = 1024 * 1024 * 1024 # 1GB in bytes
# ==================== PYDANTIC MODELS ====================
class TextAnalysisRequest(BaseModel):
text: str = Field(..., description="Text to analyze for PII")
class PatternAddRequest(BaseModel):
name: str = Field(..., description="Pattern name")
regex: str = Field(..., description="Regex pattern")
class DatabaseConnectionRequest(BaseModel):
host: str
port: str
database: str
user: str
password: str
table: str = Field(None, description="Table name (or collection for MongoDB)")
scan_mode: str = Field("full", description="full, incremental, sampling, metadata_only, metadata_and_sampling")
class AwsRdsConnectionRequest(DatabaseConnectionRequest):
engine: str = Field(..., description="E.g., postgres, mysql, aurora_mysql")
class DynamoDbConnectionRequest(BaseModel):
access_key: str
secret_key: str
region: str
table: str = Field(None, description="Table name")
class MetadataScanRequest(DatabaseConnectionRequest):
connector_type: str = Field(..., description="E.g., postgresql or mysql")
engine: Optional[str] = Field(None, description="Only for aws-rds")
access_key: Optional[str] = None
secret_key: Optional[str] = None
region: Optional[str] = None
class S3ConnectionRequest(BaseModel):
access_key: str
secret_key: str
region: str
bucket: str = Field(None, description="Bucket name")
file_key: str = Field(None, description="File key/path")
class AzureConnectionRequest(BaseModel):
connection_string: str
container: str = Field(None, description="Container name")
blob: str = Field(None, description="Blob name")
class GCSConnectionRequest(BaseModel):
credentials: Dict[str, Any]
bucket: str = Field(None, description="Bucket name")
file_name: str = Field(None, description="File name")
class GoogleDriveRequest(BaseModel):
credentials: Dict[str, Any]
file_id: str = Field(None, description="Drive file ID")
mime_type: str = Field(None, description="File MIME type")
class SlackRequest(BaseModel):
token: str
channel_id: str
class GmailRequest(BaseModel):
credentials: Dict[str, Any]
auth_type: str = Field("service_account", description="service_account or oauth2_token")
class ZendeskRequest(BaseModel):
subdomain: str
email: str
api_token: str
class SalesforceRequest(BaseModel):
instance_url: str
access_token: str
class GlueCredentials(BaseModel):
access_key: str
secret_key: str
region: str
class GlueListTablesRequest(GlueCredentials):
database_name: str
class GlueScanRequest(GlueCredentials):
database_name: str
table_name: str
class ConfluenceRequest(BaseModel):
url: str
username: str
token: str
page_id: str
class PDFPageRequest(BaseModel):
page_number: int = 0
class WelcomeEmailRequest(BaseModel):
name: str = Field(..., description="User's name")
email: EmailStr = Field(..., description="User's email address")
class RegisterRequest(BaseModel):
name: str
email: str
password: str
class LoginRequest(BaseModel):
email: str
password: str
# ==================== HELPER FUNCTIONS ====================
def validate_file_size(file: UploadFile):
"""Validate uploaded file size"""
file.file.seek(0, 2) # Seek to end
size = file.file.tell() # Get position (file size)
file.file.seek(0) # Reset to beginning
if size > MAX_FILE_SIZE:
raise HTTPException(
status_code=413,
detail=f"File size ({size} bytes) exceeds maximum allowed size (1GB)"
)
return size
def format_pii_response(df: pd.DataFrame, source_df: pd.DataFrame = None, text: str = None,
selected_models: list = None) -> Dict:
"""Format PII analysis response"""
count_df = classifier.get_pii_counts_dataframe(df, selected_models) if source_df is not None else classifier.get_pii_counts(text, selected_models)
response = {
"pii_counts": count_df.fillna("").to_dict(orient="records") if not count_df.empty else [],
"total_pii_found": int(count_df["Count"].sum()) if not count_df.empty else 0
}
# Add schema if source dataframe provided
if source_df is not None and not source_df.empty:
schema_df = classifier.get_data_schema(source_df)
response["schema"] = schema_df.fillna("").to_dict(orient="records")
# Add inspector results if text provided
if text:
inspector_df = classifier.run_full_inspection(text, selected_models)
if not inspector_df.empty:
response["inspector"] = inspector_df.fillna("").to_dict(orient="records")
return response
# ==================== MODEL CATALOGUE ENDPOINT ====================
@app.get("/api/models")
async def get_available_models():
"""
Returns the full list of available PII detection models,
separated into always-on and lazy-loaded categories.
"""
return JSONResponse(content={
"always_on": [
{"key": "regex", "label": "🛠️ Regex", "description": "Fast rule-based pattern matching (emails, phones, SSNs)"},
{"key": "nltk", "label": "🧠 NLTK", "description": "Statistical NLP chunker for names and locations"},
{"key": "spacy", "label": "🤖 SpaCy", "description": "Industrial-strength NER (en_core_web_lg)"},
{"key": "presidio", "label": "🛡️ Presidio", "description": "Microsoft Presidio — enterprise PII analyser"},
{"key": "gliner", "label": "🦅 GLiNER", "description": "Zero-shot entity extraction (urchade/gliner_small-v2.1)"},
{"key": "deberta", "label": "🚀 DeBERTa", "description": "Kaggle-winning DeBERTa V3 fine-tuned for PII"},
],
"lazy_loaded": [
{"key": "pasteproof", "label": "📋 Pasteproof", "description": "joneauxedgar/pasteproof-pii-detector-v2 — broad PII detection"},
{"key": "piiranha", "label": "🐟 Piiranha", "description": "iiiorg/piiranha-v1 — personal information specialist"},
{"key": "nvidia_gliner", "label": "⚡ NVIDIA-GLiNER", "description": "nvidia/gliner-PII — enterprise-grade zero-shot NER"},
{"key": "mmbert", "label": "🌐 mmbert32k", "description": "llm-semantic-router/mmbert32k — 32k-context document scanner"},
]
})
# ==================== FILE UPLOAD ENDPOINTS ====================
@app.post("/api/upload/csv")
async def upload_csv(file: UploadFile = File(...), mask: bool = Form(False),
selected_models: str = Form("")):
"""Upload and analyze CSV file"""
try:
validate_file_size(file)
content = await file.read()
df = pd.read_csv(io.BytesIO(content))
models = [m.strip() for m in selected_models.split(",") if m.strip()] or None
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample, models)
if mask:
masked_df = classifier.mask_dataframe(df.head(50), models)
response["data"] = masked_df.fillna("").to_dict(orient="records")
else:
highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/upload/txt")
async def upload_txt(file: UploadFile = File(...), mask: bool = Form(False),
selected_models: str = Form("")):
"""Upload and analyze TXT file"""
try:
validate_file_size(file)
content = await file.read()
text = content.decode("utf-8")
df = pd.DataFrame({"Content": [text]})
models = [m.strip() for m in selected_models.split(",") if m.strip()] or None
text_sample = text[:2000]
response = format_pii_response(df, df, text_sample, models)
if mask:
masked_df = classifier.mask_dataframe(df, models)
response["data"] = masked_df.fillna("").to_dict(orient="records")
else:
highlighted_df = classifier.scan_dataframe_with_html(df, models)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/upload/json")
async def upload_json(file: UploadFile = File(...), mask: bool = Form(False),
selected_models: str = Form("")):
"""Upload and analyze JSON file"""
try:
validate_file_size(file)
df = classifier.get_json_data(file.file)
models = [m.strip() for m in selected_models.split(",") if m.strip()] or None
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample, models)
if mask:
masked_df = classifier.mask_dataframe(df.head(50), models)
response["data"] = masked_df.fillna("").to_dict(orient="records")
else:
highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/upload/parquet")
async def upload_parquet(file: UploadFile = File(...), mask: bool = Form(False),
selected_models: str = Form("")):
"""Upload and analyze Parquet file"""
try:
validate_file_size(file)
content = await file.read()
df = classifier.get_parquet_data(content)
models = [m.strip() for m in selected_models.split(",") if m.strip()] or None
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample, models)
if mask:
masked_df = classifier.mask_dataframe(df.head(50), models)
response["data"] = masked_df.fillna("").to_dict(orient="records")
else:
highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/upload/avro")
async def upload_avro(file: UploadFile = File(...), mask: bool = Form(False),
selected_models: str = Form("")):
"""Upload and analyze Apache Avro file"""
try:
validate_file_size(file)
content = await file.read()
df = classifier.get_avro_data(content)
models = [m.strip() for m in selected_models.split(",") if m.strip()] or None
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample, models)
if mask:
masked_df = classifier.mask_dataframe(df.head(50), models)
response["data"] = masked_df.fillna("").to_dict(orient="records")
else:
highlighted_df = classifier.scan_dataframe_with_html(df.head(50), models)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/upload/pdf")
async def upload_pdf(file: UploadFile = File(...), page_number: int = Form(0)):
"""Upload and analyze PDF file (with pagination)"""
try:
validate_file_size(file)
content = await file.read()
# Get total pages and extract text from specific page
total_pages = classifier.get_pdf_total_pages(content)
text = classifier.get_pdf_page_text(content, page_number)
# Format PII response
response = format_pii_response(None, None, text)
response["total_pages"] = total_pages
response["current_page"] = page_number
# Get labeled PDF image
img = classifier.get_labeled_pdf_image(content, page_number)
if img:
import base64
from PIL import Image
# Check if img is already bytes or a PIL Image
if isinstance(img, bytes):
# Already bytes, just encode
img_str = base64.b64encode(img).decode()
elif isinstance(img, Image.Image):
# PIL Image, need to convert to bytes
buffered = io.BytesIO()
img.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode()
else:
# Unknown type, skip image
img_str = None
if img_str:
response["image"] = f"data:image/png;base64,{img_str}"
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/upload/image")
async def upload_image(file: UploadFile = File(...), mask: bool = Form(False)):
"""Upload and analyze image with OCR"""
try:
validate_file_size(file)
content = await file.read()
# Extract text via OCR
text = classifier.get_ocr_text_from_image(content)
if not text:
raise HTTPException(status_code=400, detail="No text could be extracted from the image")
df = pd.DataFrame({"Content": [text]})
response = format_pii_response(df, df, text)
if mask:
masked_df = classifier.mask_dataframe(df)
response["data"] = masked_df.fillna("").to_dict(orient="records")
else:
highlighted_df = classifier.scan_dataframe_with_html(df)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
# Return original image as base64
import base64
img_str = base64.b64encode(content).decode()
response["original_image"] = f"data:image/png;base64,{img_str}"
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# ==================== ANALYSIS ENDPOINTS ====================
@app.post("/api/analyze/text")
async def analyze_text(request: TextAnalysisRequest):
"""Analyze plain text for PII"""
try:
matches = classifier.analyze_text_hybrid(request.text)
count_df = classifier.get_pii_counts(request.text)
return JSONResponse(content={
"matches": matches,
"pii_counts": count_df.fillna("").to_dict(orient="records") if not count_df.empty else [],
"total_pii_found": len(matches)
})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/inspect")
async def inspect_text(request: TextAnalysisRequest):
"""Run full model inspection on text"""
try:
inspector_df = classifier.run_full_inspection(request.text)
if inspector_df.empty:
return JSONResponse(content={
"inspector": [],
"message": "No PII detected by any model"
})
return JSONResponse(content={
"inspector": inspector_df.fillna("").to_dict(orient="records")
})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/mask")
async def mask_text(request: TextAnalysisRequest):
"""Mask PII in text"""
try:
df = pd.DataFrame({"Content": [request.text]})
masked_df = classifier.mask_dataframe(df)
return JSONResponse(content={
"original": request.text,
"masked": masked_df.iloc[0]["Content"]
})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# ==================== PATTERN MANAGEMENT ====================
@app.get("/api/patterns")
async def get_patterns():
"""Get all regex patterns"""
try:
patterns = classifier.list_patterns()
return JSONResponse(content={
"patterns": [{"name": k, "regex": v} for k, v in patterns.items()]
})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/patterns")
async def add_pattern(request: PatternAddRequest):
"""Add a new regex pattern"""
try:
classifier.add_pattern(request.name, request.regex)
return JSONResponse(content={
"message": f"Pattern '{request.name}' added successfully",
"pattern": {"name": request.name, "regex": request.regex}
})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.delete("/api/patterns/{pattern_name}")
async def delete_pattern(pattern_name: str):
"""Remove a regex pattern"""
try:
classifier.remove_pattern(pattern_name)
return JSONResponse(content={
"message": f"Pattern '{pattern_name}' removed successfully"
})
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# ==================== DATABASE CONNECTORS ====================
@app.post("/api/connect/postgresql")
async def connect_postgresql(request: DatabaseConnectionRequest, authorization: str = Header(None)):
"""Connect to PostgreSQL, scan table, persist to Supabase, return analysis."""
uid = _uid_from_auth(authorization)
try:
limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100
df = classifier.get_postgres_data(
request.host, request.port, request.database,
request.user, request.password, request.table, limit=limit
)
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# If hybrid mode, fetch metadata and append to response
if request.scan_mode == "metadata_and_sampling":
from classifier_manager.metadata_scanner import MetadataScanner
scanner = MetadataScanner()
cols_meta = classifier.pg_handler.get_schema_metadata(
request.host, request.port, request.database,
request.user, request.password, request.table
)
cat_resp = scanner.scan_schema(cols_meta)
response["metadata"] = cat_resp.get("metadata", {})
# ── Persist to Supabase universal schema ──────────────────────────
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="postgresql",
db_name=request.database,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned,
scan_mode=request.scan_mode
)
# ─────────────────────────────────────────────────────────────────
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"PostgreSQL connection failed: {str(e)}")
@app.post("/api/connect/mysql")
async def connect_mysql(request: DatabaseConnectionRequest, authorization: str = Header(None)):
"""Connect to MySQL, scan table, persist to Supabase, return analysis."""
uid = _uid_from_auth(authorization)
try:
limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100
df = classifier.get_mysql_data(
request.host, request.port, request.database,
request.user, request.password, request.table, limit=limit
)
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned # additive: row count after 100-row cap
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# If hybrid mode, fetch metadata and append to response
if request.scan_mode == "metadata_and_sampling":
from classifier_manager.metadata_scanner import MetadataScanner
scanner = MetadataScanner()
cols_meta = classifier.mysql_handler.get_schema_metadata(
request.host, request.port, request.database,
request.user, request.password, request.table
)
cat_resp = scanner.scan_schema(cols_meta)
response["metadata"] = cat_resp.get("metadata", {})
# ── Persist to Supabase universal schema ──────────────────────────
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="mysql",
db_name=request.database,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned,
scan_mode=request.scan_mode
)
# ─────────────────────────────────────────────────────────────────
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"MySQL connection failed: {str(e)}")
@app.post("/api/connect/postgresql/list-tables")
async def list_postgresql_tables(request: DatabaseConnectionRequest):
"""List all tables in a PostgreSQL database (via SchemaInspector)."""
try:
engine = _schema_inspector.build_engine(
"postgresql", request.host, request.port,
request.database, request.user, request.password
)
tables = _schema_inspector.get_tables(engine)
return JSONResponse(content={"tables": tables})
except Exception as e:
raise HTTPException(status_code=500, detail=f"PostgreSQL list-tables failed: {str(e)}")
@app.post("/api/connect/mysql/list-tables")
async def list_mysql_tables_endpoint(request: DatabaseConnectionRequest):
"""List all tables in a MySQL database (via SchemaInspector)."""
try:
engine = _schema_inspector.build_engine(
"mysql", request.host, request.port,
request.database, request.user, request.password
)
tables = _schema_inspector.get_tables(engine)
return JSONResponse(content={"tables": tables})
except Exception as e:
raise HTTPException(status_code=500, detail=f"MySQL list-tables failed: {str(e)}")
@app.post("/api/connect/mongodb")
async def connect_mongodb(request: DatabaseConnectionRequest, authorization: str = Header(None)):
"""Connect to MongoDB, scan collection, persist to Supabase, return analysis."""
uid = _uid_from_auth(authorization)
try:
limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100
df = classifier.get_mongodb_data(
request.host, request.port, request.database,
request.user, request.password, request.table, limit=limit
)
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# ── Persist to Supabase universal schema ──────────────────────────
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="mongodb",
db_name=request.database,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned,
scan_mode=request.scan_mode
)
# ─────────────────────────────────────────────────────────────────
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"MongoDB connection failed: {str(e)}")
@app.post("/api/connect/mongodb/list-tables")
async def list_mongodb_collections_endpoint(request: DatabaseConnectionRequest):
"""List all collections in a MongoDB database."""
try:
collections = classifier.list_mongodb_collections(
request.host, request.port, request.database,
request.user, request.password
)
return JSONResponse(content={"tables": collections})
except Exception as e:
raise HTTPException(status_code=500, detail=f"MongoDB list-collections failed: {str(e)}")
@app.post("/api/connect/mariadb")
async def connect_mariadb(request: DatabaseConnectionRequest, authorization: str = Header(None)):
"""Connect to MariaDB, scan table, persist to Supabase, return analysis."""
uid = _uid_from_auth(authorization)
try:
limit = 5 if request.scan_mode in ("sampling", "metadata_and_sampling") else 100
df = classifier.get_mariadb_data(
request.host, request.port, request.database,
request.user, request.password, request.table, limit=limit
)
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
if request.scan_mode == "metadata_and_sampling":
from classifier_manager.metadata_scanner import MetadataScanner
scanner = MetadataScanner()
cols_meta = classifier.mariadb_handler.get_schema_metadata(
request.host, request.port, request.database,
request.user, request.password, request.table
)
cat_resp = scanner.scan_schema(cols_meta)
response["metadata"] = cat_resp.get("metadata", {})
# ── Persist to Supabase universal schema ──────────────────────────
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="mariadb",
db_name=request.database,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned,
scan_mode=request.scan_mode
)
# ─────────────────────────────────────────────────────────────────
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"MariaDB connection failed: {str(e)}")
@app.post("/api/connect/mariadb/list-tables")
async def list_mariadb_tables_endpoint(request: DatabaseConnectionRequest):
"""List all tables in a MariaDB database (via SchemaInspector)."""
try:
engine = _schema_inspector.build_engine(
"mariadb", request.host, request.port,
request.database, request.user, request.password
)
tables = _schema_inspector.get_tables(engine)
return JSONResponse(content={"tables": tables})
except Exception as e:
raise HTTPException(status_code=500, detail=f"MariaDB list-tables failed: {str(e)}")
@app.post("/api/connect/slack")
async def connect_slack(request: SlackRequest, authorization: str = Header(None)):
"""Connect to Slack, fetch channel messages, scan, persist to Supabase."""
uid = _uid_from_auth(authorization)
try:
df = classifier.get_slack_messages(request.token, request.channel_id)
if df.empty:
raise ValueError("No messages found or invalid channel/token.")
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# Supabase Persistence
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="slack",
db_name=request.channel_id,
table_name="Messages",
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Slack connection failed: {str(e)}")
@app.post("/api/connect/gmail")
async def connect_gmail(request: GmailRequest, authorization: str = Header(None)):
"""Connect to Gmail, fetch messages, scan, persist to Supabase."""
from api import _uid_from_auth
uid = _uid_from_auth(authorization)
if not uid:
raise HTTPException(status_code=401, detail="Authentication required to scan")
try:
df = classifier.get_gmail_data(request.credentials, request.auth_type, num_emails=10)
if df.empty:
raise ValueError("No emails found or invalid credentials.")
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# Supabase Persistence
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="gmail",
db_name="Gmail Inbox",
table_name="Emails",
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Gmail connection failed: {str(e)}")
@app.post("/api/connect/zendesk")
async def connect_zendesk(request: ZendeskRequest, authorization: str = Header(None)):
"""Connect to Zendesk, fetch tickets, scan, persist to Supabase."""
from api import _uid_from_auth
uid = _uid_from_auth(authorization)
if not uid:
raise HTTPException(status_code=401, detail="Authentication required to scan")
try:
df = classifier.get_zendesk_tickets(request.subdomain, request.email, request.api_token)
if df.empty:
raise ValueError("No tickets found or invalid credentials.")
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# Supabase Persistence
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="zendesk",
db_name=f"{request.subdomain}.zendesk.com",
table_name="Tickets",
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Zendesk connection failed: {str(e)}")
@app.post("/api/connect/salesforce")
async def connect_salesforce(request: SalesforceRequest, authorization: str = Header(None)):
"""Connect to Salesforce, fetch records, scan, persist to Supabase."""
from api import _uid_from_auth
uid = _uid_from_auth(authorization)
if not uid:
raise HTTPException(status_code=401, detail="Authentication required to scan")
try:
df = classifier.get_salesforce_records(request.instance_url, request.access_token)
if df.empty:
raise ValueError("No records found or invalid credentials.")
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# Supabase Persistence
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="salesforce",
db_name=request.instance_url.replace("https://", ""),
table_name="Records",
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Salesforce connection failed: {str(e)}")
@app.post("/api/connect/glue/list-databases")
async def list_glue_databases(request: GlueCredentials):
try:
databases = classifier.list_glue_databases(request.access_key, request.secret_key, request.region)
return {"databases": databases}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/connect/glue/list-tables")
async def list_glue_tables(request: GlueListTablesRequest):
try:
tables = classifier.list_glue_tables(request.access_key, request.secret_key, request.region, request.database_name)
return {"tables": tables}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/connect/glue")
async def connect_glue(request: GlueScanRequest, authorization: str = Header(None)):
"""Connect to AWS Glue, fetch table schema, scan for PII, and persist metadata."""
from api import _uid_from_auth
uid = _uid_from_auth(authorization)
if not uid:
raise HTTPException(status_code=401, detail="Authentication required to scan")
try:
df = classifier.get_glue_table_schema(request.access_key, request.secret_key, request.region, request.database_name, request.table_name)
if df.empty:
raise ValueError("No schema found or invalid credentials.")
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# Supabase Persistence
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="glue",
db_name=request.database_name,
table_name=request.table_name,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"AWS Glue connection failed: {str(e)}")
@app.post("/api/connect/aws-rds")
async def connect_aws_rds(request: AwsRdsConnectionRequest, authorization: str = Header(None)):
"""Connect to AWS RDS, scan table, persist to Supabase, return analysis."""
from api import _uid_from_auth
uid = _uid_from_auth(authorization)
if not uid:
raise HTTPException(status_code=401, detail="Authentication required to scan")
try:
df = classifier.get_aws_rds_data(
request.engine, request.host, request.port, request.database,
request.user, request.password, request.table
)
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# ── Persist to Supabase universal schema ──────────────────────────
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="aws-rds",
db_name=request.database,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
# ─────────────────────────────────────────────────────────────────
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"AWS RDS connection failed: {str(e)}")
@app.post("/api/connect/dynamodb")
async def connect_dynamodb(request: DynamoDbConnectionRequest, authorization: str = Header(None)):
"""Connect to DynamoDB, scan table, persist to Supabase, return analysis."""
uid = _uid_from_auth(authorization)
try:
df = classifier.get_dynamodb_data(
request.region, request.access_key, request.secret_key, request.table
)
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
# Bug #8 fix: rows_scanned was undefined — define it before persist call
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# ── Persist to Supabase universal schema ──────────────────────────
from db.supabase_client import persist_scan_results
persist_scan_results(
uid=uid,
connector_type="dynamodb",
db_name=request.region,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned
)
# ─────────────────────────────────────────────────────────────────
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"DynamoDB connection failed: {str(e)}")
@app.post("/api/connect/aws-rds/list-tables")
async def connect_aws_rds_list_tables(request: AwsRdsConnectionRequest):
try:
tables = classifier.list_aws_rds_tables(
request.engine, request.host, request.port, request.database,
request.user, request.password
)
return JSONResponse(content={"tables": tables})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to list tables: {str(e)}")
@app.post("/api/connect/dynamodb/list-tables")
async def connect_dynamodb_list_tables(request: DynamoDbConnectionRequest):
try:
tables = classifier.list_dynamodb_tables(
request.region, request.access_key, request.secret_key
)
return JSONResponse(content={"tables": tables})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to list tables: {str(e)}")
@app.post("/api/connectors/metadata-scan")
async def connect_metadata_scan(request: MetadataScanRequest, authorization: str = Header(None)):
"""
Connect to a DB (Postgres/MySQL), run a metadata-only scan via MetadataScanner,
and persist results to Supabase returning the CatalogResponse shape.
"""
uid = _uid_from_auth(authorization)
try:
from classifier_manager.metadata_scanner import MetadataScanner
from db.supabase_client import (
create_scan_session, finish_scan_session,
upsert_db_registry_entry, upsert_db_scan_state,
get_file_catalog, get_scan_sessions
)
scanner = MetadataScanner()
db_name = request.database
table_name = request.table
connector_type = request.connector_type.lower()
# Postgres schema defaults to public for deterministic file_ids
schema_name = "public" if connector_type in ["postgresql", "aws-rds"] else ""
# 1. Fetch schema metadata via unified SchemaInspector (TICKET-5)
# DynamoDB has no SQL schema — falls back to native handler
try:
if connector_type in ("postgresql", "mysql", "mariadb", "aws-rds"):
rds_engine = getattr(request, "engine", None)
engine = _schema_inspector.build_engine(
connector_type, request.host, request.port, db_name,
request.user, request.password,
**({} if not rds_engine else {"rds_engine": rds_engine})
)
raw_cols = _schema_inspector.get_columns(engine, table_name)
# MetadataScanner expects list[dict] with column_name / data_type keys
columns_metadata = raw_cols
elif connector_type == "dynamodb":
db_name = request.region
columns_metadata = classifier.dynamodb_handler.get_schema_metadata(
request.region, request.access_key, request.secret_key, table_name
)
else:
raise HTTPException(
status_code=400,
detail=f"Unsupported connector_type for metadata scan: {connector_type}"
)
except HTTPException:
raise
except Exception as e:
err_str = str(e)
if "missing read permissions" in err_str.lower():
raise HTTPException(status_code=403, detail=err_str)
raise HTTPException(status_code=500, detail=err_str)
# 2. Run MetadataScanner
scan_results = scanner.scan_schema(columns_metadata)
flagged_columns = scan_results.get("flagged_columns", [])
# 3. Generate LLM Explanations for metadata scan
if flagged_columns and classifier:
rows = [{"PII Type": c["matched_rule"], "matched_rule": c["matched_rule"], "contributing_models": ["MetadataScanner"]} for c in flagged_columns]
classifier.generate_explanations(rows)
for c, r in zip(flagged_columns, rows):
c["llm_explanation"] = r.get("llm_explanation", "")
# 4. Persist to Supabase
try:
session_id = create_scan_session(uid, connector_type, "METADATA_ONLY")
# Upsert Folder row
folder_metadata = {"scan_mode": "metadata_only"}
upsert_db_registry_entry(uid, connector_type, db_name, "", None, True, folder_metadata)
# Upsert Table row
if table_name:
table_classification = "SENSITIVE" if len(flagged_columns) > 0 else "NON-SENSITIVE"
upsert_db_registry_entry(uid, connector_type, db_name, schema_name, table_name, False, scan_results)
upsert_db_scan_state(uid, connector_type, db_name, schema_name, table_name,
table_classification, "METADATA_ONLY", session_id)
# Update folder scan state
upsert_db_scan_state(uid, connector_type, db_name, "", None,
table_classification, "METADATA_ONLY", session_id)
total_pii = len(flagged_columns)
finish_scan_session(session_id, 1, total_pii)
except Exception as persist_err:
print(f"[WARN] Metadata scan persistence failed: {persist_err}")
# 4. Return standard CatalogResponse format
files = get_file_catalog(uid, connector_type)
sessions = get_scan_sessions(uid, connector_type)
for f in files:
f["connector_type"] = connector_type
files.sort(key=lambda f: (not f.get("is_folder", False), f.get("file_name", "").lower()))
sessions.sort(key=lambda s: s.get("triggered_at", ""), reverse=True)
completed_sessions = [s for s in sessions if s.get("status") == "completed"]
last_session = completed_sessions[0] if completed_sessions else None
response_content = jsonable_encoder({"files": files, "last_session": last_session, "sessions": sessions})
return JSONResponse(content=response_content)
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=f"Metadata scan failed: {str(e)}")
# ==================== CLOUD STORAGE - AWS S3 ====================
@app.post("/api/cloud/s3/list-buckets")
async def list_s3_buckets(request: S3ConnectionRequest):
"""List S3 buckets"""
try:
buckets = classifier.get_s3_buckets(request.access_key, request.secret_key, request.region)
return JSONResponse(content={"buckets": buckets})
except Exception as e:
raise HTTPException(status_code=500, detail=f"S3 connection failed: {str(e)}")
@app.post("/api/cloud/s3/list-files")
async def list_s3_files(request: S3ConnectionRequest):
"""List files in S3 bucket"""
try:
if not request.bucket:
raise HTTPException(status_code=400, detail="Bucket name is required")
files = classifier.get_s3_files(
request.access_key, request.secret_key, request.region, request.bucket
)
return JSONResponse(content={"files": files})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to list S3 files: {str(e)}")
@app.post("/api/cloud/s3/scan")
async def scan_s3_file(request: S3ConnectionRequest):
"""Download and scan S3 file"""
try:
if not request.bucket or not request.file_key:
raise HTTPException(status_code=400, detail="Bucket and file_key are required")
content = classifier.download_s3_file(
request.access_key, request.secret_key, request.region,
request.bucket, request.file_key
)
# Assume CSV for now
df = pd.read_csv(io.BytesIO(content))
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"S3 scan failed: {str(e)}")
# ==================== CLOUD STORAGE - AZURE ====================
@app.post("/api/cloud/azure/list-containers")
async def list_azure_containers(request: AzureConnectionRequest):
"""List Azure containers"""
try:
containers = classifier.get_azure_containers(request.connection_string)
return JSONResponse(content={"containers": containers})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Azure connection failed: {str(e)}")
@app.post("/api/cloud/azure/list-blobs")
async def list_azure_blobs(request: AzureConnectionRequest):
"""List blobs in Azure container"""
try:
if not request.container:
raise HTTPException(status_code=400, detail="Container name is required")
blobs = classifier.get_azure_blobs(request.connection_string, request.container)
return JSONResponse(content={"blobs": blobs})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to list blobs: {str(e)}")
@app.post("/api/cloud/azure/scan")
async def scan_azure_blob(request: AzureConnectionRequest):
"""Download and scan Azure blob"""
try:
if not request.container or not request.blob:
raise HTTPException(status_code=400, detail="Container and blob are required")
content = classifier.download_azure_blob(
request.connection_string, request.container, request.blob
)
df = pd.read_csv(io.BytesIO(content))
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Azure scan failed: {str(e)}")
# ==================== CLOUD STORAGE - GCS ====================
@app.post("/api/cloud/gcs/list-buckets")
async def list_gcs_buckets(request: GCSConnectionRequest):
"""List GCS buckets"""
try:
buckets = classifier.get_gcs_buckets(request.credentials)
return JSONResponse(content={"buckets": buckets})
except Exception as e:
raise HTTPException(status_code=500, detail=f"GCS connection failed: {str(e)}")
@app.post("/api/cloud/gcs/list-files")
async def list_gcs_files(request: GCSConnectionRequest):
"""List files in GCS bucket"""
try:
if not request.bucket:
raise HTTPException(status_code=400, detail="Bucket name is required")
files = classifier.get_gcs_files(request.credentials, request.bucket)
return JSONResponse(content={"files": files})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Failed to list GCS files: {str(e)}")
@app.post("/api/cloud/gcs/scan")
async def scan_gcs_file(request: GCSConnectionRequest):
"""Download and scan GCS file"""
try:
if not request.bucket or not request.file_name:
raise HTTPException(status_code=400, detail="Bucket and file_name are required")
content = classifier.download_gcs_file(
request.credentials, request.bucket, request.file_name
)
df = pd.read_csv(io.BytesIO(content))
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"GCS scan failed: {str(e)}")
# ==================== CLOUD STORAGE - GOOGLE DRIVE ====================
@app.post("/api/cloud/drive/list-files")
async def list_drive_files(request: GoogleDriveRequest):
"""List Google Drive files"""
try:
files = classifier.get_google_drive_files(request.credentials)
return JSONResponse(content={"files": files})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Google Drive connection failed: {str(e)}")
@app.post("/api/cloud/drive/scan")
async def scan_drive_file(request: GoogleDriveRequest):
"""Download and scan Google Drive file"""
try:
if not request.file_id or not request.mime_type:
raise HTTPException(status_code=400, detail="file_id and mime_type are required")
content = classifier.download_drive_file(
request.file_id, request.mime_type, request.credentials
)
if isinstance(content, bytes):
try:
text = content.decode('utf-8')
df = pd.DataFrame({"Content": [text]})
response = format_pii_response(df, df, text)
highlighted_df = classifier.scan_dataframe_with_html(df)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except:
raise HTTPException(status_code=400, detail="Binary file cannot be processed")
except Exception as e:
raise HTTPException(status_code=500, detail=f"Google Drive scan failed: {str(e)}")
# ==================== ENTERPRISE CONNECTORS ====================
@app.post("/api/enterprise/gmail")
async def scan_gmail(file: UploadFile = File(...), num_emails: int = Form(10)):
"""Scan Gmail messages"""
try:
df = classifier.get_gmail_data(file.file, num_emails)
if df.empty:
raise HTTPException(status_code=400, detail="No emails fetched")
text_sample = df.iloc[0]['Content']
response = format_pii_response(df, df, text_sample)
masked_df = classifier.mask_dataframe(df)
response["data"] = masked_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Gmail scan failed: {str(e)}")
@app.post("/api/enterprise/slack")
async def scan_slack(request: SlackRequest):
"""Scan Slack messages"""
try:
df = classifier.get_slack_messages(request.token, request.channel_id)
if df.empty:
raise HTTPException(status_code=400, detail="No messages found or authentication failed")
text_sample = df.iloc[0]['Content']
response = format_pii_response(df, df, text_sample)
masked_df = classifier.mask_dataframe(df)
response["data"] = masked_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Slack scan failed: {str(e)}")
@app.post("/api/enterprise/confluence")
async def scan_confluence(request: ConfluenceRequest):
"""Scan Confluence page"""
try:
df = classifier.get_confluence_page(
request.url, request.username, request.token, request.page_id
)
if df.empty:
raise HTTPException(status_code=400, detail="Failed to fetch page")
text_sample = df.iloc[0]['Content']
response = format_pii_response(df, df, text_sample)
highlighted_df = classifier.scan_dataframe_with_html(df)
response["data"] = highlighted_df.fillna("").to_dict(orient="records")
return JSONResponse(content=response)
except Exception as e:
raise HTTPException(status_code=500, detail=f"Confluence scan failed: {str(e)}")
# ==================== EMAIL FUNCTIONALITY ====================
@app.post("/api/send-welcome")
async def send_welcome(request: WelcomeEmailRequest):
"""
Send a welcome email to a new user.
This endpoint is called by the frontend after a user submits the contact form.
"""
try:
# Send the welcome email
success = send_welcome_email(request.name, request.email)
if success:
return JSONResponse(content={
"success": True,
"message": f"Welcome email sent to {request.email}"
})
else:
raise HTTPException(
status_code=500,
detail="Failed to send welcome email. SMTP configuration may be missing."
)
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"Email sending failed: {str(e)}"
)
# ==================== EVALUATOR ENDPOINTS ====================
# Full model catalogue served to the MODEL LAB frontend
EVALUATOR_MODEL_CATALOGUE = [
{"key": "regex", "label": "Regex Engine", "hf_id": "deterministic", "type": "Rule-based", "params": "—", "f1_benchmark": 1.0, "lazy": False, "description": "Deterministic regex patterns for EMAIL, PHONE, SSN, CC, IP, URL, MAC."},
{"key": "nltk", "label": "NLTK Chunker", "hf_id": "nltk", "type": "Statistical", "params": "—", "f1_benchmark": 0.0, "lazy": False, "description": "NLTK ne_chunk: PERSON → FIRST_NAME, GPE → LOCATION."},
{"key": "spacy", "label": "SpaCy LG", "hf_id": "en_core_web_lg", "type": "Statistical", "params": "685M", "f1_benchmark": 0.0, "lazy": False, "description": "SpaCy en_core_web_lg NER model."},
{"key": "presidio", "label": "MS Presidio", "hf_id": "microsoft/presidio-analyzer", "type": "Rule+ML", "params": "—", "f1_benchmark": 0.0, "lazy": False, "description": "Microsoft Presidio enterprise PII analyzer."},
{"key": "gliner", "label": "GLiNER Small", "hf_id": "urchade/gliner_small-v2.1", "type": "GLiNER", "params": "small", "f1_benchmark": 0.850, "lazy": False, "description": "Zero-shot GLiNER small model."},
{"key": "deberta", "label": "DeBERTa PII", "hf_id": "lakshyakh93/deberta-large-finetuned-pii", "type": "NER", "params": "86M", "f1_benchmark": 0.920, "lazy": False, "description": "Kaggle-winning DeBERTa V3 fine-tuned for PII."},
{"key": "pasteproof", "label": "Pasteproof v2", "hf_id": "joneauxedgar/pasteproof-pii-detector-v2", "type": "NER", "params": "149M", "f1_benchmark": 0.970, "lazy": True, "description": "ModernBERT 149M fine-tuned PII detector."},
{"key": "piiranha", "label": "Piiranha v1", "hf_id": "iiiorg/piiranha-v1-detect-personal-information", "type": "NER", "params": "86M", "f1_benchmark": 0.931, "lazy": True, "description": "DeBERTa-based PII detector."},
{"key": "nvidia_gliner","label": "NVIDIA GLiNER", "hf_id": "nvidia/gliner-PII-0.1", "type": "GLiNER", "params": "570M", "f1_benchmark": 0.870, "lazy": True, "description": "NVIDIA GLiNER with 37-label PII vocabulary."},
{"key": "mmbert", "label": "mmbert32k", "hf_id": "llm-semantic-router/mmbert32k-pii-detector-merged", "type": "NER", "params": "307M", "f1_benchmark": 0.969, "lazy": True, "description": "ModernBERT 32k-context PII detector."},
{"key": "nerguard", "label": "NerGuard-0.3B", "hf_id": "exdsgift/NerGuard-0.3B", "type": "NER", "params": "300M", "f1_benchmark": 0.996, "lazy": True, "description": "mDeBERTa 300M — highest F1 in the registry."},
{"key": "gliner_large", "label": "GLiNER PII Large", "hf_id": "knowledgator/gliner-pii-large-v1.0", "type": "GLiNER", "params": "large","f1_benchmark": 0.833, "lazy": True, "description": "GLiNER large architecture fine-tuned for PII."},
]
class EvaluatorScanRequest(BaseModel):
text: str = Field(..., description="Document text to scan")
gt_spans: List[Dict[str, Any]] = Field(default=[], description="Ground truth spans from /api/evaluator/parse")
model_keys: List[str] = Field(default=["regex", "spacy", "deberta"], description="Model keys to run")
conf_threshold: float = Field(default=0.5, ge=0.0, le=1.0)
entropy_threshold: float = Field(default=4.5, ge=3.0, le=6.0)
@app.get("/api/evaluator/models")
async def evaluator_models():
"""Return the full model catalogue for the MODEL LAB page."""
return JSONResponse(content={"models": EVALUATOR_MODEL_CATALOGUE})
@app.post("/api/evaluator/parse")
async def evaluator_parse(
file: UploadFile = File(...),
format: str = Form("auto"),
doc_index: int = Form(0),
schema: str = Form(""),
):
"""
Parse a labeled (or unlabeled) dataset file.
Returns the document text, ground-truth spans, format detected, and doc count.
format: 'auto' | 'bigcode' | 'nemotron' | 'csv_spans' | 'json_spans' | 'unlabeled'
schema: JSON string e.g. '{"text_col":"text","spans_col":"spans"}'
"""
try:
content = await file.read()
schema_dict = json.loads(schema) if schema.strip() else None
text, gt_spans, has_gt, doc_count, fmt_detected = detect_and_parse(
content, file.filename or "", doc_index, schema_dict, format
)
return JSONResponse(content={
"text": text,
"gt_spans": gt_spans,
"has_gt": has_gt,
"format_detected": fmt_detected,
"doc_count": doc_count,
"char_count": len(text),
})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Parse error: {str(e)}")
@app.post("/api/evaluator/scan")
async def evaluator_scan(request: EvaluatorScanRequest):
"""
Run selected models on text in parallel, compare against GT spans.
Returns partial results — if a model times out or errors, other models
still return results. Each model entry includes timed_out + error fields.
"""
try:
t0 = time.time()
has_gt = bool(request.gt_spans)
# Parallel scan — new shape: {key: {detections, error, timed_out}}
raw_results = await asyncio.get_event_loop().run_in_executor(
None,
lambda: classifier.scan_with_models(request.text, request.model_keys),
)
per_model: Dict[str, Any] = {}
for model_key, payload in raw_results.items():
predictions = payload["detections"]
timed_out = payload["timed_out"]
model_error = payload["error"]
# Normalise canonical label
for p in predictions:
p["canonical"] = norm_model_out(p.get("label", ""))
if has_gt and not timed_out:
comparison = compare_spans(predictions, request.gt_spans, model_key)
metrics = compute_metrics(comparison)
coverage = get_label_coverage(request.gt_spans, model_key)
failures = analyse_failures(comparison, request.text)
else:
comparison = {"TP": [], "FP": [], "FN": []}
metrics = []
coverage = {"in_scope": [], "out_of_scope": []}
failures = {"missed": [], "false_positives": []}
per_model[model_key] = {
"predictions": predictions,
"comparison": comparison,
"metrics": metrics,
"coverage": coverage,
"failures": failures,
"timed_out": timed_out,
"error": model_error,
}
elapsed = round(time.time() - t0, 2)
any_timeout = any(v["timed_out"] for v in per_model.values())
return JSONResponse(content={
"per_model": per_model,
"has_gt": has_gt,
"elapsed": elapsed,
"any_timeout": any_timeout,
})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Evaluator scan error: {str(e)}")
@app.post("/api/evaluator/scan/stream")
async def evaluator_scan_stream(request: EvaluatorScanRequest):
"""
SSE streaming version of /api/evaluator/scan.
Emits one Server-Sent Event per model as results arrive, so the frontend
can show live per-model progress rather than waiting for all models.
Event format (JSON):
data: {"model_key": ..., "done": false, "status": "running"}
data: {"model_key": ..., "done": true, "payload": {...}, "elapsed": ...}
data: {"done": true, "all_complete": true}
"""
import concurrent.futures
has_gt = bool(request.gt_spans)
loop = asyncio.get_event_loop()
async def event_generator():
t0 = time.time()
always_on = {"regex", "nltk", "spacy", "presidio", "gliner", "deberta"}
def _run_one_key(key: str):
payload = classifier.scan_with_models(request.text, [key])
return key, payload.get(key, {"detections": [], "error": "No result", "timed_out": False})
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
future_map = {pool.submit(_run_one_key, key): key for key in request.model_keys}
# Yield start event for each model
for key in request.model_keys:
yield f'data: {{"model_key": "{key}", "done": false, "status": "running"}}\n\n'
for f in concurrent.futures.as_completed(future_map, timeout=110):
try:
key, raw = await loop.run_in_executor(None, f.result)
predictions = raw["detections"]
for p in predictions:
p["canonical"] = norm_model_out(p.get("label", ""))
if has_gt and not raw["timed_out"]:
comparison = compare_spans(predictions, request.gt_spans, key)
metrics = compute_metrics(comparison)
coverage = get_label_coverage(request.gt_spans, key)
failures = analyse_failures(comparison, request.text)
else:
comparison = {"TP": [], "FP": [], "FN": []}
metrics = []; coverage = {"in_scope": [], "out_of_scope": []}; failures = {"missed": [], "false_positives": []}
result_payload = {
"predictions": predictions, "comparison": comparison,
"metrics": metrics, "coverage": coverage, "failures": failures,
"timed_out": raw["timed_out"], "error": raw["error"],
}
event = json.dumps({"model_key": key, "done": True, "payload": result_payload, "elapsed": round(time.time() - t0, 2)})
yield f'data: {event}\n\n'
except Exception as e:
err_event = json.dumps({"model_key": future_map[f], "done": True, "error": str(e), "timed_out": True})
yield f'data: {err_event}\n\n'
yield f'data: {{"done": true, "all_complete": true, "total_elapsed": {round(time.time() - t0, 2)}}}\n\n'
return StreamingResponse(
event_generator(),
media_type="text/event-stream",
headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
)
@app.post("/api/evaluator/batch")
async def evaluator_batch(
file: UploadFile = File(...),
format: str = Form("nemotron"),
n_docs: int = Form(50),
model_keys: str = Form("deberta"),
conf_threshold: float = Form(0.5),
entropy_threshold: float = Form(4.5),
):
"""
Batch evaluation over N documents from a labeled file.
Returns aggregate metrics + per-doc breakdown for each selected model.
model_keys: comma-separated string e.g. 'deberta,piiranha'
"""
try:
content = await file.read()
keys = [k.strip() for k in model_keys.split(",") if k.strip()]
n_docs = max(1, min(n_docs, 1000))
agg: Dict[str, Dict[str, int]] = {k: {"TP": 0, "FP": 0, "FN": 0} for k in keys}
per_doc: List[Dict[str, Any]] = []
for i in range(n_docs):
try:
text, gt_spans, has_gt, doc_count, _ = detect_and_parse(
content, file.filename or "", i, None, format
)
if i >= doc_count:
break
if not has_gt:
continue
model_results = classifier.scan_with_models(text, keys)
doc_entry: Dict[str, Any] = {"doc_index": i}
for key, predictions in model_results.items():
for p in predictions:
p["canonical"] = norm_model_out(p.get("label", ""))
comp = compare_spans(predictions, gt_spans, key)
tp, fp, fn = len(comp["TP"]), len(comp["FP"]), len(comp["FN"])
agg[key]["TP"] += tp
agg[key]["FP"] += fp
agg[key]["FN"] += fn
prec = tp / (tp + fp) if (tp + fp) else 0.0
rec = tp / (tp + fn) if (tp + fn) else 0.0
f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0
doc_entry[key] = {
"f1": round(f1, 4), "precision": round(prec, 4),
"recall": round(rec, 4), "tp": tp, "fp": fp, "fn": fn,
}
per_doc.append(doc_entry)
except Exception as doc_err:
per_doc.append({"doc_index": i, "error": str(doc_err)})
# Aggregate metrics per model
aggregate: Dict[str, Any] = {}
for key, c in agg.items():
tp, fp, fn = c["TP"], c["FP"], c["FN"]
p = tp / (tp + fp) if (tp + fp) else 0.0
r = tp / (tp + fn) if (tp + fn) else 0.0
f = 2 * p * r / (p + r) if (p + r) else 0.0
aggregate[key] = {
"f1": round(f, 4), "precision": round(p, 4),
"recall": round(r, 4), "tp": tp, "fp": fp, "fn": fn,
}
return JSONResponse(content={
"aggregate": aggregate,
"per_doc": per_doc,
"n_docs_evaluated": len(per_doc),
})
except Exception as e:
raise HTTPException(status_code=500, detail=f"Batch eval error: {str(e)}")
# ==================== AUTH ENDPOINTS ====================
from db.supabase_auth import auth_sign_up, auth_sign_in, auth_get_user, auth_sign_out
def _uid_from_auth(authorization: Optional[str]) -> str:
"""Extract and verify the Supabase GoTrue user_id from a Bearer token.
Raises HTTP 401 if the token is missing or invalid."""
if not authorization or not authorization.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
token = authorization.split(" ", 1)[1]
user_data = auth_get_user(access_token=token)
# auth_get_user returns { "id": "...", "email": "...", ... } or { "error": "..." }
if "error" in user_data or "id" not in user_data:
raise HTTPException(status_code=401, detail=user_data.get("error", "Invalid token"))
return user_data["id"]
@app.post("/api/auth/register")
async def register(req: RegisterRequest):
"""Register a new Sense user via Supabase GoTrue."""
# Validate minimum password length before hitting Supabase
if len(req.password) < 6:
raise HTTPException(status_code=422, detail="Password must be at least 6 characters")
# Delegate to GoTrue
result = auth_sign_up(email=req.email, password=req.password, full_name=req.name)
# Seed credits row — non-fatal if it fails
try:
if "id" in result:
from db.credits_manager import seed_user_credits
seed_user_credits(result["id"])
except Exception as e:
print(f"[WARN] Credit seeding failed: {e}")
# Send welcome email — failure must NOT block registration
try:
send_welcome_email(name=req.name, email=req.email)
except Exception:
pass
return JSONResponse(content=result)
@app.post("/api/auth/login")
async def login(req: LoginRequest):
"""Authenticate a Sense user via Supabase GoTrue."""
result = auth_sign_in(email=req.email, password=req.password)
# Restore credits if it's Sunday (idempotent) — non-fatal
try:
if "id" in result:
from db.credits_manager import maybe_restore_credits
maybe_restore_credits(result["id"])
except Exception as e:
print(f"[WARN] Credit restore check failed: {e}")
return JSONResponse(content=result)
@app.get("/api/auth/me")
async def get_me(authorization: str = Header(None)):
"""Return the authenticated user's profile. Requires Bearer token."""
if not authorization or not authorization.startswith("Bearer "):
raise HTTPException(status_code=401, detail="Missing or invalid Authorization header")
token = authorization.split(" ", 1)[1]
user = auth_get_user(access_token=token)
return JSONResponse(content=user)
@app.post("/api/auth/logout")
async def logout(authorization: str = Header(None)):
"""Sign out the user from Supabase GoTrue (server-side token invalidation)."""
if authorization and authorization.startswith("Bearer "):
token = authorization.split(" ", 1)[1]
auth_sign_out(access_token=token)
return JSONResponse(content={"message": "Signed out"})
# ==================== CREDIT ENDPOINTS ====================
@app.get("/api/auth/credits")
async def get_credits_endpoint(authorization: str = Header(None)):
"""Return the authenticated user's credit balance."""
user_id = _uid_from_auth(authorization)
from db.credits_manager import get_credits
data = get_credits(user_id)
return JSONResponse(content=data)
@app.post("/api/auth/credits/deduct")
async def deduct_credits_endpoint(authorization: str = Header(None), amount: int = 1):
"""Atomically deduct `amount` credits. Returns 402 when out of credits."""
user_id = _uid_from_auth(authorization)
from db.credits_manager import deduct_credit
result = deduct_credit(user_id, amount)
if not result["success"]:
return JSONResponse(status_code=402, content=result)
return JSONResponse(content=result)
@app.get("/api/auth/profile-stats")
async def get_profile_stats_endpoint(authorization: str = Header(None)):
"""Return aggregated scan statistics + credits for the profile dashboard."""
user_id = _uid_from_auth(authorization)
from db.credits_manager import get_profile_stats
stats = get_profile_stats(user_id)
return JSONResponse(content=stats)
# ==================== HEALTH CHECK ====================
@app.get("/api/debug/db")
async def debug_db():
"""Debug: show recent scan sessions and registry rows (uses unified schema)."""
try:
from db.supabase_client import _get_conn
conn = _get_conn()
cur = conn.cursor()
cur.execute(
"SELECT id, uid, connector_type, scan_type, scan_mode, status, "
"triggered_at, files_scanned, total_pii_found "
"FROM sense_scan_sessions ORDER BY triggered_at DESC LIMIT 10"
)
sessions = [dict(r) for r in cur.fetchall()]
cur.execute(
"SELECT file_id, file_name, connector_type, is_folder "
"FROM sense_file_registry ORDER BY first_seen_at DESC LIMIT 20"
)
registry = [dict(r) for r in cur.fetchall()]
cur.close()
conn.close()
# Stringify datetimes
for row in sessions + registry:
for k, v in row.items():
if hasattr(v, 'isoformat'):
row[k] = v.isoformat()
return {"sessions": sessions, "registry": registry}
except Exception as e:
return {"error": str(e)}
@app.get("/")
async def root():
"""API health check"""
return {
"message": "Segmento Sense API",
"status": "operational",
"version": "1.0.0"
}
@app.get("/health")
async def health_check():
"""Detailed health check"""
return {
"status": "healthy",
"classifiers": {
"regex": True,
"nltk": True,
"spacy": True,
"presidio": True,
"gliner": True,
"deberta": True
}
}
# ==================== CONNECTOR CATALOG ENDPOINTS ====================
@app.get("/api/connector/file-catalog")
async def get_file_catalog_endpoint(connector_type: str, authorization: str = Header(None)):
uid = _uid_from_auth(authorization)
try:
from db.supabase_client import get_file_catalog, get_scan_sessions
files = get_file_catalog(uid, connector_type)
sessions = get_scan_sessions(uid, connector_type)
completed_sessions = [s for s in sessions if s.get("status") == "completed"]
last_session = completed_sessions[0] if completed_sessions else None
return {"files": files, "last_session": last_session, "sessions": sessions}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/db/catalog")
async def get_db_catalog(
connector_type: Optional[str] = None,
db_name: Optional[str] = None,
authorization: str = Header(None)
):
"""
Unified catalog read endpoint (TICKET-5).
Returns sense_file_registry + sense_file_scan_state rows for given user.
Optional filters: connector_type, db_name.
All scan modes (FULL_LOAD, SAMPLING, METADATA_ONLY, etc.) appear here.
"""
uid = _uid_from_auth(authorization)
try:
from db.supabase_client import get_file_catalog, get_scan_sessions
DB_CONNECTOR_TYPES = ["postgresql", "mysql", "mongodb", "mariadb"]
types_to_query = [connector_type] if connector_type else DB_CONNECTOR_TYPES
all_files: list = []
all_sessions: list = []
for ct in types_to_query:
# Pass db_name filter to get_file_catalog (TICKET-5)
files = get_file_catalog(uid, ct, db_name=db_name)
sessions = get_scan_sessions(uid, ct)
for f in files:
f["connector_type"] = ct
all_files.extend(files)
all_sessions.extend(sessions)
all_files.sort(key=lambda f: (not f.get("is_folder", False), f.get("file_name", "").lower()))
all_sessions.sort(key=lambda s: s.get("triggered_at", ""), reverse=True)
completed_sessions = [s for s in all_sessions if s.get("status") == "completed"]
last_session = completed_sessions[0] if completed_sessions else None
return {"files": all_files, "last_session": last_session, "sessions": all_sessions}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/connector/scan-sessions")
async def get_scan_sessions_endpoint(connector_type: str, authorization: str = Header(None)):
uid = _uid_from_auth(authorization)
try:
from db.supabase_client import get_scan_sessions
sessions = get_scan_sessions(uid, connector_type)
return {"sessions": sessions}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
# NOTE: /api/scan/async, /api/scan/status/{job_id}, /api/catalog/{db_name}
# have been removed in TICKET-5.
# Use /api/connectors/metadata-scan (sync) + GET /db/catalog?db_name=... instead.
# ==================== INCREMENTAL SCAN ENDPOINTS (TICKET-8) ====================
@app.get("/api/connectors/watermark-check")
async def watermark_check(
connector_type: str,
db_name: str,
table_name: str,
host: str,
port: str,
user: str,
password: str,
authorization: str = Header(None)
):
"""
Check whether incremental scan is possible for a given table.
Returns:
{ supported: true, watermark_column: str, last_value: str | null }
{ supported: false, reason: str }
"""
uid = _uid_from_auth(authorization)
try:
ct = connector_type.lower()
if ct not in ("postgresql", "mysql", "mariadb", "aws-rds"):
return {"supported": False, "reason": f"{ct} does not support watermark-based incremental scan."}
engine = _schema_inspector.build_engine(ct, host, port, db_name, user, password)
wm_col = _schema_inspector.detect_watermark_column(engine, table_name)
if not wm_col:
return {
"supported": False,
"reason": "No suitable watermark column found (updated_at, created_at, or auto-increment PK required)."
}
# Build the canonical file_id to look up existing watermark
schema_name = "public" if ct in ("postgresql", "aws-rds") else ""
from db.supabase_client import _build_db_file_id, check_watermark
file_id = _build_db_file_id(db_name, schema_name, table_name)
_, last_value = check_watermark(uid, ct, file_id)
return {
"supported": True,
"watermark_column": wm_col,
"last_value": last_value, # None means first-time (no prior watermark)
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/connect/{db_type}/incremental")
async def incremental_scan(
db_type: str,
request: DatabaseConnectionRequest,
authorization: str = Header(None)
):
"""
Watermark-based incremental scan (TICKET-8).
1. Detect watermark column.
2. Read last stored watermark from sense_file_scan_state.
3. Fetch only rows WHERE watermark_col > last_value.
4. Run PII scan on new rows.
5. Persist results + advance watermark.
"""
uid = _uid_from_auth(authorization)
ct = db_type.lower()
if ct not in ("postgresql", "mysql", "mariadb"):
raise HTTPException(status_code=400, detail=f"Incremental scan not supported for: {ct}")
try:
from sqlalchemy import text
from db.supabase_client import (
persist_scan_results, check_watermark, update_watermark,
_build_db_file_id
)
schema_name = "public" if ct == "postgresql" else ""
file_id = _build_db_file_id(request.database, schema_name, request.table)
# 1. Detect watermark column
engine = _schema_inspector.build_engine(
ct, request.host, request.port, request.database, request.user, request.password
)
wm_col = _schema_inspector.detect_watermark_column(engine, request.table)
if not wm_col:
raise HTTPException(
status_code=400,
detail=(
f"Table '{request.table}' has no suitable watermark column "
"(needs updated_at, created_at, or auto-increment PK). "
"Use Full Load or Sampling scan instead."
)
)
# 2. Read last watermark
_, last_value = check_watermark(uid, ct, file_id)
# 3. Fetch incremental rows
with engine.connect() as conn:
if last_value is not None:
# Quoted identifiers prevent SQL injection on col/table names
query = text(
f'SELECT * FROM "{request.table}" '
f'WHERE "{wm_col}" > :last_val ORDER BY "{wm_col}" ASC LIMIT 500'
)
rows = conn.execute(query, {"last_val": last_value}).mappings().fetchall()
else:
# First incremental scan = full load (no prior watermark)
query = text(
f'SELECT * FROM "{request.table}" ORDER BY "{wm_col}" ASC LIMIT 500'
)
rows = conn.execute(query).mappings().fetchall()
if not rows:
return JSONResponse(content={
"message": "No new rows since last scan.",
"rows_scanned": 0,
"total_pii_found": 0,
"watermark_column": wm_col,
"last_value": last_value,
})
import pandas as pd
df = pd.DataFrame([dict(r) for r in rows])
text_sample = df.head(10).to_string()
response = format_pii_response(df, df, text_sample)
rows_scanned = len(df)
response["rows_scanned"] = rows_scanned
masked_df = classifier.mask_dataframe(df.head(50))
response["data"] = masked_df.fillna("").to_dict(orient="records")
# 4. Persist results
persist_scan_results(
uid=uid,
connector_type=ct,
db_name=request.database,
table_name=request.table,
response=response,
column_count=len(df.columns),
rows_scanned=rows_scanned,
scan_mode="incremental",
)
# 5. Advance watermark
new_wm = _schema_inspector.get_max_watermark_value(engine, request.table, wm_col)
if new_wm:
update_watermark(uid, ct, file_id, wm_col, new_wm)
response["watermark_column"] = wm_col
response["last_value"] = new_wm
return JSONResponse(content=response)
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=f"Incremental scan failed: {str(e)}")
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860) # HuggingFace Spaces default port