Construindo um aplicativo Quartokus para realizar a Vector Search do MongoDB

1
@POST  
2
@Path("/generate-embeddings")  
3
public Response generateEmbeddings() {  
4
   listingService.generateAndStoreEmbeddings();  
5
   return Response.ok("Embeddings generated and stored").build();  
6
}

1
public void generateAndStoreEmbeddings() {  
2
   MongoCollection<Document> listingsCollection = mongoClient.getDatabase("sample_airbnb").getCollection("listingsAndReviews");
3
4
   int processedDocuments = 0;  
5
   long totalDocuments = listingsCollection.countDocuments();
6
7
   while (processedDocuments < totalDocuments) {  
8
       List<Document> documents = listingsCollection.find()  
9
               .skip(processedDocuments)  
10
               .limit(BATCH_SIZE)  
11
               .into(new ArrayList<>());  
12
       List<UpdateOneModel<Document>> bulkUpdates = new ArrayList<>();
13
14
       for (Document document : documents) {  
15
           String description = document.getString("description");  
16
           List<Double> embeddings = geminiClient.getEmbedding(description);
17
18
           if (embeddings != null) {  
19
               UpdateOneModel<Document> updateModel = new UpdateOneModel<>(  
20
                       Filters.eq("_id", document.getString("_id")),  
21
                       Updates.set("embeddings", embeddings)  
22
               );  
23
               bulkUpdates.add(updateModel);  
24
           }  
25
       }  
26
       if (!bulkUpdates.isEmpty()) {  
27
           listingsCollection.bulkWrite(bulkUpdates, new BulkWriteOptions().ordered(false));  
28
       }
29
30
       processedDocuments += documents.size();  
31
       System.out.println("Processed " + processedDocuments + " out of " + totalDocuments + " documents.");  
32
   }
33
34
   System.out.println("Embedding generation complete for all documents.");  
35
}

1
package com.example.vectorsearch.gateway;
2
3
import com.example.vectorsearch.config.GeminiConfig;  
4
import com.example.vectorsearch.request.GeminiRequest;  
5
import com.fasterxml.jackson.core.JsonProcessingException;  
6
import com.fasterxml.jackson.databind.ObjectMapper;  
7
import io.quarkus.rest.client.reactive.QuarkusRestClientBuilder;  
8
import jakarta.enterprise.context.ApplicationScoped;  
9
import jakarta.inject.Inject;  
10
import jakarta.ws.rs.WebApplicationException;  
11
import org.slf4j.Logger;  
12
import org.slf4j.LoggerFactory;
13
14
import java.net.URI;  
15
import java.util.List;
16
17
@ApplicationScoped  
18
public class GeminiAIGatewayImpl implements GeminiAIGateway {
19
20
   @Inject  
21
   GeminiConfig geminiAIConfig;
22
23
   private static final Logger LOGGER = LoggerFactory.getLogger(GeminiAIGatewayImpl.class);
24
25
   @Override  
26
   public List<Double> getEmbedding(String input) {  
27
       ObjectMapper objectMapper = new ObjectMapper();
28
29
       var geminiAIGateway = QuarkusRestClientBuilder.newBuilder()  
30
               .baseUri(URI.create*(geminiAIConfig.getGeminiAIUrl() + "?key=" + geminiAIConfig.getApiKey()))  
31
               .build(GeminiClient.class);
32
33
       GeminiRequest request = new GeminiRequest(  
34
               geminiAIConfig.getModel(),  
35
               new GeminiRequest.Content(List.of(new GeminiRequest.Part(input)))  
36
       );  
37
       try {  
38
           String jsonRequest = objectMapper.writeValueAsString(request);  
39
           var embedding = geminiAIGateway.embedding(jsonRequest);
40
41
           return embedding.getEmbedding().getValues();  
42
       } catch (WebApplicationException e) {  
43
           LOGGER.error("Error from Gemini AI: {}", e.getResponse().readEntity(String.class));  
44
           throw e;  
45
       } catch (JsonProcessingException e) {  
46
           throw new RuntimeException(e);  
47
       }  
48
   }  
49
}

1
 curl -X POST http://localhost:8080/api/listings/generate-embeddings

1
{  
2
 "fields": [  
3
   {  
4
     "numDimensions": 768,  
5
     "path": "embeddings",  
6
     "similarity": "euclidean",  
7
     "type": "vector"  
8
   }  
9
 ]  
10
}

