pydantic-ai-examples 0.4.10__tar.gz → 0.5.0__tar.gz

{pydantic_ai_examples-0.4.10 → pydantic_ai_examples-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-ai-examples
-Version: 0.4.10
+Version: 0.5.0
 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.5.0
+Requires-Dist: pydantic-evals==0.5.0
 Requires-Dist: python-multipart>=0.0.17
 Requires-Dist: rich>=13.9.2
 Requires-Dist: uvicorn>=0.32.0

pydantic_ai_examples-0.5.0/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 → pydantic_ai_examples-0.5.0}/pyproject.toml

@@ -60,6 +60,9 @@ dependencies = [
     "gradio>=5.9.0; python_version>'3.9'",
     "mcp[cli]>=1.4.1; python_version >= '3.10'",
     "modal>=1.0.4",
+    "duckdb>=1.3.2",
+    "datasets>=4.0.0",
+    "pandas>=2.2.3",
 ]
 
 [tool.hatch.build.targets.wheel]