본문 바로가기
WIKI 기술 지식 베이스

Milvus와 OpenAI Agents 통합: 단계별 가이드 (Milvus Integration with OpenAI Agents: A Step-by-Step Guide)

원문 보기 위키 갱신

이 노트북은 Function Calling을 통해 자연어로 Milvus를 쿼리할 수 있는 에이전트를 만드는 방법을 보여 줘요. OpenAI의 Agents 프레임워크와 Milvus의 강력한 벡터 검색 기능을 결합해 훌륭한 검색 경험을 만들어요.

출처: Milvus 문서

본문

OpenAI Agents

OpenAI Agents SDK는 매우 적은 추상화로 가볍고 사용하기 쉬운 패키지에서 에이전트형 AI 앱을 구축할 수 있게 해줘요. 이전의 에이전트용 실험인 Swarm의 프로덕션 준비 업그레이드예요. Agents SDK는 매우 작은 기본 요소 집합을 가져요.

  • Agents: 지시와 도구를 갖춘 LLM.
  • Handoffs: 에이전트가 특정 작업을 위해 다른 에이전트에게 위임할 수 있게 해줘요.
  • Guardrails: 에이전트에 대한 입력이 검증되도록 해줘요.

Python과 결합하면 이 기본 요소들은 도구와 에이전트 사이의 복잡한 관계를 표현하기에 충분히 강력하며, 가파른 학습 곡선 없이 실제 애플리케이션을 구축할 수 있게 해줘요. 또한 SDK에는 에이전트형 흐름을 시각화하고 디버깅할 수 있는 내장 추적(tracing)이 있어요. 이를 평가하고 애플리케이션에 맞게 모델을 미세 조정할 수도 있어요.

Milvus

Milvus는 노트북부터 대규모 분산 시스템까지 다양한 환경에서 효율적으로 실행되는 고성능, 고확장성 오픈소스 벡터 데이터베이스예요. 오픈소스 소프트웨어와 Cloud Offering으로 모두 제공돼요.

설정과 의존성 (Setup and Dependencies)

먼저 환경에 필요한 라이브러리를 설정하고 Jupyter 호환을 위해 asyncio를 초기화해야 해요.

$ pip install openai pymilvus pydantic nest_asyncio

Google Colab을 사용한다면 방금 설치한 의존성을 활성화하기 위해 런타임을 재시작해야 할 수 있어요(화면 상단의 "Runtime" 메뉴를 클릭하고 드롭다운에서 "Restart session"을 선택).

import asyncio
import nest_asyncio
from dotenv import load_dotenv

load_dotenv()

nest_asyncio.apply()

OpenAI의 모델을 사용할 거예요. api key OPENAI_API_KEY를 환경 변수로 준비해야 해요.

import os

os.environ["OPENAI_API_KEY"] = "sk-***********"

Milvus에 연결하고 스키마 만들기 (Connecting to Milvus and Creating a Schema)

이제 Milvus 인스턴스에 연결하고 컬렉션용 스키마를 만들어요. 이 스키마는 다음을 포함한 데이터 구조를 정의해요.

  • 기본 키인 ID 필드.
  • 문서 콘텐츠를 저장할 텍스트 필드.
  • BM25 임베딩을 저장할 스파스 벡터 필드.

Milvus 2.5의 전문 검색 (Full-Text Search in Milvus 2.5)

  • 벡터 검색과 키워드 검색을 위한 통합 시스템(통합 API).
  • 내장 스파스-BM25 알고리즘(Elasticsearch가 사용하는 방식과 유사하지만 벡터 기반).
  • 키워드 검색용 임베딩을 수동으로 생성할 필요 없음.

Docker로 Milvus 설치 (Install Milvus with Docker)

이 예시를 실행하기 전에 Milvus를 설치하고 Docker로 시작해야 해요. 문서를 확인하세요 - https://milvus.io/docs/install_standalone-docker.md

from pymilvus import DataType, FunctionType, MilvusClient

client = MilvusClient(uri="http://localhost:19530")

schema = client.create_schema()

# Simple schema that handles both text and vectors
schema.add_field(
    field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True
)
schema.add_field(
    field_name="text", datatype=DataType.VARCHAR, max_length=1000, enable_analyzer=True
)
schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)
{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>}], 'enable_dynamic_field': False}

Milvus는 BM25 함수를 통한 전문 검색을 지원해요. 여기서는 텍스트 데이터를 텍스트 검색에 최적화된 스파스 벡터 표현으로 자동 변환하는 함수를 정의해요.

