Construindo um serviço de pesquisa semântica do Atlas com Spring AI e MongoDB Atlas

1
spring.application.name=lyric-semantic-search
2
spring.ai.openai.api-key=<Your-API-key>
3
spring.ai.openai.embedding.options.model=text-embedding-ada-002
4
5
spring.data.mongodb.uri=<Your-MongoDB-connection-string>
6
spring.data.mongodb.database=lyrics
7
spring.ai.vectorstore.mongodb.indexName=vector_index
8
spring.ai.vectorstore.mongodb.collection-name=vector_store
9
spring.ai.vectorstore.mongodb.initialize-schema=true

1
package com.mongodb.lyric_semantic_search.config;
2
3
import org.springframework.ai.embedding.EmbeddingModel;
4
import org.springframework.ai.openai.OpenAiEmbeddingModel;
5
import org.springframework.ai.openai.api.OpenAiApi;
6
import org.springframework.beans.factory.annotation.Value;
7
import org.springframework.boot.SpringBootConfiguration;
8
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
9
import org.springframework.context.annotation.Bean;
10
import org.springframework.context.annotation.Configuration;
11
12
@SpringBootConfiguration
13
@EnableAutoConfiguration
14
15
@Configuration
16
public class Config {
17
18
    @Value("${spring.ai.openai.api-key}")
19
    private String openAiKey;
20
21
    @Bean
22
    public EmbeddingModel embeddingModel() {
23
        return new OpenAiEmbeddingModel(new OpenAiApi(openAiKey));
24
    }
25
}

1
package com.mongodb.lyric_semantic_search.model;
2
3
import java.util.Map;
4
5
public class DocumentRequest {
6
    private String content;
7
    private Map<String, Object> metadata;
8
9
    public DocumentRequest() {
10
    }
11
12
    public DocumentRequest(String content, Map<String, Object> metadata) {
13
        this.content = content;
14
        this.metadata = metadata;
15
    }
16
17
    public String getContent() {
18
        return content;
19
    }
20
21
    public void setContent(String content) {
22
        this.content = content;
23
    }
24
25
    public Map<String, Object> getMetadata() {
26
        return metadata;
27
    }
28
29
    public void setMetadata(Map<String, Object> metadata) {
30
        this.metadata = metadata;
31
    }
32
33
}

1
package com.mongodb.lyric_semantic_search.repository;
2
3
import java.util.List;
4
import java.util.Optional;
5
6
import org.springframework.ai.document.Document;
7
import org.springframework.ai.vectorstore.SearchRequest;
8
9
public interface LyricSearchRepository {
10
11
    void addDocuments(List<Document> docs);
12
13
    Optional<Boolean> deleteDocuments(List<String> ids);
14
15
    List<Document> semanticSearchByLyrics(SearchRequest searchRequest);
16
}

1
package com.mongodb.lyric_semantic_search.repository;
2
3
import java.util.List;
4
import java.util.Optional;
5
6
import org.springframework.ai.document.Document;
7
import org.springframework.ai.vectorstore.SearchRequest;
8
import org.springframework.ai.vectorstore.VectorStore;
9
import org.springframework.beans.factory.annotation.Autowired;
10
import org.springframework.stereotype.Repository;
11
12
@Repository
13
public class LyricSearchRepositoryImpl implements LyricSearchRepository {
14
15
    private final VectorStore vectorStore;
16
17
    @Autowired
18
    public LyricSearchRepositoryImpl(VectorStore vectorStore) {
19
        this.vectorStore = vectorStore;
20
    }
21
22
    @Override
23
    public void addDocuments(List<Document> docs) {
24
        vectorStore.add(docs);
25
    }
26
27
    @Override
28
    public Optional<Boolean> deleteDocuments(List<String> ids) {
29
        return vectorStore.delete(ids);
30
    }
31
32
    @Override
33
    public List<Document> semanticSearchByLyrics(SearchRequest searchRequest) {
34
        return vectorStore.similaritySearch(searchRequest);
35
    }
36
}

