Desbloqueando a pesquisa semântica: crie um mecanismo de pesquisa de filmes baseado em Java com o Atlas Vector Search e o Spring Boot

1
{
2
  "fields": [{
3
    "type": "vector",
4
    "path": "plot_embedding",
5
    "numDimensions": 1536,
6
    "similarity": "euclidean"
7
  }]
8
}

1
<mongodb.version>4.11.0</mongodb.version>

1
public class Movie {
2
3
    @BsonProperty("_id")
4
    private ObjectId Id;
5
    private String title;
6
    private int year;
7
    private int runtime;
8
    private Date released;
9
    private String poster;
10
    private String plot;
11
    private String fullplot;
12
    private String lastupdated;
13
    private String type;
14
    private List<String> directors;
15
    private Imdb imdb;
16
    private List<String> cast;
17
    private List<String> countries;
18
    private List<String> genres;
19
    private Tomatoes tomatoes;
20
    private int num_mflix_comments;
21
    private String plot_embeddings;
22
23
    // Getters and setters for Movie fields
24
25
}

1
public static class Imdb {
2
3
    private double rating;
4
    private int votes;
5
    private int id;
6
7
    // Getters and setters for Imdb fields
8
9
}

1
public static class Tomatoes {
2
3
    private Viewer viewer;
4
    private Date lastUpdated;
5
6
    // Getters and setters for Tomatoes fields
7
8
}

1
public static class Viewer {
2
3
    private double rating;
4
    private int numReviews; 
5
6
    // Getters and setters for Viewer fields
7
8
}

1
mongodb.uri=mongodb+srv://<username>:<password>@<cluster>.mongodb.net/
2
mongodb.database=sample_mflix

1
@Configuration
2
public class MongodbConfig {
3
4
    @Value("${mongodb.uri}")
5
    private String MONGODB_URI;
6
7
    @Value("${mongodb.database}")
8
    private String MONGODB_DATABASE;

1
    @Bean
2
    public MongoClient mongoClient() {
3
        CodecRegistry pojoCodecRegistry = CodecRegistries.fromRegistries(
4
            MongoClientSettings.getDefaultCodecRegistry(),
5
            CodecRegistries.fromProviders(
6
                PojoCodecProvider.builder().automatic(true).build()
7
            )
8
        );
9
10
        MongoClientSettings settings = MongoClientSettings.builder()
11
            .applyConnectionString(new ConnectionString(MONGODB_URI))
12
            .codecRegistry(pojoCodecRegistry)
13
            .build();
14
15
        return MongoClients.create(settings);
16
    }

1
    @Bean
2
    public MongoDatabase mongoDatabase(MongoClient mongoClient) {
3
        return mongoClient.getDatabase(MONGODB_DATABASE); 
4
    }
5
}

1
@Service
2
public class OpenAIService {
3
	
4
	private static final String OPENAI_API_URL = "https://api.openai.com";
5
	
6
	@Value("${openai.api.key}")
7
	
8
	private String OPENAI_API_KEY;
9
	
10
	private WebClient webClient;
11
	
12
	@PostConstruct
13
	void init() {
14
		this.webClient = WebClient.builder()
15
			.clientConnector(new ReactorClientHttpConnector())
16
			.baseUrl(OPENAI_API_URL)
17
			.defaultHeader("Content-Type", MediaType.APPLICATION_JSON_VALUE)
18
			.defaultHeader("Authorization", "Bearer " + OPENAI_API_KEY)
19
			.build();
20
	}
21
22
	
23
	public Mono<List<Double>> createEmbedding(String text) {
24
	Map<String, Object> body = Map.of(
25
		"model", "text-embedding-ada-002",
26
		"input", text
27
	);
28
	
29
	return webClient.post()
30
		.uri("/v1/embeddings")
31
		.bodyValue(body)
32
		.retrieve()
33
		.bodyToMono(EmbeddingResponse.class)
34
		.map(EmbeddingResponse::getEmbedding);
35
	}
36
}

