시맨틱 검색

// Define the document schema
val documentSchema = AppSearchSchema.Builder("my_schema_type")
    .addProperty(
        AppSearchSchema.StringPropertyConfig.Builder("text")
            .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
            .setIndexingType(AppSearchSchema.StringPropertyConfig.INDEXING_TYPE_PREFIXES)
            .setTokenizerType(AppSearchSchema.StringPropertyConfig.TOKENIZER_TYPE_PLAIN)
            .build()
    )
    .addProperty(
        AppSearchSchema.StringPropertyConfig.Builder("imageUri")
            .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
            .setIndexingType(AppSearchSchema.StringPropertyConfig.INDEXING_TYPE_NONE)
            .setTokenizerType(AppSearchSchema.StringPropertyConfig.TOKENIZER_TYPE_NONE)
            .build()
    )
    .addProperty(
        AppSearchSchema.EmbeddingPropertyConfig.Builder("embedding")
            .setCardinality(AppSearchSchema.PropertyConfig.CARDINALITY_OPTIONAL)
            .setIndexingType(
                AppSearchSchema.EmbeddingPropertyConfig.INDEXING_TYPE_SIMILARITY)
            .setQuantizationType(
                AppSearchSchema.EmbeddingPropertyConfig.QUANTIZATION_TYPE_8_BIT)
            .build()
    )
    .build()

suspend fun syncDatabase(documents: List<Pair<String, Uri?>>) {

    val modelSignature = model?.getModelSignature() ?: return
    val docsWithEmbedding = mutableListOf<GenericDocument>()

    for ((index, item) in documents.withIndex()) {
        val (content, imageUri) = item
        // 1. Generate embedding (text and/or image)
        val parts = mutableListOf<Part>()
        if (content.isNotBlank()) {
            parts.add(TextPart(content))
        }
        var imagePart: ImagePart? = null
        val vector = try {
            if (imageUri != null && isImageSupported) {
                imagePart = ImagePart(imageUri)
                parts.add(imagePart)
            }
            if (parts.isEmpty()) continue
            val request = EmbeddingRequest(
                Content(parts)
            )
            model?.generateEmbedding(request)?.embedding?.vector
                ?: continue
        } finally {
            imagePart?.close()
        }

        // 2. Build the AppSearch document
        val docBuilder = GenericDocument.Builder<GenericDocument.Builder<*>>(
            "my_namespace",
            "doc_$index",
            "my_schema_type"
        )
        .setPropertyString("text", content)
        // Associate the vector with the signature that generated it
        .setPropertyEmbedding("embedding", EmbeddingVector(vector, modelSignature))

        if (imageUri != null) {
            docBuilder.setPropertyString("imageUri", imageUri.toString())
        }

        docsWithEmbedding.add(docBuilder.build())
    }

    // 3. Clear the old database by force-overriding the schema
    val clearRequest = SetSchemaRequest.Builder().setForceOverride(true).build()
    session?.setSchemaAsync(clearRequest)?.await()

    val schemaRequest = SetSchemaRequest.Builder().addSchemas(documentSchema).build()
    session?.setSchemaAsync(schemaRequest)?.await()

    // 4. Index documents
    val putRequest = PutDocumentsRequest.Builder().addGenericDocuments(docsWithEmbedding).build()
    session?.putAsync(putRequest)?.await()
}

suspend fun search(
    queryText: String,
    imageUri: Uri? = null,
    threshold: Double = 0.5
): List<Triple<String, Uri?, Double>> {
    val modelSignature = model?.getModelSignature() ?: return emptyList()

    // 1. Generate query embedding (text and/or image)
    val parts = mutableListOf<Part>()
    if (queryText.isNotBlank()) {
        parts.add(TextPart(queryText))
    }
    var imagePart: ImagePart? = null
    val queryVector = try {
        if (imageUri != null && isImageSupported) {
            imagePart = ImagePart(imageUri)
            parts.add(imagePart)
        }
        if (parts.isEmpty()) return emptyList()
        val request = EmbeddingRequest(
            content = Content(parts)
        )
        model?.generateEmbedding(request)?.embedding?.vector
            ?: return emptyList()
    } finally {
        imagePart?.close()
    }

    // 2. Configure search specifications
    val searchSpec = SearchSpec.Builder()
        .addFilterSchemas("my_schema_type")
        .setDefaultEmbeddingSearchMetricType(SearchSpec.EMBEDDING_SEARCH_METRIC_TYPE_COSINE)
        .setRankingStrategy("sum(this.matchedSemanticScores(getEmbeddingParameter(0)))")
        .setListFilterQueryLanguageEnabled(true)
        .setResultCountPerPage(5)
        // Add the query vector as parameter 0
        .addEmbeddingParameters(EmbeddingVector(queryVector, modelSignature))
        .build()

    // 3. Execute search
    val queryExpression = "semanticSearch(getEmbeddingParameter(0), $threshold)"
    val results = session?.search(queryExpression, searchSpec) ?: return emptyList()

    return try {
        val page = results.nextPageAsync.await()
        page.map {
            val text = it.genericDocument.getPropertyString("text") ?: ""
            val matchedImageUri = it.genericDocument.getPropertyString("imageUri")?.let(Uri::parse)
            // rankingSignal contains the computed similarity score
            Triple(text, matchedImageUri, it.rankingSignal)
        }
    } finally {
        results.close()
    }
}