基于 Pandas 数据框的查询管道¶

这是一个简单示例，演示如何构建可在 Pandas 数据框上执行结构化操作的查询管道，通过使用大语言模型（LLM）来推断操作集合，从而满足用户查询需求。

这可以视为我们 PandasQueryEngine 的"从零开始"实现版本。

警告：该工具允许 LLM 访问 eval 函数。在运行此工具的机器上可能存在任意代码执行风险。不建议在生产环境中使用此工具，若需使用则必须进行严格的沙箱隔离或虚拟机部署。

In [ ]:

%pip install llama-index-llms-openai llama-index-experimental

%pip install llama-index-llms-openai llama-index-experimental

In [ ]:

from llama_index.core.query_pipeline import (
    QueryPipeline as QP,
    Link,
    InputComponent,
)
from llama_index.experimental.query_engine.pandas import (
    PandasInstructionParser,
)
from llama_index.llms.openai import OpenAI
from llama_index.core import PromptTemplate

from llama_index.core.query_pipeline import ( QueryPipeline as QP, Link, InputComponent, ) from llama_index.experimental.query_engine.pandas import ( PandasInstructionParser, ) from llama_index.llms.openai import OpenAI from llama_index.core import PromptTemplate

下载数据¶

此处我们加载泰坦尼克号CSV数据集。

In [ ]:

!wget 'https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/docs/examples/data/csv/titanic_train.csv' -O 'titanic_train.csv'

!wget 'https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/docs/examples/data/csv/titanic_train.csv' -O 'titanic_train.csv'

--2024-01-13 18:39:07--  https://raw.githubusercontent.com/jerryjliu/llama_index/main/docs/docs/examples/data/csv/titanic_train.csv
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 2606:50c0:8003::154, 2606:50c0:8001::154, 2606:50c0:8002::154, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|2606:50c0:8003::154|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 57726 (56K) [text/plain]
Saving to: ‘titanic_train.csv’

titanic_train.csv   100%[===================>]  56.37K  --.-KB/s    in 0.007s  

2024-01-13 18:39:07 (7.93 MB/s) - ‘titanic_train.csv’ saved [57726/57726]

In [ ]:

import pandas as pd

df = pd.read_csv("./titanic_train.csv")

import pandas as pd df = pd.read_csv("./titanic_train.csv")

定义模块¶

此处我们定义以下模块集合：

1. Pandas 提示模块：用于从用户查询中推断 pandas 操作指令
2. Pandas 输出解析器：在数据框上执行 pandas 指令并返回结果数据框
3. 响应合成提示模块：根据数据框生成最终响应
4. 大语言模型（LLM）

特别设计的 pandas 输出解析器能够安全执行 Python 代码。它包含大量安全检查机制，这些机制若从零开始编写会相当繁琐。包括仅允许从预设的受信任模块导入（例如禁止使用可能修改文件系统的模块如 os），同时确保不会调用任何私有/双下划线方法。

In [ ]:

instruction_str = (
    "1. Convert the query to executable Python code using Pandas.\n"
    "2. The final line of code should be a Python expression that can be called with the `eval()` function.\n"
    "3. The code should represent a solution to the query.\n"
    "4. PRINT ONLY THE EXPRESSION.\n"
    "5. Do not quote the expression.\n"
)

pandas_prompt_str = (
    "You are working with a pandas dataframe in Python.\n"
    "The name of the dataframe is `df`.\n"
    "This is the result of `print(df.head())`:\n"
    "{df_str}\n\n"
    "Follow these instructions:\n"
    "{instruction_str}\n"
    "Query: {query_str}\n\n"
    "Expression:"
)
response_synthesis_prompt_str = (
    "Given an input question, synthesize a response from the query results.\n"
    "Query: {query_str}\n\n"
    "Pandas Instructions (optional):\n{pandas_instructions}\n\n"
    "Pandas Output: {pandas_output}\n\n"
    "Response: "
)

pandas_prompt = PromptTemplate(pandas_prompt_str).partial_format(
    instruction_str=instruction_str, df_str=df.head(5)
)
pandas_output_parser = PandasInstructionParser(df)
response_synthesis_prompt = PromptTemplate(response_synthesis_prompt_str)
llm = OpenAI(model="gpt-3.5-turbo")

