Metadata-Version: 2.4
Name: pandaspipe-llm
Version: 0.1.0
Summary: Safe DataFrame query pipeline for LLM agents. Generates instructions, validates pipelines, executes operations.
Author: Alex Moura
License: MIT
Project-URL: Homepage, https://github.com/alexmoura/pandaspipe
Project-URL: Repository, https://github.com/alexmoura/pandaspipe
Project-URL: Issues, https://github.com/alexmoura/pandaspipe/issues
Keywords: pandas,llm,pipeline,data,query,agent,safe
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Typing :: Typed
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: pandas>=2.1.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: pydantic>=2.5.0
Provides-Extra: dev
Requires-Dist: pytest>=7.0.0; extra == "dev"
Requires-Dist: faker>=20.0.0; extra == "dev"
Dynamic: license-file

# pandaspipe

A safe DataFrame query pipeline for LLM agents. **pandaspipe does not call LLMs.** It generates structured instructions that tell an LLM how to build a JSON pipeline, then validates and executes the pipeline the LLM returns. Zero LLM dependencies.

The library sits between your LLM and your data: it provides the prompt context, enforces a strict query language, and runs only allowlisted operations against a pandas DataFrame.

## Install

```bash
pip install pandaspipe
```

Or for development:

```bash
git clone https://github.com/your-org/pandaspipe.git
cd pandaspipe
pip install -e ".[dev]"
```

**Requirements:** Python 3.11+, pandas >= 2.1, pydantic >= 2.5

## Quick Start

```python
import pandas as pd
import pandaspipe

df = pd.read_csv("sales.csv")

# 1. Generate a prompt to send to your LLM
prompt = pandaspipe.generate_prompt(df, instructions="Total revenue by region, top 5")

# 2. Send the prompt to any LLM (you handle this part)
pipeline_json = your_llm_call(prompt)

# 3. Validate and execute in one call
result = pandaspipe.run(df, pipeline_json)

print(result.result_df)
print(result.code_markdown)  # equivalent pandas code
```

## Two Modes: Questions and Transforms

pandaspipe handles two types of instructions:

**Analytical questions** produce summary results (aggregations, rankings, comparisons):

```python
prompt = pandaspipe.generate_prompt(df, instructions="What is the average revenue by region for VIP customers?")
```

**Transform instructions** modify the DataFrame itself (add columns, reshape, clean, split-and-recombine):

```python
prompt = pandaspipe.generate_prompt(df, instructions=(
    "For VIP customers, set discount to 25%. "
    "For Regular customers, set discount to 10%. "
    "For New customers, set discount to 5%. "
    "Combine all groups back into one table. "
    "Add a column adjusted_revenue = units * price * (1 - discount / 100). "
    "Sort by adjusted_revenue descending."
))
```

The second example demonstrates a **split-process-recombine** pattern: the LLM splits the DataFrame into three branches (one per customer type), applies different logic to each, then concatenates them back into one table. This is fully supported by the pipeline DAG.

Other transform examples:
- "Add a profit_margin column equal to (revenue - cost) / revenue"
- "Fill null values in discount with 0, then drop rows where revenue is negative"
- "Rename qty to quantity and unit_cost to cost, then drop the notes column"
- "Pivot the table: regions as rows, categories as columns, sum of revenue as values"
- "Melt the table so that Q1, Q2, Q3, Q4 columns become a single quarter column with a revenue value column"

## Core Functions

pandaspipe exposes three core functions and one convenience wrapper.

### `generate_prompt(df, instructions) -> str`

Inspects the DataFrame schema (column names, types, cardinality, sample values, null info) and assembles a complete prompt with the full query language reference, rules, and worked examples. Send this string to any capable LLM.

The `instructions` parameter accepts both **analytical questions** and **transform instructions**:

```python
# Analytical question
prompt = pandaspipe.generate_prompt(df, instructions="What are the top 10 products by revenue?")

# Transform instruction
prompt = pandaspipe.generate_prompt(df, instructions="Add a margin column (revenue - cost), fill null discounts with 0, sort by margin descending")
```

