How to pass custom prompt to create_sql_query_chain?

[Mourdhwaj]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

Hi Team,

I'm facing an issue with the create_sql_query_chain function, similar to the error described in this discussion https://github.com/langchain-ai/langchain/discussions/26966?sort=new 

but I didn't find a proper solution there.
Below is my code 

##few shot examples

examples = [
    {
        "input": "Give me records of items stuck in carton error",
        "query": "SELECT  * FROM t_order t (NOLOCK) WHERE status = 'CRTNERROR' ORDER BY t.imported_date DESC"
    },
    {
        "input": "Orders with no t_pick_detail_cartonize record",
        "query": "SELECT * FROM t_order ord (NOLOCK) LEFT JOIN t_pick_detail_cartonize pdc (NOLOCK) ON ord.order_number = pdc.order_number WHERE imported_date < DATEADD(mi, -5, GETDATE()) AND pdc.order_number IS NULL"
    },
    {
        "input": "Give me records where picked inventory is zero",
        "query": "SELECT * FROM t_stored_item sto (NOLOCK) WHERE actual_qty = 0 AND type <> 0"
    },
]

vectorstore = Chroma()
vectorstore.delete_collection()

example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    HuggingFaceEndpointEmbeddings(),
    vectorstore,
    k=2,
    input_keys=["input"],
)

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, FewShotChatMessagePromptTemplate, PromptTemplate
example_selector.select_examples({"input": "Give me LPN details"})

example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}\nSQLQuery:"),
        ("ai", "{query}"),
    ]
)

few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    example_selector=example_selector,
    input_variables=["input"],
)

print(few_shot_prompt.format(input="Give me employee details which is logged into device"))

final_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "You are MSSQL Expert. Given the input question, create a syntactically correct MSSQL query to execute. Unless otherwise specified.\nHere is the relevant table info: {table_info}\n\nBelow are number of questions and their corresponding SQL queries"),
        few_shot_prompt,
        ("human", "{input}"),
    ]
)

print(final_prompt.format(input="Give me employee details which are logged into devices", table_info="Show some table info"))

from langchain.chains import create_sql_query_chain

# Assuming llm is your language model, db is your database, and final_prompt is your prompt template
generate_query = create_sql_query_chain(llm, db, final_prompt)

Description

I am getting below error for my case. When i use default prompt it automatically fetch table info from my DB but when i am customizing it, i am getting error.

ValueError: Prompt must have input variables: 'input', 'top_k', 'table_info'. Received prompt with input variables: ['input', 'table_info']

can anyone help me?

System Info

input_variables=['input', 'table_info'] input_types={} partial_variables={} messages=[SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=['table_info'], input_types={}, partial_variables={}, template='You are MSSQL Expert. Given the input question, create a syntactically correct MSSQL query to execute. Unless otherwise specified.\nHere is the relevant table info :{table_info}\n\nBelow are number of questions and thier correspoding SQL queries'), additional_kwargs={}), FewShotChatMessagePromptTemplate(example_selector=SemanticSimilarityExampleSelector(vectorstore=<langchain_community.vectorstores.chroma.Chroma object at 0x000001F74899DAF0>, k=2, example_keys=None, input_keys=['input'], vectorstore_kwargs=None), input_variables=['input'], input_types={}, partial_variables={}, example_prompt=ChatPromptTemplate(input_variables=['input', 'query'], input_types={}, partial_variables={}, messages=[HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], input_types={}, partial_variables={}, template='{input}\nSQLQuery:'), additional_kwargs={}), AIMessagePromptTemplate(prompt=PromptTemplate(input_variables=['query'], input_types={}, partial_variables={}, template='{query}'), additional_kwargs={})])), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], input_types={}, partial_variables={}, template='{input}'), additional_kwargs={})]