Just the Gist: Snowflake Cortex LLM with Langchain LLM
Dated: Jan-2024
Snowflake Cortex is a fully managed service from Snowflake, which gives you instant access to industry-leading generative large language models (LLMs), including Meta’s LLaMA 2 model. It also offers models that Snowflake has fine-tuned for specific use cases.
Of the various different functions, one function that is getting noticed is the “Complete.” With the “Complete” function, users can pass a prompt select the LLM they want to use and get the response. For the private preview, users will be able to choose between the three model sizes (7B, 13B, and 70B) of Llama 2.
Langchain offers a rich set of LLM integrations, using which LLM applications are currently built. Currently, there is no implementation for Snowflake Cortex yet. While I do believe a community integration would be developed.
For now, here is the Gist of a sample prototype implementation, that you could adopt and get started on the Snowflake + Langchain journey.
Snowflake Cortex LLM
Langchain offers the ability to implement your own LLM, as explained in this doc: Custom LLM.
from typing import Any, List, Mapping, Optional
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.language_models.llms import LLM
class SnowflakeCortexLLM(LLM):
sp_session: Session = None
"""Snowpark session. It is assumed database, role, warehouse etc.. are set before invoking the LLM"""
model: str = 'llama2-7b-chat'
'''The Snowflake cortex hosted LLM model name. Defaulted to :llama2-7b-chat. Refer to doc for other options. '''
cortex_function: str = 'complete'
'''The cortex function to use, defaulted to complete. for other types refer to doc'''
@property
def _llm_type(self) -> str:
return "snowflake_cortex"
def _call(
self,
prompt: str,
stop: Optional[List[str]] = None,
run_manager: Optional[CallbackManagerForLLMRun] = None,
**kwargs: Any,
) -> str:
prompt_text = prompt
sql_stmt = f'''
select snowflake.cortex.{self.cortex_function}(
'{self.model}'
,'{prompt_text}') as llm_reponse;'''
l_rows = self.sp_session.sql(sql_stmt).collect()
llm_response = l_rows[0]['LLM_REPONSE']
return llm_response
@property
def _identifying_params(self) -> Mapping[str, Any]:
"""Get the identifying parameters."""
return {
"model": self.model
,"cortex_function" : self.cortex_function
,"snowpark_session": self.sp_session.session_id
}To interact with Snowflake Cortex you would need to establish a session. In the code above, the LLM implementation class is initialized with a Snowpark Session. It is assumed that the session has already been context set with database, role, warehouse, etc.
Sample Execution
Here is a sample execution of using the LLM
# Example chain definition
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.llms.vllm import VLLMOpenAI
# Define the prompt
instruction_prompt = '''
You are an expert Snowflake data analyst, Your task is generate Snowflake compliant SQL statement that help in
retreiving data related to the user question {user_query} .
Understand the structure of the table and its columns, which can be found between the tags [df_structure] and [/df_structure].
[DF_STRUCTURE]
The table TBL_3W contains time series measurements recorded from oil wells, it has the columns:
- the columnn measured_at_t indicates the date and time when the measurements was made
- the columnn p_pdg indicates pressure at the pdg
- the columnn p_tpt indicates pressure at the tpt
- the columnn t_tpt tindicates temperature at the tpt
- the columnn p_mon_ckp indicates pressure upstream of the pck
- the columnn t_jus_ckp indicates temperature downstream of the pck
- the columnn p_jus_ckgl indicates pressure downstream of the glck
- the columnn t_jus_ckgl indicates temperature downstream of the glck
- the columnn qgl indicates gas lift flow rate and is measured in standard cubic metres per second
- the column well indicates the well identifier
[/DF_STRUCTURE]
'''
l_prompt = PromptTemplate.from_template(instruction_prompt)
# Define the chain
llm = SnowflakeCortexLLM(sp_session = sp_session)
chain = l_prompt | llm | StrOutputParser()
# process the user request
user_query = '''Get the average tpt pressure for the well WELL-00002, do not add any filters like dates or timestamp'''
llm_response = chain.invoke({'user_query' : user_query})
print(llm_response, sep='\n')And here is the output of the interaction:
So what's next…
With Snowflake Cortex, it will enable customers to develop solutions using the data and with Snowflake handling many of the complexities just around hosting an LLM, this functionality would become beneficial.
There is a lot more to come from Snowflake Cortex and also, equally, Langchain.
So Get Inspired -> Learn -> Develop -> Rock-on!!!
