NeuralQuerySpec.hs

{-# LANGUAGE OverloadedStrings #-}

-- | Pure-JSON tests for the OpenSearch @neural@ query clause.
--
-- These tests do not require a running OpenSearch instance; they verify
-- that 'NeuralQuery' (and the 'QueryNeuralQuery' arm of 'Query') serialize
-- to the JSON shape documented at
-- <https://docs.opensearch.org/latest/search-plugins/neural-search/#using-the-neural-query-clause>
-- and round-trip through 'ToJSON' \/ 'FromJSON'.
module Test.NeuralQuerySpec (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.Maybe (fromJust)
import Data.Text (Text)
import Database.Bloodhound.Types
  ( FieldName (..),
    NeuralQuery (..),
    NeuralQueryInput (..),
    NeuralQueryTermination (..),
    Query (..),
    mkNeuralQuery,
    mkNeuralQueryByMaxDistance,
    mkNeuralQueryByMinScore,
  )
import Test.Hspec (Spec, describe, it, shouldBe, shouldNotSatisfy, shouldSatisfy)

textField :: FieldName
textField = FieldName "passage_embedding"

imageField :: FieldName
imageField = FieldName "image_embedding"

modelId :: Text
modelId = "aVeifm3sBaaSkm0rKKxQ"

-- | 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 NeuralQuery produces the OS-correct shape
-- @{"neural": {<field>: {...}}}@ with the field name as a dict key.
hasNeuralNestedUnderField :: FieldName -> Value -> Bool
hasNeuralNestedUnderField (FieldName f) (Object obj) =
  case KM.lookup "neural" obj of
    Just (Object neuralObj) -> KM.member (AK.fromText f) neuralObj
    _ -> False
hasNeuralNestedUnderField _ _ = False

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

spec :: Spec
spec = describe "Neural query clause (OpenSearch)" $ do
  describe "NeuralQuery body shape (text input)" $ do
    it "serializes as neural.<field> with the field name as the dict key" $ do
      let clause = mkNeuralQuery textField (NeuralTextInput "hello world") modelId 5
      let json = encode clause
      json `shouldSatisfy` hasKey "neural"
      json `shouldSatisfy` hasNeuralNestedUnderFieldBS textField

    it "emits query_text (not query_image) for text input" $ do
      let clause = mkNeuralQuery textField (NeuralTextInput "hello world") modelId 5
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode clause)))
      KM.lookup "query_text" body `shouldBe` Just (String "hello world")
      body `shouldNotSatisfy` KM.member "query_image"

    it "omits boost when not set" $ do
      let clause = mkNeuralQuery textField (NeuralTextInput "hello world") modelId 5
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode clause)))
      body `shouldNotSatisfy` KM.member "boost"

    it "includes boost when set" $ do
      let clause = (mkNeuralQuery textField (NeuralTextInput "hello world") modelId 5) {neuralBoost = Just 1.5}
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode clause)))
      KM.lookup "boost" body `shouldBe` Just (toJSON (1.5 :: Double))

  describe "NeuralQuery body shape (image input)" $ do
    it "emits query_image (not query_text) for image input" $ do
      let clause = mkNeuralQuery imageField (NeuralImageInput "iVBORw0KGgo...") modelId 3
      let Just (Object body) = walk ["neural", AK.fromText "image_embedding"] (fromJust (decode (encode clause)))
      KM.lookup "query_image" body `shouldBe` Just (String "iVBORw0KGgo...")
      body `shouldNotSatisfy` KM.member "query_text"

  describe "NeuralQuery radial-search variants" $ do
    it "max_distance variant emits max_distance and no k in the clause body" $ do
      let clause = mkNeuralQueryByMaxDistance textField (NeuralTextInput "hello") modelId 1.5
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode clause)))
      KM.lookup "max_distance" body `shouldBe` Just (toJSON (1.5 :: Double))
      body `shouldNotSatisfy` KM.member "k"
      body `shouldNotSatisfy` KM.member "min_score"

    it "min_score variant emits min_score and no k in the clause body" $ do
      let clause = mkNeuralQueryByMinScore textField (NeuralTextInput "hello") modelId 0.8
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode clause)))
      KM.lookup "min_score" body `shouldBe` Just (toJSON (0.8 :: Double))
      body `shouldNotSatisfy` KM.member "k"
      body `shouldNotSatisfy` KM.member "max_distance"

    it "k variant still emits k (regression)" $ do
      let clause = mkNeuralQuery textField (NeuralTextInput "hello") modelId 5
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode clause)))
      KM.lookup "k" body `shouldBe` Just (toJSON (5 :: Int))
      body `shouldNotSatisfy` KM.member "max_distance"
      body `shouldNotSatisfy` KM.member "min_score"

  describe "NeuralQuery round-trip" $ do
    it "round-trips a minimal text+k query through ToJSON/FromJSON" $ do
      let clause = mkNeuralQuery textField (NeuralTextInput "hello") modelId 5
      decode (encode clause) `shouldBe` Just clause

    it "round-trips an image+k query through ToJSON/FromJSON" $ do
      let clause = mkNeuralQuery imageField (NeuralImageInput "iVBORw0KGgo...") modelId 3
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a fully-populated text+k query (with boost)" $ do
      let clause =
            (mkNeuralQuery textField (NeuralTextInput "hello") modelId 5)
              { neuralBoost = Just 1.5
              }
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a radial max_distance query" $ do
      let clause = mkNeuralQueryByMaxDistance textField (NeuralTextInput "hello") modelId 1.5
      decode (encode clause) `shouldBe` Just clause

    it "round-trips a radial min_score query" $ do
      let clause = mkNeuralQueryByMinScore textField (NeuralTextInput "hello") modelId 0.8
      decode (encode clause) `shouldBe` Just clause

  describe "NeuralQueryInput and NeuralQueryTermination" $ do
    it "NeuralTextInput round-trips" $ do
      let input = NeuralTextInput "hello"
      decode (encode input) `shouldBe` Just input

    it "NeuralImageInput round-trips" $ do
      let input = NeuralImageInput "iVBORw0KGgo..."
      decode (encode input) `shouldBe` Just input

    it "NeuralTerminationByK round-trips" $ do
      let term = NeuralTerminationByK 5
      decode (encode term) `shouldBe` Just term

    it "NeuralTerminationByMaxDistance round-trips" $ do
      let term = NeuralTerminationByMaxDistance 1.5
      decode (encode term) `shouldBe` Just term

    it "NeuralTerminationByMinScore round-trips" $ do
      let term = NeuralTerminationByMinScore 0.8
      decode (encode term) `shouldBe` Just term

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

    it "rejects NeuralQueryTermination with no k/max_distance/min_score" $ do
      let bad = object ["unrelated" .= Number 1]
      (decode (encode bad) :: Maybe NeuralQueryTermination) `shouldBe` Nothing

  describe "Query sum type integration" $ do
    it "QueryNeuralQuery wraps the clause under the \"neural\" key" $ do
      let inner = mkNeuralQuery textField (NeuralTextInput "hello") modelId 5
      let wrapped = QueryNeuralQuery inner
      let json = encode wrapped
      json `shouldSatisfy` \bs -> case decode bs of
        Just (Object obj) -> KM.member "neural" obj
        _ -> False

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

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

    it "QueryNeuralQuery preserves boost in the body shape" $ do
      let inner = (mkNeuralQuery textField (NeuralTextInput "hello") modelId 5) {neuralBoost = Just 2.0}
      let Just (Object body) = walk ["neural", AK.fromText "passage_embedding"] (fromJust (decode (encode (QueryNeuralQuery inner))))
      KM.lookup "boost" body `shouldBe` Just (toJSON (2.0 :: Double))

    it "QueryNeuralQuery round-trips through the parent Query sum type" $ do
      let inner = mkNeuralQuery textField (NeuralTextInput "hello") modelId 5
      let wrapped = QueryNeuralQuery inner
      decode (encode wrapped) `shouldBe` Just wrapped

    it "QueryNeuralQuery round-trips with boost set" $ do
      let inner =
            (mkNeuralQuery textField (NeuralTextInput "hello") modelId 5)
              { neuralBoost = Just 2.0
              }
      let wrapped = QueryNeuralQuery 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 wire JSON" $ do
      let wire = "{\"neural\":{\"passage_embedding\":{\"query_text\":\"hello\",\"model_id\":\"aVeifm3sBaaSkm0rKKxQ\",\"k\":5}}}"
      let expected = QueryNeuralQuery (mkNeuralQuery textField (NeuralTextInput "hello") modelId 5)
      (decode wire :: Maybe Query) `shouldBe` Just expected

    it "distinguishes neural from neural_sparse in the parser" $ do
      let inner = mkNeuralQuery textField (NeuralTextInput "hello") modelId 5
      let wrapped = QueryNeuralQuery inner
      let Just (Object obj) = decode (encode wrapped)
      obj `shouldSatisfy` KM.member "neural"
      obj `shouldNotSatisfy` KM.member "neural_sparse"

    it "rejects neural clause with multiple field keys" $ do
      let bad =
            object
              [ "neural"
                  .= object
                    [ "field_a" .= object ["query_text" .= String "x", "model_id" .= String "m", "k" .= Number 1],
                      "field_b" .= object ["query_text" .= String "y", "model_id" .= String "m", "k" .= Number 1]
                    ]
              ]
      (decode (encode bad) :: Maybe NeuralQuery) `shouldBe` Nothing

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

    it "rejects neural clause with no termination field" $ do
      let bad =
            object
              [ "neural"
                  .= object
                    [ "passage_embedding"
                        .= object
                          [ "query_text" .= String "x",
                            "model_id" .= String "m"
                          ]
                    ]
              ]
      (decode (encode bad) :: Maybe NeuralQuery) `shouldBe` Nothing

-- | 'hasKey' helper mirroring 'KnnOs3Spec' style: takes a 'Key' and a raw
-- encoded 'LBS.ByteString', decoding internally.
hasKey :: Key -> LBS.ByteString -> Bool
hasKey k bs = case decode bs of
  Just (Object obj) -> k `KM.member` obj
  _ -> False