Revamp hybrid search algorithm

.github/workflows/ci.yml

@@ -15,10 +15,16 @@ jobs:
       pull-requests: write
     services:
       elasticsearch:
-        image: docker.elastic.co/elasticsearch/elasticsearch:7.17.29
-        # Wait for elasticsearch to report healthy before continuing.
+        image: docker.elastic.co/elasticsearch/elasticsearch:8.18.3
+        # These options should mimic those set in ElasticSearchDockerBase.scala, because this is
+        # the CI version of the TestContainers approach we use for running tests locally.
+        # `xpack.security.enabled=false` is needed since v8 for tests to use plain
+        # unauthenticated HTTP, because in v7 => v8 they changed the default to `true`:
+        # https://www.elastic.co/guide/en/elasticsearch/reference/8.18/release-notes-8.0.0.html#:~:text=Set%20xpack.security.enabled%20to%20true%20for%20all%20licenses
+        #
+        # `--health-` settings make it wait for es to report healthy before continuing.
         # see https://github.com/actions/example-services/blob/master/.github/workflows/postgres-service.yml#L28
-        options: -e "discovery.type=single-node" --expose 9200 --health-cmd "curl localhost:9200/_cluster/health" --health-interval 10s --health-timeout 5s --health-retries 10
+        options: -e "discovery.type=single-node" -e "xpack.security.enabled=false" --expose 9200 --health-cmd "curl localhost:9200/_cluster/health" --health-interval 10s --health-timeout 5s --health-retries 10
         ports:
           - 9200:9200
       localstack:

common-lib/src/main/scala/com/gu/mediaservice/lib/VectorUtils.scala

@@ -0,0 +1,80 @@
+package com.gu.mediaservice.lib
+
+object VectorUtils {
+  def dotProduct(vectorOne: List[Double], vectorTwo: List[Double]): Double = {
+    require(vectorOne.length == vectorTwo.length, "Vectors must be the same dimensionality")
+    vectorOne.zip(vectorTwo).map(component => component._1 * component._2).sum
+  }
+
+  def magnitude(vector: List[Double]): Double = {
+    math.sqrt(vector.map(math.pow(_, 2)).sum)
+  }
+
+  def cosineSimilarity(vectorOne: List[Double], vectorTwo: List[Double]): Option[Double] = {
+    // Cosine similarity is undefined when the vectors differ in dimensionality
+    // (the dot product is meaningless) or when either has zero magnitude (the
+    // angle to the origin is meaningless), so we return None in those cases and
+    // let the caller choose a domain-appropriate fallback.
+    if (vectorOne.length != vectorTwo.length) None
+    else {
+      val magnitudeProduct = magnitude(vectorOne) * magnitude(vectorTwo)
+      if (magnitudeProduct == 0.0) None
+      else Some(dotProduct(vectorOne, vectorTwo) / magnitudeProduct)
+    }
+  }
+
+  // ===============================================================================
+  // The below functions are primarily for constructing artificial vectors for tests
+  // ===============================================================================
+  private def oneHotVector(dims: Int, hotIndex: Int): List[Double] =
+    List.tabulate(dims)(i => if (i == hotIndex) 1.0 else 0.0)
+
+  // The first standard basis vector e0: (1, 0, 0, ...).
+  def firstBasisVector(dims: Int): List[Double] = oneHotVector(dims, 0)
+
+  private def vectorWithComponents(dims: Int, components: (Int, Double)*): List[Double] = {
+    val arr = Array.fill(dims)(0.0)
+    components.foreach { case (i, v) => arr(i) = v }
+    arr.toList
+  }
+
+  // Builds a vector whose cosine similarity with the first basis vector equals
+  // the requested `similarity` (in [-1, 1], where 1 is identical and 0 is orthogonal).
+  // This is specifically for use in tests where we need to create test fixtures with
+  // specific semantic scores.
+  def vectorWithCosineSimilarity(dims: Int, similarity: Double): List[Double] = {
+    require(dims > 1, "Dimensions must be at least 2")
+    require(similarity >= -1.0 && similarity <= 1.0, "Cosine similarity must be between -1 and 1")
+    if (similarity == 0.0) {
+      // A cosine similarity of 0 means orthogonal to the first basis vector.
+      // That vector is one-hot at index 0, so any vector with a zero
+      // first component is orthogonal to it. (0, 1, 0, 0, ...) is the simplest choice.
+      oneHotVector(dims, 1)
+    } else {
+      val absVec2D = {
+        // Let's pretend we're in just 2 dimensions and think about triangles.
+        //
+        //            /|
+        //           / |
+        //        n /  | opposite = sqrt(n^2 - 1)
+        //         /   |
+        //        /θ___|
+        //          1
+        //       (adjacent)
+        //
+        // Our queryEmbedding lies flat on the x axis.
+        // cos(θ) = adjacent / hypotenuse = 1 / n
+        // If we want cos θ to be 1/n, e.g. 1/2,
+        // then the adjacent side must be 1, and the hypotenuse n.
+        // We want to find the opposite side
+        // n^2 = 1^2 + adj^2
+        // opposite = sqrt(n^2 - 1)
+        // The two components of our vector are therefore (1, sqrt(n^2 - 1))
+        val n = 1 / Math.abs(similarity)
+        (1, Math.sqrt(Math.pow(n, 2) - 1))
+      }
+      val vec2D = if (similarity < 0.0) (-absVec2D._1, -absVec2D._2) else absVec2D
+      vectorWithComponents(dims, 0 -> vec2D._1, 1 -> vec2D._2)
+    }
+  }
+}

