fix to add metadata plus tests - Signed-off-by: Theo van Kraay <[email protected]>

Signed-off-by: Theo van Kraay <[email protected]>

Author: TheovanKraay

vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java

@@ -18,7 +18,9 @@
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
@@ -351,16 +353,32 @@ public List<Document> doSimilaritySearch(SearchRequest request) {
 
 		CosmosPagedFlux<JsonNode> pagedFlux = this.container.queryItems(sqlQuerySpec, options, JsonNode.class);
 
+
 		logger.info("Executing similarity search query: {}", query);
 		try {
 			// Collect documents from the paged flux
 			List<JsonNode> documents = pagedFlux.byPage()
 				.flatMap(page -> Flux.fromIterable(page.getResults()))
 				.collectList()
 				.block();
+			
+			// Collect metadata fields from the documents
+			Map<String, Object> docFields = new HashMap<>();
+			for (var doc : documents) {
+				JsonNode metadata = doc.get("metadata");
+				metadata.fieldNames().forEachRemaining(field -> {
+					JsonNode value = metadata.get(field);
+					Object parsedValue = value.isTextual() ? value.asText() :
+							value.isNumber()  ? value.numberValue() :
+									value.isBoolean() ? value.booleanValue() :
+											value.toString();
+					docFields.put(field, parsedValue);
+				});
+			}
+
 			// Convert JsonNode to Document
 			List<Document> docs = documents.stream()
-				.map(doc -> Document.builder().id(doc.get("id").asText()).text(doc.get("content").asText()).build())
+				.map(doc -> Document.builder().id(doc.get("id").asText()).text(doc.get("content").asText()).metadata(docFields).build())
 				.collect(Collectors.toList());
 
 			return docs != null ? docs : List.of();

vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStoreIT.java

@@ -141,6 +141,12 @@ void testSimilaritySearchWithFilter() {
 
 		assertThat(results).hasSize(2);
 		assertThat(results).extracting(Document::getId).containsExactlyInAnyOrder("1", "2");
+		for (Document doc : results) {
+			assertThat(doc.getMetadata().get("country")).isIn("UK", "NL");
+			assertThat(doc.getMetadata().get("year")).isIn(2021, 2022);
+			assertThat(doc.getMetadata().get("city")).isIn("London", "Amsterdam")
+				.isNotEqualTo("Sofia");
+		}
 
 		List<Document> results2 = this.vectorStore.similaritySearch(SearchRequest.builder()
 			.query("The World")
@@ -152,6 +158,8 @@ void testSimilaritySearchWithFilter() {
 		assertThat(results2).hasSize(1);
 		assertThat(results2).extracting(Document::getId).containsExactlyInAnyOrder("1");
 
+
+
 		List<Document> results3 = this.vectorStore.similaritySearch(SearchRequest.builder()
 			.query("The World")
 			.topK(10)