airow 0.1.0__py3-none-any.whl

airow/__init__.py

@@ -0,0 +1,11 @@
+"""airow package public API and version information.
+
+Exposes `Airow` for batched, row-wise LLM inference over pandas DataFrames
+and `OutputColumn` for declaring structured outputs.
+"""
+__version__ = "0.1.0"
+
+from .airow import Airow
+from .schemas import OutputColumn
+
+__all__ = ["Airow", "OutputColumn"]

airow/agent.py

@@ -0,0 +1,71 @@
+"""Agent utilities for executing pydantic-ai models with structured outputs."""
+
+from typing import Iterable
+
+from pydantic import BaseModel, Field, create_model
+from pydantic_ai import Agent
+from pydantic_ai.models import Model
+from loguru import logger
+
+from . import schemas
+
+
+class AirowAgent:
+    """Wrapper around `pydantic_ai.Agent` that builds structured output models."""
+    def __init__(
+        self,
+        model: Model | str,
+        system_prompt: str,
+        retries: int = 3,
+    ):
+        """Initialize the agent.
+
+        Args:
+            model: The underlying model used by pydantic-ai.
+            system_prompt: System prompt applied to all runs.
+            retries: Number of retries for a run.
+        """
+        self.model = model
+        self.system_prompt = system_prompt
+        self.agent = Agent(model=model, system_prompt=self.system_prompt, retries=retries)
+
+    async def run(
+        self,
+        prompt: str,
+        output_columns: Iterable[schemas.OutputColumn],
+    ) -> dict[str, object]:
+        """Run the agent with the given prompt and expected outputs.
+
+        Args:
+            prompt: User prompt to pass to the model.
+            output_columns: Iterable of expected output columns specifications.
+
+        Returns:
+            A dictionary mapping output column names to parsed values. Returns
+            an empty dictionary when the underlying model call fails.
+        """
+        output_columns_fields = self.build_agent_output_type(output_columns)
+        try:
+            result = await self.agent.run(prompt, output_type=output_columns_fields)
+        except Exception as e:
+            logger.error(f"{e=}")
+            return {}
+        return result.output.model_dump()
+
+    def build_agent_output_type(
+        self,
+        output_columns: Iterable[schemas.OutputColumn],
+    ) -> type[BaseModel]:
+        """Create a `pydantic.BaseModel` for the requested output columns.
+
+        Args:
+            output_columns: Iterable of output column specifications.
+
+        Returns:
+            A dynamically created `BaseModel` subclass with fields per column.
+        """
+        fields = {
+            col.name: (col.type, Field(..., description=col.description))
+            for col in output_columns
+        }
+        return create_model("OutputColumns", **fields)

airow/airow.py

