MilvusとOpenAIエージェントの統合：ステップバイステップガイド

このノートブックでは、関数呼び出しを通して自然言語を使ってMilvusに問い合わせができるエージェントを作成する方法を紹介します。OpenAIのエージェントフレームワークとMilvusの強力なベクトル検索機能を組み合わせて、素敵な検索体験を作りましょう。

OpenAIエージェント

OpenAI Agents SDKは、エージェント型AIアプリを軽量で使いやすいパッケージで構築することを可能にします。このSDKは、OpenAIのエージェント向け実験製品であるSwarmを製品化できるようにアップグレードしたものです。エージェントSDKには、非常に小さなプリミティブのセットがある：

• エージェントは、命令とツールを備えたLLMである。
• ハンドオフ：エージェントは、特定のタスクを他のエージェントに委任することができる。
• エージェントへの入力を検証するガードレール

Pythonと組み合わせることで、これらのプリミティブはツールとエージェント間の複雑な関係を表現するのに十分強力であり、急な学習曲線なしに実世界のアプリケーションを構築することができます。さらに、SDKには組み込みのトレーシング機能があり、エージェントフローを可視化してデバッグしたり、評価したり、アプリケーションのモデルを微調整したりすることができます。

Milvus

Milvusは、ラップトップから大規模な分散システムまで、幅広い環境で効率的に動作する、高性能で拡張性の高いオープンソースのベクトルデータベースです。オープンソースソフトウェアとしても、クラウドオファリングとしてもご利用いただけます。

セットアップと依存関係

まず、Jupyterとの互換性を保つために、必要なライブラリとasyncioを初期化して環境をセットアップする必要がある。

$ pip install openai pymilvus pydantic nest_asyncio

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への接続とスキーマの作成

Milvusインスタンスに接続し、コレクションのスキーマを作成します。このスキーマは以下のようなデータの構造を定義します：

• 主キーとしてのIDフィールド
• ドキュメントの内容を格納するテキストフィールド
• BM25埋め込みを格納するスパースベクトルフィールド

Milvus 2.5での全文検索

• ベクトル検索とキーワード検索のための統一されたシステム（統一されたAPI）
• 組み込みのスパースBM25アルゴリズム（Elasticsearchと同様、ベクトルベース）
• キーワード検索のための埋め込みを手動で生成する必要がない

Dockerによるmilvusのインストール

この例を実行する前に、Milvusをインストールし、Dockerで起動することを確認してください。

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.