Introducing Function RAG for SQL Generation

Vanna.ai is excited to introduce a groundbreaking feature in SQL generation: Function RAG (Retrieval-Augmented Generation). This new optional experimental feature is designed to enhance the consistency and determinism of SQL generation, addressing a multitude of requests from our users.

What is Function RAG ?

Function RAG transforms traditional Question-SQL training pairs into invokable templates (aka functions/tools). These templates are then invoked by the Large Language Model (LLM) to generate SQL queries and any relevant post-processing code like chart code. Critically, when using Function RAG, the LLM is only deciding on the SQL template to use and any parameters that need to be given. This method not only ensures more consistent outputs but also significantly speeds up the SQL generation process.

We are releasing Function RAG API under a different domain and branding to distinguish it from the core open-source Python package.

Key Features of Function RAG

Template-Based SQL Generation : By converting training pairs into templates, Function RAG ensures that the SQL generated is both accurate and relevant to the user's query.
Enhanced Security : Function RAG reduces prompt injection and prompt escape, ensuring that the SQL generation process is secure from external manipulations.
User-Specific Queries : Users can now pass information like a user ID within their queries without the risk of this data being overridden. This allows for personalized questions such as "what are my last 10 orders?" to be executed securely and efficiently.
Integrated Chart Code Generation : Alongside SQL, Function RAG also handles the generation of chart code, making it possible to generate both SQL and its corresponding visualization code in a single request.
Multi-Language Support : The feature is accessible via a GraphQL API, making it usable from any programming language, not just Python. This opens up possibilities for integration into various backends, including frameworks like Ruby on Rails, .NET and others.

See it in action

As an example, these are some sample functions:

When you ask a question, vn.get_function(question=...) will find the most appropriate function and use the LLM to fill in the necessary parameters.

If you want to scope a question to a specific user, you can use the vn.get_function(question=..., additional_data={"user_id": ...}) method. This will ensure that the user_id parameter is set deterministically.

To create a new function, you can use the vn.create_function(...) method if you're doing it manually, or you can use the built-in web app to do it.

Manual creation

vn.create_function(question=..., sql=..., plotly_code=...)

Web app function creation

The web app has a new button that allows you to select whether you want to train the results as a raw Question-SQL pair or whether you want to train the results as a function. Function RAG will automatically extract things that look like parameters from the question and make them available as function parameters.

When to use it

If you have end users who ask variants of similar questions, and you only want users to be able to run specific types of analyses that have been manually approved by your engineering team, then Function RAG is a great way to ensure that users can only run approved analyses.

Restricting to approved analyses helps end users have confidence that the data that they're seeing is correct and that it's the right way to view the data. This enables you to put this in the hands of internal business users or end-users if you're integrating this into a SaaS app.

Since Function RAG offers speed improvements, any user can benefit from it. However, for data analysts, you will likely want to fall back to the rest of Vanna functionality for SQL generation when you're performing a lot of novel analysis that doesn't exist in the training data. The fallback is automatically part of the built-in web app with the parameter function_generation=True if you're using the Vanna-hosted vector database.

GraphQL API Integration

In addition to the Python package, Function RAG is callable using a GraphQL API, allowing it to be used in other languages and frameworks. The most common request we've had is to be able to use it in Ruby on Rails.

Availability

Function RAG is available as part of the Vanna Free, Premium, and Enterprise plans . You can access it in Vanna v0.6.0 and later.

In the Vanna Package

The Vanna package has a new VannaAdvanced abstract class with an equivalent implementation here .

Future Plans

This functionality will enable new user interfaces. The parameters to the functions are set by the user's question and in the future we will eventually enable the user to see a full dashboard widget where they can modify the parameters via dropdowns without having to re-ask the question.

Additionally we will be adding additional functionality around detecting the available values for parameters so that they are always up to date and the LLM knows what the available values are in order to generate the answer.