github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-06-24 15:43:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

sql return direct (#864)

Browse Source
This commit is contained in:
Harrison Chase 2023-02-02 22:07:41 -08:00 committed by GitHub
parent 483441d305
commit 364b771743
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 14 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
langchain/chains/sql_database/base.py
View File

@ -35,6 +35,9 @@ class SQLDatabaseChain(Chain, BaseModel):
input_key: str = "query" #: :meta private:
output_key: str = "result" #: :meta private:
return_intermediate_steps: bool = False
"""Whether or not to return the intermediate steps along with the final answer."""
return_direct: bool = False
"""Whether or not to return the result of querying the SQL table directly."""
class Config:
"""Configuration for this pydantic object."""
@ -83,11 +86,17 @@ class SQLDatabaseChain(Chain, BaseModel):
intermediate_steps.append(result)
self.callback_manager.on_text("\nSQLResult: ", verbose=self.verbose)
self.callback_manager.on_text(result, color="yellow", verbose=self.verbose)
self.callback_manager.on_text("\nAnswer:", verbose=self.verbose)
input_text += f"{sql_cmd}\nSQLResult: {result}\nAnswer:"
llm_inputs["input"] = input_text
final_result = llm_chain.predict(**llm_inputs)
self.callback_manager.on_text(final_result, color="green", verbose=self.verbose)
# If return direct, we just set the final result equal to the sql query
if self.return_direct:
final_result = result
else:
self.callback_manager.on_text("\nAnswer:", verbose=self.verbose)
input_text += f"{sql_cmd}\nSQLResult: {result}\nAnswer:"
llm_inputs["input"] = input_text
final_result = llm_chain.predict(**llm_inputs)
self.callback_manager.on_text(
final_result, color="green", verbose=self.verbose
)
chain_result: Dict[str, Any] = {self.output_key: final_result}
if self.return_intermediate_steps:
chain_result["intermediate_steps"] = intermediate_steps