src/api/schemas.py

from pydantic import BaseModel

# Just the ID is sufficient, as the API itself handles fetching the data from the NVD
class CVERequest(BaseModel):
    cve_id: str

class PredictionResponse(BaseModel):
    cve_id: str
    probability: float
    prediction: int
    vendor: str | None = None
    cwe: str | None = None