create_sql_query_chain alternative with a custom database access layer

[HemanthVikash]

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

# The chain
sql_gen = create_sql_query_chain(
            llm=llm,
            db=self.snowflake_engine,
            prompt = self.construct_system_prompt(),
        )


# The proposed solution
def get_dialect_and_table_info(self, input):
  # API call that returns dialect and table_info based on input
  pass

Description

I have this piece of code that is my sql generation chain. The system prompt expects 'input', 'history', 'dialect', and 'table_info'.

I am trying to completely get rid of the db=self.snowflake_engine line and replace it with a secure local api call that doesn't have complete access to the database.

With that in mind, how do I properly import the dialect and table_info variables into the system prompt. Ideally, a ChatPromptTemplate message will take in the input. Right now this is my final_prompt:

final_prompt = ChatPromptTemplate.from_messages(
            [
                ('system', system_prompt_template),
                few_shot_prompt,
                MessagesPlaceholder(variable_name="history"),
                ('user', '{input}'),
            ],
        )

Just like few_shot_prompt, I would like to add a message that can include dialect and table_info.

Is this the right approach or am I thinking about it entirely wrong? If there is no way to pass this through the system prompt, I could always just send the information right when I invoke. But I would rather do it the right way if there is one for this kind of problem.

System Info

System Information

OS: Darwin
OS Version: Darwin Kernel Version 23.5.0: Wed May 1 20:13:18 PDT 2024; root:xnu-10063.121.3~5/RELEASE_ARM64_T6030
Python Version: 3.11.9 (main, Apr 19 2024, 11:43:47) [Clang 14.0.6 ]

Package Information

langchain_core: 0.2.10
langchain: 0.2.6
langchain_community: 0.2.5
langsmith: 0.1.82
langchain_aws: 0.1.8
langchain_openai: 0.1.10
langchain_text_splitters: 0.2.1

Packages not installed (Not Necessarily a Problem)

The following packages were not found:

langgraph
langserve

[dosubot[bot]]

To properly import the dialect and table_info variables into the system prompt using a ChatPromptTemplate message that takes in the input, you can use the ChatPromptTemplate class from langchain_core.prompts. Here is an example of how to create such a prompt:

from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessage

PROMPT_SUFFIX = """Only use the following tables:
{table_info}

Question: {input}"""

_VECTOR_SQL_DEFAULT_TEMPLATE = """You are a {dialect} expert. Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer to the input question.
{dialect} queries has a vector distance function called `DISTANCE(column, array)` to compute relevance to the user's question and sort the feature array column by the relevance. 
When the query is asking for {top_k} closest row, you have to use this distance function to calculate distance to entity's array on vector column and order by the distance to retrieve relevant rows.

*NOTICE*: `DISTANCE(column, array)` only accept an array column as its first argument and a `NeuralArray(entity)` as its second argument. You also need a user defined function called `NeuralArray(entity)` to retrieve the entity's array. 

Unless the user specifies in the question a specific number of examples to obtain, query for at most {top_k} results using the LIMIT clause as per {dialect}. You should only order according to the distance function.
Never query for all columns from a table. You must query only the columns that are needed to answer the question. Wrap each column name in double quotes (") to denote them as delimited identifiers.
Pay attention to use only the column names you can see in the tables below. Be careful to not query for columns that do not exist. Also, pay attention to which column is in which table.
Pay attention to use today() function to get the current date, if the question involves "today". `ORDER BY` clause should always be after `WHERE` clause. DO NOT add semicolon to the end of SQL. Pay attention to the comment in table schema.

Use the following format:

Question: "Question here"
SQLQuery: "SQL Query to run"
SQLResult: "Result of the SQLQuery"
Answer: "Final answer here"
"""

system_message = SystemMessage(content=_VECTOR_SQL_DEFAULT_TEMPLATE + PROMPT_SUFFIX)

human_prompt = HumanMessagePromptTemplate(
    prompt=PromptTemplate(
        template="Based on the following example, extract entities and relations from the provided text.\n\n"
                 "Use the following entity types, don't use other entity that is not defined below:\n"
                 "# ENTITY TYPES:\n{node_labels}\n\n"
                 "Use the following relation types, don't use other relation that is not defined below:\n"
                 "# RELATION TYPES:\n{rel_types}\n\n"
                 "Below are a number of examples of text and their extracted entities and relationships.\n"
                 "{examples}\n\n"
                 "For the following text, extract entities and relations as in the provided example.\n"
                 "{format_instructions}\nText: {input}",
        input_variables=["input"],
        partial_variables={
            "format_instructions": "Format instructions here",
            "node_labels": "Node labels here",
            "rel_types": "Relation types here",
            "examples": "Examples here",
        },
    )
)

VECTOR_SQL_PROMPT = ChatPromptTemplate.from_messages([system_message, human_prompt])

This template includes placeholders for dialect, table_info, input, and top_k, which will be filled in when the prompt is used ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.