def _get_env_output(
      self,
      query: str,
      original_query: environment_pb2.GetQueryResponse,
      documents: Optional[Sequence[environment_pb2.Document]] = None
  ) -> types.HistoryEntry:
    """Query the environment with a new query.

    Args:
      query:  Query (reformulation) issued by the agent.
      original_query:  The original query (and its gold answers) that the agent
        is playing.
      documents:  If set, does not actually query the environment.  This is used
        to avoid repeating requests for pre-training data for which we already
        computed the results at creation time.

    Returns:
      A `HistoryEntry` summarizing the environment's response for `query` for
      the agent.
    """
    if documents is None:
      try:
        response = self.nq_server.get_documents(
            query,
            original_query=original_query.query,
            num_documents=common_flags.NUM_DOCUMENTS_TO_RETRIEVE.value,
            num_ir_documents=common_flags.NUM_IR_DOCUMENTS_TO_RETRIEVE.value)
      except grpc.RpcError as rpc_exception:
        raise mzcore.RLEnvironmentError from rpc_exception

      docs = response.documents
    else:
      docs = documents
    ranked_docs = sorted(
        docs, key=lambda doc: doc.answer.mr_score, reverse=True)

    entry = types.HistoryEntry(
        query=query, original_query=original_query, documents=ranked_docs)

    return entry
Query the environment with a new query.

    Args:
      query:  Query (reformulation) issued by the agent.
      original_query:  The original query (and its gold answers) that the agent
        is playing.
      documents:  If set, does not actually query the environment.  This is used
        to avoid repeating requests for pre-training data for which we already
        computed the results at creation time.

    Returns:
      A `HistoryEntry` summarizing the environment's response for `query` for
      the agent.
    
language/search_agents/muzero/env.py