pydantic-ai-examples 0.4.10__py3-none-any.whl → 0.4.11__py3-none-any.whl

pydantic_ai_examples/data_analyst.py

@@ -0,0 +1,107 @@
+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)
+
+    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-4o',
+    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 {}  # pyright: ignore[reportUnknownMemberType]
+    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())  # pyright: ignore[reportUnknownMemberType]
+    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)

{pydantic_ai_examples-0.4.10.dist-info → pydantic_ai_examples-0.4.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-examples
-Version: 0.4.10
+Version: 0.4.11
 Summary: Examples of how to use Pydantic AI and what it can do.
 Author-email: Samuel Colvin <samuel@pydantic.dev>, Marcelo Trylesinski <marcelotryle@gmail.com>, David Montague <david@pydantic.dev>, Alex Hall <alex@pydantic.dev>, Douwe Maan <douwe@pydantic.dev>
 License-Expression: MIT
@@ -26,14 +26,17 @@ Classifier: Topic :: Internet
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
 Requires-Dist: asyncpg>=0.30.0
+Requires-Dist: datasets>=4.0.0
 Requires-Dist: devtools>=0.12.2
+Requires-Dist: duckdb>=1.3.2
 Requires-Dist: fastapi>=0.115.4
 Requires-Dist: gradio>=5.9.0; python_version > '3.9'
 Requires-Dist: logfire[asyncpg,fastapi,httpx,sqlite3]>=2.6
 Requires-Dist: mcp[cli]>=1.4.1; python_version >= '3.10'
 Requires-Dist: modal>=1.0.4
-Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,groq,openai,vertexai]==0.4.10
-Requires-Dist: pydantic-evals==0.4.10
+Requires-Dist: pandas>=2.2.3
+Requires-Dist: pydantic-ai-slim[ag-ui,anthropic,groq,openai,vertexai]==0.4.11
+Requires-Dist: pydantic-evals==0.4.11
 Requires-Dist: python-multipart>=0.0.17
 Requires-Dist: rich>=13.9.2
 Requires-Dist: uvicorn>=0.32.0

{pydantic_ai_examples-0.4.10.dist-info → pydantic_ai_examples-0.4.11.dist-info}/RECORD

@@ -3,6 +3,7 @@ pydantic_ai_examples/bank_support.py,sha256=vJL2zLEq19OztP1fUGG7_6cYHllvxvzkafFM
 pydantic_ai_examples/chat_app.html,sha256=90XhxrpDAT09mPVTn9edEn8PqAD-tHxWkeeMz9r_okQ,2580
 pydantic_ai_examples/chat_app.py,sha256=zrurOhwbjPDTYi9FkYZF3wMxc2AT9PFckM-xngQ5188,7105
 pydantic_ai_examples/chat_app.ts,sha256=2KfZ2rJU2o0iCPjelyqEi5sH6vfemzWaa5Evx_VcAE4,3307
+pydantic_ai_examples/data_analyst.py,sha256=vSpfIWpxNgaZqKx2DJD721ZA3QgSjaZkTixy5OHz9xY,3753
 pydantic_ai_examples/flight_booking.py,sha256=hwz1ct0laIZQCb0cuhFYyrR-PBDUL2zc4CS39ytOwn0,7454
 pydantic_ai_examples/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pydantic_ai_examples/pydantic_model.py,sha256=EQsHBig2bvb0PG_2XDgP9Le5xJ4n8eZJzQDGQYhDykg,775
@@ -43,7 +44,7 @@ pydantic_ai_examples/slack_lead_qualifier/modal.py,sha256=f464AaeyP-n3UIfvEVVc4D
 pydantic_ai_examples/slack_lead_qualifier/models.py,sha256=WTp6D2WCASXqrjPVT3vGgTSYATLPBM3_cjq9wvXMRao,1586
 pydantic_ai_examples/slack_lead_qualifier/slack.py,sha256=VJVfMeUXYClWUJBLHNuaW8PB2sxjNzpTC-O_AJwcxQ4,833
 pydantic_ai_examples/slack_lead_qualifier/store.py,sha256=04vB4eZWKk_Tx0b9K4QuVI1U24JEyJyBS4X76cui7OI,896
-pydantic_ai_examples-0.4.10.dist-info/METADATA,sha256=_TGzxgOYzsyl5TYmqWnq4gBjzIK9DzVyFVlnM6ugmm8,2757
-pydantic_ai_examples-0.4.10.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydantic_ai_examples-0.4.10.dist-info/licenses/LICENSE,sha256=vA6Jc482lEyBBuGUfD1pYx-cM7jxvLYOxPidZ30t_PQ,1100
-pydantic_ai_examples-0.4.10.dist-info/RECORD,,
+pydantic_ai_examples-0.4.11.dist-info/METADATA,sha256=Docxd8tyYNI3U9WesiQTEEx2-IjM2VEK1mK5iZFS2xI,2846
+pydantic_ai_examples-0.4.11.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydantic_ai_examples-0.4.11.dist-info/licenses/LICENSE,sha256=vA6Jc482lEyBBuGUfD1pYx-cM7jxvLYOxPidZ30t_PQ,1100
+pydantic_ai_examples-0.4.11.dist-info/RECORD,,