import numpy as np
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
client = OpenAI()   # 임베딩 + 생성 모두 OpenAI


# 1) 정답 — OpenAI API로 임베딩
def get_embedding(text):
    response = client.embeddings.create(input=text, model="text-embedding-ada-002")
    return np.array(response.data[0].embedding, dtype=np.float32)


documents = [
    "Python은 강력한 프로그래밍 언어입니다.",
    "OpenAI는 AI 연구를 선도하는 기업입니다.",
    "FAISS는 벡터 검색을 위한 라이브러리입니다.",
]

# 2) 정답 — 문서들을 임베딩해 둔다 (faiss 없이 numpy 배열로 보관)
doc_embeddings = np.array([get_embedding(doc) for doc in documents], dtype=np.float32)


# 3) 정답 — 검색 + 생성(RAG)
def rag_query(user_query):
    # ① 질문도 '같은 모델'로 임베딩 → 코사인 유사도가 가장 큰 문서 (numpy)
    query_vec = get_embedding(user_query)
    sims = doc_embeddings @ query_vec / (
        np.linalg.norm(doc_embeddings, axis=1) * np.linalg.norm(query_vec))
    retrieved_doc = documents[int(np.argmax(sims))]

    # ② 검색된 문서를 프롬프트에 끼워넣어 LLM에게 전달 = RAG의 핵심
    prompt = f"참고 정보: {retrieved_doc}\n\n질문: {user_query}\n답변:"
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "당신은 친절한 AI 도우미입니다."},
            {"role": "user", "content": prompt},
        ],
    )
    return retrieved_doc, response.choices[0].message.content


if __name__ == "__main__":
    q = "FAISS가 뭐야?"
    doc, answer = rag_query(q)
    print("질문:", q)
    print("검색된 문서:", doc)
    print("답변:", answer)
