Integrasi Milvus dengan Agen OpenAI: Panduan Langkah-demi-Langkah

Buku catatan ini menunjukkan cara membuat agen yang dapat menanyakan Milvus menggunakan bahasa alami melalui Pemanggilan Fungsi. Kami akan menggabungkan kerangka kerja Agents OpenAI dengan kemampuan pencarian vektor yang kuat dari Milvus untuk menciptakan pengalaman pencarian yang bagus.

Agen OpenAI

OpenAI Agents SDK memungkinkan Anda untuk membangun aplikasi AI agen dalam paket yang ringan dan mudah digunakan dengan abstraksi yang sangat sedikit. Ini adalah peningkatan siap produksi dari eksperimen mereka sebelumnya untuk agen, Swarm. Agents SDK memiliki sekumpulan primitif yang sangat kecil:

• Agen, yang merupakan LLM yang dilengkapi dengan instruksi dan alat
• Handoff, yang memungkinkan agen untuk mendelegasikan ke agen lain untuk tugas-tugas tertentu
• Pagar, yang memungkinkan input ke agen untuk divalidasi

Jika dikombinasikan dengan Python, primitif-primitif ini cukup kuat untuk mengekspresikan hubungan yang kompleks antara alat dan agen, dan memungkinkan Anda untuk membangun aplikasi dunia nyata tanpa kurva pembelajaran yang curam. Selain itu, SDK ini dilengkapi dengan penelusuran bawaan yang memungkinkan Anda memvisualisasikan dan men-debug alur agen Anda, serta mengevaluasinya dan bahkan menyempurnakan model untuk aplikasi Anda.

Milvus

Milvus adalah database vektor Open-Source berkinerja tinggi dan sangat skalabel yang berjalan secara efisien di berbagai lingkungan, mulai dari laptop hingga sistem terdistribusi berskala besar. Milvus tersedia sebagai perangkat lunak sumber terbuka dan Penawaran Cloud.

Pengaturan dan Ketergantungan

Pertama, kita perlu menyiapkan lingkungan kita dengan pustaka yang diperlukan dan menginisialisasi asyncio untuk kompatibilitas Jupyter.

$ pip install openai pymilvus pydantic nest_asyncio

import asyncio
import nest_asyncio
from dotenv import load_dotenv

load_dotenv()

nest_asyncio.apply()

Kita akan menggunakan model dari OpenAI. Anda harus menyiapkan kunci api OPENAI_API_KEY sebagai variabel lingkungan.

import os

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

Menghubungkan ke Milvus dan Membuat Skema

Sekarang kita akan menghubungkan ke instance Milvus dan membuat skema untuk koleksi kita. Skema ini akan mendefinisikan struktur data kita, termasuk:

• Bidang ID sebagai kunci utama
• Bidang teks untuk menyimpan konten dokumen kita
• Bidang vektor jarang untuk menyimpan sematan BM25

Pencarian Teks Lengkap di Milvus 2.5

• Sistem terpadu untuk pencarian vektor dan kata kunci (API terpadu)
• Algoritme BM25 jarang bawaan (serupa dengan yang digunakan Elasticsearch tetapi berbasis vektor)
• Tidak perlu membuat sematan secara manual untuk pencarian kata kunci

Instal Milvus dengan Docker

Sebelum menjalankan contoh ini, pastikan untuk menginstal Milvus dan memulainya dengan Docker, lihat dokumentasi kami - 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}

Menyiapkan BM25 untuk Pencarian Teks Lengkap

Milvus mendukung pencarian teks lengkap melalui fungsi BM25. Di sini kita mendefinisikan sebuah fungsi yang secara otomatis akan mengubah data teks kita menjadi representasi vektor jarang yang dioptimalkan untuk pencarian teks.

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

Membuat Koleksi dan Memuat Data Sampel

Sekarang kita akan membuat koleksi kita dengan skema dan parameter indeks, lalu memuat beberapa data sampel tentang pencarian informasi dan 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}

Mendefinisikan Jenis Keluaran untuk Hasil Terstruktur

Untuk membuat hasil pencarian kita lebih terstruktur dan lebih mudah digunakan, kita akan mendefinisikan model Pydantic yang menentukan format hasil pencarian kita.

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

Membuat Alat Pencarian Khusus

Selanjutnya, kita akan membuat alat fungsi khusus yang dapat digunakan oleh agen kita untuk mencari database Milvus. Alat ini akan:

1. Menerima nama koleksi, teks kueri, dan parameter batas
2. Menjalankan pencarian BM25 terhadap koleksi Milvus
3. Mengembalikan hasilnya dalam format terstruktur

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

Membangun Agen

Sekarang kita akan membuat sebuah agen yang dapat menggunakan alat bantu pencarian kita. Kita akan memberinya instruksi tentang bagaimana menangani permintaan pencarian dan menentukan bahwa ia harus mengembalikan hasil dalam format terstruktur.

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.