使用大型語言模型產生 Milvus 查詢篩選表達式

在本教程中,我們將示範如何使用大型語言模型 (LLM) 自動從自然語言查詢生成 Milvus 過濾表達式。這種方法允許使用者以純英文表達複雜的篩選條件,然後將其轉換為適當的 Milvus 語法,使向量資料庫查詢更容易使用。

Milvus 支援複雜的篩選功能,包括

  • 基本運算符號:比較運算符如==,!=,>,<,>=<=
  • 布林運算符:邏輯運算符如and,or,not 用於複雜的條件
  • 字串運算:使用like 和其他字串函數進行模式匹配
  • 陣列操作:使用array_contains,array_length 等處理陣列欄位。
  • JSON 操作:使用專門的運算符查詢 JSON 欄位

透過整合 LLM 與 Milvus 文件,我們可以建立一個智慧型系統,它能理解自然語言查詢,並產生語法正確的篩選表達式。本教學將介紹建立此系統的過程,並強調其在各種過濾情境中的有效性。

依賴與環境

$ pip install --upgrade pymilvus openai requests docling beautifulsoup4
print("Environment setup complete!")

設定環境變數

配置您的 OpenAI API 認證,以啟用嵌入生成和基於 LLM 的篩選表達式創建。以您實際的 OpenAI API 密鑰取代'your_openai_api_key'

import os
import openai

os.environ["OPENAI_API_KEY"] = "your_openai_api_key"
api_key = os.getenv("OPENAI_API_KEY")

if not api_key:
    raise ValueError("Please set the OPENAI_API_KEY environment variable!")

openai.api_key = api_key
print("API key loaded.")

建立樣本集合

現在讓我們建立一個包含使用者資料的 Milvus 樣本集合。此集合將包含標量欄位 (用於篩選) 和向量嵌入 (用於語意搜尋)。我們將使用 OpenAI 的文字嵌入模型來產生使用者資訊的向量表示。

from pymilvus import MilvusClient, FieldSchema, CollectionSchema, DataType
import os
from openai import OpenAI
import uuid

client = MilvusClient(uri="http://localhost:19530")
openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
embedding_model = "text-embedding-3-small"
embedding_dim = 1536

fields = [
    FieldSchema(
        name="pk",
        dtype=DataType.VARCHAR,
        is_primary=True,
        auto_id=False,
        max_length=100,
    ),
    FieldSchema(name="name", dtype=DataType.VARCHAR, max_length=128),
    FieldSchema(name="age", dtype=DataType.INT64),
    FieldSchema(name="city", dtype=DataType.VARCHAR, max_length=128),
    FieldSchema(name="hobby", dtype=DataType.VARCHAR, max_length=128),
    FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
]
schema = CollectionSchema(fields=fields, description="User data embedding example")
collection_name = "user_data_collection"

if client.has_collection(collection_name):
    client.drop_collection(collection_name)
# Strong consistency waits for all loads to complete, adding latency with large datasets
# client.create_collection(
#     collection_name=collection_name, schema=schema, consistency_level="Strong"
# )
client.create_collection(collection_name=collection_name, schema=schema)

index_params = client.prepare_index_params()
index_params.add_index(
    field_name="embedding",
    index_type="IVF_FLAT",
    metric_type="COSINE",
    params={"nlist": 128},
)
client.create_index(collection_name=collection_name, index_params=index_params)

data_to_insert = [
    {"name": "John", "age": 23, "city": "Shanghai", "hobby": "Drinking coffee"},
    {"name": "Alice", "age": 29, "city": "New York", "hobby": "Reading books"},
    {"name": "Bob", "age": 31, "city": "London", "hobby": "Playing chess"},
    {"name": "Eve", "age": 27, "city": "Paris", "hobby": "Painting"},
    {"name": "Charlie", "age": 35, "city": "Tokyo", "hobby": "Cycling"},
    {"name": "Grace", "age": 22, "city": "Berlin", "hobby": "Photography"},
    {"name": "David", "age": 40, "city": "Toronto", "hobby": "Watching movies"},
    {"name": "Helen", "age": 30, "city": "Sydney", "hobby": "Cooking"},
    {"name": "Frank", "age": 28, "city": "Beijing", "hobby": "Hiking"},
    {"name": "Ivy", "age": 26, "city": "Seoul", "hobby": "Dancing"},
    {"name": "Tom", "age": 33, "city": "Madrid", "hobby": "Writing"},
]


def get_embeddings(texts):
    return [
        rec.embedding
        for rec in openai_client.embeddings.create(
            input=texts, model=embedding_model, dimensions=embedding_dim
        ).data
    ]


texts = [
    f"{item['name']} from {item['city']} is {item['age']} years old and likes {item['hobby']}."
    for item in data_to_insert
]
embeddings = get_embeddings(texts)

insert_data = []
for item, embedding in zip(data_to_insert, embeddings):
    item_with_embedding = {
        "pk": str(uuid.uuid4()),
        "name": item["name"],
        "age": item["age"],
        "city": item["city"],
        "hobby": item["hobby"],
        "embedding": embedding,
    }
    insert_data.append(item_with_embedding)

client.insert(collection_name=collection_name, data=insert_data)

print(f"Collection '{collection_name}' has been created and data has been inserted.")

上面的程式碼會以下列結構建立一個 Milvus 套件:

  • pk:主鍵欄位 (VARCHAR)
  • name:使用者名稱 (VARCHAR)
  • 年齡:使用者年齡 (INT64)
  • 城市: 使用者城市 (VARCHAR)使用者城市 (VARCHAR)
  • 興趣使用者的興趣 (VARCHAR)
  • embedding:向量嵌入 (FLOAT_VECTOR, 1536 維度)

