add new prediction service

This commit is contained in:
heyethereum
2024-08-13 02:34:47 +08:00
parent 53f9acd922
commit c8cfe610a6
45 changed files with 167 additions and 82 deletions

View File

@@ -0,0 +1,63 @@
package com.safeqr.app.prediction.service;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.safeqr.app.prediction.model.URLFeatures;
import com.safeqr.app.qrcode.model.URLModel;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import org.springframework.http.ResponseEntity;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import static com.safeqr.app.constants.APIConstants.PREDICTION_API_URL;
@Service
public class PredictionService {
private static final Logger logger = LoggerFactory.getLogger(PredictionService.class);
private final RestTemplate restTemplate;
private final ObjectMapper objectMapper;
public PredictionService(RestTemplate restTemplate, ObjectMapper objectMapper) {
this.restTemplate = restTemplate;
this.objectMapper = objectMapper;
}
public String predict(URLModel urlModel) {
// Convert URLModel to URLFeatures
URLFeatures features = URLFeatures.fromEntity(urlModel);
logger.info("Prediction request: {}", features);
// Prepare the HTTP headers
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON);
// Create the HTTP entity containing the features and headers
HttpEntity<URLFeatures> requestEntity = new HttpEntity<>(features, headers);
// Make the HTTP POST request to the FastAPI prediction endpoint
ResponseEntity<String> response = restTemplate.exchange(
PREDICTION_API_URL,
HttpMethod.POST,
requestEntity,
String.class
);
// Use ObjectMapper to deserialize the response and automatically remove quotes
String prediction = response.getBody();
try {
prediction = objectMapper.readValue(prediction, String.class);
} catch (Exception e) {
logger.error("Failed to parse prediction response", e);
prediction = "Unknown";
}
logger.info("Prediction response: {}", prediction);
// Return the prediction
return prediction;
}
}