Skip to content

Commit b94c4be

Browse files
author
RN
committed
feat(sql_database): support content_and_artifact in SQLDatabase.run
1 parent 8ee81e1 commit b94c4be

File tree

2 files changed

+41
-8
lines changed

2 files changed

+41
-8
lines changed

libs/community/langchain_community/utilities/sql_database.py

Lines changed: 21 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from __future__ import annotations
44

55
import re
6-
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Union
6+
from typing import Any, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union
77

88
import sqlalchemy
99
from langchain_core._api import deprecated
@@ -544,11 +544,22 @@ def run(
544544
*,
545545
parameters: Optional[Dict[str, Any]] = None,
546546
execution_options: Optional[Dict[str, Any]] = None,
547-
) -> Union[str, Sequence[Dict[str, Any]], Result[Any]]:
548-
"""Execute a SQL command and return a string representing the results.
547+
**kwargs: Any,
548+
) -> Union[str, Sequence[Dict[str, Any]], Result[Any], Tuple[str, Any]]:
549+
"""Execute a SQL command and return its result.
549550
550-
If the statement returns rows, a string of the results is returned.
551-
If the statement returns no rows, an empty string is returned.
551+
Args:
552+
command: The SQL command to execute.
553+
fetch: The number of rows to fetch. Can be "one", "all", or "cursor".
554+
include_columns: Whether to include column names in the result.
555+
parameters: A dictionary of parameters to pass to the SQL command.
556+
execution_options: A dictionary of execution options for the engine.
557+
**kwargs: Additional keyword arguments, checked for 'response_format'.
558+
559+
Returns:
560+
A string representation of the results, or a cursor object. If
561+
`response_format` is 'content_and_artifact', returns a tuple of
562+
(string_representation, processed_result_list).
552563
"""
553564
result = self._execute(
554565
command, fetch, parameters=parameters, execution_options=execution_options
@@ -568,10 +579,12 @@ def run(
568579
if not include_columns:
569580
res = [tuple(row.values()) for row in res] # type: ignore[misc]
570581

571-
if not res:
572-
return ""
582+
string_result = "" if not res else str(res)
583+
584+
if kwargs.get("response_format") == "content_and_artifact":
585+
return string_result, res
573586
else:
574-
return str(res)
587+
return string_result
575588

576589
def get_table_info_no_throw(self, table_names: Optional[List[str]] = None) -> str:
577590
"""Get information about specified tables.
Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
from langchain_community.utilities.sql_database import SQLDatabase
2+
3+
4+
def test_sql_database_run_content_and_artifact():
5+
db = SQLDatabase.from_uri("sqlite+pysqlite:///:memory:")
6+
db.run("CREATE TABLE test (id INTEGER, name TEXT);")
7+
db.run("INSERT INTO test (id, name) VALUES (1, 'foo');")
8+
db.run("INSERT INTO test (id, name) VALUES (2, 'bar');")
9+
10+
# Test content_and_artifact format
11+
string_result, artifact = db.run(
12+
"SELECT * FROM test", response_format="content_and_artifact"
13+
)
14+
assert "foo" in string_result and "bar" in string_result
15+
assert (1, "foo") in artifact and (2, "bar") in artifact
16+
17+
# Test default format (string only)
18+
only_string = db.run("SELECT * FROM test")
19+
assert isinstance(only_string, str)
20+
assert "foo" in only_string and "bar" in only_string

0 commit comments

Comments
 (0)