class SQLTableRetrieverQueryEngine(BaseSQLTableQueryEngine):
"""SQL Table retriever query engine."""
def init(
self,
sql_database: SQLDatabase,
table_retriever: ObjectRetriever[SQLTableSchema],
rows_retrievers: Optional[dict[str, BaseRetriever]] = None,
llm: Optional[LLM] = None,
text_to_sql_prompt: Optional[BasePromptTemplate] = None,
context_query_kwargs: Optional[dict] = None,
synthesize_response: bool = True,
response_synthesis_prompt: Optional[BasePromptTemplate] = None,
refine_synthesis_prompt: Optional[BasePromptTemplate] = None,
context_str_prefix: Optional[str] = None,
sql_only: bool = False,
callback_manager: Optional[CallbackManager] = None,
kwargs: Any, ) -> None: """Initialize params.""" self._sql_retriever = NLSQLRetriever( sql_database, llm=llm, text_to_sql_prompt=text_to_sql_prompt, context_query_kwargs=context_query_kwargs, table_retriever=table_retriever, rows_retrievers=rows_retrievers, context_str_prefix=context_str_prefix, sql_only=sql_only, callback_manager=callback_manager, verbose=kwargs.get("verbose", False), ) super().init( synthesize_response=synthesize_response, response_synthesis_prompt=response_synthesis_prompt, refine_synthesis_prompt=refine_synthesis_prompt, llm=llm, callback_manager=callback_manager, kwargs,
)
@property
def sql_retriever(self) -> NLSQLRetriever:
"""Get SQL retriever."""
return self._sql_retriever
