def as_dict()

in google/generativeai/notebook/lib/llmfn_outputs.py [0:0]

• 34 lines of code
• 18 McCabe index (conditional complexity)

    def as_dict(self) -> Mapping[str, Sequence[Any]]:
        """Formats returned results as dictionary."""

        # `data` is a table in column order, with the columns listed from left to
        # right.
        data = {
            ColumnNames.PROMPT_NUM: [],
            ColumnNames.INPUT_NUM: [],
            # RESULT_NUM is special: each LLMFnOutputRow in self._outputs is
            # expected to have a RESULT_NUM key.
            ColumnNames.RESULT_NUM: [],
        }
        if _has_model_input_field(self._outputs):
            data[ColumnNames.MODEL_INPUT] = []

        if not self._outputs:
            return data

        # Add column names of added data.
        # The last key in LLMFnOutputRow is special as it is considered
        # the result. To preserve order in the (unlikely) event of inconsistent
        # keys across rows, we first add all-but-the-last key to `total_keys_set`,
        # then the last key.
        # Note: `total_keys_set` is a Python dictionary instead of a Python set
        # because Python dictionaries preserve the order in which entries are
        # added, whereas Python sets do not.
        total_keys_set: dict[str, None] = {k: None for k in data.keys()}
        for output in self._outputs:
            for result in output.output_rows:
                for key in list(result.keys())[:-1]:
                    total_keys_set[key] = None
        for output in self._outputs:
            for result in output.output_rows:
                total_keys_set[list(result.keys())[-1]] = None

        # `data` represents the table as a dictionary of:
        #   column names -> list of values
        for key in total_keys_set:
            data[key] = []

        next_num_rows = 1
        for output in self._outputs:
            for result in output.output_rows:
                data[ColumnNames.PROMPT_NUM].append(output.prompt_num)
                data[ColumnNames.INPUT_NUM].append(output.input_num)
                if ColumnNames.MODEL_INPUT in data:
                    data[ColumnNames.MODEL_INPUT].append(output.model_input)

                for key, value in result.items():
                    data[key].append(value)

                # Look for empty cells and pad them with None.
                for column in data.values():
                    if len(column) < next_num_rows:
                        column.append(None)

                next_num_rows += 1

        return data