class PGVectorSQLQueryEngine(BaseSQLTableQueryEngine):
"""PGvector SQL查询引擎。
这是一个修改过的普通文本到SQL查询引擎的版本,因为我们可以在SQL查询中推断嵌入向量。
注意:这是一个测试版功能。
注意:任何文本到SQL的应用都应该意识到执行任意SQL查询可能存在安全风险。建议采取必要的预防措施,比如使用受限角色、只读数据库、沙盒等。"""
def __init__(
self,
sql_database: SQLDatabase,
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,
tables: Optional[Union[List[str], List[Table]]] = None,
service_context: Optional[ServiceContext] = None,
context_str_prefix: Optional[str] = None,
sql_only: bool = False,
callback_manager: Optional[CallbackManager] = None,
**kwargs: Any,
) -> None:
"""初始化参数。"""
text_to_sql_prompt = text_to_sql_prompt or DEFAULT_TEXT_TO_SQL_PGVECTOR_PROMPT
self._sql_retriever = NLSQLRetriever(
sql_database,
llm=llm,
text_to_sql_prompt=text_to_sql_prompt,
context_query_kwargs=context_query_kwargs,
tables=tables,
sql_parser_mode=SQLParserMode.PGVECTOR,
context_str_prefix=context_str_prefix,
service_context=service_context,
sql_only=sql_only,
callback_manager=callback_manager,
)
super().__init__(
synthesize_response=synthesize_response,
response_synthesis_prompt=response_synthesis_prompt,
refine_synthesis_prompt=refine_synthesis_prompt,
llm=llm,
service_context=service_context,
callback_manager=callback_manager,
**kwargs,
)
@property
def sql_retriever(self) -> NLSQLRetriever:
"""获取SQL检索器。"""
return self._sql_retriever