하이브리드 검색

from pinecone import Pinecone, ServerlessSpec
from langchain_openai import OpenAIEmbeddings
from langchain_community.retrievers import PineconeHybridSearchRetriever
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pinecone_text.sparse import BM25Encoder
import tiktoken

# 1. Pinecone 하이브리드 검색 설정
def setup_hybrid_search():
    """Pinecone 하이브리드 검색 인덱스 생성"""
    pc = Pinecone()

    # 하이브리드 검색용 인덱스 (sparse + dense 지원)
    index_name = "hybrid-search-rag"

    if index_name not in pc.list_indexes().names():
        pc.create_index(
            name=index_name,
            dimension=1536,  # OpenAI ada-002
            metric="dotproduct",  # 하이브리드 검색에 필수
            spec=ServerlessSpec(
                cloud="aws",
                region="us-east-1"
            )
        )

    return pc.Index(index_name)

# 2. 문서 인덱싱 (Dense + Sparse 벡터)
def index_documents(documents: list[str], index):
    """문서를 벡터화하고 인덱싱"""
    # Dense 임베딩 (시맨틱)
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

    # Sparse 임베딩 (키워드 - BM25)
    bm25_encoder = BM25Encoder()
    bm25_encoder.fit(documents)  # 코퍼스에 맞게 학습
    bm25_encoder.dump("bm25_params.json")  # 저장

    # 텍스트 분할
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50
    )

    vectors = []
    for i, doc in enumerate(documents):
        chunks = text_splitter.split_text(doc)
        for j, chunk in enumerate(chunks):
            # Dense 벡터
            dense_vec = embeddings.embed_query(chunk)
            # Sparse 벡터 (BM25)
            sparse_vec = bm25_encoder.encode_documents([chunk])[0]

            vectors.append({
                "id": f"doc_{i}_chunk_{j}",
                "values": dense_vec,
                "sparse_values": sparse_vec,
                "metadata": {"text": chunk, "doc_id": i}
            })

    # 배치 업서트
    index.upsert(vectors=vectors, namespace="documents")
    return bm25_encoder

# 3. 하이브리드 검색 수행
def hybrid_search(
    query: str,
    index,
    bm25_encoder: BM25Encoder,
    top_k: int = 5,
    alpha: float = 0.7  # dense 가중치 (1-alpha = sparse)
):
    """하이브리드 검색 실행"""
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

    # Dense 쿼리 벡터
    dense_query = embeddings.embed_query(query)
    # Sparse 쿼리 벡터
    sparse_query = bm25_encoder.encode_queries([query])[0]

    # 하이브리드 검색
    results = index.query(
        vector=dense_query,
        sparse_vector=sparse_query,
        top_k=top_k,
        include_metadata=True,
        namespace="documents"
    )

    return [(match.metadata["text"], match.score) for match in results.matches]

# 4. RRF (Reciprocal Rank Fusion) 직접 구현
def reciprocal_rank_fusion(
    vector_results: list,
    keyword_results: list,
    k: int = 60  # RRF 상수
) -> list:
    """두 검색 결과를 RRF로 결합"""
    fused_scores = {}

    # 벡터 검색 결과 스코어링
    for rank, (doc_id, _) in enumerate(vector_results):
        fused_scores[doc_id] = fused_scores.get(doc_id, 0) + 1 / (k + rank + 1)

    # 키워드 검색 결과 스코어링
    for rank, (doc_id, _) in enumerate(keyword_results):
        fused_scores[doc_id] = fused_scores.get(doc_id, 0) + 1 / (k + rank + 1)

    # 점수 기준 정렬
    sorted_results = sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
    return sorted_results

# 5. LangChain 하이브리드 Retriever 사용
def create_langchain_hybrid_retriever(index, bm25_encoder):
    """LangChain 하이브리드 검색 리트리버"""
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

    retriever = PineconeHybridSearchRetriever(
        embeddings=embeddings,
        sparse_encoder=bm25_encoder,
        index=index,
        top_k=5,
        alpha=0.7  # dense 비중
    )
    return retriever

# 사용 예시
if __name__ == "__main__":
    # 인덱스 설정
    index = setup_hybrid_search()

    # 샘플 문서 인덱싱
    documents = [
        "iPhone 15 Pro Max는 티타늄 프레임과 A17 Pro 칩을 탑재한 애플의 플래그십 스마트폰입니다.",
        "갤럭시 S24 울트라는 삼성의 최신 프리미엄 스마트폰으로 AI 기능이 강화되었습니다.",
        "스마트폰 가격은 브랜드, 스펙, 출시 시기에 따라 크게 달라집니다."
    ]
    bm25 = index_documents(documents, index)

    # 하이브리드 검색 테스트
    query = "아이폰 15 프로 맥스 가격"
    results = hybrid_search(query, index, bm25)
    print("하이브리드 검색 결과:")
    for text, score in results:
        print(f"  [{score:.3f}] {text[:50]}...")