github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-04 04:28:58 +00:00
Code Issues Packages Projects Releases Wiki Activity

community[minor]: SQLDatabase Add fetch mode cursor, query parameters, query by selectable, expose execution options, and documentation (#17191)

Browse Source 
- **Description:** Improve `SQLDatabase` adapter component to promote
code re-use, see
[suggestion](https://github.com/langchain-ai/langchain/pull/16246#pullrequestreview-1846590962).
  - **Needed by:** GH-16246
  - **Addressed to:** @baskaryan, @cbornet 

## Details
- Add `cursor` fetch mode
- Accept SQL query parameters
- Accept both `str` and SQLAlchemy selectables as query expression
- Expose `execution_options`
- Documentation page (notebook) about `SQLDatabase` [^1]
See [About
SQLDatabase](https://github.com/langchain-ai/langchain/blob/c1c7b763/docs/docs/integrations/tools/sql_database.ipynb).

[^1]: Apparently there hasn't been any yet?

---------

Co-authored-by: Andreas Motl <andreas.motl@crate.io>
This commit is contained in:
Eugene Yurtsev
2024-02-07 22:23:43 -05:00
committed by GitHub
parent 7e4b676d53
commit 780e84ae79
6 changed files with 600 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

440
docs/docs/integrations/tools/sql_database.ipynb Normal file
View File

@@ -0,0 +1,440 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"# SQL Database\n",
"\n",
"::: {.callout-note}\n",
"The `SQLDatabase` adapter utility is a wrapper around a database connection.\n",
"\n",
"For talking to SQL databases, it uses the [SQLAlchemy] Core API .\n",
":::\n",
"\n",
"\n",
"This notebook shows how to use the utility to access an SQLite database.\n",
"It uses the example [Chinook Database], and demonstrates those features:\n",
"\n",
"- Query using SQL\n",
"- Query using SQLAlchemy selectable\n",
"- Fetch modes `cursor`, `all`, and `one`\n",
"- Bind query parameters\n",
"\n",
"[Chinook Database]: https://github.com/lerocha/chinook-database\n",
"[SQLAlchemy]: https://www.sqlalchemy.org/\n",
"\n",
"\n",
"You can use the `Tool` or `@tool` decorator to create a tool from this utility.\n",
"\n",
"\n",
"::: {.callout-caution}\n",
"If creating a tool from the SQLDatbase utility and combining it with an LLM or exposing it to an end user\n",
"remember to follow good security practices.\n",
"\n",
"See security information: https://python.langchain.com/docs/security\n",
":::"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"collapsed": true,
"jupyter": {
"outputs_hidden": true
}
},
"outputs": [],
"source": [
"!wget 'https://github.com/lerocha/chinook-database/releases/download/v1.4.2/Chinook_Sqlite.sql'"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1|AC/DC\r\n",
"2|Accept\r\n",
"3|Aerosmith\r\n",
"4|Alanis Morissette\r\n",
"5|Alice In Chains\r\n",
"6|Antônio Carlos Jobim\r\n",
"7|Apocalyptica\r\n",
"8|Audioslave\r\n",
"9|BackBeat\r\n",
"10|Billy Cobham\r\n",
"11|Black Label Society\r\n",
"12|Black Sabbath\r\n"
]
}
],
"source": [
"!sqlite3 -bail -cmd '.read Chinook_Sqlite.sql' -cmd 'SELECT * FROM Artist LIMIT 12;' -cmd '.quit'"
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [],
"source": [
"!sqlite3 -bail -cmd '.read Chinook_Sqlite.sql' -cmd '.save Chinook.db' -cmd '.quit'"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Initialize Database"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [],
"source": [
"from pprint import pprint\n",
"\n",
"import sqlalchemy as sa\n",
"from langchain.sql_database import SQLDatabase\n",
"\n",
"db = SQLDatabase.from_uri(\"sqlite:///Chinook.db\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query as cursor\n",
"\n",
"The fetch mode `cursor` returns results as SQLAlchemy's\n",
"`CursorResult` instance."
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<class 'sqlalchemy.engine.cursor.CursorResult'>\n",
"[{'ArtistId': 1, 'Name': 'AC/DC'},\n",
" {'ArtistId': 2, 'Name': 'Accept'},\n",
" {'ArtistId': 3, 'Name': 'Aerosmith'},\n",
" {'ArtistId': 4, 'Name': 'Alanis Morissette'},\n",
" {'ArtistId': 5, 'Name': 'Alice In Chains'},\n",
" {'ArtistId': 6, 'Name': 'Antônio Carlos Jobim'},\n",
" {'ArtistId': 7, 'Name': 'Apocalyptica'},\n",
" {'ArtistId': 8, 'Name': 'Audioslave'},\n",
" {'ArtistId': 9, 'Name': 'BackBeat'},\n",
" {'ArtistId': 10, 'Name': 'Billy Cobham'},\n",
" {'ArtistId': 11, 'Name': 'Black Label Society'},\n",
" {'ArtistId': 12, 'Name': 'Black Sabbath'}]\n"
]
}
],
"source": [
"result = db.run(\"SELECT * FROM Artist LIMIT 12;\", fetch=\"cursor\")\n",
"print(type(result))\n",
"pprint(list(result.mappings()))"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query as string payload\n",
"\n",
"The fetch modes `all` and `one` return results in string format."
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<class 'str'>\n",
"[(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'), (11, 'Black Label Society'), (12, 'Black Sabbath')]\n"
]
}
],
"source": [
"result = db.run(\"SELECT * FROM Artist LIMIT 12;\", fetch=\"all\")\n",
"print(type(result))\n",
"print(result)"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"<class 'str'>\n",
"[(1, 'AC/DC')]\n"
]
}
],
"source": [
"result = db.run(\"SELECT * FROM Artist LIMIT 12;\", fetch=\"one\")\n",
"print(type(result))\n",
"print(result)"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query with parameters\n",
"\n",
"In order to bind query parameters, use the optional `parameters` argument."
]
},
{
"cell_type": "code",
"execution_count": 41,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'ArtistId': 35, 'Name': 'Pedro Luís & A Parede'},\n",
" {'ArtistId': 115, 'Name': 'Page & Plant'},\n",
" {'ArtistId': 116, 'Name': 'Passengers'},\n",
" {'ArtistId': 117, 'Name': \"Paul D'Ianno\"},\n",
" {'ArtistId': 118, 'Name': 'Pearl Jam'},\n",
" {'ArtistId': 119, 'Name': 'Peter Tosh'},\n",
" {'ArtistId': 120, 'Name': 'Pink Floyd'},\n",
" {'ArtistId': 121, 'Name': 'Planet Hemp'},\n",
" {'ArtistId': 186, 'Name': 'Pedro Luís E A Parede'},\n",
" {'ArtistId': 256, 'Name': 'Philharmonia Orchestra & Sir Neville Marriner'},\n",
" {'ArtistId': 275, 'Name': 'Philip Glass Ensemble'}]\n"
]
}
],
"source": [
"result = db.run(\n",
" \"SELECT * FROM Artist WHERE Name LIKE :search;\",\n",
" parameters={\"search\": \"p%\"},\n",
" fetch=\"cursor\",\n",
")\n",
"pprint(list(result.mappings()))"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query with SQLAlchemy selectable\n",
"\n",
"Other than plain-text SQL statements, the adapter also accepts SQLAlchemy selectables."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[{'ArtistId': 35, 'Name': 'Pedro Luís & A Parede'},\n",
" {'ArtistId': 115, 'Name': 'Page & Plant'},\n",
" {'ArtistId': 116, 'Name': 'Passengers'},\n",
" {'ArtistId': 117, 'Name': \"Paul D'Ianno\"},\n",
" {'ArtistId': 118, 'Name': 'Pearl Jam'},\n",
" {'ArtistId': 119, 'Name': 'Peter Tosh'},\n",
" {'ArtistId': 120, 'Name': 'Pink Floyd'},\n",
" {'ArtistId': 121, 'Name': 'Planet Hemp'},\n",
" {'ArtistId': 186, 'Name': 'Pedro Luís E A Parede'},\n",
" {'ArtistId': 256, 'Name': 'Philharmonia Orchestra & Sir Neville Marriner'},\n",
" {'ArtistId': 275, 'Name': 'Philip Glass Ensemble'}]\n"
]
}
],
"source": [
"# In order to build a selectable on SA's Core API, you need a table definition.\n",
"metadata = sa.MetaData()\n",
"artist = sa.Table(\n",
" \"Artist\",\n",
" metadata,\n",
" sa.Column(\"ArtistId\", sa.INTEGER, primary_key=True),\n",
" sa.Column(\"Name\", sa.TEXT),\n",
")\n",
"\n",
"# Build a selectable with the same semantics of the recent query.\n",
"query = sa.select(artist).where(artist.c.Name.like(\"p%\"))\n",
"result = db.run(query, fetch=\"cursor\")\n",
"pprint(list(result.mappings()))"
]
},
{
"cell_type": "markdown",
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"source": [
"## Query with execution options\n",
"\n",
"It is possible to augment the statement invocation with custom execution options.\n",
"For example, when applying a schema name translation, subsequent statements will\n",
"fail, because they try to hit a non-existing table."
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {
"collapsed": false,
"jupyter": {
"outputs_hidden": false
}
},
"outputs": [
{
"ename": "OperationalError",
"evalue": "(sqlite3.OperationalError) no such table: bar.Artist\n[SQL: SELECT bar.\"Artist\".\"ArtistId\", bar.\"Artist\".\"Name\" \nFROM bar.\"Artist\" \nWHERE bar.\"Artist\".\"Name\" LIKE ?]\n[parameters: ('p%',)]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1969\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[0;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[1;32m 1968\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m evt_handled:\n\u001b[0;32m-> 1969\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1970\u001b[0m \u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[1;32m 1971\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1973\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_has_events \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_has_events:\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/default.py:922\u001b[0m, in \u001b[0;36mDefaultDialect.do_execute\u001b[0;34m(self, cursor, statement, parameters, context)\u001b[0m\n\u001b[1;32m 921\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdo_execute\u001b[39m(\u001b[38;5;28mself\u001b[39m, cursor, statement, parameters, context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 922\u001b[0m cursor\u001b[38;5;241m.\u001b[39mexecute(statement, parameters)\n",
"\u001b[0;31mOperationalError\u001b[0m: no such table: bar.Artist",
"\nThe above exception was the direct cause of the following exception:\n",
"\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)",
"Cell \u001b[0;32mIn[6], line 3\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Build a selectable with the same semantics of the recent query.\u001b[39;00m\n\u001b[1;32m 2\u001b[0m query \u001b[38;5;241m=\u001b[39m sa\u001b[38;5;241m.\u001b[39mselect(artist)\u001b[38;5;241m.\u001b[39mwhere(artist\u001b[38;5;241m.\u001b[39mc\u001b[38;5;241m.\u001b[39mName\u001b[38;5;241m.\u001b[39mlike(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mp\u001b[39m\u001b[38;5;124m%\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[0;32m----> 3\u001b[0m \u001b[43mdb\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[43mquery\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfetch\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mresult\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m{\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mschema_translate_map\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[43m{\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m:\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mbar\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m}\u001b[49m\u001b[43m}\u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/libs/community/langchain_community/utilities/sql_database.py:484\u001b[0m, in \u001b[0;36mSQLDatabase.run\u001b[0;34m(self, command, fetch, parameters, execution_options, include_columns)\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mrun\u001b[39m(\n\u001b[1;32m 472\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 473\u001b[0m command: Union[\u001b[38;5;28mstr\u001b[39m, Executable],\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 477\u001b[0m include_columns: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m 478\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Union[\u001b[38;5;28mstr\u001b[39m, Result]:\n\u001b[1;32m 479\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Execute a SQL command and return a string representing the results.\u001b[39;00m\n\u001b[1;32m 480\u001b[0m \n\u001b[1;32m 481\u001b[0m \u001b[38;5;124;03m If the statement returns rows, a string of the results is returned.\u001b[39;00m\n\u001b[1;32m 482\u001b[0m \u001b[38;5;124;03m If the statement returns no rows, an empty string is returned.\u001b[39;00m\n\u001b[1;32m 483\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 484\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 485\u001b[0m \u001b[43m \u001b[49m\u001b[43mcommand\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfetch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mexecution_options\u001b[49m\n\u001b[1;32m 486\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 488\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m fetch \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mresult\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n\u001b[1;32m 489\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/libs/community/langchain_community/utilities/sql_database.py:450\u001b[0m, in \u001b[0;36mSQLDatabase._execute\u001b[0;34m(self, command, fetch, parameters, execution_options)\u001b[0m\n\u001b[1;32m 448\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 449\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mQuery expression has unknown type: \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m\u001b[38;5;28mtype\u001b[39m(command)\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m--> 450\u001b[0m cursor \u001b[38;5;241m=\u001b[39m \u001b[43mconnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 451\u001b[0m \u001b[43m \u001b[49m\u001b[43mcommand\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 452\u001b[0m \u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 453\u001b[0m \u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mexecution_options\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 454\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 456\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cursor\u001b[38;5;241m.\u001b[39mreturns_rows:\n\u001b[1;32m 457\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m fetch \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mall\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1416\u001b[0m, in \u001b[0;36mConnection.execute\u001b[0;34m(self, statement, parameters, execution_options)\u001b[0m\n\u001b[1;32m 1414\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mObjectNotExecutableError(statement) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01merr\u001b[39;00m\n\u001b[1;32m 1415\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1416\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mmeth\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1417\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1418\u001b[0m \u001b[43m \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1419\u001b[0m \u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mNO_OPTIONS\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1420\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/sql/elements.py:516\u001b[0m, in \u001b[0;36mClauseElement._execute_on_connection\u001b[0;34m(self, connection, distilled_params, execution_options)\u001b[0m\n\u001b[1;32m 514\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m TYPE_CHECKING:\n\u001b[1;32m 515\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(\u001b[38;5;28mself\u001b[39m, Executable)\n\u001b[0;32m--> 516\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mconnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute_clauseelement\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 517\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdistilled_params\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\n\u001b[1;32m 518\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 519\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 520\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc\u001b[38;5;241m.\u001b[39mObjectNotExecutableError(\u001b[38;5;28mself\u001b[39m)\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1639\u001b[0m, in \u001b[0;36mConnection._execute_clauseelement\u001b[0;34m(self, elem, distilled_parameters, execution_options)\u001b[0m\n\u001b[1;32m 1627\u001b[0m compiled_cache: Optional[CompiledCacheType] \u001b[38;5;241m=\u001b[39m execution_options\u001b[38;5;241m.\u001b[39mget(\n\u001b[1;32m 1628\u001b[0m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcompiled_cache\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_compiled_cache\n\u001b[1;32m 1629\u001b[0m )\n\u001b[1;32m 1631\u001b[0m compiled_sql, extracted_params, cache_hit \u001b[38;5;241m=\u001b[39m elem\u001b[38;5;241m.\u001b[39m_compile_w_cache(\n\u001b[1;32m 1632\u001b[0m dialect\u001b[38;5;241m=\u001b[39mdialect,\n\u001b[1;32m 1633\u001b[0m compiled_cache\u001b[38;5;241m=\u001b[39mcompiled_cache,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1637\u001b[0m linting\u001b[38;5;241m=\u001b[39m\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdialect\u001b[38;5;241m.\u001b[39mcompiler_linting \u001b[38;5;241m|\u001b[39m compiler\u001b[38;5;241m.\u001b[39mWARN_LINTING,\n\u001b[1;32m 1638\u001b[0m )\n\u001b[0;32m-> 1639\u001b[0m ret \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_execute_context\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1640\u001b[0m \u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1641\u001b[0m \u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mexecution_ctx_cls\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_init_compiled\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1642\u001b[0m \u001b[43m \u001b[49m\u001b[43mcompiled_sql\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1643\u001b[0m \u001b[43m \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1644\u001b[0m \u001b[43m \u001b[49m\u001b[43mexecution_options\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1645\u001b[0m \u001b[43m \u001b[49m\u001b[43mcompiled_sql\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1646\u001b[0m \u001b[43m \u001b[49m\u001b[43mdistilled_parameters\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1647\u001b[0m \u001b[43m \u001b[49m\u001b[43melem\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1648\u001b[0m \u001b[43m \u001b[49m\u001b[43mextracted_params\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1649\u001b[0m \u001b[43m \u001b[49m\u001b[43mcache_hit\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcache_hit\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 1650\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1651\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m has_events:\n\u001b[1;32m 1652\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mafter_execute(\n\u001b[1;32m 1653\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 1654\u001b[0m elem,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1658\u001b[0m ret,\n\u001b[1;32m 1659\u001b[0m )\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1848\u001b[0m, in \u001b[0;36mConnection._execute_context\u001b[0;34m(self, dialect, constructor, statement, parameters, execution_options, *args, **kw)\u001b[0m\n\u001b[1;32m 1843\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exec_insertmany_context(\n\u001b[1;32m 1844\u001b[0m dialect,\n\u001b[1;32m 1845\u001b[0m context,\n\u001b[1;32m 1846\u001b[0m )\n\u001b[1;32m 1847\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1848\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_exec_single_context\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1849\u001b[0m \u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstatement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparameters\u001b[49m\n\u001b[1;32m 1850\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1988\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[0;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[1;32m 1985\u001b[0m result \u001b[38;5;241m=\u001b[39m context\u001b[38;5;241m.\u001b[39m_setup_result_proxy()\n\u001b[1;32m 1987\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m-> 1988\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_handle_dbapi_exception\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1989\u001b[0m \u001b[43m \u001b[49m\u001b[43me\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[1;32m 1990\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1992\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m result\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:2343\u001b[0m, in \u001b[0;36mConnection._handle_dbapi_exception\u001b[0;34m(self, e, statement, parameters, cursor, context, is_sub_exec)\u001b[0m\n\u001b[1;32m 2341\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m should_wrap:\n\u001b[1;32m 2342\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m sqlalchemy_exception \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 2343\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m sqlalchemy_exception\u001b[38;5;241m.\u001b[39mwith_traceback(exc_info[\u001b[38;5;241m2\u001b[39m]) \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m 2344\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 2345\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_info[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/base.py:1969\u001b[0m, in \u001b[0;36mConnection._exec_single_context\u001b[0;34m(self, dialect, context, statement, parameters)\u001b[0m\n\u001b[1;32m 1967\u001b[0m \u001b[38;5;28;01mbreak\u001b[39;00m\n\u001b[1;32m 1968\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m evt_handled:\n\u001b[0;32m-> 1969\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdo_execute\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 1970\u001b[0m \u001b[43m \u001b[49m\u001b[43mcursor\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstr_statement\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43meffective_parameters\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcontext\u001b[49m\n\u001b[1;32m 1971\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1973\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_has_events \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine\u001b[38;5;241m.\u001b[39m_has_events:\n\u001b[1;32m 1974\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdispatch\u001b[38;5;241m.\u001b[39mafter_cursor_execute(\n\u001b[1;32m 1975\u001b[0m \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m 1976\u001b[0m cursor,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1980\u001b[0m context\u001b[38;5;241m.\u001b[39mexecutemany,\n\u001b[1;32m 1981\u001b[0m )\n",
"File \u001b[0;32m~/dev/crate/ecosystem/langchain/.venv/lib/python3.11/site-packages/sqlalchemy/engine/default.py:922\u001b[0m, in \u001b[0;36mDefaultDialect.do_execute\u001b[0;34m(self, cursor, statement, parameters, context)\u001b[0m\n\u001b[1;32m 921\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdo_execute\u001b[39m(\u001b[38;5;28mself\u001b[39m, cursor, statement, parameters, context\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m):\n\u001b[0;32m--> 922\u001b[0m cursor\u001b[38;5;241m.\u001b[39mexecute(statement, parameters)\n",
"\u001b[0;31mOperationalError\u001b[0m: (sqlite3.OperationalError) no such table: bar.Artist\n[SQL: SELECT bar.\"Artist\".\"ArtistId\", bar.\"Artist\".\"Name\" \nFROM bar.\"Artist\" \nWHERE bar.\"Artist\".\"Name\" LIKE ?]\n[parameters: ('p%',)]\n(Background on this error at: https://sqlalche.me/e/20/e3q8)"
]
}
],
"source": [
"query = sa.select(artist).where(artist.c.Name.like(\"p%\"))\n",
"db.run(query, fetch=\"cursor\", execution_options={\"schema_translate_map\": {None: \"bar\"}})"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.2"
}
},
"nbformat": 4,
"nbformat_minor": 4
}

