Chroma

# pip install chromadb openai
import chromadb
from chromadb.utils import embedding_functions

# Chroma 클라이언트 생성 (영구 저장소 사용)
client = chromadb.PersistentClient(path="./chroma_db")

# OpenAI 임베딩 함수 설정
openai_ef = embedding_functions.OpenAIEmbeddingFunction(
    api_key="your-openai-api-key",
    model_name="text-embedding-3-small"  # 또는 text-embedding-3-large
)

# 컬렉션 생성 또는 가져오기
collection = client.get_or_create_collection(
    name="documents",
    embedding_function=openai_ef,
    metadata={"hnsw:space": "cosine"}  # 유사도 메트릭 설정
)

# 문서 추가
documents = [
    "인공지능은 인간의 학습능력과 추론능력을 컴퓨터로 구현하는 기술입니다.",
    "머신러닝은 데이터로부터 패턴을 학습하여 예측하는 AI의 하위 분야입니다.",
    "딥러닝은 인공 신경망을 사용하여 복잡한 패턴을 학습하는 기법입니다.",
    "자연어 처리(NLP)는 컴퓨터가 인간의 언어를 이해하고 생성하는 기술입니다.",
    "RAG는 검색과 생성을 결합하여 LLM의 응답 품질을 높이는 기법입니다."
]

# 메타데이터와 함께 문서 추가
collection.add(
    documents=documents,
    metadatas=[
        {"category": "AI", "level": "basic"},
        {"category": "ML", "level": "intermediate"},
        {"category": "DL", "level": "advanced"},
        {"category": "NLP", "level": "intermediate"},
        {"category": "RAG", "level": "advanced"}
    ],
    ids=[f"doc_{i}" for i in range(len(documents))]
)

print(f"컬렉션에 {collection.count()}개의 문서가 저장되었습니다.")

# 시맨틱 검색 수행
def semantic_search(query: str, n_results: int = 3, filter_dict: dict = None):
    """
    시맨틱 검색을 수행하는 함수

    Args:
        query: 검색 쿼리
        n_results: 반환할 결과 수
        filter_dict: 메타데이터 필터 (예: {"category": "AI"})

    Returns:
        검색 결과
    """
    results = collection.query(
        query_texts=[query],
        n_results=n_results,
        where=filter_dict,  # 메타데이터 필터링
        include=["documents", "metadatas", "distances"]
    )

    return results

# 검색 예시 1: 일반 시맨틱 검색
print("\n=== 검색: '신경망과 학습' ===")
results = semantic_search("신경망과 학습")
for doc, meta, dist in zip(
    results['documents'][0],
    results['metadatas'][0],
    results['distances'][0]
):
    print(f"[유사도: {1-dist:.4f}] {doc[:50]}... ({meta['category']})")

# 검색 예시 2: 메타데이터 필터링과 함께
print("\n=== 검색: 'AI 기술' (고급 수준만) ===")
results = semantic_search(
    "AI 기술",
    n_results=2,
    filter_dict={"level": "advanced"}
)
for doc, meta in zip(results['documents'][0], results['metadatas'][0]):
    print(f"[{meta['level']}] {doc[:50]}...")

# LangChain 통합 예시
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings

# LangChain용 Chroma 벡터스토어
vectorstore = Chroma(
    collection_name="langchain_docs",
    embedding_function=OpenAIEmbeddings(),
    persist_directory="./langchain_chroma_db"
)

# 리트리버로 사용
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 3}
)