RAG 검색 증강 생성

pip install openai chromadb tiktoken langchain langchain-openai
pip install pypdf python-docx unstructured

from openai import OpenAI
import numpy as np

client = OpenAI()

# ---- 단일 텍스트 벡터화 ----
response = client.embeddings.create(
    model="text-embedding-3-small",   # 추천, 저렴하고 성능 우수
    input="Python은 프로그래밍 언어입니다"
)
vector = response.data[0].embedding
print(f"벡터 차원: {len(vector)}")      # 1536
print(f"처음 5개 값: {vector[:5]}")

# ---- 배치 벡터화 ----
texts = [
    "Python은 프로그래밍 언어입니다",
    "Java도 프로그래밍 언어입니다",
    "오늘 날씨가 정말 좋네요",
]
response = client.embeddings.create(
    model="text-embedding-3-small",
    input=texts
)
vectors = [d.embedding for d in response.data]

# ---- 유사도 계산 ----
def cosine_similarity(a, b):
    a, b = np.array(a), np.array(b)
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

print(f"Python vs Java: {cosine_similarity(vectors[0], vectors[1]):.4f}")
# → 0.89 (매우 유사! 모두 프로그래밍 언어)

print(f"Python vs 날씨: {cosine_similarity(vectors[0], vectors[2]):.4f}")
# → 0.21 (유사하지 않음, 주제가 완전히 다름)

# ---- TXT 읽기 ----
def load_txt(path):
    with open(path, "r", encoding="utf-8") as f:
        return f.read()

# ---- PDF 읽기 ----
from pypdf import PdfReader

def load_pdf(path):
    reader = PdfReader(path)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n"
    return text

# ---- Word 읽기 ----
from docx import Document

def load_docx(path):
    doc = Document(path)
    return "\n".join([p.text for p in doc.paragraphs])

# ---- 통합 로딩 인터페이스 ----
def load_document(path):
    if path.endswith(".pdf"):
        return load_pdf(path)
    elif path.endswith(".docx"):
        return load_docx(path)
    elif path.endswith(".txt") or path.endswith(".md"):
        return load_txt(path)
    else:
        raise ValueError(f"지원하지 않는 파일 형식: {path}")

# ==== 방법 1: 고정 길이 + 중복 (가장 일반적) ====
def chunk_by_size(text, chunk_size=500, overlap=100):
    """문자 수로 분할, 인접 블록에 중복 영역"""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap   # 중복 부분으로 컨텍스트 연속성 유지
    return chunks

# ==== 방법 2: 단락별 분할 ====
def chunk_by_paragraph(text, max_size=800):
    """단락별로 분할, 짧은 단락은 병합, 긴 단락은 분리"""
    paragraphs = text.split("\n\n")
    chunks = []
    current = ""
    for para in paragraphs:
        if len(current) + len(para) < max_size:
            current += para + "\n\n"
        else:
            if current:
                chunks.append(current.strip())
            current = para + "\n\n"
    if current.strip():
        chunks.append(current.strip())
    return chunks

# ==== 방법 3: LangChain Splitter 사용 (추천) ====
from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,        # 블록당 최대 문자 수
    chunk_overlap=100,     # 중복 문자 수
    separators=["\n\n", "\n", "。", "！", "？", "，", " "]
    # 단락 우선, 그 다음 문장, 마지막으로 문자
)

text = load_document("product_manual.pdf")
chunks = splitter.split_text(text)
print(f"총 {len(chunks)}개 블록으로 분할")
print(f"평균 길이: {sum(len(c) for c in chunks) / len(chunks):.0f} 문자")

def create_documents(file_path, chunks):
    """각 chunk에 메타데이터 추가"""
    documents = []
    for i, chunk in enumerate(chunks):
        doc = {
            "id": f"{file_path}_chunk_{i}",
            "text": chunk,
            "metadata": {
                "source": file_path,
                "chunk_index": i,
                "total_chunks": len(chunks),
                "char_count": len(chunk),
            }
        }
        documents.append(doc)
    return documents

import chromadb
from openai import OpenAI

client = OpenAI()

# ========== 1. 벡터 데이터베이스 생성/연결 ==========
chroma_client = chromadb.PersistentClient(path="./my_vectordb")
collection = chroma_client.get_or_create_collection(
    name="knowledge_base",
    metadata={"description": "기업 지식베이스"}
)

# ========== 2. 벡터화 함수 ==========
def get_embeddings(texts):
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    return [d.embedding for d in response.data]