**Parameters:**
- `df` (pd.DataFrame) -- The DataFrame the pipeline will operate on.
- `instructions` (str) -- Natural-language question or transformation instructions.
- `include_examples` (bool, default True) -- Include worked examples in the prompt (covers both query and transform patterns).

### `validate_pipeline(pipeline_json, df) -> ValidationResult`

Parses the JSON, checks every operation name, verifies column references against the DataFrame, validates filter operators, aggregation functions, expression safety, step ID uniqueness, and DAG references.

```python
result = pandaspipe.validate_pipeline(pipeline_json, df)
if not result.is_valid:
    for error in result.errors:
        print(f"[{error.step_id}] {error.message}")
else:
    steps = result.pipeline  # list of PipelineStep objects
```

**Accepts:** JSON string, `dict` with a `"pipeline"` key, or a `list` of step dicts. Automatically strips markdown code fences if present.

### `execute_pipeline(df, pipeline, generate_code=True) -> PipelineResult`

Runs a validated pipeline (list of `PipelineStep` objects) against a DataFrame. Returns the result DataFrame, step-by-step execution logs, and optionally the equivalent pandas Python code.

```python
result = pandaspipe.execute_pipeline(df, validated.pipeline)
print(result.result_df)       # final DataFrame
print(result.steps_log)       # per-step row counts and columns
print(result.code_markdown)   # reproducible pandas code
```

### `run(df, pipeline_json, generate_code=True) -> PipelineResult`

Convenience function that calls `validate_pipeline` then `execute_pipeline`. Raises `ValueError` if validation fails.

```python
result = pandaspipe.run(df, '[{"id": "s1", "op": "head", "n": 10}]')
```

## Pipeline Format

A pipeline is a JSON array of step objects. Every step must have:

- `"id"` -- unique string identifier (e.g., `"s1"`, `"s2"`)
- `"op"` -- operation name from the supported set
- Additional parameters specific to the operation

Steps execute sequentially by default: each step receives the output of the previous step. To create branching pipelines, set `"input"` to a specific step ID or `"source"` (the original DataFrame).

```json
[
  {"id": "s1", "op": "filter", "column": "region", "operator": "==", "value": "South"},
  {"id": "s2", "op": "groupby", "by": "category", "agg": {"revenue": "sum"}},
  {"id": "s3", "op": "sort", "by": "revenue", "ascending": false},
  {"id": "s4", "op": "head", "n": 10}
]
```

## Operation Reference

pandaspipe supports 25 operations organized into 9 categories.

### Row Selection

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `filter` | Subset rows by condition | `column`, `operator`, `value` (simple) or `logic`, `conditions` (compound) |
| `limit` | Keep first N rows | `n` (required, int) |
| `head` | Keep first N rows | `n` (optional, default 5) |
| `tail` | Keep last N rows | `n` (optional, default 5) |
| `sample` | Random sample of rows | `n` or `frac`, `random_state` |
| `drop_duplicates` | Remove duplicate rows | `subset` (list), `keep` ("first", "last", or false) |

### Column Operations

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `select` | Keep only specified columns | `columns` (required, list) |
| `drop` | Remove specified columns | `columns` (required, list) |
| `rename` | Rename columns | `mapping` (required, dict of old -> new) |
| `compute` | New column from arithmetic expression | `target` (required, str), `expression` (required, str) |
| `assign` | Add columns with constant values | `columns` (required, dict of name -> value) |
| `astype` | Cast column types | `mapping` (required, dict of column -> dtype) |

### Sorting

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `sort` | Sort by one or more columns | `by` (required, str or list), `ascending` (optional, bool or list) |

### Aggregation

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `groupby` | Group and aggregate | `by` (required, str or list), `agg` (required, dict of column -> function) |
| `agg` | Aggregate entire DataFrame (single row) | `agg` (required, dict of column -> function) |
| `value_counts` | Count unique values | `column` (required), `normalize` (optional), `top_n` (optional) |

