Imperva optimiza la generación de SQL a partir de lenguaje natural utilizando Amazon Bedrock

Esta es una publicación invitada coescrita con Ori Nakar de Imperva.

Imperva Cloud WAF protege cientos de miles de sitios web contra amenazas cibernéticas y bloquea miles de millones de eventos de seguridad todos los días. Los contadores y la información basados ​​en eventos de seguridad se calculan diariamente y los utilizan usuarios de varios departamentos. Diariamente se agregan millones de contadores, junto con 20 millones de conocimientos que se actualizan diariamente para detectar patrones de amenazas.

Nuestro objetivo era mejorar la experiencia del usuario de una aplicación existente utilizada para explorar los contadores y los datos de información. Los datos se almacenan en un lago de datos y SQL los recupera utilizando Atenea amazónica.

Como parte de nuestra solución, reemplazamos múltiples campos de búsqueda con un único campo de texto libre. Utilizamos un modelo de lenguaje grande (LLM) con ejemplos de consulta para que la búsqueda funcione utilizando el lenguaje utilizado por los usuarios internos de Imperva (analistas de negocios).

La siguiente figura muestra una consulta de búsqueda traducida a SQL y ejecutada. Posteriormente, la aplicación formateó los resultados como un gráfico. Disponemos de muchos tipos de información: información global, industrial y a nivel de cliente utilizada por múltiples departamentos, como marketing, soporte e investigación. Los datos se pusieron a disposición de nuestros usuarios a través de una experiencia de usuario simplificada impulsada por un LLM.

Figura 1: Búsqueda de insights por lenguaje natural

Roca Amazónica es un servicio totalmente administrado que ofrece una selección de modelos básicos (FM) de alto rendimiento de empresas líderes en inteligencia artificial (IA), como AI21 Labs, Anthropic, Cohere, Meta, Mistral, Stability AI y Amazon dentro de una única API, junto con con un amplio conjunto de capacidades que necesita para crear aplicaciones de IA generativa con seguridad, privacidad e IA responsable. Amazon Bedrock Studio es una nueva interfaz web habilitada para inicio de sesión único (SSO) que proporciona una manera para que los desarrolladores de una organización experimenten con LLM y otros FM, colaboren en proyectos e iteren en aplicaciones de IA generativa. Ofrece un entorno de creación rápida de prototipos y agiliza el acceso a múltiples FM y herramientas de desarrollo en Amazon Bedrock.

Lea más para conocer el problema y cómo obtuvimos resultados de calidad utilizando Amazon Bedrock para nuestra experimentación e implementación.

El problema

Hacer que los datos sean accesibles para los usuarios a través de aplicaciones siempre ha sido un desafío. Los datos normalmente se almacenan en bases de datos y se pueden consultar utilizando el lenguaje de consulta más común, SQL. Las aplicaciones utilizan diferentes componentes de la interfaz de usuario para permitir a los usuarios filtrar y consultar los datos. Hay aplicaciones con decenas de filtros diferentes y otras opciones, todas creadas para que los datos sean accesibles.

Consultar bases de datos a través de aplicaciones no puede ser tan flexible como ejecutar consultas SQL en un esquema conocido. Dar más poder al usuario se debe a una experiencia de usuario simple (UX). El lenguaje natural puede resolver este problema: es posible admitir consultas en lenguaje natural complejas pero legibles sin conocimientos de SQL. En los cambios de esquema, la UX y el código de la aplicación siguen siendo los mismos, o con cambios menores, lo que ahorra tiempo de desarrollo y mantiene estable la interfaz de usuario (UI) de la aplicación para los usuarios.

Construir consultas SQL a partir de lenguaje natural no es una tarea sencilla. Las consultas SQL deben ser precisas tanto sintáctica como lógicamente. Utilizar un LLM con los ejemplos correctos puede hacer que esta tarea sea menos difícil.

Figura 2: Acceso a la base de datos de alto nivel mediante un flujo de LLM

El reto

Un LLM puede construir consultas SQL basadas en lenguaje natural. El desafío es asegurar la calidad. El usuario puede ingresar cualquier texto y la aplicación construye una consulta basada en él. No existe una opción, como en las aplicaciones tradicionales, para cubrir todas las opciones y asegurarse de que la aplicación funcione correctamente. Agregar un LLM a una aplicación agrega otra capa de complejidad. La respuesta del LLM no es determinista. Los ejemplos enviados al LLM se basan en los datos de la base de datos, lo que hace aún más difícil controlar las solicitudes enviadas al LLM y garantizar la calidad.

La solución: un enfoque de ciencia de datos

En la ciencia de datos, es común desarrollar un modelo y ajustarlo mediante experimentación. La idea es utilizar métricas para comparar experimentos durante el desarrollo. Los experimentos pueden diferir entre sí en muchos aspectos, como la entrada enviada al modelo, el tipo de modelo y otros parámetros. La capacidad de comparar diferentes experimentos permite avanzar. Es posible saber cómo contribuye cada cambio al modelo.

Un conjunto de prueba es un conjunto estático de registros que incluye un resultado de predicción para cada registro. La ejecución de predicciones en el conjunto de pruebas registra los resultados con las métricas necesarias para comparar experimentos. Una métrica común es la precisión, que es el porcentaje de resultados correctos.

En nuestro caso los resultados generados por el LLM son sentencias SQL. Las declaraciones SQL generadas por el LLM no son deterministas y son difíciles de medir; sin embargo, ejecutar declaraciones SQL en una base de datos de prueba estática es determinista y se puede medir. Utilizamos una base de datos de prueba y una lista de preguntas con respuestas conocidas como conjunto de prueba. Nos permitió realizar experimentos y ajustar nuestra aplicación basada en LLM.

