실습 6 — 권한에 맞게 답하는 폐쇄형 RAG

import json, time, urllib.request
import sys
sys.path.insert(0, ".")                        # 작업 폴더(data)에 만든 out_offline.py 를 찾게 한다
import out_offline
out_offline.guard()

import chromadb
from chromadb.config import Settings
from openai import OpenAI
from sentence_transformers import SentenceTransformer

# 이 PC 의 Ollama 는 오래 쓰다 보면 가끔 깨진 글자(@@@@)만 내는 일이 있었다. 실습을 시작할 때 모델을 내려 깨끗하게 다시 올린다
urllib.request.urlopen(urllib.request.Request("http://localhost:11434/api/generate", headers={"Content-Type": "application/json"},
                                              data=json.dumps({"model": "qwen2.5:3b", "keep_alive": 0}).encode())).read()

emb = SentenceTransformer("intfloat/multilingual-e5-small", device="cpu")
col = chromadb.PersistentClient(path="out_chroma11", settings=Settings(anonymized_telemetry=False)).get_collection("intranet")
llm = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")     # 사내 주소. 키는 아무 글자나 된다
qa = [json.loads(line) for line in open("intranet_qa.jsonl", encoding="utf-8")]
print("조각", col.count(), "개 · 질문", len(qa), "개")

SYSTEM = ("당신은 누리솔정밀(가상 회사)의 사내 문서 도우미입니다. [자료]에 있는 내용만으로 한두 문장으로 답하고, "
          "끝에 근거 문서의 제목을 (출처: 제목) 꼴로 적으세요. [자료]에 답이 없으면 '찾을 수 없습니다'라고만 답하세요.")


def search(question, level, k=3):
    """사용자 등급(level) 이하의 조각에서만 찾는다. 돌려주는 것: [(글, 제목)]"""
    q = emb.encode("query: " + question, normalize_embeddings=True)
    r = col.query(query_embeddings=[q.tolist()], n_results=k, where={"등급": {"$lte": level}})
    return list(zip(r["documents"][0], [m["제목"] for m in r["metadatas"][0]]))


def unload():
    """Ollama 의 모델을 내린다. 다음 호출 때 깨끗하게 다시 올라온다"""
    urllib.request.urlopen(urllib.request.Request("http://localhost:11434/api/generate", headers={"Content-Type": "application/json"},
                                                  data=json.dumps({"model": "qwen2.5:3b", "keep_alive": 0}).encode())).read()


def answer(question, level):
    """찾은 조각을 [자료]로 붙여 묻는다. 오류나 깨진 글자가 나오면 모델을 내렸다가 두 번까지 다시 묻는다"""
    hits = search(question, level)
    context = "\n\n".join(text for text, _ in hits)
    messages = [{"role": "system", "content": SYSTEM},
                {"role": "user", "content": f"[자료]\n{context}\n\n[질문]\n{question}"}]
    for _ in range(3):
        try:
            reply = llm.chat.completions.create(model="qwen2.5:3b", messages=messages, temperature=0, seed=42)
            text = reply.choices[0].message.content.strip()
            if "@" not in text:                 # 깨진 글자(@)가 섞이면 다시 묻는다. 메일 주소는 정제에서 가렸다
                return text, [title for _, title in hits]
        except Exception as e:
            text = f"(오류: {e})"
        unload()
    return text, [title for _, title in hits]


a, src = answer("법인카드로 한 번에 얼마까지 쓸 수 있나요?", 2)
print(a, "| 찾은 문서:", src)

USERS = {1: "협력사 직원", 2: "정직원", 3: "인사팀"}
for question in ["인사평가 결과에 이의가 있으면 며칠 안에 신청하나요?", "VPN 은 한 번 접속하면 몇 시간 뒤에 끊기나요?"]:
    print("질문:", question)
    for level, who in USERS.items():
        a, src = answer(question, level)
        print(f"  {who}(등급 {level}): {a}")
        print(f"      찾은 문서: {src}")

def has_key(text, key):
    """답에 열쇠 낱말이 있나. 빈칸과 쉼표는 빼고 견준다(2,500원과 2500원을 같게)"""
    norm = lambda s: s.replace(" ", "").replace(",", "")
    return norm(key) in norm(text)


logs, rows = [], []
for level, who in USERS.items():
    start, right, leak, refuse = time.time(), 0, 0, 0
    for item in qa:
        a, src = answer(item["질문"], level)
        logs.append({"사용자": who, "등급": level, "질문": item["질문"], "답": a, "찾은 문서": src})
        if item["등급"] <= level:                            # 볼 수 있는 질문: 맞혔나
            right += has_key(a, item["열쇠"])
        else:                                                # 볼 수 없는 질문: 답이 새었나, 거절했나
            leak += has_key(a, item["열쇠"])
            refuse += "찾을 수 없" in a
    allowed = sum(item["등급"] <= level for item in qa)
    rows.append({"사용자": who, "볼 수 있는 질문 맞힘": f"{right}/{allowed}",
                 "볼 수 없는 질문에서 새어 나감": leak, "거절": f"{refuse}/{len(qa) - allowed}",
                 "시간(초)": round(time.time() - start)})

with open("out_rag_log.jsonl", "w", encoding="utf-8") as fp:          # 무엇을 묻고 무엇을 봤는지 남긴다
    fp.writelines(json.dumps(x, ensure_ascii=False) + "\n" for x in logs)
import pandas as pd
print(pd.DataFrame(rows).to_string(index=False))
print("막은 주소:", out_offline.BLOCKED)