Milvus-Integration mit OpenAI-Agenten: Eine Schritt-für-Schritt-Anleitung

Dieses Notizbuch zeigt, wie man einen Agenten erstellt, der Milvus mit natürlicher Sprache durch Funktionsaufrufe abfragen kann. Wir kombinieren das OpenAI Agents-Framework mit den leistungsstarken Vektorsuchfunktionen von Milvus, um eine schöne Sucherfahrung zu schaffen.

OpenAI-Agenten

Mit dem OpenAI Agents SDK können Sie agentenbasierte KI-Anwendungen in einem leichtgewichtigen, benutzerfreundlichen Paket mit nur wenigen Abstraktionen erstellen. Es handelt sich um ein produktionsreifes Upgrade ihres früheren Experiments für Agenten, Swarm. Das Agents SDK verfügt über einen sehr kleinen Satz von Primitiven:

• Agenten, die LLMs sind, die mit Anweisungen und Werkzeugen ausgestattet sind
• Handoffs, die es Agenten ermöglichen, bestimmte Aufgaben an andere Agenten zu delegieren
• Guardrails, die es ermöglichen, die Eingaben der Agenten zu validieren

In Kombination mit Python sind diese Primitive leistungsfähig genug, um komplexe Beziehungen zwischen Tools und Agenten auszudrücken, und ermöglichen es Ihnen, reale Anwendungen ohne eine steile Lernkurve zu erstellen. Darüber hinaus verfügt das SDK über eine integrierte Ablaufverfolgung, mit der Sie Ihre Agentenflüsse visualisieren und debuggen sowie auswerten und sogar Modelle für Ihre Anwendung feinabstimmen können.

Milvus

Milvus ist eine leistungsstarke, hoch skalierbare Open-Source-Vektordatenbank, die effizient in einer Vielzahl von Umgebungen läuft, von einem Laptop bis zu großen verteilten Systemen. Sie ist sowohl als Open-Source-Software als auch als Cloud-Angebot erhältlich.

Einrichtung und Abhängigkeiten

Zunächst müssen wir unsere Umgebung mit den erforderlichen Bibliotheken einrichten und asyncio für die Jupyter-Kompatibilität initialisieren.

$ pip install openai pymilvus pydantic nest_asyncio

import asyncio
import nest_asyncio
from dotenv import load_dotenv

load_dotenv()

nest_asyncio.apply()

Wir werden die Modelle von OpenAI verwenden. Sie sollten den api-Schlüssel OPENAI_API_KEY als Umgebungsvariable vorbereiten.

import os

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

Verbindung zu Milvus herstellen und ein Schema erstellen

Nun werden wir uns mit unserer Milvus-Instanz verbinden und ein Schema für unsere Sammlung erstellen. Dieses Schema wird die Struktur unserer Daten definieren, einschließlich:

• Ein ID-Feld als Primärschlüssel
• Ein Textfeld zum Speichern des Dokumentinhalts
• Ein spärliches Vektorfeld zum Speichern der BM25-Einbettungen

Volltextsuche in Milvus 2.5

• Einheitliches System für die Vektor- und Schlagwortsuche (einheitliche APIs)
• Eingebauter sparse-BM25-Algorithmus (ähnlich wie Elasticsearch, aber vektorbasiert)
• Keine Notwendigkeit der manuellen Erzeugung von Einbettungen für die Schlagwortsuche

Milvus mit Docker installieren

Bevor Sie dieses Beispiel ausführen, stellen Sie sicher, dass Sie Milvus installieren und mit Docker starten. Schauen Sie in unsere Dokumentation - 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 für die Volltextsuche einrichten

Milvus unterstützt die Volltextsuche durch BM25-Funktionen. Hier definieren wir eine Funktion, die unsere Textdaten automatisch in spärliche Vektordarstellungen konvertiert, die für die Textsuche optimiert sind.

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

Erstellen der Sammlung und Laden von Beispieldaten

Nun erstellen wir unsere Sammlung mit den Schema- und Indexparametern und laden dann einige Beispieldaten über Information Retrieval und 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}

Definition von Ausgabetypen für strukturierte Ergebnisse

Um unsere Suchergebnisse besser zu strukturieren und die Arbeit damit zu erleichtern, werden wir Pydantic-Modelle definieren, die das Format unserer Suchergebnisse festlegen.

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

Erstellen eines benutzerdefinierten Suchwerkzeugs

Als Nächstes erstellen wir ein benutzerdefiniertes Funktionswerkzeug, das unser Agent zum Durchsuchen der Milvus-Datenbank verwenden kann. Dieses Werkzeug wird:

1. einen Sammlungsnamen, einen Abfragetext und einen Grenzparameter akzeptieren
2. eine BM25-Suche in der Milvus-Sammlung durchführen
3. Rückgabe der Ergebnisse in einem strukturierten Format

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

Erstellen des Agenten

Nun werden wir einen Agenten erstellen, der unser Suchwerkzeug verwenden kann. Wir geben ihm Anweisungen, wie er Suchanfragen behandeln soll, und legen fest, dass er die Ergebnisse in unserem strukturierten Format zurückgeben soll.

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.