Milvus와 OpenAI 에이전트 통합: 단계별 가이드

이 노트북에서는 함수 호출을 통해 자연어를 사용하여 Milvus를 쿼리할 수 있는 에이전트를 만드는 방법을 보여드립니다. OpenAI의 에이전트 프레임워크와 Milvus의 강력한 벡터 검색 기능을 결합하여 멋진 검색 환경을 만들어 보겠습니다.

OpenAI 에이전트

OpenAI 에이전트 SDK를 사용하면 추상화가 거의 없는 가볍고 사용하기 쉬운 패키지로 에이전트 AI 앱을 구축할 수 있습니다. 이는 에이전트를 위한 이전 실험 버전인 Swarm을 프로덕션용으로 업그레이드한 것입니다. 에이전트 SDK에는 매우 작은 기본 요소 집합이 있습니다:

• 에이전트: 지침과 도구를 갖춘 LLM인 에이전트
• 핸드오프: 상담원이 특정 작업을 다른 상담원에게 위임할 수 있는 기능
• 에이전트에 대한 입력의 유효성을 검사할 수 있는 가드레일

이러한 기본 요소는 Python과 함께 사용하면 도구와 에이전트 간의 복잡한 관계를 표현할 수 있을 만큼 강력하며, 가파른 학습 곡선 없이 실제 애플리케이션을 구축할 수 있습니다. 또한 SDK에는 에이전트 플로우를 시각화하고 디버깅할 수 있는 추적 기능이 내장되어 있어 이를 평가하고 애플리케이션에 맞게 모델을 미세 조정할 수도 있습니다.

Milvus

Milvus는 노트북부터 대규모 분산 시스템에 이르기까지 다양한 환경에서 효율적으로 실행되는 고성능, 확장성 높은 오픈 소스 벡터 데이터베이스입니다. 오픈 소스 소프트웨어와 클라우드 제품으로 모두 제공됩니다.

설정 및 종속성

먼저 필요한 라이브러리로 환경을 설정하고 Jupyter 호환성을 위해 비동기화를 초기화해야 합니다.

$ pip install openai pymilvus pydantic nest_asyncio

import asyncio
import nest_asyncio
from dotenv import load_dotenv

load_dotenv()

nest_asyncio.apply()

OpenAI의 모델을 사용합니다. 환경 변수로 OPENAI_API_KEY API 키를 준비해야 합니다.

import os

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

Milvus에 연결하고 스키마 만들기

이제 Milvus 인스턴스에 연결하여 컬렉션에 대한 스키마를 생성하겠습니다. 이 스키마는 다음을 포함하여 데이터의 구조를 정의합니다:

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

Milvus 2.5의 전체 텍스트 검색

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

Docker로 Milvus 설치

이 예제를 실행하기 전에, 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}

전체 텍스트 검색을 위한 BM25 설정하기

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': {}}]}

컬렉션 생성 및 샘플 데이터 로드

이제 스키마와 인덱스 매개변수를 사용하여 컬렉션을 생성한 다음 정보 검색 및 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}

구조화된 결과를 위한 출력 유형 정의

검색 결과를 보다 체계적이고 작업하기 쉽게 만들기 위해 검색 결과의 형식을 지정하는 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

사용자 정의 검색 도구 만들기

다음으로, 상담원이 Milvus 데이터베이스를 검색하는 데 사용할 수 있는 사용자 지정 함수 도구를 만들겠습니다. 이 도구는 다음과 같습니다:

1. 컬렉션 이름, 쿼리 텍스트 및 제한 매개변수를 수락합니다.
2. Milvus 컬렉션에 대해 BM25 검색을 실행합니다.
3. 결과를 구조화된 형식으로 반환합니다.

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)}"

에이전트 구축하기

이제 검색 도구를 사용할 수 있는 에이전트를 만들어 보겠습니다. 검색 요청을 처리하는 방법에 대한 지침을 제공하고 구조화된 형식으로 결과를 반환하도록 지정할 것입니다.

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.