1
public void performVectorSearch(String query) {  
2
   MongoCollection<Document> collection = mongoClient.getDatabase("sample_airbnb").getCollection("listingsAndReviews");
3
4
   List<Double> queryEmbeddings = geminiClient.getEmbedding(query);  
5
   String indexName = "vector_index";  
6
   int numCandidates = 150;  
7
   int limit = 10;
8
9
   List<Document> pipeline = Arrays.asList(  
10
           new Document("$vectorSearch",  
11
                   new Document("index", indexName)  
12
                           .append("path", "embeddings")  
13
                           .append("queryVector", queryEmbeddings)  
14
                           .append("numCandidates", numCandidates)  
15
                           .append("limit", limit)  
16
           ),  
17
           new Document("$project",  
18
                   new Document("_id", 0)  
19
                           .append("name", 1)  
20
                           .append("listing_url", 1)  
21
                           .append("description", 1)  
22
                           .append("price", 1)  
23
           ),  
24
           new Document("$limit", 5));
25
26
   collection.aggregate(pipeline)  
27
           .forEach(doc -> System.out.println(doc.toJson()));  
28
}

1
curl -X GET "http://localhost:8080/api/listings/perform-vector-search?query=Hotels%20that%20are%20recommended%20for%20romantic%20stay" | jq

1
[
2
  {
3
    "listing_url": "https://www.airbnb.com/rooms/15266254",
4
    "name": "The Manhattan Club in the heart of midtown!!!!",
5
    "description": "My place is good for couples, solo adventurers, and business travelers.",
6
    "price": 305.00
7
  },
8
  {
9
    "listing_url": "https://www.airbnb.com/rooms/14261122",
10
    "name": "Cute Condo in Mile-End, sleeps 2",
11
    "description": "You’ll love my place because of the location. My place is good for couples, solo adventurers, and business travellers.",
12
    "price": 70.00
13
  },
14
  {
15
    "listing_url": "https://www.airbnb.com/rooms/13609188",
16
    "name": "Habitacion al lado de Sagrada Familia",
17
    "description": "Lugares de interés: La Paradeta Sagarada Familia, FCB Official Point, Basílica de la Sagrada Família, Starbucks, . Mi alojamiento es bueno para parejas y aventureros.",
18
    "price": 45.00
19
  },
20
  {
21
    "listing_url": "https://www.airbnb.com/rooms/13748059",
22
    "name": "Studio for the 2016 Olympics",
23
    "description": "My space is good for couples, business travelers and individual adventures.",
24
    "price": 298.00
25
  },
26
  {
27
    "listing_url": "https://www.airbnb.com/rooms/11130156",
28
    "name": "Double Room",
29
    "description": "Literally just across the street is the world famous Cagaloglu Hammam featured in 1,000 Places to See Before You Die by Patricia Schultz. The hotel is tucked away on a quiet street and is just a few minutes walk to some of the world’s famous sites,",
30
    "price": 359.00
31
  }
32
]

1
 curl -X GET "http://localhost:8080/api/listings/perform-vector-search?query=Stays%20that%20the%20kids%20would%20enjoy" | jq

1
[
2
  {
3
    "listing_url": "https://www.airbnb.com/rooms/13997910",
4
    "name": "Apartamento de luxo em Copacabana - 4 quartos",
5
    "description": "Meu espaço é bom para casais, viajantes de negócios e famílias (com crianças).",
6
    "price": 11190.00
7
  },
8
  {
9
    "listing_url": "https://www.airbnb.com/rooms/14571224",
10
    "name": "Alugo Quarto Para Temporada de Olimpiadas",
11
    "description": "Meu espaço é bom para casais, viajantes de negócios e famílias (com crianças).",
12
    "price": 798.00
13
  },
14
  {
15
    "listing_url": "https://www.airbnb.com/rooms/13927230",
16
    "name": "Casa completa p olimpíadas com serviços incluído",
17
    "description": "Você vai amar meu espaço por causa de a área externa, da piscina, o bairro com vários bares e diversões, o condomínio seguro com porteiro, a iluminação, a cama confortável e a cozinha grande. Meu espaço é bom para famílias (com crianças) e grandes grupos. Podendo se contratar junto também serviços de cozinheira, garçons e bartender. Podendo se contratar também com alimentação. Tudo para se sentir em casa e não se preocupar com nada. Consultar valores extras.",
18
    "price": 601.00
19
  },
20
  {
21
    "listing_url": "https://www.airbnb.com/rooms/13617872",
22
    "name": "Amor do Rio",
23
    "description": "Você vai amar meu espaço por causa de o aconchego e a localização. Meu espaço é bom para casais, aventuras individuais, viajantes de negócios, famílias (com crianças) e grandes grupos.",
24
    "price": 116.00
25
  },
26
  {
27
    "listing_url": "https://www.airbnb.com/rooms/13993059",
28
    "name": "Ventanas Nature Resort  ambiente familiar",
29
    "description": "Meu espaço é perto de transporte público, restaurantes e refeições, atividades para famílias, , praia. Você vai amar meu espaço por causa de a localização, as pessoas, o ambiente e o bairro. Meu espaço é bom para casais, famílias (com crianças), grandes grupos e amigos peludos (animais de estimação).",
30
    "price": 899.00
31
  }
32
]