Pinecone

from pinecone import Pinecone, ServerlessSpec
from openai import OpenAI

# 초기화
pc = Pinecone(api_key="your-pinecone-api-key")
openai_client = OpenAI(api_key="your-openai-api-key")

# 인덱스 생성 (최초 1회)
index_name = "my-rag-index"
if index_name not in pc.list_indexes().names():
    pc.create_index(
        name=index_name,
        dimension=1536,  # text-embedding-3-small 차원
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1")
    )

# 인덱스 연결
index = pc.Index(index_name)

# 임베딩 생성 함수
def get_embedding(text):
    response = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return response.data[0].embedding

# 문서 업서트 (삽입/업데이트)
documents = [
    {"id": "doc1", "text": "Pinecone은 벡터 데이터베이스입니다.", "category": "database"},
    {"id": "doc2", "text": "RAG는 검색 증강 생성을 의미합니다.", "category": "llm"},
    {"id": "doc3", "text": "임베딩은 텍스트를 벡터로 변환합니다.", "category": "ml"}
]

vectors = []
for doc in documents:
    embedding = get_embedding(doc["text"])
    vectors.append({
        "id": doc["id"],
        "values": embedding,
        "metadata": {"text": doc["text"], "category": doc["category"]}
    })

index.upsert(vectors=vectors, namespace="documents")

# 유사도 검색
query = "벡터 DB란 무엇인가요?"
query_embedding = get_embedding(query)

results = index.query(
    vector=query_embedding,
    top_k=3,
    include_metadata=True,
    namespace="documents"
)

for match in results.matches:
    print(f"Score: {match.score:.4f} | {match.metadata['text']}")

# 메타데이터 필터링 검색
filtered_results = index.query(
    vector=query_embedding,
    top_k=3,
    include_metadata=True,
    filter={"category": {"$eq": "database"}},
    namespace="documents"
)

# 인덱스 통계 확인
stats = index.describe_index_stats()
print(f"총 벡터 수: {stats.total_vector_count}")