Cree una solución sólida de texto a SQL que genere consultas complejas, se autocorrija y consulte diversas fuentes de datos

El lenguaje de consulta estructurado (SQL) es un lenguaje complejo que requiere comprensión de bases de datos y metadatos. Hoy en día, la IA generativa puede ayudar a personas sin conocimientos de SQL. Esta tarea generativa de IA se llama texto a SQL, que genera consultas SQL a partir del procesamiento del lenguaje natural (NLP) y convierte el texto en SQL semánticamente correcto. La solución de esta publicación tiene como objetivo llevar las operaciones de análisis empresarial al siguiente nivel al acortar el camino hacia sus datos utilizando lenguaje natural.

Con la aparición de grandes modelos de lenguaje (LLM), la generación de SQL basada en PNL ha experimentado una transformación significativa. Al demostrar un rendimiento excepcional, los LLM ahora son capaces de generar consultas SQL precisas a partir de descripciones en lenguaje natural. Sin embargo, aún persisten desafíos. En primer lugar, el lenguaje humano es inherentemente ambiguo y depende del contexto, mientras que SQL es preciso, matemático y estructurado. Esta brecha puede resultar en una conversión inexacta de las necesidades del usuario en el SQL que se genera. En segundo lugar, es posible que necesite crear funciones de texto a SQL para cada base de datos porque los datos a menudo no se almacenan en un único destino. Es posible que deba volver a crear la capacidad de cada base de datos para permitir a los usuarios la generación de SQL basada en NLP. En tercer lugar, a pesar de la mayor adopción de soluciones de análisis centralizadas, como lagos y almacenes de datos, la complejidad aumenta con los diferentes nombres de tablas y otros metadatos que se requieren para crear el SQL para las fuentes deseadas. Por lo tanto, la recopilación de metadatos completos y de alta calidad también sigue siendo un desafío. Para obtener más información sobre las mejores prácticas y patrones de diseño de texto a SQL, consulte Generación de valor a partir de datos empresariales: mejores prácticas para Text2SQL e IA generativa.

Nuestra solución tiene como objetivo abordar esos desafíos utilizando Amazon Bedrock y AWS Analytics Services. Usamos Anthropic Claude v2.1 en Amazon Bedrock como nuestro LLM. Para abordar los desafíos, nuestra solución primero incorpora los metadatos de las fuentes de datos dentro del catálogo de datos de AWS Glue para aumentar la precisión de la consulta SQL generada. El flujo de trabajo también incluye un ciclo de evaluación y corrección final, en caso de que Amazon Athena, que se utiliza en sentido descendente como motor SQL, identifique algún problema de SQL. Athena también nos permite utilizar una multitud de terminales y conectores compatibles para cubrir un gran conjunto de fuentes de datos.

Después de seguir los pasos para crear la solución, presentamos los resultados de algunos escenarios de prueba con distintos niveles de complejidad de SQL. Finalmente, analizamos cómo es sencillo incorporar diferentes fuentes de datos a sus consultas SQL.

Descripción general de la solución

Hay tres componentes críticos en nuestra arquitectura: recuperación de generación aumentada (RAG) con metadatos de base de datos, un bucle de autocorrección de varios pasos y Athena como nuestro motor SQL.

Usamos el método RAG para recuperar las descripciones de las tablas y las descripciones de los esquemas (columnas) del metastore de AWS Glue para garantizar que la solicitud esté relacionada con la tabla y los conjuntos de datos correctos. En nuestra solución, creamos los pasos individuales para ejecutar un marco RAG con el catálogo de datos de AWS Glue con fines de demostración. Sin embargo, también puede utilizar las bases de conocimientos de Amazon Bedrock para crear soluciones RAG rápidamente.

El componente de varios pasos permite al LLM corregir la precisión de la consulta SQL generada. Aquí, el SQL generado se envía en caso de errores de sintaxis. Usamos mensajes de error de Athena para enriquecer nuestro mensaje para el LLM y obtener correcciones más precisas y efectivas en el SQL generado.

Puede considerar los mensajes de error que ocasionalmente provienen de Athena como comentarios. Las implicaciones de costos de un paso de corrección de errores son insignificantes en comparación con el valor entregado. Incluso puede incluir estos pasos correctivos como ejemplos de aprendizaje reforzado supervisado para perfeccionar sus LLM. Sin embargo, no cubrimos este flujo en nuestra publicación por motivos de simplicidad.

