Notebooks
M
Milvus Openai Agents Milvus

Openai Agents Milvus

image-searchvector-databasesemantic-searchIntegrationmilvusembeddingsunstructured-dataquestion-answeringLLMmilvus-bootcampdeep-learningimage-recognitionimage-classificationaudio-searchPythonragNLP
alph-notebooks/milvus-bootcamp / openai_agents_milvus.ipynb

Milvus Integration with OpenAI Agents: A Step-by-Step Guide

This notebook shows how to create an agent that can query Milvus using natural language through Function Calling. We'll combine OpenAI's Agents framework with Milvus's powerful vector search capabilities to create a nice search experience.

OpenAI Agents

The OpenAI Agents SDK enables you to build agentic AI apps in a lightweight, easy-to-use package with very few abstractions. It's a production-ready upgrade of their previous experimentation for agents, Swarm. The Agents SDK has a very small set of primitives:

  • Agents, which are LLMs equipped with instructions and tools
  • Handoffs, which allow agents to delegate to other agents for specific tasks
  • Guardrails, which enable the inputs to agents to be validated

In combination with Python, these primitives are powerful enough to express complex relationships between tools and agents, and allow you to build real-world applications without a steep learning curve. In addition, the SDK comes with built-in tracing that lets you visualize and debug your agentic flows, as well as evaluate them and even fine-tune models for your application.

Milvus

Milvus is a high-performance, highly scalable Open-Source vector database that runs efficiently across a wide range of environments, from a laptop to large-scale distributed systems. It is available as both open-source software and a Cloud Offering.

Setup and Dependencies

First, we need to set up our environment with the necessary libraries and initialize asyncio for Jupyter compatibility.

[1]

If you are using Google Colab, to enable dependencies just installed, you may need to restart the runtime (click on the "Runtime" menu at the top of the screen, and select "Restart session" from the dropdown menu).

[1]

We will use the models from OpenAI. You should prepare the api key OPENAI_API_KEY as an environment variable.

[2]

Connecting to Milvus and Creating a Schema

Now we'll connect to our Milvus instance and create a schema for our collection. This schema will define the structure of our data, including:

  • An ID field as the primary key
  • A text field to store our document content
  • A sparse vector field to store the BM25 embeddings

Full-Text Search in Milvus 2.5

  • Unified system for both vector and keyword search (unified APIs)
  • Built-in sparse-BM25 algorithm (similar as Elasticsearch use but vector based)
  • No need to manually generate embeddings for keyword search
img

Install Milvus with Docker

Before running this example, make sure to install Milvus and start it with Docker, have a look at our documentation - https://milvus.io/docs/install_standalone-docker.md

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

Setting Up BM25 for Full-Text Search

Milvus supports full-text search through BM25 functions. Here we define a function that will automatically convert our text data into sparse vector representations optimized for text search.

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

Creating the Collection and Loading Sample Data

Now we'll create our collection with the schema and index parameters, then load some sample data about information retrieval and Milvus.

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

Defining Output Types for Structured Results

To make our search results more structured and easier to work with, we'll define Pydantic models that specify the format of our search results.

[5]

Creating a Custom Search Tool

Next, we'll create a custom function tool that our agent can use to search the Milvus database. This tool will:

  1. Accept a collection name, query text, and limit parameter
  2. Execute a BM25 search against the Milvus collection
  3. Return the results in a structured format
[6]

Building the Agent

Now we'll create an agent that can use our search tool. We'll give it instructions on how to handle search requests and specify that it should return results in our structured format.

[7]
[8]
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.

⭐️ Github

We hope you liked this tutorial showcasing how to use Milvus with OpenAI Agents. If you liked it and our project, please give us a star on Github! ⭐

image.png

🤝 Add me on Linkedin!

If you have some questions related to Milvus, GenAI, etc, I am Stephen Batifol, you can add me on LinkedIn and I'll gladly help you.

image.png

💬 Join our Discord

If you're interested in learning more about Milvus or you wanna share some feedback, feel free to join our Discord channel.