본문 바로가기
스터디/RAG

[코드리뷰] VectorStore - 리팩토링2

by gnosis90 2026. 4. 20.


app.py로부터 호출되는 메서드 부터 시작한다.

# 5.2. 벡터 DB 검색(42 line)
        results = store.search(query_vector, k=3)

 


1. 주요 리팩토링 포인트

1 설정값 하드코딩 제거(config.py) 활용

현재 chroma_db 경로, 컬렉션 이름 등이 코드 내부에 박혀 있는 상태.

서비스 규모가 커지면 운영 환경(Dev, Prod)에 따라 이 값들이 달라져야 하므로 외부 설정을 주입받아야 한다.

 

2. 대량 데이터 처리를 위한 배치 로직

add_do

 

3. 예외 처리 및 로깅

실제 서비스에서는 DB 연결 실패나 검색 오류 시 시스템이 뻗지 않도록 try-except와 로깅이 필수적이다.

 

 

 

하위의 코드는 리팩토링 하기 전 코드

# 검색
    # : 유사도 임계값(threshold)를 설정함.
    def search(self, query_vector, k=3, threshold=0.4):  # threshold 추가
        
        results = self.collection.query(
              query_embeddings=[query_vector]
            , n_results=k
        )

        output = []

        for i in range(len(results["documents"][0])):
            
            distance = results["distances"][0][i]
            similarity = 1 - distance
            
            if similarity < threshold:  # 임계값 이하 제거
                continue
            
            output.append({
                "text": results["documents"][0][i],
                "page": results["metadatas"][0][i]["page"],
                "chunk": results["metadatas"][0][i]["chunk"],
                "distance": results["distances"][0][i]  # 추가
            })

        return output
import chromadb

class VectorStore:

    def __init__(self):
        
        # 유사도가 음수 
        self.client = chromadb.PersistentClient(path="chroma_db")
        self.collection = self.client.get_or_create_collection(  name="documents"
                                                               , metadata={"hnsw:space": "cosine"}  # 추가
                                                               )
        self.current_id = self.collection.count()

    def add_documents(self, vectors, docs):
        ids = []
        embeddings = []
        documents = []
        metadatas = []

        for i, doc in enumerate(docs):
            ids.append(str(self.current_id))
            embeddings.append(vectors[i])
            documents.append(doc["text"])
            metadatas.append({
                "page": doc["page"],
                "chunk": doc["chunk"]
            })
            self.current_id += 1

        self.collection.add(
            ids=ids,
            embeddings=embeddings,
            documents=documents,
            metadatas=metadatas
        )

    def save(self):
        pass

    def print_documents(self):
        data = self.collection.get()

        print("\n===== VECTOR COUNT =====")
        print(self.collection.count())

        print("\n===== DOCUMENTS =====\n")

        for i in range(len(data["ids"])):
            print(f"ID: {data['ids'][i]}")
            print(f"PAGE: {data['metadatas'][i]['page']} | CHUNK: {data['metadatas'][i]['chunk']}")
            print(data["documents"][i])
            print("-" * 60)

    # 검색
    # : 유사도 임계값(threshold)를 설정함.
    def search(self, query_vector, k=3, threshold=0.4):  # threshold 추가
        
        results = self.collection.query(
              query_embeddings=[query_vector]
            , n_results=k
        )

        output = []

        for i in range(len(results["documents"][0])):
            
            distance = results["distances"][0][i]
            similarity = 1 - distance
            
            if similarity < threshold:  # 임계값 이하 제거
                continue
            
            output.append({
                "text": results["documents"][0][i],
                "page": results["metadatas"][0][i]["page"],
                "chunk": results["metadatas"][0][i]["chunk"],
                "distance": results["distances"][0][i]  # 추가
            })

        return output

리팩토링 한 후 코드

def search(self, query_vector: List[float], k: int = 3, threshold: float = 0.4) -> List[Dict[str, Any]]:
        """
        검색 로직 - 결과가 없을 경우의 처리와 데이터 구조 최적화
        """
        try:
            results = self.collection.query(
                query_embeddings=[query_vector],
                n_results=k
            )

            # 결과가 아예 없는 경우 방어 코드
            if not results or not results["documents"] or not results["documents"][0]:
                return []

            output = []
            # results["distances"][0] 등의 리스트를 zip으로 묶으면 훨씬 '파이썬'스럽고 깔끔합니다.
            for doc, meta, dist in zip(results["documents"][0], results["metadatas"][0], results["distances"][0]):
                
                similarity = 1 - dist
                if similarity >= threshold:
                    output.append({
                        "text": doc,
                        "page": meta.get("page"),
                        "chunk": meta.get("chunk"),
                        "similarity": round(similarity, 4)
                    })
            
            # 유사도 높은 순으로 정렬 (이미 되어있겠지만 명시적 확인)
            return sorted(output, key=lambda x: x['similarity'], reverse=True)

        except Exception as e:
            logger.error(f"검색 중 오류 발생: {e}")
            return []

 

 

for doc, meta, dist in zip(results["documents"][0], results["metadatas"][0], results["distances"][0]):

 

zip

  • 문장 리스트: ["안녕하세요", "반갑습니다", "질문입니다"]
  • 메타데이터 리스트: [{"page":1}, {"page":2}, {"page":3}]
  • 거리(유사도) 리스트: [0.1, 0.2, 0.3]

이 세 리스트는 인덱스 번호(0번, 1번, 2번...)가 같으면 서로 한 세트인 데이터이다.

 

2. 자바식 접근 (불편함)

for (int i = 0; i < docs.size(); i++) {

    String doc  = docs.get(i);
    Map meta    = metas.get(i);
    Double dist = dists.get(i);
    
    // 여기서 doc, meta, dist를 합쳐서 결과 리스트에 담음
};

 

파이썬은 어차피 인덱스 똑같은 것 끼리 묶어서 쓸거니 내가 지퍼(Zip)처럼 딱 맞물리게 인덱스가 같은 것 끼리 묶어줄게라고 제안한것. zip(리스트A, 리스트B, 리스트C)를 실행하면서, 각 리스트의 0번끼리 1세트, 1번끼리 1세트... 이렇게 묶인 꾸러미를 차례대로 던져준다.


similarity = 1 - dist

 

왜 하필 1 - dist 인가?

유사도는 비슷한 정도이기 때문에 차이라 볼 수 있는 거리를 빼면 이게 유사도가 된다.

달리 말하면, 거리와 유사도는 반비례 관계이기 때문이다.

4. 왜 이게 중요할까? (Threshold의 기준)

아까 코드에서 if similarity < threshold: continue 라고 하셨죠?

  • 만약 threshold = 0.4라면:
  • **"유사도가 40% 미만인 놈들은 검색 결과에서 빼버려!"**라는 뜻이 됩니다.
  • 즉, 거리가 너무 먼($0.6$ 이상인) 데이터는 신뢰할 수 없으니 답변 생성에 쓰지 않겠다는 아주 중요한 방어 로직인 셈입니다.