@@ -0,0 +1,110 @@
+"""High-level API to run LLMs over DataFrame rows in batches."""
+
+import asyncio
+from typing import Iterable
+
+import pandas as pd
+from pydantic_ai.models import Model
+from tqdm import tqdm
+
+from . import schemas
+from .agent import AirowAgent
+
+
+class Airow:
+    """Apply an LLM to each row of a DataFrame and write results.
+
+    Uses `AirowAgent` internally and supports parallel row processing per batch.
+    """
+    def __init__(
+        self,
+        *,
+        model: Model | str,
+        system_prompt: str,
+        batch_size: int = 1,
+        retries: int = 3,
+    ):
+        """Configure the runner.
+
+        Args:
+            model: The pydantic-ai model to use.
+            system_prompt: System prompt applied to each request.
+            batch_size: Number of DataFrame rows to process concurrently.
+            retries: Number of retries for a run.
+        """
+        self.model = model
+        self.system_prompt = system_prompt
+        self.batch_size = batch_size
+        self.agent = AirowAgent(self.model, self.system_prompt, retries)
+
+    async def run(
+        self,
+        df: pd.DataFrame,
+        *,
+        prompt: str,
+        input_columns: Iterable[str],
+        output_columns: schemas.OutputColumn | Iterable[schemas.OutputColumn],
+        show_progress: bool = False,
+    ) -> pd.DataFrame:
+        """Run the model across the DataFrame and return results.
+
+        For each row, the values from `input_columns` are appended to `prompt`
+        as labeled lines and the model is asked to produce `output_columns`.
+        Results are written into the original DataFrame.
+
+        Args:
+            df: Input DataFrame.
+            prompt: Base prompt text provided to the model.
+            input_columns: Columns whose values are passed as context to the model.
+            output_columns: One or more output column specifications.
+            show_progress: Whether to display a progress bar.
+
+        Returns:
+            The input DataFrame with new output columns populated.
+        """
+        if isinstance(output_columns, schemas.OutputColumn):
+            output_columns = [output_columns]
+
+        # Convert to list for easier handling
+        input_columns = list(input_columns)
+
+        # Work with a copy to avoid SettingWithCopyWarning
+        df = df.copy()
+
+        # Split dataframe into batches
+        total_rows = df.shape[0]
+        batche_ranges = [
+            (i, i + self.batch_size)
+            for i in range(0, total_rows, self.batch_size)
+        ]
+        if show_progress:
+            batche_ranges = tqdm(batche_ranges)
+
+        for batch_range in batche_ranges:
+            # Process each row in the batch in parallel
+            tasks = []
+            row_indices = []
+            
+            # Get row indices for this batch
+            start_idx = batch_range[0]
+            end_idx = min(batch_range[1], total_rows)
+            
+            for row_idx in range(start_idx, end_idx):
+                row = df.iloc[row_idx]
+                input_data = {col: row[col] for col in input_columns}
+                input_data_str = "\n".join([f"Column: {k}, Value: {v}" for k, v in input_data.items()])
+                full_prompt = f"{prompt}\n\n{input_data_str}"
+                task = self.agent.run(full_prompt, output_columns)
+                tasks.append(task)
+                row_indices.append(row_idx)
+
+            # Run all tasks in parallel
+            results = await asyncio.gather(*tasks)
+
+            # Add results to dataframe
+            for i, result in enumerate(results):
+                row_idx = row_indices[i]
+                for col_name, value in result.items():
+                    df.loc[row_idx, col_name] = value
+
+        return df

airow/schemas.py

@@ -0,0 +1,25 @@
+"""Data structures for declaring model outputs."""
+
+from dataclasses import dataclass
+from typing import Any, Type
+
+from pydantic import BaseModel
+
+
+@dataclass
+class OutputColumn:
+    """
+    Output column for the AI model.
+
+    Args:
+        name: Name of the output column.
+        type: Type of the output column.
+
+    Examples:
+        >>> OutputColumn(name="output_column", type=str)
+        >>> OutputColumn(name="output_column", type=int)
+    """
+
+    name: str
+    type: Type[Any]
+    description: str

airow-0.1.0.dist-info/METADATA

