Fixed app v2

app.py

@@ -1,63 +1,70 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from typing import List
+from typing import List, Dict, Any
 from transformers import pipeline
 
 # Initialize FastAPI app
 app = FastAPI()
 
-
+# Load the custom CyBERTuned model directly from Hugging Face
 unmasker = pipeline("fill-mask", model="s2w-ai/CyBERTuned-SecurityLLM")
 
-# Define request model
+# Define request model for multiple log entries
 class LogRequest(BaseModel):
-    log: str
+    logs: List[str]  # Expecting a list of log entries
 
 # Define response model
 class ThreatResponse(BaseModel):
     log: str
     prompt: str
-    #threat_level_predictions: List[str]
-    #threat_type_predictions: List[str]
-    #detected_threat_level: str
-    #detected_threat_type: str
-    pred : List[object]
-
-# Function to predict masked words for threat level and type
-def predict_threat(log: str, unmasker, topk=5) -> List[List[object]]:
-    # Create prompt with masked tokens for threat level and threat type
-    prompt = f"{log} Threat level most likely is: <mask>. Most likely the Threat type: <mask>."
+    pred: Dict[str, List[Dict[str, Any]]]  # Dictionary with structured lists for predictions
+
+# Function to predict masked words for threat level, type, SRC IP, DEST IP, and Protocol
+def predict_threat(log: str, unmasker, topk=5) -> Dict[str, List[Dict[str, Any]]]:
+    # Create prompt with masked tokens for each category
+    prompt = (
+        f"{log} Threat level: <mask>. Threat type: <mask>. "
+        f"Attack type: <mask>. SRC IP: <mask>. DEST IP: <mask>. Protocol: <mask>."
+    )
     
     # Predict top options for each <mask>
     predictions = unmasker(prompt, top_k=topk)
     
-    # Extract top predictions for each <mask>
-    #threat_level_predictions = [pred["token_str"].strip() for pred in predictions[:topk]]
-    #threat_type_predictions = [pred["token_str"].strip() for pred in predictions[topk:2*topk]]
-    
-    return predictions
+    # Separate predictions for each masked category
+    threat_level_predictions = predictions[:topk]
+    threat_type_predictions = predictions[topk:2*topk]
+    attack_type_predictions = predictions[2*topk:3*topk]
+    src_ip_predictions = predictions[3*topk:4*topk]
+    dest_ip_predictions = predictions[4*topk:5*topk]
+    protocol_predictions = predictions[5*topk:6*topk]
+
+    # Structure the response in a dictionary
+    return {
+        "threat_level_predictions": threat_level_predictions,
+        "threat_type_predictions": threat_type_predictions,
+        "attack_type_predictions": attack_type_predictions,
+        "src_ip_predictions": src_ip_predictions,
+        "dest_ip_predictions": dest_ip_predictions,
+        "protocol_predictions": protocol_predictions
+    }
 
-# FastAPI endpoint for detecting threat level and type
-@app.post("/detect_threat", response_model=ThreatResponse)
+# FastAPI endpoint for detecting threat level and type for multiple logs
+@app.post("/detect_threat", response_model=List[ThreatResponse])
 async def detect_threat(log_request: LogRequest):
-    log = log_request.log
+    responses = []
+    for log in log_request.logs:
+        # Predict the threat level and type for each log entry
+        predictions = predict_threat(log, unmasker)
 
-    # Predict the threat level and type for the given log entry
-    predictions = predict_threat(log, unmasker)
-    
-    # Extract top predictions for threat level and type
-    ##threat_level_predictions = predictions[0] if len(predictions) > 0 else ["Unknown"]
-   ## threat_type_predictions = predictions[1] if len(predictions) > 1 else ["Unknown"]
-
-    # Use the top prediction as the most likely threat level and type
-    ##detected_threat_level = threat_level_predictions[0] if threat_level_predictions else "Unknown"
-    #detected_threat_type = threat_type_predictions[0] if threat_type_predictions else "Unknown"
-
-    # Prepare response
-    response = ThreatResponse(
-        log=log,
-        prompt=f"{log} Threat level: <mask>. Threat type: <mask>.",
-        pred=predictions
-    )
+        # Prepare response for each log entry
+        response = ThreatResponse(
+            log=log,
+            prompt=(
+                f"{log} Threat level: <mask>. Threat type: <mask>. "
+                f"Attack type: <mask>. SRC IP: <mask>. DEST IP: <mask>. Protocol: <mask>."
+            ),
+            pred=predictions
+        )
+        responses.append(response)
 
-    return response
+    return responses