1
public class EmbeddingData {
2
	private List<Double> embedding;
3
	
4
	public List<Double> getEmbedding() {
5
		return embedding;
6
	}
7
	
8
	public void setEmbedding(List<Double> embedding) {
9
		this.embedding = embedding;
10
	}
11
}

1
public class EmbeddingResponse {
2
	private List<EmbeddingData> data;
3
	
4
	public List<Double> getEmbedding() {
5
		return data.get(0).getEmbedding();
6
	}
7
	
8
	public List<EmbeddingData> getData() {
9
		return data;
10
	}
11
	
12
	public void setData(List<EmbeddingData> data) {
13
		this.data = data;
14
	}
15
}

1
public interface MovieRepository {
2
    Flux<Movie> findMoviesByVector(List<Double> embedding);
3
}

1
@Repository
2
public class MovieRepositoryImpl implements MovieRepository {
3
4
    private final MongoDatabase mongoDatabase;
5
6
    public MovieRepositoryImpl(MongoDatabase mongoDatabase) {
7
        this.mongoDatabase = mongoDatabase;
8
    }
9
10
    private MongoCollection<Movie> getMovieCollection() {
11
        return mongoDatabase.getCollection("embedded_movies", Movie.class);
12
    }
13
14
    @Override
15
    public Flux<Movie> findMoviesByVector(List<Double> embedding) {
16
        String indexName = "PlotVectorSearch";
17
        int numCandidates = 100;
18
        int limit = 5;
19
20
        List<Bson> pipeline = asList(
21
            vectorSearch(
22
                fieldPath("plot_embedding"),
23
                embedding,
24
                indexName,
25
                numCandidates,
26
                limit));
27
28
        return Flux.from(getMovieCollection().aggregate(pipeline, Movie.class));
29
    }
30
}

1
@Service
2
public class MovieService {
3
4
    private final MovieRepository movieRepository;
5
    private final OpenAIService embedder;
6
7
    @Autowired
8
    public MovieService(MovieRepository movieRepository, OpenAIService embedder) {
9
        this.movieRepository = movieRepository;
10
        this.embedder = embedder;
11
    }
12
13
    public Mono<List<Movie>> getMoviesSemanticSearch(String plotDescription) {
14
        return embedder.createEmbedding(plotDescription)
15
                .flatMapMany(movieRepository::findMoviesByVector)
16
                .collectList();
17
    }
18
}

1
@RestController
2
public class MovieController {
3
4
	private final MovieService movieService;
5
	
6
	@Autowired
7
	public MovieController(MovieService movieService) {
8
		this.movieService = movieService;
9
	}
10
	
11
	@GetMapping("/movies/semantic-search")
12
	public Mono<List<Movie>> performSemanticSearch(@RequestParam("plotDescription") String plotDescription) {
13
		return movieService.getMoviesSemanticSearch(plotDescription);
14
	}
15
}

1
curl -X GET "http://localhost:8080/movies/semantic-search?plotDescription=A%20cop%20from%20china%20and%20cop%20from%20america%20save%20kidnapped%20girl"

1
Movie title: Rush Hour
2
Plot: Two cops team up to get back a kidnapped daughter.
3
4
Movie title: Police Story 3: Supercop
5
Plot: A Hong Kong detective teams up with his female Red Chinese counterpart to stop a Chinese drug czar.
6
  
7
Movie title: Fuk sing go jiu
8
Plot: Two Hong-Kong cops are sent to Tokyo to catch an ex-cop who stole a large amount of money in diamonds. After one is captured by the Ninja-gang protecting the rogue cop, the other one gets ...
9
  
10
Movie title: Motorway
11
Plot: A rookie cop takes on a veteran escape driver in a death defying final showdown on the streets of Hong Kong.
12
  
13
Movie title: The Corruptor
14
Plot: With the aid from a NYC policeman, a top immigrant cop tries to stop drug-trafficking and corruption by immigrant Chinese Triads, but things complicate when the Triads try to bribe the policeman.