我們已插入 11 個樣本使用者的個人資訊,並產生嵌入資料以提供語意搜尋功能。在嵌入之前,每個使用者的資訊都會轉換成描述性文字,以擷取他們的姓名、所在地、年齡和興趣。讓我們透過查詢一些樣本記錄來驗證我們的收集是否已成功建立,並包含預期的資料。

from pymilvus import MilvusClient
import os
from openai import OpenAI

client = MilvusClient(uri="http://localhost:19530")
collection_name = "user_data_collection"

client.load_collection(collection_name=collection_name)

result = client.query(
    collection_name=collection_name,
    filter="",
    output_fields=["name", "age", "city", "hobby"],
    limit=3,
)

for record in result:
    print(record)

收集 Milvus 過濾表達說明文件

為了幫助大型語言模型更好地理解 Milvus 的篩選表達語法,我們需要為它提供相關的官方文件。我們將使用docling 函式庫從 Milvus 官方網站搜刮幾個關鍵頁面。

這些頁面包含以下的詳細資訊

  • 布林運算元and,or,not 用於複雜的邏輯條件
  • 基本運算符號:比較運算符如==,!=,>,<,>=<=
  • 過濾模板:進階過濾模式和語法
  • 字串匹配:使用like 和其他字串操作進行模式匹配

這些文件將作為我們 LLM 的知識基礎,以產生精確的過濾表達式。

import docling
from docling.document_converter import DocumentConverter

converter = DocumentConverter()
docs = [
    converter.convert(url)
    for url in [
        "https://milvus.io/docs/boolean.md",
        "https://milvus.io/docs/basic-operators.md",
        "https://milvus.io/docs/filtering-templating.md",
    ]
]

for doc in docs[:3]:
    print(doc.document.export_to_markdown())

文件搜刮提供 Milvus 過濾語法的全面涵蓋。此知識庫可讓我們的 LLM 瞭解篩選表達式建構的細微差異,包括正確的運算符使用、欄位參照以及複雜的條件組合。

由 LLM 驅動的篩選程式產生

現在我們有了文件上下文,讓我們設定 LLM 系統來產生篩選表達式。我們將建立一個結構化的提示,結合刮取的文件與使用者查詢,產生語法正確的 Milvus 過濾器表達式。

我們的過濾器產生系統使用精心製作的提示,它可以

  1. 提供上下文:包括完整的 Milvus 文檔作為參考資料
  2. 設定限制:確保 LLM 只使用文件中的語法和功能
  3. 強制準確性:要求語法表達正確
  4. 保持焦點:僅傳回過濾表達式,不提供說明

讓我們用自然語言查詢來測試一下,看看 LLM 的表現如何。

from openai import OpenAI
import json
from IPython.display import display, Markdown

context = "\n".join([doc.document.export_to_markdown() for doc in docs])

prompt = f"""
You are an expert Milvus vector database engineer. Your task is to convert a user's natural language query into a valid Milvus filter expression, using the provided Milvus documentation as your knowledge base.

Follow these rules strictly:
1. Only use the provided documents as your source of knowledge.
2. Ensure the generated filter expression is syntactically correct.
3. If there isn't enough information in the documents to create an expression, state that directly.
4. Only return the final filter expression. Do not include any explanations or extra text.

---
**Milvus Documentation Context:**
{context}

---
**User Query:**
{user_query}

---
**Filter Expression:**
"""

client = OpenAI()


def generate_filter_expr(user_query):
    """
    Generates a Milvus filter expression from a user query using GPT-4o-mini.
    """
    completion = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": prompt},
            {"role": "user", "content": user_query},
        ],
        temperature=0.0,
    )
    return completion.choices[0].message.content


user_query = "Find people older than 30 who live in London, Tokyo, or Toronto"

filter_expr = generate_filter_expr(user_query)

print(f"Generated filter expression: {filter_expr}")

LLM 成功產生一個結合多個條件的篩選表達式:

  • 年齡比較使用>
  • 使用in 運算符進行多個城市比對
  • 正確的欄位引用和語法

這展示了提供全面的文件上下文來引導 LLM 過濾器產生的威力。

測試產生的篩選程式

現在讓我們在實際的 Milvus 搜尋作業中使用它來測試我們所產生的篩選表達。我們將結合語意搜尋與精確篩選,找出符合查詢意圖與特定條件的使用者。

from pymilvus import MilvusClient
from openai import OpenAI
import os

client = MilvusClient(uri="http://localhost:19530")
openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))

clean_filter = (
    filter_expr.replace("```", "").replace('filter="', "").replace('"', "").strip()
)
print(f"Using filter: {clean_filter}")

query_embedding = (
    openai_client.embeddings.create(
        input=[user_query], model="text-embedding-3-small", dimensions=1536
    )
    .data[0]
    .embedding
)

search_results = client.search(
    collection_name="user_data_collection",
    data=[query_embedding],
    limit=10,
    filter=clean_filter,
    output_fields=["pk", "name", "age", "city", "hobby"],
    search_params={
        "metric_type": "COSINE",
        "params": {"nprobe": 10},
    },
)

print("Search results:")
for i, hits in enumerate(search_results):
    print(f"Query {i}:")
    for hit in hits:
        print(f"  - {hit}")
    print()

結果分析

搜尋結果顯示 LLM 產生的篩選器與 Milvus 向量搜尋的成功整合。篩選器能正確找出符合下列條件的使用者

  • 年齡超過 30 歲
  • 住在倫敦、東京或多倫多
  • 符合查詢的語意上下文

此方法結合了結構化篩選的精確度與自然語言輸入的彈性,讓不熟悉特定查詢語法的使用者更容易使用向量資料庫。