Tenga en cuenta que siempre existe un riesgo inherente de tener imprecisiones, lo que naturalmente viene con las soluciones de IA generativa. Incluso si los mensajes de error de Athena son muy eficaces para mitigar este riesgo, puede agregar más controles y vistas, como comentarios humanos o consultas de ejemplo para realizar ajustes, para minimizar aún más dichos riesgos.

Athena no solo nos permite corregir las consultas SQL, sino que también nos simplifica el problema general porque sirve como centro, donde los radios son múltiples fuentes de datos. La gestión del acceso, la sintaxis SQL y más se gestionan a través de Athena.

El siguiente diagrama ilustra la arquitectura de la solución.

El flujo del proceso incluye los siguientes pasos:

1. Cree el catálogo de datos de AWS Glue mediante un rastreador de AWS Glue (o un método diferente).
2. Utilizando el modelo Titan-Text-Embeddings en Amazon Bedrock, convierta los metadatos en incrustaciones y guárdelos en un almacén de vectores sin servidor de Amazon OpenSearch, que sirve como nuestra base de conocimientos en nuestro marco RAG.

En esta etapa el proceso está listo para recibir la consulta en lenguaje natural. Los pasos 7 a 9 representan un ciclo de corrección, si corresponde.

1. El usuario ingresa su consulta en lenguaje natural. Puede utilizar cualquier aplicación web para proporcionar la interfaz de usuario del chat. Por lo tanto, no cubrimos los detalles de la interfaz de usuario en nuestra publicación.
2. La solución aplica un marco RAG mediante búsqueda por similitud, que agrega contexto adicional de los metadatos de la base de datos vectorial. Esta tabla se utiliza para encontrar la tabla, la base de datos y los atributos correctos.
3. La consulta se fusiona con el contexto y se envía a Anthropic Claude v2.1 en Amazon Bedrock.
4. El modelo obtiene la consulta SQL generada y se conecta a Athena para validar la sintaxis.
5. Si Athena proporciona un mensaje de error que menciona que la sintaxis es incorrecta, el modelo utiliza el texto de error de la respuesta de Athena.
6. El nuevo mensaje agrega la respuesta de Athena.
7. El modelo crea el SQL corregido y continúa el proceso. Esta iteración se puede realizar varias veces.
8. Finalmente, ejecutamos SQL usando Athena y generamos resultados. Aquí, el resultado se presenta al usuario. Por motivos de simplicidad arquitectónica, no mostramos este paso.

Requisitos previos

Para esta publicación, debe completar los siguientes requisitos previos:

1. Tener una cuenta de AWS.
2. Instale la interfaz de línea de comandos de AWS (AWS CLI).
3. Configurar el SDK para Python (Boto3).
4. Cree el catálogo de datos de AWS Glue mediante un rastreador de AWS Glue (o un método diferente).
5. Utilizando el modelo Titan-Text-Embeddings en Amazon Bedrock, convierta los metadatos en incrustaciones y guárdelos en un almacén de vectores OpenSearch Serverless.

Implementar la solución

Puedes usar lo siguiente cuaderno jupyter, que incluye todos los fragmentos de código proporcionados en esta sección, para crear la solución. Recomendamos utilizar Amazon SageMaker Studio para abrir este cuaderno con una instancia ml.t3.medium con el kernel Python 3 (Ciencia de datos). Para obtener instrucciones, consulte Entrenar un modelo de aprendizaje automático. Complete los siguientes pasos para configurar la solución:

1. Cree la base de conocimientos en OpenSearch Service para el marco RAG:

   def add_documnets(self,index_name: str,file_name:str):
   
   documents = JSONLoader(file_path=file_name, jq_schema=".", text_content=False, json_lines=False).load()
   docs = OpenSearchVectorSearch.from_documents(embedding=self.embeddings, opensearch_url=self.opensearch_domain_endpoint, http_auth=self.http_auth, documents=documents, index_name=index_name, engine="faiss")
   index_exists = self.check_if_index_exists(index_name,aws_region,opensearch_domain_endpoint,http_auth)
   if not index_exists :
   logger.info(f'index :{index_name} is not existing ')
   sys.exit(-1)
   else:
   logger.info(f'index :{index_name} Got created')