### Reshaping

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `pivot_table` | Create pivot table (with aggregation) | `index`, `columns`, `values` (all required), `aggfunc`, `fill_value` |
| `pivot` | Reshape long to wide (no aggregation) | `index`, `columns`, `values` (all required) |
| `melt` | Reshape wide to long (unpivot) | `id_vars` (required), `value_vars`, `var_name`, `value_name` |

### Statistics

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `describe` | Summary statistics | `columns` (optional), `percentiles` (optional) |
| `correlation` | Pairwise correlation matrix | `columns` (optional), `method` ("pearson", "kendall", "spearman") |

### Combining

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `concat` | Concatenate pipeline branches | `inputs` (required, list of step IDs), `axis` (0 or 1), `ignore_index` |
| `merge` | Join two pipeline branches | `right` (required, step ID), `on`/`left_on`/`right_on`, `how` |

### Missing Data

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `fillna` | Fill missing values | `column` (optional), `value` or `method` ("ffill", "bfill", "mean", "median") |

### Utility

| Operation | Description | Key Parameters |
|-----------|-------------|----------------|
| `reset_index` | Reset DataFrame index | `drop` (optional, bool) |

## Branching Pipelines

Pipelines support DAG-style branching. Use the `"input"` field to point a step at any prior step instead of the default (previous step). Use `"source"` to reference the original DataFrame.

```json
[
  {"id": "s1", "op": "filter", "column": "region", "operator": "==", "value": "North"},
  {"id": "s2", "op": "agg", "agg": {"revenue": "sum"}, "input": "s1"},

  {"id": "s3", "op": "filter", "column": "region", "operator": "==", "value": "South", "input": "source"},
  {"id": "s4", "op": "agg", "agg": {"revenue": "sum"}, "input": "s3"},

  {"id": "s5", "op": "concat", "inputs": ["s2", "s4"]}
]
```

This pipeline computes total revenue for North and South independently, then stacks the results. Steps s1-s2 form one branch, s3-s4 form another (starting from `"source"`), and s5 combines them.

### Split-Process-Recombine

A common transform pattern: split the DataFrame by a categorical column, apply different logic to each segment, then recombine.

```json
[
  {"id": "vip", "op": "filter", "input": "source", "column": "type", "operator": "==", "value": "VIP"},
  {"id": "vip_mod", "op": "assign", "input": "vip", "columns": {"discount": 25}},

  {"id": "reg", "op": "filter", "input": "source", "column": "type", "operator": "==", "value": "Regular"},
  {"id": "reg_mod", "op": "assign", "input": "reg", "columns": {"discount": 10}},

  {"id": "combined", "op": "concat", "inputs": ["vip_mod", "reg_mod"]},
  {"id": "result", "op": "compute", "input": "combined", "target": "net_revenue",
   "expression": "revenue * (1 - discount / 100)"}
]
```

This splits by customer type, applies different discount rates, recombines, and computes a derived column on the combined result. The output has the same number of rows as the original (VIP + Regular).

For joins, use `merge` with the `right` parameter pointing at a step ID:

```json
[
  {"id": "s1", "op": "groupby", "by": "customer_id", "agg": {"revenue": "sum"}},
  {"id": "s2", "op": "groupby", "by": "customer_id", "agg": {"orders": "count"}, "input": "source"},
  {"id": "s3", "op": "merge", "input": "s1", "right": "s2", "on": "customer_id", "how": "inner"}
]
```

## Filter Operators

The `filter` operation supports 14 comparison operators:

| Operator | Description | Value |
|----------|-------------|-------|
| `==` | Equal | scalar |
| `!=` | Not equal | scalar |
| `>` | Greater than | scalar |
| `<` | Less than | scalar |
| `>=` | Greater than or equal | scalar |
| `<=` | Less than or equal | scalar |
| `in` | Value is in a list | list |
| `not_in` | Value is not in a list | list |
| `contains` | String contains substring (case-insensitive) | string |
| `startswith` | String starts with prefix | string |
| `endswith` | String ends with suffix | string |
| `isnull` | Value is null | (none needed) |
| `notnull` | Value is not null | (none needed) |
| `between` | Value is between low and high (inclusive) | `[low, high]` |

