Merge pull request #4758 from bbc/t3214-mappings-exclude-vectors

T3214 mappings exclude vectors

Author: andrew-nowak

common-lib/src/main/scala/com/gu/mediaservice/lib/config/CommonConfigWithElastic.scala

@@ -13,7 +13,8 @@ class CommonConfigWithElastic(resources: GridConfigResources) extends CommonConf
     ),
     url = string("es6.url"),
     shards = string("es6.shards").toInt,
-    replicas = string("es6.replicas").toInt
+    replicas = string("es6.replicas").toInt,
+    includeDenseVectorMappings = booleanOpt("es.includeDenseVectorMappings").getOrElse(true)
   )
 
   private val persistenceIdentifier = string("persistence.identifier")

common-lib/src/main/scala/com/gu/mediaservice/lib/elasticsearch/ElasticSearchClient.scala

@@ -37,6 +37,7 @@ trait ElasticSearchClient extends ElasticSearchExecutions with GridLogging {
 
   def shards: Int
   def replicas: Int
+  def includeDenseVectorMappings: Boolean
 
   lazy val client = {
     logger.info("Connecting to Elastic 8: " + url)
@@ -146,7 +147,7 @@ trait ElasticSearchClient extends ElasticSearchExecutions with GridLogging {
         "to remove the need for these: " + nonRecommendenedIndexSettingOverrides)
 
       createIndex(index).
-        mapping(Mappings.imageMapping).
+        mapping(Mappings.imageMapping(includeDenseVectorMappings)).
         analysis(IndexSettings.analysis).
         settings(nonRecommendenedIndexSettingOverrides).
         shards(shards).

common-lib/src/main/scala/com/gu/mediaservice/lib/elasticsearch/ElasticSearchConfig.scala

@@ -1,3 +1,3 @@
 package com.gu.mediaservice.lib.elasticsearch
 
-case class ElasticSearchConfig(aliases: ElasticSearchAliases, url: String, shards: Int, replicas: Int)
+case class ElasticSearchConfig(aliases: ElasticSearchAliases, url: String, shards: Int, replicas: Int, includeDenseVectorMappings: Boolean = true)

common-lib/src/main/scala/com/gu/mediaservice/lib/elasticsearch/Mappings.scala

@@ -8,7 +8,7 @@ import play.api.libs.json.{JsObject, Json}
 
 object Mappings {
 
-  val imageMapping: MappingDefinition = {
+  def imageMapping(includeDenseVectorMappings: Boolean = true): MappingDefinition = {
 
     // Non indexed fields stored as keywords can still participate in exists / has queries
     def filemetaDataStringsAsKeyword: DynamicTemplateRequest = {
@@ -75,8 +75,7 @@ object Mappings {
         leasesMapping("leases"),
         collectionMapping("collections"),
         esInfoMapping("esInfo"),
-        embeddingMapping("embedding"),
-      )
+      ) ++ (if (includeDenseVectorMappings) Seq(embeddingMapping("embedding")) else Seq.empty)
     )
   }
 

dev/script/generate-config/service-config.js

@@ -146,7 +146,7 @@ function getKahunaConfig(config){
         |]`;
 
     const permissionsConfig = stripMargin`
-      |usePermissionsFilter=true
+      |usePermissionsFilter=false
       |permissionsDefault="allPermissions"
       `;
 
@@ -191,6 +191,7 @@ function getMediaApiConfig(config) {
         |es6.url="${config.es6.url}"
         |es6.shards=${config.es6.shards}
         |es6.replicas=${config.es6.replicas}
+        |es.includeDenseVectorMappings=true
         |quota.store.key="rcs-quota.json"
         |security.cors.allowedOrigins="${getCorsAllowedOriginString(config)}"
         |search.fuzziness={
@@ -225,6 +226,7 @@ function getThrallConfig(config) {
         |es6.url="${config.es6.url}"
         |es6.shards=${config.es6.shards}
         |es6.replicas=${config.es6.replicas}
+        |es.includeDenseVectorMappings=true
         |metrics.request.enabled=false
         |`;
 }

media-api/app/lib/elasticsearch/ElasticSearch.scala

@@ -75,6 +75,7 @@ class ElasticSearch(
   lazy val url = elasticConfig.url
   lazy val shards = elasticConfig.shards
   lazy val replicas = elasticConfig.replicas
+  lazy val includeDenseVectorMappings = elasticConfig.includeDenseVectorMappings
 
   private val SearchQueryTimeout = FiniteDuration(10, TimeUnit.SECONDS)
   // there is 15 seconds timeout set on cluster level as well
@@ -180,18 +181,23 @@ class ElasticSearch(
 
   def knnSearch(queryEmbedding: List[Float], k: Int, numCandidates: Int, filterOpt: Option[Query])
                (implicit ex: ExecutionContext, logMarker: LogMarker): Future[SearchResults] = {
-    val knn = Knn("embedding.cohereEmbedV4.image", filter = filterOpt)
-      .queryVector(queryEmbedding.map(_.toDouble))
-      .k(k)
-      .numCandidates(numCandidates)
+    if (!includeDenseVectorMappings) {
+      logger.warn(logMarker, "knnSearch called but includeDenseVectorMappings=false, returning empty results")
+      Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
+    } else {
+      val knn = Knn("embedding.cohereEmbedV4.image", filter = filterOpt)
+        .queryVector(queryEmbedding.map(_.toDouble))
+        .k(k)
+        .numCandidates(numCandidates)
 
-    val searchRequest = ElasticDsl.search(imagesCurrentAlias)
-      .knn(knn)
-      .size(k)
+      val searchRequest = ElasticDsl.search(imagesCurrentAlias)
+        .knn(knn)
+        .size(k)
 
-    executeAndLog(withSearchQueryTimeout(searchRequest), "knn search").map { r =>
-      val imageHits = r.result.hits.hits.map(resolveHit).toSeq.flatten.map(i => (i.instance.id, i))
-      SearchResults(hits = imageHits, total = imageHits.length, extraCounts = None)
+      executeAndLog(withSearchQueryTimeout(searchRequest), "knn search").map { r =>
+        val imageHits = r.result.hits.hits.map(resolveHit).toSeq.flatten.map(i => (i.instance.id, i))
+        SearchResults(hits = imageHits, total = imageHits.length, extraCounts = None)
+      }
     }
   }
 
@@ -271,15 +277,20 @@ class ElasticSearch(
     implicit ex: ExecutionContext,
     logMarker: LogMarker
   ): Future[SearchResults] = {
-    val queryEmbeddingDouble: List[Double] = queryEmbedding.map(_.toDouble)
-
-    for {
-      maxScore <- fetchMaxBm25Score(query, filterOpt)
-      searchRequest = makeHybridSearchRequest(query, queryEmbeddingDouble, k, numCandidates, vecWeight, maxScore, filterOpt)
-      result <- executeAndLog(withSearchQueryTimeout(searchRequest), "hybrid search")
-    } yield {
-      val imageHits = result.result.hits.hits.map(resolveHit).toSeq.flatten.map(i => (i.instance.id, i))
-      SearchResults(hits = imageHits, total = imageHits.length, extraCounts = None)
+    if (!includeDenseVectorMappings) {
+      logger.warn(logMarker, "hybridSearch called but includeDenseVectorMappings=false, returning empty results")
+      Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
+    } else {
+      val queryEmbeddingDouble: List[Double] = queryEmbedding.map(_.toDouble)
+
+      for {
+        maxScore <- fetchMaxBm25Score(query, filterOpt)
+        searchRequest = makeHybridSearchRequest(query, queryEmbeddingDouble, k, numCandidates, vecWeight, maxScore, filterOpt)
+        result <- executeAndLog(withSearchQueryTimeout(searchRequest), "hybrid search")
+      } yield {
+        val imageHits = result.result.hits.hits.map(resolveHit).toSeq.flatten.map(i => (i.instance.id, i))
+        SearchResults(hits = imageHits, total = imageHits.length, extraCounts = None)
+      }
     }
   }
 

scripts/src/main/scala/com/gu/mediaservice/scripts/EsScript.scala

@@ -197,15 +197,19 @@ object Reindex extends EsScript {
 object UpdateMapping extends EsScript {
 
   def run(esUrl: String, extraArgs: List[String]): Unit = {
+    val includeDenseVectors = !extraArgs.contains("--no-dense-vectors")
+    val remainingArgs = extraArgs.filterNot(_ == "--no-dense-vectors")
 
-    object MappingsClient extends EsClient(esUrl) {
+    println(s"includeDenseVectorMappings: $includeDenseVectors")
+
+    object MappingsClient extends EsClient(esUrl, includeDenseVectors) {
       def updateMappings(specifiedIndex: Option[String]): Unit = {
         val index = specifiedIndex.getOrElse(imagesCurrentAlias)
         println(s"Updating mapping on index: $index")
 
         val result = client.execute {
           putMapping(Indexes(index)) as {
-            Mappings.imageMapping.properties
+            Mappings.imageMapping(includeDenseVectors).properties
           }
         }.await
 
@@ -219,11 +223,11 @@ object UpdateMapping extends EsScript {
 
     }
 
-    MappingsClient.updateMappings(extraArgs.headOption)
+    MappingsClient.updateMappings(remainingArgs.headOption)
   }
 
   def usageError: Nothing = {
-    System.err.println("Usage: UpdateMapping <ES_URL>")
+    System.err.println("Usage: UpdateMapping <ES_URL> [<index_name>] [--no-dense-vectors]")
     sys.exit(1)
   }
 }
@@ -422,7 +426,7 @@ abstract class EsScript {
     run(esUrl, extraArgs)
   }
 
-  class EsClient(val url: String) extends ElasticSearchClient {
+  class EsClient(val url: String, val includeDenseVectorMappings: Boolean = true) extends ElasticSearchClient {
     override def imagesCurrentAlias = esImagesAlias
     override def shards = esShards
     override def replicas = esReplicas

thrall/app/lib/elasticsearch/ElasticSearch.scala

@@ -41,6 +41,7 @@ class ElasticSearch(
   lazy val url: String = config.url
   lazy val shards: Int = config.shards
   lazy val replicas: Int = config.replicas
+  lazy val includeDenseVectorMappings: Boolean = config.includeDenseVectorMappings
 
 
   def migrationAwareUpdater[REQUEST, RESPONSE](
@@ -78,9 +79,13 @@ class ElasticSearch(
     executeAndLog(request, s"Setting migration info on image id: ${imageId}")
   }
 
+  private def prepareImageJson(json: JsObject): JsObject =
+    if (includeDenseVectorMappings) json
+    else json - "embedding"
+
   def directInsert(image: Image, indexName: String)(implicit ex: ExecutionContext, logMarker: LogMarker): Future[ElasticSearchInsertResponse] =
     executeAndLog(
-      indexInto(indexName).id(image.id).source(Json.stringify(Json.toJson(image))),
+      indexInto(indexName).id(image.id).source(Json.stringify(prepareImageJson(Json.toJson(image).as[JsObject]))),
       s"ES6 indexing image ${image.id} into index '$indexName'"
     ).map(indexResponse =>
       ElasticSearchInsertResponse(indexResponse.result.index)
@@ -92,14 +97,14 @@ class ElasticSearch(
 
     // On insert, we know we will not have a lastModified to consider, so we always take the one we get
     val insertImage = image.copy(lastModified = Some(lastModified))
-    val insertImageAsJson = Json.toJson(insertImage)
+    val insertImageAsJson = prepareImageJson(Json.toJson(insertImage).as[JsObject])
 
     def runUpsertIntoIndex(indexAlias: String, maybeEsInfo: Option[JsObject]) = {
       val esInfo = maybeEsInfo.getOrElse(JsObject.empty)
 
       // On update, we do not want to take the one we have been given unless it is newer - see updateLastModifiedScript script
       val updateImage = image.copy(lastModified = None)
-      val upsertImageAsJson = Json.toJson(updateImage)
+      val upsertImageAsJson = prepareImageJson(Json.toJson(updateImage).as[JsObject])
 
       val painlessSource =
         // If there are old identifiers, then merge any new identifiers into old and use the merged results as the new identifiers
@@ -214,18 +219,24 @@ class ElasticSearch(
   def updateEmbedding(id: String, embedding: Embedding, lastModified: DateTime)
                      (implicit ex: ExecutionContext, logMarker: LogMarker): List[Future[ElasticSearchUpdateResponse]] = {
 
-    val replaceEmbeddingScript = "ctx._source.embedding = params.embedding;"
+    if (!includeDenseVectorMappings) {
+      logger.info(s"Skipping embedding update for image $id: dense vector mappings are disabled")
+      List(Future.successful(ElasticSearchUpdateResponse()))
+    } else {
 
-    val scriptSource = loadPainless(replaceEmbeddingScript)
+      val replaceEmbeddingScript = "ctx._source.embedding = params.embedding;"
 
-    val embeddingParameter = asNestedMap(Json.toJson(embedding))
+      val scriptSource = loadPainless(replaceEmbeddingScript)
 
-    val eventualUpdateResponse = migrationAwareUpdater(
-      requestFromIndexName = indexName => prepareUpdateRequest(indexName, id, scriptSource, lastModified, ("embedding", embeddingParameter)),
-      logMessageFromIndexName = indexName => s"ES6 updating embedding on image $id in index $indexName"
-    ).incrementOnFailure(metrics.map(_.failedEmbeddingUpdates)) { case _ => true }
+      val embeddingParameter = asNestedMap(Json.toJson(embedding))
 
-    List(eventualUpdateResponse.map(_ => ElasticSearchUpdateResponse()))
+      val eventualUpdateResponse = migrationAwareUpdater(
+        requestFromIndexName = indexName => prepareUpdateRequest(indexName, id, scriptSource, lastModified, ("embedding", embeddingParameter)),
+        logMessageFromIndexName = indexName => s"ES6 updating embedding on image $id in index $indexName"
+      ).incrementOnFailure(metrics.map(_.failedEmbeddingUpdates)) { case _ => true }
+
+      List(eventualUpdateResponse.map(_ => ElasticSearchUpdateResponse()))
+    }
   }
 
   def applyImageMetadataOverride(id: String, metadata: Edits, lastModified: DateTime)