Интеграция Milvus с агентами OpenAI: Пошаговое руководство

В этом блокноте показано, как создать агента, который может запрашивать Milvus на естественном языке с помощью вызова функций. Мы объединим фреймворк OpenAI Agents с мощными возможностями векторного поиска Milvus для создания удобного поиска.

Агенты OpenAI

OpenAI Agents SDK позволяет создавать агентурные приложения ИИ в легком, простом в использовании пакете с очень небольшим количеством абстракций. Это готовое к производству обновление их предыдущей экспериментальной разработки для агентов, Swarm. Agents SDK имеет очень небольшой набор примитивов:

• Агенты, которые представляют собой LLM, оснащенные инструкциями и инструментами.
• Handoffs, которые позволяют агентам делегировать другим агентам выполнение определенных задач
• Guardrails, которые позволяют проверять входные данные агентов.

В сочетании с Python эти примитивы достаточно мощны, чтобы выразить сложные отношения между инструментами и агентами, и позволяют создавать реальные приложения без сложного обучения. Кроме того, SDK поставляется со встроенной трассировкой, которая позволяет визуализировать и отлаживать агентские потоки, а также оценивать их и даже настраивать модели для вашего приложения.

Milvus

Milvus - это высокопроизводительная и масштабируемая векторная база данных с открытым исходным кодом, которая эффективно работает в самых разных средах, от ноутбука до крупных распределенных систем. Она доступна как в виде программного обеспечения с открытым исходным кодом, так и в виде облачного предложения.

Настройка и зависимости

Сначала нам нужно настроить среду с необходимыми библиотеками и инициализировать asyncio для совместимости с Jupyter.

$ pip install openai pymilvus pydantic nest_asyncio

import asyncio
import nest_asyncio
from dotenv import load_dotenv

load_dotenv()

nest_asyncio.apply()

Мы будем использовать модели из OpenAI. Вам необходимо подготовить api ключ OPENAI_API_KEY в качестве переменной окружения.

import os

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

Подключение к Milvus и создание схемы

Теперь мы подключимся к нашему экземпляру Milvus и создадим схему для нашей коллекции. Эта схема будет определять структуру наших данных, включая:

• поле ID в качестве первичного ключа
• текстовое поле для хранения содержимого документа
• Разреженное векторное поле для хранения вкраплений BM25.

Полнотекстовый поиск в Milvus 2.5

• Единая система для векторного поиска и поиска по ключевым словам (единые API)
• Встроенный алгоритм sparse-BM25 (аналогичен использованию Elasticsearch, но основан на векторах)
• Нет необходимости вручную генерировать вкрапления для поиска по ключевым словам

Установка Milvus с помощью 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}

Настройка 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. Выполнять поиск BM25 по коллекции Milvus
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.