2. Construya el mensaje (final_question) combinando la entrada del usuario en lenguaje natural (user_query), los metadatos relevantes del almacén de vectores (vector_search_match), y nuestras instrucciones (details):

   def userinput(user_query):
   logger.info(f'Searching metadata from vector store')
   
   # vector_search_match=rqst.getEmbeddding(user_query)
   vector_search_match = rqst.getOpenSearchEmbedding(index_name,user_query)
   
   # print(vector_search_match)
   details = "It is important that the SQL query complies with Athena syntax. \
   During join if column name are same please use alias ex llm.customer_id \
   in select statement. It is also important to respect the type of columns: \
   if a column is string, the value should be enclosed in quotes. \
   If you are writing CTEs then include all the required columns. \
   While concatenating a non string column, make sure cast the column to string. \
   For date columns comparing to string , please cast the string input."
   final_question = "\n\nHuman:"+details + vector_search_match + user_query+ "n\nAssistant:"
   answer = rqst.generate_sql(final_question)
   return answer

3. Invoque Amazon Bedrock para el LLM (Claude v2) y solicite que genere la consulta SQL. En el siguiente código, realiza varios intentos para ilustrar el paso de autocorrección:x

   try:
   logger.info(f'we are in Try block to generate the sql and count is :{attempt + 1}')
   generated_sql = self.llm.predict(prompt)
   query_str = generated_sql.split("```")(1)
   query_str = " ".join(query_str.split("\n")).strip()
   sql_query = query_str(3:) if query_str.startswith("sql") else query_str
   
   # return sql_query
   syntaxcheckmsg=rqstath.syntax_checker(sql_query)
   if syntaxcheckmsg=='Passed':
   logger.info(f'syntax checked for query passed in attempt number :{attempt + 1}')
   return sql_query

4. Si se recibe algún problema con la consulta SQL generada ({sqlgenerated}) de la respuesta de Atenas ({syntaxcheckmsg}), el nuevo mensaje (prompt) se genera en función de la respuesta y el modelo intenta nuevamente generar el nuevo SQL:

   else:
   prompt = f"""{prompt} \
   This is syntax error: {syntaxcheckmsg}.
   To correct this, please generate an alternative SQL query which will correct the syntax error. The updated query should take care of all the syntax issues encountered. Follow the instructions mentioned above to remediate the error.
   Update the below SQL query to resolve the issue:
   {sqlgenerated}
   Make sure the updated SQL query aligns with the requirements provided in the initial question."""
   prompts.append(prompt)

5. Una vez generado el SQL, se invoca al cliente Athena para ejecutarlo y generar el resultado:

   query_execution = self.athena_client.start_query_execution(
   QueryString=query_string,
   ResultConfiguration=result_config,
   QueryExecutionContext=query_execution_context, )
   execution_id = query_execution("QueryExecutionId")

Prueba la solución

En esta sección, ejecutamos nuestra solución con diferentes escenarios de ejemplo para probar diferentes niveles de complejidad de consultas SQL.

Para probar nuestro texto a SQL, utilizamos dos conjuntos de datos disponibles en IMDB. Subconjuntos de datos de IMDb están disponibles para uso personal y no comercial. Puede descargar los conjuntos de datos y almacenarlos en Amazon Simple Storage Service (Amazon S3). Puede utilizar el siguiente fragmento de Spark SQL para crear tablas en AWS Glue. Para este ejemplo, utilizamos title_ratings y title:

source_title_ratings3_path="s3://llm-athena-output/input_data/title.ratings.tsv"
target_title_s3_path="s3://llm-athena-output/output_data/imdb_stg/title_ratings"
source_titleratingdf=spark.read.csv(source_title_ratings3_path,sep="\t",header=True)
source_titleratingdf.write.mode('overwrite').format('parquet').option('path', target_title_s3_path).saveAsTable('imdb_stg.title_ratings')

Almacene datos en Amazon S3 y metadatos en AWS Glue

En este escenario, nuestro conjunto de datos se almacena en un depósito de S3. Athena tiene un conector S3 que le permite utilizar Amazon S3 como fuente de datos que se puede consultar.

Para nuestra primera consulta, proporcionamos la entrada “Soy nuevo en esto. ¿Puedes ayudarme a ver todas las tablas y columnas en el esquema imdb?

