NeuralSparseQuerySpec.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Pure-JSON tests for the OpenSearch @neural_sparse@ query clause.
--
-- These tests do not require a running OpenSearch instance; they verify
-- that 'NeuralSparseQuery' (and the 'QueryNeuralSparseQuery' arm of
-- 'Query') serialize to the JSON shape documented at
-- <https://docs.opensearch.org/latest/query-dsl/specialized/neural-sparse/>
-- and round-trip through 'ToJSON' \/ 'FromJSON'.
module Test.NeuralSparseQuerySpec (spec) where

import Data.Aeson
import Data.Aeson.Key (Key)
import Data.Aeson.Key qualified as AK
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy qualified as LBS
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as M
import Data.Maybe (fromJust)
import Data.Text (Text)
import Database.Bloodhound.Types
  ( FieldName (..),
    NeuralSparseInput (..),
    NeuralSparseMethodParameters (..),
    NeuralSparseModelSource (..),
    NeuralSparseQuery (..),
    Query (..),
    mkNeuralSparseQueryText,
    mkNeuralSparseQueryTextAnalyzer,
    mkNeuralSparseQueryTokens,
    mkNeuralSparseQueryTokensAnalyzer,
  )
import Test.Hspec (Spec, describe, it, shouldBe, shouldNotSatisfy, shouldSatisfy)

textField :: FieldName
textField = FieldName "passage_embedding"

sparseVectorField :: FieldName
sparseVectorField = FieldName "sparse_embedding"

modelId :: Text
modelId = "aP2Q8ooBpBj3wT4HVS8a"

analyzer :: Text
analyzer = "bert-uncased"

sampleTokens :: Map Text Double
sampleTokens =
  M.fromList
    [ ("hello", 3.3210192),
      ("world", 2.6087382)
    ]

-- | Walk into a nested dict path, returning the inner Value if reachable.
walk :: [Key] -> Value -> Maybe Value
walk [] v = Just v
walk (k : ks) (Object obj) = case KM.lookup k obj of
  Just v -> walk ks v
  Nothing -> Nothing
walk _ _ = Nothing

-- | True iff the encoded NeuralSparseQuery produces the OS-correct shape
-- @{"neural_sparse": {<field>: {...}}}@ with the field name as a dict key.
hasNeuralSparseNestedUnderField :: FieldName -> Value -> Bool
hasNeuralSparseNestedUnderField (FieldName f) (Object obj) =
  case KM.lookup "neural_sparse" obj of
    Just (Object neuralObj) -> KM.member (AK.fromText f) neuralObj
    _ -> False
hasNeuralSparseNestedUnderField _ _ = False

