Add vendor and cwe fields to prediction response

src/api/main.py

@@ -94,5 +94,7 @@ def predict(request: CVERequest):
     return PredictionResponse(
         cve_id=request.cve_id,
         probability=float(proba),
-        prediction=int(proba >= THRESHOLD)
+        prediction=int(proba >= THRESHOLD),
+        vendor=features.get("vendor"),
+    cwe=features.get("cwe"),
     )

src/api/schemas.py

@@ -1,5 +1,4 @@
 from pydantic import BaseModel
-# Definition of data contracts
 
 # Just the ID is sufficient, as the API itself handles fetching the data from the NVD
 class CVERequest(BaseModel):
@@ -8,4 +7,6 @@ class CVERequest(BaseModel):
 class PredictionResponse(BaseModel):
     cve_id: str
     probability: float
-    prediction: int
+    prediction: int
+    vendor: str | None = None
+    cwe: str | None = None