from pymilvus import Function

# Milvus handles tokenization and BM25 conversion
bm25_function = Function(
    name="text_bm25_emb",  # Function name
    input_field_names=["text"],  # Name of the VARCHAR field containing raw text data
    output_field_names=[
        "sparse"
    ],  # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
    function_type=FunctionType.BM25,
)

schema.add_function(bm25_function)
{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>, 'is_function_output': True}], 'enable_dynamic_field': False, 'functions': [{'name': 'text_bm25_emb', 'description': '', 'type': <FunctionType.BM25: 1>, 'input_field_names': ['text'], 'output_field_names': ['sparse'], 'params': {}}]}

컬렉션 만들고 샘플 데이터 로드 (Creating the Collection and Loading Sample Data)

이제 스키마와 인덱스 파라미터로 컬렉션을 만든 뒤 정보 검색과 Milvus에 대한 몇 가지 샘플 데이터를 로드해요.

index_params = client.prepare_index_params()

index_params.add_index(field_name="sparse", index_type="AUTOINDEX", metric_type="BM25")

if client.has_collection("demo"):
    client.drop_collection("demo")

client.create_collection(
    collection_name="demo",
    schema=schema,
    index_params=index_params,
)

## 3. Loading Test Data
client.insert(
    "demo",
    [
        {
            "text": "Information retrieval helps users find relevant documents in large datasets."
        },
        {
            "text": "Search engines use information retrieval techniques to index and rank web pages."
        },
        {
            "text": "The core of IR is matching user queries with the most relevant content."
        },
        {
            "text": "Vector search is revolutionising modern information retrieval systems."
        },
        {
            "text": "Machine learning improves ranking algorithms in information retrieval."
        },
        {
            "text": "IR techniques include keyword-based search, semantic search, and vector search."
        },
        {
            "text": "Boolean retrieval is one of the earliest information retrieval methods."
        },
        {"text": "TF-IDF is a classic method used to score document relevance in IR."},
        {
            "text": "Modern IR systems integrate deep learning for better contextual understanding."
        },
        {
            "text": "Milvus is an open-source vector database designed for AI-powered search."
        },
        {
            "text": "Milvus enables fast and scalable similarity search on high-dimensional data."
        },
        {
            "text": "With Milvus, developers can build applications that support image, text, and video retrieval."
        },
        {
            "text": "Milvus integrates well with deep learning frameworks like PyTorch and TensorFlow."
        },
        {
            "text": "The core of Milvus is optimised for approximate nearest neighbour (ANN) search."
        },
        {
            "text": "Milvus supports hybrid search combining structured and unstructured data."
        },
        {
            "text": "Large-scale AI applications rely on Milvus for efficient vector retrieval."
        },
        {"text": "Milvus makes it easy to perform high-speed similarity searches."},
        {"text": "Cloud-native by design, Milvus scales effortlessly with demand."},
        {
            "text": "Milvus powers applications in recommendation systems, fraud detection, and genomics."
        },
        {
            "text": "The latest version of Milvus introduces faster indexing and lower latency."
        },
        {"text": "Milvus supports HNSW, IVF_FLAT, and other popular ANN algorithms."},
        {
            "text": "Vector embeddings from models like OpenAI's CLIP can be indexed in Milvus."
        },
        {
            "text": "Milvus has built-in support for multi-tenancy in enterprise use cases."
        },
        {
            "text": "The Milvus community actively contributes to improving its performance."
        },
        {
            "text": "Milvus integrates with data pipelines like Apache Kafka for real-time updates."
        },
        {
            "text": "Using Milvus, companies can enhance search experiences with vector search."
        },
        {
            "text": "Milvus plays a crucial role in powering AI search in medical research."
        },
        {"text": "Milvus integrates with LangChain for advanced RAG pipelines."},
        {
            "text": "Open-source contributors continue to enhance Milvus' search performance."
        },
        {
            "text": "Multi-modal search in Milvus enables applications beyond text and images."
        },
        {"text": "Milvus has an intuitive REST API for easy integration."},
        {"text": "Milvus' FAISS and HNSW backends provide flexibility in indexing."},
        {
            "text": "The architecture of Milvus ensures fault tolerance and high availability."
        },
        {"text": "Milvus integrates seamlessly with LLM-based applications."},
        {"text": "Startups leverage Milvus to build next-gen AI-powered products."},
        {"text": "Milvus Cloud offers a managed solution for vector search at scale."},
        {
            "text": "The future of AI search is being shaped by Milvus and similar vector databases."
        },
    ],
)
{'insert_count': 37, 'ids': [456486814660619140, 456486814660619141, 456486814660619142, 456486814660619143, 456486814660619144, 456486814660619145, 456486814660619146, 456486814660619147, 456486814660619148, 456486814660619149, 456486814660619150, 456486814660619151, 456486814660619152, 456486814660619153, 456486814660619154, 456486814660619155, 456486814660619156, 456486814660619157, 456486814660619158, 456486814660619159, 456486814660619160, 456486814660619161, 456486814660619162, 456486814660619163, 456486814660619164, 456486814660619165, 456486814660619166, 456486814660619167, 456486814660619168, 456486814660619169, 456486814660619170, 456486814660619171, 456486814660619172, 456486814660619173, 456486814660619174, 456486814660619175, 456486814660619176], 'cost': 0}

