Thrall embedding

build.sbt

@@ -52,7 +52,7 @@ val commonSettings = Seq(
     "org.scalatestplus" %% "mockito-3-4" % "3.1.4.0" % Test,
     "org.mockito" % "mockito-core" % "2.18.0" % Test,
     "org.scalamock" %% "scalamock" % "5.1.0" % Test,
-    "org.testcontainers" % "localstack" % "1.21.4" % Test
+    "org.testcontainers" % "testcontainers-localstack" % "2.0.5" % Test
   ),
   dependencyOverrides ++= jacksonOverrides,
 
@@ -124,10 +124,10 @@ lazy val commonLib = project("common-lib").settings(
     "software.amazon.awssdk" % "cloudwatch" % awsSdkV2Version,
     // declare explicit dependency on desired version of aws sdk v2 bedrock runtime
     "software.amazon.awssdk" % "bedrockruntime" % awsSdkV2Version,
-    "software.amazon.awssdk" % "s3vectors" % awsSdkV2Version,
     "com.adobe.xmp" % "xmpcore" % "6.1.11",
     ws,
-    "org.testcontainers" % "testcontainers-elasticsearch" % "2.0.2" % Test,
+    "com.google.genai" % "google-genai" % "1.53.0" exclude("com.fasterxml.jackson.module", "jackson-module-kotlin"),
+    "org.testcontainers" % "testcontainers-elasticsearch" % "2.0.5" % Test,
   ),
   dependencyOverrides += "ch.qos.logback" % "logback-classic" % "1.2.13" % Test
 )
@@ -185,10 +185,14 @@ lazy val thrall = playProject("thrall", 9002)
       "com.gu" %% "kcl-pekko-stream" % "0.1.2",
       "org.testcontainers" % "testcontainers-elasticsearch" % "2.0.2" % Test,
       "com.google.protobuf" % "protobuf-java" % "3.19.6",
-      "software.amazon.awssdk" % "sqs" % awsSdkV2Version
+      "software.amazon.awssdk" % "sqs" % awsSdkV2Version,
+      "org.apache.pekko" %% "pekko-connectors-sqs" % "1.0.2"
     ),
     dependencyOverrides ++= Seq(
-      "org.apache.pekko" %% "pekko-stream" % "1.0.3"
+      "org.apache.pekko" %% "pekko-stream" % "1.0.3",
+      "org.apache.pekko" %% "pekko-http" % "1.0.1",
+      "org.apache.pekko" %% "pekko-http-core" % "1.0.1",
+      "org.apache.pekko" %% "pekko-parsing" % "1.0.1"
     )
   )
 

cloudbuild.yaml

@@ -1,5 +1,6 @@
 options:
   machineType: 'N1_HIGHCPU_8'
+timeout: '1200s'
 steps:
   - name: 'node:24-alpine'
     entrypoint: 'npm'
@@ -15,9 +16,11 @@ steps:
     args: [ 'run', 'dist' ]
 
   - name: 'gcr.io/$PROJECT_ID/scala-sbt:1.11.7-jdk-25'
-    args: ['docker:publishLocal']
+    args: ['compile', 'collections/test', 'leases/test', 'media-api/test', 'docker:publishLocal']
     env:
-    - 'DOCKER_API_VERSION=1.41'
+     - 'DOCKER_HOST=unix:///var/run/docker.sock'
+     - 'DOCKER_API_VERSION=1.41'
+
   - name: 'gcr.io/cloud-builders/docker'
     args: ['tag', 'auth:0.1', 'eu.gcr.io/$PROJECT_ID/auth:$BRANCH_NAME']
   - name: 'gcr.io/cloud-builders/docker'

collections/test/store/CollectionsStoreTest.scala

@@ -7,8 +7,7 @@ import org.scalatest.concurrent.ScalaFutures
 import org.scalatest.funspec.AnyFunSpec
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.time.{Millis, Seconds, Span}
-import org.testcontainers.containers.localstack.LocalStackContainer
-import org.testcontainers.containers.localstack.LocalStackContainer.Service.DYNAMODB
+import org.testcontainers.localstack.LocalStackContainer
 import org.testcontainers.utility.DockerImageName
 import software.amazon.awssdk.auth.credentials.{AwsBasicCredentials, StaticCredentialsProvider}
 import software.amazon.awssdk.regions.Region
