Wren AI Text-to-SQL: API — the good stuff

Wren AI is an AI pipeline application that translates natural language user queries to SQL also known as Text-to-SQL (also known as Natural Language to SQL or NL2SQL), so you can chat with your database. It features both a UI and an API, allowing you to generate SQL queries based on user prompts and retrieve data using those queries. This significantly simplifies working with structured tabular data in your RAG/Agent applications.

Take a look at their demo https://demo.getwren.ai/

For more details on AI RAG and Agent pipeline applications, please see my previous article.

Evolution of the AI Bots: Harnessing the Power of Agents, RAG, and LLM Models

How Wren AI Works:

Deploying Wren AI is straightforward using the official Wren AI Docker Compose locally. You will need an OpenAI API key and to select the model you plan to use. Alternatively, you can use other OpenAI API-compatible inference engines like LocalAI by setting the OPENAI_API_BASE environment variable for the container. The Wren AI team has also implemented Ollama inference and recommends using Llama3 70b-instruct model. To utilize other than OpenAI models, you’ll need to create ~/.wrenai/.env.ai file following this example. Supported databases currently include BigQuery, DuckDB, PostgreSQL, MySQL, MS SQL, Clickhouse, and even Microsoft Excel Add-in. You can vote for additional database support [here].

## Example: Ollama inference
LLM_PROVIDER=ollama

## ollama. URL should be reachable from Docker Container!!!
OLLAMA_URL=http://host.docker.internal:11434
## https://ollama.com/library/llama3:70b-instruct-q8_0
GENERATION_MODEL=llama3:70b-instruct-q8_0
## supported embedding models providers by qdrant: https://qdrant.tech/documentation/embeddings/
## https://ollama.com/library/mxbai-embed-large:335m-v1-fp16
EMBEDDING_MODEL=mxbai-embed-large:335m-v1-fp16
EMBEDDING_MODEL_DIMENSION=1024
## DOCUMENT_STORE
DOCUMENT_STORE_PROVIDER=qdrant
QDRANT_HOST=qdrant

For Kubernetes deployment, refer to my article here:

WrenAI in Kubernetes: Text-to-SQL

Upon completing installing the Wren AI app you can access it via the URL.

Wren AI: Text-to-SQL. Connect a data source: BigQuery, DuckDB, PostgreSQL, MySQL

Once connected to your database or after uploading a sample dataset, such as the NBA playground, you’ll enter the home page. If you’d like to see the data in the sample dataset you can find it here just replace the version in the URL with the DATASET_VERSION variable listed at the same location.

Configuring your schema and adding descriptive text is crucial in the Modeling menu.

Modeling aids both you and your Large Language Model (LLM) in understanding the data and metadata, resulting in improved query outcomes. To provide better results you must setup relations and provide descriptions since API names of the columns in your tables might not be enough.

With your schema with relations configured and descriptions added in the Modeling menu, head to the Home menu to generate your first Text-to-SQL queries. The system attempts to produce three different SQL queries for you, with the results visible.

The good stuff: API

Moreover, you can access the same functionality via an API, making it an ideal platform for integrating structured tabular data into your RAG/Agent pipeline applications simply by consuming Wren AI as a service. Unfortunately, currently, there’s no mechanism for sending real-time updates from Wren AI back to the client such as WebSocket, Server-Sent-Events, or WebHooks. You’ll have to use Long Polling initiated by the client to constantly check for updates if Wren AI gets the answers ready for you, but the team is working on it.

While there isn’t currently public API documentation available, you can explore the functionality through the following steps by utilizing the Chrome Dev Tools -> Network to observe API interactions.

1. Create an asynchronous task by submitting a question to AI using the `createAskingTask` mutation in GraphQL.

2. Poll the task’s status change using the `askingTask` query.

3. Upon receiving a `FINISHED` status, retrieve three candidates in the payload.

4. Call the `createThread` mutation with the payload obtained.

5. A thread comprises multiple responses, each containing a question and an answer. Query a thread with the `thread` query.

6. The `responses` field in each thread contains an array of responses. Subsequent questions append to this array.

7. Continuously poll the `threadResponse` API to monitor status changes.

8. Upon the status turning `FINISHED`, view the detailed answer in the `detail` field.

9. The `detail` field provides a step-by-step answer seen in the UI.

Currently, Wren AI doesn’t offer a direct way to download or export results with data. As a workaround, consider copying the native SQL, executable within your database, and perform exporting from the database itself.

Semantic Layer

By the way, Wren AI is a good example of a RAG application pipeline built using Semantic Layer, where instead of a static instruction that is fed as a context to the LLM, the instruction is dynamically inserted from a vector database. It is very simple to check: running the application locally, there will be a container with a Qdrant vector database, if you open port 6333 and go here http://localhost:6333/dashboard#/collections/Document in the vector database, dynamic prompts for LLM created for the selected database with its fields and descriptions from “Modeling” page will be listed. Semantic Layer allows the app to get better quality output SQL from LLM and reduce hallucinations.

/* {"description":"A model containing information about review of orders.","alias":"reviews"} */ CREATE TABLE reviews ( -- {"description":"The date when the answer was provided.","alias":"AnswerTimestamp"} AnswerTimestamp TIMESTAMP, -- {"description":"The date when the review was created.","alias":"CreationTimestamp"} CreationTimestamp TIMESTAMP, -- {"description":"A unique identifier assigned to each review entry.","alias":"Id"} Id VARCHAR PRIMARY KEY, -- {"description":"The order id of the order which the review belongs to.","alias":"OrderId"} OrderId VARCHAR, -- {"description":"The score associated with each review entry.","alias":"Score"} Score BIGINT, -- {"condition": "orders".OrderId = "reviews".OrderId, "joinType": ONE_TO_MANY} FOREIGN KEY (OrderId) REFERENCES orders(OrderId) );

Semantic Layer: Dynamic LLM instructions in a Qdrant Database, the Document collection of the Wren AI app.

Summary

Wren AI platform makes it extremely easy for you or your RAG/Agent AI application to work with structured data in your Relational Databases, converting user prompts into valid SQL. UI with a modeling menu is important to provide additional metadata descriptions and schemas for your LLM to better understand your database and its data. Once modeling is done, Wren AI is ready to serve as a middleman in front of your DB. In short, you just speak to your database. This is unheard-of simplicity unachievable before, which makes work with DB highly effective in translating business language into SQL queries directly.

Enjoyed This Story?

If you like this article and you want to support me:

1. Clap 👏 my article 10 times; that will help me out
2. Follow me on Medium to get my latest articles 🫶
3. Follow Wren AI on Medium
4. Share this article on social media ➡️🌐
5. Give me feedback in the comments 💬 below. It’ll help me to better understand that this work was useful, even a simple “thanks” or “+” will do. Give me good, give me bad, whatever you think as long as you tell me place to improve and how.
6. Connect with me or follow me on LinkedIn.