Snowflake Cortex LLM: Self-Querying in Snowflake with Cortex AI

Exploring Self-Querying Retrievers in Snowflake Using LLMs and the latest Cortex AI functionality

Self-querying in a retrieval system. It shows different components involved in translating a natural language query into a search query that can be executed against a vector store or database. Source: LangChain https://python.langchain.com/docs/modules/data_connection/retrievers/self_query/

The concept of a self-querying retriever is gaining traction in the realm of natural language processing and database querying. Imagine a process that understands natural language queries, constructs structured queries, applies them to its underlying VectorStore or Databases, and answers users' questions with more specific information.

In this blogpost, we delve into the fascinating world of self-querying retrievers in Snowflake, leveraging the power of the latest Cortex AI functionality.

Understanding Self-Querying

At its core, a self-querying retriever possesses the unique ability to query itself. This means that when presented with a natural language query, the retriever utilizes a query-constructing LLM to create a structured query from the user’s input and then uses this information as metadata filters to extract the most relevant information only. Let's see an example.

Imagine that the user's query is as follows:

“Will the weather be suitable for sightseeing in New York this week?”.

To provide an answer, the retriever must first extract the location and date details from the user's query. This information will be used as context to determine if the weather conditions in New York next week are conducive to sightseeing. Something like this:

user_query = "Will the weather be suitable for sightseeing in New York this week?"
# some magic happens
parameters = {
"user_query": "Will the weather be suitable for sightseeing in New York this week?",
"location": "New York",
"date": "2024-05-23" # Next week
}

The process of extracting metadata from a user's initial query is called self-querying.

Let’s see how we can use the latest Snowflake Cortex AI to incorporate this functionality into your applications. But first, why Snowflake Cortex AI?

Snowflake Cortex AI

Snowflake Cortex is an intelligent, fully managed service designed to facilitate the quick analysis of data and the development of AI applications directly within the Snowflake ecosystem. Snowflake Cortex empowers users of varying skill levels to harness the capabilities of advanced artificial intelligence models, including large language models (LLMs) and vector search functionality. The platform offers a seamless integration of these technologies, enabling all users to access dynamic insights and enhance productivity by leveraging their enterprise data.

Advantages of Implementing Self-Querying with Snowflake Cortex AI

  1. Accessibility for Non-experts: Snowflake Cortex is designed to be user-friendly, allowing individuals without deep AI expertise to effectively use generative AI technologies.
  2. Serverless Functions: The platform provides a variety of serverless functions that simplify analytics and AI application development.
  3. Cost Efficiency: By offering LLM-based models and ML-based models as managed services, Snowflake Cortex reduces the need for costly GPU infrastructure.
  4. Versatility in AI Models: Snowflake Cortex includes models for a wide range of tasks from sentiment detection to forecasting and anomaly detection.
  5. Rapid Application Development: The platform enables quick development of powerful AI applications, like chatbots, through minimal coding, thereby reducing the development time from days or weeks to just minutes or hours.
  6. Built-in Advanced AI Capabilities: Snowflake Cortex includes state-of-the-art LLMs and advanced search capabilities, which are fully hosted and managed.
  7. Integration with Snowflake Ecosystem: The service is deeply integrated with the Snowflake environment, which ensures that users have a cohesive and seamless experience when handling data and developing applications within the same ecosystem.
  8. Continuous Innovation and Updates: Being in private preview, Snowflake Cortex is continually updated with the latest AI advancements and functionalities.
  9. Scalable AI Deployment: The platform's scalable architecture supports organizations' growing data and processing needs.

Implementation

The implementation process is relatively straightforward. The first step is to define the prompt to extract the location and time information from the user query.

self_query_prompt='''
Your goal is to extract the parameters `query`, `date`, and `location` from the user`s query and return a JSON structured response.
Apply the rules provided below.
<< Rules >>
Make sure the returned information is formatted as a JSON object.
Make sure that the JSON structured response contains the user`s query exactly as provided, without any modifications or additional content.
Make sure that no date parameter is returned in the JSON structured response if the user`s query does not reference any specific time element.
Make sure that the date parameter in the JSON structured response is formatted using the format YYYY-MM-DD.
Make sure that when the user`s query is analyzed, anytime today is referenced its value is {today_date}.

<< Example 1 >>
User Query: Hey buddy, I was just wondering if it would be right for me to hit the waves and go surfing at Dominical tomorrow. What do you think?
JSON Structured Response:
```json
{{
"query": "Hey buddy, I was just wondering if it would be alright for me to hit the waves and go surfing at Dominical tomorrow. What do you think?",
"location": "Dominical",
"date": "2024-03-22"
}}
```

User Query: {user_query}
JSON Structured Response:
'''

