<commit_before>package jobs

import java.util.{ Date, UUID }

import akka.actor.Actor
import akka.stream.ActorMaterializer
import com.google.inject.Inject
import jobs.KafkaJob.OnStart
import models.Prediction
import models.daos.{ ImageDAO, PredictionDAO }
import models.services.Kafka
import play.Logger
import play.api.libs.json.{ JsObject, JsValue, Json }

import scala.concurrent.Future
import scala.util.{ Failure, Success }

class KafkaJob @Inject() (
  kafka: Kafka,
  imageDAO: ImageDAO,
  predictionDAO: PredictionDAO
) extends Actor {

  import scala.concurrent.ExecutionContext.Implicits.global
  implicit val materializer = ActorMaterializer()

  override def receive: Receive = {
    case OnStart =>
      Logger.info("KafkaJob: Starting")

      kafka.source("classification-status") match {
        case Failure(e) =>
          Logger.info(s"Could not connect to Kafka! Error: $e")
        case Success(source) =>
          source.runForeach(x => {
            Logger.info(s"Received Status Update from Kafka: ${x.toString}")
            val json = Json.parse(x.value())

            (json \ "status").asOpt[String] match {
              case Some("CLASSIFICATION_STARTING") =>
                for {
                  status <- imageDAO.setClassificationStart(UUID.fromString(x.key()), new Date)
                } yield sender() ! true

              case Some(other) =>
                Logger.info(s"Got message from Kafka that is unhandled: $other")

              case None =>
                Logger.warn(s"Invalid Message from Kafka: ${x.value()}")
            }
          })
      }

      kafka.source("predictions") match {
        case Failure(e) =>
          Logger.info(s"Could not connect to Kafka! Error: $e")
        case Success(source) =>
          source.runForeach(x => {
            Logger.info("Received Prediction from Kafka: " + x.toString)
            val json = Json.parse(x.value())

            val preds = kafkaResponseToPredictions(json, x.key())
            val duration = getClassificationDuration(json)

            Logger.debug(s"Setting Classification duration of $duration for id=${x.key()}")<commit_msg>Add setStatus() call for other status updates<commit_after>package jobs

import java.util.{ Date, UUID }

import akka.actor.Actor
import akka.stream.ActorMaterializer
import com.google.inject.Inject
import jobs.KafkaJob.OnStart
import models.Prediction
import models.daos.{ ImageDAO, PredictionDAO }
import models.services.Kafka
import play.Logger
import play.api.libs.json.{ JsObject, JsValue, Json }

import scala.concurrent.Future
import scala.util.{ Failure, Success }

class KafkaJob @Inject() (
  kafka: Kafka,
  imageDAO: ImageDAO,
  predictionDAO: PredictionDAO
) extends Actor {

  import scala.concurrent.ExecutionContext.Implicits.global
  implicit val materializer = ActorMaterializer()

  override def receive: Receive = {
    case OnStart =>
      Logger.info("KafkaJob: Starting")

      kafka.source("classification-status") match {
        case Failure(e) =>
          Logger.info(s"Could not connect to Kafka! Error: $e")
        case Success(source) =>
          source.runForeach(x => {
            Logger.info(s"Received Status Update from Kafka: ${x.toString}")
            val json = Json.parse(x.value())

            (json \ "status").asOpt[String] match {
              case Some("CLASSIFICATION_STARTING") =>
                for {
                  status <- imageDAO.setClassificationStart(UUID.fromString(x.key()), new Date)
                } yield sender ! status

              case Some(other) =>
                for {
                  status <- imageDAO.setStatus(UUID.fromString(x.key()), Some(other))
                } yield sender ! status

              case None =>
                Logger.warn(s"Invalid Message from Kafka: ${x.value()}")
            }
          })
      }

      kafka.source("predictions") match {
        case Failure(e) =>
          Logger.info(s"Could not connect to Kafka! Error: $e")
        case Success(source) =>
          source.runForeach(x => {
            Logger.info("Received Prediction from Kafka: " + x.toString)
            val json = Json.parse(x.value())

            val preds = kafkaResponseToPredictions(json, x.key())
            val duration = getClassificationDuration(json)

            Logger.debug(s"Setting Classification duration of $duration for id=${x.key()}")
{'fraction_non_alphanumeric': 0.11179277436946149, 'fraction_numerical': 0.0, 'mean_word_length': 2.4688731284475964, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2229202', 'n_tokens_mistral': 1275, 'n_tokens_neox': 1243, 'n_words': 302}