Pandas ai

PandasAIReader

Bases: BaseReader

Pandas AI阅读器。

对https://github.com/gventuri/pandas-ai的轻量封装。

Parameters:

Name	Type	Description	Default
llm	Optional[llm]	要使用的LLM。默认为None。	None
concat_rows	bool	是否将所有行连接成一个文档。 如果设置为False，则将为每一行创建一个文档。 默认为True。	True
col_joiner	str	用于连接每行列的分隔符。 默认设置为", "。	', '
row_joiner	str	用于连接每行的分隔符。 仅在concat_rows=True时使用。 默认设置为"\n"。	'\n'
pandas_config	dict	pandas.read_csv函数调用的选项。 有关更多信息，请参考https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html。 默认设置为空字典，这意味着pandas将尝试自行确定分隔符、表头等。	{}

Source code in llama_index/readers/pandas_ai/base.py

class PandasAIReader(BaseReader):
    r"""Pandas AI阅读器。

    对https://github.com/gventuri/pandas-ai的轻量封装。

    Args:
        llm (Optional[pandas.llm]): 要使用的LLM。默认为None。
        concat_rows (bool): 是否将所有行连接成一个文档。
            如果设置为False，则将为每一行创建一个文档。
            默认为True。

        col_joiner (str): 用于连接每行列的分隔符。
            默认设置为", "。

        row_joiner (str): 用于连接每行的分隔符。
            仅在`concat_rows=True`时使用。
            默认设置为"\n"。

        pandas_config (dict): `pandas.read_csv`函数调用的选项。
            有关更多信息，请参考https://pandas.pydata.org/docs/reference/api/pandas.read_csv.html。
            默认设置为空字典，这意味着pandas将尝试自行确定分隔符、表头等。"""

    def __init__(
        self,
        llm: Optional[Any] = None,
        concat_rows: bool = True,
        col_joiner: str = ", ",
        row_joiner: str = "\n",
        pandas_config: dict = {},
    ) -> None:
        """初始化参数。"""
        try:
            from pandasai import PandasAI
            from pandasai.llm.openai import OpenAI
        except ImportError:
            raise ImportError("Please install pandasai to use this reader.")

        self._llm = llm or OpenAI()
        self._pandas_ai = PandasAI(llm)

        self._concat_rows = concat_rows
        self._col_joiner = col_joiner
        self._row_joiner = row_joiner
        self._pandas_config = pandas_config

    def run_pandas_ai(
        self,
        initial_df: pd.DataFrame,
        query: str,
        is_conversational_answer: bool = False,
    ) -> Any:
        """加载数据框。"""
        return self._pandas_ai.run(
            initial_df, prompt=query, is_conversational_answer=is_conversational_answer
        )

    def load_data(
        self,
        initial_df: pd.DataFrame,
        query: str,
        is_conversational_answer: bool = False,
    ) -> List[Document]:
        """解析文件。"""
        result = self.run_pandas_ai(
            initial_df, query, is_conversational_answer=is_conversational_answer
        )
        if is_conversational_answer:
            return [Document(text=result)]
        else:
            if isinstance(result, (np.generic)):
                result = pd.Series(result)
            elif isinstance(result, (pd.Series, pd.DataFrame)):
                pass
            else:
                raise ValueError(f"Unexpected type for result: {type(result)}")
            # if not conversational answer, use Pandas CSV Reader
            reader = PandasCSVReader(
                concat_rows=self._concat_rows,
                col_joiner=self._col_joiner,
                row_joiner=self._row_joiner,
                pandas_config=self._pandas_config,
            )

            with TemporaryDirectory() as tmpdir:
                outpath = Path(tmpdir) / "out.csv"
                with outpath.open("w") as f:
                    # TODO: add option to specify index=False
                    result.to_csv(f, index=False)

                return reader.load_data(outpath)

run_pandas_ai

run_pandas_ai(
    initial_df: DataFrame,
    query: str,
    is_conversational_answer: bool = False,
) -> Any

加载数据框。

Source code in llama_index/readers/pandas_ai/base.py

def run_pandas_ai(
    self,
    initial_df: pd.DataFrame,
    query: str,
    is_conversational_answer: bool = False,
) -> Any:
    """加载数据框。"""
    return self._pandas_ai.run(
        initial_df, prompt=query, is_conversational_answer=is_conversational_answer
    )

load_data

load_data(
    initial_df: DataFrame,
    query: str,
    is_conversational_answer: bool = False,
) -> List[Document]

解析文件。

Source code in llama_index/readers/pandas_ai/base.py

def load_data(
    self,
    initial_df: pd.DataFrame,
    query: str,
    is_conversational_answer: bool = False,
) -> List[Document]:
    """解析文件。"""
    result = self.run_pandas_ai(
        initial_df, query, is_conversational_answer=is_conversational_answer
    )
    if is_conversational_answer:
        return [Document(text=result)]
    else:
        if isinstance(result, (np.generic)):
            result = pd.Series(result)
        elif isinstance(result, (pd.Series, pd.DataFrame)):
            pass
        else:
            raise ValueError(f"Unexpected type for result: {type(result)}")
        # if not conversational answer, use Pandas CSV Reader
        reader = PandasCSVReader(
            concat_rows=self._concat_rows,
            col_joiner=self._col_joiner,
            row_joiner=self._row_joiner,
            pandas_config=self._pandas_config,
        )

        with TemporaryDirectory() as tmpdir:
            outpath = Path(tmpdir) / "out.csv"
            with outpath.open("w") as f:
                # TODO: add option to specify index=False
                result.to_csv(f, index=False)

            return reader.load_data(outpath)