| | from __future__ import annotations |
| |
|
| | from typing import TYPE_CHECKING, Any, Dict, List, Optional, TypedDict, Union |
| |
|
| | from langchain_core.language_models import BaseLanguageModel |
| | from langchain_core.output_parsers import StrOutputParser |
| | from langchain_core.prompts import BasePromptTemplate |
| | from langchain_core.runnables import Runnable, RunnablePassthrough |
| |
|
| | from langchain.chains.sql_database.prompt import PROMPT, SQL_PROMPTS |
| |
|
| | if TYPE_CHECKING: |
| | from langchain_community.utilities.sql_database import SQLDatabase |
| |
|
| |
|
| | def _strip(text: str) -> str: |
| | return text.strip() |
| |
|
| |
|
| | class SQLInput(TypedDict): |
| | """Input for a SQL Chain.""" |
| |
|
| | question: str |
| |
|
| |
|
| | class SQLInputWithTables(TypedDict): |
| | """Input for a SQL Chain.""" |
| |
|
| | question: str |
| | table_names_to_use: List[str] |
| |
|
| |
|
| | def create_sql_query_chain( |
| | llm: BaseLanguageModel, |
| | db: SQLDatabase, |
| | prompt: Optional[BasePromptTemplate] = None, |
| | k: int = 5, |
| | ) -> Runnable[Union[SQLInput, SQLInputWithTables, Dict[str, Any]], str]: |
| | """Create a chain that generates SQL queries. |
| | |
| | *Security Note*: This chain generates SQL queries for the given database. |
| | |
| | The SQLDatabase class provides a get_table_info method that can be used |
| | to get column information as well as sample data from the table. |
| | |
| | To mitigate risk of leaking sensitive data, limit permissions |
| | to read and scope to the tables that are needed. |
| | |
| | Optionally, use the SQLInputWithTables input type to specify which tables |
| | are allowed to be accessed. |
| | |
| | Control access to who can submit requests to this chain. |
| | |
| | See https://python.langchain.com/docs/security for more information. |
| | |
| | Args: |
| | llm: The language model to use. |
| | db: The SQLDatabase to generate the query for. |
| | prompt: The prompt to use. If none is provided, will choose one |
| | based on dialect. Defaults to None. See Prompt section below for more. |
| | k: The number of results per select statement to return. Defaults to 5. |
| | |
| | Returns: |
| | A chain that takes in a question and generates a SQL query that answers |
| | that question. |
| | |
| | Example: |
| | |
| | .. code-block:: python |
| | |
| | # pip install -U langchain langchain-community langchain-openai |
| | from langchain_openai import ChatOpenAI |
| | from langchain.chains import create_sql_query_chain |
| | from langchain_community.utilities import SQLDatabase |
| | |
| | db = SQLDatabase.from_uri("sqlite:///Chinook.db") |
| | llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0) |
| | chain = create_sql_query_chain(llm, db) |
| | response = chain.invoke({"question": "How many employees are there"}) |
| | |
| | Prompt: |
| | If no prompt is provided, a default prompt is selected based on the SQLDatabase dialect. If one is provided, it must support input variables: |
| | * input: The user question plus suffix "\nSQLQuery: " is passed here. |
| | * top_k: The number of results per select statement (the `k` argument to |
| | this function) is passed in here. |
| | * table_info: Table definitions and sample rows are passed in here. If the |
| | user specifies "table_names_to_use" when invoking chain, only those |
| | will be included. Otherwise, all tables are included. |
| | * dialect (optional): If dialect input variable is in prompt, the db |
| | dialect will be passed in here. |
| | |
| | Here's an example prompt: |
| | |
| | .. code-block:: python |
| | |
| | from langchain_core.prompts import PromptTemplate |
| | |
| | template = '''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. |
| | Use the following format: |
| | |
| | Question: "Question here" |
| | SQLQuery: "SQL Query to run" |
| | SQLResult: "Result of the SQLQuery" |
| | Answer: "Final answer here" |
| | |
| | Only use the following tables: |
| | |
| | {table_info}. |
| | |
| | Question: {input}''' |
| | prompt = PromptTemplate.from_template(template) |
| | """ |
| | if prompt is not None: |
| | prompt_to_use = prompt |
| | elif db.dialect in SQL_PROMPTS: |
| | prompt_to_use = SQL_PROMPTS[db.dialect] |
| | else: |
| | prompt_to_use = PROMPT |
| | if {"input", "top_k", "table_info"}.difference( |
| | prompt_to_use.input_variables + list(prompt_to_use.partial_variables) |
| | ): |
| | raise ValueError( |
| | f"Prompt must have input variables: 'input', 'top_k', " |
| | f"'table_info'. Received prompt with input variables: " |
| | f"{prompt_to_use.input_variables}. Full prompt:\n\n{prompt_to_use}" |
| | ) |
| | if "dialect" in prompt_to_use.input_variables: |
| | prompt_to_use = prompt_to_use.partial(dialect=db.dialect) |
| |
|
| | inputs = { |
| | "input": lambda x: x["question"] + "\nSQLQuery: ", |
| | "table_info": lambda x: db.get_table_info( |
| | table_names=x.get("table_names_to_use") |
| | ), |
| | } |
| | return ( |
| | RunnablePassthrough.assign(**inputs) |
| | | ( |
| | lambda x: { |
| | k: v |
| | for k, v in x.items() |
| | if k not in ("question", "table_names_to_use") |
| | } |
| | ) |
| | | prompt_to_use.partial(top_k=str(k)) |
| | | llm.bind(stop=["\nSQLResult:"]) |
| | | StrOutputParser() |
| | | _strip |
| | ) |
| |
|