La siguiente es la consulta generada:

WITH tables AS (
SELECT table_name
FROM information_schema.tables
WHERE table_schema="imdb_stg"),
columns AS (
SELECT
c.table_name,
c.column_name,
c.data_type,
c.is_nullable,
c.column_default,
c.ordinal_position
FROM information_schema.columns c
WHERE c.table_schema="imdb_stg")
SELECT
t.table_name,
c.column_name,
c.data_type,
c.is_nullable,
c.column_default,
c.ordinal_position
FROM tables t
INNER JOIN columns c
ON t.table_name = c.table_name
ORDER BY
t.table_name,
c.ordinal_position
LIMIT 10;

La siguiente captura de pantalla y el código muestran nuestro resultado.

table_name column_name data_type
0 title titleid varchar
1 title ordering integer
2 title title varchar
3 title region varchar
4 title language varchar

Para nuestra segunda consulta, preguntamos “Muéstreme todos los títulos y detalles en la región de EE. UU. cuya calificación sea superior a 9,5”.

La siguiente es nuestra consulta generada:

WITH us_titles AS (
SELECT t.title, t.region, tr.averageRating, tr.numVotes
FROM imdb_stg.title t
INNER JOIN imdb_stg.title_ratings tr
ON t.titleId = tr.tconst
WHERE t.region = 'US' AND cast(tr.averageRating as varchar) > '9.5'
)
SELECT title, region, averageRating, numVotes
FROM us_titles
LIMIT 100;

La respuesta es la siguiente.

title region averageRating numVotes
0 The Way You Saw Me US 9.7 8
1 The Brother Side of the Wake US 9.6 20
2 Ignis Fatuus US 9.6 11
3 Love and Hip Hop Atlanta US 9.9 11
4 ronny/lily US 9.7 14781

Para nuestra tercera consulta, ingresamos “¡Excelente respuesta! Ahora muéstrame todos los títulos de tipo original que tengan calificaciones superiores a 7,5 y que no estén en la región de EE. UU.”.

Se genera la siguiente consulta:

WITH titles AS (
SELECT t.titleId,
t.title,
t.types,
t.isOriginalTitle,
cast(tr.averageRating as decimal(3,1)) as averageRating,
tr.numVotes,
t.region
FROM imdb_stg.title t
INNER JOIN imdb_stg.title_ratings tr
ON t.titleId = tr.tconst
WHERE t.isOriginalTitle="1"
AND cast(tr.averageRating as decimal(3,1)) > 7.5
AND t.region != 'US')
SELECT *
FROM titles
LIMIT 100;

Obtenemos los siguientes resultados.

titleId title types isOriginalTitle averageRating numVotes region
0 tt0986264 Taare Zameen Par original 1 8.3 203760 XWW

Generar SQL autocorregido

Este escenario simula una consulta SQL que tiene problemas de sintaxis. Aquí, el SQL generado se autocorregirá según la respuesta de Athena. En la siguiente respuesta, Atenea dio una COLUMN_NOT_FOUND error y menciono que table_description no se puede resolver:

Status : {'State': 'FAILED', 'StateChangeReason': "COLUMN_NOT_FOUND: line 1:50: Column 'table_description' \
cannot be resolved or requester is not authorized to access requested resources",
'SubmissionDateTime': datetime.datetime(2024, 1, 14, 14, 38, 57, 501000, tzinfo=tzlocal()),
'CompletionDateTime': datetime.datetime(2024, 1, 14, 14, 38, 57, 778000, tzinfo=tzlocal()),
'AthenaError': {'ErrorCategory': 2, 'ErrorType': 1006, 'Retryable': False, 'ErrorMessage': "COLUMN_NOT_FOUND: \
line 1:50: Column 'table_description' cannot be resolved or requester is not authorized to \
access requested resources"}}
COLUMN_NOT_FOUND: line 1:50: Column 'table_description' cannot be resolved or requester is not authorized to access requested resources
Try Count: 2
2024-01-14 14:39:02,521,llm_execute,MainProcess,INFO,Try Count: 2
we are in Try block to generate the sql and count is :2
2024-01-14 14:39:02,521,llm_execute,MainProcess,INFO,we are in Try block to generate the sql and count is :2
Executing: Explain WITH tables AS ( SELECT table_name FROM information_schema.tables WHERE table_schema="imdb_stg" ), columns AS ( SELECT c.table_name, c.column_name, c.data_type, c.is_nullable, c.column_default, c.ordinal_position FROM information_schema.columns c WHERE c.table_schema="imdb_stg" ) SELECT t.table_name, c.column_name, c.data_type, c.is_nullable, c.column_default, c.ordinal_position FROM tables t INNER JOIN columns c ON t.table_name = c.table_name ORDER BY t.table_name, c.ordinal_position LIMIT 10;
I am checking the syntax here
execution_id: 904857c3-b7ac-47d0-8e7e-6b9d0456099b
Status : {'State': 'SUCCEEDED', 'SubmissionDateTime': datetime.datetime(2024, 1, 14, 14, 39, 29, 537000, tzinfo=tzlocal()), 'CompletionDateTime': datetime.datetime(2024, 1, 14, 14, 39, 30, 183000, tzinfo=tzlocal())}
syntax checked for query passed in tries number :2

