KnnEs8Spec.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.KnnEs8Spec (spec) where

import Data.Aeson
import Data.Aeson.Key (Key)
import Data.Aeson.KeyMap qualified as KM
import Data.ByteString.Lazy.Char8 qualified as LBS
import Data.Maybe (isJust)
import Database.Bloodhound.Client.Cluster (BackendType (..))
import Database.Bloodhound.ElasticSearch8.Client
import TestsUtils.Common
import TestsUtils.Import

vectorField :: FieldName
vectorField = FieldName "vector"

queryVector :: [Double]
queryVector = [0.1, 0.2, 0.3, 0.4, 0.5]

hasKey :: Key -> LBS.ByteString -> Bool
hasKey k bs = case decode bs of
  Just (Object obj) -> k `KM.member` obj
  _ -> False

-- | True iff the encoded object has a "knn" key whose value is itself an
-- object containing another "knn" key. Detects the pre-fix wrapper shape
-- @{"index": ..., "knn": {"field": ...}}@.
hasNestedKnnKey :: LBS.ByteString -> Bool
hasNestedKnnKey bs = case decode bs of
  Just (Object obj) -> case KM.lookup "knn" obj of
    Just (Object inner) -> "knn" `KM.member` inner
    _ -> False
  _ -> False

spec :: Spec
spec = describe "k-NN Search API - Elasticsearch 8" $ do
  describe "request body shape" $ do
    it "serializes a KnnQuery as a bare ES clause (no index wrapper, no nested knn key)" $ do
      let clause =
            (mkKnnQuery vectorField queryVector (Just 3))
              { knnNumCandidates = Just 100
              }
      let json = encode clause
      json `shouldSatisfy` hasKey "field"
      json `shouldSatisfy` hasKey "query_vector"
      json `shouldSatisfy` hasKey "k"
      json `shouldSatisfy` hasKey "num_candidates"
      json `shouldNotSatisfy` hasKey "index"
      json `shouldNotSatisfy` hasNestedKnnKey

    it "omits Nothing optional fields from the clause body" $ do
      let clause = mkKnnQuery vectorField queryVector (Just 3)
      let json = encode clause
      json `shouldNotSatisfy` hasKey "filter"
      json `shouldNotSatisfy` hasKey "similarity"
      json `shouldNotSatisfy` hasKey "boost"
      json `shouldNotSatisfy` hasKey "query_vector_builder"
      json `shouldNotSatisfy` hasKey "inner_hits"

    it "includes filter, similarity, boost when set" $ do
      let clause =
            (mkKnnQuery vectorField queryVector (Just 5))
              { knnNumCandidates = Just 50,
                knnSimilarity = Just 0.9,
                knnBoost = Just 1.2,
                knnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing))
              }
      let json = encode clause
      json `shouldSatisfy` hasKey "filter"
      json `shouldSatisfy` hasKey "similarity"
      json `shouldSatisfy` hasKey "boost"

    it "serializes a query_vector_builder with text_embedding shape" $ do
      let clause =
            (mkKnnQuery vectorField [] (Just 5))
              { knnQueryVectorBuilder =
                  Just $
                    KnnQueryVectorBuilderTextEmbedding
                      { knnTextEmbeddingModelId = "model-x",
                        knnTextEmbeddingModelText = "hello world"
                      }
              }
      let Just (Object obj) = decode (encode clause)
      Just (Object builder) <- pure $ KM.lookup "query_vector_builder" obj
      Just (Object te) <- pure $ KM.lookup "text_embedding" builder
      KM.lookup "model_id" te `shouldBe` Just "model-x"
      KM.lookup "model_text" te `shouldBe` Just "hello world"

    it "round-trips a KnnQuery through ToJSON/FromJSON" $ do
      let clause =
            (mkKnnQuery vectorField queryVector (Just 3))
              { knnNumCandidates = Just 100
              }
      decode (encode clause) `shouldBe` Just clause

    it "serializes a Search with knnBody as a top-level knn array" $ do
      let clause = mkKnnQuery vectorField queryVector (Just 3)
      let search = (mkSearch Nothing Nothing) {knnBody = Just [clause]}
      let json = encode search
      json `shouldSatisfy` hasKey "knn"
      json `shouldNotSatisfy` hasNestedKnnKey
      let Just (Object obj) = decode json
      let Just (Array knnArr) = KM.lookup "knn" obj
      length knnArr `shouldBe` 1

  describe "spec compliance (bloodhound-bqe)" $ do
    it "never emits the non-spec window_bytes key" $ do
      let clause =
            (mkKnnQuery vectorField queryVector (Just 5))
              { knnNumCandidates = Just 100,
                knnSimilarity = Just 0.9,
                knnBoost = Just 1.2,
                knnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing))
              }
      let json = encode clause
      json `shouldNotSatisfy` hasKey "window_bytes"

    it "omits k when knnK = Nothing (k is optional per _types.KnnSearch)" $ do
      let clause = (mkKnnQuery vectorField queryVector Nothing) {knnNumCandidates = Just 100}
      let json = encode clause
      json `shouldNotSatisfy` hasKey "k"
      json `shouldSatisfy` hasKey "field"

    it "round-trips a KnnQuery with knnK = Nothing" $ do
      let clause = (mkKnnQuery vectorField queryVector Nothing) {knnNumCandidates = Just 100}
      decode (encode clause) `shouldBe` Just clause

    it "decodes a clause body that omits the optional k" $ do
      let bare = "{\"field\":\"vector\",\"query_vector\":[0.1,0.2,0.3,0.4,0.5]}"
      decode bare `shouldBe` Just (mkKnnQuery vectorField queryVector Nothing)

  describe "endpoint shape" $ do
    it "POSTs to /{index}/_search (not /_knn_search)" $ do
      let req = knnSearch @KnnDoc knnTestIndex (mkKnnQuery vectorField queryVector (Just 3)) (mkSearch Nothing Nothing)
      getRawEndpoint (bhRequestEndpoint req)
        `shouldBe` [unIndexName knnTestIndex, "_search"]

    it "carries no query string (the old bug shoved knn JSON into ?index=)" $ do
      let req = knnSearch @KnnDoc knnTestIndex (mkKnnQuery vectorField queryVector (Just 3)) (mkSearch Nothing Nothing)
      getRawEndpointQueries (bhRequestEndpoint req) `shouldBe` []

    it "embeds the knn clause in the request body, not as a query param" $ do
      let req = knnSearch @KnnDoc knnTestIndex (mkKnnQuery vectorField queryVector (Just 3)) (mkSearch Nothing Nothing)
      let body = bhRequestBody req
      body `shouldSatisfy` isJust
      let Just bodyBytes = body
      hasKey "knn" bodyBytes `shouldBe` True

  describe "live integration (requires ES8+ backend with dense_vector)" $
    backendSpecific [ElasticSearch8, ElasticSearch9] $ do
      it "returns the nearest neighbour document" $
        withTestEnv $ do
          insertKnnData
          let search = mkSearch Nothing Nothing
              clause = mkKnnQuery vectorField queryVector (Just 3)
          result <- tryPerformBHRequest $ knnSearch @KnnDoc knnTestIndex clause search
          case result of
            Left e -> liftIO $ expectationFailure $ "knnSearch failed: " <> show e
            Right searchResult -> do
              let topHit = headMay $ hits $ searchHits searchResult
              liftIO $ (hitDocId <$> topHit) `shouldBe` Just (Just (DocId "a"))

      it "supports optional filter and similarity fields" $
        withTestEnv $ do
          insertKnnData
          let search = mkSearch Nothing Nothing
              clause =
                (mkKnnQuery vectorField queryVector (Just 5))
                  { knnNumCandidates = Just 50,
                    knnSimilarity = Just 0.0,
                    knnFilter = Just (Filter (TermQuery (Term "label" "alpha") Nothing))
                  }
          result <- tryPerformBHRequest $ knnSearch @KnnDoc knnTestIndex clause search
          liftIO $ result `shouldSatisfy` isRight