From 364b771743cdeab320a9cc233a3ef8754b9fb822 Mon Sep 17 00:00:00 2001 From: Harrison Chase Date: Thu, 2 Feb 2023 22:07:41 -0800 Subject: [PATCH] sql return direct (#864) --- langchain/chains/sql_database/base.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/langchain/chains/sql_database/base.py b/langchain/chains/sql_database/base.py index 78eec071c6f..ccb3416f200 100644 --- a/langchain/chains/sql_database/base.py +++ b/langchain/chains/sql_database/base.py @@ -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