6
libs/community/langchain_community/tools/sql_database/tool.py
View File

@@ -1,6 +1,8 @@
# flake8: noqa # flake8: noqa
"""Tools for interacting with a SQL database.""" """Tools for interacting with a SQL database."""
from typing import Any, Dict, Optional, Type from typing import Any, Dict, Optional, Sequence, Type, Union
from sqlalchemy import Result
from langchain_core.pydantic_v1 import BaseModel, Field, root_validator from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
@@ -42,7 +44,7 @@ class QuerySQLDataBaseTool(BaseSQLDatabaseTool, BaseTool):
self, self,
query: str, query: str,
run_manager: Optional[CallbackManagerForToolRun] = None, run_manager: Optional[CallbackManagerForToolRun] = None,
) -> str: ) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
"""Execute the query, return the results or an error message.""" """Execute the query, return the results or an error message."""
return self.db.run_no_throw(query) return self.db.run_no_throw(query)

102
libs/community/langchain_community/utilities/sql_database.py
View File

@@ -1,12 +1,21 @@
"""SQLAlchemy wrapper around a database.""" """SQLAlchemy wrapper around a database."""
from __future__ import annotations from __future__ import annotations
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Union
import sqlalchemy import sqlalchemy
from langchain_core._api import deprecated from langchain_core._api import deprecated
from langchain_core.utils import get_from_env from langchain_core.utils import get_from_env
from sqlalchemy import MetaData, Table, create_engine, inspect, select, text from sqlalchemy import (
Executable,
MetaData,
Result,
Table,
create_engine,
inspect,
select,
text,
)
from sqlalchemy.engine import Engine from sqlalchemy.engine import Engine
from sqlalchemy.exc import ProgrammingError, SQLAlchemyError from sqlalchemy.exc import ProgrammingError, SQLAlchemyError
from sqlalchemy.schema import CreateTable from sqlalchemy.schema import CreateTable
@@ -373,67 +382,113 @@ class SQLDatabase:
def _execute( def _execute(
self, self,
command: str, command: Union[str, Executable],
fetch: Literal["all", "one"] = "all", fetch: Literal["all", "one", "cursor"] = "all",
) -> Sequence[Dict[str, Any]]: *,
parameters: Optional[Dict[str, Any]] = None,
execution_options: Optional[Dict[str, Any]] = None,
) -> Union[Sequence[Dict[str, Any]], Result]:
""" """
Executes SQL command through underlying engine. Executes SQL command through underlying engine.
If the statement returns no rows, an empty list is returned. If the statement returns no rows, an empty list is returned.
""" """
parameters = parameters or {}
execution_options = execution_options or {}
with self._engine.begin() as connection: # type: Connection # type: ignore[name-defined] with self._engine.begin() as connection: # type: Connection # type: ignore[name-defined]
if self._schema is not None: if self._schema is not None:
if self.dialect == "snowflake": if self.dialect == "snowflake":
connection.exec_driver_sql( connection.exec_driver_sql(
"ALTER SESSION SET search_path = %s", (self._schema,) "ALTER SESSION SET search_path = %s",
(self._schema,),
execution_options=execution_options,
) )
elif self.dialect == "bigquery": elif self.dialect == "bigquery":
connection.exec_driver_sql("SET @@dataset_id=?", (self._schema,)) connection.exec_driver_sql(
"SET @@dataset_id=?",
(self._schema,),
execution_options=execution_options,
)
elif self.dialect == "mssql": elif self.dialect == "mssql":
pass pass
elif self.dialect == "trino": elif self.dialect == "trino":
connection.exec_driver_sql("USE ?", (self._schema,)) connection.exec_driver_sql(
"USE ?",
(self._schema,),
execution_options=execution_options,
)
elif self.dialect == "duckdb": elif self.dialect == "duckdb":
# Unclear which parameterized argument syntax duckdb supports. # Unclear which parameterized argument syntax duckdb supports.
# The docs for the duckdb client say they support multiple, # The docs for the duckdb client say they support multiple,
# but `duckdb_engine` seemed to struggle with all of them: # but `duckdb_engine` seemed to struggle with all of them:
# https://github.com/Mause/duckdb_engine/issues/796 # https://github.com/Mause/duckdb_engine/issues/796
connection.exec_driver_sql(f"SET search_path TO {self._schema}") connection.exec_driver_sql(
f"SET search_path TO {self._schema}",
execution_options=execution_options,
)
elif self.dialect == "oracle": elif self.dialect == "oracle":
connection.exec_driver_sql( connection.exec_driver_sql(
f"ALTER SESSION SET CURRENT_SCHEMA = {self._schema}" f"ALTER SESSION SET CURRENT_SCHEMA = {self._schema}",
execution_options=execution_options,
) )
elif self.dialect == "sqlany": elif self.dialect == "sqlany":
# If anybody using Sybase SQL anywhere database then it should not # If anybody using Sybase SQL anywhere database then it should not
# go to else condition. It should be same as mssql. # go to else condition. It should be same as mssql.
pass pass
elif self.dialect == "postgresql": # postgresql elif self.dialect == "postgresql": # postgresql
connection.exec_driver_sql("SET search_path TO %s", (self._schema,)) connection.exec_driver_sql(
"SET search_path TO %s",
(self._schema,),
execution_options=execution_options,
)
if isinstance(command, str):
command = text(command)
elif isinstance(command, Executable):
pass
else:
raise TypeError(f"Query expression has unknown type: {type(command)}")
cursor = connection.execute(
command,
parameters,
execution_options=execution_options,
)
cursor = connection.execute(text(command))
if cursor.returns_rows: if cursor.returns_rows:
if fetch == "all": if fetch == "all":
result = [x._asdict() for x in cursor.fetchall()] result = [x._asdict() for x in cursor.fetchall()]
elif fetch == "one": elif fetch == "one":
first_result = cursor.fetchone() first_result = cursor.fetchone()
result = [] if first_result is None else [first_result._asdict()] result = [] if first_result is None else [first_result._asdict()]
elif fetch == "cursor":
return cursor
else: else:
raise ValueError("Fetch parameter must be either 'one' or 'all'") raise ValueError(
"Fetch parameter must be either 'one', 'all', or 'cursor'"
)
return result return result
return [] return []
def run( def run(
self, self,
command: str, command: Union[str, Executable],
fetch: Literal["all", "one"] = "all", fetch: Literal["all", "one", "cursor"] = "all",
include_columns: bool = False, include_columns: bool = False,
) -> str: *,
parameters: Optional[Dict[str, Any]] = None,
execution_options: Optional[Dict[str, Any]] = None,
) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
"""Execute a SQL command and return a string representing the results. """Execute a SQL command and return a string representing the results.
If the statement returns rows, a string of the results is returned. If the statement returns rows, a string of the results is returned.
If the statement returns no rows, an empty string is returned. If the statement returns no rows, an empty string is returned.
""" """
result = self._execute(command, fetch) result = self._execute(
command, fetch, parameters=parameters, execution_options=execution_options
)
if fetch == "cursor":
return result
res = [ res = [
{ {
@@ -472,7 +527,10 @@ class SQLDatabase:
command: str, command: str,
fetch: Literal["all", "one"] = "all", fetch: Literal["all", "one"] = "all",
include_columns: bool = False, include_columns: bool = False,
) -> str: *,
parameters: Optional[Dict[str, Any]] = None,
execution_options: Optional[Dict[str, Any]] = None,
) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
"""Execute a SQL command and return a string representing the results. """Execute a SQL command and return a string representing the results.
If the statement returns rows, a string of the results is returned. If the statement returns rows, a string of the results is returned.
@@ -481,7 +539,13 @@ class SQLDatabase:
If the statement throws an error, the error message is returned. If the statement throws an error, the error message is returned.
""" """
try: try:
return self.run(command, fetch, include_columns) return self.run(
command,
fetch,
parameters=parameters,
execution_options=execution_options,
include_columns=include_columns,
)
except SQLAlchemyError as e: except SQLAlchemyError as e:
"""Format the error message""" """Format the error message"""
return f"Error: {e}" return f"Error: {e}"

2
libs/experimental/langchain_experimental/sql/base.py
View File

@@ -173,7 +173,7 @@ class SQLDatabaseChain(Chain):
sql_cmd = checked_sql_command sql_cmd = checked_sql_command
_run_manager.on_text("\nSQLResult: ", verbose=self.verbose) _run_manager.on_text("\nSQLResult: ", verbose=self.verbose)
_run_manager.on_text(result, color="yellow", verbose=self.verbose) _run_manager.on_text(str(result), color="yellow", verbose=self.verbose)
# If return direct, we just set the final result equal to # If return direct, we just set the final result equal to
# the result of the sql query result, otherwise try to get a human readable # the result of the sql query result, otherwise try to get a human readable
# final answer # final answer

1
libs/experimental/langchain_experimental/sql/vector_sql.py
View File

@@ -78,6 +78,7 @@ class VectorSQLRetrieveAllOutputParser(VectorSQLOutputParser):
def get_result_from_sqldb(db: SQLDatabase, cmd: str) -> Sequence[Dict[str, Any]]: def get_result_from_sqldb(db: SQLDatabase, cmd: str) -> Sequence[Dict[str, Any]]:
result = db._execute(cmd, fetch="all") result = db._execute(cmd, fetch="all")
assert isinstance(result, Sequence)
return result return result

75
libs/langchain/tests/unit_tests/test_sql_database.py
View File

@@ -1,16 +1,19 @@
# flake8: noqa=E501 # flake8: noqa: E501
"""Test SQL database wrapper.""" """Test SQL database wrapper."""
import pytest
import sqlalchemy as sa
from langchain_community.utilities.sql_database import SQLDatabase, truncate_word from langchain_community.utilities.sql_database import SQLDatabase, truncate_word
from sqlalchemy import ( from sqlalchemy import (
Column, Column,
Integer, Integer,
MetaData, MetaData,
Result,
String, String,
Table, Table,
Text, Text,
create_engine, create_engine,
insert, insert,
select,
) )
metadata_obj = MetaData() metadata_obj = MetaData()
@@ -108,8 +111,8 @@ def test_table_info_w_sample_rows() -> None:
assert sorted(output.split()) == sorted(expected_output.split()) assert sorted(output.split()) == sorted(expected_output.split())
def test_sql_database_run() -> None: def test_sql_database_run_fetch_all() -> None:
"""Test that commands can be run successfully and returned in correct format.""" """Verify running SQL expressions returning results as strings."""
engine = create_engine("sqlite:///:memory:") engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine) metadata_obj.create_all(engine)
stmt = insert(user).values( stmt = insert(user).values(
@@ -131,6 +134,52 @@ def test_sql_database_run() -> None:
assert full_output == expected_full_output assert full_output == expected_full_output
def test_sql_database_run_fetch_result() -> None:
"""Verify running SQL expressions returning results as SQLAlchemy `Result` instances."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn:
conn.execute(stmt)
db = SQLDatabase(engine)
command = "select user_id, user_name, user_bio from user where user_id = 17"
result = db.run(command, fetch="cursor", include_columns=True)
expected = [{"user_id": 17, "user_name": "hwchase", "user_bio": None}]
assert isinstance(result, Result)
assert result.mappings().fetchall() == expected
def test_sql_database_run_with_parameters() -> None:
"""Verify running SQL expressions with query parameters."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn:
conn.execute(stmt)
db = SQLDatabase(engine)
command = "select user_id, user_name, user_bio from user where user_id = :user_id"
full_output = db.run(command, parameters={"user_id": 17}, include_columns=True)
expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]"
assert full_output == expected_full_output
def test_sql_database_run_sqlalchemy_selectable() -> None:
"""Verify running SQL expressions using SQLAlchemy selectable."""
engine = create_engine("sqlite:///:memory:")
metadata_obj.create_all(engine)
stmt = insert(user).values(user_id=17, user_name="hwchase")
with engine.begin() as conn:
conn.execute(stmt)
db = SQLDatabase(engine)
command = select(user).where(user.c.user_id == 17)
full_output = db.run(command, include_columns=True)
expected_full_output = "[{'user_id': 17, 'user_name': 'hwchase', 'user_bio': None}]"
assert full_output == expected_full_output
def test_sql_database_run_update() -> None: def test_sql_database_run_update() -> None:
"""Test commands which return no rows return an empty string.""" """Test commands which return no rows return an empty string."""
engine = create_engine("sqlite:///:memory:") engine = create_engine("sqlite:///:memory:")
@@ -145,6 +194,24 @@ def test_sql_database_run_update() -> None:
assert output == expected_output assert output == expected_output
def test_sql_database_schema_translate_map() -> None:
"""Verify using statement-specific execution options."""
engine = create_engine("sqlite:///:memory:")
db = SQLDatabase(engine)
# Define query using SQLAlchemy selectable.
command = select(user).where(user.c.user_id == 17)
# Define statement-specific execution options.
execution_options = {"schema_translate_map": {None: "bar"}}
# Verify the schema translation is applied.
with pytest.raises(sa.exc.OperationalError) as ex:
db.run(command, execution_options=execution_options, fetch="cursor")
assert ex.match("no such table: bar.user")
def test_truncate_word() -> None: def test_truncate_word() -> None:
assert truncate_word("Hello World", length=5) == "He..." assert truncate_word("Hello World", length=5) == "He..."
assert truncate_word("Hello World", length=0) == "Hello World" assert truncate_word("Hello World", length=0) == "Hello World"