Harrison/sql db chain (#641)

Co-authored-by: Bruno Bornsztein <bruno.bornsztein@gmail.com>
This commit is contained in:
Harrison Chase 2023-01-17 22:32:28 -08:00 committed by GitHub
parent b550f57912
commit 6be5f4e4c4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 19 additions and 5 deletions

View File

@ -266,7 +266,7 @@
"metadata": {}, "metadata": {},
"outputs": [], "outputs": [],
"source": [ "source": [
"from langchain.chains.sql_database.base import SQLDatabaseSequentialChain" "from langchain.chains import SQLDatabaseSequentialChain"
] ]
}, },
{ {
@ -293,14 +293,22 @@
"\n", "\n",
"\u001b[1m> Entering new SQLDatabaseSequentialChain chain...\u001b[0m\n", "\u001b[1m> Entering new SQLDatabaseSequentialChain chain...\u001b[0m\n",
"Table names to use:\n", "Table names to use:\n",
"\u001b[33;1m\u001b[1;3m['Employee', 'Customer']\u001b[0m\n", "\u001b[33;1m\u001b[1;3m['Customer', 'Employee']\u001b[0m\n",
"\n",
"\u001b[1m> Entering new SQLDatabaseChain chain...\u001b[0m\n",
"How many employees are also customers? \n",
"SQLQuery:\u001b[32;1m\u001b[1;3m SELECT COUNT(*) FROM Customer c INNER JOIN Employee e ON c.SupportRepId = e.EmployeeId;\u001b[0m\n",
"SQLResult: \u001b[33;1m\u001b[1;3m[(59,)]\u001b[0m\n",
"Answer:\u001b[32;1m\u001b[1;3m There are 59 employees who are also customers.\u001b[0m\n",
"\u001b[1m> Finished chain.\u001b[0m\n",
"\n",
"\u001b[1m> Finished chain.\u001b[0m\n" "\u001b[1m> Finished chain.\u001b[0m\n"
] ]
}, },
{ {
"data": { "data": {
"text/plain": [ "text/plain": [
"' 0 employees are also customers.'" "' There are 59 employees who are also customers.'"
] ]
}, },
"execution_count": 5, "execution_count": 5,

View File

@ -12,7 +12,10 @@ from langchain.chains.pal.base import PALChain
from langchain.chains.qa_with_sources.base import QAWithSourcesChain from langchain.chains.qa_with_sources.base import QAWithSourcesChain
from langchain.chains.qa_with_sources.vector_db import VectorDBQAWithSourcesChain from langchain.chains.qa_with_sources.vector_db import VectorDBQAWithSourcesChain
from langchain.chains.sequential import SequentialChain, SimpleSequentialChain from langchain.chains.sequential import SequentialChain, SimpleSequentialChain
from langchain.chains.sql_database.base import SQLDatabaseChain from langchain.chains.sql_database.base import (
SQLDatabaseChain,
SQLDatabaseSequentialChain,
)
from langchain.chains.transform import TransformChain from langchain.chains.transform import TransformChain
from langchain.chains.vector_db_qa.base import VectorDBQA from langchain.chains.vector_db_qa.base import VectorDBQA
@ -35,4 +38,5 @@ __all__ = [
"TransformChain", "TransformChain",
"MapReduceChain", "MapReduceChain",
"OpenAIModerationChain", "OpenAIModerationChain",
"SQLDatabaseSequentialChain",
] ]

View File

@ -109,7 +109,9 @@ class SQLDatabaseSequentialChain(Chain, BaseModel):
**kwargs: Any, **kwargs: Any,
) -> SQLDatabaseSequentialChain: ) -> SQLDatabaseSequentialChain:
"""Load the necessary chains.""" """Load the necessary chains."""
sql_chain = SQLDatabaseChain(llm=llm, database=database, prompt=query_prompt) sql_chain = SQLDatabaseChain(
llm=llm, database=database, prompt=query_prompt, **kwargs
)
decider_chain = LLMChain( decider_chain = LLMChain(
llm=llm, prompt=decider_prompt, output_key="table_names" llm=llm, prompt=decider_prompt, output_key="table_names"
) )