Генерация выражений фильтрации запросов 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. Замените 'your_openai_api_key' на ваш реальный ключ OpenAI API.

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

Пример данных для печати 3

Приведенный выше код создает коллекцию Milvus со следующей структурой:

• pk: Поле первичного ключа (VARCHAR)
• имя: имя пользователя (VARCHAR)
• age: Возраст пользователя (INT64)
• city: Город пользователя (VARCHAR)
• хобби: Хобби пользователя (VARCHAR)
• встраивание: Векторное встраивание (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 лет
• Живут в Лондоне, Токио или Торонто
• Соответствуют семантическому контексту запроса.

Этот подход сочетает в себе точность структурированной фильтрации и гибкость естественного языка, что делает векторные базы данных более доступными для пользователей, которые не знакомы со специфическим синтаксисом запросов.