go back to unpaginated

ellenmuller · ellenmuller · commit 8cf74cfcb463 · 2026-05-01T10:12:54.000+01:00
diff --git a/kahuna/public/js/search/results.js b/kahuna/public/js/search/results.js
@@ -229,55 +229,11 @@ results.controller('SearchResultsCtrl', [
             lastSearchFirstResultTime = undefined;
         }
 
-        // Detect whether this results page is using semantic search.
-        // We branch the initial loading behaviour because AI search should
-        // bootstrap with a real first page of results rather than the
-        // normal `length: 1` probe used by lexical search.
-        const isAiSearch = !!$stateParams.useAISearch;
-
-        // Shared helper for writing returned images into both the sparse
-        // backing array (`imagesAll`) and the reactive `results` list.
-        // `start` is the offset of the returned page in the overall result set.
-        function applyLoadedImages(images, start = 0) {
-          images.data.forEach((image, index) => {
-            const position = index + start;
-            const imageId = image.data.id;
-
-            // If image already present in results at a
-            // different position (result set shifted due to
-            // items being spliced in or deleted?), get rid of
-            // item at its previous position to avoid
-            // duplicates
-            const existingPosition = imagesPositions.get(imageId);
-            if (angular.isDefined(existingPosition) &&
-              existingPosition !== position) {
-              $log.info(`Detected duplicate image ${imageId}, ` +
-                    `old ${existingPosition}, new ${position}`);
-              delete ctrl.imagesAll[existingPosition];
-
-              results.set(existingPosition, undefined);
-            }
-
-            ctrl.imagesAll[position] = image;
-            imagesPositions.set(imageId, position);
-
-            results.set(position, image);
-          });
-          // images should not contain any 'holes'
-          ctrl.images = compact(ctrl.imagesAll);
-        }
-
         // Initial search to find upper `until` boundary of result set
         // (i.e. the uploadTime of the newest result in the set)
 
         // TODO: avoid this initial search (two API calls to init!)
-        // For lexical search we keep the old lightweight probe (`length: 1`) so
-        // we can establish totals cheaply before lazy-loading visible ranges.
-        // For AI search that probe is harmful because it makes the backend ask
-        // KNN for only one neighbour, which in turn poisons the displayed total.
-        // So AI search performs a normal first-page request here instead.
-        const initialSearch = isAiSearch ? search() : search({length: 1, orderBy: 'newest'});
-        ctrl.searched = initialSearch.then(function(images) {
+        ctrl.searched = search({length: 1, orderBy: 'newest'}).then(function(images) {
             ctrl.totalResults = images.total;
             // FIXME: https://github.com/argo-rest/theseus has forced us to co-opt the actions field for this
             ctrl.tickerCounts = images.$response?.$$state?.value?.actions?.tickerCounts;
@@ -301,16 +257,8 @@ results.controller('SearchResultsCtrl', [
 
             notificationMessages(ctrl.extendedSortProps, images.total);
 
-            // Initialise the map before we merge in either the initial AI page or
-            // any later lazy-loaded ranges.
             imagesPositions = new Map();
 
-            if (isAiSearch) {
-              // Seed the sparse result arrays with the first AI page immediately.
-              // This keeps the match count and initially rendered images aligned.
-              applyLoadedImages(images);
-            }
-
             checkForNewImages();
 
             // Keep track of time of the latest result for all
@@ -334,13 +282,35 @@ results.controller('SearchResultsCtrl', [
         });
 
         ctrl.loadRange = function(start, end) {
-          // Lazy-table gives us an inclusive start/end pair, so convert it to
-          // the page size expected by the API.
             const length = end - start + 1;
             search({offset: start, length: length, countAll: false}).then(images => {
-            // Merge the fetched page into the results array at the correct offset,
-            // deduplicating any images that may have shifted position since the initial load.
-            applyLoadedImages(images, start);
+            // Update imagesAll with newly loaded images
+                images.data.forEach((image, index) => {
+                    const position = index + start;
+                    const imageId = image.data.id;
+
+                    // If image already present in results at a
+                    // different position (result set shifted due to
+                    // items being spliced in or deleted?), get rid of
+                    // item at its previous position to avoid
+                    // duplicates
+                    const existingPosition = imagesPositions.get(imageId);
+                    if (angular.isDefined(existingPosition) &&
+                        existingPosition !== position) {
+                        $log.info(`Detected duplicate image ${imageId}, ` +
+                                  `old ${existingPosition}, new ${position}`);
+                        delete ctrl.imagesAll[existingPosition];
+
+                        results.set(existingPosition, undefined);
+                    }
+
+                    ctrl.imagesAll[position] = image;
+                    imagesPositions.set(imageId, position);
+
+                    results.set(position, image);
+                });
+                // images should not contain any 'holes'
+                ctrl.images = compact(ctrl.imagesAll);
             });
         };
 
diff --git a/media-api/app/controllers/MediaApi.scala b/media-api/app/controllers/MediaApi.scala
@@ -586,31 +586,21 @@ class MediaApi(
         .getOrElse(TextSearch(query))
     }
 
-    def emptyAiSearchResponse(searchParams: SearchParams) =
-      Future.successful(
-        respondCollection(
-          List.empty[EmbeddedEntity[JsValue]],
-          Some(searchParams.offset),
-          Some(0),
-          None,
-          List()
-        )
-      )
+    def emptyAiSearchResponse =
+      Future.successful(respondCollection(List.empty[EmbeddedEntity[JsValue]], Some(0), Some(0), None, List()))
 
-    def aiSearchResponseFromResults(searchParams: SearchParams, searchResults: SearchResults): Result = {
+    def aiSearchResponseFromResults(searchResults: SearchResults): Result = {
       val imageEntities = searchResults.hits map (hitToImageEntity _).tupled
       respondCollection(
         data = imageEntities,
-        offset = Some(searchParams.offset),
+        offset = Some(0),
         total = Some(searchResults.total),
         maybeExtraCounts = None,
         links = Nil
       )
     }
 
-    def semanticSearchByImage(searchParams: SearchParams, imageId: String): Future[SearchResults] = {
-      val semanticRetrievalWindow = Math.max(config.aiSearchSemanticRetrievalWindow, searchParams.length)
-      val countAll = searchParams.countAll.getOrElse(true)
+    def semanticSearchByImage(imageId: String, k: Int): Future[SearchResults] = {
       for {
         maybeImage <- elasticSearch.getImageById(imageId)
         maybeEmbedding = maybeImage
@@ -621,23 +611,14 @@ 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 = semanticRetrievalWindow,
-              offset = searchParams.offset,
-              length = searchParams.length,
-              countAll = countAll,
-              numCandidates = Math.max(semanticRetrievalWindow * 2, 100)
-            )
+            elasticSearch.knnSearch(embedding, k = k, numCandidates = Math.max(k * 2, 100))
           case None =>
             Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
         }
       } yield searchResults
     }
 
-    def semanticSearchByText(searchParams: SearchParams, query: String): Future[SearchResults] = {
-      val semanticRetrievalWindow = Math.max(config.aiSearchSemanticRetrievalWindow, searchParams.length)
-      val countAll = searchParams.countAll.getOrElse(true)
+    def semanticSearchByText(query: String, k: Int): Future[SearchResults] = {
       // Normalise key so that "Dogs" and "dogs " share a cache entry.
       val cacheKey = query.trim.toLowerCase
 
@@ -655,56 +636,47 @@ class MediaApi(
 
       for {
         embedding <- embeddingFuture
-        searchResults <- elasticSearch.knnSearch(
-          embedding,
-          k = semanticRetrievalWindow,
-          offset = searchParams.offset,
-          length = searchParams.length,
-          countAll = countAll,
-          numCandidates = Math.max(semanticRetrievalWindow * 2, 100)
-        )
+        searchResults <- elasticSearch.knnSearch(embedding, k = k, numCandidates = Math.max(k * 2, 100))
       } yield searchResults
     }
 
-    def performAiSearchAndRespond(searchParams: SearchParams, query: String): Future[Result] = {
+    def performAiSearchAndRespond(query: String): Future[Result] = {
+      val k = 200
       val searchResultsFuture = parseAiSearchMode(query) match {
-        case SimilarSearch(imageId) => semanticSearchByImage(searchParams, imageId)
-        case TextSearch(textQuery) => semanticSearchByText(searchParams, textQuery)
+        case SimilarSearch(imageId) => semanticSearchByImage(imageId, k)
+        case TextSearch(textQuery) => semanticSearchByText(textQuery, k)
       }
 
-      searchResultsFuture.map(aiSearchResponseFromResults(searchParams, _))
+      searchResultsFuture.map(aiSearchResponseFromResults)
     }
 
-    val searchParams = if(canViewDeletedImages) {
-      _searchParams.copy(uploadedBy = Some(Authentication.getIdentity(request.user)))
-    }
-    else {
-      _searchParams
-    }
-
-    SearchParams.validate(searchParams).fold(
-      errors => Future.successful(respondError(UnprocessableEntity, InvalidUriParams.errorKey,
-        errors.map(_.message).mkString(", "))
-      ),
-      params => {
-        if (params.useAISearch.contains(true)) {
-          params.query match {
-            // Short-circuit polling requests (length=0) and empty queries to avoid unnecessary Bedrock/KNN calls
-            case _ if params.length == 0 =>
-              emptyAiSearchResponse(params)
-            case Some(q) if !q.isBlank =>
-              performAiSearchAndRespond(params, q)
-            // 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.
-            case _ =>
-              emptyAiSearchResponse(params)
-          }
-        } else {
-          performSearchAndRespond(params)
-        }
+    if (_searchParams.useAISearch.contains(true)) {
+      _searchParams.query match {
+        // Short-circuit polling requests (length=0) and empty queries to avoid unnecessary Bedrock/KNN calls
+        case _ if _searchParams.length == 0 =>
+          emptyAiSearchResponse
+        case Some(q) if !q.isBlank =>
+          performAiSearchAndRespond(q)
+        // 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.
+        case _ =>
+          emptyAiSearchResponse
       }
-    )
+    } else {
+      val searchParams = if (canViewDeletedImages) {
+        _searchParams.copy(uploadedBy = Some(Authentication.getIdentity(request.user)))
+      } else {
+        _searchParams
+      }
+      SearchParams.validate(searchParams).fold(
+        // TODO: respondErrorCollection?
+        errors => Future.successful(respondError(UnprocessableEntity, InvalidUriParams.errorKey,
+          errors.map(_.message).mkString(", "))
+        ),
+        params => performSearchAndRespond(params)
+      )
+    }
   }
 
   private def getImageResponseFromES(
diff --git a/media-api/app/lib/MediaApiConfig.scala b/media-api/app/lib/MediaApiConfig.scala
@@ -76,9 +76,6 @@ class MediaApiConfig(resources: GridConfigResources) extends CommonConfigWithEla
 
   val queueUrl: String = stringOpt("sqs.embedder.queue.url").getOrElse("")
 
-  // Semantic search retrieves from a bounded nearest-neighbour window.
-  // This keeps KNN query cost predictable and makes result totals explicitly capped.
-  val aiSearchSemanticRetrievalWindow: Int = intOpt("ai.search.semanticRetrievalWindow").getOrElse(1000)
   val aiSearchEmbeddingCacheMaxSize: Int = intOpt("ai.search.embeddingCache.maxSize").getOrElse(500)
 
   val maybeAgencyPickQuery: Option[Query] = agencyPicksIngredients.map { ingredients =>
diff --git a/media-api/app/lib/elasticsearch/ElasticSearch.scala b/media-api/app/lib/elasticsearch/ElasticSearch.scala
@@ -174,40 +174,20 @@ class ElasticSearch(
       }
   }
 
-  def knnSearch(queryEmbedding: List[Float], k: Int, offset: Int, length: Int, countAll: Boolean, numCandidates: Int)
+  def knnSearch(queryEmbedding: List[Float], k: Int, numCandidates: Int)
                (implicit ex: ExecutionContext, logMarker: LogMarker): Future[SearchResults] = {
-    if (length <= 0 || k <= 0) {
-      Future.successful(SearchResults(Nil, total = 0, extraCounts = None))
-    } else {
-      val knn = Knn("embedding.cohereEmbedV4.image")
+    val knn = Knn("embedding.cohereEmbedV4.image")
         .queryVector(queryEmbedding.map(_.toDouble))
         .k(k)
         .numCandidates(numCandidates)
 
-      val searchRequest = ElasticDsl.search(imagesCurrentAlias)
-        .trackTotalHits(countAll)
-        .knn(knn)
-        .from(offset)
-        .size(length)
+    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))
-        // In semantic search we expose a capped/approximate total based on the
-        // configured retrieval window (`k`) rather than an exact global corpus size.
-        // This aligns pagination semantics with the bounded nearest-neighbour set.
-        // Note: when trackTotalHits(false), ES may not populate totalHits.
-        val total = if (countAll) {
-          val totalHits = Option(r.result.totalHits).getOrElse(0L)
-          Math.min(totalHits, k.toLong)
-        } else {
-          0L
-        }
-        SearchResults(
-          hits = imageHits,
-          total = total,
-          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 = r.result.totalHits, extraCounts = None)
     }
   }
 
