MLTrainSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.MLTrainSpec (spec) where

import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude

-- | Sync train response (@POST /_plugins/_ml/_train/{algorithm}@):
-- @{model_id, status}@. The async shape (@{task_id, status}@) is the
-- same 'MLTaskAck' envelope; both decode with one parser.
sampleTrainSync :: LBS.ByteString
sampleTrainSync = "{\"model_id\":\"lblVmX8BO5w8y8RaYYvN\",\"status\":\"COMPLETED\"}"

-- | Async train response: @{task_id, status}@.
sampleTrainAsync :: LBS.ByteString
sampleTrainAsync = "{\"task_id\":\"lrlamX8BO5w8y8Ra2otd\",\"status\":\"CREATED\"}"

-- | Train-predict response: @{status, prediction_result}@. The
-- prediction payload is algorithm-specific, so it is kept opaque.
sampleTrainPredict :: LBS.ByteString
sampleTrainPredict =
  "{\
  \  \"status\": \"COMPLETED\",\
  \  \"prediction_result\": {\
  \    \"column_metas\": [{\"name\": \"ClusterID\", \"column_type\": \"INTEGER\"}],\
  \    \"rows\": [{\"values\": [{\"column_type\": \"INTEGER\", \"value\": 1}]}]\
  \  }\
  \}"

spec :: Spec
spec = describe "ML Commons train / train-predict APIs" $ do
  describe "MLTaskAck (train responses)" $ do
    it "decodes the sync train response (model_id + status)" $
      decode sampleTrainSync
        `shouldBe` Just
          MLTaskAck
            { mlTaskAckTaskId = Nothing,
              mlTaskAckStatus = Just "COMPLETED",
              mlTaskAckModelId = Just "lblVmX8BO5w8y8RaYYvN",
              mlTaskAckTaskType = Nothing
            }
    it "decodes the async train response (task_id + status)" $
      decode sampleTrainAsync
        `shouldBe` Just
          MLTaskAck
            { mlTaskAckTaskId = Just "lrlamX8BO5w8y8Ra2otd",
              mlTaskAckStatus = Just "CREATED",
              mlTaskAckModelId = Nothing,
              mlTaskAckTaskType = Nothing
            }

  describe "TrainModelRequest" $ do
    it "encodes parameters + input_query + input_index, omitting absent input_data" $ do
      let req =
            TrainModelRequest
              { trainModelParameters =
                  object
                    [ "centroids" .= (3 :: Int),
                      "iterations" .= (10 :: Int)
                    ],
                trainModelInputQuery = Just (object ["size" .= (10000 :: Int)]),
                trainModelInputIndex = Just ["iris_data"],
                trainModelInputData = Nothing
              }
      encode req `shouldBe` "{\"input_index\":[\"iris_data\"],\"input_query\":{\"size\":10000},\"parameters\":{\"centroids\":3,\"iterations\":10}}"
    it "round-trips a request built from inline input_data" $ do
      let req =
            TrainModelRequest
              { trainModelParameters = object ["k" .= (2 :: Int)],
                trainModelInputQuery = Nothing,
                trainModelInputIndex = Nothing,
                trainModelInputData = Just (object ["rows" .= ([] :: [Value])])
              }
      decode (encode req) `shouldBe` Just req

  describe "TrainAndPredictResponse" $ do
    it "decodes the doc sample (status + opaque prediction_result)" $ do
      let decoded = decode sampleTrainPredict :: Maybe TrainAndPredictResponse
      case decoded of
        Just r -> do
          trainAndPredictResponseStatus r `shouldBe` Just "COMPLETED"
          trainAndPredictResponsePredictionResult r `shouldSatisfy` isJust
        Nothing -> expectationFailure "expected TrainAndPredictResponse decode"
    it "round-trips a minimal response (status only)" $ do
      let r = TrainAndPredictResponse {trainAndPredictResponseStatus = Just "COMPLETED", trainAndPredictResponsePredictionResult = Nothing}
      decode (encode r) `shouldBe` Just r