@@ -25,11 +24,11 @@ class CollectionsStoreTest extends AnyFunSpec with Matchers with ScalaFutures wi
 
   implicit val defaultPatience: PatienceConfig = PatienceConfig(timeout = Span(2, Seconds), interval = Span(100, Millis))
 
-  private val dynamoContainer = new LocalStackContainer(DockerImageName.parse("localstack/localstack:1.4.0")).withServices(DYNAMODB)
+  private val dynamoContainer = new LocalStackContainer(DockerImageName.parse("localstack/localstack:1.4.0")).withServices("dynamodb")
   dynamoContainer.start()
 
   private val dynamoClient = DynamoDbAsyncClient.builder().
-    endpointOverride(dynamoContainer.getEndpointOverride(DYNAMODB)).
+    endpointOverride(dynamoContainer.getEndpoint).
     region(Region.of(dynamoContainer.getRegion)).
     credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create(dynamoContainer.getAccessKey, dynamoContainer.getSecretKey))).build()
 

collections/test/store/ImageCollectionsStoreTest.scala

@@ -7,8 +7,7 @@ import org.scalatest.concurrent.ScalaFutures
 import org.scalatest.funspec.AnyFunSpec
 import org.scalatest.matchers.should.Matchers
 import org.scalatest.time.{Millis, Seconds, Span}
-import org.testcontainers.containers.localstack.LocalStackContainer
-import org.testcontainers.containers.localstack.LocalStackContainer.Service.DYNAMODB
+import org.testcontainers.localstack.LocalStackContainer
 import org.testcontainers.utility.DockerImageName
 import software.amazon.awssdk.auth.credentials.{AwsBasicCredentials, StaticCredentialsProvider}
 import software.amazon.awssdk.regions.Region
@@ -22,11 +21,11 @@ class ImageCollectionsStoreTest extends AnyFunSpec with Matchers with ScalaFutur
 
   implicit val defaultPatience: PatienceConfig = PatienceConfig(timeout = Span(5, Seconds), interval = Span(500, Millis))
 
-  private val dynamoContainer = new LocalStackContainer(DockerImageName.parse("localstack/localstack:1.4.0")).withServices(DYNAMODB)
+  private val dynamoContainer = new LocalStackContainer(DockerImageName.parse("localstack/localstack:1.4.0")).withServices("dynamodb")
   dynamoContainer.start()
 
   private val dynamoClient = DynamoDbAsyncClient.builder().
-    endpointOverride(dynamoContainer.getEndpointOverride(DYNAMODB)).
+    endpointOverride(dynamoContainer.getEndpoint).
     region(Region.of(dynamoContainer.getRegion)).
     credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create(dynamoContainer.getAccessKey, dynamoContainer.getSecretKey))).build()
 

common-lib/src/main/resources/application.conf