Usar la solución con otras fuentes de datos

Para utilizar la solución con otras fuentes de datos, Athena se encarga del trabajo por usted. Para ello, Athena utiliza conectores de fuentes de datos que se pueden utilizar con consultas federadas. Puede considerar un conector como una extensión del motor de consultas de Athena. Existen conectores de fuentes de datos Athena prediseñados para fuentes de datos como Amazon CloudWatch Logs, Amazon DynamoDB, Amazon DocumentDB (con compatibilidad con MongoDB) y Amazon Relational Database Service (Amazon RDS), y fuentes de datos relacionales compatibles con JDBC como MySQL y PostgreSQL en la licencia Apache 2.0. Después de configurar una conexión a cualquier fuente de datos, puede utilizar el código base anterior para ampliar la solución. Para obtener más información, consulte Consultar cualquier origen de datos con la nueva consulta federada de Amazon Athena.

Limpiar

Para limpiar los recursos, puede comenzar limpiando su depósito S3 donde residen los datos. A menos que su aplicación invoque Amazon Bedrock, no generará ningún costo. En aras de las mejores prácticas de gestión de infraestructura, recomendamos eliminar los recursos creados en esta demostración.

Conclusión

En esta publicación, presentamos una solución que le permite utilizar PNL para generar consultas SQL complejas con una variedad de recursos habilitados por Athena. También aumentamos la precisión de las consultas SQL generadas a través de un ciclo de evaluación de varios pasos basado en mensajes de error de procesos posteriores. Además, utilizamos los metadatos del catálogo de datos de AWS Glue para considerar los nombres de las tablas solicitados en la consulta a través del marco RAG. Luego probamos la solución en varios escenarios realistas con diferentes niveles de complejidad de consultas. Finalmente, discutimos cómo aplicar esta solución a diferentes fuentes de datos respaldadas por Athena.

Amazon Bedrock está en el centro de esta solución. Amazon Bedrock puede ayudarle a crear muchas aplicaciones de IA generativa. Para comenzar con Amazon Bedrock, recomendamos seguir el inicio rápido a continuación repositorio de GitHub y familiarizarse con la creación de aplicaciones de IA generativa. También puede probar las bases de conocimientos en Amazon Bedrock para crear rápidamente este tipo de soluciones RAG.

Sobre los autores

Panda Sanjeeb es ingeniero de datos y aprendizaje automático en Amazon. Con experiencia en IA/ML, ciencia de datos y Big Data, Sanjeeb diseña y desarrolla datos innovadores y soluciones de ML que resuelven desafíos técnicos complejos y logran objetivos estratégicos para los vendedores 3P globales que administran sus negocios en Amazon. Además de su trabajo como ingeniero de datos y aprendizaje automático en Amazon, Sanjeeb Panda es un ávido entusiasta de la comida y la música.

Burak Gozluklu es un arquitecto principal de soluciones especializado en IA/ML ubicado en Boston, MA. Ayuda a clientes estratégicos a adoptar tecnologías de AWS y específicamente soluciones de IA generativa para lograr sus objetivos comerciales. Burak tiene un doctorado en Ingeniería Aeroespacial de METU, una maestría en Ingeniería de Sistemas y un postdoctorado en dinámica de sistemas del MIT en Cambridge, MA. Burak sigue siendo un afiliado de investigación en el MIT. Burak es un apasionado del yoga y la meditación.