KnnTrainSpec.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE QuasiQuotes #-}

module Test.KnnTrainSpec (spec) where

import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Char8 qualified as BS
import Data.ByteString.Lazy.Char8 qualified as LBS
import Database.Bloodhound.OpenSearch3.Requests qualified as OS3Requests
import Database.Bloodhound.OpenSearch3.Types
import TestsUtils.Import
import Prelude

-- | A representative @KnnTrainingRequest@ used by the endpoint-shape tests
-- in both the 'trainKnnModel' and 'trainKnnModelWith' describe blocks.
sampleReq :: KnnTrainingRequest
sampleReq =
  KnnTrainingRequest
    { knnTrainingRequestIndex = [qqIndexName|train-index-name|],
      knnTrainingRequestField = FieldName "train-field-name",
      knnTrainingRequestDimension = 16,
      knnTrainingRequestMethod =
        KnnTrainingMethod
          { knnTrainingMethodName = "ivf",
            knnTrainingMethodEngine = KnnEngineFaiss,
            knnTrainingMethodSpaceType = Just "l2",
            knnTrainingMethodParameters =
              Just (object ["nlist" .= (128 :: Int)])
          },
      knnTrainingRequestSpaceType = Nothing,
      knnTrainingRequestMaxTrainingVectorCount = Nothing,
      knnTrainingRequestSearchSize = Nothing,
      knnTrainingRequestDescription = Just "My model"
    }

-- | A representative @method@ sub-object with full @parameters@, drawn from
-- the OpenSearch vector-search train-model example. Only IVF-style methods
-- (FAISS) require training; @parameters@ is deeply nested and
-- engine-specific.
sampleMethodJson :: LBS.ByteString
sampleMethodJson =
  "{\
  \  \"name\": \"ivf\",\
  \  \"engine\": \"faiss\",\
  \  \"parameters\": {\
  \    \"nlist\": 128,\
  \    \"encoder\": {\
  \      \"name\": \"pq\",\
  \      \"parameters\": { \"code_size\": 8 }\
  \    }\
  \  }\
  \}"

-- | A fully-populated train request body. Mirrors the example in the
-- OpenSearch vector-search docs.
sampleFullRequest :: LBS.ByteString
sampleFullRequest =
  "{\
  \  \"training_index\": \"train-index-name\",\
  \  \"training_field\": \"train-field-name\",\
  \  \"dimension\": 16,\
  \  \"max_training_vector_count\": 1200,\
  \  \"search_size\": 100,\
  \  \"description\": \"My model\",\
  \  \"space_type\": \"l2\",\
  \  \"method\": {\
  \    \"name\": \"ivf\",\
  \    \"engine\": \"faiss\",\
  \    \"parameters\": {\
  \      \"nlist\": 128,\
  \      \"encoder\": {\
  \        \"name\": \"pq\",\
  \        \"parameters\": { \"code_size\": 8 }\
  \      }\
  \    }\
  \  }\
  \}"

-- | A minimal request with only the four required fields. Used to verify
-- @omitNulls@ does not leak @null@ for the optional fields.
sampleMinimalRequest :: LBS.ByteString
sampleMinimalRequest =
  "{\
  \  \"training_index\": \"train-index-name\",\
  \  \"training_field\": \"train-field-name\",\
  \  \"dimension\": 4,\
  \  \"method\": {\
  \    \"name\": \"ivf\",\
  \    \"engine\": \"faiss\"\
  \  }\
  \}"

-- | The train endpoint's response is a bare object echoing the @model_id@.
sampleResponse :: LBS.ByteString
sampleResponse = "{ \"model_id\": \"dcdwscddscsad\" }"

