mirror of
https://github.com/langchain-ai/docs.git
synced 2026-08-27 02:41:59 -04:00
npentrel
106 lines
4.2 KiB
Plaintext
106 lines
4.2 KiB
Plaintext
```python
|
|
import sqlite3
|
|
from langchain.tools import tool
|
|
|
|
# Below are minimal tools for demonstration purposes.
|
|
# They are not intended to be secure or for production use.
|
|
|
|
@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")
|
|
```
|