common-lib/src/test/scala/com/gu/mediaservice/lib/VectorUtilsTest.scala

@@ -0,0 +1,143 @@
+package com.gu.mediaservice.lib
+
+import org.scalatest.Inspectors
+import org.scalatest.OptionValues
+import org.scalatest.funspec.AnyFunSpec
+import org.scalatest.matchers.should.Matchers
+import org.scalacheck.Gen
+import VectorUtils.{dotProduct, magnitude, cosineSimilarity, firstBasisVector, vectorWithCosineSimilarity}
+
+
+class VectorUtilsTest extends AnyFunSpec with Matchers with Inspectors with OptionValues {
+  val tolerance = 1e-9
+
+  describe("dotProduct") {
+    it ("should compute the dot product of two vectors") {
+      dotProduct(List(1.0, 2.0, 3.0), List(4.0, 5.0, 6.0)) shouldBe 32.0
+    }
+
+    it ("should compute a dot product of zero for orthogonal vectors") {
+      dotProduct(List(1.0, 0.0), List(0.0, 1.0)) shouldBe 0.0
+    }
+
+    it ("should return zero for the dot product of empty vectors") {
+      dotProduct(List.empty, List.empty) shouldBe 0.0
+    }
+
+    it ("should reject vectors of differing dimensionality") {
+      an [IllegalArgumentException] should be thrownBy dotProduct(List(1.0, 2.0), List(1.0, 2.0, 3.0))
+    }
+  }
+
+  describe("magnitude") {
+    it ("should compute the magnitude of a vector") {
+      magnitude(List(3.0, 4.0)) shouldBe 5.0
+    }
+
+    it ("should return a magnitude of zero for a zero vector") {
+      magnitude(List(0.0, 0.0, 0.0)) shouldBe 0.0
+    }
+  }
+
+  describe("cosineSimilarity") {
+    it ("should return a cosine similarity of 1 for identical vectors") {
+      cosineSimilarity(List(1.0, 2.0, 3.0), List(1.0, 2.0, 3.0)).value shouldBe 1.0 +- tolerance
+    }
+
+    it ("should return a cosine similarity of 1 for parallel vectors") {
+      cosineSimilarity(List(1.0, 2.0, 3.0), List(2.0, 4.0, 6.0)).value shouldBe 1.0 +- tolerance
+    }
+
+    it ("should return a cosine similarity of 0 for orthogonal vectors") {
+      cosineSimilarity(List(1.0, 0.0), List(0.0, 1.0)).value shouldBe 0.0 +- tolerance
+    }
+
+    it ("should return a cosine similarity of -1 for opposite vectors") {
+      cosineSimilarity(List(1.0, 2.0, 3.0), List(-1.0, -2.0, -3.0)).value shouldBe -1.0 +- tolerance
+    }
+
+    it ("should be undefined (None) when the first vector has zero magnitude") {
+      cosineSimilarity(List(0.0, 0.0, 0.0), List(1.0, 2.0, 3.0)) shouldBe None
+    }
+
+    it ("should be undefined (None) when the second vector has zero magnitude") {
+      cosineSimilarity(List(1.0, 2.0, 3.0), List(0.0, 0.0, 0.0)) shouldBe None
+    }
+
+    it ("should be undefined (None) when both vectors have zero magnitude") {
+      cosineSimilarity(List(0.0, 0.0, 0.0), List(0.0, 0.0, 0.0)) shouldBe None
+    }
+
+    it ("should be undefined (None) when the vectors differ in dimensionality") {
+      cosineSimilarity(List(1.0, 2.0), List(1.0, 2.0, 3.0)) shouldBe None
+    }
+  }
+
+  // This function is specifically for use in tests,
+  // where we need to create test fixtures with specific semantic scores.
+  describe("vectorWithCosineSimilarity") {
+    val dims = 256
+
+    it ("should create a vector at the requested cosine similarity across a fine sweep of the valid range") {
+      val similarities = (-1000 to 1000).map(_ / 1000.0)
+      forAll(similarities) { s =>
+        cosineSimilarity(firstBasisVector(dims), vectorWithCosineSimilarity(dims, s)).value shouldBe s +- tolerance
+      }
+    }
+
+    it ("should create a vector at the requested cosine similarity for randomly generated similarities") {
+      val similarityGen = Gen.chooseNum(-1.0, 1.0)
+      val samples = Gen.listOfN(500, similarityGen).sample.getOrElse(Nil)
+      samples should not be empty
+      forAll(samples) { s =>
+        cosineSimilarity(firstBasisVector(dims), vectorWithCosineSimilarity(dims, s)).value shouldBe s +- tolerance
+      }
+    }
+
+    it ("should work across a range of dimensions") {
+      forAll(List(2, 3, 8, 16, 128, 256, 1024)) { n =>
+        cosineSimilarity(firstBasisVector(n), vectorWithCosineSimilarity(n, 0.42)).value shouldBe 0.42 +- tolerance
+      }
+    }
+
+    it ("should return a vector identical to the basis vector for a similarity of 1") {
+      vectorWithCosineSimilarity(dims, 1.0) shouldBe firstBasisVector(dims)
+      cosineSimilarity(firstBasisVector(dims), vectorWithCosineSimilarity(dims, 1.0)).value shouldBe 1.0 +- tolerance
+    }
+
+    it ("should return a vector opposite to the basis vector for a similarity of -1") {
+      cosineSimilarity(firstBasisVector(dims), vectorWithCosineSimilarity(dims, -1.0)).value shouldBe -1.0 +- tolerance
+    }
+
+    it ("should return a vector orthogonal to the basis vector for a similarity of 0") {
+      cosineSimilarity(firstBasisVector(dims), vectorWithCosineSimilarity(dims, 0.0)).value shouldBe 0.0 +- tolerance
+    }
+
+    it ("should negate the vector when the similarity is negated") {
+      val positive = vectorWithCosineSimilarity(dims, 0.25)
+      val negative = vectorWithCosineSimilarity(dims, -0.25)
+      forAll(positive.zip(negative)) { case (p, n) => n shouldBe (-p) +- tolerance }
+    }
+
+    it ("should produce only finite components for valid similarities") {
+      forAll(List(-1.0, -0.5, -1e-3, 0.0, 1e-3, 0.5, 1.0)) { s =>
+        forAll(vectorWithCosineSimilarity(dims, s)) { component =>
+          component.isNaN shouldBe false
+          component.isInfinite shouldBe false
+        }
+      }
+    }
+
+    it ("should reject a similarity greater than 1") {
+      an [IllegalArgumentException] should be thrownBy vectorWithCosineSimilarity(dims, 1.0001)
+    }
+
+    it ("should reject a similarity less than -1") {
+      an [IllegalArgumentException] should be thrownBy vectorWithCosineSimilarity(dims, -1.0001)
+    }
+
+    it ("should reject fewer than 2 dimensions") {
+      an [IllegalArgumentException] should be thrownBy vectorWithCosineSimilarity(1, 0.5)
+    }
+  }
+}

