!pip install --upgrade --quiet  langchain langchain-community langchain-experimental langchain-openai

import getpass
import os

def _set_if_undefined(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"Please provide your {var}")

_set_if_undefined("OPENAI_API_KEY")
_set_if_undefined("LANGCHAIN_API_KEY")

# Optional, add tracing in LangSmith
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "Database Agent"

from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db", sample_rows_in_table_info=3)

print(db.get_usable_table_names())

['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']

print(db.run("SELECT * FROM Artist LIMIT 10;"))

[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]

print(db.run("SELECT * FROM Album LIMIT 10;"))

[(1, 'For Those About To Rock We Salute You', 1), (2, 'Balls to the Wall', 2), (3, 'Restless and Wild', 2), (4, 'Let There Be Rock', 1), (5, 'Big Ones', 3), (6, 'Jagged Little Pill', 4), (7, 'Facelift', 5), (8, 'Warner 25 Anos', 6), (9, 'Plays Metallica By Four Cellos', 7), (10, 'Audioslave', 8)]

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-3.5-turbo")

agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)

print(agent_executor)

AgentExecutor(name='SQL Agent Executor', verbose=True, agent=RunnableMultiActionAgent(runnable=RunnableAssign(mapper={
  agent_scratchpad: RunnableLambda(lambda x: format_to_openai_tool_messages(x['intermediate_steps']))
})
| ChatPromptTemplate(input_variables=['agent_scratchpad', 'input'], input_types={'agent_scratchpad': typing.List[typing.Union[langchain_core.messages.ai.AIMessage, langchain_core.messages.human.HumanMessage, langchain_core.messages.chat.ChatMessage, langchain_core.messages.system.SystemMessage, langchain_core.messages.function.FunctionMessage, langchain_core.messages.tool.ToolMessage]]}, messages=[SystemMessage(content='You are an agent designed to interact with a SQL database.\nGiven an input question, create a syntactically correct sqlite query to run, then look at the results of the query and return the answer.\nUnless the user specifies a specific number of examples they wish to obtain, always limit your query to at most 10 results.\nYou can order the results by a relevant column to return the most interesting examples in the database.\nNever query for all the columns from a specific table, only ask for the relevant columns given the question.\nYou have access to tools for interacting with the database.\nOnly use the below tools. Only use the information returned by the below tools to construct your final answer.\nYou MUST double check your query before executing it. If you get an error while executing a query, rewrite the query and try again.\n\nDO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database.\n\nIf the question does not seem related to the database, just return "I don\'t know" as the answer.\n'), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], template='{input}')), AIMessage(content='I should look at the tables in the database to see what I can query.  Then I should query the schema of the most relevant tables.'), MessagesPlaceholder(variable_name='agent_scratchpad')])
| RunnableBinding(bound=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x7ffa4aa815d0>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x7ffa4aa853d0>, root_client=<openai.OpenAI object at 0x7ffa4b263ed0>, root_async_client=<openai.AsyncOpenAI object at 0x7ffa4aa81790>, temperature=0.0, openai_api_key=SecretStr('**********'), openai_proxy=''), kwargs={'tools': [{'type': 'function', 'function': {'name': 'sql_db_query', 'description': "Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.", 'parameters': {'type': 'object', 'properties': {'query': {'description': 'A detailed and correct SQL query.', 'type': 'string'}}, 'required': ['query']}}}, {'type': 'function', 'function': {'name': 'sql_db_schema', 'description': 'Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3', 'parameters': {'type': 'object', 'properties': {'table_names': {'description': "A comma-separated list of the table names for which to return the schema. Example input: 'table1, table2, table3'", 'type': 'string'}}, 'required': ['table_names']}}}, {'type': 'function', 'function': {'name': 'sql_db_list_tables', 'description': 'Input is an empty string, output is a comma-separated list of tables in the database.', 'parameters': {'type': 'object', 'properties': {'tool_input': {'description': 'An empty string', 'default': '', 'type': 'string'}}}}}, {'type': 'function', 'function': {'name': 'sql_db_query_checker', 'description': 'Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!', 'parameters': {'type': 'object', 'properties': {'query': {'description': 'A detailed and SQL query to be checked.', 'type': 'string'}}, 'required': ['query']}}}]})
| OpenAIToolsAgentOutputParser(), input_keys_arg=['input'], return_keys_arg=['output'], stream_runnable=True), tools=[QuerySQLDataBaseTool(description="Input to this tool is a detailed and correct SQL query, output is a result from the database. If the query is not correct, an error message will be returned. If an error is returned, rewrite the query, check the query, and try again. If you encounter an issue with Unknown column 'xxxx' in 'field list', use sql_db_schema to query the correct table fields.", db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x7ffa60b65350>), InfoSQLDatabaseTool(description='Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling sql_db_list_tables first! Example Input: table1, table2, table3', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x7ffa60b65350>), ListSQLDatabaseTool(db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x7ffa60b65350>), QuerySQLCheckerTool(description='Use this tool to double check if your query is correct before executing it. Always use this tool before executing a query with sql_db_query!', db=<langchain_community.utilities.sql_database.SQLDatabase object at 0x7ffa60b65350>, llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x7ffa4aa815d0>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x7ffa4aa853d0>, root_client=<openai.OpenAI object at 0x7ffa4b263ed0>, root_async_client=<openai.AsyncOpenAI object at 0x7ffa4aa81790>, temperature=0.0, openai_api_key=SecretStr('**********'), openai_proxy=''), llm_chain=LLMChain(prompt=PromptTemplate(input_variables=['dialect', 'query'], template='\n{query}\nDouble check the {dialect} query above for common mistakes, including:\n- Using NOT IN with NULL values\n- Using UNION when UNION ALL should have been used\n- Using BETWEEN for exclusive ranges\n- Data type mismatch in predicates\n- Properly quoting identifiers\n- Using the correct number of arguments for functions\n- Casting to the correct data type\n- Using the proper columns for joins\n\nIf there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.\n\nOutput the final SQL query only.\n\nSQL Query: '), llm=ChatOpenAI(client=<openai.resources.chat.completions.Completions object at 0x7ffa4aa815d0>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x7ffa4aa853d0>, root_client=<openai.OpenAI object at 0x7ffa4b263ed0>, root_async_client=<openai.AsyncOpenAI object at 0x7ffa4aa81790>, temperature=0.0, openai_api_key=SecretStr('**********'), openai_proxy='')))])

