[side-1] Milvus Collection에 적재된 데이터에서 유사한 값 검색하기

개요

지난번 글에서 vector_store 부분 수정해서 검색결과가 나오지 않았는데, 이 부분이 잘 되지 않아서,

우선 Milvus Client로 통신하여 구축해보았다.

 

환경

Docker Compose로 모두 컨테이너화하였으며, 사용자의 입력값을 받기 위한 FastAPI 컨테이너에서 Milvus 컨테이너 연결해 데이터 조회.

 

작업

임베딩은 OpenAI 모델 사용해서 진행할 예정이고, Milvus Client 사용해 통신하기 위해 모듈을 가져온다.

from openai import OpenAI
from langchain_openai import OpenAIEmbeddings
from pymilvus import connections, MilvusClient

 

.env에 저장해둔 환경 변수 가져오기

from dotenv import load_dotenv
import os

load_dotenv()

# 환경 변수 가져오기
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
MILVUS_HOST = os.getenv("MILVUS_HOST")
MILVUS_PORT = os.getenv("MILVUS_PORT")

 

Milvus 연결 확인하기

try:
    connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT)
    print("Milvus에 연결 성공")
except Exception as e:
    print(f"Milvus 연결 실패: {e}")
    exit(1)

 

임베딩은 text embedding large 사용하였다.

embeddings = OpenAIEmbeddings(model='text-embedding-3-large')

 

사용자 입력값으로 queries에 검색할 내용을 넣어주고, 해당 쿼리 역시 임베딩한다.

여기서 dimensions를 설정해주지 않은 경우 오류가 발생하기 때문에 지정해주었다.

데이터 적재 시, 차원을 384로 저장했었기 때문에 맞춰주었다.

queries = ["크로스백"] # 검색할 분실물 항목
openai_client = OpenAI(api_key=OPENAI_API_KEY)
milvus_client = MilvusClient(uri="http://{컨테이너 서비스명}:{포트}")
query_vectors = [
    vec.embedding
    for vec in openai_client.embeddings.create(input=queries, model='text-embedding-3-large',dimensions=384).data
]

 

이제 Milvus에 저장된 값 중 유사한 값을 뽑아본다.

3개만 출력해보았다.

그리고 output_fields로는 저장한 필드값 다 보려고 다 설정해주었다.

보고자 하는 항목만 기재해주어도 된다.

res = milvus_client.search(
    collection_name=COLLECTION_NAME,  # target collection
    data=query_vectors,
    limit=3,
    output_fields=["item_id", "name", "description"],
)

for q in queries:
    print("Query:", q)
    for result in res:
        print(result)
    print("\n")

 

결과

Milvus에 성공적으로 연결되었습니다.
Query: 크로스백
[{'id': 'F2024121000001173-1', 'distance': 0.4028303027153015, 'entity': {'item_id': 'F2024121000001173-1', 'name': '크로스백', 'description': '내용 거제경찰서에서는 [2024.12.01]  [크로스백(블랙(검정)색)]을 습득/보관 하였습니다. 분실하신 분께서는 본인을 증명할 수 있는 서류를 지참하시어 보관중으로 기재되어 있는 기관에 방문하시어 보관물품을 수령하시기 바랍니다. 특이사항 : 없음'}}, 
{'id': 'F2024121000002011-1', 'distance': 0.18288661539554596, 'entity': {'item_id': 'F2024121000002011-1', 'name': '검정색 핸드백(루이비통)', 'description': '동대구역(한국철도공사)에서는 [24.12.01]  [검정색 핸드백(루이비통)(블랙(검정)색)]을 습득/보관 하였습니다. 분실하신 분께서는 본인을 증명할 수 있는 서류를 지참하시어 보관중으로 기재되어 있는 기관에 방문하시어 보관물품을 수령하시기 바랍니다. 특이사항 : 없음'}}, 
{'id': 'F2024121000001484-1', 'distance': 0.1518131047487259, 'entity': {'item_id': 'F2024121000001484-1', 'name': '여성용 가방', 'description': '수영역(부산교통공사)에서는 [24.12.01]  [여성용 가방(블랙(검정)색)]을 습득/보관 하였습니다. 분실하신 분께서는 본인을 증명할 수 있는 서류를 지참하시어 보관중으로 기재되어 있는 기관에 방문하시어 보관물품을 수령하시기 바랍니다. 특이사항 : 없음'}}]

 

쿼리문을 단순히 크로스백으로만 넣어서 그렇긴 한데, 크로스백 제대로 잘 찾아왔다!

 

이후에는 langchain_milvus 사용해서 통신 테스트하고 웹에서 사용자 쿼리문 받아서 보여줄 수 있도록 작업하면 된다.

 

최종 코드

from langchain_openai import OpenAIEmbeddings
from openai import OpenAI
from pymilvus import connections, MilvusClient
from dotenv import load_dotenv
import os

load_dotenv()

# 환경 변수 가져오기
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
MILVUS_HOST = os.getenv("MILVUS_HOST")
MILVUS_PORT = os.getenv("MILVUS_PORT")
COLLECTION_NAME = "lost_items"

# Milvus 연결 테스트
try:
    connections.connect("default", host=MILVUS_HOST, port=MILVUS_PORT)
    print("Milvus에 연결 성공.")
except Exception as e:
    print(f"Milvus 연결 실패: {e}")
    exit(1)

embeddings = OpenAIEmbeddings(model='text-embedding-3-large')

queries = ["크로스백"]
openai_client = OpenAI(api_key=OPENAI_API_KEY)
milvus_client = MilvusClient(uri="http://{컨테이너 서비스명}:{포트}")
query_vectors = [
    vec.embedding
    for vec in openai_client.embeddings.create(input=queries, model='text-embedding-3-large',dimensions=384).data
]

res = milvus_client.search(
    collection_name=COLLECTION_NAME,
    data=query_vectors,
    limit=3,
    output_fields=["item_id", "name", "description"],
)

for q in queries:
    print("Query:", q)
    for result in res:
        print(result)
    print("\n")