Can you please provide your expertise on why this ends up as 2 stack traces. I want
retriever
to be inside
chain
. Kinda lost on this 😦
class CustomSxQLRetriever(BaseRetriever):
def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
"""Retrieve nodes given query."""
with self.callback_manager.as_trace("query"):
with self.callback_manager.event(
CBEventType.RETRIEVE,
payload={EventPayload.QUERY_STR: query_bundle.query_str},
) as retrieve_event:
retrieved_nodes, _ = self.retrieve_with_metadata(query_bundle)
retrieve_event.on_end(
payload={EventPayload.NODES: retrieved_nodes},
)
return retrieved_nodes
engine = create_engine(CONN_STR)
sql_database = SQLDatabase(engine)
sql_retriever = CustomSQLRetriever(sql_database, return_raw=False)
class VeryCustomQueryEngine(CustomQueryEngine):
"""SQL Search but fallbacks to Vector Search Query Engine"""
sql_retriever: CustomSQLRetriever
vector_retriever: BaseRetriever
response_synthesizer: BaseSynthesizer
variable_data: Dict[str, Any]
def custom_query(self, query_str: str):
nodes = self.sql_retriever.retrieve(query)
if len(nodes) > 0:
print("GOT SQL RESPONSE SO PROCEED")
response_obj = self.response_synthesizer.synthesize(query_str, nodes)
return response_obj
synthesizer = get_response_synthesizer(response_mode="compact")
query_engine = VeryCustomQueryEngine(sql_retriever=sql_retriever, vector_retriever=vector_retriever, response_synthesizer=synthesizer, variable_data=variable_data, callback_manager=callback_manager)