Spaces:

Alvlt
/

CVE-exploit-prediction

Sleeping

CVE-exploit-prediction / src /api /schemas.py

Add vendor and cwe fields to prediction response

631793e 4 months ago

316 Bytes

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

	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