구조화된 결과를 위한 출력 타입 정의 (Defining Output Types for Structured Results)

검색 결과를 더 구조화하고 다루기 쉽게 만들려면 검색 결과 형식을 지정하는 Pydantic 모델을 정의해요.

from pydantic import BaseModel

# Simplified output model for search results
class MilvusSearchResult(BaseModel):
    id: int
    text: str

class MilvusSearchResults(BaseModel):
    results: list[MilvusSearchResult]
    query: str

사용자 지정 검색 도구 만들기 (Creating a Custom Search Tool)

다음으로 에이전트가 Milvus 데이터베이스를 검색하는 데 사용할 수 있는 사용자 지정 함수 도구를 만들어요. 이 도구는 다음을 수행해요.

  • 컬렉션 이름, 쿼리 텍스트, limit 파라미터를 받아요.
  • Milvus 컬렉션에 대해 BM25 검색을 실행해요.
  • 결과를 구조화된 형식으로 반환해요.
import json
from typing import Any
from pymilvus import MilvusClient
from agents import function_tool, RunContextWrapper

@function_tool
async def search_milvus_text(
    ctx: RunContextWrapper[Any], collection_name: str, query_text: str, limit: int
) -> str:
    """Search for text documents in a Milvus collection using full text search.

    Args:
        collection_name: Name of the Milvus collection to search.
        query_text: The text query to search for.
        limit: Maximum number of results to return.
    """
    try:
        # Initialize Milvus client
        client = MilvusClient()

        # Prepare search parameters for BM25
        search_params = {"metric_type": "BM25", "params": {"drop_ratio_search": 0.2}}

        # Execute search with text query
        results = client.search(
            collection_name=collection_name,
            data=[query_text],
            anns_field="sparse",
            limit=limit,
            search_params=search_params,
            output_fields=["text"],
        )
        return json.dumps(
            {"results": results, "query": query_text, "collection": collection_name}
        )

    except Exception as e:
        print(f"Exception is: {e}")
        return f"Error searching Milvus: {str(e)}"

에이전트 구축 (Building the Agent)

이제 검색 도구를 사용할 수 있는 에이전트를 만들어요. 검색 요청을 처리하는 방법에 대한 지시를 주고 결과를 구조화된 형식으로 반환하도록 지정해요.

from agents import Agent, Runner, WebSearchTool, trace

async def main():
    agent = Agent(
        name="Milvus Searcher",
        instructions="""
        You are a helpful agent that can search through Milvus vector database using full text search. Return the results in a structured format.
        """,
        tools=[
            WebSearchTool(user_location={"type": "approximate", "city": "New York"}),
            search_milvus_text,
        ],
        output_type=MilvusSearchResults,
    )

    with trace("Milvus search example"):
        result = await Runner.run(
            agent,
            "Find documents in the 'demo' collection that are similar to this concept: 'information retrieval'",
        )
        # print(result.final_output.results)
        formatted_results = "\n".join(
            f"{i+1}. ID: {res.id}, Text: {res.text}"
            for i, res in enumerate(result.final_output.results)
        )
        print(f"Search results:\n{formatted_results}")
asyncio.run(main())
Search results:
1. ID: 456486814660619146, Text: Boolean retrieval is one of the earliest information retrieval methods.
2. ID: 456486814660619144, Text: Machine learning improves ranking algorithms in information retrieval.
3. ID: 456486814660619143, Text: Vector search is revolutionising modern information retrieval systems.
4. ID: 456486814660619140, Text: Information retrieval helps users find relevant documents in large datasets.
5. ID: 456486814660619141, Text: Search engines use information retrieval techniques to index and rank web pages.

더 알아보기 (Learn more)