code: """Base query engine."""
  
  import logging
  from abc import abstractmethod
  from typing import Any, Dict, List, Optional, Sequence
  
  from llama_index.legacy.bridge.pydantic import Field
  from llama_index.legacy.callbacks.base import CallbackManager
  from llama_index.legacy.core.query_pipeline.query_component import (
      ChainableMixin,
      InputKeys,
      OutputKeys,
      QueryComponent,
      validate_and_convert_stringable,
  )
  from llama_index.legacy.core.response.schema import RESPONSE_TYPE
  from llama_index.legacy.prompts.mixin import PromptDictType, PromptMixin
  from llama_index.legacy.schema import NodeWithScore, QueryBundle, QueryType
  
  logger = logging.getLogger(__name__)
  
  
  class BaseQueryEngine(ChainableMixin, PromptMixin):
      """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."""
  
      def query(self, str_or_query_bundle: QueryType) -> RESPONSE_TYPE:
          with self.callback_manager.as_trace("query"):
              if isinstance(str_or_query_bundle, str):
                  str_or_query_bundle = QueryBundle(str_or_query_bundle)
              return self._query(str_or_query_bundle)
  
      async def aquery(self, str_or_query_bundle: QueryType) -> RESPONSE_TYPE:
          with self.callback_manager.as_trace("query"):
              if isinstance(str_or_query_bundle, str):
                  str_or_query_bundle = QueryBundle(str_or_query_bundle)
              return await self._aquery(str_or_query_bundle)
  
      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
  
      def _as_query_component(self, **kwargs: Any) -> QueryComponent:
          """Return a query component."""
          return QueryEngineComponent(query_engine=self)
  
  
  class QueryEngineComponent(QueryComponent):
      """Query engine component."""
  
      query_engine: BaseQueryEngine = Field(..., description="Query engine")
  
      class Config:
          arbitrary_types_allowed = True
  
      def set_callback_manager(self, callback_manager: CallbackManager) -> None:
          """Set callback manager."""
          self.query_engine.callback_manager = callback_manager
  
      def _validate_component_inputs(self, input: Dict[str, Any]) -> Dict[str, Any]:
          """Validate component inputs during run_component."""
          # make sure input is a string
          input["input"] = validate_and_convert_stringable(input["input"])
          return input
  
      def _run_component(self, **kwargs: Any) -> Any:
          """Run component."""
          output = self.query_engine.query(kwargs["input"])
          return {"output": output}
  
      async def _arun_component(self, **kwargs: Any) -> Any:
          """Run component."""
          output = await self.query_engine.aquery(kwargs["input"])
          return {"output": output}
  
      @property
      def input_keys(self) -> InputKeys:
          """Input keys."""
          return InputKeys.from_keys({"input"})
  
      @property
      def output_keys(self) -> OutputKeys:
          """Output keys."""
          return OutputKeys.from_keys({"output"})
  
apis: llama_index.legacy.core.query_pipeline.query_component.OutputKeys.from_keys llama_index.legacy.callbacks.base.CallbackManager llama_index.legacy.schema.QueryBundle llama_index.legacy.core.query_pipeline.query_component.InputKeys.from_keys llama_index.legacy.bridge.pydantic.Field llama_index.legacy.core.query_pipeline.query_component.validate_and_convert_stringable
extract_api: [((647, 674), 'logging.getLogger', 'logging.getLogger', (['__name__'], {}), '(__name__)\n', (664, 674), False, 'import logging\n'), ((3066, 3104), 'llama_index.legacy.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""Query engine"""'}), "(..., description='Query engine')\n", (3071, 3104), False, 'from llama_index.legacy.bridge.pydantic import Field\n'), ((3550, 3597), 'llama_index.legacy.core.query_pipeline.query_component.validate_and_convert_stringable', 'validate_and_convert_stringable', (["input['input']"], {}), "(input['input'])\n", (3581, 3597), False, 'from llama_index.legacy.core.query_pipeline.query_component import ChainableMixin, InputKeys, OutputKeys, QueryComponent, validate_and_convert_stringable\n'), ((4076, 4106), 'llama_index.legacy.core.query_pipeline.query_component.InputKeys.from_keys', 'InputKeys.from_keys', (["{'input'}"], {}), "({'input'})\n", (4095, 4106), False, 'from llama_index.legacy.core.query_pipeline.query_component import ChainableMixin, InputKeys, OutputKeys, QueryComponent, validate_and_convert_stringable\n'), ((4205, 4237), 'llama_index.legacy.core.query_pipeline.query_component.OutputKeys.from_keys', 'OutputKeys.from_keys', (["{'output'}"], {}), "({'output'})\n", (4225, 4237), False, 'from llama_index.legacy.core.query_pipeline.query_component import ChainableMixin, InputKeys, OutputKeys, QueryComponent, validate_and_convert_stringable\n'), ((888, 907), 'llama_index.legacy.callbacks.base.CallbackManager', 'CallbackManager', (['[]'], {}), '([])\n', (903, 907), False, 'from llama_index.legacy.callbacks.base import CallbackManager\n'), ((1311, 1343), 'llama_index.legacy.schema.QueryBundle', 'QueryBundle', (['str_or_query_bundle'], {}), '(str_or_query_bundle)\n', (1322, 1343), False, 'from llama_index.legacy.schema import NodeWithScore, QueryBundle, QueryType\n'), ((1619, 1651), 'llama_index.legacy.schema.QueryBundle', 'QueryBundle', (['str_or_query_bundle'], {}), '(str_or_query_bundle)\n', (1630, 1651), False, 'from llama_index.legacy.schema import NodeWithScore, QueryBundle, QueryType\n')]