core/embedder.py

# core/embedder.py
"""임베딩 생성"""
from openai import OpenAI
from typing import List, Dict
from config.settings import OPENAI_API_KEY, EMBEDDING_MODEL
import time

client = OpenAI(api_key=OPENAI_API_KEY)


def embed_chunks(chunks: List[Dict]) -> List[Dict]:
    """
    청크 리스트를 임베딩
    
    Args:
        chunks: 청크 리스트
    
    Returns:
        List[Dict]: 임베딩이 추가된 청크
            [
                {
                    "chunk_id": "...",
                    "text": "...",
                    "embedding": [0.1, 0.2, ...],
                    ...
                },
                ...
            ]
    """
    print(f"🔢 임베딩 시작 ({len(chunks)}개 청크)")
    start_time = time.time()
    
    # 배치 처리 (OpenAI는 한 번에 여러 개 가능)
    texts = [chunk["text"] for chunk in chunks]
    
    response = client.embeddings.create(
        model=EMBEDDING_MODEL,
        input=texts
    )
    
    # 임베딩 추가
    for i, chunk in enumerate(chunks):
        chunk["embedding"] = response.data[i].embedding
    
    elapsed = time.time() - start_time
    print(f"✅ 임베딩 완료 ({elapsed:.2f}초)")
    print(f"   - 속도: {len(chunks)/elapsed:.2f} chunks/sec")
    
    return chunks


if __name__ == "__main__":
    # 테스트
    test_chunks = [
        {"chunk_id": "chunk_0", "text": "테스트 문서입니다."}
    ]
    
    embedded = embed_chunks(test_chunks)
    print(f"임베딩 차원: {len(embedded[0]['embedding'])}")