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,
 
@@ -125,10 +125,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
 )
@@ -186,10 +186,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"
     )
   )
 

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/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.bucket} / $instanceSpecificKey")
+    storeImage(embeddingSourceBucket, instanceSpecificKey, storableImage.file, Some(storableImage.mimeType),
+      overwrite = true)
+  }
+
+  def getEmbeddingStoreImage(key: String): ResponseInputStream[GetObjectResponse] = {
+    getObjectV2(embeddingSourceBucket, key)
+  }
+
+  def storeEmbedding(key: String, embedding: Embedding): Unit = {
+    logger.info(s"Storing embedding to key: ${embeddingsBucket.bucket} / $key")
+    putString(embeddingsBucket, key, Json.stringify(Json.toJson(embedding)))
+  }
+
   private def bulkDeleteV2(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) = bulkDeleteV2(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) = bulkDeleteV2(imageBucket, ids.map(id => optimisedPngKeyFromId(id)).toList)
+  def deleteEmbeddings(ids: Set[String])(implicit instance: Instance) = bulkDeleteV2(embeddingsBucket, ids.map(id => embeddingKeyFromId(id)).toList)
 
   def doesOriginalExistV2(id: String)(implicit instance: Instance): Boolean = {
     this.doesObjectExistV2(imageBucket, fileKeyFromId(id))
@@ -122,7 +145,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
@@ -136,5 +168,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/auth/KeyStore.scala

@@ -18,7 +18,7 @@ class KeyStore(bucket: S3Bucket, s3: S3)(implicit ec: ExecutionContext)
 
   private def fetchAll: Map[String, ApiAccessor] = {
     val objects = Await.result(s3.listV2(bucket, ""), 10.seconds)
-    val keys = objects.map(_.uri.getPath.stripPrefix("/"))
+    val keys = objects.map(_.uri.getPath.stripPrefix("/"))  // TODO use unction!
     keys.flatMap(k => getS3Object(k).map(k -> ApiAccessor(_))).toMap
   }
 }

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 AwsClientV2BuilderUtils {
+  extends EmbeddingImplementation with AwsClientV2BuilderUtils {
 
   // 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/Embedder.scala

@@ -1,33 +1,38 @@
 package com.gu.mediaservice.lib.aws
+import com.gu.mediaservice.lib.embeddings.{EmbeddingImplementation, EmbeddingSourceImageFormat}
 import com.gu.mediaservice.lib.logging.{GridLogging, LogMarker}
-import com.gu.mediaservice.model.{Jpeg, MimeType, Png, Tiff}
+import com.gu.mediaservice.model.{Embedding, ImageMetadata, MimeType}
 import play.api.libs.json.{Json, OFormat}
-import software.amazon.awssdk.services.s3vectors.model.QueryVectorsResponse
-import software.amazon.awssdk.services.s3vectors.model.{QueryOutputVector, QueryVectorsResponse, VectorData}
 import software.amazon.awssdk.services.sqs.model.SendMessageResponse
 
-import java.nio.file.{Files, Path}
 import scala.concurrent.{ExecutionContext, Future}
-import scala.jdk.CollectionConverters.CollectionHasAsScala
 
-case class EmbedderMessage(imageId: String, fileType: String, s3Bucket: String, s3Key: String, instance: String)
+case class EmbedderMessage(imageId: String, s3Bucket: String, s3Key: String, instance: String, imageMetadata: Option[ImageMetadata])
 
 object EmbedderMessage {
   implicit val format: OFormat[EmbedderMessage] = Json.format[EmbedderMessage]
 }
 
-class Embedder(bedrock: Bedrock, sqs: SimpleSqsMessageConsumer)(implicit ec: ExecutionContext) extends GridLogging {
+class Embedder(embedding: EmbeddingImplementation, sqs: SimpleSqsMessageConsumer)(implicit ec: ExecutionContext) extends GridLogging {
 
   def createQueryEmbedding(query: String)(implicit logMarker: LogMarker): Future[List[Double]] = {
     logger.info(logMarker, s"Creating text embedding for query: $query")
     for {
-      embedding <- bedrock.createTextEmbedding(query)
+      embedding <- embedding.createTextEmbedding(query)
     } yield embedding
   }
 
-  def queueImageToEmbed(message: EmbedderMessage)(implicit logMarker: LogMarker) = {
+  def createImageEmbedding(source: Array[Byte], mimeType: MimeType, maybeMetadata: Option[ImageMetadata])(implicit logMarker: LogMarker): Future[Embedding] = {
+    logger.info(logMarker, s"Creating image embedding")
+    embedding.createImageEmbeddings(source, mimeType, maybeMetadata)
+  }
+
+  def queueImageToEmbed(message: EmbedderMessage)(implicit logMarker: LogMarker): Unit = {
     val messageBody = Json.stringify(Json.toJson(message))
     val result: SendMessageResponse = sqs.sendMessage(messageBody)
     logger.info(logMarker, s"Queued image for embedding with message ID: ${result.messageId()}")
   }
+
+  def embeddingSourceImageFormat(): EmbeddingSourceImageFormat = embedding.embeddingSourceImageFormat()
+
 }