Compound filters combine multiple conditions with `"logic": "and"` or `"logic": "or"`:

```json
{
  "id": "s1", "op": "filter",
  "logic": "and",
  "conditions": [
    {"column": "age", "operator": ">=", "value": 18},
    {"column": "status", "operator": "==", "value": "active"}
  ]
}
```

## Aggregation Functions

The following functions are available in `groupby` and `agg` operations:

`sum`, `mean`, `median`, `min`, `max`, `count`, `nunique`, `std`, `var`, `first`, `last`

Multiple aggregations per column are supported by passing a list:

```json
{"id": "s1", "op": "groupby", "by": "region", "agg": {"revenue": ["sum", "mean"], "orders": "count"}}
```

## Compute Expressions

The `compute` operation creates new columns from arithmetic expressions on existing columns.

**Allowed:** column names, numeric literals, arithmetic operators (`+`, `-`, `*`, `/`, `//`, `%`, `**`), unary minus, parentheses.

**Blocked:** function calls, attribute access, imports, string operations, `eval`, `exec`, `__dunder__` access, and all other Python constructs.

Expressions are validated at two levels:
1. The AST is parsed and every node is checked against an allowlist of safe types.
2. A regex scan rejects patterns like `__`, `import`, `exec`, `eval`, `open`, `getattr`, etc.
3. At execution time, `eval` runs with `__builtins__` set to an empty dict, and only column Series are available in the namespace.

```json
{"id": "s1", "op": "compute", "target": "profit_margin", "expression": "(revenue - cost) / revenue"}
```

## Code Generation

Every `PipelineResult` includes a `code_markdown` field containing the equivalent pandas Python code as a markdown code block. This lets users inspect, audit, or reuse the generated logic.

```python
result = pandaspipe.run(df, pipeline_json)
print(result.code_markdown)
```

Output:

````
```python
import pandas as pd

# df = pd.read_csv('your_data.csv')  # load your DataFrame

# Step: s1 (filter)
df_s1 = df[(df['region'] == 'South')]

# Step: s2 (groupby)
df_s2 = df_s1.groupby('category', as_index=False).agg({'revenue': 'sum'})

# Step: s3 (sort)
df_s3 = df_s2.sort_values(by='revenue', ascending=False)

result = df_s3
```
````

Set `generate_code=False` in `run()` or `execute_pipeline()` to skip code generation.

## Safety

pandaspipe is designed to be safe for use in production LLM agent systems:

- **No eval/exec on untrusted code.** The `compute` operation uses AST analysis to allow only arithmetic, then evaluates with an empty `__builtins__` namespace.
- **Allowlisted operations.** Only the 25 documented operations are accepted. Unknown operations are rejected at validation time.
- **Expression validation.** Compute expressions are parsed into an AST. Only `BinOp`, `UnaryOp`, `Constant`, and `Name` nodes are permitted. Function calls, attribute access, imports, and comprehensions are all rejected.
- **Forbidden pattern scanning.** A regex layer blocks `__`, `import`, `exec`, `eval`, `open`, `getattr`, `setattr`, `globals`, `locals`, `compile`, `breakpoint`, and similar patterns.
- **Column name validation.** Every column reference in every operation is checked against the actual DataFrame schema. Invented column names are caught before execution.
- **Step ID and DAG validation.** Duplicate step IDs and references to nonexistent steps are caught during validation.
- **Row limits.** The `head`, `tail`, and `limit` operations are capped at 10,000 rows to prevent accidental memory issues.
- **JSON fence stripping.** The validator automatically strips markdown code fences from LLM responses before parsing.

## Integration Examples

### With Claude (Anthropic SDK)

```python
import anthropic
import pandas as pd
import pandaspipe

client = anthropic.Anthropic()
df = pd.read_csv("sales.csv")

prompt = pandaspipe.generate_prompt(df, instructions="Top 10 products by revenue")

response = client.messages.create(
    model="claude-sonnet-4-20250514",
    max_tokens=1024,
    messages=[{"role": "user", "content": prompt}],
)

pipeline_json = response.content[0].text
result = pandaspipe.run(df, pipeline_json)
print(result.result_df)
```