@@ -0,0 +1,121 @@
+Metadata-Version: 2.4
+Name: airow
+Version: 0.1.0
+Summary: AI-powered DataFrame processing made simple
+Author-email: Dmitrii K <dmitriik@protonmail.com>
+Maintainer-email: Dmitrii K <dmitriik@protonmail.com>
+License: MIT
+Project-URL: Homepage, https://github.com/dmitriiweb/airow
+Project-URL: Repository, https://github.com/dmitriiweb/airow
+Project-URL: Documentation, https://github.com/dmitriiweb/airow
+Project-URL: Bug Tracker, https://github.com/dmitriiweb/airow/issues
+Keywords: ai,ai-agent,dataframe,pandas,pydantic-ai,async,data-processing
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Topic :: Text Processing
+Classifier: Topic :: Database
+Classifier: Typing :: Typed
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: loguru>=0.7.3
+Requires-Dist: pandas>=2.3.2
+Requires-Dist: pydantic>=2.11.7
+Requires-Dist: pydantic-ai>=0.8.1
+Requires-Dist: tqdm>=4.67.1
+Provides-Extra: dev
+Requires-Dist: mypy>=1.17.1; extra == "dev"
+Requires-Dist: pytest>=8.4.2; extra == "dev"
+Requires-Dist: pytest-asyncio>=1.1.0; extra == "dev"
+Requires-Dist: pytest-cov>=6.3.0; extra == "dev"
+Requires-Dist: ruff>=0.12.12; extra == "dev"
+Dynamic: license-file
+
+# Airow
+
+**AI-powered DataFrame processing made simple**
+
+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.
+
+## Features
+
+- 🚀 **Async processing** with batch support for high performance
+- 🔒 **Type-safe outputs** using Pydantic models
+- 📊 **Progress tracking** with built-in progress bars
+- 🔄 **Automatic retries** with configurable retry logic
+- 🤖 **Flexible AI models** - works with OpenAI, Ollama, Anthropic, and more
+- ⚡ **Parallel processing** within batches for maximum throughput
+- 📝 **Structured outputs** with defined schemas and validation
+
+## Installation
+
+```bash
+# Using pip
+pip install airow
+
+# Using uv (recommended)
+uv add airow
+
+# Using conda
+conda install -c conda-forge airow
+```
+
+## Quick Start
+
+```python
+import pandas as pd
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.ollama import OllamaProvider
+from airow import Airow, OutputColumn
+import asyncio
+
+async def main():
+    # Setup your AI model
+    model = OpenAIChatModel(
+        model_name="llama3.2:latest",
+        provider=OllamaProvider(base_url="http://localhost:11434/v1"),
+    )
+    # or use strings:
+    model = "openai:gpt-5"
+    model = "anthropic:claude-sonnet-4-0"
+    
+    # Create Airow instance
+    airow = Airow(
+        model=model,
+        system_prompt="You are an expert in wine tasting and selection.",
+    )
+    
+    # Load your data
+    df = pd.read_csv("wine_data.csv")
+
+    output_columns = [
+        OutputColumn(name="sentiment", type=str, description="Positive, negative, or neutral sentiment"),
+        OutputColumn(name="confidence", type=float, description="Confidence score between 0 and 1"),
+        OutputColumn(name="keywords", type=list, description="List of key terms extracted"),
+    ]
+    
+    # Process with AI
+    result_df = await airow.run(
+        df,
+        prompt="Analyze the wine description and provide sentiment analysis, confidence score, and extract key terms.",
+        input_columns=["description"],
+        output_columns=output_columns,
+        show_progress=True,
+    )
+    
+    print(result_df.head())
+
+if __name__ == "__main__":
+    asyncio.run(main())
+```

airow-0.1.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+airow/__init__.py,sha256=SOIZ4vEvMRj1fnoEpFOqH9WJyqXexjVwNYqJPbxsoag,305
+airow/agent.py,sha256=Qp77KCzNwF84dwkbqbDonmw7Yqp-IU_OJ9rwOVMUuE8,2331
+airow/airow.py,sha256=ykrIizSVy_bEFKUgvebCYRqw0nx_9dlY0hyyAVWkEaY,3723
+airow/schemas.py,sha256=2qCW5DR4tcDcY3M8IQihuR9rqtfiYwkp2z_xEOB6vA0,510
+airow-0.1.0.dist-info/licenses/LICENSE,sha256=vTx38XxI12rNKWp4ebsj-Yc9eke56CBetDgHW2L2Cfg,1066
+airow-0.1.0.dist-info/METADATA,sha256=jO7g4QxGJGLGkKULWiI60UVXMHUGzqXksCH_PGrn7y0,4196
+airow-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+airow-0.1.0.dist-info/top_level.txt,sha256=Ux5WLwH3HuMGPpvHKYkBIbrbmxWQd_z-aIR2R_mHG5U,6
+airow-0.1.0.dist-info/RECORD,,

airow-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

airow-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Dmitrii K
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

airow-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+airow