spec :: Spec
spec = describe "k-NN Train Model API" $ do
  describe "KnnTrainingMethod JSON" $ do
    it "decodes a method with full parameters" $ do
      let Just decoded = decode sampleMethodJson :: Maybe KnnTrainingMethod
      knnTrainingMethodName decoded `shouldBe` "ivf"
      knnTrainingMethodEngine decoded `shouldBe` KnnEngineFaiss
      knnTrainingMethodSpaceType decoded `shouldBe` Nothing
      knnTrainingMethodParameters decoded `shouldSatisfy` isJust

    it "decodes a method with space_type at the method level" $ do
      let Just decoded =
            decode
              "{ \"name\": \"ivf\", \"engine\": \"faiss\", \"space_type\": \"innerproduct\" }" ::
              Maybe KnnTrainingMethod
      knnTrainingMethodSpaceType decoded `shouldBe` Just "innerproduct"

    it "requires name and engine" $ do
      let decoded = decode "{ \"name\": \"ivf\" }" :: Maybe KnnTrainingMethod
      decoded `shouldBe` Nothing

    it "rejects a top-level array" $ do
      let decoded = decode "[]" :: Maybe KnnTrainingMethod
      decoded `shouldBe` Nothing

    it "round-trips a full method through ToJSON/FromJSON" $ do
      let Just decoded = decode sampleMethodJson :: Maybe KnnTrainingMethod
      (decode . encode) decoded `shouldBe` Just decoded

    it "ToJSON omits Nothing method fields (no null leakage)" $ do
      let method =
            KnnTrainingMethod
              { knnTrainingMethodName = "ivf",
                knnTrainingMethodEngine = KnnEngineFaiss,
                knnTrainingMethodSpaceType = Nothing,
                knnTrainingMethodParameters = Nothing
              }
          encoded = LBS.toStrict (encode method)
      encoded `shouldSatisfy` not . BS.isInfixOf "\"space_type\""
      encoded `shouldSatisfy` not . BS.isInfixOf "\"parameters\""
      encoded `shouldSatisfy` BS.isInfixOf "\"name\":\"ivf\""
      encoded `shouldSatisfy` BS.isInfixOf "\"engine\":\"faiss\""

    it "ToJSON emits method parameters when present" $ do
      let Just decoded = decode sampleMethodJson :: Maybe KnnTrainingMethod
          encoded = LBS.toStrict (encode decoded)
      encoded `shouldSatisfy` BS.isInfixOf "\"parameters\""

  describe "KnnTrainingRequest JSON" $ do
    it "decodes a fully-populated request" $ do
      let Just decoded = decode sampleFullRequest :: Maybe KnnTrainingRequest
      unIndexName (knnTrainingRequestIndex decoded) `shouldBe` "train-index-name"
      unFieldName (knnTrainingRequestField decoded) `shouldBe` "train-field-name"
      knnTrainingRequestDimension decoded `shouldBe` 16
      knnTrainingRequestSpaceType decoded `shouldBe` Just "l2"
      knnTrainingRequestMaxTrainingVectorCount decoded `shouldBe` Just 1200
      knnTrainingRequestSearchSize decoded `shouldBe` Just 100
      knnTrainingRequestDescription decoded `shouldBe` Just "My model"
      (knnTrainingMethodName . knnTrainingRequestMethod) decoded `shouldBe` "ivf"

    it "decodes a minimal (required-only) request" $ do
      let Just decoded = decode sampleMinimalRequest :: Maybe KnnTrainingRequest
      knnTrainingRequestDimension decoded `shouldBe` 4
      knnTrainingRequestSpaceType decoded `shouldBe` Nothing
      knnTrainingRequestMaxTrainingVectorCount decoded `shouldBe` Nothing
      knnTrainingRequestSearchSize decoded `shouldBe` Nothing
      knnTrainingRequestDescription decoded `shouldBe` Nothing

    it "requires the four required fields" $ do
      decode "{ \"training_index\": \"x\", \"training_field\": \"y\", \"dimension\": 4 }"
        `shouldBe` (Nothing :: Maybe KnnTrainingRequest)

    it "tolerates unknown fields (forward compatibility)" $ do
      let Just decoded =
            decode
              "{\
              \  \"training_index\": \"ix\",\
              \  \"training_field\": \"fl\",\
              \  \"dimension\": 2,\
              \  \"method\": { \"name\": \"ivf\", \"engine\": \"faiss\" },\
              \  \"future_field\": 42\
              \}" ::
              Maybe KnnTrainingRequest
      unIndexName (knnTrainingRequestIndex decoded) `shouldBe` "ix"

    it "rejects a top-level array" $ do
      let decoded = decode "[]" :: Maybe KnnTrainingRequest
      decoded `shouldBe` Nothing

    it "rejects malformed JSON" $ do
      let decoded = decode "{ not json" :: Maybe KnnTrainingRequest
      decoded `shouldBe` Nothing

    it "round-trips a full request through ToJSON/FromJSON" $ do
      let Just decoded = decode sampleFullRequest :: Maybe KnnTrainingRequest
      (decode . encode) decoded `shouldBe` Just decoded

    it "round-trips a minimal request through ToJSON/FromJSON" $ do
      let Just decoded = decode sampleMinimalRequest :: Maybe KnnTrainingRequest
      (decode . encode) decoded `shouldBe` Just decoded

    it "ToJSON omits Nothing request fields (no null leakage)" $ do
      let Just decoded = decode sampleMinimalRequest :: Maybe KnnTrainingRequest
          encoded = LBS.toStrict (encode decoded)
      encoded `shouldSatisfy` not . BS.isInfixOf "\"space_type\""
      encoded `shouldSatisfy` not . BS.isInfixOf "\"max_training_vector_count\""
      encoded `shouldSatisfy` not . BS.isInfixOf "\"search_size\""
      encoded `shouldSatisfy` not . BS.isInfixOf "\"description\""

    it "ToJSON serializes training_index and training_field as bare strings (not wrapped objects)" $ do
      let Just decoded = decode sampleMinimalRequest :: Maybe KnnTrainingRequest
          encoded = LBS.toStrict (encode decoded)
      encoded `shouldSatisfy` BS.isInfixOf "\"training_index\":\"train-index-name\""
      encoded `shouldSatisfy` BS.isInfixOf "\"training_field\":\"train-field-name\""

    it "ToJSON emits optional fields when present" $ do
      let Just decoded = decode sampleFullRequest :: Maybe KnnTrainingRequest
          encoded = LBS.toStrict (encode decoded)
      encoded `shouldSatisfy` BS.isInfixOf "\"space_type\":\"l2\""
      encoded `shouldSatisfy` BS.isInfixOf "\"max_training_vector_count\":1200"
      encoded `shouldSatisfy` BS.isInfixOf "\"search_size\":100"

  describe "KnnTrainResponse JSON" $ do
    it "decodes the documented response body" $ do
      let Just decoded = decode sampleResponse :: Maybe KnnTrainResponse
      knnTrainResponseModelId decoded `shouldBe` "dcdwscddscsad"

    it "requires the model_id field" $ do
      let decoded = decode "{}" :: Maybe KnnTrainResponse
      decoded `shouldBe` Nothing

    it "rejects a top-level array" $ do
      let decoded = decode "[]" :: Maybe KnnTrainResponse
      decoded `shouldBe` Nothing

    it "round-trips through ToJSON/FromJSON" $ do
      let Just decoded = decode sampleResponse :: Maybe KnnTrainResponse
      (decode . encode) decoded `shouldBe` Just decoded

  describe "trainKnnModel endpoint shape" $ do
    -- Pure checks against the BHRequest shape; no live backend needed.
    it "POSTs to /_plugins/_knn/models/{model_id}/_train for the given ModelId" $ do
      let req = OS3Requests.trainKnnModel (ModelId "test-model") sampleReq
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "test-model", "_train"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "uses the model id as a single opaque path segment" $ do
      let req = OS3Requests.trainKnnModel (ModelId "faiss-ivf-pq_42") sampleReq
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "faiss-ivf-pq_42", "_train"]

    it "uses a distinct id in the path when given a different ModelId" $ do
      let req = OS3Requests.trainKnnModel (ModelId "other-model") sampleReq
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "other-model", "_train"]

    it "uses the POST method" $ do
      let req = OS3Requests.trainKnnModel (ModelId "test-model") sampleReq
      bhRequestMethod req `shouldBe` "POST"

    it "attaches a JSON body (POST semantics)" $ do
      let req = OS3Requests.trainKnnModel (ModelId "test-model") sampleReq
      bhRequestBody req `shouldSatisfy` isJust

    it "embeds training_index and method in the request body" $ do
      let req = OS3Requests.trainKnnModel (ModelId "test-model") sampleReq
          Just bodyBytes = bhRequestBody req
      case decode bodyBytes of
        Just (Object obj) -> do
          KM.member "training_index" obj `shouldBe` True
          KM.member "training_field" obj `shouldBe` True
          KM.member "dimension" obj `shouldBe` True
          KM.member "method" obj `shouldBe` True
        _ -> expectationFailure "body did not decode to Object"

    it "does not leak null for the unset optional fields" $ do
      let req = OS3Requests.trainKnnModel (ModelId "test-model") sampleReq
          Just bodyBytes = bhRequestBody req
          encoded = LBS.toStrict bodyBytes
      encoded `shouldSatisfy` not . BS.isInfixOf "\"space_type\":null"
      encoded `shouldSatisfy` not . BS.isInfixOf "\"max_training_vector_count\":null"
      encoded `shouldSatisfy` not . BS.isInfixOf "\"search_size\":null"

  describe "trainKnnModelWith endpoint shape" $ do
    it "trainKnnModel is equivalent to trainKnnModelWith Nothing" $ do
      let base = OS3Requests.trainKnnModel (ModelId "test-model") sampleReq
          withNothing = OS3Requests.trainKnnModelWith Nothing (ModelId "test-model") sampleReq
      getRawEndpoint (bhRequestEndpoint base)
        `shouldBe` getRawEndpoint (bhRequestEndpoint withNothing)
      getRawEndpointQueries (bhRequestEndpoint base)
        `shouldBe` getRawEndpointQueries (bhRequestEndpoint withNothing)
      bhRequestMethod base `shouldBe` bhRequestMethod withNothing

    it "trainKnnModelWith Nothing produces no query parameters" $ do
      let req = OS3Requests.trainKnnModelWith Nothing (ModelId "test-model") sampleReq
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "test-model", "_train"]
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "trainKnnModelWith (Just node) attaches ?preference=<node_id>" $ do
      let req = OS3Requests.trainKnnModelWith (Just (KnnNodeId "node-abc")) (ModelId "test-model") sampleReq
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "test-model", "_train"]
      getRawEndpointQueries (bhRequestEndpoint req)
        `shouldBe` [("preference", Just "node-abc")]

    it "trainKnnModelWith preserves the path for an opaque node id" $ do
      let req = OS3Requests.trainKnnModelWith (Just (KnnNodeId "ZcRWj1BLTEq9OMxK3wEuwA")) (ModelId "faiss-ivf-pq_42") sampleReq
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` ["_plugins", "_knn", "models", "faiss-ivf-pq_42", "_train"]
      getRawEndpointQueries (bhRequestEndpoint req)
        `shouldBe` [("preference", Just "ZcRWj1BLTEq9OMxK3wEuwA")]

    it "trainKnnModelWith uses the POST method and attaches a JSON body" $ do
      let req = OS3Requests.trainKnnModelWith (Just (KnnNodeId "node-1")) (ModelId "test-model") sampleReq
      bhRequestMethod req `shouldBe` "POST"
      bhRequestBody req `shouldSatisfy` isJust