Skip to content
Merged
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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).
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down Expand Up @@ -75,8 +75,7 @@ object Mappings {
leasesMapping("leases"),
collectionMapping("collections"),
esInfoMapping("esInfo"),
embeddingMapping("embedding"),
)
) ++ (if (includeDenseVectorMappings) Seq(embeddingMapping("embedding")) else Seq.empty)
)
}

Expand Down
4 changes: 3 additions & 1 deletion dev/script/generate-config/service-config.js
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,7 @@ function getKahunaConfig(config){
|]`;

const permissionsConfig = stripMargin`
|usePermissionsFilter=true
|usePermissionsFilter=false
Comment thread
AndyKilmory marked this conversation as resolved.
|permissionsDefault="allPermissions"
`;

Expand Down Expand Up @@ -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={
Expand Down Expand Up @@ -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
|`;
}
Expand Down
49 changes: 30 additions & 19 deletions media-api/app/lib/elasticsearch/ElasticSearch.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
}
}
}

Expand Down Expand Up @@ -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)
}
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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)
}
}
Expand Down Expand Up @@ -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
Expand Down
33 changes: 22 additions & 11 deletions thrall/app/lib/elasticsearch/ElasticSearch.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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](
Expand Down Expand Up @@ -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)
Expand All @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Loading