Merge pull request #4785 from guardian/js-add-aggregations-and-filter-feedback

Add GNM-owned, agency picks and filtered matches to AI search

Author: joelochlann

kahuna/public/js/search/results.html

@@ -22,7 +22,8 @@
 
             <div class="results-toolbar-item results-toolbar-item--static
                     image-results-count">
-                {{ctrl.totalResults | toLocaleString}} matches
+                <span ng-if="ctrl.isAiSearch">Best {{ctrl.aiResultsShown | toLocaleString}} of {{ctrl.totalResults | toLocaleString}} matches</span>
+                <span ng-if="! ctrl.isAiSearch">{{ctrl.totalResults | toLocaleString}} matches</span>
                 <button ng-if="ctrl.newImagesCount > 0"
                         ng-click="ctrl.revealNewImages()"
                         class="image-results-count__new"

kahuna/public/js/search/results.js

@@ -248,6 +248,9 @@ results.controller('SearchResultsCtrl', [
         function initialiseAiResults(images) {
           const totalLength = images.data.length;
           ctrl.imagesAll = new Array(totalLength);
+          // Number of results actually shown (the top k), so we can display
+          // "Best k of N matches" where N is ctrl.totalResults.
+          ctrl.aiResultsShown = totalLength;
 
           // AI search returns a single fixed result set rather than a paged/lazy-loaded one,
           // so we populate the full backing array up front to avoid placeholder rows.
@@ -293,6 +296,8 @@ results.controller('SearchResultsCtrl', [
             checkForNewImages();
           }
 
+          ctrl.isAiSearch = isAiSearch;
+
           updateLastSearchBoundary();
 
           return images;

media-api/app/controllers/MediaApi.scala

@@ -611,7 +611,7 @@ class MediaApi(
         data = imageEntities,
         offset = Some(0),
         total = Some(searchResults.total),
-        maybeExtraCounts = None,
+        maybeExtraCounts = searchResults.extraCounts,
         links = Nil
       )
     }

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

@@ -341,7 +341,16 @@ class ElasticSearch(
         case _ => fusedLexicalAndSemanticSearch(query, queryEmbedding, k, numCandidates, vecWeight, filterOpt)
       }
 
-      searchResults.andThen { case _ =>
+      // Run in parallel: how many images match the filters in total (so we can
+      // show "Best k of N matches"). The ticker count badges are computed over
+      // the top-k results we actually return, not the whole filtered set.
+      val filterTotal = countMatchingFilter(filterOpt)
+
+      (for {
+        results <- searchResults
+        total <- filterTotal
+        extraCounts <- extraCountsForIds(results.hits.map(_._1))
+      } yield results.copy(total = total, extraCounts = Some(extraCounts))).andThen { case _ =>
         val elapsed = stopwatch.elapsed
         logger.info(
           combineMarkers(logMarker, elapsed),
@@ -351,6 +360,32 @@ class ElasticSearch(
     }
   }
 
+  // How many images match the active filters in total, so we can show
+  // "Best k of N matches".
+  private def countMatchingFilter(filterOpt: Option[Query])(implicit ex: ExecutionContext, logMarker: LogMarker): Future[Long] = {
+    val countRequest = ElasticDsl.count(imagesCurrentAlias).query(filterOpt.getOrElse(matchAllQuery()))
+
+    executeAndLog(countRequest, "hybrid AI search filter count").map(_.result.count)
+  }
+
+  // The ticker count badges restricted to the given (top-k) result ids.
+  // In principle, we could compute this here, without an extra request.
+  // But it would require awkward duplication of the search logic
+  // from `maybeOrgOwnedExtraCount` and `maybeAgencyPicksExtraCount`
+  private def extraCountsForIds(ids: Seq[String])(implicit ex: ExecutionContext, logMarker: LogMarker): Future[ExtraCounts] = {
+    if (ids.isEmpty) Future.successful(ExtraCounts(tickerCounts = Map.empty))
+    else {
+      val searchRequest = ElasticDsl.search(imagesCurrentAlias)
+        .query(filters.ids(ids.toList))
+        .size(0)
+        .aggregations(extraCountAggregations)
+
+      executeAndLog(withSearchQueryTimeout(searchRequest), "hybrid AI search ticker counts").map { r =>
+        extraCountsFrom(r.result.aggregations)
+      }
+    }
+  }
+
   def search(params: SearchParams)(implicit ex: ExecutionContext, request: AuthenticatedRequest[AnyContent, Principal], logMarker: LogMarker = MarkerMap()): Future[SearchResults] = {
     val query: Query = queryBuilder.makeQuery(params.structuredQuery)
 
@@ -391,10 +426,7 @@ class ElasticSearch(
       .runtimeMappings(runtimeMappings)
       .storedFields("_source") // this needs to be explicit when using script fields
       .scriptfields(graphicImagesScriptFields)
-      .aggregations(aggregationsNameToSearchClauseMap.map {
-        case (name, ExtraCountConfig(searchClause, _, maybeSubAggregation)) =>
-          filterAgg(name, queryBuilder.makeQuery(Parser.run(searchClause))).subAggregations(maybeSubAggregation)
-      })
+      .aggregations(extraCountAggregations)
       .from(params.offset)
       .size(params.length)
       .sortBy(sort)
@@ -408,27 +440,36 @@ class ElasticSearch(
       SearchResults(
         hits = imageHits,
         total = if (trackTotalHits) r.result.totalHits else 0,
-        extraCounts = Some(ExtraCounts(
-          tickerCounts = aggregationsNameToSearchClauseMap.map {
-            case (name, ExtraCountConfig(searchClause, backgroundColour, maybeSubAggregation)) =>
-              val aggResult = r.result.aggregations.filter(name)
-              val maybeSubAggResult = maybeSubAggregation.map(_.name).map(aggResult.result[Terms])
-              name -> ExtraCount(
-                value = aggResult.docCount,
-                searchClause,
-                backgroundColour,
-                subCounts = maybeSubAggResult.map { termsAgg =>
-                  ListMap(termsAgg.buckets.sortBy(_.docCount).reverse.map { bucket =>
-                    (bucket.key, bucket.docCount)
-                  }: _*) + ("other" -> termsAgg.otherDocCount)
-                }.filter(_.exists { case (_, count) => count > 0 })
-              )
-          }
-        ))
+        extraCounts = Some(extraCountsFrom(r.result.aggregations))
       )
     }
   }
 
+  // The aggregations used to compute the ticker count badges ("GNM-owned",
+  // "agency picks" etc) shown above search results.
+  private def extraCountAggregations = aggregationsNameToSearchClauseMap.map {
+    case (name, ExtraCountConfig(searchClause, _, maybeSubAggregation)) =>
+      filterAgg(name, queryBuilder.makeQuery(Parser.run(searchClause))).subAggregations(maybeSubAggregation)
+  }
+
+  private def extraCountsFrom(aggregations: Aggregations): ExtraCounts = ExtraCounts(
+    tickerCounts = aggregationsNameToSearchClauseMap.map {
+      case (name, ExtraCountConfig(searchClause, backgroundColour, maybeSubAggregation)) =>
+        val aggResult = aggregations.filter(name)
+        val maybeSubAggResult = maybeSubAggregation.map(_.name).map(aggResult.result[Terms])
+        name -> ExtraCount(
+          value = aggResult.docCount,
+          searchClause,
+          backgroundColour,
+          subCounts = maybeSubAggResult.map { termsAgg =>
+            ListMap(termsAgg.buckets.sortBy(_.docCount).reverse.map { bucket =>
+              (bucket.key, bucket.docCount)
+            }: _*) + ("other" -> termsAgg.otherDocCount)
+          }.filter(_.exists { case (_, count) => count > 0 })
+        )
+    }
+  )
+
   def usageForSupplier(id: String, numDays: Int)(implicit ex: ExecutionContext, logMarker: LogMarker): Future[SupplierUsageSummary] = {
     val supplier = Agencies.get(id)
     val supplierName = supplier.supplier