# ========== 3. 문서 저장 ==========
documents = [
    "Python은 인터프리터 방식의 객체지향 고급 프로그래밍 언어입니다.",
    "Pandas는 Python에서 가장 인기 있는 데이터 처리 라이브러리입니다.",
    "RAG는 검색 증강 생성의 약자로, AI가 외부 지식에 접근하는 데 사용됩니다.",
    "벡터 데이터베이스는 텍스트의 수학적 표현을 저장하여 시맨틱 검색을 지원합니다.",
    "ChromaDB는 경량 오픈소스 벡터 데이터베이스입니다.",
]
ids = [f"doc_{i}" for i in range(len(documents))]
metadatas = [{"source": "tutorial", "index": i} for i in range(len(documents))]

embeddings = get_embeddings(documents)

collection.add(
    ids=ids,
    documents=documents,
    embeddings=embeddings,
    metadatas=metadatas
)
print(f"{collection.count()}개 문서 저장 완료")

# ========== 4. 시맨틱 검색 ==========
def search(query, top_k=3):
    query_embedding = get_embeddings([query])[0]
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )
    return results

# 검색 테스트
results = search("데이터를 어떻게 처리하나요?")
for i, doc in enumerate(results["documents"][0]):
    dist = results["distances"][0][i]
    print(f"[유사도: {1-dist:.3f}] {doc}")
# → [유사도: 0.856] Pandas는 Python에서 가장 인기 있는 데이터 처리 라이브러리입니다.

import os, glob

def index_documents(folder_path, collection, chunk_size=500, overlap=100):
    """폴더 내 모든 문서의 인덱스 생성"""
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size, chunk_overlap=overlap,
        separators=["\n\n", "\n", "。", "！", "？", " "]
    )

    all_chunks, all_ids, all_metas = [], [], []
    doc_count = 0

    for filepath in glob.glob(os.path.join(folder_path, "*")):
        try:
            text = load_document(filepath)
            chunks = splitter.split_text(text)
            filename = os.path.basename(filepath)

            for i, chunk in enumerate(chunks):
                all_chunks.append(chunk)
                all_ids.append(f"{filename}_{i}")
                all_metas.append({
                    "source": filename,
                    "chunk_index": i
                })

            doc_count += 1
            print(f"✅ {filename}: {len(chunks)}개 블록")
        except Exception as e:
            print(f"❌ {filepath}: {e}")

    # 배치 벡터화 및 저장
    if all_chunks:
        embeddings = get_embeddings(all_chunks)
        collection.add(
            ids=all_ids,
            documents=all_chunks,
            embeddings=embeddings,
            metadatas=all_metas
        )

    print(f"\n완료! 총 {doc_count}개 파일, {len(all_chunks)}개 블록 처리")

# 사용
index_documents("./docs", collection)

from openai import OpenAI
import chromadb

client = OpenAI()
chroma = chromadb.PersistentClient(path="./my_vectordb")
collection = chroma.get_collection("knowledge_base")

RAG_PROMPT = """당신은 스마트 Q&A 어시스턴트입니다. 아래 참고 자료를 기반으로 사용자의 질문에 답변하세요.

규칙:
- 참고 자료에 있는 정보만으로 답변하고, 정보를 지어내지 마세요
- 참고 자료에 관련 정보가 없으면 "현재 자료로는 이 질문에 답변할 수 없습니다"라고 말하세요
- 답변 끝에 정보 출처를 표시하세요

참고 자료:
---
{context}
---

사용자 질문: {question}"""

def rag_query(question, top_k=3):
    """완전한 RAG Q&A 흐름"""

    # Step 1: 관련 문서 검색
    query_emb = get_embeddings([question])[0]
    results = collection.query(
        query_embeddings=[query_emb],
        n_results=top_k,
        include=["documents", "metadatas", "distances"]
    )

    # Step 2: 컨텍스트 조립
    context_parts = []
    sources = []
    for i, doc in enumerate(results["documents"][0]):
        source = results["metadatas"][0][i].get("source", "알 수 없음")
        context_parts.append(f"[출처: {source}]\n{doc}")
        sources.append(source)

    context = "\n\n".join(context_parts)

    # Step 3: 대규모 모델로 답변 생성
    prompt = RAG_PROMPT.format(context=context, question=question)
    response = client.chat.completions.create(
        model="gpt-4o",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.3    # RAG 시나리오는 낮은 temperature 권장, 더 충실
    )
    answer = response.choices[0].message.content

    return {
        "answer": answer,
        "sources": list(set(sources)),
        "context_used": context_parts
    }

# ---- 사용 ----
result = rag_query("RAG란 무엇인가요?")
print(f"답변: {result['answer']}")
print(f"출처: {result['sources']}")

