Data Analyst

Data Analyst

에이전트 워크플로에서는 때로 도구 출력의 정확한 내용이 아니라, 그 출력을 특정 방식으로 처리한 결과만 필요한 경우가 있어요. 특히 데이터 분석에서 흔한 패턴이죠. 에이전트는 쿼리 도구의 결과가 특정 이름의 컬럼을 가진 DataFrame이라는 것만 알면 되지, 모든 행의 내용까지 알 필요는 없어요.

Pydantic AI에서는 의존성 객체에 한 도구의 결과를 저장해 두고, 다른 도구에서 그걸 사용할 수 있어요. 이 예시는 그 패턴을 보여주면서, Cornell의 Rotten Tomatoes 영화 리뷰 데이터셋을 분석하는 에이전트를 만들어요.

출처: 공식문서

예시가 보여주는 것

예시 실행하기

의존성을 설치하고 환경 변수를 설정했다면, 실행해요.

python -m pydantic_ai_examples.data_analyst

uv run -m pydantic_ai_examples.data_analyst

출력 (debug):

Based on my analysis of the Cornell Movie Review dataset (rotten_tomatoes), there are 4,265 negative comments in the training split. These are the reviews labeled as 'neg' (represented by 0 in the dataset).

예시 코드

data_analyst.py

from dataclasses import dataclass, field

import datasets
import duckdb
import pandas as pd

from pydantic_ai import Agent, ModelRetry, RunContext


@dataclass
class AnalystAgentDeps:
    output: dict[str, pd.DataFrame] = field(default_factory=dict[str, pd.DataFrame])

    def store(self, value: pd.DataFrame) -> str:
        """Store the output in deps and return the reference such as Out[1] to be used by the LLM."""
        ref = f'Out[{len(self.output) + 1}]'
        self.output[ref] = value
        return ref

    def get(self, ref: str) -> pd.DataFrame:
        if ref not in self.output:
            raise ModelRetry(
                f'Error: {ref} is not a valid variable reference. Check the previous messages and try again.'
            )
        return self.output[ref]


analyst_agent = Agent(
    'openai:gpt-5.2',
    deps_type=AnalystAgentDeps,
    instructions='You are a data analyst and your job is to analyze the data according to the user request.',
)


@analyst_agent.tool
def load_dataset(
    ctx: RunContext[AnalystAgentDeps],
    path: str,
    split: str = 'train',
) -> str:
    """Load the `split` of dataset `dataset_name` from huggingface.

    Args:
        ctx: Pydantic AI agent RunContext
        path: name of the dataset in the form of `<user_name>/<dataset_name>`
        split: load the split of the dataset (default: "train")
    """
    # begin load data from hf
    builder = datasets.load_dataset_builder(path)  # pyright: ignore[reportUnknownMemberType]
    splits: dict[str, datasets.SplitInfo] = builder.info.splits or {}
    if split not in splits:
        raise ModelRetry(
            f'{split} is not valid for dataset {path}. Valid splits are {",".join(splits.keys())}'
        )

    builder.download_and_prepare()  # pyright: ignore[reportUnknownMemberType]
    dataset = builder.as_dataset(split=split)
    assert isinstance(dataset, datasets.Dataset)
    dataframe = dataset.to_pandas()
    assert isinstance(dataframe, pd.DataFrame)
    # end load data from hf

    # store the dataframe in the deps and get a ref like "Out[1]"
    ref = ctx.deps.store(dataframe)
    # construct a summary of the loaded dataset
    output = [
        f'Loaded the dataset as `{ref}`.',
        f'Description: {dataset.info.description}'
        if dataset.info.description
        else None,
        f'Features: {dataset.info.features!r}' if dataset.info.features else None,
    ]
    return '\n'.join(filter(None, output))


@analyst_agent.tool
def run_duckdb(ctx: RunContext[AnalystAgentDeps], dataset: str, sql: str) -> str:
    """Run DuckDB SQL query on the DataFrame.

    Note that the virtual table name used in DuckDB SQL must be `dataset`.

    Args:
        ctx: Pydantic AI agent RunContext
        dataset: reference string to the DataFrame
        sql: the query to be executed using DuckDB
    """
    data = ctx.deps.get(dataset)
    result = duckdb.query_df(df=data, virtual_table_name='dataset', sql_query=sql)
    # pass the result as ref (because DuckDB SQL can select many rows, creating another huge dataframe)
    ref = ctx.deps.store(result.df())
    return f'Executed SQL, result is `{ref}`'


@analyst_agent.tool
def display(ctx: RunContext[AnalystAgentDeps], name: str) -> str:
    """Display at most 5 rows of the dataframe."""
    dataset = ctx.deps.get(name)
    return dataset.head().to_string()  # pyright: ignore[reportUnknownMemberType]


if __name__ == '__main__':
    deps = AnalystAgentDeps()
    result = analyst_agent.run_sync(
        user_prompt='Count how many negative comments are there in the dataset `cornell-movie-review-data/rotten_tomatoes`',
        deps=deps,
    )
    print(result.output)

부록

모델 선택하기

이 예시는 DuckDB SQL을 이해하는 모델이 필요해요. clai로 확인할 수 있어요.

> clai -m bedrock:us.anthropic.claude-sonnet-4-5-20250929-v1:0
clai - Pydantic AI CLI v0.0.1.dev920+41dd069 with bedrock:us.anthropic.claude-sonnet-4-5-20250929-v1:0
clai ➤ do you understand duckdb sql?
# DuckDB SQL

Yes, I understand DuckDB SQL. DuckDB is an in-process analytical SQL database
that uses syntax similar to PostgreSQL. It specializes in analytical queries
and is designed for high-performance analysis of structured data.

Some key features of DuckDB SQL include:

 • OLAP (Online Analytical Processing) optimized
 • Columnar-vectorized query execution
 • Standard SQL support with PostgreSQL compatibility
 • Support for complex analytical queries
 • Efficient handling of CSV/Parquet/JSON files

I can help you with DuckDB SQL queries, schema design, optimization, or other
DuckDB-related questions.

이 예시의 핵심 아이디어

이 예시가 보여주는 가장 중요한 패턴은 의존성 객체를 "노트북 변수 저장소"처럼 쓰는 것이에요. AnalystAgentDeps.output은 Out[1], Out[2] 같은 참조 문자열을 키로, DataFrame을 값으로 보관해요.

흐름은 이렇게 진행돼요. load_dataset 도구가 데이터셋을 불러와 ctx.deps.store(dataframe)으로 저장하고, "Loaded the dataset as Out[1]" 같은 요약을 반환해요. 그러면 LLM은 Out[1]이라는 참조를 알고, run_duckdb 도구에 dataset='Out[1]'을 넘겨 DuckDB SQL을 실행해요. 그 결과도 다시 저장되어 Out[2]가 돼요.

store 대신 get을 부를 때, 없는 참조를 요청하면 ModelRetry를 던져 LLM이 이전 메시지를 다시 확인하도록 유도해요. 이렇게 하면 에이전트가 거대한 DataFrame 전체를 컨텍스트에 넣지 않고도, 그 결과를 참조하며 여러 단계로 분석을 이어갈 수 있어요.

더 알아보기 (Learn more)