import numpy as np
from sentence_transformers import SentenceTransformer

data = np.load("index.npz", allow_pickle=True)
vecs, slugs, texts = data["vecs"], data["slugs"], data["texts"]

model = SentenceTransformer("intfloat/multilingual-e5-small")


def search(question, k=3):
    """질문과 가까운 조각 k 개를 찾는다."""
    qv = model.encode([question], normalize_embeddings=True)[0]
    scores = vecs @ qv

    # 17강과 다른 곳은 여기 한 줄이다. 글 단위로 묶지 않고
    # 점수 높은 순서대로 조각을 그냥 k 개 집는다
    out = []
    for i in scores.argsort()[::-1][:k]:
        out.append((float(scores[i]), str(slugs[i]), str(texts[i])))
    return out


for score, slug, text in search("조각을 겹치게 자르는 이유가 뭔가"):
    print(f"{score:.3f}  {slug}")
    print(f"        {text[:40]}")

# (21_rag_search.py 의 위쪽을 그대로 복사해 온다 — data, model, search)

question = "조각을 겹치게 자르는 이유가 뭔가"
hits = search(question, k=3)

# 조각 셋을 빈 줄 하나로 이어 붙인다
context = "\n\n".join(text for score, slug, text in hits)

prompt = f"""아래 자료를 읽고 질문에 답해라.

[자료]
{context}

[질문]
{question}
"""

print(f"자료 {len(context)}자 / 프롬프트 전체 {len(prompt)}자")
print("-" * 40)
print(prompt)

# (21 · 22 의 위쪽을 그대로 복사해 온다 — search, question, hits, context, prompt)
from openai import OpenAI

client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")

res = client.chat.completions.create(
    model="qwen2.5:3b",
    messages=[{"role": "user", "content": prompt}],   # 질문 대신 프롬프트 전체
)

print(res.choices[0].message.content)

import numpy as np
from sentence_transformers import SentenceTransformer
from openai import OpenAI

data = np.load("index.npz", allow_pickle=True)
vecs, slugs, texts = data["vecs"], data["slugs"], data["texts"]

model = SentenceTransformer("intfloat/multilingual-e5-small")
client = OpenAI(base_url="http://localhost:11434/v1", api_key="ollama")


def search(question, k=3):
    qv = model.encode([question], normalize_embeddings=True)[0]
    scores = vecs @ qv
    return [
        (float(scores[i]), str(slugs[i]), str(texts[i]))
        for i in scores.argsort()[::-1][:k]
    ]


def ask(question, k=3):
    """검색 -> 조립 -> 생성. 답과 근거를 같이 돌려준다."""
    hits = search(question, k)
    context = "\n\n".join(text for score, slug, text in hits)

    prompt = f"""아래 자료를 읽고 질문에 답해라.

[자료]
{context}

[질문]
{question}
"""
    res = client.chat.completions.create(
        model="qwen2.5:3b",
        messages=[{"role": "user", "content": prompt}],
    )
    # 답만 주면 어디서 왔는지 알 수 없다. 근거 글 이름도 같이 돌려준다
    return res.choices[0].message.content, [slug for score, slug, text in hits]


for q in ["브랜치는 왜 따로 만드나", "엑셀 파일을 파이썬으로 어떻게 여나"]:
    answer, sources = ask(q)
    print(f"\n[질문] {q}")
    print(f"[근거] {sources}")
    print(f"[답] {answer[:80]}")