1
package com.mongodb.lyric_semantic_search.service;
2
3
import org.springframework.stereotype.Service;
4
5
@Service
6
public class LyricSearchService {
7
8
}

1
package com.mongodb.lyric_semantic_search.controller;
2
3
import org.springframework.web.bind.annotation.RestController;
4
5
@RestController
6
public class LyricSearchController {
7
8
}

1
    private static final int MAX_TOKENS = (int) (8192 * 0.80); // OpenAI model's maximum content length + BUFFER for when one word > 1 token
2
3
    @Autowired
4
    LyricSearchRepository lyricSearchRepository;
5
6
    public List<Document> addDocuments(List<DocumentRequest> documents) {
7
        if (documents == null || documents.isEmpty()) {
8
            return Collections.emptyList();
9
        }
10
11
        List<Document> docs = documents.stream()
12
            .filter(doc -> doc != null && doc.getContent() != null && !doc.getContent()
13
                .trim()
14
                .isEmpty())
15
            .map(doc -> new Document(doc.getContent(), doc.getMetadata()))
16
            .filter(doc -> {
17
                int wordCount = doc.getContent()
18
                    .split("\\s+").length;
19
                return wordCount <= MAX_TOKENS;
20
            })
21
            .collect(Collectors.toList());
22
23
        if (!docs.isEmpty()) {
24
            lyricSearchRepository.addDocuments(docs);
25
        }
26
27
        return docs;
28
    }

1
public List<String> deleteDocuments(List<String> ids) {
2
    if (ids == null || ids.isEmpty()) {
3
        return Collections.emptyList(); // Nothing to delete
4
    }
5
6
    Optional<Boolean> result = lyricSearchRepository.deleteDocuments(ids);
7
    if (result.isPresent() && result.get()) {
8
        return ids; // Return the list of successfully deleted IDs
9
    } else {
10
        return Collections.emptyList(); // Return empty list if deletion was unsuccessful
11
    }
12
}

1
import java.util.Collections;
2
import java.util.List;
3
import java.util.Optional;
4
import java.util.stream.Collectors;
5
6
import org.springframework.ai.document.Document;
7
import org.springframework.beans.factory.annotation.Autowired;
8
9
import com.mongodb.lyric_semantic_search.model.DocumentRequest;
10
import com.mongodb.lyric_semantic_search.repository.LyricSearchRepository;

1
    @Autowired
2
    private LyricSearchService lyricSearchService;
3
4
    @PostMapping("/addDocuments")
5
    public List<Map<String, Object>> addDocuments(@RequestBody List<DocumentRequest> documents) {
6
        return lyricSearchService.addDocuments(documents).stream()
7
            .map(doc -> Map.of("content", doc.getContent(), "metadata", doc.getMetadata()))
8
            .collect(Collectors.toList());
9
    }
10
11
    @DeleteMapping("/delete")
12
    public List<String> deleteDocuments(@RequestBody List<String> ids) {
13
        return lyricSearchService.deleteDocuments(ids);
14
    }

1
import java.util.List;
2
import java.util.Map;
3
import java.util.stream.Collectors;
4
5
import org.springframework.beans.factory.annotation.Autowired;
6
import org.springframework.web.bind.annotation.DeleteMapping;
7
import org.springframework.web.bind.annotation.PostMapping;
8
import org.springframework.web.bind.annotation.RequestBody;
9
10
import com.mongodb.lyric_semantic_search.model.DocumentRequest;
11
import com.mongodb.lyric_semantic_search.service.LyricSearchService;

1
curl -X POST "http://localhost:8080/addDocuments" \
2
     -H "Content-Type: application/json" \          
3
     -d '[
4
           {"content": "Twinkle, twinkle, little star, How I wonder what you are! Up above the world so high, Like a diamond in the sky.", "metadata": {"title": "Twinkle Twinkle Little Star", "artist": "Jane Taylor", "year": "1806"}},
