근거 보여 주기와 모르면 모른다고 답하기

# 준비: "FAISS 벡터DB 에 저장하고 검색하기" 절과 같은 연습용 임베딩과 조항 네 개짜리 벡터DB
import zlib

from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings


class BigramEmbeddings(Embeddings):
    """글자 두 개 묶음(바이그램)을 세어 숫자 목록을 만드는 연습용 임베딩"""

    def __init__(self, size=512):
        self.size = size

    def _vector(self, text):
        text = text.replace(" ", "").replace("\n", "")
        vec = [0.0] * self.size
        for i in range(len(text) - 1):
            # 글자 두 개를 crc32 로 0~511 사이의 칸 번호로 바꿔 그 칸을 1 올린다
            vec[zlib.crc32(text[i:i + 2].encode("utf-8")) % self.size] += 1.0
        norm = sum(v * v for v in vec) ** 0.5 or 1.0
        return [v / norm for v in vec]

    def embed_documents(self, texts):
        return [self._vector(t) for t in texts]

    def embed_query(self, text):
        return self._vector(text)


emb = BigramEmbeddings()
docs = [
    Document("제5조 출발 10일 전까지 취소하면 전액 돌려준다.", metadata={"article": "제5조"}),
    Document("제10조 국내선 상품은 15kg 짐 한 개가 포함된다.", metadata={"article": "제10조"}),
    Document("제12조 반려동물은 7kg 이하의 개와 고양이만 된다.", metadata={"article": "제12조"}),
    Document("제14조 여권의 남은 기간이 6개월 이상이어야 한다.", metadata={"article": "제14조"}),
]
db = FAISS.from_documents(docs, emb)

THRESHOLD = 1.75     # 가장 가까운 조각의 거리가 이보다 크면 "찾지 못함"


def check(question):
    doc, score = db.similarity_search_with_score(question, k=1)[0]
    score = round(float(score), 3)
    if score > THRESHOLD:
        return f"찾지 못함 (가장 가까운 {doc.metadata['article']}, 거리 {score})"
    return f"모델에 물음 (근거 {doc.metadata['article']}, 거리 {score})"


for q in ["출발 10일 전에 취소하면 돌려받나요?", "환율이 얼마예요?", "강아지 데려가도 돼요?"]:
    print(q, "→", check(q))