feat: add result size

Angular2Guy · May 12, 2024 · 6f877c2 · 6f877c2
1 parent ba79df2
commit 6f877c2
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/backend/src/main/java/ch/xxx/aidoclibchat/usecase/service/ImageService.java b/backend/src/main/java/ch/xxx/aidoclibchat/usecase/service/ImageService.java
@@ -30,6 +30,7 @@
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.chat.prompt.Prompt;
 import org.springframework.ai.document.Document;
+import org.springframework.beans.factory.annotation.Value;
 import org.springframework.stereotype.Service;
 import org.springframework.util.MimeType;
 
@@ -50,6 +51,8 @@ public class ImageService {
 	private ChatClient chatClient;
 	private ImageRepository imageRepository;
 	private DocumentVsRepository documentVsRepository;
+	@Value("${image.result-size:20}")
+	private Long resultSize;
 
 	private record ResultData(String answer, ImageQueryDto imageQueryDto) {
 	}
@@ -82,11 +85,11 @@ public List<ImageDto> queryImage(ImageQueryDto imageDto) {
 				.toList();
 		var imageMap = this.imageRepository
 				.findAllById(aiDocuments.stream().map(myDoc -> (String) myDoc.getMetadata().get(MetaData.ID))
-						.map(myUuid -> UUID.fromString(myUuid)).limit(20).toList())
+						.map(myUuid -> UUID.fromString(myUuid)).toList())
 				.stream().collect(Collectors.toMap(myDoc -> myDoc.getId(), myDoc -> myDoc));
 		record Container(Document document, Image image, Float distance) {
-		}		
-		var containers = imageMap.entrySet().stream().map(myEntry -> new Container(
+		}
+		return imageMap.entrySet().stream().map(myEntry -> new Container(
 				aiDocuments.stream()
 						.filter(myDoc -> myEntry.getKey().toString()
 								.equals((String) myDoc.getMetadata().get(MetaData.ID)))
@@ -96,13 +99,11 @@ record Container(Document document, Image image, Float distance) {
 						.filter(myDoc -> myEntry.getKey().toString()
 								.equals((String) myDoc.getMetadata().get(MetaData.ID)))
 						.map(myDoc -> (Float) myDoc.getMetadata().get(MetaData.DISTANCE)).findFirst().orElseThrow()))
-				.toList();
-		return containers.stream()
 				.sorted((containerA, containerB) -> containerA.distance().compareTo(containerB.distance()))
 				.map(myContainer -> new ImageDto(myContainer.document().getContent(),
 						Base64.getEncoder().encodeToString(myContainer.image().getImageContent()),
-						myContainer.image().getImageType()))				
-				.toList();
+						myContainer.image().getImageType()))
+				.limit(this.resultSize).toList();
 	}
 
 	private ResultData createAIResult(ImageQueryDto imageDto) {

diff --git a/backend/src/main/resources/application.properties b/backend/src/main/resources/application.properties
@@ -38,4 +38,5 @@ management.endpoint.health.status.http-mapping.out_of_service=503
 management.endpoint.health.show-details=always
 
 embedding-token-limit=2000
-document-token-limit=2000
+document-token-limit=2000
+image.result-size=20