Add chips and filtering to AI Search

kahuna/public/js/components/gr-chips/gr-chips.js

@@ -27,7 +27,7 @@ grChips.controller('grChipsCtrl', ['$scope', function($scope) {
     $grChipsCtrl.configureNgModel = function(ngModelCtrl, onChangeExpr, autoCompleteExpr,
                                              validKeysExpr, autofocus, placeholder) {
         $grChipsCtrl.defaultPlaceholder = placeholder || 'Search for images... (type + for advanced search)';
-        $grChipsCtrl.aiSearchPlaceholder = 'Search for conceptual images using AI search... (no filters available)';
+        $grChipsCtrl.aiSearchPlaceholder = 'Search for conceptual images using AI search...';
 
 
         $grChipsCtrl.onChange = () => onChangeExpr($scope, {$chips: $grChipsCtrl.items});

kahuna/public/js/search/query.html

@@ -30,7 +30,7 @@
         <gr-icon>info_outline</gr-icon>
     </button>
 
-    <div ng-if="!searchQuery.usePermissionsFilter && !searchQuery.useAISearch" class="search__modifier-container">
+    <div ng-if="!searchQuery.usePermissionsFilter" class="search__modifier-container">
         <button class="search__modifier-toggle"
                 type="button"
                 ng-click="filtersShown = !filtersShown ">

kahuna/public/js/search/results.html

@@ -221,7 +221,7 @@
         <div ng-if="ctrl.totalResults == 0"
              class="image-no-results">
             <span ng-if="!ctrl.needsQuery">No matches, sorry! Try altering your search.</span>
-            <span ng-if="ctrl.needsQuery">Please enter a query</span>
+            <span ng-if="ctrl.needsQuery">Please enter a query. You can include filters, but must include either some text or a similar:image-id term.</span>
         </div>
     </div>
 </div>

kahuna/public/js/search/results.js

@@ -1,7 +1,6 @@
 import angular from 'angular';
 import Rx from 'rx';
 import moment from 'moment';
-
 import '../services/scroll-position';
 import '../services/panel';
 import '../util/async';

media-api/app/controllers/MediaApi.scala

@@ -16,6 +16,7 @@ import com.gu.mediaservice.lib.play.RequestLoggingFilter
 import com.gu.mediaservice.model._
 import com.gu.mediaservice.syntax.MessageSubjects
 import com.gu.mediaservice.{GridClient, JsonDiff}
+import com.sksamuel.elastic4s.requests.searches.queries.Query
 import lib._
 import lib.elasticsearch._
 import org.apache.http.entity.ContentType
@@ -575,16 +576,26 @@ class MediaApi(
     val canViewDeletedImages = _searchParams.query.contains("is:deleted") && !hasDeletePermission
 
     sealed trait AiSearchMode
-    case class TextSearch(query: String) extends AiSearchMode
+    case object TextSearch extends AiSearchMode
     case class SimilarSearch(imageId: String) extends AiSearchMode
 
-    def parseAiSearchMode(query: String): AiSearchMode = {
-      query
-        .split("\\s+")
-        .find(_.startsWith("similar:"))
-        .flatMap(_.split(":", 2).lift(1).filter(_.nonEmpty))
-        .map(SimilarSearch)
-        .getOrElse(TextSearch(query))
+    def parseAiSearchMode(params: SearchParams): AiSearchMode =
+      params.aiQueryParts.similarImageId.map(SimilarSearch).getOrElse(TextSearch)
+
+    def buildAiFilter(params: SearchParams): Option[Query] = {
+      val filterConditions = params.aiQueryParts.filterConditions
+
+      val chipFilterOpt =
+        if (filterConditions.nonEmpty) Some(elasticSearch.queryBuilder.makeQuery(filterConditions))
+        else None
+
+      val requestFilterOpt = elasticSearch.queryBuilder.buildFilterOpt(
+        params,
+        elasticSearch.searchFilters,
+        elasticSearch.syndicationFilter
+      )
+
+      elasticSearch.searchFilters.filterAndFilter(chipFilterOpt, requestFilterOpt)
     }
 
     def emptyAiSearchResponse =
@@ -601,7 +612,9 @@ class MediaApi(
       )
     }
 
-    def semanticSearchByImage(imageId: String, k: Int): Future[SearchResults] = {
+    def semanticSearchByImage(imageId: String, k: Int, params: SearchParams): Future[SearchResults] = {
+      val filterOpt = buildAiFilter(params)
+
       for {
         maybeImage <- elasticSearch.getImageById(imageId)
         maybeEmbedding = maybeImage
@@ -612,52 +625,68 @@ class MediaApi(
         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))
+            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
     }
 
-    def semanticSearchByText(query: String, k: Int, vecWeight: Option[Double]): Future[SearchResults] = {
-      // Normalise key so that "Dogs" and "dogs " share a cache entry.
-      val cacheKey = query.trim.toLowerCase
+    def semanticSearchByText(k: Int, params: SearchParams): Future[SearchResults] = {
+      // Separate the chips from the main query text
+      // So that we can embed just the query text
+      params.aiQueryParts.semanticQuery match {
+        case None =>
+          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 markers = logMarker ++ Map("query" -> query)
+          val weight = params.vecWeight.getOrElse(1.0)
 
-      if (embeddingCache.getIfPresent(cacheKey).isDefined) {
-        logger.info(markers, s"AI search embedding cache hit query=$query")
-      } else {
-        logger.info(markers, s"AI search embedding cache miss query=$query")
-      }
+          // 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 weight = vecWeight.getOrElse(1.0)
+          val filterOpt = buildAiFilter(params)
 
-      // 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)
-
-      logger.info(markers, s"vecWeight for query '$query' is $weight")
-      for {
-        embedding <- embeddingFuture
-        searchResults <- elasticSearch.hybridSearch(
-          query = query,
-          queryEmbedding = embedding,
-          k = k,
-          numCandidates = Math.max(k * 2, 100),
-          vecWeight = weight,
-        )
-      } yield searchResults
+          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
+      }
     }
 
-    def performAiSearchAndRespond(query: String, vecWeight: Option[Double]): Future[Result] = {
-      val k = config.aiSearchResultLimit
-      val searchResultsFuture = parseAiSearchMode(query) match {
-        case SimilarSearch(imageId) => semanticSearchByImage(imageId, k)
-        case TextSearch(textQuery) => semanticSearchByText(textQuery, k, vecWeight)
-      }
+    def performAiSearchAndRespond(params: SearchParams): Future[Result] = {
+      params.aiQueryParts.validationError match {
+        case Left(errorMessage) =>
+          logger.info(logMarker, s"Invalid AI search query: $errorMessage")
+          Future.successful(respondError(UnprocessableEntity, "invalid-ai-search", errorMessage))
+        case scala.util.Right(_) =>
+          val k = config.aiSearchResultLimit
+          val searchResultsFuture = parseAiSearchMode(params) match {
+            case SimilarSearch(imageId) => semanticSearchByImage(imageId, k, params)
+            case TextSearch => semanticSearchByText(k, params)
+          }
 
-      searchResultsFuture.map(aiSearchResponseFromResults)
+          searchResultsFuture.map(aiSearchResponseFromResults)
+      }
     }
 
     if (_searchParams.useAISearch.contains(true)) {
@@ -666,7 +695,7 @@ class MediaApi(
         case _ if _searchParams.length == 0 =>
           emptyAiSearchResponse
         case Some(q) if !q.isBlank =>
-          performAiSearchAndRespond(q, _searchParams.vecWeight)
+          performAiSearchAndRespond(_searchParams)
         // Empty queries do not make sense for AI search as we can
         // only rank results once we have a meaningful vector to compare with.
         // So return 0 results if the query was empty.

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

@@ -178,12 +178,12 @@ class ElasticSearch(
       }
   }
 
-  def knnSearch(queryEmbedding: List[Float], k: Int, numCandidates: Int)
+  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")
-        .queryVector(queryEmbedding.map(_.toDouble))
-        .k(k)
-        .numCandidates(numCandidates)
+    val knn = Knn("embedding.cohereEmbedV4.image", filter = filterOpt)
+      .queryVector(queryEmbedding.map(_.toDouble))
+      .k(k)
+      .numCandidates(numCandidates)
 
     val searchRequest = ElasticDsl.search(imagesCurrentAlias)
       .knn(knn)
@@ -211,9 +211,12 @@ class ElasticSearch(
   // than cosine similarity (knn). So we get the max BM25 score for the query and use that to calculate
   // the scaling factor for the lexical part of the query, so that BM25 and knn scores are both between 0-1 scale
   // and can be effectively combined in a hybrid query.
-  private def fetchMaxBm25Score(query: String)(implicit ex: ExecutionContext, logMarker: LogMarker): Future[Double] = {
+  private def fetchMaxBm25Score(query: String, filterOpt: Option[Query])(implicit ex: ExecutionContext, logMarker: LogMarker): Future[Double] = {
+    val baseQuery = createMultiMatchQuery(query)
+    val filteredQuery = filterOpt.map(filter => boolQuery().must(baseQuery).filter(filter)).getOrElse(baseQuery)
+
     val maxScoreRequest = ElasticDsl.search(imagesCurrentAlias)
-      .query(createMultiMatchQuery(query))
+      .query(filteredQuery)
 
     executeAndLog(withSearchQueryTimeout(maxScoreRequest), "max BM25 score").map { r =>
       logger.info(logMarker, s"Max BM25 score for query '$query' is ${r.result.hits.maxScore}")
@@ -227,9 +230,10 @@ class ElasticSearch(
     k: Int,
     numCandidates: Int,
     vecWeight: Double,
-    maxScore: Double
+    maxScore: Double,
+    filterOpt: Option[Query]
   )(implicit logMarker: LogMarker): SearchRequest = {
-    val knn = Knn("embedding.cohereEmbedV4.image")
+    val knn = Knn("embedding.cohereEmbedV4.image", filter = filterOpt)
       .queryVector(queryEmbedding)
       .k(k)
       .numCandidates(numCandidates)
@@ -242,17 +246,17 @@ class ElasticSearch(
     // BM25 score to bring it to the same range as the cosine similarity.
     val scalingFactor = if (maxScore > 0.0) 1.0 / maxScore else 1.0
 
-    //    We want to apply only one boost if we can help it, so we scale the
-    //    multi_match boost to be in line with the max_score and the desired
-    //    lexical_weight/vec_weight balance
+    // We want to apply only one boost if we can help it, so we scale the
+    // multi_match boost to be in line with the max_score and the desired
+    // lexical_weight/vec_weight balance
     val multiMatchBoost = if (vecWeight > 0.0) (lexicalWeight / vecWeight) * scalingFactor else 1.0
 
     logger.info(logMarker, s"Scaling factor for BM25 score is $scalingFactor, multi-match boost is $multiMatchBoost")
 
     val multiMatchQuery = createMultiMatchQuery(query, boost = Some(multiMatchBoost))
 
     ElasticDsl.search(imagesCurrentAlias)
-      .bool(BoolQuery().should(Seq(multiMatchQuery, knn)))
+      .bool(BoolQuery().should(Seq(multiMatchQuery, knn)).filter(filterOpt))
       .size(k)
   }
 
@@ -262,15 +266,16 @@ class ElasticSearch(
     k: Int,
     numCandidates: Int,
     vecWeight: Double,
+    filterOpt: Option[Query]
   )(
     implicit ex: ExecutionContext,
     logMarker: LogMarker
   ): Future[SearchResults] = {
     val queryEmbeddingDouble: List[Double] = queryEmbedding.map(_.toDouble)
 
     for {
-      maxScore <- fetchMaxBm25Score(query)
-      searchRequest = makeHybridSearchRequest(query, queryEmbeddingDouble, k, numCandidates, vecWeight, maxScore)
+      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))
@@ -281,84 +286,7 @@ class ElasticSearch(
   def search(params: SearchParams)(implicit ex: ExecutionContext, request: AuthenticatedRequest[AnyContent, Principal], logMarker: LogMarker = MarkerMap()): Future[SearchResults] = {
     val query: Query = queryBuilder.makeQuery(params.structuredQuery)
 
-    val uploadTimeFilter = filters.date("uploadTime", params.since, params.until)
-    val lastModTimeFilter = filters.date("lastModified", params.modifiedSince, params.modifiedUntil)
-    val takenTimeFilter = filters.date("metadata.dateTaken", params.takenSince, params.takenUntil)
-    // we only inject filters if there are actual date parameters
-    val dateFilterList = List(uploadTimeFilter, lastModTimeFilter, takenTimeFilter).flatten.toNel
-    val dateFilter = dateFilterList.map(dateFilters => filters.and(dateFilters.list.toList: _*))
-
-    val idsFilter = params.ids.map(filters.ids)
-    val labelFilter = params.labels.toNel.map(filters.terms("labels", _))
-    val metadataFilter = params.hasMetadata.map(metadataField).toNel.map(filters.exists)
-    val archivedFilter = params.archived.map(filters.existsOrMissing(editsField("archived"), _))
-    val hasExports = params.hasExports.map(filters.existsOrMissing("exports", _))
-    val hasIdentifier = params.hasIdentifier.map(idName => filters.exists(NonEmptyList(identifierField(idName))))
-    val missingIdentifier = params.missingIdentifier.map(idName => filters.missing(NonEmptyList(identifierField(idName))))
-    val uploadedByFilter = params.uploadedBy.map(uploadedBy => filters.terms("uploadedBy", NonEmptyList(uploadedBy)))
-    val simpleCostFilter = params.free.flatMap(free => if (free) searchFilters.freeFilter else searchFilters.nonFreeFilter)
-    val costFilter = params.payType match {
-      case Some(PayType.Free) => searchFilters.freeFilter
-      case Some(PayType.MaybeFree) => searchFilters.maybeFreeFilter
-      case Some(PayType.Pay) => searchFilters.nonFreeFilter
-      case _ => None
-    }
-
-    val printUsageFilter = params.printUsageFilters.map(searchFilters.printUsageFilters)
-
-    val hasRightsCategory = params.hasRightsCategory.filter(_ == true).map(_ => searchFilters.hasRightsCategoryFilter)
-
-    val validityFilter = params.valid.map(valid => if (valid) searchFilters.validFilter else searchFilters.invalidFilter)
-
-    val persistFilter = params.persisted map {
-      case true => searchFilters.persistedFilter
-      case false => searchFilters.nonPersistedFilter
-    }
-
-    val usageFilter: Iterable[Query] =
-      params.usageStatus.toNel.map(status => filters.terms("usagesStatus", status.map(_.toString))).toOption ++
-        params.usagePlatform.toNel.map(filters.terms("usagesPlatform", _)).toOption
-
-    val syndicationStatusFilter = params.syndicationStatus.map(status => syndicationFilter.statusFilter(status))
-
-    // Port of special case code in elastic1 sorts. Using the dateAddedToCollection sort implies an additional filter for reasons unknown
-    val dateAddedToCollectionFilter = {
-      params.orderBy match {
-        case Some("dateAddedToCollection") => {
-          val pathHierarchyOpt = params.structuredQuery.flatMap {
-            case Match(HierarchyField, Phrase(value)) => Some(value)
-            case _ => None
-          }.headOption
-
-          pathHierarchyOpt.map { pathHierarchy =>
-            termQuery("collections.pathHierarchy", pathHierarchy)
-          }
-        }
-        case _ => None
-      }
-    }
-
-    val filterOpt = (
-      metadataFilter.toOption.toList
-        ++ persistFilter
-        ++ labelFilter.toOption
-        ++ archivedFilter
-        ++ uploadedByFilter
-        ++ idsFilter
-        ++ validityFilter
-        ++ simpleCostFilter
-        ++ costFilter
-        ++ hasExports
-        ++ hasIdentifier
-        ++ missingIdentifier
-        ++ dateFilter.toOption
-        ++ usageFilter
-        ++ hasRightsCategory
-        ++ searchFilters.tierFilter(params.tier)
-        ++ syndicationStatusFilter
-        ++ dateAddedToCollectionFilter
-        ++ printUsageFilter
-      ).toNel.map(filter => filter.list.toList.reduceLeft(filters.and(_, _))).toOption
+    val filterOpt: Option[Query] = queryBuilder.buildFilterOpt(params, searchFilters, syndicationFilter)
 
     val withFilter = filterOpt.map { f =>
       boolQuery() must (query) filter f