5
           {"content": "The itsy bitsy spider climbed up the waterspout. Down came the rain and washed the spider out. Out came the sun and dried up all the rain and the itsy bitsy spider climbed up the spout again.", "metadata": {"title": "Itsy Bitsy Spider", "artist": "Traditional", "year": "1910"}},
6
           {"content": "Humpty Dumpty sat on a wall, Humpty Dumpty had a great fall. All the kings horses and all the kings men couldnt put Humpty together again.", "metadata": {"title": "Humpty Dumpty", "artist": "Mother Goose", "year": "1797"}}
7
         ]'

1
    public List<Map<String, Object>> searchDocuments(String query, int topK, double similarityThreshold) {
2
        SearchRequest searchRequest = SearchRequest.query(query)
3
            .withTopK(topK)
4
            .withSimilarityThreshold(similarityThreshold);
5
6
        List<Document> results = lyricSearchRepository.semanticSearchByLyrics(searchRequest);
7
8
        return results.stream()
9
            .map(doc -> Map.of("content", doc.getContent(), "metadata", doc.getMetadata()))
10
            .collect(Collectors.toList());
11
    }

1
import java.util.Map;
2
import org.springframework.ai.vectorstore.SearchRequest;

1
    @GetMapping("/search")
2
    public List<Map<String, Object>> searchDocuments(@RequestParam String query, @RequestParam int topK, @RequestParam double similarityThreshold
3
        ) {
4
        return lyricSearchService.searchDocuments(query, topK, similarityThreshold);
5
        
6
    }

1
import org.springframework.web.bind.annotation.GetMapping;
2
import org.springframework.web.bind.annotation.RequestParam;

1
curl -X GET "http://localhost:8080/search?query=small%20celestial%20bodie&topK=5&similarityThreshold=0.8"

1
[{
2
	"metadata":{
3
		"title":"Twinkle Twinkle Little Star",
4
		"artist":"Jane Taylor",
5
		"year":"1806"
6
	},
7
	"content":"Twinkle, twinkle, little star,..."
8
},
9
...

1
{
2
  "fields": [
3
    {
4
      "numDimensions": 1536,
5
      "path": "embedding",
6
      "similarity": "cosine",
7
      "type": "vector"
8
    },
9
    {
10
      "path": "metadata.artist",
11
      "type": "filter"
12
    }
13
  ]
14
}

1
    public List<Map<String, Object>> searchDocumentsWithFilter(String query, int topK, double similarityThreshold, String artist) {
2
        FilterExpressionBuilder filterBuilder = new FilterExpressionBuilder();
3
        Expression filterExpression = filterBuilder.eq("artist", artist)
4
            .build();
5
6
        SearchRequest searchRequest = SearchRequest.query(query)
7
            .withTopK(topK)
8
            .withSimilarityThreshold(similarityThreshold)
9
            .withFilterExpression(filterExpression);
10
11
        List<Document> results = lyricSearchRepository.semanticSearchByLyrics(searchRequest);
12
13
        return results.stream()
14
            .map(doc -> Map.of("content", doc.getContent(), "metadata", doc.getMetadata()))
15
            .collect(Collectors.toList());
16
    }

1
import org.springframework.ai.vectorstore.filter.Filter.Expression;
2
import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder;

1
    @GetMapping("/searchWithFilter")
2
    public List<Map<String, Object>> searchDocumentsWithFilter(@RequestParam String query, @RequestParam int topK, @RequestParam double similarityThreshold, @RequestParam String artist) {
3
        return lyricSearchService.searchDocumentsWithFilter(query, topK, similarityThreshold, artist);
4
    }

1
curl -X GET "http://localhost:8080/searchWithFilter?query=little%20star&topK=5&similarityThreshold=0.8&artist=Jane%20Taylor"

1
[{
2
	"metadata":{
3
		"title":"Twinkle Twinkle Little Star",
4
		"artist":"Jane Taylor",
5
		"year":"1806"
6
	},
7
	"content":"Twinkle, twinkle, little star,..."
8
}]