mirror of
https://github.com/langchain-ai/docs.git
synced 2026-08-27 10:52:15 -04:00
161 lines
5.9 KiB
Plaintext
161 lines
5.9 KiB
Plaintext
```python
|
|
# sql_agent.py for studio
|
|
import pathlib
|
|
import sqlite3
|
|
|
|
import requests
|
|
from langchain.agents import create_agent
|
|
from langchain.chat_models import init_chat_model
|
|
from langchain.tools import tool
|
|
|
|
# Initialize an LLM
|
|
model = init_chat_model("gpt-5.5")
|
|
|
|
# Get the database, store it locally
|
|
url = "https://storage.googleapis.com/benchmarks-artifacts/chinook/Chinook.db"
|
|
local_path = pathlib.Path("Chinook.db")
|
|
|
|
if local_path.exists():
|
|
print(f"{local_path} already exists, skipping download.")
|
|
else:
|
|
response = requests.get(url, timeout=60)
|
|
if response.status_code == 200:
|
|
local_path.write_bytes(response.content)
|
|
print(f"File downloaded and saved as {local_path}")
|
|
else:
|
|
print(f"Failed to download the file. Status code: {response.status_code}")
|
|
|
|
# Below are minimal tools for demonstration purposes.
|
|
|
|
@tool
|
|
def sql_db_list_tables() -> str:
|
|
"""Input is an empty string, output is a comma-separated list of tables in the database."""
|
|
con = sqlite3.connect("Chinook.db")
|
|
try:
|
|
cursor = con.cursor()
|
|
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
|
|
tables = [row[0] for row in cursor.fetchall() if not row[0].startswith("sqlite_")]
|
|
return ", ".join(tables)
|
|
finally:
|
|
con.close()
|
|
|
|
@tool
|
|
def sql_db_schema(table_names: str) -> str:
|
|
"""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"""
|
|
con = sqlite3.connect("Chinook.db")
|
|
try:
|
|
cursor = con.cursor()
|
|
cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
|
|
valid_tables = {row[0] for row in cursor.fetchall() if not row[0].startswith("sqlite_")}
|
|
results = []
|
|
for table in table_names.split(","):
|
|
table = table.strip()
|
|
if table not in valid_tables:
|
|
results.append(f"Error: table_names {{{table!r}}} not found in database")
|
|
continue
|
|
cursor.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name=?;", (table,))
|
|
schema_row = cursor.fetchone()
|
|
if schema_row:
|
|
results.append(schema_row[0])
|
|
try:
|
|
quoted_table = '"' + table.replace('"', '""') + '"'
|
|
cursor.execute(f"SELECT * FROM {quoted_table} LIMIT 3;")
|
|
rows = cursor.fetchall()
|
|
if rows:
|
|
col_names = [description[0] for description in cursor.description]
|
|
results.append(
|
|
f"/*\n3 rows from {table} table:\n"
|
|
+ "\t".join(col_names)
|
|
+ "\n"
|
|
+ "\n".join("\t".join(str(x) for x in row) for row in rows)
|
|
+ "\n*/"
|
|
)
|
|
except Exception as e:
|
|
results.append(f"Error fetching sample rows: {e}")
|
|
return "\n\n".join(results)
|
|
finally:
|
|
con.close()
|
|
|
|
@tool
|
|
def sql_db_query(query: str) -> str:
|
|
"""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."""
|
|
con = sqlite3.connect("Chinook.db")
|
|
try:
|
|
cursor = con.cursor()
|
|
cursor.execute(query)
|
|
res = cursor.fetchall()
|
|
return str(res)
|
|
except Exception as e:
|
|
return f"Error: {e}"
|
|
finally:
|
|
con.close()
|
|
|
|
@tool
|
|
def sql_db_query_checker(query: str) -> str:
|
|
"""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!"""
|
|
trigger_prompt = """{query}
|
|
Double check the sqlite query above for common mistakes, including:
|
|
- Using NOT IN with NULL values
|
|
- Using UNION when UNION ALL should have been used
|
|
- Using BETWEEN for exclusive ranges
|
|
- Data type mismatch in predicates
|
|
- Properly quoting identifiers
|
|
- Using the correct number of arguments for functions
|
|
- Casting to the correct data type
|
|
- Using the proper columns for joins
|
|
|
|
If there are any of the above mistakes, rewrite the query. If there are no mistakes, just reproduce the original query.
|
|
|
|
Output the final SQL query only.
|
|
|
|
SQL Query: """.format(query=query)
|
|
|
|
response = model.invoke(trigger_prompt)
|
|
return response.text.strip()
|
|
|
|
tools = [sql_db_list_tables, sql_db_schema, sql_db_query, sql_db_query_checker]
|
|
|
|
# Use a distinct loop variable so it does not shadow the `tool` decorator.
|
|
for t in tools:
|
|
print(f"{t.name}: {t.description}\n")
|
|
|
|
# Use create_agent
|
|
system_prompt = """
|
|
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 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.
|
|
|
|
To start you should ALWAYS look at the tables in the database to see what you
|
|
can query. Do NOT skip this step.
|
|
|
|
Then you should query the schema of the most relevant tables.
|
|
""".format(
|
|
dialect="sqlite",
|
|
top_k=5,
|
|
)
|
|
|
|
agent = create_agent(
|
|
model,
|
|
tools,
|
|
system_prompt=system_prompt,
|
|
)
|
|
```
|