feat: add query

Angular2Guy · May 12, 2024 · ba79df2 · ba79df2
1 parent 7a3c580
commit ba79df2
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 8 deletions.
diff --git a/backend/src/main/java/ch/xxx/aidoclibchat/adapter/controller/ImageController.java b/backend/src/main/java/ch/xxx/aidoclibchat/adapter/controller/ImageController.java
@@ -12,6 +12,8 @@
  */
 package ch.xxx.aidoclibchat.adapter.controller;
 
+import java.util.List;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.web.bind.annotation.PostMapping;
@@ -37,7 +39,7 @@ public ImageController(ImageMapper imageMapper, ImageService imageService) {
 	}
 
 	@PostMapping("/query")
-	public ImageDto postImageQuery(@RequestParam("query") String query,@RequestParam("type") String type, @RequestParam("file") MultipartFile imageQuery) {		
+	public List<ImageDto> postImageQuery(@RequestParam("query") String query,@RequestParam("type") String type, @RequestParam("file") MultipartFile imageQuery) {		
 		var result = this.imageService.queryImage(this.imageMapper.map(imageQuery, query));		
 		return result;
 	}
@@ -47,4 +49,5 @@ public ImageDto postImportImage(@RequestParam("query") String query,@RequestPara
 		var result = this.imageService.importImage(this.imageMapper.map(imageQuery, query), this.imageMapper.map(imageQuery));		
 		return result;
 	}
+
 }
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
@@ -18,6 +18,8 @@
 import java.io.IOException;
 import java.util.Base64;
 import java.util.List;
+import java.util.UUID;
+import java.util.stream.Collectors;
 
 import javax.imageio.ImageIO;
 
@@ -32,6 +34,7 @@
 import org.springframework.util.MimeType;
 
 import ch.xxx.aidoclibchat.domain.common.MetaData;
+import ch.xxx.aidoclibchat.domain.common.MetaData.DataType;
 import ch.xxx.aidoclibchat.domain.common.MetaData.ImageType;
 import ch.xxx.aidoclibchat.domain.model.dto.ImageDto;
 import ch.xxx.aidoclibchat.domain.model.dto.ImageQueryDto;
@@ -51,14 +54,15 @@ public class ImageService {
 	private record ResultData(String answer, ImageQueryDto imageQueryDto) {
 	}
 
-	public ImageService(ChatClient chatClient, ImageRepository imageRepository, DocumentVsRepository documentVsRepository) {
+	public ImageService(ChatClient chatClient, ImageRepository imageRepository,
+			DocumentVsRepository documentVsRepository) {
 		this.chatClient = chatClient;
 		this.imageRepository = imageRepository;
 		this.documentVsRepository = documentVsRepository;
 	}
 
 	public ImageDto importImage(ImageQueryDto imageDto, Image image) {
-		var resultData = createAIResult(imageDto);
+		var resultData = this.createAIResult(imageDto);
 		image.setImageContent(resultData.imageQueryDto().getImageContent());
 		var myImage = this.imageRepository.save(image);
 		var aiDocument = new Document(resultData.answer());
@@ -70,11 +74,35 @@ public ImageDto importImage(ImageQueryDto imageDto, Image image) {
 				resultData.imageQueryDto().getImageType());
 	}
 
-	public ImageDto queryImage(ImageQueryDto imageDto) {
-		var resultData = createAIResult(imageDto);
-		return new ImageDto(resultData.answer(),
-				Base64.getEncoder().encodeToString(resultData.imageQueryDto().getImageContent()),
-				resultData.imageQueryDto().getImageType());
+	public List<ImageDto> queryImage(ImageQueryDto imageDto) {
+		var aiDocuments = this.documentVsRepository.retrieve(imageDto.getQuery(), MetaData.DataType.IMAGE).stream()
+				.filter(myDoc -> myDoc.getMetadata().get(MetaData.DATATYPE).equals(DataType.IMAGE.toString()))
+				.sorted((myDocA, myDocB) -> ((Float) myDocA.getMetadata().get(MetaData.DISTANCE))
+						.compareTo(((Float) myDocB.getMetadata().get(MetaData.DISTANCE))))
+				.toList();
+		var imageMap = this.imageRepository
+				.findAllById(aiDocuments.stream().map(myDoc -> (String) myDoc.getMetadata().get(MetaData.ID))
+						.map(myUuid -> UUID.fromString(myUuid)).limit(20).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(
+				aiDocuments.stream()
+						.filter(myDoc -> myEntry.getKey().toString()
+								.equals((String) myDoc.getMetadata().get(MetaData.ID)))
+						.findFirst().orElseThrow(),
+				myEntry.getValue(),
+				aiDocuments.stream()
+						.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();
 	}
 
 	private ResultData createAIResult(ImageQueryDto imageDto) {