module DataRobot.PredictResponse
( PredictError(..)
, PredictResult(..)
, PredictionValue(..)
, parseResponse
, predictionValue
) where
import Control.Applicative ((<|>))
import Control.Monad.Catch (Exception)
import Data.Aeson (FromJSON(..), ToJSON, Value(..), decode, defaultOptions, genericParseJSON, withObject, (.:), eitherDecode)
import Data.Aeson.Types (Options(..), typeMismatch)
import Data.List (find)
import Data.Maybe (fromMaybe, maybe)
import Data.Text (Text)
import Data.String.Conversions (cs)
import Data.Typeable (Typeable)
import GHC.Generics (Generic)
import Safe (headMay)
import Lens.Micro ((^.))
import Network.Wreq (Response, responseBody, responseHeader)
import Data.ByteString.Lazy (ByteString)
underscorePrefixOptions :: Options
underscorePrefixOptions =
defaultOptions { fieldLabelModifier = dropWhile (== '_') }
type Code = Int
data PredictError
= APIError Code Text
| MissingPrediction
deriving (Typeable, Show, Generic)
instance Exception PredictError
instance ToJSON PredictError
data ResponseSuccess = ResponseSuccess
{ _data :: [Prediction]
} deriving (Eq, Show, Generic)
instance FromJSON ResponseSuccess where
parseJSON = genericParseJSON underscorePrefixOptions
data ResponseFailure = ResponseFailure
{ _message :: Text
} deriving (Show, Eq, Generic)
instance FromJSON ResponseFailure where
parseJSON = genericParseJSON underscorePrefixOptions
data ResponseData
= ResponseData (Either ResponseFailure ResponseSuccess)
deriving (Eq, Show)
instance FromJSON ResponseData where
parseJSON v = ResponseData <$>
((Right <$> parseJSON v) <|> (Left <$> parseJSON v))
data PredictionValue = PredictionValue
{ label :: Text
, value :: Float
} deriving (Eq, Show, Generic)
instance ToJSON PredictionValue
instance FromJSON PredictionValue where
parseJSON = withObject "prediction_value" $ \o -> do
value' <- o .: "value"
label' <- labelText =<< o .: "label"
return $ PredictionValue label' value'
where
labelText (Number n) = pure $ (cs .show) n
labelText (String s) = pure s
labelText invalid = typeMismatch "label" invalid
data Prediction = Prediction
{ _prediction :: Value
, _predictionValues :: Maybe [PredictionValue]
} deriving (Eq, Show, Generic)
instance ToJSON Prediction
instance FromJSON Prediction where
parseJSON = genericParseJSON underscorePrefixOptions
data PredictResult = PredictResult
{ prediction :: Value
, predictionTimeMs :: Float
, predictionValues :: Maybe [PredictionValue]
} deriving (Show, Eq, Generic)
instance ToJSON PredictResult
handleResponse :: Float -> ResponseData -> Either PredictError PredictResult
handleResponse et (ResponseData (Right rs)) =
maybe (Left MissingPrediction) Right $ do
p <- headMay (_data rs)
pure PredictResult
{ prediction = _prediction p
, predictionValues = _predictionValues p
, predictionTimeMs = et
}
handleResponse _ (ResponseData (Left err)) =
responseFailure $ _message err
responseFailure :: Text -> Either PredictError PredictResult
responseFailure e = Left $ APIError 422 e
parseResponse :: Response ByteString -> Either PredictError PredictResult
parseResponse r = do
either (responseFailure . cs) (handleResponse tm) $ eitherDecode b
where
b = r ^. responseBody
et = r ^. responseHeader "X-DataRobot-Execution-Time"
tm = fromMaybe 0.0 $ decode (cs et)
predictionValue :: Text -> PredictResult -> Maybe Float
predictionValue c r = do
ps <- predictionValues r
pd <- find ((== c) . label) ps
pure $ value pd