feat: qdrant 정확성 수정

Author: yongho9064

backend/src/main/java/com/ai/lawyer/global/qdrant/initializer/QdrantInitializer.java

@@ -0,0 +1,46 @@
+package com.ai.lawyer.global.qdrant.initializer;
+
+import io.qdrant.client.QdrantClient;
+import jakarta.annotation.PostConstruct;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.stereotype.Component;
+
+import java.util.concurrent.ExecutionException;
+
+@Slf4j
+@Component
+@RequiredArgsConstructor
+public class QdrantInitializer {
+
+    private final QdrantClient qdrantClient;
+
+    @Value("${spring.ai.vectorstore.qdrant.collection-name}")
+    private String collectionName;
+
+    @Value("${spring.ai.vectorstore.qdrant.vector-size}")
+    private Long vectorSize;
+
+    @PostConstruct
+    private void existQdrantCollection() throws InterruptedException, ExecutionException {
+        var collections = qdrantClient.listCollectionsAsync().get();
+        boolean collectionExists = collections.stream()
+                .anyMatch(collection -> collection.equals(collectionName));
+
+        if (!collectionExists) {
+            log.info("'{}' 컬렉션이 존재하지 않아 새로 생성 중", collectionName);
+            qdrantClient.createCollectionAsync(
+                    collectionName,
+                    io.qdrant.client.grpc.Collections.VectorParams.newBuilder()
+                            .setSize(vectorSize.intValue())
+                            .setDistance(io.qdrant.client.grpc.Collections.Distance.Cosine)
+                            .build()
+            ).get();
+            log.info("'{}' 컬렉션 생성 완료", collectionName);
+        } else {
+            log.info("'{}' 컬렉션이 이미 존재합니다.", collectionName);
+        }
+    }
+
+}

backend/src/main/java/com/ai/lawyer/global/qdrant/service/QdrantService.java

@@ -7,35 +7,64 @@
 import org.springframework.ai.vectorstore.filter.Filter;
 import org.springframework.stereotype.Service;
 
+import java.util.ArrayList;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.List;
+import java.util.stream.Collectors;
 
 @Service
 @RequiredArgsConstructor
 public class QdrantService {
 
     private final VectorStore vectorStore;
 
-    public List<Document> searchDocument(String query, String key, String value, int topK) {
+    public List<Document> searchDocument(String query, String key, String value) {
 
-        SearchRequest caseSearchRequest = SearchRequest.builder()
-                .query(query)
-                .topK(topK)
+        SearchRequest findCaseNumberRequest = SearchRequest.builder()
+                .query(query).topK(1)
                 .filterExpression(new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key(key), new Filter.Value(value)))
                 .build();
-        List<Document> similarCaseDocuments = vectorStore.similaritySearch(caseSearchRequest);
-
-        if (caseSearchRequest == null) {
-            return Collections.singletonList(
-                    Document.builder()
-                            .text("더미")
-                            .metadata(key, value)
-                            .score(0.0)
-                            .build()
-            );
+        List<Document> mostSimilarDocuments = vectorStore.similaritySearch(findCaseNumberRequest);
+
+
+        if (mostSimilarDocuments.isEmpty()) {
+            return Collections.emptyList();
+        }
+        String targetCaseNumber = (String) mostSimilarDocuments.get(0).getMetadata().get("caseNumber");
+        if (targetCaseNumber == null) {
+            return mostSimilarDocuments;
+        }
+
+        SearchRequest fetchAllChunksRequest = SearchRequest.builder()
+                .query(query).topK(100)
+                .filterExpression(new Filter.Expression(Filter.ExpressionType.EQ, new Filter.Key("caseNumber"), new Filter.Value(targetCaseNumber)))
+                .build();
+        List<Document> allChunksOfCase = new ArrayList<>(vectorStore.similaritySearch(fetchAllChunksRequest));
+
+        if (allChunksOfCase.isEmpty()) {
+            return Collections.emptyList();
         }
 
-        return similarCaseDocuments;
+        allChunksOfCase.sort(Comparator.comparingInt(doc ->
+                ((Number) doc.getMetadata().get("chunkIndex")).intValue()
+        ));
+
+        String mergedContent = allChunksOfCase.stream()
+                .map(Document::getText)
+                .collect(Collectors.joining(""));
+
+        Document bestScoringDoc = allChunksOfCase.stream()
+                .max(Comparator.comparing(Document::getScore))
+                .orElse(allChunksOfCase.get(0));
+
+        Document finalDocument = Document.builder()
+                .text(mergedContent)
+                .metadata(bestScoringDoc.getMetadata())
+                .score(bestScoringDoc.getScore())
+                .build();
+
+        return Collections.singletonList(finalDocument);
     }
 
 }