# Vector Database 활용 예제 - Pinecone + OpenAI
# pip install pinecone-client openai
import os
from typing import List, Dict, Any
import hashlib
# Pinecone 클라이언트 (v3.0+)
from pinecone import Pinecone, ServerlessSpec
# OpenAI 임베딩
from openai import OpenAI
class VectorDBManager:
"""벡터 데이터베이스 관리 클래스"""
def __init__(self, pinecone_api_key: str, openai_api_key: str):
# 클라이언트 초기화
self.pc = Pinecone(api_key=pinecone_api_key)
self.openai = OpenAI(api_key=openai_api_key)
self.embedding_model = "text-embedding-3-small"
self.dimension = 1536 # text-embedding-3-small 차원
def create_index(self, index_name: str):
"""인덱스 생성"""
# 기존 인덱스 확인
existing_indexes = [idx.name for idx in self.pc.list_indexes()]
if index_name not in existing_indexes:
self.pc.create_index(
name=index_name,
dimension=self.dimension,
metric="cosine", # 코사인 유사도
spec=ServerlessSpec(
cloud="aws",
region="us-east-1"
)
)
print(f"Index '{index_name}' created successfully")
else:
print(f"Index '{index_name}' already exists")
return self.pc.Index(index_name)
def get_embedding(self, text: str) -> List[float]:
"""텍스트를 임베딩 벡터로 변환"""
response = self.openai.embeddings.create(
input=text,
model=self.embedding_model
)
return response.data[0].embedding
def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
"""배치 임베딩 생성"""
response = self.openai.embeddings.create(
input=texts,
model=self.embedding_model
)
return [item.embedding for item in response.data]
def upsert_documents(self, index, documents: List[Dict[str, Any]]):
"""
문서를 벡터 DB에 삽입/업데이트
Args:
index: Pinecone 인덱스
documents: [{"id": str, "text": str, "metadata": dict}, ...]
"""
# 임베딩 생성
texts = [doc["text"] for doc in documents]
embeddings = self.get_embeddings_batch(texts)
# 벡터 준비
vectors = []
for doc, embedding in zip(documents, embeddings):
# ID 생성 (없으면 텍스트 해시)
doc_id = doc.get("id") or hashlib.md5(doc["text"].encode()).hexdigest()
vectors.append({
"id": doc_id,
"values": embedding,
"metadata": {
"text": doc["text"][:1000], # 메타데이터 크기 제한
**doc.get("metadata", {})
}
})
# 배치 업서트 (1000개씩)
batch_size = 100
for i in range(0, len(vectors), batch_size):
batch = vectors[i:i + batch_size]
index.upsert(vectors=batch)
print(f"Upserted {len(vectors)} documents")
def search(self, index, query: str, top_k: int = 5,
filter: Dict = None) -> List[Dict]:
"""
유사도 검색
Args:
index: Pinecone 인덱스
query: 검색 쿼리
top_k: 반환할 결과 수
filter: 메타데이터 필터
Returns:
검색 결과 리스트
"""
# 쿼리 임베딩 생성
query_embedding = self.get_embedding(query)
# 검색 실행
results = index.query(
vector=query_embedding,
top_k=top_k,
include_metadata=True,
filter=filter
)
# 결과 정리
formatted_results = []
for match in results.matches:
formatted_results.append({
"id": match.id,
"score": match.score,
"text": match.metadata.get("text", ""),
"metadata": match.metadata
})
return formatted_results
def hybrid_search(self, index, query: str, keyword_filter: str = None,
top_k: int = 5) -> List[Dict]:
"""
하이브리드 검색 (벡터 + 키워드)
"""
# 메타데이터 필터 구성
filter_dict = None
if keyword_filter:
filter_dict = {
"text": {"$contains": keyword_filter}
}
return self.search(index, query, top_k, filter_dict)
class RAGPipeline:
"""RAG (Retrieval-Augmented Generation) 파이프라인"""
def __init__(self, vector_manager: VectorDBManager, index):
self.vector_manager = vector_manager
self.index = index
def add_documents(self, texts: List[str], metadata_list: List[Dict] = None):
"""문서 추가"""
documents = []
for i, text in enumerate(texts):
doc = {
"text": text,
"metadata": metadata_list[i] if metadata_list else {}
}
documents.append(doc)
self.vector_manager.upsert_documents(self.index, documents)
def retrieve_context(self, query: str, top_k: int = 3) -> str:
"""쿼리와 관련된 컨텍스트 검색"""
results = self.vector_manager.search(self.index, query, top_k)
context_parts = []
for i, result in enumerate(results):
context_parts.append(f"[{i+1}] {result['text']}")
return "\n\n".join(context_parts)
def generate_answer(self, query: str, top_k: int = 3) -> str:
"""RAG 기반 답변 생성"""
# 1. 관련 문서 검색
context = self.retrieve_context(query, top_k)
# 2. 프롬프트 구성
prompt = f"""다음 컨텍스트를 바탕으로 질문에 답변해주세요.
컨텍스트:
{context}
질문: {query}
답변:"""
# 3. LLM으로 답변 생성
response = self.vector_manager.openai.chat.completions.create(
model="gpt-4o-mini",
messages=[{"role": "user", "content": prompt}],
temperature=0.7
)
return response.choices[0].message.content
# 사용 예제
if __name__ == "__main__":
# 환경 변수에서 API 키 로드
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY", "your-api-key")
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY", "your-api-key")
# 매니저 초기화
manager = VectorDBManager(PINECONE_API_KEY, OPENAI_API_KEY)
# 인덱스 생성
index = manager.create_index("knowledge-base")
# 문서 추가
documents = [
{
"text": "벡터 데이터베이스는 고차원 벡터를 저장하고 유사도 검색을 수행합니다.",
"metadata": {"category": "database", "lang": "ko"}
},
{
"text": "RAG는 검색과 생성을 결합하여 LLM의 답변 품질을 향상시킵니다.",
"metadata": {"category": "ai", "lang": "ko"}
},
{
"text": "Pinecone은 완전 관리형 벡터 데이터베이스 서비스입니다.",
"metadata": {"category": "database", "lang": "ko"}
}
]
manager.upsert_documents(index, documents)
# 검색 테스트
print("\n=== 유사도 검색 ===")
results = manager.search(index, "벡터 DB란 무엇인가요?", top_k=2)
for r in results:
print(f"Score: {r['score']:.4f} | {r['text'][:50]}...")
# RAG 파이프라인 테스트
print("\n=== RAG 답변 생성 ===")
rag = RAGPipeline(manager, index)
answer = rag.generate_answer("RAG가 뭔가요?")
print(f"Answer: {answer}")