code: """Base retriever."""
  
  from abc import abstractmethod
  from typing import Any, Dict, List, Optional
  
  from llama_index.core.base.base_query_engine import BaseQueryEngine
  from llama_index.core.base.query_pipeline.query import (
      ChainableMixin,
      InputKeys,
      OutputKeys,
      QueryComponent,
      validate_and_convert_stringable,
  )
  from llama_index.core.bridge.pydantic import Field
  from llama_index.core.callbacks.base import CallbackManager
  from llama_index.core.callbacks.schema import CBEventType, EventPayload
  from llama_index.core.prompts.mixin import (
      PromptDictType,
      PromptMixin,
      PromptMixinType,
  )
  from llama_index.core.schema import (
      BaseNode,
      IndexNode,
      NodeWithScore,
      QueryBundle,
      QueryType,
      TextNode,
  )
  from llama_index.core.service_context import ServiceContext
  from llama_index.core.settings import Settings
  from llama_index.core.utils import print_text
  from llama_index.core.instrumentation.events.retrieval import (
      RetrievalEndEvent,
      RetrievalStartEvent,
  )
  import llama_index.core.instrumentation as instrument
  
  dispatcher = instrument.get_dispatcher(__name__)
  
  
  class BaseRetriever(ChainableMixin, PromptMixin):
      """Base retriever."""
  
      def __init__(
          self,
          callback_manager: Optional[CallbackManager] = None,
          object_map: Optional[Dict] = None,
          objects: Optional[List[IndexNode]] = None,
          verbose: bool = False,
      ) -> None:
          self.callback_manager = callback_manager or CallbackManager()
  
          if objects is not None:
              object_map = {obj.index_id: obj.obj for obj in objects}
  
          self.object_map = object_map or {}
          self._verbose = verbose
  
      def _check_callback_manager(self) -> None:
          """Check callback manager."""
          if not hasattr(self, "callback_manager"):
              self.callback_manager = Settings.callback_manager
  
      def _get_prompts(self) -> PromptDictType:
          """Get prompts."""
          return {}
  
      def _get_prompt_modules(self) -> PromptMixinType:
          """Get prompt modules."""
          return {}
  
      def _update_prompts(self, prompts: PromptDictType) -> None:
          """Update prompts."""
  
      def _retrieve_from_object(
          self,
          obj: Any,
          query_bundle: QueryBundle,
          score: float,
      ) -> List[NodeWithScore]:
          """Retrieve nodes from object."""
          if self._verbose:
              print_text(
                  f"Retrieving from object {obj.__class__.__name__} with query {query_bundle.query_str}\n",
                  color="llama_pink",
              )
          if isinstance(obj, NodeWithScore):
              return [obj]
          elif isinstance(obj, BaseNode):
              return [NodeWithScore(node=obj, score=score)]
          elif isinstance(obj, BaseQueryEngine):
              response = obj.query(query_bundle)
              return [
                  NodeWithScore(
                      node=TextNode(text=str(response), metadata=response.metadata or {}),
                      score=score,
                  )
              ]
          elif isinstance(obj, BaseRetriever):
              return obj.retrieve(query_bundle)
          elif isinstance(obj, QueryComponent):
              component_keys = obj.input_keys.required_keys
              if len(component_keys) > 1:
                  raise ValueError(
                      f"QueryComponent {obj} has more than one input key: {component_keys}"
                  )
              elif len(component_keys) == 0:
                  component_response = obj.run_component()
              else:
                  kwargs = {next(iter(component_keys)): query_bundle.query_str}
                  component_response = obj.run_component(**kwargs)
  
              result_output = str(next(iter(component_response.values())))
              return [NodeWithScore(node=TextNode(text=result_output), score=score)]
          else:
              raise ValueError(f"Object {obj} is not retrievable.")
  
      async def _aretrieve_from_object(
          self,
          obj: Any,
          query_bundle: QueryBundle,
          score: float,
      ) -> List[NodeWithScore]:
          """Retrieve nodes from object."""
          if isinstance(obj, NodeWithScore):
              return [obj]
          elif isinstance(obj, BaseNode):
              return [NodeWithScore(node=obj, score=score)]
          elif isinstance(obj, BaseQueryEngine):
              response = await obj.aquery(query_bundle)
              return [NodeWithScore(node=TextNode(text=str(response)), score=score)]
          elif isinstance(obj, BaseRetriever):
              return await obj.aretrieve(query_bundle)
          elif isinstance(obj, QueryComponent):
              component_keys = obj.input_keys.required_keys
              if len(component_keys) > 1:
                  raise ValueError(
                      f"QueryComponent {obj} has more than one input key: {component_keys}"
                  )
              elif len(component_keys) == 0:
                  component_response = await obj.arun_component()
              else:
                  kwargs = {next(iter(component_keys)): query_bundle.query_str}
                  component_response = await obj.arun_component(**kwargs)
  
              result_output = str(next(iter(component_response.values())))
              return [NodeWithScore(node=TextNode(text=result_output), score=score)]
          else:
              raise ValueError(f"Object {obj} is not retrievable.")
  
      def _handle_recursive_retrieval(
          self, query_bundle: QueryBundle, nodes: List[NodeWithScore]
      ) -> List[NodeWithScore]:
          retrieved_nodes: List[NodeWithScore] = []
          for n in nodes:
              node = n.node
              score = n.score or 1.0
              if isinstance(node, IndexNode):
                  obj = node.obj or self.object_map.get(node.index_id, None)
                  if obj is not None:
                      if self._verbose:
                          print_text(
                              f"Retrieval entering {node.index_id}: {obj.__class__.__name__}\n",
                              color="llama_turquoise",
                          )
                      retrieved_nodes.extend(
                          self._retrieve_from_object(
                              obj, query_bundle=query_bundle, score=score
                          )
                      )
                  else:
                      retrieved_nodes.append(n)
              else:
                  retrieved_nodes.append(n)
  
          seen = set()
          return [
              n
              for n in retrieved_nodes
              if not (n.node.hash in seen or seen.add(n.node.hash))  # type: ignore[func-returns-value]
          ]
  
      async def _ahandle_recursive_retrieval(
          self, query_bundle: QueryBundle, nodes: List[NodeWithScore]
      ) -> List[NodeWithScore]:
          retrieved_nodes: List[NodeWithScore] = []
          for n in nodes:
              node = n.node
              score = n.score or 1.0
              if isinstance(node, IndexNode):
                  obj = node.obj or self.object_map.get(node.index_id, None)
                  if obj is not None:
                      if self._verbose:
                          print_text(
                              f"Retrieval entering {node.index_id}: {obj.__class__.__name__}\n",
                              color="llama_turquoise",
                          )
                      # TODO: Add concurrent execution via `run_jobs()` ?
                      retrieved_nodes.extend(
                          await self._aretrieve_from_object(
                              obj, query_bundle=query_bundle, score=score
                          )
                      )
                  else:
                      retrieved_nodes.append(n)
              else:
                  retrieved_nodes.append(n)
  
          # remove any duplicates based on hash
          seen = set()
          return [
              n
              for n in retrieved_nodes
              if not (n.node.hash in seen or seen.add(n.node.hash))  # type: ignore[func-returns-value]
          ]
  
      @dispatcher.span
      def retrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
          """Retrieve nodes given query.
  
          Args:
              str_or_query_bundle (QueryType): Either a query string or
                  a QueryBundle object.
  
          """
          self._check_callback_manager()
          dispatcher.event(RetrievalStartEvent())
          if isinstance(str_or_query_bundle, str):
              query_bundle = QueryBundle(str_or_query_bundle)
          else:
              query_bundle = str_or_query_bundle
          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:
                  nodes = self._retrieve(query_bundle)
                  nodes = self._handle_recursive_retrieval(query_bundle, nodes)
                  retrieve_event.on_end(
                      payload={EventPayload.NODES: nodes},
                  )
          dispatcher.event(RetrievalEndEvent())
          return nodes
  
      @dispatcher.span
      async def aretrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
          self._check_callback_manager()
          dispatcher.event(RetrievalStartEvent())
          if isinstance(str_or_query_bundle, str):
              query_bundle = QueryBundle(str_or_query_bundle)
          else:
              query_bundle = str_or_query_bundle
          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:
                  nodes = await self._aretrieve(query_bundle=query_bundle)
                  nodes = await self._ahandle_recursive_retrieval(
                      query_bundle=query_bundle, nodes=nodes
                  )
                  retrieve_event.on_end(
                      payload={EventPayload.NODES: nodes},
                  )
          dispatcher.event(RetrievalEndEvent())
          return nodes
  
      @abstractmethod
      def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
          """Retrieve nodes given query.
  
          Implemented by the user.
  
          """
  
      # TODO: make this abstract
      # @abstractmethod
      async def _aretrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
          """Asynchronously retrieve nodes given query.
  
          Implemented by the user.
  
          """
          return self._retrieve(query_bundle)
  
      def get_service_context(self) -> Optional[ServiceContext]:
          """Attempts to resolve a service context.
          Short-circuits at self.service_context, self._service_context,
          or self._index.service_context.
          """
          if hasattr(self, "service_context"):
              return self.service_context
          if hasattr(self, "_service_context"):
              return self._service_context
          elif hasattr(self, "_index") and hasattr(self._index, "service_context"):
              return self._index.service_context
          return None
  
      def _as_query_component(self, **kwargs: Any) -> QueryComponent:
          """Return a query component."""
          return RetrieverComponent(retriever=self)
  
  
  class RetrieverComponent(QueryComponent):
      """Retriever component."""
  
      retriever: BaseRetriever = Field(..., description="Retriever")
  
      class Config:
          arbitrary_types_allowed = True
  
      def set_callback_manager(self, callback_manager: CallbackManager) -> None:
          """Set callback manager."""
          self.retriever.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.retriever.retrieve(kwargs["input"])
          return {"output": output}
  
      async def _arun_component(self, **kwargs: Any) -> Any:
          """Run component."""
          output = await self.retriever.aretrieve(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.core.instrumentation.events.retrieval.RetrievalStartEvent llama_index.core.instrumentation.events.retrieval.RetrievalEndEvent llama_index.core.callbacks.base.CallbackManager llama_index.core.base.query_pipeline.query.InputKeys.from_keys llama_index.core.utils.print_text llama_index.core.bridge.pydantic.Field llama_index.core.instrumentation.get_dispatcher llama_index.core.base.query_pipeline.query.OutputKeys.from_keys llama_index.core.schema.TextNode llama_index.core.schema.QueryBundle llama_index.core.schema.NodeWithScore llama_index.core.base.query_pipeline.query.validate_and_convert_stringable
extract_api: [((1092, 1127), 'llama_index.core.instrumentation.get_dispatcher', 'instrument.get_dispatcher', (['__name__'], {}), '(__name__)\n', (1117, 1127), True, 'import llama_index.core.instrumentation as instrument\n'), ((11319, 11354), 'llama_index.core.bridge.pydantic.Field', 'Field', (['...'], {'description': '"""Retriever"""'}), "(..., description='Retriever')\n", (11324, 11354), False, 'from llama_index.core.bridge.pydantic import Field\n'), ((11797, 11844), 'llama_index.core.base.query_pipeline.query.validate_and_convert_stringable', 'validate_and_convert_stringable', (["input['input']"], {}), "(input['input'])\n", (11828, 11844), False, 'from llama_index.core.base.query_pipeline.query import ChainableMixin, InputKeys, OutputKeys, QueryComponent, validate_and_convert_stringable\n'), ((12323, 12353), 'llama_index.core.base.query_pipeline.query.InputKeys.from_keys', 'InputKeys.from_keys', (["{'input'}"], {}), "({'input'})\n", (12342, 12353), False, 'from llama_index.core.base.query_pipeline.query import ChainableMixin, InputKeys, OutputKeys, QueryComponent, validate_and_convert_stringable\n'), ((12452, 12484), 'llama_index.core.base.query_pipeline.query.OutputKeys.from_keys', 'OutputKeys.from_keys', (["{'output'}"], {}), "({'output'})\n", (12472, 12484), False, 'from llama_index.core.base.query_pipeline.query import ChainableMixin, InputKeys, OutputKeys, QueryComponent, validate_and_convert_stringable\n'), ((1491, 1508), 'llama_index.core.callbacks.base.CallbackManager', 'CallbackManager', ([], {}), '()\n', (1506, 1508), False, 'from llama_index.core.callbacks.base import CallbackManager\n'), ((2409, 2539), 'llama_index.core.utils.print_text', 'print_text', (['f"""Retrieving from object {obj.__class__.__name__} with query {query_bundle.query_str}\n"""'], {'color': '"""llama_pink"""'}), "(\n    f'Retrieving from object {obj.__class__.__name__} with query {query_bundle.query_str}\\n'\n    , color='llama_pink')\n", (2419, 2539), False, 'from llama_index.core.utils import print_text\n'), ((8276, 8297), 'llama_index.core.instrumentation.events.retrieval.RetrievalStartEvent', 'RetrievalStartEvent', ([], {}), '()\n', (8295, 8297), False, 'from llama_index.core.instrumentation.events.retrieval import RetrievalEndEvent, RetrievalStartEvent\n'), ((8375, 8407), 'llama_index.core.schema.QueryBundle', 'QueryBundle', (['str_or_query_bundle'], {}), '(str_or_query_bundle)\n', (8386, 8407), False, 'from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle, QueryType, TextNode\n'), ((8984, 9003), 'llama_index.core.instrumentation.events.retrieval.RetrievalEndEvent', 'RetrievalEndEvent', ([], {}), '()\n', (9001, 9003), False, 'from llama_index.core.instrumentation.events.retrieval import RetrievalEndEvent, RetrievalStartEvent\n'), ((9198, 9219), 'llama_index.core.instrumentation.events.retrieval.RetrievalStartEvent', 'RetrievalStartEvent', ([], {}), '()\n', (9217, 9219), False, 'from llama_index.core.instrumentation.events.retrieval import RetrievalEndEvent, RetrievalStartEvent\n'), ((9297, 9329), 'llama_index.core.schema.QueryBundle', 'QueryBundle', (['str_or_query_bundle'], {}), '(str_or_query_bundle)\n', (9308, 9329), False, 'from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle, QueryType, TextNode\n'), ((9990, 10009), 'llama_index.core.instrumentation.events.retrieval.RetrievalEndEvent', 'RetrievalEndEvent', ([], {}), '()\n', (10007, 10009), False, 'from llama_index.core.instrumentation.events.retrieval import RetrievalEndEvent, RetrievalStartEvent\n'), ((2705, 2741), 'llama_index.core.schema.NodeWithScore', 'NodeWithScore', ([], {'node': 'obj', 'score': 'score'}), '(node=obj, score=score)\n', (2718, 2741), False, 'from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle, QueryType, TextNode\n'), ((4246, 4282), 'llama_index.core.schema.NodeWithScore', 'NodeWithScore', ([], {'node': 'obj', 'score': 'score'}), '(node=obj, score=score)\n', (4259, 4282), False, 'from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle, QueryType, TextNode\n'), ((5852, 5958), 'llama_index.core.utils.print_text', 'print_text', (['f"""Retrieval entering {node.index_id}: {obj.__class__.__name__}\n"""'], {'color': '"""llama_turquoise"""'}), "(f'Retrieval entering {node.index_id}: {obj.__class__.__name__}\\n',\n    color='llama_turquoise')\n", (5862, 5958), False, 'from llama_index.core.utils import print_text\n'), ((7079, 7185), 'llama_index.core.utils.print_text', 'print_text', (['f"""Retrieval entering {node.index_id}: {obj.__class__.__name__}\n"""'], {'color': '"""llama_turquoise"""'}), "(f'Retrieval entering {node.index_id}: {obj.__class__.__name__}\\n',\n    color='llama_turquoise')\n", (7089, 7185), False, 'from llama_index.core.utils import print_text\n'), ((3794, 3822), 'llama_index.core.schema.TextNode', 'TextNode', ([], {'text': 'result_output'}), '(text=result_output)\n', (3802, 3822), False, 'from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle, QueryType, TextNode\n'), ((5240, 5268), 'llama_index.core.schema.TextNode', 'TextNode', ([], {'text': 'result_output'}), '(text=result_output)\n', (5248, 5268), False, 'from llama_index.core.schema import BaseNode, IndexNode, NodeWithScore, QueryBundle, QueryType, TextNode\n')]