instruction_str = ( "1. Convert the query to executable Python code using Pandas.\n" "2. The final line of code should be a Python expression that can be called with the `eval()` function.\n" "3. The code should represent a solution to the query.\n" "4. PRINT ONLY THE EXPRESSION.\n" "5. Do not quote the expression.\n" ) pandas_prompt_str = ( "You are working with a pandas dataframe in Python.\n" "The name of the dataframe is `df`.\n" "This is the result of `print(df.head())`:\n" "{df_str}\n\n" "Follow these instructions:\n" "{instruction_str}\n" "Query: {query_str}\n\n" "Expression:" ) response_synthesis_prompt_str = ( "Given an input question, synthesize a response from the query results.\n" "Query: {query_str}\n\n" "Pandas Instructions (optional):\n{pandas_instructions}\n\n" "Pandas Output: {pandas_output}\n\n" "Response: " ) pandas_prompt = PromptTemplate(pandas_prompt_str).partial_format( instruction_str=instruction_str, df_str=df.head(5) ) pandas_output_parser = PandasInstructionParser(df) response_synthesis_prompt = PromptTemplate(response_synthesis_prompt_str) llm = OpenAI(model="gpt-3.5-turbo")

构建查询处理流水线¶

处理流程如下所示： 输入 query_str -> pandas_prompt -> llm1 -> pandas_output_parser -> response_synthesis_prompt -> llm2

与 response_synthesis_prompt 的额外连接：llm1 -> pandas_instructions，以及 pandas_output_parser -> pandas_output。

In [ ]:

qp = QP(
    modules={
        "input": InputComponent(),
        "pandas_prompt": pandas_prompt,
        "llm1": llm,
        "pandas_output_parser": pandas_output_parser,
        "response_synthesis_prompt": response_synthesis_prompt,
        "llm2": llm,
    },
    verbose=True,
)
qp.add_chain(["input", "pandas_prompt", "llm1", "pandas_output_parser"])
qp.add_links(
    [
        Link("input", "response_synthesis_prompt", dest_key="query_str"),
        Link(
            "llm1", "response_synthesis_prompt", dest_key="pandas_instructions"
        ),
        Link(
            "pandas_output_parser",
            "response_synthesis_prompt",
            dest_key="pandas_output",
        ),
    ]
)
# add link from response synthesis prompt to llm2
qp.add_link("response_synthesis_prompt", "llm2")

qp = QP( modules={ "input": InputComponent(), "pandas_prompt": pandas_prompt, "llm1": llm, "pandas_output_parser": pandas_output_parser, "response_synthesis_prompt": response_synthesis_prompt, "llm2": llm, }, verbose=True, ) qp.add_chain(["input", "pandas_prompt", "llm1", "pandas_output_parser"]) qp.add_links( [ Link("input", "response_synthesis_prompt", dest_key="query_str"), Link( "llm1", "response_synthesis_prompt", dest_key="pandas_instructions" ), Link( "pandas_output_parser", "response_synthesis_prompt", dest_key="pandas_output", ), ] ) # add link from response synthesis prompt to llm2 qp.add_link("response_synthesis_prompt", "llm2")

运行查询¶

In [ ]:

response = qp.run(
    query_str="What is the correlation between survival and age?",
)

response = qp.run( query_str="What is the correlation between survival and age?", )

> Running module input with input: 
query_str: What is the correlation between survival and age?

> Running module pandas_prompt with input: 
query_str: What is the correlation between survival and age?

> Running module llm1 with input: 
messages: You are working with a pandas dataframe in Python.
The name of the dataframe is `df`.
This is the result of `print(df.head())`:
   survived  pclass                                               name  ...

> Running module pandas_output_parser with input: 
input: assistant: df['survived'].corr(df['age'])

> Running module response_synthesis_prompt with input: 
query_str: What is the correlation between survival and age?
pandas_instructions: assistant: df['survived'].corr(df['age'])
pandas_output: -0.07722109457217755

> Running module llm2 with input: 
messages: Given an input question, synthesize a response from the query results.
Query: What is the correlation between survival and age?

Pandas Instructions (optional):
df['survived'].corr(df['age'])

Pandas ...

In [ ]:

print(response.message.content)

print(response.message.content)

The correlation between survival and age is -0.0772. This indicates a weak negative correlation, suggesting that as age increases, the likelihood of survival slightly decreases.