airow 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- airow-0.1.0/LICENSE +21 -0
- airow-0.1.0/PKG-INFO +121 -0
- airow-0.1.0/README.md +77 -0
- airow-0.1.0/airow/__init__.py +11 -0
- airow-0.1.0/airow/agent.py +71 -0
- airow-0.1.0/airow/airow.py +110 -0
- airow-0.1.0/airow/schemas.py +25 -0
- airow-0.1.0/airow.egg-info/PKG-INFO +121 -0
- airow-0.1.0/airow.egg-info/SOURCES.txt +13 -0
- airow-0.1.0/airow.egg-info/dependency_links.txt +1 -0
- airow-0.1.0/airow.egg-info/requires.txt +12 -0
- airow-0.1.0/airow.egg-info/top_level.txt +1 -0
- airow-0.1.0/pyproject.toml +60 -0
- airow-0.1.0/setup.cfg +4 -0
- airow-0.1.0/tests/test_agents.py +189 -0
airow-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Dmitrii K
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
airow-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: airow
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: AI-powered DataFrame processing made simple
|
|
5
|
+
Author-email: Dmitrii K <dmitriik@protonmail.com>
|
|
6
|
+
Maintainer-email: Dmitrii K <dmitriik@protonmail.com>
|
|
7
|
+
License: MIT
|
|
8
|
+
Project-URL: Homepage, https://github.com/dmitriiweb/airow
|
|
9
|
+
Project-URL: Repository, https://github.com/dmitriiweb/airow
|
|
10
|
+
Project-URL: Documentation, https://github.com/dmitriiweb/airow
|
|
11
|
+
Project-URL: Bug Tracker, https://github.com/dmitriiweb/airow/issues
|
|
12
|
+
Keywords: ai,ai-agent,dataframe,pandas,pydantic-ai,async,data-processing
|
|
13
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
14
|
+
Classifier: Intended Audience :: Developers
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Operating System :: OS Independent
|
|
18
|
+
Classifier: Programming Language :: Python :: 3
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
23
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
24
|
+
Classifier: Topic :: Scientific/Engineering :: Information Analysis
|
|
25
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
26
|
+
Classifier: Topic :: Text Processing
|
|
27
|
+
Classifier: Topic :: Database
|
|
28
|
+
Classifier: Typing :: Typed
|
|
29
|
+
Requires-Python: >=3.10
|
|
30
|
+
Description-Content-Type: text/markdown
|
|
31
|
+
License-File: LICENSE
|
|
32
|
+
Requires-Dist: loguru>=0.7.3
|
|
33
|
+
Requires-Dist: pandas>=2.3.2
|
|
34
|
+
Requires-Dist: pydantic>=2.11.7
|
|
35
|
+
Requires-Dist: pydantic-ai>=0.8.1
|
|
36
|
+
Requires-Dist: tqdm>=4.67.1
|
|
37
|
+
Provides-Extra: dev
|
|
38
|
+
Requires-Dist: mypy>=1.17.1; extra == "dev"
|
|
39
|
+
Requires-Dist: pytest>=8.4.2; extra == "dev"
|
|
40
|
+
Requires-Dist: pytest-asyncio>=1.1.0; extra == "dev"
|
|
41
|
+
Requires-Dist: pytest-cov>=6.3.0; extra == "dev"
|
|
42
|
+
Requires-Dist: ruff>=0.12.12; extra == "dev"
|
|
43
|
+
Dynamic: license-file
|
|
44
|
+
|
|
45
|
+
# Airow
|
|
46
|
+
|
|
47
|
+
**AI-powered DataFrame processing made simple**
|
|
48
|
+
|
|
49
|
+
Airow is a Python library that combines the power of pandas DataFrames with AI models to process structured data at scale. Built on top of `pydantic-ai`, it provides type-safe, async processing of DataFrames using any AI model.
|
|
50
|
+
|
|
51
|
+
## Features
|
|
52
|
+
|
|
53
|
+
- 🚀 **Async processing** with batch support for high performance
|
|
54
|
+
- 🔒 **Type-safe outputs** using Pydantic models
|
|
55
|
+
- 📊 **Progress tracking** with built-in progress bars
|
|
56
|
+
- 🔄 **Automatic retries** with configurable retry logic
|
|
57
|
+
- 🤖 **Flexible AI models** - works with OpenAI, Ollama, Anthropic, and more
|
|
58
|
+
- ⚡ **Parallel processing** within batches for maximum throughput
|
|
59
|
+
- 📝 **Structured outputs** with defined schemas and validation
|
|
60
|
+
|
|
61
|
+
## Installation
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
# Using pip
|
|
65
|
+
pip install airow
|
|
66
|
+
|
|
67
|
+
# Using uv (recommended)
|
|
68
|
+
uv add airow
|
|
69
|
+
|
|
70
|
+
# Using conda
|
|
71
|
+
conda install -c conda-forge airow
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Quick Start
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
import pandas as pd
|
|
78
|
+
from pydantic_ai.models.openai import OpenAIChatModel
|
|
79
|
+
from pydantic_ai.providers.ollama import OllamaProvider
|
|
80
|
+
from airow import Airow, OutputColumn
|
|
81
|
+
import asyncio
|
|
82
|
+
|
|
83
|
+
async def main():
|
|
84
|
+
# Setup your AI model
|
|
85
|
+
model = OpenAIChatModel(
|
|
86
|
+
model_name="llama3.2:latest",
|
|
87
|
+
provider=OllamaProvider(base_url="http://localhost:11434/v1"),
|
|
88
|
+
)
|
|
89
|
+
# or use strings:
|
|
90
|
+
model = "openai:gpt-5"
|
|
91
|
+
model = "anthropic:claude-sonnet-4-0"
|
|
92
|
+
|
|
93
|
+
# Create Airow instance
|
|
94
|
+
airow = Airow(
|
|
95
|
+
model=model,
|
|
96
|
+
system_prompt="You are an expert in wine tasting and selection.",
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
# Load your data
|
|
100
|
+
df = pd.read_csv("wine_data.csv")
|
|
101
|
+
|
|
102
|
+
output_columns = [
|
|
103
|
+
OutputColumn(name="sentiment", type=str, description="Positive, negative, or neutral sentiment"),
|
|
104
|
+
OutputColumn(name="confidence", type=float, description="Confidence score between 0 and 1"),
|
|
105
|
+
OutputColumn(name="keywords", type=list, description="List of key terms extracted"),
|
|
106
|
+
]
|
|
107
|
+
|
|
108
|
+
# Process with AI
|
|
109
|
+
result_df = await airow.run(
|
|
110
|
+
df,
|
|
111
|
+
prompt="Analyze the wine description and provide sentiment analysis, confidence score, and extract key terms.",
|
|
112
|
+
input_columns=["description"],
|
|
113
|
+
output_columns=output_columns,
|
|
114
|
+
show_progress=True,
|
|
115
|
+
)
|
|
116
|
+
|
|
117
|
+
print(result_df.head())
|
|
118
|
+
|
|
119
|
+
if __name__ == "__main__":
|
|
120
|
+
asyncio.run(main())
|
|
121
|
+
```
|
airow-0.1.0/README.md
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
# Airow
|
|
2
|
+
|
|
3
|
+
**AI-powered DataFrame processing made simple**
|
|
4
|
+
|
|
5
|
+
Airow is a Python library that combines the power of pandas DataFrames with AI models to process structured data at scale. Built on top of `pydantic-ai`, it provides type-safe, async processing of DataFrames using any AI model.
|
|
6
|
+
|
|
7
|
+
## Features
|
|
8
|
+
|
|
9
|
+
- 🚀 **Async processing** with batch support for high performance
|
|
10
|
+
- 🔒 **Type-safe outputs** using Pydantic models
|
|
11
|
+
- 📊 **Progress tracking** with built-in progress bars
|
|
12
|
+
- 🔄 **Automatic retries** with configurable retry logic
|
|
13
|
+
- 🤖 **Flexible AI models** - works with OpenAI, Ollama, Anthropic, and more
|
|
14
|
+
- ⚡ **Parallel processing** within batches for maximum throughput
|
|
15
|
+
- 📝 **Structured outputs** with defined schemas and validation
|
|
16
|
+
|
|
17
|
+
## Installation
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
# Using pip
|
|
21
|
+
pip install airow
|
|
22
|
+
|
|
23
|
+
# Using uv (recommended)
|
|
24
|
+
uv add airow
|
|
25
|
+
|
|
26
|
+
# Using conda
|
|
27
|
+
conda install -c conda-forge airow
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Quick Start
|
|
31
|
+
|
|
32
|
+
```python
|
|
33
|
+
import pandas as pd
|
|
34
|
+
from pydantic_ai.models.openai import OpenAIChatModel
|
|
35
|
+
from pydantic_ai.providers.ollama import OllamaProvider
|
|
36
|
+
from airow import Airow, OutputColumn
|
|
37
|
+
import asyncio
|
|
38
|
+
|
|
39
|
+
async def main():
|
|
40
|
+
# Setup your AI model
|
|
41
|
+
model = OpenAIChatModel(
|
|
42
|
+
model_name="llama3.2:latest",
|
|
43
|
+
provider=OllamaProvider(base_url="http://localhost:11434/v1"),
|
|
44
|
+
)
|
|
45
|
+
# or use strings:
|
|
46
|
+
model = "openai:gpt-5"
|
|
47
|
+
model = "anthropic:claude-sonnet-4-0"
|
|
48
|
+
|
|
49
|
+
# Create Airow instance
|
|
50
|
+
airow = Airow(
|
|
51
|
+
model=model,
|
|
52
|
+
system_prompt="You are an expert in wine tasting and selection.",
|
|
53
|
+
)
|
|
54
|
+
|
|
55
|
+
# Load your data
|
|
56
|
+
df = pd.read_csv("wine_data.csv")
|
|
57
|
+
|
|
58
|
+
output_columns = [
|
|
59
|
+
OutputColumn(name="sentiment", type=str, description="Positive, negative, or neutral sentiment"),
|
|
60
|
+
OutputColumn(name="confidence", type=float, description="Confidence score between 0 and 1"),
|
|
61
|
+
OutputColumn(name="keywords", type=list, description="List of key terms extracted"),
|
|
62
|
+
]
|
|
63
|
+
|
|
64
|
+
# Process with AI
|
|
65
|
+
result_df = await airow.run(
|
|
66
|
+
df,
|
|
67
|
+
prompt="Analyze the wine description and provide sentiment analysis, confidence score, and extract key terms.",
|
|
68
|
+
input_columns=["description"],
|
|
69
|
+
output_columns=output_columns,
|
|
70
|
+
show_progress=True,
|
|
71
|
+
)
|
|
72
|
+
|
|
73
|
+
print(result_df.head())
|
|
74
|
+
|
|
75
|
+
if __name__ == "__main__":
|
|
76
|
+
asyncio.run(main())
|
|
77
|
+
```
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
"""airow package public API and version information.
|
|
2
|
+
|
|
3
|
+
Exposes `Airow` for batched, row-wise LLM inference over pandas DataFrames
|
|
4
|
+
and `OutputColumn` for declaring structured outputs.
|
|
5
|
+
"""
|
|
6
|
+
__version__ = "0.1.0"
|
|
7
|
+
|
|
8
|
+
from .airow import Airow
|
|
9
|
+
from .schemas import OutputColumn
|
|
10
|
+
|
|
11
|
+
__all__ = ["Airow", "OutputColumn"]
|
|
@@ -0,0 +1,71 @@
|
|
|
1
|
+
"""Agent utilities for executing pydantic-ai models with structured outputs."""
|
|
2
|
+
|
|
3
|
+
from typing import Iterable
|
|
4
|
+
|
|
5
|
+
from pydantic import BaseModel, Field, create_model
|
|
6
|
+
from pydantic_ai import Agent
|
|
7
|
+
from pydantic_ai.models import Model
|
|
8
|
+
from loguru import logger
|
|
9
|
+
|
|
10
|
+
from . import schemas
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class AirowAgent:
|
|
14
|
+
"""Wrapper around `pydantic_ai.Agent` that builds structured output models."""
|
|
15
|
+
def __init__(
|
|
16
|
+
self,
|
|
17
|
+
model: Model | str,
|
|
18
|
+
system_prompt: str,
|
|
19
|
+
retries: int = 3,
|
|
20
|
+
):
|
|
21
|
+
"""Initialize the agent.
|
|
22
|
+
|
|
23
|
+
Args:
|
|
24
|
+
model: The underlying model used by pydantic-ai.
|
|
25
|
+
system_prompt: System prompt applied to all runs.
|
|
26
|
+
retries: Number of retries for a run.
|
|
27
|
+
"""
|
|
28
|
+
self.model = model
|
|
29
|
+
self.system_prompt = system_prompt
|
|
30
|
+
self.agent = Agent(model=model, system_prompt=self.system_prompt, retries=retries)
|
|
31
|
+
|
|
32
|
+
async def run(
|
|
33
|
+
self,
|
|
34
|
+
prompt: str,
|
|
35
|
+
output_columns: Iterable[schemas.OutputColumn],
|
|
36
|
+
) -> dict[str, object]:
|
|
37
|
+
"""Run the agent with the given prompt and expected outputs.
|
|
38
|
+
|
|
39
|
+
Args:
|
|
40
|
+
prompt: User prompt to pass to the model.
|
|
41
|
+
output_columns: Iterable of expected output columns specifications.
|
|
42
|
+
|
|
43
|
+
Returns:
|
|
44
|
+
A dictionary mapping output column names to parsed values. Returns
|
|
45
|
+
an empty dictionary when the underlying model call fails.
|
|
46
|
+
"""
|
|
47
|
+
output_columns_fields = self.build_agent_output_type(output_columns)
|
|
48
|
+
try:
|
|
49
|
+
result = await self.agent.run(prompt, output_type=output_columns_fields)
|
|
50
|
+
except Exception as e:
|
|
51
|
+
logger.error(f"{e=}")
|
|
52
|
+
return {}
|
|
53
|
+
return result.output.model_dump()
|
|
54
|
+
|
|
55
|
+
def build_agent_output_type(
|
|
56
|
+
self,
|
|
57
|
+
output_columns: Iterable[schemas.OutputColumn],
|
|
58
|
+
) -> type[BaseModel]:
|
|
59
|
+
"""Create a `pydantic.BaseModel` for the requested output columns.
|
|
60
|
+
|
|
61
|
+
Args:
|
|
62
|
+
output_columns: Iterable of output column specifications.
|
|
63
|
+
|
|
64
|
+
Returns:
|
|
65
|
+
A dynamically created `BaseModel` subclass with fields per column.
|
|
66
|
+
"""
|
|
67
|
+
fields = {
|
|
68
|
+
col.name: (col.type, Field(..., description=col.description))
|
|
69
|
+
for col in output_columns
|
|
70
|
+
}
|
|
71
|
+
return create_model("OutputColumns", **fields)
|
|
@@ -0,0 +1,110 @@
|
|
|
1
|
+
"""High-level API to run LLMs over DataFrame rows in batches."""
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
from typing import Iterable
|
|
5
|
+
|
|
6
|
+
import pandas as pd
|
|
7
|
+
from pydantic_ai.models import Model
|
|
8
|
+
from tqdm import tqdm
|
|
9
|
+
|
|
10
|
+
from . import schemas
|
|
11
|
+
from .agent import AirowAgent
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class Airow:
|
|
15
|
+
"""Apply an LLM to each row of a DataFrame and write results.
|
|
16
|
+
|
|
17
|
+
Uses `AirowAgent` internally and supports parallel row processing per batch.
|
|
18
|
+
"""
|
|
19
|
+
def __init__(
|
|
20
|
+
self,
|
|
21
|
+
*,
|
|
22
|
+
model: Model | str,
|
|
23
|
+
system_prompt: str,
|
|
24
|
+
batch_size: int = 1,
|
|
25
|
+
retries: int = 3,
|
|
26
|
+
):
|
|
27
|
+
"""Configure the runner.
|
|
28
|
+
|
|
29
|
+
Args:
|
|
30
|
+
model: The pydantic-ai model to use.
|
|
31
|
+
system_prompt: System prompt applied to each request.
|
|
32
|
+
batch_size: Number of DataFrame rows to process concurrently.
|
|
33
|
+
retries: Number of retries for a run.
|
|
34
|
+
"""
|
|
35
|
+
self.model = model
|
|
36
|
+
self.system_prompt = system_prompt
|
|
37
|
+
self.batch_size = batch_size
|
|
38
|
+
self.agent = AirowAgent(self.model, self.system_prompt, retries)
|
|
39
|
+
|
|
40
|
+
async def run(
|
|
41
|
+
self,
|
|
42
|
+
df: pd.DataFrame,
|
|
43
|
+
*,
|
|
44
|
+
prompt: str,
|
|
45
|
+
input_columns: Iterable[str],
|
|
46
|
+
output_columns: schemas.OutputColumn | Iterable[schemas.OutputColumn],
|
|
47
|
+
show_progress: bool = False,
|
|
48
|
+
) -> pd.DataFrame:
|
|
49
|
+
"""Run the model across the DataFrame and return results.
|
|
50
|
+
|
|
51
|
+
For each row, the values from `input_columns` are appended to `prompt`
|
|
52
|
+
as labeled lines and the model is asked to produce `output_columns`.
|
|
53
|
+
Results are written into the original DataFrame.
|
|
54
|
+
|
|
55
|
+
Args:
|
|
56
|
+
df: Input DataFrame.
|
|
57
|
+
prompt: Base prompt text provided to the model.
|
|
58
|
+
input_columns: Columns whose values are passed as context to the model.
|
|
59
|
+
output_columns: One or more output column specifications.
|
|
60
|
+
show_progress: Whether to display a progress bar.
|
|
61
|
+
|
|
62
|
+
Returns:
|
|
63
|
+
The input DataFrame with new output columns populated.
|
|
64
|
+
"""
|
|
65
|
+
if isinstance(output_columns, schemas.OutputColumn):
|
|
66
|
+
output_columns = [output_columns]
|
|
67
|
+
|
|
68
|
+
# Convert to list for easier handling
|
|
69
|
+
input_columns = list(input_columns)
|
|
70
|
+
|
|
71
|
+
# Work with a copy to avoid SettingWithCopyWarning
|
|
72
|
+
df = df.copy()
|
|
73
|
+
|
|
74
|
+
# Split dataframe into batches
|
|
75
|
+
total_rows = df.shape[0]
|
|
76
|
+
batche_ranges = [
|
|
77
|
+
(i, i + self.batch_size)
|
|
78
|
+
for i in range(0, total_rows, self.batch_size)
|
|
79
|
+
]
|
|
80
|
+
if show_progress:
|
|
81
|
+
batche_ranges = tqdm(batche_ranges)
|
|
82
|
+
|
|
83
|
+
for batch_range in batche_ranges:
|
|
84
|
+
# Process each row in the batch in parallel
|
|
85
|
+
tasks = []
|
|
86
|
+
row_indices = []
|
|
87
|
+
|
|
88
|
+
# Get row indices for this batch
|
|
89
|
+
start_idx = batch_range[0]
|
|
90
|
+
end_idx = min(batch_range[1], total_rows)
|
|
91
|
+
|
|
92
|
+
for row_idx in range(start_idx, end_idx):
|
|
93
|
+
row = df.iloc[row_idx]
|
|
94
|
+
input_data = {col: row[col] for col in input_columns}
|
|
95
|
+
input_data_str = "\n".join([f"Column: {k}, Value: {v}" for k, v in input_data.items()])
|
|
96
|
+
full_prompt = f"{prompt}\n\n{input_data_str}"
|
|
97
|
+
task = self.agent.run(full_prompt, output_columns)
|
|
98
|
+
tasks.append(task)
|
|
99
|
+
row_indices.append(row_idx)
|
|
100
|
+
|
|
101
|
+
# Run all tasks in parallel
|
|
102
|
+
results = await asyncio.gather(*tasks)
|
|
103
|
+
|
|
104
|
+
# Add results to dataframe
|
|
105
|
+
for i, result in enumerate(results):
|
|
106
|
+
row_idx = row_indices[i]
|
|
107
|
+
for col_name, value in result.items():
|
|
108
|
+
df.loc[row_idx, col_name] = value
|
|
109
|
+
|
|
110
|
+
return df
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
"""Data structures for declaring model outputs."""
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass
|
|
4
|
+
from typing import Any, Type
|
|
5
|
+
|
|
6
|
+
from pydantic import BaseModel
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
@dataclass
|
|
10
|
+
class OutputColumn:
|
|
11
|
+
"""
|
|
12
|
+
Output column for the AI model.
|
|
13
|
+
|
|
14
|
+
Args:
|
|
15
|
+
name: Name of the output column.
|
|
16
|
+
type: Type of the output column.
|
|
17
|
+
|
|
18
|
+
Examples:
|
|
19
|
+
>>> OutputColumn(name="output_column", type=str)
|
|
20
|
+
>>> OutputColumn(name="output_column", type=int)
|
|
21
|
+
"""
|
|
22
|
+
|
|
23
|
+
name: str
|
|
24
|
+
type: Type[Any]
|
|
25
|
+
description: str
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: airow
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: AI-powered DataFrame processing made simple
|
|
5
|
+
Author-email: Dmitrii K <dmitriik@protonmail.com>
|
|
6
|
+
Maintainer-email: Dmitrii K <dmitriik@protonmail.com>
|
|
7
|
+
License: MIT
|
|
8
|
+
Project-URL: Homepage, https://github.com/dmitriiweb/airow
|
|
9
|
+
Project-URL: Repository, https://github.com/dmitriiweb/airow
|
|
10
|
+
Project-URL: Documentation, https://github.com/dmitriiweb/airow
|
|
11
|
+
Project-URL: Bug Tracker, https://github.com/dmitriiweb/airow/issues
|
|
12
|
+
Keywords: ai,ai-agent,dataframe,pandas,pydantic-ai,async,data-processing
|
|
13
|
+
Classifier: Development Status :: 5 - Production/Stable
|
|
14
|
+
Classifier: Intended Audience :: Developers
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Operating System :: OS Independent
|
|
18
|
+
Classifier: Programming Language :: Python :: 3
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
23
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
24
|
+
Classifier: Topic :: Scientific/Engineering :: Information Analysis
|
|
25
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
26
|
+
Classifier: Topic :: Text Processing
|
|
27
|
+
Classifier: Topic :: Database
|
|
28
|
+
Classifier: Typing :: Typed
|
|
29
|
+
Requires-Python: >=3.10
|
|
30
|
+
Description-Content-Type: text/markdown
|
|
31
|
+
License-File: LICENSE
|
|
32
|
+
Requires-Dist: loguru>=0.7.3
|
|
33
|
+
Requires-Dist: pandas>=2.3.2
|
|
34
|
+
Requires-Dist: pydantic>=2.11.7
|
|
35
|
+
Requires-Dist: pydantic-ai>=0.8.1
|
|
36
|
+
Requires-Dist: tqdm>=4.67.1
|
|
37
|
+
Provides-Extra: dev
|
|
38
|
+
Requires-Dist: mypy>=1.17.1; extra == "dev"
|
|
39
|
+
Requires-Dist: pytest>=8.4.2; extra == "dev"
|
|
40
|
+
Requires-Dist: pytest-asyncio>=1.1.0; extra == "dev"
|
|
41
|
+
Requires-Dist: pytest-cov>=6.3.0; extra == "dev"
|
|
42
|
+
Requires-Dist: ruff>=0.12.12; extra == "dev"
|
|
43
|
+
Dynamic: license-file
|
|
44
|
+
|
|
45
|
+
# Airow
|
|
46
|
+
|
|
47
|
+
**AI-powered DataFrame processing made simple**
|
|
48
|
+
|
|
49
|
+
Airow is a Python library that combines the power of pandas DataFrames with AI models to process structured data at scale. Built on top of `pydantic-ai`, it provides type-safe, async processing of DataFrames using any AI model.
|
|
50
|
+
|
|
51
|
+
## Features
|
|
52
|
+
|
|
53
|
+
- 🚀 **Async processing** with batch support for high performance
|
|
54
|
+
- 🔒 **Type-safe outputs** using Pydantic models
|
|
55
|
+
- 📊 **Progress tracking** with built-in progress bars
|
|
56
|
+
- 🔄 **Automatic retries** with configurable retry logic
|
|
57
|
+
- 🤖 **Flexible AI models** - works with OpenAI, Ollama, Anthropic, and more
|
|
58
|
+
- ⚡ **Parallel processing** within batches for maximum throughput
|
|
59
|
+
- 📝 **Structured outputs** with defined schemas and validation
|
|
60
|
+
|
|
61
|
+
## Installation
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
# Using pip
|
|
65
|
+
pip install airow
|
|
66
|
+
|
|
67
|
+
# Using uv (recommended)
|
|
68
|
+
uv add airow
|
|
69
|
+
|
|
70
|
+
# Using conda
|
|
71
|
+
conda install -c conda-forge airow
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
## Quick Start
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
import pandas as pd
|
|
78
|
+
from pydantic_ai.models.openai import OpenAIChatModel
|
|
79
|
+
from pydantic_ai.providers.ollama import OllamaProvider
|
|
80
|
+
from airow import Airow, OutputColumn
|
|
81
|
+
import asyncio
|
|
82
|
+
|
|
83
|
+
async def main():
|
|
84
|
+
# Setup your AI model
|
|
85
|
+
model = OpenAIChatModel(
|
|
86
|
+
model_name="llama3.2:latest",
|
|
87
|
+
provider=OllamaProvider(base_url="http://localhost:11434/v1"),
|
|
88
|
+
)
|
|
89
|
+
# or use strings:
|
|
90
|
+
model = "openai:gpt-5"
|
|
91
|
+
model = "anthropic:claude-sonnet-4-0"
|
|
92
|
+
|
|
93
|
+
# Create Airow instance
|
|
94
|
+
airow = Airow(
|
|
95
|
+
model=model,
|
|
96
|
+
system_prompt="You are an expert in wine tasting and selection.",
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
# Load your data
|
|
100
|
+
df = pd.read_csv("wine_data.csv")
|
|
101
|
+
|
|
102
|
+
output_columns = [
|
|
103
|
+
OutputColumn(name="sentiment", type=str, description="Positive, negative, or neutral sentiment"),
|
|
104
|
+
OutputColumn(name="confidence", type=float, description="Confidence score between 0 and 1"),
|
|
105
|
+
OutputColumn(name="keywords", type=list, description="List of key terms extracted"),
|
|
106
|
+
]
|
|
107
|
+
|
|
108
|
+
# Process with AI
|
|
109
|
+
result_df = await airow.run(
|
|
110
|
+
df,
|
|
111
|
+
prompt="Analyze the wine description and provide sentiment analysis, confidence score, and extract key terms.",
|
|
112
|
+
input_columns=["description"],
|
|
113
|
+
output_columns=output_columns,
|
|
114
|
+
show_progress=True,
|
|
115
|
+
)
|
|
116
|
+
|
|
117
|
+
print(result_df.head())
|
|
118
|
+
|
|
119
|
+
if __name__ == "__main__":
|
|
120
|
+
asyncio.run(main())
|
|
121
|
+
```
|
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
README.md
|
|
3
|
+
pyproject.toml
|
|
4
|
+
airow/__init__.py
|
|
5
|
+
airow/agent.py
|
|
6
|
+
airow/airow.py
|
|
7
|
+
airow/schemas.py
|
|
8
|
+
airow.egg-info/PKG-INFO
|
|
9
|
+
airow.egg-info/SOURCES.txt
|
|
10
|
+
airow.egg-info/dependency_links.txt
|
|
11
|
+
airow.egg-info/requires.txt
|
|
12
|
+
airow.egg-info/top_level.txt
|
|
13
|
+
tests/test_agents.py
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
airow
|
|
@@ -0,0 +1,60 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "airow"
|
|
3
|
+
dynamic = ["version"]
|
|
4
|
+
description = "AI-powered DataFrame processing made simple"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = {text = "MIT"}
|
|
7
|
+
requires-python = ">=3.10"
|
|
8
|
+
authors = [
|
|
9
|
+
{name = "Dmitrii K", email = "dmitriik@protonmail.com"},
|
|
10
|
+
]
|
|
11
|
+
maintainers = [
|
|
12
|
+
{name = "Dmitrii K", email = "dmitriik@protonmail.com"},
|
|
13
|
+
]
|
|
14
|
+
keywords = ["ai", "ai-agent", "dataframe", "pandas", "pydantic-ai", "async", "data-processing"]
|
|
15
|
+
classifiers = [
|
|
16
|
+
"Development Status :: 5 - Production/Stable",
|
|
17
|
+
"Intended Audience :: Developers",
|
|
18
|
+
"Intended Audience :: Science/Research",
|
|
19
|
+
"License :: OSI Approved :: MIT License",
|
|
20
|
+
"Operating System :: OS Independent",
|
|
21
|
+
"Programming Language :: Python :: 3",
|
|
22
|
+
"Programming Language :: Python :: 3.10",
|
|
23
|
+
"Programming Language :: Python :: 3.11",
|
|
24
|
+
"Programming Language :: Python :: 3.12",
|
|
25
|
+
"Programming Language :: Python :: 3.13",
|
|
26
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
27
|
+
"Topic :: Scientific/Engineering :: Information Analysis",
|
|
28
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
29
|
+
"Topic :: Text Processing",
|
|
30
|
+
"Topic :: Database",
|
|
31
|
+
"Typing :: Typed",
|
|
32
|
+
]
|
|
33
|
+
dependencies = [
|
|
34
|
+
"loguru>=0.7.3",
|
|
35
|
+
"pandas>=2.3.2",
|
|
36
|
+
"pydantic>=2.11.7",
|
|
37
|
+
"pydantic-ai>=0.8.1",
|
|
38
|
+
"tqdm>=4.67.1",
|
|
39
|
+
]
|
|
40
|
+
|
|
41
|
+
[project.urls]
|
|
42
|
+
Homepage = "https://github.com/dmitriiweb/airow"
|
|
43
|
+
Repository = "https://github.com/dmitriiweb/airow"
|
|
44
|
+
Documentation = "https://github.com/dmitriiweb/airow"
|
|
45
|
+
"Bug Tracker" = "https://github.com/dmitriiweb/airow/issues"
|
|
46
|
+
|
|
47
|
+
[project.optional-dependencies]
|
|
48
|
+
dev = [
|
|
49
|
+
"mypy>=1.17.1",
|
|
50
|
+
"pytest>=8.4.2",
|
|
51
|
+
"pytest-asyncio>=1.1.0",
|
|
52
|
+
"pytest-cov>=6.3.0",
|
|
53
|
+
"ruff>=0.12.12",
|
|
54
|
+
]
|
|
55
|
+
|
|
56
|
+
[tool.setuptools.dynamic]
|
|
57
|
+
version = {attr = "airow.__version__"}
|
|
58
|
+
|
|
59
|
+
[tool.pytest.ini_options]
|
|
60
|
+
asyncio_mode = "auto"
|
airow-0.1.0/setup.cfg
ADDED
|
@@ -0,0 +1,189 @@
|
|
|
1
|
+
import pytest
|
|
2
|
+
from pydantic import BaseModel, ValidationError
|
|
3
|
+
|
|
4
|
+
from airow.agent import AirowAgent
|
|
5
|
+
from airow.schemas import OutputColumn
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def test_single_output_column():
|
|
9
|
+
"""Test creating a model with a single output column."""
|
|
10
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
11
|
+
|
|
12
|
+
output_columns = [
|
|
13
|
+
OutputColumn(name="result", type=str, description="The result of processing")
|
|
14
|
+
]
|
|
15
|
+
|
|
16
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
17
|
+
|
|
18
|
+
# Verify it's a BaseModel subclass
|
|
19
|
+
assert issubclass(model_class, BaseModel)
|
|
20
|
+
|
|
21
|
+
# Verify the model name
|
|
22
|
+
assert model_class.__name__ == "OutputColumns"
|
|
23
|
+
|
|
24
|
+
# Verify field exists
|
|
25
|
+
assert "result" in model_class.model_fields
|
|
26
|
+
|
|
27
|
+
# Verify field configuration
|
|
28
|
+
field_info = model_class.model_fields["result"]
|
|
29
|
+
assert field_info.annotation is str
|
|
30
|
+
assert field_info.description == "The result of processing"
|
|
31
|
+
assert field_info.is_required()
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def test_multiple_output_columns():
|
|
35
|
+
"""Test creating a model with multiple output columns."""
|
|
36
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
37
|
+
|
|
38
|
+
output_columns = [
|
|
39
|
+
OutputColumn(name="name", type=str, description="Person's name"),
|
|
40
|
+
OutputColumn(name="age", type=int, description="Person's age"),
|
|
41
|
+
OutputColumn(name="is_active", type=bool, description="Whether person is active"),
|
|
42
|
+
]
|
|
43
|
+
|
|
44
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
45
|
+
|
|
46
|
+
# Verify all fields exist
|
|
47
|
+
assert "name" in model_class.model_fields
|
|
48
|
+
assert "age" in model_class.model_fields
|
|
49
|
+
assert "is_active" in model_class.model_fields
|
|
50
|
+
|
|
51
|
+
# Verify field types and descriptions
|
|
52
|
+
name_field = model_class.model_fields["name"]
|
|
53
|
+
assert name_field.annotation is str
|
|
54
|
+
assert name_field.description == "Person's name"
|
|
55
|
+
|
|
56
|
+
age_field = model_class.model_fields["age"]
|
|
57
|
+
assert age_field.annotation is int
|
|
58
|
+
assert age_field.description == "Person's age"
|
|
59
|
+
|
|
60
|
+
is_active_field = model_class.model_fields["is_active"]
|
|
61
|
+
assert is_active_field.annotation is bool
|
|
62
|
+
assert is_active_field.description == "Whether person is active"
|
|
63
|
+
|
|
64
|
+
|
|
65
|
+
def test_different_data_types():
|
|
66
|
+
"""Test creating a model with various data types."""
|
|
67
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
68
|
+
|
|
69
|
+
output_columns = [
|
|
70
|
+
OutputColumn(name="text", type=str, description="Text field"),
|
|
71
|
+
OutputColumn(name="number", type=int, description="Integer field"),
|
|
72
|
+
OutputColumn(name="float_val", type=float, description="Float field"),
|
|
73
|
+
OutputColumn(name="flag", type=bool, description="Boolean field"),
|
|
74
|
+
OutputColumn(name="items", type=list, description="List field"),
|
|
75
|
+
OutputColumn(name="metadata", type=dict, description="Dictionary field"),
|
|
76
|
+
]
|
|
77
|
+
|
|
78
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
79
|
+
|
|
80
|
+
# Verify all fields exist with correct types
|
|
81
|
+
assert model_class.model_fields["text"].annotation is str
|
|
82
|
+
assert model_class.model_fields["number"].annotation is int
|
|
83
|
+
assert model_class.model_fields["float_val"].annotation is float
|
|
84
|
+
assert model_class.model_fields["flag"].annotation is bool
|
|
85
|
+
assert model_class.model_fields["items"].annotation is list
|
|
86
|
+
assert model_class.model_fields["metadata"].annotation is dict
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
def test_empty_output_columns():
|
|
90
|
+
"""Test creating a model with no output columns."""
|
|
91
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
92
|
+
|
|
93
|
+
output_columns = []
|
|
94
|
+
|
|
95
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
96
|
+
|
|
97
|
+
# Should still create a valid model class
|
|
98
|
+
assert issubclass(model_class, BaseModel)
|
|
99
|
+
assert model_class.__name__ == "OutputColumns"
|
|
100
|
+
assert len(model_class.model_fields) == 0
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
def test_model_validation():
|
|
104
|
+
"""Test that the created model validates data correctly."""
|
|
105
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
106
|
+
|
|
107
|
+
output_columns = [
|
|
108
|
+
OutputColumn(name="name", type=str, description="Person's name"),
|
|
109
|
+
OutputColumn(name="age", type=int, description="Person's age"),
|
|
110
|
+
]
|
|
111
|
+
|
|
112
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
113
|
+
|
|
114
|
+
# Test valid data
|
|
115
|
+
valid_data = {"name": "John Doe", "age": 30}
|
|
116
|
+
instance = model_class(**valid_data)
|
|
117
|
+
assert instance.name == "John Doe"
|
|
118
|
+
assert instance.age == 30
|
|
119
|
+
|
|
120
|
+
# Test invalid data (missing required field)
|
|
121
|
+
with pytest.raises(ValidationError):
|
|
122
|
+
model_class(name="John Doe") # Missing age
|
|
123
|
+
|
|
124
|
+
# Test invalid data (wrong type)
|
|
125
|
+
with pytest.raises(ValidationError):
|
|
126
|
+
model_class(name="John Doe", age="thirty") # age should be int
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
def test_model_serialization():
|
|
130
|
+
"""Test that the created model can be serialized."""
|
|
131
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
132
|
+
|
|
133
|
+
output_columns = [
|
|
134
|
+
OutputColumn(name="result", type=str, description="Processing result"),
|
|
135
|
+
OutputColumn(name="score", type=float, description="Processing score"),
|
|
136
|
+
]
|
|
137
|
+
|
|
138
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
139
|
+
|
|
140
|
+
instance = model_class(result="success", score=0.95)
|
|
141
|
+
|
|
142
|
+
# Test model_dump
|
|
143
|
+
data = instance.model_dump()
|
|
144
|
+
expected = {"result": "success", "score": 0.95}
|
|
145
|
+
assert data == expected
|
|
146
|
+
|
|
147
|
+
# Test model_dump_json
|
|
148
|
+
json_data = instance.model_dump_json()
|
|
149
|
+
assert '"result":"success"' in json_data
|
|
150
|
+
assert '"score":0.95' in json_data
|
|
151
|
+
|
|
152
|
+
|
|
153
|
+
def test_field_descriptions_preserved():
|
|
154
|
+
"""Test that field descriptions are properly preserved."""
|
|
155
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
156
|
+
|
|
157
|
+
output_columns = [
|
|
158
|
+
OutputColumn(
|
|
159
|
+
name="complex_field",
|
|
160
|
+
type=str,
|
|
161
|
+
description="This is a complex field with special characters: @#$%^&*()"
|
|
162
|
+
),
|
|
163
|
+
]
|
|
164
|
+
|
|
165
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
166
|
+
|
|
167
|
+
field_info = model_class.model_fields["complex_field"]
|
|
168
|
+
assert field_info.description == "This is a complex field with special characters: @#$%^&*()"
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
def test_duplicate_field_names():
|
|
172
|
+
"""Test behavior with duplicate field names (should overwrite)."""
|
|
173
|
+
agent = AirowAgent(model=None, system_prompt="test")
|
|
174
|
+
|
|
175
|
+
output_columns = [
|
|
176
|
+
OutputColumn(name="field", type=str, description="First field"),
|
|
177
|
+
OutputColumn(name="field", type=int, description="Second field"), # Duplicate name
|
|
178
|
+
]
|
|
179
|
+
|
|
180
|
+
model_class = agent.build_agent_output_type(output_columns)
|
|
181
|
+
|
|
182
|
+
# Should have only one field (last one wins)
|
|
183
|
+
assert len(model_class.model_fields) == 1
|
|
184
|
+
assert "field" in model_class.model_fields
|
|
185
|
+
|
|
186
|
+
# Should use the last definition
|
|
187
|
+
field_info = model_class.model_fields["field"]
|
|
188
|
+
assert field_info.annotation is int
|
|
189
|
+
assert field_info.description == "Second field"
|