# ==== 최적화 1: Query 재작성 ====
# 사용자 질문이 구어체인 경우, AI에게 먼저 검색에 적합한 형태로 재작성하게 함
def rewrite_query(original_query):
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{
            "role": "user",
            "content": f"""다음 사용자 질문을 지식베이스 검색에 더 적합한 형태로 재작성하세요.
다른 각도의 검색 쿼리 3개를 출력하세요, 줄당 하나씩.

사용자 질문: {original_query}"""
        }],
        temperature=0.3
    )
    queries = response.choices[0].message.content.strip().split("\n")
    return [q.strip() for q in queries if q.strip()]

# 예: "이거 어떻게 반품해요" → ["반품 절차", "반품/교환 정책", "반품 신청 방법"]

# ==== 최적화 2: 하이브리드 검색 ====
# 시맨틱 검색 + 키워드 검색을 동시 사용, 합집합
def hybrid_search(query, collection, top_k=5):
    # 시맨틱 검색
    emb = get_embeddings([query])[0]
    semantic_results = collection.query(
        query_embeddings=[emb], n_results=top_k
    )

    # 키워드 검색 (Chroma는 where_document 필터 지원)
    keyword_results = collection.query(
        query_texts=[query], n_results=top_k  # 내장 키워드 매칭
    )

    # 병합 후 중복 제거
    seen = set()
    combined = []
    for source in [semantic_results, keyword_results]:
        for doc, doc_id in zip(source["documents"][0], source["ids"][0]):
            if doc_id not in seen:
                seen.add(doc_id)
                combined.append(doc)
    return combined[:top_k]

# ==== 최적화 3: Rerank 재정렬 ====
# 10개를 검색한 후, AI에게 가장 관련 있는 3개를 선택하게 함
def rerank(query, documents, top_k=3):
    doc_list = "\n".join([f"[{i}] {d[:200]}" for i, d in enumerate(documents)])
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": f"""다음 문서 중 질문과 가장 관련 있는 것은?
가장 관련 있는 {top_k}개 문서 번호를 쉼표로 구분하여 반환하세요.

질문: {query}

문서 목록:
{doc_list}"""}],
        temperature=0
    )
    indices = [int(x.strip()) for x in resp.choices[0].message.content.split(",")]
    return [documents[i] for i in indices if i < len(documents)]

# 환각 방지 Prompt 템플릿
ANTI_HALLUCINATION_PROMPT = """다음 참고 자료를 기반으로 질문에 답변하세요.

중요 규칙:
1. 참고 자료의 정보만 사용하여 답변
2. 자료에 답이 없으면 명확히 "자료에 언급되지 않음"이라고 표시
3. 원문의 핵심 문장을 직접 인용, 「」로 표시
4. 신뢰도 제시 (높음/중간/낮음)

참고 자료:
{context}

질문: {question}

