실습 4 — 벡터DB 세 가지 견주기

import json, os, shutil, time
from pathlib import Path
import numpy as np
import sys
sys.path.insert(0, ".")                        # 작업 폴더(data)에 만든 out_offline.py 를 찾게 한다
import out_offline
out_offline.guard()                               # 라이브러리를 부르기 전에 막는다(실습 1의 파일)

import faiss
import chromadb
from chromadb.config import Settings
from qdrant_client import QdrantClient, models
from sentence_transformers import SentenceTransformer

chunks = [json.loads(line) for line in open("out_chunks.jsonl", encoding="utf-8")]
emb = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")
E = emb.encode(["passage: " + c["글"] for c in chunks], normalize_embeddings=True)
question = "법인카드로 한 번에 얼마까지 쓸 수 있나요?"
q = emb.encode("query: " + question, normalize_embeddings=True)
print("조각", E.shape, "· 질문", q.shape)

ids = np.arange(len(chunks), dtype="int64")             # FAISS 와 Qdrant 는 정수 id 를 쓴다
timing = {}

# FAISS: IndexIDMap 으로 감싸야 정해 준 id 로 넣고 지울 수 있다
t = time.time()
fx = faiss.IndexIDMap(faiss.IndexFlatIP(E.shape[1]))
fx.add_with_ids(E, ids)
faiss.write_index(fx, "out_intranet.faiss")
timing["FAISS"] = time.time() - t

# Chroma: 폴더에 저장한다. 다시 돌릴 때를 위해 같은 이름의 컬렉션이 있으면 지운다
t = time.time()
cc = chromadb.PersistentClient(path="out_chroma11", settings=Settings(anonymized_telemetry=False))
if "intranet" in [c.name for c in cc.list_collections()]:
    cc.delete_collection("intranet")
col = cc.create_collection("intranet", metadata={"hnsw:space": "cosine"})       # 코사인 거리로 찾는다
col.add(ids=[c["id"] for c in chunks], embeddings=E.tolist(), documents=[c["글"] for c in chunks],
        metadatas=[{k: c[k] for k in ("경로", "제목", "등급", "버전")} for c in chunks])
timing["Chroma"] = time.time() - t

# Qdrant: path 를 주면 서버 없이 폴더에 저장한다
t = time.time()
qc = QdrantClient(path="out_qdrant11")
if qc.collection_exists("intranet"):
    qc.delete_collection("intranet")
qc.create_collection("intranet", vectors_config=models.VectorParams(size=E.shape[1], distance=models.Distance.COSINE))
qc.upsert("intranet", points=[models.PointStruct(id=int(i), vector=E[i].tolist(), payload=chunks[i]) for i in ids])
timing["Qdrant"] = time.time() - t

# 같은 질문의 상위 3개를 셋에서 찾아 견준다
_, I = fx.search(q[None], 3)
r_c = col.query(query_embeddings=[q.tolist()], n_results=3)
r_q = qc.query_points("intranet", query=q.tolist(), limit=3).points
print("FAISS :", [chunks[i]["id"] for i in I[0]])
print("Chroma:", r_c["ids"][0])
print("Qdrant:", [p.payload["id"] for p in r_q])
print("넣는 데 걸린 시간:", {k: f"{v:.2f}초" for k, v in timing.items()})

level = 1                                                # 협력사 직원은 공개(1) 문서만 본다
allowed = ids[[c["등급"] <= level for c in chunks]]     # FAISS 는 볼 수 있는 id 목록을 우리가 만든다

_, I = fx.search(q[None], 3, params=faiss.SearchParameters(sel=faiss.IDSelectorBatch(allowed)))
r_c = col.query(query_embeddings=[q.tolist()], n_results=3, where={"등급": {"$lte": level}})
r_q = qc.query_points("intranet", query=q.tolist(), limit=3,
                      query_filter=models.Filter(must=[models.FieldCondition(key="등급", range=models.Range(lte=level))])).points
print("FAISS :", [chunks[i]["제목"] for i in I[0]])
print("Chroma:", [m["제목"] for m in r_c["metadatas"][0]])
print("Qdrant:", [p.payload["제목"] for p in r_q])

target = "it/VPN_접속안내.txt"
del_ids = ids[[c["경로"] == target for c in chunks]]
print("지울 조각:", len(del_ids), "개")

# 셋에서 같은 문서를 지운다. FAISS 는 id 로, Chroma 와 Qdrant 는 메타데이터 조건으로 지운다
fx.remove_ids(del_ids)
col.delete(where={"경로": target})
qc.delete("intranet", points_selector=models.FilterSelector(filter=models.Filter(
    must=[models.FieldCondition(key="경로", match=models.MatchValue(value=target))])))
print("지운 뒤 개수:", fx.ntotal, col.count(), qc.count("intranet").count)

# 원래대로 다시 넣는다(실습 6이 이 Chroma 를 쓴다)
back = [int(i) for i in del_ids]
fx.add_with_ids(E[back], del_ids)
col.add(ids=[chunks[i]["id"] for i in back], embeddings=E[back].tolist(), documents=[chunks[i]["글"] for i in back],
        metadatas=[{k: chunks[i][k] for k in ("경로", "제목", "등급", "버전")} for i in back])
qc.upsert("intranet", points=[models.PointStruct(id=i, vector=E[i].tolist(), payload=chunks[i]) for i in back])
print("다시 넣은 뒤:", fx.ntotal, col.count(), qc.count("intranet").count)
faiss.write_index(fx, "out_intranet.faiss")
qc.close()                                              # Qdrant 로컬 모드는 폴더를 잠가 두므로 닫는다


def size_kb(path):
    p = Path(path)
    files = [p] if p.is_file() else [f for f in p.rglob("*") if f.is_file()]
    return round(sum(f.stat().st_size for f in files) / 1024)


print("저장 크기(KB):", {"FAISS": size_kb("out_intranet.faiss"), "Chroma": size_kb("out_chroma11"),
                       "Qdrant": size_kb("out_qdrant11")})
print("막은 주소:", out_offline.BLOCKED)