The prompt contains the following key components for self-querying and generating a JSON-structured response:

  1. Goal: The goal is to extract the parameters query, date, and location from the user's query and return a JSON structured response.
  2. Rules: Rules on how the user query should be processed. Today’s date is provided as a parameter to give enough information to the LLM to answer questions about relative time components (today, tomorrow, yesterday)
  3. Example: An example is provided, showing the expected JSON structured response for a given user query.
  4. User Query: The actual user query is provided as a variable, which should be analyzed to extract the required parameters and generate the JSON structured response according to the rules.

The second step is to provide the parameters to the prompt and use Snowflake Cortex AI to extract the answer.

from datetime import datetime

today_date = datetime.now().strftime("%B %d, %Y")
user_query = "Will the weather be suitable for sightseeing in New York this week?"
session = # Snowpark session

# Build the prompt
self_query_prompt = self_query_prompt.format(
user_query=user_query,
today_date=today_date
)

# Use Cortex AI to get the metadata parameters from the user query
# We are using Mistral 7B in this case, but Snowflake provides others models that
# you can easily try https://docs.snowflake.com/en/user-guide/snowflake-cortex/llm-functions#availability
metadata_query_model_call = f'''
select snowflake.cortex.complete(
'mistral-7b',
'{self_query_prompt}'
) as response;
'''
df_query = session.sql(metadata_query_model_call).to_pandas()
llm_response = df_query['RESPONSE'][0]

Note: Your Snowflake role must have access to SNOWFLAKE.CORTEX_USER for the code above to work. To enable Snowflake Cortex AI functionality, grant the following role.

GRANT DATABASE ROLE SNOWFLAKE.CORTEX_USER TO ROLE analyst;

The third step involves extracting the JSON object from the LLM response. We recommend using a regex pattern as shown below:

import re
import json

# Extract the JSON object from the response
pattern = r'{[\n\s\S\w\W:\d\D]+}'
match = re.search(pattern, llm_response, re.DOTALL)
if match:
query_response = json.loads(match.group())
else:
print('The model has not returned the response with the desired JSON schema')
print(llm_response)
print(query_response)

>>>
{
"user_query": "Will the weather be suitable for sightseeing in New York this week?",
"location": "New York",
"date": "2024-05-23" # Next week
}
<<<

Finally, with the parameters in the JSON object, we can query our weather table in the database to retrieve information for a specific location and time and provide context for another LLM to answer the question, as shown below

# Extract the weather conditions for the date and location
weather_condition = (
session.table(f'{Your Database}.{Your Schema}.WEATHER_DESCRIPTION')
.filter(col('DATE')==query_response['date'])
.filter(col('LOCATION')==query_response['location']))
.select(col('DESCRIPTION'))
.to_pandas()
)
weather_condition_description = weather_condition['DESCRIPTION'][0]

# Build the prompt used to answer the user's query
query_prompt='''
Your goal is to analyze the <<Weather Condition>>information and answer the user`s query.

<<Weather Condition>>
{weather_condition_description}

Apply the following rules to your answer.
<< Rules >>
-Make sure that the user`s query is analyzed, considering the current date is {today_date}.
-Make sure to use bullet points to justify your answer.
-Make sure to provide a summary of your answer at the end.

User Query: {query}
Answer:
'''
# Use the weather conditions as context
query_prompt = query_prompt.format(
weather_condition_description=weather_condition_description,
today_date=today_date,
query=query
)

# Call the Snowflake Cortex AI functionality
# We are using Mistral 7B in this case, but Snowflake provides others models that
# you can easily try https://docs.snowflake.com/en/user-guide/snowflake-cortex/llm-functions#availability
user_query_model_call = f'''
select snowflake.cortex.complete(
'{model}',
'{query_prompt}'
) as response;
'''
df_query = session.sql(user_query_model_call).to_pandas()
print(df_query['RESPONSE'][0])

With these simple steps, you can create a fully functional Q&A process capable of processing user queries that involve metadata parameters. Quite cool, it is not? Give me feedback on your thoughts!

Conclusion

In conclusion, Snowflake Cortex AI is revolutionizing how organizations interact with and leverage their data. By enabling self-querying capabilities, Snowflake Cortex not only simplifies the process of data retrieval and analysis but also democratizes access to powerful AI tools, making them accessible to users regardless of their technical expertise. Find out how Infostrux Solutions has ready-made solutions to help you implement this quicker.” in the word ready-made, add a link to this https://www.infostrux.com/ai-ml/secure-ai-powered-chatbots

I’m Fabian Hernandez, Senior Data Architect at Infostrux Solutions. Thanks for reading my blog post, and I hope you enjoy the content! You can follow me on LinkedIn and subscribe to Infostrux Medium Blogs for the most interesting Data Engineering and Snowflake news. Please let us know your thoughts about this approach in the comment section.

--

--