media-api/app/controllers/MediaApi.scala

@@ -613,23 +613,32 @@ class MediaApi(
     }
 
     def semanticSearchByImage(imageId: String, k: Int, params: SearchParams): Future[SearchResults] = {
-      val filterOpt = buildAiFilter(params)
-
-      for {
-        maybeImage <- elasticSearch.getImageById(imageId)
-        maybeEmbedding = maybeImage
-          .filter(image => isVisibleToAccessor(request.user, image))
-          .flatMap(_.embedding)
-          .flatMap(_.cohereEmbedV4)
-          .map(_.image)
-        searchResults <- maybeEmbedding match {
-          // If we have an embedding, perform the KNN search. If not, return an empty result set.
-          case Some(embedding) =>
-            elasticSearch.knnSearch(embedding, k = k, numCandidates = Math.max(k * 2, 100), filterOpt = filterOpt)
-          case None =>
-            Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
-        }
-      } yield searchResults
+      val outerMarker = logMarker
+
+      {
+        implicit val logMarker: LogMarker = outerMarker ++ Map(
+          "imageId" -> imageId,
+          "aiSearchType" -> "image"
+        )
+
+        val filterOpt = buildAiFilter(params)
+
+        for {
+          maybeImage <- elasticSearch.getImageById(imageId)
+          maybeEmbedding = maybeImage
+            .filter(image => isVisibleToAccessor(request.user, image))
+            .flatMap(_.embedding)
+            .flatMap(_.cohereEmbedV4)
+            .map(_.image)
+          searchResults <- maybeEmbedding match {
+            // If we have an embedding, perform the KNN search. If not, return an empty result set.
+            case Some(embedding) =>
+              elasticSearch.semanticSearch(embedding, k = k, numCandidates = Math.max(k * 2, 100), filterOpt = filterOpt)
+            case None =>
+              Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
+          }
+        } yield searchResults
+      }
     }
 
     def semanticSearchByText(k: Int, params: SearchParams): Future[SearchResults] = {
@@ -640,36 +649,42 @@ class MediaApi(
           logger.info(logMarker, s"No semantic query found in structured query; returning no AI results")
           Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
         case Some(semanticQuery) =>
-          // Normalise key so that "Dogs" and "dogs " share a cache entry.
-          val cacheKey = semanticQuery.trim.toLowerCase
-          val markers = logMarker ++ Map("query" -> semanticQuery)
-
-          if (embeddingCache.getIfPresent(cacheKey).isDefined) {
-            logger.info(markers, s"AI search embedding cache hit query=$semanticQuery")
-          } else {
-            logger.info(markers, s"AI search embedding cache miss query=$semanticQuery")
-          }
-
-          val weight = params.vecWeight.getOrElse(0.8)
-
-          // cache.get(key) is atomic: if two requests race on the same key, only one
-          // load fires and both callers receive the same Future.
-          val embeddingFuture = embeddingCache.get(cacheKey)
+          val vecWeight = params.vecWeight.getOrElse(0.8)
+          val outerMarker = logMarker
+
+          {
+            implicit val logMarker: LogMarker = outerMarker ++ Map(
+              "query" -> semanticQuery,
+              "vecWeight" -> vecWeight,
+              "aiSearchType" -> "hybrid"
+            )
 
-          val filterOpt = buildAiFilter(params)
+            // Normalise key so that "Dogs" and "dogs " share a cache entry.
+            val cacheKey = semanticQuery.trim.toLowerCase
+            if (embeddingCache.getIfPresent(cacheKey).isDefined) {
+              logger.info(logMarker, s"AI search embedding cache hit query=$semanticQuery")
+            } else {
+              logger.info(logMarker, s"AI search embedding cache miss query=$semanticQuery")
+            }
 
-          logger.info(markers, s"vecWeight for query '$semanticQuery' is $weight")
-          for {
-            embedding <- embeddingFuture
-            searchResults <- elasticSearch.hybridSearch(
-              query = semanticQuery,
-              queryEmbedding = embedding,
-              k = k,
-              numCandidates = Math.max(k * 2, 100),
-              vecWeight = weight,
-              filterOpt = filterOpt
-            )
-          } yield searchResults
+            // cache.get(key) is atomic: if two requests race on the same key, only one
+            // load fires and both callers receive the same Future.
+            val embeddingFuture = embeddingCache.get(cacheKey)
+
+            val filterOpt = buildAiFilter(params)
+
+            for {
+              embedding <- embeddingFuture
+              searchResults <- elasticSearch.hybridSearch(
+                query = semanticQuery,
+                queryEmbedding = embedding,
+                k = k,
+                numCandidates = Math.max(k * 2, 100),
+                vecWeight = vecWeight,
+                filterOpt = filterOpt
+              )
+            } yield searchResults
+          }
       }
     }