agent_executor.invoke("アーティストは何人登録されていますか？日本語で回答してください。")

> Entering new SQL Agent Executor chain...

Invoking: `sql_db_list_tables` with `{}`


Album, Artist, Customer, Employee, Genre, Invoice, InvoiceLine, MediaType, Playlist, PlaylistTrack, Track
Invoking: `sql_db_query` with `{'query': 'SELECT COUNT(*) AS total_artists FROM Artist'}`


[(275,)]データベースには275人のアーティストが登録されています。

> Finished chain.
{'input': 'アーティストは何人登録されていますか？日本語で回答してください。',
 'output': 'データベースには275人のアーティストが登録されています。'}

print(db.run("SELECT count(*) FROM artist;"))
[(275,)]

import getpass
import os

def _set_if_undefined(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"Please provide your {var}")

_set_if_undefined("OPENAI_API_KEY")
_set_if_undefined("LANGCHAIN_API_KEY")

# Optional, add tracing in LangSmith
os.environ["LANGCHAIN_TRACING_V2"] = "true"
os.environ["LANGCHAIN_PROJECT"] = "Database Agent"

from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db", sample_rows_in_table_info=3)

from langchain_community.agent_toolkits import create_sql_agent
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)
agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=True)

examples = [
    {"input": "すべてのアーティストをリストアップしてください。", "query": "SELECT * FROM Artist;"},
    {
        "input": "アーティスト 'AC/DC' のすべてのアルバムを見つけてください。",
        "query": "SELECT * FROM Album WHERE ArtistId = (SELECT ArtistId FROM Artist WHERE Name = 'AC/DC');",
    },
    {
        "input": "ジャンルが 'Rock' のすべてのトラックをリストアップしてください。",
        "query": "SELECT * FROM Track WHERE GenreId = (SELECT GenreId FROM Genre WHERE Name = 'Rock');",
    },
    {
        "input": "すべてのトラックの総再生時間を求めてください。",
        "query": "SELECT SUM(Milliseconds) FROM Track;",
    },
    {
        "input": "IDが5のアルバムにはトラックが何曲ありますか？",
        "query": "SELECT COUNT(*) FROM Track WHERE AlbumId = 5;",
    },
    {
        "input": "5分以上のトラックをリストアップしてください。",
        "query": "SELECT * FROM Track WHERE Milliseconds > 300000;",
    },
    {
        "input": "2000年のアルバムはどれですか？",
        "query": "SELECT * FROM Album WHERE strftime('%Y', ReleaseDate) = '2000';",
    },
]

from langchain_community.vectorstores import FAISS
from langchain_core.example_selectors import SemanticSimilarityExampleSelector
from langchain_openai import OpenAIEmbeddings

example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    OpenAIEmbeddings(),
    FAISS,
    k=3,
    input_keys=["input"],
)

from langchain_core.prompts import (
    ChatPromptTemplate,
    FewShotPromptTemplate,
    MessagesPlaceholder,
    PromptTemplate,
    SystemMessagePromptTemplate,
)

system_prefix = """
You are an agent designed to interact with a SQL database. Given an input question, create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Unless the user specifies a specific number of examples they wish to obtain, always limit your query to at most {top_k} results. You can order the results by a relevant column to return the most interesting examples in the database.
Never query for all the columns from a specific table, only ask for the relevant columns given the question.
You have access to tools for interacting with the database. Only use the given tools. Only use the information returned by the tools to construct your final answer.
You MUST double check your query before executing it.
If you get an error while executing a query, rewrite the query and try again.
DO NOT make any DML statements (INSERT, UPDATE, DELETE, DROP etc.) to the database. If the question does not seem related to the database, just return "I don't know" as the answer.
Here are some examples of user inputs and their corresponding SQL queries:
"""

few_shot_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=PromptTemplate.from_template(
        "User input: {input}\nSQL query: {query}"
    ),
    input_variables=["input", "dialect", "top_k"],
    prefix=system_prefix,
    suffix="",
)

full_prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate(prompt=few_shot_prompt),
        ("human", "{input}"),
        MessagesPlaceholder("agent_scratchpad"),
    ]
)

agent.invoke({"input": "10枚以上のアルバムをリリースしているアーティストの名前をリストしてください。"})

> Entering new SQL Agent Executor chain...

Invoking: `sql_db_query` with `{'query': 'SELECT Artist.Name FROM Artist JOIN Album ON Artist.ArtistId = Album.ArtistId GROUP BY Artist.Name HAVING COUNT(Album.AlbumId) >= 10'}`


[('Deep Purple',), ('Iron Maiden',), ('Led Zeppelin',), ('Metallica',), ('U2',)]10枚以上のアルバムをリリースしているアーティストの名前は、Deep Purple、Iron Maiden、Led Zeppelin、Metallica、U2です。

> Finished chain.
{'input': '10枚以上のアルバムをリリースしているアーティストの名前をリストしてください。',
 'output': '10枚以上のアルバムをリリースしているアーティストの名前は、Deep Purple、Iron Maiden、Led Zeppelin、Metallica、U2です。'}