Acceso a la base de datos mediante LLM: flujo de pregunta a respuesta

Ante una pregunta definimos el siguiente flujo. La pregunta se envía a través de un proceso de generación aumentada de recuperación (RAG), que encuentra documentos similares. Cada documento contiene una pregunta de ejemplo e información al respecto. Los documentos relevantes se crean como un mensaje y se envían al LLM, que genera una declaración SQL. Este flujo se utiliza tanto para el desarrollo como para el tiempo de ejecución de la aplicación:

Figura 3: Flujo de pregunta a respuesta

Como ejemplo, considere un esquema de base de datos con dos tablas: pedidos y artículos. La siguiente figura es una pregunta para el flujo de ejemplo de SQL:

Figura 4: Ejemplo de flujo de pregunta a respuesta

Acceso a la base de datos mediante LLM: proceso de desarrollo

Para desarrollar y ajustar la aplicación, creamos los siguientes conjuntos de datos:

• Una base de datos de prueba estática: contiene las tablas relevantes y una copia de muestra de los datos.
• Un conjunto de pruebas: incluye preguntas y respuestas de los resultados de la base de datos de pruebas.
• Ejemplos de preguntas a SQL: un conjunto con preguntas y traducción a SQL. En algunos ejemplos, los datos devueltos se incluyen para permitir hacer preguntas sobre los datos y no solo sobre el esquema.

El desarrollo de la aplicación se realiza agregando nuevas preguntas y actualizando los diferentes conjuntos de datos, como se muestra en la siguiente figura.

Figura 5: Agregar una nueva pregunta

Se realiza un seguimiento de los conjuntos de datos y otras actualizaciones de parámetros como parte de la adición de nuevas preguntas y el ajuste de la aplicación. Utilizamos una herramienta de seguimiento para rastrear información sobre los experimentos, como:

• Parámetros como el número de preguntas, el número de ejemplos, el tipo de LLM, el método de búsqueda RAG
• Métricas como la precisión y la tasa de errores SQL
• Artefactos como una lista de resultados incorrectos que incluyen SQL generado, datos devueltos y más

Figura 6: Flujo del experimento

Utilizando una herramienta de seguimiento, pudimos avanzar comparando experimentos. La siguiente figura muestra las métricas de precisión y tasa de error de los diferentes experimentos que hicimos:

Figura 7: Precisión y tasa de error a lo largo del tiempo

Cuando hay un error o un error, se realiza un análisis de los resultados falsos y los detalles del experimento para comprender el origen del error y solucionarlo.

Experimente e implemente con Amazon Bedrock

Amazon Bedrock es un servicio administrado que ofrece una variedad de modelos de base de alto rendimiento. Puede experimentar y evaluar los principales FM para su caso de uso y personalizarlos con sus datos.

Al utilizar Amazon Bedrock, pudimos cambiar fácilmente entre modelos y opciones de integración. El siguiente es un código de ejemplo que utiliza la biblioteca Python de LangChain, que permite utilizar diferentes modelos e incrustaciones:

import boto3
from langchain_community.llms.bedrock import Bedrock
from langchain_community.embeddings import BedrockEmbeddings

def get_llm(model_id: str, args: dict):
   return Bedrock(model_id=model_id,
                  model_kwargs=args,
                  client=boto3.client("bedrock-runtime"))

def get_embeddings(model_id: str):
   return BedrockEmbeddings(model_id=model_id, 
                            client=boto3.client("bedrock-runtime"))

Conclusión

Utilizamos el mismo enfoque utilizado en proyectos de ciencia de datos para construir consultas SQL a partir de lenguaje natural. La solución mostrada se puede aplicar a otras aplicaciones basadas en LLM, y no solo para construir SQL. Por ejemplo, se puede utilizar para acceder a API, crear datos JSON y más. La clave es crear un conjunto de pruebas junto con resultados medibles y progreso mediante la experimentación.

Amazon Bedrock le permite usar diferentes modelos y alternar entre ellos para encontrar el adecuado para su caso de uso. Puede comparar diferentes modelos, incluidos los pequeños, para obtener un mejor rendimiento y costos. Como Amazon Bedrock no tiene servidor, no es necesario administrar ninguna infraestructura. Pudimos probar varios modelos rápidamente y finalmente integrar e implementar capacidades de IA generativa en nuestra aplicación.

Puede comenzar a experimentar con lenguaje natural para SQL ejecutando los ejemplos de código en este repositorio de GitHub. Este taller está dividido en módulos, cada uno de los cuales se basa en el anterior e introduce una nueva técnica para resolver este problema. Muchos de estos enfoques se basan en un trabajo existente de la comunidad y se citan en consecuencia.

Sobre los autores

Ori Nakar es investigador principal de ciberseguridad, ingeniero de datos y científico de datos en el grupo Imperva Threat Research.

Eitan Sela es arquitecto de soluciones especializado en inteligencia artificial generativa y aprendizaje automático en AWS. Trabaja con clientes de AWS para brindarles orientación y asistencia técnica, ayudándolos a crear y operar soluciones de inteligencia artificial generativa y aprendizaje automático en AWS. En su tiempo libre, a Eitan le gusta correr y leer los últimos artículos sobre aprendizaje automático.

Elad Eizner es arquitecto de soluciones en Amazon Web Services. Trabaja con clientes empresariales de AWS para ayudarlos a diseñar y crear soluciones en la nube y lograr sus objetivos.