@@ -33,7 +33,10 @@ authentication.providers {
   machine {
     className = "com.gu.mediaservice.lib.auth.provider.ApiKeyAuthenticationProvider"
     config {
-      # authKeyStoreBucket = <s3-bucket-with-api-keys>
+      # authKeyStoreBucket.name = <s3-bucket-with-api-keys>
+      # authKeyStoreBucket.endpoint = <optional-s3-endpoint-override>
+      # authKeyStoreBucket.pathStyleUrls = false
+      # authKeyStoreBucket.region = <optional-region-override>
     }
   }
   # TODO: short term we put panda here for backwards compatibility but the default provider should be something better

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

@@ -150,8 +150,9 @@ class GridClient(services: Services, originDomain: Instance => String)(implicit
 
   def getImageLoaderProjection(mediaId: String, imageLoaderEndpoint: String, authFn: WSRequest => WSRequest)
                               (implicit ec: ExecutionContext): Future[Option[Image]] = {
-    logger.info("attempt to get image projection from image-loader")
-    val url = new URL(s"$imageLoaderEndpoint/images/project/$mediaId")
+    val projectUrl = s"$imageLoaderEndpoint/images/project/$mediaId"
+    logger.info(s"attempt to get image projection from image-loader: $projectUrl")
+    val url = new URL(projectUrl)
     makeGetRequestAsync(url, authFn, requestTimeout = Some(300.seconds)) map {
       case Found(json, _) => Some(json.as[Image])
       case NotFound(_, _) => None

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

@@ -1,12 +1,14 @@
 package com.gu.mediaservice.lib
 
-
+import _root_.play.api.libs.json._
 import com.gu.mediaservice.lib.aws.{S3Bucket, S3Object}
 import com.gu.mediaservice.lib.config.CommonConfig
 import com.gu.mediaservice.lib.logging.LogMarker
-import com.gu.mediaservice.model.{Instance, MimeType}
+import com.gu.mediaservice.model.{Embedding, Instance, MimeType}
 import com.typesafe.scalalogging.StrictLogging
 import org.joda.time.DateTime
+import software.amazon.awssdk.core.ResponseInputStream
+import software.amazon.awssdk.services.s3.model.GetObjectResponse
 
 import java.io.File
 import scala.concurrent.Future
@@ -16,19 +18,22 @@ object ImageIngestOperations {
 
   def optimisedPngKeyFromId(id: String)(implicit instance: Instance): String = instance.id + "/" + "optimised/" + snippetForId(id: String)
 
+  def embeddingKeyFromId(id: String)(implicit instance: Instance): String = instance.id + "/" + snippetForId(id)
+
   private def snippetForId(id: String) = id.take(6).mkString("/") + "/" + id
 }
 
-class ImageIngestOperations(imageBucket: S3Bucket, thumbnailBucket: S3Bucket, config: CommonConfig, isVersionedS3: Boolean = false)
+class ImageIngestOperations(imageBucket: S3Bucket, thumbnailBucket: S3Bucket, embeddingSourceBucket: S3Bucket, embeddingsBucket: S3Bucket, config: CommonConfig, isVersionedS3: Boolean = false)
   extends S3ImageStorage(config) with StrictLogging {
 
-  import ImageIngestOperations.{fileKeyFromId, optimisedPngKeyFromId}
+  import ImageIngestOperations.{embeddingKeyFromId, fileKeyFromId, optimisedPngKeyFromId}
 
   def store(storableImage: StorableImage)
            (implicit logMarker: LogMarker): Future[S3Object] = storableImage match {
     case s:StorableOriginalImage => storeOriginalImage(s)
     case s:StorableThumbImage => storeThumbnailImage(s)
     case s:StorableOptimisedImage => storeOptimisedImage(s)
+    case s:StorableEmbeddingSourceImage => storeEmbeddingSourceImage(s)
   }
 
   private def storeOriginalImage(storableImage: StorableOriginalImage)
@@ -55,6 +60,23 @@ class ImageIngestOperations(imageBucket: S3Bucket, thumbnailBucket: S3Bucket, co
       overwrite = true)
   }
 
+  private def storeEmbeddingSourceImage(storableImage: StorableEmbeddingSourceImage)
+                                       (implicit logMarker: LogMarker): Future[S3Object] = {
+    val instanceSpecificKey = fileKeyFromId(storableImage.id)(storableImage.instance)
+    logger.info(s"Storing embedding source to instance specific key: ${embeddingSourceBucket.name} / $instanceSpecificKey")
+    storeImage(embeddingSourceBucket, instanceSpecificKey, storableImage.file, Some(storableImage.mimeType),
+      overwrite = true)
+  }
+
+  def getEmbeddingStoreImage(key: String): ResponseInputStream[GetObjectResponse] = {
+    getObject(embeddingSourceBucket, key)
+  }
+
+  def storeEmbedding(key: String, embedding: Embedding): Unit = {
+    logger.info(s"Storing embedding source to key: ${embeddingsBucket.name} / $key")
+    putString(embeddingsBucket, key, Json.stringify(Json.toJson(embedding)))
+  }
+
   private def bulkDelete(bucket: S3Bucket, keys: List[String]): Future[Map[String, Boolean]] = keys match {
     case Nil => Future.successful(Map.empty)
     case _ => Future {
@@ -68,6 +90,7 @@ class ImageIngestOperations(imageBucket: S3Bucket, thumbnailBucket: S3Bucket, co
   def deleteThumbnails(ids: Set[String])(implicit instance: Instance) = bulkDelete(thumbnailBucket, ids.map(id => fileKeyFromId(id)).toList)
   def deletePNG(id: String)(implicit logMarker: LogMarker, instance: Instance): Future[Unit] = deleteImage(imageBucket, optimisedPngKeyFromId(id))
   def deletePNGs(ids: Set[String])(implicit instance: Instance) = bulkDelete(imageBucket, ids.map(id => optimisedPngKeyFromId(id)).toList)
+  def deleteEmbeddings(ids: Set[String])(implicit instance: Instance) = bulkDelete(embeddingsBucket, ids.map(id => embeddingKeyFromId(id)).toList)
 
   def doesOriginalExist(id: String)(implicit instance: Instance): Boolean =
     this.doesObjectExist(imageBucket, fileKeyFromId(id))
@@ -121,7 +144,16 @@ case class StorableOptimisedImage(id: String, file: File, mimeType: MimeType, me
     meta = meta
   )
 }
-
+case class StorableEmbeddingSourceImage(id: String, file: File, mimeType: MimeType, meta: Map[String, String] = Map.empty, instance: Instance) extends StorableImage {
+  override def toProjectedS3Object(embeddingSourcesBucket: S3Bucket): S3Object = S3Object(
+    embeddingSourcesBucket,
+    ImageIngestOperations.fileKeyFromId(id)(instance),
+    file,
+    Some(mimeType),
+    lastModified = None,
+    meta = meta
+  )
+}
 
 /**
   * @param id
@@ -135,5 +167,6 @@ case class StorableOptimisedImage(id: String, file: File, mimeType: MimeType, me
 case class BrowserViewableImage(id: String, file: File, mimeType: MimeType, meta: Map[String, String] = Map.empty, isTransformedFromSource: Boolean = false, instance: Instance) extends ImageWrapper {
   def asStorableOptimisedImage = StorableOptimisedImage(id, file, mimeType, meta, instance)
   def asStorableThumbImage = StorableThumbImage(id, file, mimeType, meta, instance)
+  def asStorableEmbeddingSourceImage = StorableEmbeddingSourceImage(id, file, mimeType, meta, instance)
 }
 

common-lib/src/main/scala/com/gu/mediaservice/lib/aws/Bedrock.scala

@@ -1,16 +1,17 @@
 package com.gu.mediaservice.lib.aws
 
-import software.amazon.awssdk.services.bedrockruntime.model._
-import software.amazon.awssdk.services.bedrockruntime._
 import com.gu.mediaservice.lib.config.CommonConfig
-import play.api.libs.json.Json
-import software.amazon.awssdk.core.SdkBytes
-
-import java.net.URI
+import com.gu.mediaservice.lib.embeddings.{EmbeddingImplementation, EmbeddingSourceImageFormat}
 import com.gu.mediaservice.lib.logging.LogMarker
+import com.gu.mediaservice.model.{CohereV4Embedding, Embedding, ImageMetadata, Jpeg, MimeType}
+import org.apache.commons.codec.binary.Base64
 import play.api.libs.json.OFormat.oFormatFromReadsAndOWrites
 import play.api.libs.json._
+import software.amazon.awssdk.core.SdkBytes
+import software.amazon.awssdk.services.bedrockruntime._
+import software.amazon.awssdk.services.bedrockruntime.model._
 
+import java.net.URI
 import scala.concurrent.{ExecutionContext, Future}
 
 object Bedrock {
@@ -22,10 +23,19 @@ object Bedrock {
   )
 
   private implicit val bedrockTextRequestFormat: OFormat[BedrockTextRequest] = Json.format[BedrockTextRequest]
+
+  case class BedrockImageRequest(
+    input_type: String,
+    embedding_types: List[String],
+    images: List[String],
+    output_dimension: Int
+  )
+
+  private implicit val bedrockImageRequestFormat: OFormat[BedrockImageRequest] = Json.format[BedrockImageRequest]
 }
 
 class Bedrock(config: CommonConfig)
-  extends AwsClientBuilderUtils {
+  extends EmbeddingImplementation with AwsClientBuilderUtils {
 
   // TODO: figure out what the more usual pattern for turning off localstack behaviour is
   override def awsLocalEndpointUri: Option[URI] = None
@@ -37,7 +47,7 @@ class Bedrock(config: CommonConfig)
       .build()
   }
 
-  private def createRequestBody(inputData: String): InvokeModelRequest = {
+  private def createSearchQueryRequestBody(inputData: String): InvokeModelRequest = {
     val body = Bedrock.BedrockTextRequest(
       input_type = "search_query",
       embedding_types = List("float"),
@@ -58,6 +68,30 @@ class Bedrock(config: CommonConfig)
     request
   }
 
+  private def createImageSearchDocumentRequestBody(base64Image: String, imageMimeType: MimeType): InvokeModelRequest = {
+    val body = Bedrock.BedrockImageRequest(
+      input_type = "search_document",
+      embedding_types = List("float"),
+      images = List(
+        s"`data:${imageMimeType.name};base64,$base64Image`"
+      ),
+      output_dimension = 1536
+    )
+
+    val jsonBody = Json.toJson(body).toString()
+
+    val request: InvokeModelRequest = {
+      InvokeModelRequest
+        .builder()
+        .accept("*/*")
+        .body(SdkBytes.fromUtf8String(jsonBody))
+        .contentType("application/json")
+        .modelId("global.cohere.embed-v4:0")
+        .build()
+    }
+    request
+  }
+
   private def sendBedrockEmbeddingRequest(requestBody: InvokeModelRequest)(
     implicit logMarker: LogMarker
   ): InvokeModelResponse = {
@@ -77,7 +111,7 @@ class Bedrock(config: CommonConfig)
   }
 
   def createTextEmbedding(inputData: String)(implicit ec: ExecutionContext, logMarker: LogMarker): Future[List[Double]] = {
-    val requestBody = createRequestBody(inputData)
+    val requestBody = createSearchQueryRequestBody(inputData)
     val bedrockFuture = Future { sendBedrockEmbeddingRequest(requestBody) }
     bedrockFuture.map { response =>
       val responseBody = response.body().asUtf8String()
@@ -91,4 +125,32 @@ class Bedrock(config: CommonConfig)
       embedding
     }
   }
+
+  override def createImageEmbeddings(source: Array[Byte], mimeType: MimeType, maybeMetadata: Option[ImageMetadata])(implicit ec: ExecutionContext, logMarker: LogMarker): Future[Embedding] = {
+    val base64ImageData = Base64.encodeBase64String(source)
+    val requestBody = createImageSearchDocumentRequestBody(
+      base64ImageData, embeddingSourceImageFormat().format
+    )
+    val bedrockFuture = Future {
+      sendBedrockEmbeddingRequest(requestBody)
+    }
+    bedrockFuture.map { response =>
+      val responseBody = response.body().asUtf8String()
+      val json = Json.parse(responseBody)
+      // Extract the embedding array (first element since it's an array of arrays)
+      val embeddings = (json \ "embeddings" \ "float")(0).as[List[Double]]
+      logger.info(
+        logMarker,
+        s"Successfully created image embedding. Vector size: ${embeddings.size}"
+      )
+      embeddings
+    }.map { embeddings =>
+      Embedding(
+        cohereEmbedV4 = Some(CohereV4Embedding(embeddings))
+      )
+    }
+  }
+
+  override def embeddingSourceImageFormat(): EmbeddingSourceImageFormat = EmbeddingSourceImageFormat(longestAxis = 3000, format = Jpeg, letterBox = false)
+
 }

common-lib/src/main/scala/com/gu/mediaservice/lib/aws/ContentDisposition.scala

@@ -19,14 +19,16 @@ trait ContentDisposition extends GridLogging {
     getContentDisposition(filename, fallbackLatin1Filename(image, extension))
   }
 
-  def getContentDisposition(image: Image, crop: Crop, asset: Asset, shortenDownloadFilename: Boolean): String = {
-    val cropId: String = crop.id.map(id => s"($id)").getOrElse("")
+  def getContentDisposition(image: Image, crop: Crop, asset: Asset): String = {
     val extension: String = getExtension(image, asset)
-    val dimensions: String = asset.dimensions.map(dims => s"(${dims.width} x ${dims.height})").getOrElse("")
-    val filenameSuffix: String = s"(${image.id})$cropId$dimensions$extension"
-    val filename = getBaseFilename(image, filenameSuffix, shortenDownloadFilename)
+    val filename = image.uploadInfo.filename match {
+      case Some(filename) => filename
+      case _ => image.id
+    }
 
-    getContentDisposition(filename, fallbackLatin1Filename(image, extension))
+    // Drop original file's extension and replace with the crops actual extension
+    val withCropsExtension = removeExtension(filename) + extension
+    getContentDisposition(withCropsExtension, fallbackLatin1Filename(image, extension))
   }
 
   private def getExtension(image: Image, asset: Asset): String = asset.mimeType match {