Milvus

# pip install pymilvus
from pymilvus import MilvusClient, DataType

# 1. Milvus 연결 (Milvus Lite - 로컬 개발용)
client = MilvusClient("./milvus_demo.db")

# 2. 컬렉션 생성 (스키마 정의)
# 방법 1: 자동 스키마 (간편)
client.create_collection(
    collection_name="documents",
    dimension=1536,  # OpenAI embedding 차원
    metric_type="COSINE"  # 유사도 메트릭: COSINE, L2, IP
)

# 방법 2: 상세 스키마 (프로덕션 권장)
from pymilvus import CollectionSchema, FieldSchema

fields = [
    FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
    FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
    FieldSchema(name="content", dtype=DataType.VARCHAR, max_length=65535),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=1536),
]
schema = CollectionSchema(fields, description="문서 검색용 컬렉션")

# 3. 데이터 삽입
import openai

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

documents = [
    {"title": "RAG 소개", "content": "RAG는 검색 증강 생성 기법입니다..."},
    {"title": "벡터 DB란", "content": "벡터 데이터베이스는 임베딩을 저장..."},
    {"title": "Milvus 가이드", "content": "Milvus는 대규모 벡터 검색에 특화..."},
]

data = []
for doc in documents:
    data.append({
        "title": doc["title"],
        "content": doc["content"],
        "embedding": get_embedding(doc["content"])
    })

client.insert(collection_name="documents", data=data)
print(f"✅ {len(data)}개 문서 삽입 완료")

# 4. 벡터 유사도 검색
query = "벡터 데이터베이스의 활용 방법"
query_vector = get_embedding(query)

results = client.search(
    collection_name="documents",
    data=[query_vector],
    limit=3,  # Top-K
    output_fields=["title", "content"],
    search_params={"metric_type": "COSINE", "params": {"nprobe": 10}}
)

print("\n🔍 검색 결과:")
for hits in results:
    for hit in hits:
        print(f"- [{hit['distance']:.4f}] {hit['entity']['title']}")

# 5. 하이브리드 검색 (벡터 + 필터)
results = client.search(
    collection_name="documents",
    data=[query_vector],
    filter='title like "%Milvus%"',  # 메타데이터 필터
    limit=3,
    output_fields=["title", "content"]
)

# 6. 인덱스 생성 (대규모 데이터용)
index_params = {
    "index_type": "HNSW",  # 고성능 그래프 기반 인덱스
    "metric_type": "COSINE",
    "params": {"M": 16, "efConstruction": 256}
}
client.create_index(
    collection_name="documents",
    field_name="embedding",
    index_params=index_params
)
print("✅ HNSW 인덱스 생성 완료")