-- | Bytestring variant for direct use with @\`shouldSatisfy\`@ on @encode@ output.
hasNeuralSparseNestedUnderFieldBS :: FieldName -> LBS.ByteString -> Bool
hasNeuralSparseNestedUnderFieldBS field bs =
  case decode bs of
    Just v -> hasNeuralSparseNestedUnderField field v
    Nothing -> False

-- | Lookup the inner clause body for a given field name.
clauseBody :: FieldName -> Value -> Maybe Value
clauseBody (FieldName f) v =
  walk ["neural_sparse", AK.fromText f] v

-- | True iff the encoded value is an Object containing the given key.
hasKey :: Key -> LBS.ByteString -> Bool
hasKey k bs = case decode bs of
  Just (Object obj) -> k `KM.member` obj
  _ -> False

spec :: Spec
spec = describe "Neural Sparse query clause (OpenSearch)" $ do
  describe "NeuralSparseQuery body shape (text input, model_id)" $ do
    it "serializes as neural_sparse.<field> with the field name as the dict key" $ do
      let clause = mkNeuralSparseQueryText textField "Hi world" modelId
      let json = encode clause
      json `shouldSatisfy` hasKey "neural_sparse"
      json `shouldSatisfy` hasNeuralSparseNestedUnderFieldBS textField

    it "emits query_text for text input" $ do
      let clause = mkNeuralSparseQueryText textField "Hi world" modelId
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      KM.lookup "query_text" body `shouldBe` Just (String "Hi world")
      body `shouldNotSatisfy` KM.member "query_tokens"

    it "emits model_id when set" $ do
      let clause = mkNeuralSparseQueryText textField "Hi world" modelId
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      KM.lookup "model_id" body `shouldBe` Just (String modelId)
      body `shouldNotSatisfy` KM.member "analyzer"

    it "omits method_parameters when not set" $ do
      let clause = mkNeuralSparseQueryText textField "Hi world" modelId
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      body `shouldNotSatisfy` KM.member "method_parameters"

    it "never emits the deprecated max_token_score field" $ do
      let clause = mkNeuralSparseQueryText textField "Hi world" modelId
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      body `shouldNotSatisfy` KM.member "max_token_score"

  describe "NeuralSparseQuery body shape (text input, analyzer)" $ do
    it "emits analyzer when set" $ do
      let clause = mkNeuralSparseQueryTextAnalyzer textField "Hi world" analyzer
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      KM.lookup "analyzer" body `shouldBe` Just (String analyzer)
      body `shouldNotSatisfy` KM.member "model_id"

  describe "NeuralSparseQuery body shape (tokens input)" $ do
    it "emits query_tokens (not query_text) for raw token input" $ do
      let clause = mkNeuralSparseQueryTokens textField sampleTokens modelId
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      case KM.lookup "query_tokens" body of
        Just (Object _) -> pure ()
        other -> fail $ "expected query_tokens object, got " ++ show other
      body `shouldNotSatisfy` KM.member "query_text"

    it "serializes the token weights as a string-keyed map" $ do
      let clause = mkNeuralSparseQueryTokens textField sampleTokens modelId
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      case KM.lookup "query_tokens" body of
        Just (Object tokensObj) -> do
          KM.lookup "hello" tokensObj `shouldBe` Just (toJSON (3.3210192 :: Double))
          KM.lookup "world" tokensObj `shouldBe` Just (toJSON (2.6087382 :: Double))
        other -> fail $ "expected query_tokens object, got " ++ show other

    it "supports analyzer source with token input" $ do
      let clause = mkNeuralSparseQueryTokensAnalyzer textField sampleTokens analyzer
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      KM.lookup "analyzer" body `shouldBe` Just (String analyzer)
      body `shouldNotSatisfy` KM.member "model_id"

  describe "NeuralSparseQuery body shape (default analyzer)" $ do
    it "omits both model_id and analyzer when no source is set" $ do
      let clause =
            (mkNeuralSparseQueryText textField "Hi world" modelId)
              { neuralSparseModelSource = Nothing
              }
      let Just (Object body) = clauseBody textField (fromJust (decode (encode clause)))
      body `shouldNotSatisfy` KM.member "model_id"
      body `shouldNotSatisfy` KM.member "analyzer"

  describe "NeuralSparseQuery method_parameters" $ do
    it "emits method_parameters when set on sparse_vector fields" $ do
      let clause =
            (mkNeuralSparseQueryText sparseVectorField "Hi world" modelId)
              { neuralSparseMethodParameters =
                  Just
                    NeuralSparseMethodParameters
                      { neuralSparseMethodParametersTopN = Just 10,
                        neuralSparseMethodParametersHeapFactor = Just 1.0,
                        neuralSparseMethodParametersK = Just 10
                      }
              }
      let Just (Object body) = clauseBody sparseVectorField (fromJust (decode (encode clause)))
      case KM.lookup "method_parameters" body of
        Just (Object mp) -> do
          KM.lookup "top_n" mp `shouldBe` Just (toJSON (10 :: Int))
          KM.lookup "heap_factor" mp `shouldBe` Just (toJSON (1.0 :: Double))
          KM.lookup "k" mp `shouldBe` Just (toJSON (10 :: Int))
        other -> fail $ "expected method_parameters object, got " ++ show other

    it "omits null sub-fields of method_parameters" $ do
      let clause =
            (mkNeuralSparseQueryText sparseVectorField "Hi world" modelId)
              { neuralSparseMethodParameters =
                  Just
                    NeuralSparseMethodParameters
                      { neuralSparseMethodParametersTopN = Just 10,
                        neuralSparseMethodParametersHeapFactor = Nothing,
                        neuralSparseMethodParametersK = Nothing
                      }
              }
      let Just (Object body) = clauseBody sparseVectorField (fromJust (decode (encode clause)))
      case KM.lookup "method_parameters" body of
        Just (Object mp) -> do
          KM.lookup "top_n" mp `shouldBe` Just (toJSON (10 :: Int))
          mp `shouldNotSatisfy` KM.member "heap_factor"
          mp `shouldNotSatisfy` KM.member "k"
        other -> fail $ "expected method_parameters object, got " ++ show other

  describe "NeuralSparseQuery round-trip" $ do
    it "round-trips a text + model_id query through ToJSON/FromJSON" $ do
      let clause = mkNeuralSparseQueryText textField "Hi world" modelId
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a text + analyzer query through ToJSON/FromJSON" $ do
      let clause = mkNeuralSparseQueryTextAnalyzer textField "Hi world" analyzer
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a text query with no model source (default analyzer)" $ do
      let clause =
            (mkNeuralSparseQueryText textField "Hi world" modelId)
              { neuralSparseModelSource = Nothing
              }
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a tokens + model_id query through ToJSON/FromJSON" $ do
      let clause = mkNeuralSparseQueryTokens textField sampleTokens modelId
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a query with method_parameters" $ do
      let clause =
            (mkNeuralSparseQueryText sparseVectorField "Hi world" modelId)
              { neuralSparseMethodParameters =
                  Just
                    NeuralSparseMethodParameters
                      { neuralSparseMethodParametersTopN = Just 10,
                        neuralSparseMethodParametersHeapFactor = Just 1.0,
                        neuralSparseMethodParametersK = Just 10
                      }
              }
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a query with partially-populated method_parameters" $ do
      let clause =
            (mkNeuralSparseQueryText sparseVectorField "Hi world" modelId)
              { neuralSparseMethodParameters =
                  Just
                    NeuralSparseMethodParameters
                      { neuralSparseMethodParametersTopN = Just 5,
                        neuralSparseMethodParametersHeapFactor = Nothing,
                        neuralSparseMethodParametersK = Just 7
                      }
              }
      decode (encode clause) `shouldBe` Just clause

  describe "NeuralSparseInput and NeuralSparseModelSource" $ do
    it "NeuralSparseTextInput round-trips" $ do
      let input = NeuralSparseTextInput "hello"
      decode (encode input) `shouldBe` Just input

    it "NeuralSparseTokensInput round-trips" $ do
      let input = NeuralSparseTokensInput sampleTokens
      decode (encode input) `shouldBe` Just input

    it "NeuralSparseByModelId round-trips" $ do
      let src = NeuralSparseByModelId modelId
      decode (encode src) `shouldBe` Just src

    it "NeuralSparseByAnalyzer round-trips" $ do
      let src = NeuralSparseByAnalyzer analyzer
      decode (encode src) `shouldBe` Just src

    it "NeuralSparseMethodParameters round-trips with all fields set" $ do
      let params =
            NeuralSparseMethodParameters
              { neuralSparseMethodParametersTopN = Just 10,
                neuralSparseMethodParametersHeapFactor = Just 1.0,
                neuralSparseMethodParametersK = Just 5
              }
      decode (encode params) `shouldBe` Just params

    it "rejects NeuralSparseInput with neither query_text nor query_tokens" $ do
      let bad = object ["something_else" .= String "x"]
      (decode (encode bad) :: Maybe NeuralSparseInput) `shouldBe` Nothing

    it "rejects NeuralSparseModelSource with neither model_id nor analyzer" $ do
      let bad = object ["unrelated" .= Number 1]
      (decode (encode bad) :: Maybe NeuralSparseModelSource) `shouldBe` Nothing

  describe "Query sum type integration" $ do
    it "QueryNeuralSparseQuery wraps the clause under the \"neural_sparse\" key" $ do
      let inner = mkNeuralSparseQueryText textField "Hi world" modelId
      let wrapped = QueryNeuralSparseQuery inner
      let json = encode wrapped
      json `shouldSatisfy` \bs -> case decode bs of
        Just (Object obj) -> KM.member "neural_sparse" obj
        _ -> False

    -- Regression guard for the double-wrap bug: the Query wrapper must NOT
    -- re-wrap the leaf (which already self-wraps under "neural_sparse").
    -- The wire shape must be exactly {"neural_sparse": {<field>: {...}}},
    -- not {"neural_sparse": {"neural_sparse": {<field>: {...}}}}.
    it "QueryNeuralSparseQuery does NOT double-wrap the leaf neural_sparse key" $ do
      let inner = mkNeuralSparseQueryText textField "Hi world" modelId
      let json = encode (QueryNeuralSparseQuery inner)
      let Just (Object outer) = decode json
          Just (Object neuralSparseObj) = KM.lookup "neural_sparse" outer
      neuralSparseObj `shouldNotSatisfy` KM.member "neural_sparse"
      neuralSparseObj `shouldSatisfy` KM.member (AK.fromText "passage_embedding")

    -- Walk-into-body assertion: the body shape under obj.neural_sparse.<field>
    -- must match the OS-correct clause body. This is the test that would
    -- have caught the double-wrap bug.
    it "QueryNeuralSparseQuery produces OS-correct body shape under neural_sparse.<field>" $ do
      let inner = mkNeuralSparseQueryText textField "Hi world" modelId
      let Just (Object body) = walk ["neural_sparse", AK.fromText "passage_embedding"] (fromJust (decode (encode (QueryNeuralSparseQuery inner))))
      KM.lookup "query_text" body `shouldBe` Just (String "Hi world")
      KM.lookup "model_id" body `shouldBe` Just (String modelId)

    it "QueryNeuralSparseQuery round-trips through the parent Query sum type" $ do
      let inner = mkNeuralSparseQueryText textField "Hi world" modelId
      let wrapped = QueryNeuralSparseQuery inner
      decode (encode wrapped) `shouldBe` Just wrapped

    it "QueryNeuralSparseQuery round-trips with tokens input" $ do
      let inner = mkNeuralSparseQueryTokens textField sampleTokens modelId
      let wrapped = QueryNeuralSparseQuery inner
      decode (encode wrapped) `shouldBe` Just wrapped

    it "QueryNeuralSparseQuery round-trips with method_parameters set" $ do
      let inner =
            (mkNeuralSparseQueryText sparseVectorField "Hi world" modelId)
              { neuralSparseMethodParameters =
                  Just
                    NeuralSparseMethodParameters
                      { neuralSparseMethodParametersTopN = Just 10,
                        neuralSparseMethodParametersHeapFactor = Just 1.0,
                        neuralSparseMethodParametersK = Just 10
                      }
              }
      let wrapped = QueryNeuralSparseQuery inner
      decode (encode wrapped) `shouldBe` Just wrapped

    -- End-to-end: an OS-shaped wire JSON decodes to the expected Query.
    -- Pre-fix, this returned Nothing because the decoder expected the
    -- double-wrapped shape; post-fix, it succeeds.
    it "decodes a real OS-shaped neural_sparse wire JSON" $ do
      let wire = "{\"neural_sparse\":{\"passage_embedding\":{\"query_text\":\"Hi world\",\"model_id\":\"aP2Q8ooBpBj3wT4HVS8a\"}}}"
      let expected = QueryNeuralSparseQuery (mkNeuralSparseQueryText textField "Hi world" modelId)
      (decode wire :: Maybe Query) `shouldBe` Just expected

    it "distinguishes neural_sparse from neural in the parser" $ do
      let inner = mkNeuralSparseQueryText textField "Hi world" modelId
      let wrapped = QueryNeuralSparseQuery inner
      let Just (Object obj) = decode (encode wrapped)
      obj `shouldNotSatisfy` KM.member "neural"
      obj `shouldSatisfy` KM.member "neural_sparse"

  describe "NeuralSparseQuery parser rejects malformed clauses" $ do
    it "rejects neural_sparse clause with multiple field keys" $ do
      let bad =
            object
              [ "neural_sparse"
                  .= object
                    [ "field_a" .= object ["query_text" .= String "x", "model_id" .= String "m"],
                      "field_b" .= object ["query_text" .= String "y", "model_id" .= String "m"]
                    ]
              ]
      (decode (encode bad) :: Maybe NeuralSparseQuery) `shouldBe` Nothing

    it "rejects neural_sparse clause with no field key" $ do
      let bad = object ["neural_sparse" .= object []]
      (decode (encode bad) :: Maybe NeuralSparseQuery) `shouldBe` Nothing

    it "rejects neural_sparse clause with no input field" $ do
      let bad =
            object
              [ "neural_sparse"
                  .= object
                    [ "passage_embedding"
                        .= object
                          ["model_id" .= String "m"]
                    ]
              ]
      (decode (encode bad) :: Maybe NeuralSparseQuery) `shouldBe` Nothing

    it "accepts neural_sparse clause with model_id and analyzer absent (default)" $ do
      let ok =
            object
              [ "neural_sparse"
                  .= object
                    [ "passage_embedding"
                        .= object
                          ["query_text" .= String "Hi world"]
                    ]
              ]
      let expected =
            NeuralSparseQuery
              { neuralSparseField = textField,
                neuralSparseInput = NeuralSparseTextInput "Hi world",
                neuralSparseModelSource = Nothing,
                neuralSparseMethodParameters = Nothing
              }
      (decode (encode ok) :: Maybe NeuralSparseQuery) `shouldBe` Just expected