Upload converters/pipeline/step5_rag.py

03.Code/업로드용/converters/pipeline/step5_rag.py

@@ -0,0 +1,138 @@
+# -*- coding: utf-8 -*-
+from dotenv import load_dotenv
+load_dotenv()
+
+"""
+build_rag.py
+
+기능:
+- chunk_and_summary.py에서 생성된 output/rag/*_chunks.json 파일들을 읽어들임
+- 각 청크의 text + summary를 활용하여 임베딩(text-embedding-3-small) 생성
+- FAISS Index (IndexFlatIP)를 구성하여 저장
+- 최종적으로 rag/faiss.index, meta.json, vectors.npy 파일을 생성함
+"""
+import os
+import sys
+import json
+from pathlib import Path
+import numpy as np
+import faiss
+from openai import OpenAI
+
+
+# ===== OpenAI 설정 (구조 유지) =====
+OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+GPT_MODEL      = "gpt-5-2025-08-07"
+EMBED_MODEL    = "text-embedding-3-small"
+client = OpenAI(api_key=OPENAI_API_KEY)
+
+
+def log(msg: str):
+    print(msg, flush=True)
+    with (LOG_DIR / "build_rag_log.txt").open("a", encoding="utf-8") as f:
+        f.write(msg + "\n")
+
+
+def embed_texts(texts):
+    if not texts:
+        return np.zeros((0, 1536), dtype="float32")
+    embs = []
+    B = 96
+    for i in range(0, len(texts), B):
+        batch = texts[i:i+B]
+        resp = client.embeddings.create(model=EMBED_MODEL, input=batch)
+        for d in resp.data:
+            embs.append(np.array(d.embedding, dtype="float32"))
+    return np.vstack(embs)
+
+
+def _build_embed_input(u: dict) -> str:
+    """
+    text + summary를 합쳐서 임베딩용 텍스트를 만듭니다.
+    - text, summary 중 내용이 있는 것 사용
+    - 공백 정리
+    - 최대 길이 제한
+    """
+    sum_ = (u.get("summary") or "").strip()
+    txt  = (u.get("text") or "").strip()
+
+    if txt and sum_:
+        merged = txt + "\n\n요약: " + sum_[:1000]
+    else:
+        merged = txt or sum_
+
+    merged = " ".join(merged.split())
+    if not merged:
+        return ""
+    if len(merged) > 4000:
+        merged = merged[:4000]
+    return merged
+
+
+def build_faiss_index():
+    docs = []
+    metas = []
+
+    rag_files = list(RAG_DIR.glob("*_chunks.json"))
+    if not rag_files:
+        log("RAG 대상 파일이 없습니다. 먼저 chunk_and_summary.py를 실행하십시오.")
+        sys.exit(1)
+
+    for f in rag_files:
+        try:
+            units = json.loads(f.read_text(encoding="utf-8", errors="ignore"))
+        except Exception as e:
+            log(f"[WARN] RAG 로드 실패: {f.name} | {e}")
+            continue
+
+        for u in units:
+            embed_input = _build_embed_input(u)
+            if not embed_input:
+                continue
+            if len(embed_input) < 40:
+                continue
+            docs.append(embed_input)
+            metas.append({
+                "source": u.get("source", ""),
+                "chunk": int(u.get("chunk", 0)),
+                "folder_context": u.get("folder_context", "")
+            })
+
+    if not docs:
+        log("임베딩할 텍스트 내용이 없습니다.")
+        sys.exit(1)
+
+    log(f"총 {len(docs)}개 청크 임베딩 시작...")
+    vectors = embed_texts(docs)
+
+    # 정규화 (Inner Product 사용 시 필수)
+    faiss.normalize_L2(vectors)
+
+    # FAISS 인덱스 생성
+    d = 1536
+    index = faiss.IndexFlatIP(d)
+    index.add(vectors)
+
+    # 결과 저장
+    faiss.write_index(index, str(RAG_DIR / "faiss.index"))
+    (RAG_DIR / "meta.json").write_text(json.dumps(metas, ensure_ascii=False, indent=2), encoding="utf-8")
+    np.save(str(RAG_DIR / "vectors.npy"), vectors)
+
+    log(f"FAISS 인덱스 저장 완료: {RAG_DIR / 'faiss.index'}")
+
+
+def main(input_dir, output_dir):
+    global OUTPUT_ROOT, RAG_DIR, LOG_DIR
+    OUTPUT_ROOT = Path(output_dir)
+    RAG_DIR     = OUTPUT_ROOT / "rag"
+    LOG_DIR     = OUTPUT_ROOT / "logs"
+    for d in [RAG_DIR, LOG_DIR]:
+        d.mkdir(parents=True, exist_ok=True)
+
+    log("=== RAG 벡터 인덱스 구축 시작 ===")
+    build_faiss_index()
+    log("=== FAISS RAG 인덱스 구축 완료 ===")
+
+
+if __name__ == "__main__":
+    main()