def get_query(
      self,
      index: Optional[int] = None,
      dataset_type: str = 'TRAIN') -> environment_pb2.GetQueryResponse:
    """Get query by index.

    Args:
      index: int, Index of the query.
      dataset_type: str, Dataset to choose from in ['TRAIN', 'DEV', 'TEST'].

    Returns:
      The corresponding query as enviornment_pb2.Query.
    """
    dataset = environment_pb2.DataSet.Value(common_flags.DATASET.value)
    if index:
      req = environment_pb2.GetQueryRequest(
          index=index, dataset=dataset, dataset_type=dataset_type)
    else:
      req = environment_pb2.GetQueryRequest(
          dataset=dataset, dataset_type=dataset_type)
    query = self._call_rpc(self._stub.GetQuery, req)
    return query
Get query by index.

    Args:
      index: int, Index of the query.
      dataset_type: str, Dataset to choose from in ['TRAIN', 'DEV', 'TEST'].

    Returns:
      The corresponding query as enviornment_pb2.Query.
    
language/search_agents/muzero/server.py