RAG 파이프라인 만들기

!pip install -q openai chromadb

from pathlib import Path
import numpy as np
from openai import OpenAI

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=Path("openrouter-key.txt").read_text().strip(),
)
model = "openai/gpt-4.1-mini"
embedding_model = "openai/text-embedding-3-small"

def chunk_with_heading(text):  # 청크마다 절 제목과 본문을 담은 사전
    chunks = []
    for section in text.split("\n## ")[1:]:
        title, body = section.split("\n", 1)
        chunks += [{"section": title, "text": f"[{title}] {p.strip()}"}
                   for p in body.split("\n\n") if p.strip()]
    return chunks

from mindmall_guide import guide
chunks = chunk_with_heading(guide)
len(chunks)

def embed(texts):
    response = client.embeddings.create(model=embedding_model, input=texts)
    return [d.embedding for d in response.data]

import chromadb

chroma = chromadb.Client()
collection = chroma.get_or_create_collection(
    name="mindmall",
    configuration={"hnsw": {"space": "cosine"}},  # 코사인 거리 사용
)

texts = [c["text"] for c in chunks]
collection.add(
    ids=[f"chunk-{i}" for i in range(len(chunks))],  # 청크마다 고유한 id
    documents=texts,                                 # 원문
    embeddings=embed(texts),                         # 임베딩
    metadatas=[{"section": c["section"]} for c in chunks],  # 부가 정보
)
collection.count()

def retrieve(question, k=3):  # 가까운 청크 k개의 원문·절 제목·거리
    r = collection.query(query_embeddings=embed([question]), n_results=k)
    return [{"text": doc, "section": meta["section"], "distance": round(dist, 3)}  # 질문이 하나이므로 [0]
            for doc, meta, dist in zip(r["documents"][0], r["metadatas"][0], r["distances"][0])]

for hit in retrieve("골드 회원인데 제주도로 받으면 배송비가 있나요?"):
    print(hit["distance"], hit["text"][:60])

def build_prompt(question, hits):
    sources = "\n\n".join(f"[{i}] {hit['text']}" for i, hit in enumerate(hits, start=1))
    return f"""당신은 마인드몰 고객센터 직원입니다. 아래 자료만 근거로 답하세요.
- 답변의 각 문장 끝에 근거가 된 자료 번호를 [1]처럼 다세요.
- 자료에 없는 내용은 추측하지 말고 "안내서에서 찾을 수 없습니다"라고 답하세요.

[자료]
{sources}

[질문]
{question}"""

def answer(question, k=3):  # 검색 → 프롬프트 구성 → 생성 → 출처 목록 붙이기
    hits = retrieve(question, k)
    response = client.chat.completions.create(
        model=model, messages=[{"role": "user", "content": build_prompt(question, hits)}])
    sources = "\n".join(f"[{i}] {hit['section']}" for i, hit in enumerate(hits, start=1))
    return f"{response.choices[0].message.content}\n\n출처\n{sources}"

print(answer("골드 회원인데 제주도로 받으면 배송비가 있나요?"))

print(answer("해외로도 배송해 주나요?"))