Geração de expressões de filtro de consulta Milvus com modelos de linguagem grandes

Neste tutorial, demonstraremos como usar Modelos de Linguagem Grandes (LLMs) para gerar automaticamente expressões de filtro Milvus a partir de consultas em linguagem natural. Esta abordagem torna a consulta de bases de dados vectoriais mais acessível, permitindo que os utilizadores expressem condições de filtragem complexas em inglês simples, que são depois convertidas para a sintaxe Milvus adequada.

O Milvus suporta capacidades de filtragem sofisticadas, incluindo

• Operadores básicos: Operadores de comparação como ==, !=, >, <, >=, <=
• Operadores booleanos: Operadores lógicos como and, or, not para condições complexas
• Operações de cadeia de caracteres: Correspondência de padrões com like e outras funções de cadeia de caracteres
• Operações com matrizes: Trabalhar com campos de matriz utilizando array_contains, array_length, etc.
• Operações JSON: Consulta de campos JSON com operadores especializados

Ao integrar os LLMs com a documentação Milvus, podemos criar um sistema inteligente que compreende as consultas em linguagem natural e gera expressões de filtro sintaticamente corretas. Este tutorial irá percorrer o processo de configuração deste sistema, destacando a sua eficácia em vários cenários de filtragem.

Dependências e ambiente

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

Configurar variáveis de ambiente

Configure as credenciais da API do OpenAI para permitir a geração de incorporação e a criação de expressões de filtro baseadas em LLM. Substitua 'your_openai_api_key' pela sua chave de API OpenAI real.

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

Criar uma coleção de amostras

Agora vamos criar uma coleção de amostra do Milvus com dados do usuário. Esta coleção conterá campos escalares (para filtragem) e embeddings vectoriais (para pesquisa semântica). Utilizaremos o modelo de incorporação de texto do OpenAI para gerar representações vectoriais da informação do utilizador.

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

Imprimir 3 dados de amostra

O código acima cria uma coleção Milvus com a seguinte estrutura:

• pk: Campo de chave primária (VARCHAR)
• name: Nome do utilizador (VARCHAR)
• idade: Idade do utilizador (INT64)
• cidade: Cidade do utilizador (VARCHAR)
• hobby: Passatempo do utilizador (VARCHAR)
• embedding: Incorporação vetorial (FLOAT_VECTOR, 1536 dimensões)

Inserimos 11 utilizadores de amostra com as suas informações pessoais e geramos embeddings para capacidades de pesquisa semântica. As informações de cada utilizador são convertidas num texto descritivo que capta o seu nome, localização, idade e interesses antes de serem incorporadas. Vamos verificar se a nossa coleção foi criada com êxito e contém os dados esperados, consultando alguns registos de amostra.

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)

Coleta de documentação de expressão de filtro do Milvus

Para ajudar o modelo de linguagem grande a entender melhor a sintaxe da expressão de filtro do Milvus, precisamos fornecer a ele a documentação oficial relevante. Usaremos a biblioteca docling para extrair várias páginas-chave do site oficial do Milvus.

Estas páginas contêm informações detalhadas sobre:

• Operadores booleanos: and, or, not para condições lógicas complexas
• Operadores básicos: Operadores de comparação como ==, !=, >, <, >=, <=
• Modelos de filtragem: Padrões e sintaxe de filtragem avançados
• Correspondênciade strings: correspondência de padrões com like e outras operações de strings

Esta documentação servirá de base de conhecimento para o nosso LLM gerar expressões de filtragem exactas.

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

A recolha de documentação fornece uma cobertura completa da sintaxe de filtragem do Milvus. Esta base de conhecimentos permitirá ao nosso LLM compreender as nuances da construção de expressões de filtros, incluindo a utilização correta de operadores, a referência a campos e combinações complexas de condições.

Geração de filtros com LLM

Agora que temos o contexto da documentação, vamos configurar o sistema LLM para gerar expressões de filtro. Vamos criar um prompt estruturado que combina a documentação recolhida com as consultas do utilizador para produzir expressões de filtro do Milvus sintaticamente corretas.

O nosso sistema de geração de filtros utiliza uma linha de comandos cuidadosamente elaborada que:

1. Fornece contexto: Inclui a documentação completa do Milvus como material de referência
2. Estabelece restrições: Assegura que o LLM utiliza apenas a sintaxe e as caraterísticas documentadas
3. Impõe precisão: Requer expressões sintaticamente corretas
4. Mantém o foco: Retorna apenas a expressão do filtro sem explicações

Vamos testar isso com uma consulta em linguagem natural e ver o desempenho do 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}")

O LLM gerou com sucesso uma expressão de filtro que combina várias condições:

• Comparação de idades usando >
• Correspondência de várias cidades usando o operador in
• Referência a campos e sintaxe corretas

Isso demonstra o poder de fornecer um contexto de documentação abrangente para orientar a geração de filtros do LLM.

Testar o filtro gerado

Agora vamos testar a nossa expressão de filtro gerada, utilizando-a numa operação de pesquisa Milvus real. Vamos combinar a pesquisa semântica com uma filtragem precisa para encontrar utilizadores que correspondam à intenção da consulta e aos critérios específicos.

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

Análise dos resultados

Os resultados da pesquisa demonstram o sucesso da integração dos filtros gerados pelo LLM na pesquisa vetorial Milvus. O filtro identificou corretamente os utilizadores que:

• Têm mais de 30 anos
• Vivem em Londres, Tóquio ou Toronto
• Correspondem ao contexto semântico da consulta

Esta abordagem combina a precisão da filtragem estruturada com a flexibilidade da entrada em linguagem natural, tornando as bases de dados vectoriais mais acessíveis aos utilizadores que podem não estar familiarizados com a sintaxe específica da consulta.