Generazione di espressioni di filtro per le query Milvus con i Large Language Models

In questo tutorial dimostreremo come utilizzare i Large Language Models (LLM) per generare automaticamente espressioni di filtro Milvus da query in linguaggio naturale. Questo approccio rende più accessibile l'interrogazione di database vettoriali, consentendo agli utenti di esprimere condizioni di filtraggio complesse in inglese semplice, che vengono poi convertite nella sintassi corretta di Milvus.

Milvus supporta sofisticate capacità di filtraggio, tra cui:

• Operatori di base: Operatori di confronto come ==, !=, >, <, >=, <=
• Operatori booleani: Operatori logici come and, or, not per condizioni complesse.
• Operazioni sulle stringhe: Corrispondenza di pattern con like e altre funzioni per le stringhe
• Operazioni con gli array: Lavorare con campi di array usando array_contains, array_length, ecc.
• Operazioni JSON: Interrogazione di campi JSON con operatori specializzati

Integrando gli LLM con la documentazione di Milvus, possiamo creare un sistema intelligente che comprende le query in linguaggio naturale e genera espressioni di filtro sintatticamente corrette. Questo tutorial illustra il processo di impostazione di questo sistema, evidenziandone l'efficacia in vari scenari di filtraggio.

Dipendenze e ambiente

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

Impostare le variabili d'ambiente

Configurare le credenziali API di OpenAI per abilitare la generazione di embedding e la creazione di espressioni di filtro basate su LLM. Sostituire 'your_openai_api_key' con la propria chiave API OpenAI.

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

Creare una collezione di esempi

Ora creiamo una raccolta Milvus di esempio con i dati dell'utente. Questa raccolta conterrà sia campi scalari (per il filtraggio) sia embeddings vettoriali (per la ricerca semantica). Utilizzeremo il modello di incorporazione del testo di OpenAI per generare rappresentazioni vettoriali delle informazioni dell'utente.

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

Stampa 3 dati di esempio

Il codice precedente crea una collezione Milvus con la seguente struttura:

• pk: Campo chiave primaria (VARCHAR)
• name: Nome utente (VARCHAR)
• età: Età dell'utente (INT64)
• città: Città dell'utente (VARCHAR)
• hobby: Hobby dell'utente (VARCHAR)
• incorporazione: Incorporamento vettoriale (FLOAT_VECTOR, 1536 dimensioni)

Abbiamo inserito 11 utenti campione con le loro informazioni personali e generato embedding per le funzionalità di ricerca semantica. Le informazioni di ciascun utente vengono convertite in un testo descrittivo che ne cattura il nome, la posizione, l'età e gli interessi prima di essere incorporato. Verifichiamo che la nostra raccolta sia stata creata con successo e che contenga i dati previsti, interrogando alcuni record di esempio.

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)

Raccolta della documentazione dell'espressione del filtro Milvus

Per aiutare il modello linguistico di grandi dimensioni a comprendere meglio la sintassi delle espressioni di filtro di Milvus, dobbiamo fornirgli la documentazione ufficiale pertinente. Utilizzeremo la libreria docling per raccogliere diverse pagine chiave dal sito ufficiale di Milvus.

Queste pagine contengono informazioni dettagliate su:

• Operatori booleani: and, or, not per le condizioni logiche complesse
• Operatori di base: Operatori di confronto come ==, !=, >, <, >=, <=
• Modelli di filtraggio: Modelli e sintassi di filtraggio avanzati
• Corrispondenza di stringhe: Corrispondenza di pattern con like e altre operazioni sulle stringhe

Questa documentazione servirà come base di conoscenza per il nostro LLM per generare espressioni di filtro accurate.

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())

Lo scraping della documentazione fornisce una copertura completa della sintassi dei filtri di Milvus. Questa base di conoscenze consentirà al nostro LLM di comprendere le sfumature della costruzione delle espressioni di filtro, compreso l'uso corretto degli operatori, il riferimento ai campi e le combinazioni di condizioni complesse.

Generazione dei filtri con l'LLM

Ora che abbiamo il contesto della documentazione, impostiamo il sistema LLM per generare espressioni di filtro. Creeremo un prompt strutturato che combina la documentazione scraped con le query dell'utente per produrre espressioni di filtro Milvus sintatticamente corrette.

Il nostro sistema di generazione dei filtri utilizza un prompt accuratamente realizzato che:

1. Fornisce un contesto: Include la documentazione completa di Milvus come materiale di riferimento.
2. Imposta vincoli: Assicura che l'LLM utilizzi solo la sintassi e le caratteristiche documentate.
3. Impone l'accuratezza: Richiede espressioni sintatticamente corrette
4. Mantiene l'attenzione: Restituisce solo l'espressione del filtro senza spiegazioni

Verifichiamo questo con una query in linguaggio naturale e vediamo come si comporta l'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}")

L'LLM ha generato con successo un'espressione di filtro che combina più condizioni:

• Confronto dell'età utilizzando >
• Corrispondenza di più città con l'operatore in
• Riferimenti di campo e sintassi corretti

Questo dimostra la potenza di fornire un contesto di documentazione completo per guidare la generazione di filtri da parte di LLM.

Testare il filtro generato

Ora testiamo l'espressione del filtro generato utilizzandola in una vera operazione di ricerca Milvus. Combineremo la ricerca semantica con un filtro preciso per trovare gli utenti che corrispondono sia all'intento della query che ai criteri specifici.

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()

Analisi dei risultati

I risultati della ricerca dimostrano il successo dell'integrazione dei filtri generati da LLM con la ricerca vettoriale di Milvus. Il filtro ha identificato correttamente gli utenti che:

• hanno più di 30 anni
• vivono a Londra, Tokyo o Toronto
• Corrispondono al contesto semantico della query

Questo approccio combina la precisione del filtraggio strutturato con la flessibilità dell'input in linguaggio naturale, rendendo i database vettoriali più accessibili agli utenti che potrebbero non avere familiarità con la sintassi specifica delle query.