다음 형식으로 출력하세요:
답변: ...
인용: 「원문 문장」
신뢰도: 높음/중간/낮음
출처: 파일명"""

# 유사도 임계값 필터링 — 관련성이 너무 낮은 결과는 직접 폐기
def filtered_search(query, collection, top_k=5, threshold=0.3):
    emb = get_embeddings([query])[0]
    results = collection.query(
        query_embeddings=[emb], n_results=top_k,
        include=["documents", "distances", "metadatas"]
    )
    # 임계값 이하의 결과만 유지
    filtered = []
    for doc, dist, meta in zip(
        results["documents"][0],
        results["distances"][0],
        results["metadatas"][0]
    ):
        if dist < threshold:
            filtered.append({"doc": doc, "distance": dist, "meta": meta})
    return filtered

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.chains import RetrievalQA

# ---- 1. 문서 로딩 ----
loader = DirectoryLoader("./docs", glob="**/*.pdf", loader_cls=PyPDFLoader)
documents = loader.load()
print(f"{len(documents)}페이지 문서 로드 완료")

# ---- 2. 분할 ----
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
chunks = splitter.split_documents(documents)
print(f"{len(chunks)}개 블록으로 분할")

# ---- 3. 벡터화 + 저장 (한 줄로 완료!) ----
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectordb = Chroma.from_documents(chunks, embeddings, persist_directory="./chroma_db")

# ---- 4. RAG Q&A 체인 생성 ----
llm = ChatOpenAI(model="gpt-4o", temperature=0.3)
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=vectordb.as_retriever(search_kwargs={"k": 3}),
    return_source_documents=True
)

# ---- 5. 질문 ----
result = qa_chain.invoke({"query": "RAG란 무엇인가요?"})
print(f"답변: {result['result']}")
for doc in result["source_documents"]:
    print(f"  출처: {doc.metadata['source']} (page {doc.metadata.get('page', '?')})")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# ---- 이 세 줄이면 끝! ----
documents = SimpleDirectoryReader("./docs").load_data()
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

# ---- 질문 ----
response = query_engine.query("RAG란 무엇인가요?")
print(response)

import os, json, glob
from openai import OpenAI
import chromadb
from langchain.text_splitter import RecursiveCharacterTextSplitter
from pypdf import PdfReader
from dotenv import load_dotenv

load_dotenv()
client = OpenAI()

class KnowledgeBaseQA:
    """기업 지식베이스 Q&A 시스템"""

    def __init__(self, db_path="./knowledge_db"):
        self.chroma = chromadb.PersistentClient(path=db_path)
        self.collection = self.chroma.get_or_create_collection("enterprise_kb")
        self.splitter = RecursiveCharacterTextSplitter(
            chunk_size=500, chunk_overlap=100,
            separators=["\n\n", "\n", "。", "！", "？", " "]
        )
        self.chat_history = []

    # ========== 문서 관리 ==========
    def add_document(self, file_path):
        """지식베이스에 단일 문서 추가"""
        text = self._load_file(file_path)
        chunks = self.splitter.split_text(text)
        filename = os.path.basename(file_path)

        ids = [f"{filename}_{i}" for i in range(len(chunks))]
        metas = [{"source": filename, "chunk_idx": i} for i in range(len(chunks))]
        embs = self._embed(chunks)

        self.collection.add(ids=ids, documents=chunks,
                           embeddings=embs, metadatas=metas)
        print(f"✅ {filename}: {len(chunks)}개 블록 저장 완료")

    def add_folder(self, folder_path):
        """폴더 내 문서 일괄 추가"""
        for f in glob.glob(os.path.join(folder_path, "*")):
            try:
                self.add_document(f)
            except Exception as e:
                print(f"❌ {f}: {e}")

    # ========== 스마트 Q&A ==========
    def ask(self, question, top_k=3):
        """RAG Q&A"""
        # 1. 검색
        emb = self._embed([question])[0]
        results = self.collection.query(
            query_embeddings=[emb], n_results=top_k,
            include=["documents", "metadatas", "distances"]
        )

        # 2. 저품질 결과 필터링
        context_parts, sources = [], []
        for doc, dist, meta in zip(
            results["documents"][0],
            results["distances"][0],
            results["metadatas"][0]
        ):
            if dist < 0.5:
                context_parts.append(doc)
                sources.append(meta["source"])

        if not context_parts:
            return {"answer": "죄송합니다, 지식베이스에서 관련 정보를 찾지 못했습니다.", "sources": []}

        # 3. 답변 생성
        context = "\n---\n".join(context_parts)
        prompt = f"""다음 참고 자료를 기반으로 질문에 답변하세요. 자료만 기반으로 답변하고, 답변할 수 없으면 설명하세요.

참고 자료:
{context}

질문: {question}"""

        response = client.chat.completions.create(
            model="gpt-4o",
            messages=[
                {"role": "system", "content": "당신은 기업 지식베이스 어시스턴트로, 제공된 자료를 기반으로 정확하게 답변합니다."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2
        )

        return {
            "answer": response.choices[0].message.content,
            "sources": list(set(sources)),
            "tokens": response.usage.total_tokens
        }

    def stats(self):
        """지식베이스 통계 조회"""
        return {"total_chunks": self.collection.count()}

    # ========== 내부 메서드 ==========
    def _embed(self, texts):
        resp = client.embeddings.create(
            model="text-embedding-3-small", input=texts)
        return [d.embedding for d in resp.data]

    def _load_file(self, path):
        if path.endswith(".pdf"):
            reader = PdfReader(path)
            return "\n".join(p.extract_text() for p in reader.pages)
        else:
            with open(path, "r", encoding="utf-8") as f:
                return f.read()

# ========== 메인 프로그램 ==========
if __name__ == "__main__":
    kb = KnowledgeBaseQA()

    # 문서 추가
    kb.add_folder("./docs")
    print(f"\n지식베이스 상태: {kb.stats()}")

    # 대화형 Q&A
    print("\n기업 지식베이스 Q&A 시스템이 시작되었습니다! quit를 입력하면 종료합니다.\n")
    while True:
        q = input("📌 질문: ")
        if q.lower() == "quit": break
        result = kb.ask(q)
        print(f"\n💬 답변: {result['answer']}")
        print(f"📎 출처: {', '.join(result['sources'])}")
        print(f"💰 Token: {result.get('tokens', '?')}\n")