### With Any LLM (Generic)

```python
import pandaspipe

df = load_your_data()
prompt = pandaspipe.generate_prompt(df, instructions="Average salary by department")

# Use whatever LLM client you have
pipeline_json = call_your_llm(prompt)

# Validate first if you want to handle errors gracefully
validation = pandaspipe.validate_pipeline(pipeline_json, df)
if not validation.is_valid:
    # Re-prompt the LLM with the error messages
    error_feedback = "\n".join(e.message for e in validation.errors)
    pipeline_json = call_your_llm(prompt + f"\n\nPrevious attempt had errors:\n{error_feedback}")
    validation = pandaspipe.validate_pipeline(pipeline_json, df)

if validation.is_valid:
    result = pandaspipe.execute_pipeline(df, validation.pipeline)
    print(result.result_df)
```

### As a Tool in an Agent Framework

```python
import pandaspipe

def query_dataframe(df, instructions: str) -> dict:
    """Tool function for an LLM agent to query or transform a DataFrame."""
    prompt = pandaspipe.generate_prompt(df, instructions=instructions)

    # Agent sends prompt to LLM and gets pipeline JSON back
    pipeline_json = agent.call_llm(prompt)

    try:
        result = pandaspipe.run(df, pipeline_json)
        return {
            "data": result.result_df.to_dict(orient="records"),
            "code": result.code_markdown,
            "rows": len(result.result_df),
            "success": result.success,
        }
    except ValueError as e:
        return {"error": str(e), "success": False}
```

## API Reference

### `PipelineResult`

Returned by `run()` and `execute_pipeline()`.

| Field | Type | Description |
|-------|------|-------------|
| `result_df` | pd.DataFrame | The final result DataFrame |
| `code_markdown` | str | Equivalent pandas code as a markdown code block |
| `steps_log` | list[StepLog] | Per-step execution log (rows in/out, columns) |
| `success` | bool | Whether all steps completed without error |
| `error` | str or None | Error message if a step failed |

On partial failure, `result_df` contains the output of the last successful step.

### `ValidationResult`

Returned by `validate_pipeline()`.

| Field | Type | Description |
|-------|------|-------------|
| `is_valid` | bool | Whether the pipeline passed all checks |
| `errors` | list[PipelineValidationError] | List of validation errors |
| `pipeline` | list[PipelineStep] | Parsed pipeline steps (populated even if invalid) |

### `PipelineStep`

A single step in a validated pipeline.

| Field | Type | Description |
|-------|------|-------------|
| `id` | str | Unique step identifier |
| `op` | str | Operation name |
| `input` | str or None | Input step ID (None = previous step, "source" = original DataFrame) |
| `params` | dict | All operation-specific parameters |

### `PipelineValidationError`

A single validation error.

| Field | Type | Description |
|-------|------|-------------|
| `step_id` | str or None | Which step caused the error |
| `field` | str or None | Which field is invalid |
| `message` | str | Human-readable error description |

### `StepLog`

Execution log entry for a single step.

| Field | Type | Description |
|-------|------|-------------|
| `step_id` | str | Step identifier |
| `op` | str | Operation name |
| `rows_in` | int | Number of input rows |
| `rows_out` | int | Number of output rows |
| `columns_out` | list[str] | Column names in the output |

## Allowed Cast Types

For the `astype` operation, these target types are accepted:

`int`, `int32`, `int64`, `float`, `float32`, `float64`, `str`, `string`, `bool`, `datetime64[ns]`, `category`

## Demo Notebooks

Two notebooks are provided in the `notebooks/` directory:

| Notebook | Description |
|----------|-------------|
| `01_demo.ipynb` | Mock LLM responses. Tests 11 scenarios (queries, transforms, branching, split-process-recombine) with assertion checks against expected pandas results. No API key needed. |
| `02_live_anthropic.ipynb` | Live Anthropic API calls. Same scenarios as the mock demo, but with real Claude responses. Each cell includes verification assertions. Requires an Anthropic API key. |

## License

MIT
