실습 5 — 접근 제어가 있는 검색

import json
from pathlib import Path

import numpy as np
import pandas as pd
from sentence_transformers import SentenceTransformer

docs = [json.loads(line) for line in Path("secure_docs.jsonl").read_text(encoding="utf-8").splitlines()]
USERS = {
    "인턴": {"등급": 1, "부서": "개발팀"},
    "개발팀원": {"등급": 2, "부서": "개발팀"},
    "인사팀장": {"등급": 3, "부서": "인사팀"},
}
emb = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")
D = emb.encode(["passage: " + d["제목"] + ": " + d["내용"] for d in docs], normalize_embeddings=True)


def can_read(user, doc):
    if user["등급"] < doc["등급"]:
        return False
    return doc["등급"] < 3 or doc["부서"] in ("공용", user["부서"])


# 결과: 사용자마다 볼 수 있는 문서 수
print({name: sum(can_read(u, d) for d in docs) for name, u in USERS.items()})

def search(question, user, how, k=3):
    """how: 'none' 걸러 내지 않음 / 'after' 검색 뒤 걸러 내기 / 'before' 검색 전에 걸러 내기"""
    q = emb.encode(["query: " + question], normalize_embeddings=True)[0]
    order = list(np.argsort(-(D @ q)))
    if how == "before":
        order = [i for i in order if can_read(user, docs[i])]
    top = order[:k]
    if how == "after":
        top = [i for i in top if can_read(user, docs[i])]
    return [docs[i] for i in top]


QUESTIONS = ["내년 연봉 인상률이 얼마예요?", "개발 서버 관리자 비밀번호는 얼마마다 바꿔요?",
             "보안 사고가 나면 누구에게 알려요?", "구조 조정 계획이 있나요?", "재택 근무는 며칠까지 돼요?"]
rows = []
for name, user in USERS.items():
    for how in ["none", "after", "before"]:
        got = [search(q, user, how) for q in QUESTIONS]
        rows.append({"사용자": name, "방법": how,
                     "볼 수 없는 문서가 넘어간 수": sum(not can_read(user, d) for g in got for d in g),
                     "넘긴 조각 평균": np.mean([len(g) for g in got])})
# 결과: 사용자와 방법마다 LLM 에 넘어간 '볼 수 없는 문서' 수
print(pd.DataFrame(rows).round(2).to_string(index=False))

q = "재택 근무는 며칠까지 돼요?"
for how in ["after", "before"]:
    # 결과: 인턴에게 넘어가는 조각의 제목과 등급
    print(how, [(d["제목"], d["등급"]) for d in search(q, USERS["인턴"], how)])
print("개발팀원 before", [(d["제목"], d["등급"]) for d in search(q, USERS["개발팀원"], "before")])

import sys

sys.path.insert(0, ".")
import out_secbot as bot

SYSTEM = "사내 규정 도우미입니다. 자료로만 한 문장으로 답하고, 자료에 없으면 '자료에 없습니다'라고 답하세요."
q = "내년 연봉 인상률이 얼마예요?"
for how in ["none", "before"]:
    ctx = search(q, USERS["인턴"], how)
    answer = bot.chat(SYSTEM, "자료:\n" + "\n".join(f"{d['제목']}: {d['내용']}" for d in ctx) + "\n\n질문: " + q)
    # 결과: 걸러 내지 않으면 인턴에게 비밀(인상률)이 그대로 나간다
    print(f"{how:6s} → {answer}")