class BaseQueryEngine(PromptMixin, DispatcherSpanMixin):
"""Base query engine."""
def __init__(
self,
callback_manager: Optional[CallbackManager],
) -> None:
self.callback_manager = callback_manager or CallbackManager([])
def _get_prompts(self) -> Dict[str, Any]:
"""Get prompts."""
return {}
def _update_prompts(self, prompts: PromptDictType) -> None:
"""Update prompts."""
@dispatcher.span
def query(self, str_or_query_bundle: QueryType) -> RESPONSE_TYPE:
dispatcher.event(QueryStartEvent(query=str_or_query_bundle))
with self.callback_manager.as_trace("query"):
if isinstance(str_or_query_bundle, str):
str_or_query_bundle = QueryBundle(str_or_query_bundle)
query_result = self._query(str_or_query_bundle)
dispatcher.event(
QueryEndEvent(query=str_or_query_bundle, response=query_result)
)
return query_result
@dispatcher.span
async def aquery(self, str_or_query_bundle: QueryType) -> RESPONSE_TYPE:
dispatcher.event(QueryStartEvent(query=str_or_query_bundle))
with self.callback_manager.as_trace("query"):
if isinstance(str_or_query_bundle, str):
str_or_query_bundle = QueryBundle(str_or_query_bundle)
query_result = await self._aquery(str_or_query_bundle)
dispatcher.event(
QueryEndEvent(query=str_or_query_bundle, response=query_result)
)
return query_result
def retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
raise NotImplementedError(
"This query engine does not support retrieve, use query directly"
)
def synthesize(
self,
query_bundle: QueryBundle,
nodes: List[NodeWithScore],
additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
) -> RESPONSE_TYPE:
raise NotImplementedError(
"This query engine does not support synthesize, use query directly"
)
async def asynthesize(
self,
query_bundle: QueryBundle,
nodes: List[NodeWithScore],
additional_source_nodes: Optional[Sequence[NodeWithScore]] = None,
) -> RESPONSE_TYPE:
raise NotImplementedError(
"This query engine does not support asynthesize, use aquery directly"
)
@abstractmethod
def _query(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
pass
@abstractmethod
async def _aquery(self, query_bundle: QueryBundle) -> RESPONSE_TYPE:
pass