Polars 查询引擎

!pip install llama-index llama-index-experimental polars

import logging
import sys
from IPython.display import Markdown, display

import polars as pl
from llama_index.experimental.query_engine import PolarsQueryEngine


logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

# Test on some sample data
df = pl.DataFrame(
    {
        "city": ["Toronto", "Tokyo", "Berlin"],
        "population": [2930000, 13960000, 3645000],
    }
)

query_engine = PolarsQueryEngine(df=df, verbose=True)

response = query_engine.query(
    "What is the city with the highest population?",
)

display(Markdown(f"<b>{response}</b>"))

# get polars python instructions
print(response.metadata["polars_instruction_str"])

query_engine = PolarsQueryEngine(df=df, verbose=True, synthesize_response=True)
response = query_engine.query(
    "What is the city with the highest population? Give both the city and population",
)
print(str(response))

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

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

query_engine = PolarsQueryEngine(df=df, verbose=True)

response = query_engine.query(
    "What is the correlation between survival and age?",
)

display(Markdown(f"<b>{response}</b>"))

from llama_index.core import PromptTemplate

query_engine = PolarsQueryEngine(df=df, verbose=True)
prompts = query_engine.get_prompts()
print(prompts["polars_prompt"].template)

print(prompts["response_synthesis_prompt"].template)

new_prompt = PromptTemplate(
    """\
You are working with a polars dataframe in Python.
The name of the dataframe is `df`.
This is the result of `print(df.head())`:
{df_str}

Follow these instructions:
{instruction_str}
Query: {query_str}

Expression: """
)

query_engine.update_prompts({"polars_prompt": new_prompt})

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