openaivec.pandas_ext

Pandas Series / DataFrame extension for OpenAI.

Setup

from openai import OpenAI
from openaivec import pandas_ext

# Set up the OpenAI client to use with pandas_ext
# Option 1: Use an existing client instance
# pandas_ext.use(OpenAI())

# Option 2: Use environment variables (OPENAI_API_KEY or Azure variables)
# (No explicit setup needed if variables are set)

# Option 3: Provide API key directly
pandas_ext.use_openai("YOUR_API_KEY")

# Option 4: Use Azure OpenAI credentials
# pandas_ext.use_azure_openai(
#     api_key="YOUR_AZURE_KEY",
#     endpoint="YOUR_AZURE_ENDPOINT",
#     api_version="YOUR_API_VERSION"
# )

# Set up the model_name for responses and embeddings (optional, defaults shown)
pandas_ext.responses_model("gpt-4o-mini")
pandas_ext.embeddings_model("text-embedding-3-small")

This module provides .ai and .aio accessors for pandas Series and DataFrames to easily interact with OpenAI APIs for tasks like generating responses or embeddings.

AsyncOpenAIVecDataFrameAccessor

pandas DataFrame accessor (.aio) that adds OpenAI helpers.

Source code in src/openaivec/pandas_ext.py

@pd.api.extensions.register_dataframe_accessor("aio")
class AsyncOpenAIVecDataFrameAccessor:
    """pandas DataFrame accessor (``.aio``) that adds OpenAI helpers."""

    def __init__(self, df_obj: pd.DataFrame):
        self._obj = df_obj

    async def responses(
        self,
        instructions: str,
        response_format: Type[T] = str,
        batch_size: int = 128,
        temperature: float = 0.0,
        top_p: float = 1.0,
        max_concurrency: int = 8,
    ) -> pd.Series:
        """Generate a response for each row after serialising it to JSON (asynchronously).

        Example:
            ```python
            df = pd.DataFrame([
                {\"name\": \"cat\", \"legs\": 4},
                {\"name\": \"dog\", \"legs\": 4},
                {\"name\": \"elephant\", \"legs\": 4},
            ])
            # Must be awaited
            results = await df.aio.responses(\"what is the animal\'s name?\")
            ```
            This method returns a Series of strings, each containing the
            assistant's response to the corresponding input.
            Each row is serialised to JSON before being sent to the assistant.
            The model used is set by the `responses_model` function.
            The default model is `gpt-4o-mini`.

        Args:
            instructions (str): System prompt for the assistant.
            response_format (Type[T], optional): Desired Python type of the
                responses. Defaults to ``str``.
            batch_size (int, optional): Number of requests sent in one batch.
                Defaults to ``128``.
            temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
            top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.
            max_concurrency (int, optional): Maximum number of concurrent
                requests. Defaults to ``8``.

        Returns:
            pandas.Series: Responses aligned with the DataFrame’s original index.

        Note:
            This is an asynchronous method and must be awaited.
        """
        series_of_json = self._obj.pipe(
            lambda df: (
                pd.Series(df.to_dict(orient="records"), index=df.index, name="record").map(
                    lambda x: json.dumps(x, ensure_ascii=False)
                )
            )
        )
        # Await the call to the async Series method using .aio
        return await series_of_json.aio.responses(
            instructions=instructions,
            response_format=response_format,
            batch_size=batch_size,
            temperature=temperature,
            top_p=top_p,
            max_concurrency=max_concurrency,
        )

    async def pipe(self, func: Callable[[pd.DataFrame], Awaitable[T] | T]) -> T:
        """
        Apply a function to the DataFrame, supporting both synchronous and asynchronous functions.

        This method allows chaining operations on the DataFrame, similar to pandas' `pipe` method,
        but with support for asynchronous functions.

        Args:
            func (Callable[[pd.DataFrame], Awaitable[T] | T]): A function that takes a DataFrame
                as input and returns either a result or an awaitable result.

        Returns:
            T: The result of applying the function, either directly or after awaiting it.

        Note:
            This is an asynchronous method and must be awaited if the function returns an awaitable.
        """
        result = func(self._obj)
        if inspect.isawaitable(result):
            return await result
        else:
            return result

    async def assign(self, **kwargs):
        """Asynchronously assign new columns to the DataFrame, evaluating sequentially.

        This method extends pandas' `assign` method by supporting asynchronous
        functions as column values and evaluating assignments sequentially, allowing
        later assignments to refer to columns created earlier in the same call.

        For each key-value pair in `kwargs`:
        - If the value is a callable, it is invoked with the current state of the DataFrame
          (including columns created in previous steps of this `assign` call).
          If the result is awaitable, it is awaited; otherwise, it is used directly.
        - If the value is not callable, it is assigned directly to the new column.

        Example:
            ```python
            async def compute_column(df):
                # Simulate an asynchronous computation
                await asyncio.sleep(1)
                return df["existing_column"] * 2

            async def use_new_column(df):
                # Access the column created in the previous step
                await asyncio.sleep(1)
                return df["new_column"] + 5


            df = pd.DataFrame({"existing_column": [1, 2, 3]})
            # Must be awaited
            df = await df.aio.assign(
                new_column=compute_column,
                another_column=use_new_column
            )
            ```

        Args:
            **kwargs: Column names as keys and either static values or callables
                (synchronous or asynchronous) as values.

        Returns:
            pandas.DataFrame: A new DataFrame with the assigned columns.

        Note:
            This is an asynchronous method and must be awaited.
        """
        df_current = self._obj.copy()
        for key, value in kwargs.items():
            if callable(value):
                result = value(df_current)
                if inspect.isawaitable(result):
                    column_data = await result
                else:
                    column_data = result
            else:
                column_data = value

            df_current[key] = column_data

        return df_current

assign(**kwargs) async

Asynchronously assign new columns to the DataFrame, evaluating sequentially.

This method extends pandas' assign method by supporting asynchronous functions as column values and evaluating assignments sequentially, allowing later assignments to refer to columns created earlier in the same call.

For each key-value pair in kwargs: - If the value is a callable, it is invoked with the current state of the DataFrame (including columns created in previous steps of this assign call). If the result is awaitable, it is awaited; otherwise, it is used directly. - If the value is not callable, it is assigned directly to the new column.

Example

async def compute_column(df):
    # Simulate an asynchronous computation
    await asyncio.sleep(1)
    return df["existing_column"] * 2

async def use_new_column(df):
    # Access the column created in the previous step
    await asyncio.sleep(1)
    return df["new_column"] + 5


df = pd.DataFrame({"existing_column": [1, 2, 3]})
# Must be awaited
df = await df.aio.assign(
    new_column=compute_column,
    another_column=use_new_column
)

Parameters:

Name	Type	Description	Default
**kwargs		Column names as keys and either static values or callables (synchronous or asynchronous) as values.	{}

Returns:

Type	Description
	pandas.DataFrame: A new DataFrame with the assigned columns.

Note

This is an asynchronous method and must be awaited.

Source code in src/openaivec/pandas_ext.py

async def assign(self, **kwargs):
    """Asynchronously assign new columns to the DataFrame, evaluating sequentially.

    This method extends pandas' `assign` method by supporting asynchronous
    functions as column values and evaluating assignments sequentially, allowing
    later assignments to refer to columns created earlier in the same call.

    For each key-value pair in `kwargs`:
    - If the value is a callable, it is invoked with the current state of the DataFrame
      (including columns created in previous steps of this `assign` call).
      If the result is awaitable, it is awaited; otherwise, it is used directly.
    - If the value is not callable, it is assigned directly to the new column.

    Example:
        ```python
        async def compute_column(df):
            # Simulate an asynchronous computation
            await asyncio.sleep(1)
            return df["existing_column"] * 2

        async def use_new_column(df):
            # Access the column created in the previous step
            await asyncio.sleep(1)
            return df["new_column"] + 5


        df = pd.DataFrame({"existing_column": [1, 2, 3]})
        # Must be awaited
        df = await df.aio.assign(
            new_column=compute_column,
            another_column=use_new_column
        )
        ```

    Args:
        **kwargs: Column names as keys and either static values or callables
            (synchronous or asynchronous) as values.

    Returns:
        pandas.DataFrame: A new DataFrame with the assigned columns.

    Note:
        This is an asynchronous method and must be awaited.
    """
    df_current = self._obj.copy()
    for key, value in kwargs.items():
        if callable(value):
            result = value(df_current)
            if inspect.isawaitable(result):
                column_data = await result
            else:
                column_data = result
        else:
            column_data = value

        df_current[key] = column_data

    return df_current

pipe(func) async

Apply a function to the DataFrame, supporting both synchronous and asynchronous functions.

This method allows chaining operations on the DataFrame, similar to pandas' pipe method, but with support for asynchronous functions.

Parameters:

Name	Type	Description	Default
func	Callable[[DataFrame], Awaitable[T] | T]	A function that takes a DataFrame as input and returns either a result or an awaitable result.	required

Returns:

Name	Type	Description
T	T	The result of applying the function, either directly or after awaiting it.

Note

This is an asynchronous method and must be awaited if the function returns an awaitable.

Source code in src/openaivec/pandas_ext.py

async def pipe(self, func: Callable[[pd.DataFrame], Awaitable[T] | T]) -> T:
    """
    Apply a function to the DataFrame, supporting both synchronous and asynchronous functions.

    This method allows chaining operations on the DataFrame, similar to pandas' `pipe` method,
    but with support for asynchronous functions.

    Args:
        func (Callable[[pd.DataFrame], Awaitable[T] | T]): A function that takes a DataFrame
            as input and returns either a result or an awaitable result.

    Returns:
        T: The result of applying the function, either directly or after awaiting it.

    Note:
        This is an asynchronous method and must be awaited if the function returns an awaitable.
    """
    result = func(self._obj)
    if inspect.isawaitable(result):
        return await result
    else:
        return result

responses(instructions, response_format=str, batch_size=128, temperature=0.0, top_p=1.0, max_concurrency=8) async

Generate a response for each row after serialising it to JSON (asynchronously).

Example

df = pd.DataFrame([
    {"name": "cat", "legs": 4},
    {"name": "dog", "legs": 4},
    {"name": "elephant", "legs": 4},
])
# Must be awaited
results = await df.aio.responses("what is the animal's name?")

This method returns a Series of strings, each containing the assistant's response to the corresponding input. Each row is serialised to JSON before being sent to the assistant. The model used is set by the responses_model function. The default model is gpt-4o-mini.

Parameters:

Name	Type	Description	Default
instructions	str	System prompt for the assistant.	required
response_format	Type[T]	Desired Python type of the responses. Defaults to str.	str
batch_size	int	Number of requests sent in one batch. Defaults to 128.	128
temperature	float	Sampling temperature. Defaults to 0.0.	0.0
top_p	float	Nucleus sampling parameter. Defaults to 1.0.	1.0
max_concurrency	int	Maximum number of concurrent requests. Defaults to 8.	8

Returns:

Type	Description
Series	pandas.Series: Responses aligned with the DataFrame’s original index.

Note

This is an asynchronous method and must be awaited.

Source code in src/openaivec/pandas_ext.py

async def responses(
    self,
    instructions: str,
    response_format: Type[T] = str,
    batch_size: int = 128,
    temperature: float = 0.0,
    top_p: float = 1.0,
    max_concurrency: int = 8,
) -> pd.Series:
    """Generate a response for each row after serialising it to JSON (asynchronously).

    Example:
        ```python
        df = pd.DataFrame([
            {\"name\": \"cat\", \"legs\": 4},
            {\"name\": \"dog\", \"legs\": 4},
            {\"name\": \"elephant\", \"legs\": 4},
        ])
        # Must be awaited
        results = await df.aio.responses(\"what is the animal\'s name?\")
        ```
        This method returns a Series of strings, each containing the
        assistant's response to the corresponding input.
        Each row is serialised to JSON before being sent to the assistant.
        The model used is set by the `responses_model` function.
        The default model is `gpt-4o-mini`.

    Args:
        instructions (str): System prompt for the assistant.
        response_format (Type[T], optional): Desired Python type of the
            responses. Defaults to ``str``.
        batch_size (int, optional): Number of requests sent in one batch.
            Defaults to ``128``.
        temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
        top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.
        max_concurrency (int, optional): Maximum number of concurrent
            requests. Defaults to ``8``.

    Returns:
        pandas.Series: Responses aligned with the DataFrame’s original index.

    Note:
        This is an asynchronous method and must be awaited.
    """
    series_of_json = self._obj.pipe(
        lambda df: (
            pd.Series(df.to_dict(orient="records"), index=df.index, name="record").map(
                lambda x: json.dumps(x, ensure_ascii=False)
            )
        )
    )
    # Await the call to the async Series method using .aio
    return await series_of_json.aio.responses(
        instructions=instructions,
        response_format=response_format,
        batch_size=batch_size,
        temperature=temperature,
        top_p=top_p,
        max_concurrency=max_concurrency,
    )

AsyncOpenAIVecSeriesAccessor

pandas Series accessor (.aio) that adds OpenAI helpers.

Source code in src/openaivec/pandas_ext.py

@pd.api.extensions.register_series_accessor("aio")
class AsyncOpenAIVecSeriesAccessor:
    """pandas Series accessor (``.aio``) that adds OpenAI helpers."""

    def __init__(self, series_obj: pd.Series):
        self._obj = series_obj

    async def responses(
        self,
        instructions: str,
        response_format: Type[T] = str,
        batch_size: int = 128,
        temperature: float = 0.0,
        top_p: float = 1.0,
        max_concurrency: int = 8,
    ) -> pd.Series:
        """Call an LLM once for every Series element (asynchronously).

        Example:
            ```python
            animals = pd.Series(["cat", "dog", "elephant"])
            # Must be awaited
            results = await animals.aio.responses("translate to French")
            ```
            This method returns a Series of strings, each containing the
            assistant's response to the corresponding input.
            The model used is set by the `responses_model` function.
            The default model is `gpt-4o-mini`.

        Args:
            instructions (str): System prompt prepended to every user message.
            response_format (Type[T], optional): Pydantic model or built‑in
                type the assistant should return. Defaults to ``str``.
            batch_size (int, optional): Number of prompts grouped into a single
                request. Defaults to ``128``.
            temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
            top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.
            max_concurrency (int, optional): Maximum number of concurrent
                requests. Defaults to ``8``.

        Returns:
            pandas.Series: Series whose values are instances of ``response_format``.

        Note:
            This is an asynchronous method and must be awaited.
        """
        client: AsyncBatchResponses = AsyncBatchResponses(
            client=_get_async_openai_client(),
            model_name=_RESPONSES_MODEL_NAME,
            system_message=instructions,
            response_format=response_format,
            temperature=temperature,
            top_p=top_p,
            max_concurrency=max_concurrency,
        )

        # Await the async operation
        results = await client.parse(self._obj.tolist(), batch_size=batch_size)

        return pd.Series(
            results,
            index=self._obj.index,
            name=self._obj.name,
        )

    async def embeddings(self, batch_size: int = 128, max_concurrency: int = 8) -> pd.Series:
        """Compute OpenAI embeddings for every Series element (asynchronously).

        Example:
            ```python
            animals = pd.Series(["cat", "dog", "elephant"])
            # Must be awaited
            embeddings = await animals.aio.embeddings()
            ```
            This method returns a Series of numpy arrays, each containing the
            embedding vector for the corresponding input.
            The embedding model is set by the `embeddings_model` function.
            The default embedding model is `text-embedding-3-small`.

        Args:
            batch_size (int, optional): Number of inputs grouped into a
                single request. Defaults to ``128``.
            max_concurrency (int, optional): Maximum number of concurrent
                requests. Defaults to ``8``.

        Returns:
            pandas.Series: Series whose values are ``np.ndarray`` objects
                (dtype ``float32``).

        Note:
            This is an asynchronous method and must be awaited.
        """
        client: AsyncBatchEmbeddings = AsyncBatchEmbeddings(
            client=_get_async_openai_client(),
            model_name=_EMBEDDINGS_MODEL_NAME,
            max_concurrency=max_concurrency,
        )

        # Await the async operation
        results = await client.create(self._obj.tolist(), batch_size=batch_size)

        return pd.Series(
            results,
            index=self._obj.index,
            name=self._obj.name,
        )

embeddings(batch_size=128, max_concurrency=8) async

Compute OpenAI embeddings for every Series element (asynchronously).

Example

animals = pd.Series(["cat", "dog", "elephant"])
# Must be awaited
embeddings = await animals.aio.embeddings()

This method returns a Series of numpy arrays, each containing the embedding vector for the corresponding input. The embedding model is set by the embeddings_model function. The default embedding model is text-embedding-3-small.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of inputs grouped into a single request. Defaults to 128.	128
max_concurrency	int	Maximum number of concurrent requests. Defaults to 8.	8

Returns:

Type	Description
Series	pandas.Series: Series whose values are np.ndarray objects (dtype float32).

Note

This is an asynchronous method and must be awaited.

Source code in src/openaivec/pandas_ext.py

async def embeddings(self, batch_size: int = 128, max_concurrency: int = 8) -> pd.Series:
    """Compute OpenAI embeddings for every Series element (asynchronously).

    Example:
        ```python
        animals = pd.Series(["cat", "dog", "elephant"])
        # Must be awaited
        embeddings = await animals.aio.embeddings()
        ```
        This method returns a Series of numpy arrays, each containing the
        embedding vector for the corresponding input.
        The embedding model is set by the `embeddings_model` function.
        The default embedding model is `text-embedding-3-small`.

    Args:
        batch_size (int, optional): Number of inputs grouped into a
            single request. Defaults to ``128``.
        max_concurrency (int, optional): Maximum number of concurrent
            requests. Defaults to ``8``.

    Returns:
        pandas.Series: Series whose values are ``np.ndarray`` objects
            (dtype ``float32``).

    Note:
        This is an asynchronous method and must be awaited.
    """
    client: AsyncBatchEmbeddings = AsyncBatchEmbeddings(
        client=_get_async_openai_client(),
        model_name=_EMBEDDINGS_MODEL_NAME,
        max_concurrency=max_concurrency,
    )

    # Await the async operation
    results = await client.create(self._obj.tolist(), batch_size=batch_size)

    return pd.Series(
        results,
        index=self._obj.index,
        name=self._obj.name,
    )

responses(instructions, response_format=str, batch_size=128, temperature=0.0, top_p=1.0, max_concurrency=8) async

Call an LLM once for every Series element (asynchronously).

Example

animals = pd.Series(["cat", "dog", "elephant"])
# Must be awaited
results = await animals.aio.responses("translate to French")

This method returns a Series of strings, each containing the assistant's response to the corresponding input. The model used is set by the responses_model function. The default model is gpt-4o-mini.

Parameters:

Name	Type	Description	Default
instructions	str	System prompt prepended to every user message.	required
response_format	Type[T]	Pydantic model or built‑in type the assistant should return. Defaults to str.	str
batch_size	int	Number of prompts grouped into a single request. Defaults to 128.	128
temperature	float	Sampling temperature. Defaults to 0.0.	0.0
top_p	float	Nucleus sampling parameter. Defaults to 1.0.	1.0
max_concurrency	int	Maximum number of concurrent requests. Defaults to 8.	8

Returns:

Type	Description
Series	pandas.Series: Series whose values are instances of response_format.

Note

This is an asynchronous method and must be awaited.

Source code in src/openaivec/pandas_ext.py

async def responses(
    self,
    instructions: str,
    response_format: Type[T] = str,
    batch_size: int = 128,
    temperature: float = 0.0,
    top_p: float = 1.0,
    max_concurrency: int = 8,
) -> pd.Series:
    """Call an LLM once for every Series element (asynchronously).

    Example:
        ```python
        animals = pd.Series(["cat", "dog", "elephant"])
        # Must be awaited
        results = await animals.aio.responses("translate to French")
        ```
        This method returns a Series of strings, each containing the
        assistant's response to the corresponding input.
        The model used is set by the `responses_model` function.
        The default model is `gpt-4o-mini`.

    Args:
        instructions (str): System prompt prepended to every user message.
        response_format (Type[T], optional): Pydantic model or built‑in
            type the assistant should return. Defaults to ``str``.
        batch_size (int, optional): Number of prompts grouped into a single
            request. Defaults to ``128``.
        temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
        top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.
        max_concurrency (int, optional): Maximum number of concurrent
            requests. Defaults to ``8``.

    Returns:
        pandas.Series: Series whose values are instances of ``response_format``.

    Note:
        This is an asynchronous method and must be awaited.
    """
    client: AsyncBatchResponses = AsyncBatchResponses(
        client=_get_async_openai_client(),
        model_name=_RESPONSES_MODEL_NAME,
        system_message=instructions,
        response_format=response_format,
        temperature=temperature,
        top_p=top_p,
        max_concurrency=max_concurrency,
    )

    # Await the async operation
    results = await client.parse(self._obj.tolist(), batch_size=batch_size)

    return pd.Series(
        results,
        index=self._obj.index,
        name=self._obj.name,
    )

OpenAIVecDataFrameAccessor

pandas DataFrame accessor (.ai) that adds OpenAI helpers.

Source code in src/openaivec/pandas_ext.py

@pd.api.extensions.register_dataframe_accessor("ai")
class OpenAIVecDataFrameAccessor:
    """pandas DataFrame accessor (``.ai``) that adds OpenAI helpers."""

    def __init__(self, df_obj: pd.DataFrame):
        self._obj = df_obj

    def extract(self, column: str) -> pd.DataFrame:
        """Flatten one column of Pydantic models/dicts into top‑level columns.

        Example:
            ```python
            df = pd.DataFrame([
                {"animal": {"name": "cat", "legs": 4}},
                {"animal": {"name": "dog", "legs": 4}},
                {"animal": {"name": "elephant", "legs": 4}},
            ])
            df.ai.extract("animal")
            ```
            This method returns a DataFrame with the same index as the original,
            where each column corresponds to a key in the dictionaries.
            The source column is dropped.

        Args:
            column (str): Column to expand.

        Returns:
            pandas.DataFrame: Original DataFrame with the extracted columns; the source column is dropped.
        """
        if column not in self._obj.columns:
            raise ValueError(f"Column '{column}' does not exist in the DataFrame.")

        return (
            self._obj.pipe(lambda df: df.reset_index(drop=True))
            .pipe(lambda df: df.join(df[column].ai.extract()))
            .pipe(lambda df: df.set_index(self._obj.index))
            .pipe(lambda df: df.drop(columns=[column], axis=1))
        )

    def responses(
        self,
        instructions: str,
        response_format: Type[T] = str,
        batch_size: int = 128,
        temperature: float = 0.0,
        top_p: float = 1.0,
    ) -> pd.Series:
        """Generate a response for each row after serialising it to JSON.

        Example:
            ```python
            df = pd.DataFrame([
                {"name": "cat", "legs": 4},
                {"name": "dog", "legs": 4},
                {"name": "elephant", "legs": 4},
            ])
            df.ai.responses("what is the animal's name?")
            ```
            This method returns a Series of strings, each containing the
            assistant's response to the corresponding input.
            Each row is serialised to JSON before being sent to the assistant.
            The model used is set by the `responses_model` function.
            The default model is `gpt-4o-mini`.

        Args:
            instructions (str): System prompt for the assistant.
            response_format (Type[T], optional): Desired Python type of the
                responses. Defaults to ``str``.
            batch_size (int, optional): Number of requests sent in one batch.
                Defaults to ``128``.
            temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
            top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.

        Returns:
            pandas.Series: Responses aligned with the DataFrame’s original index.
        """
        return self._obj.pipe(
            lambda df: (
                df.pipe(lambda df: pd.Series(df.to_dict(orient="records"), index=df.index, name="record"))
                .map(lambda x: json.dumps(x, ensure_ascii=False))
                .ai.responses(
                    instructions=instructions,
                    response_format=response_format,
                    batch_size=batch_size,
                    temperature=temperature,
                    top_p=top_p,
                )
            )
        )

    def similarity(self, col1: str, col2: str) -> pd.Series:
        return self._obj.apply(
            lambda row: np.dot(row[col1], row[col2]) / (np.linalg.norm(row[col1]) * np.linalg.norm(row[col2])),
            axis=1,
        ).rename("similarity")

extract(column)

Flatten one column of Pydantic models/dicts into top‑level columns.

Example

df = pd.DataFrame([
    {"animal": {"name": "cat", "legs": 4}},
    {"animal": {"name": "dog", "legs": 4}},
    {"animal": {"name": "elephant", "legs": 4}},
])
df.ai.extract("animal")

This method returns a DataFrame with the same index as the original, where each column corresponds to a key in the dictionaries. The source column is dropped.

Parameters:

Name	Type	Description	Default
column	str	Column to expand.	required

Returns:

Type	Description
DataFrame	pandas.DataFrame: Original DataFrame with the extracted columns; the source column is dropped.

Source code in src/openaivec/pandas_ext.py

def extract(self, column: str) -> pd.DataFrame:
    """Flatten one column of Pydantic models/dicts into top‑level columns.

    Example:
        ```python
        df = pd.DataFrame([
            {"animal": {"name": "cat", "legs": 4}},
            {"animal": {"name": "dog", "legs": 4}},
            {"animal": {"name": "elephant", "legs": 4}},
        ])
        df.ai.extract("animal")
        ```
        This method returns a DataFrame with the same index as the original,
        where each column corresponds to a key in the dictionaries.
        The source column is dropped.

    Args:
        column (str): Column to expand.

    Returns:
        pandas.DataFrame: Original DataFrame with the extracted columns; the source column is dropped.
    """
    if column not in self._obj.columns:
        raise ValueError(f"Column '{column}' does not exist in the DataFrame.")

    return (
        self._obj.pipe(lambda df: df.reset_index(drop=True))
        .pipe(lambda df: df.join(df[column].ai.extract()))
        .pipe(lambda df: df.set_index(self._obj.index))
        .pipe(lambda df: df.drop(columns=[column], axis=1))
    )

responses(instructions, response_format=str, batch_size=128, temperature=0.0, top_p=1.0)

Generate a response for each row after serialising it to JSON.

Example

df = pd.DataFrame([
    {"name": "cat", "legs": 4},
    {"name": "dog", "legs": 4},
    {"name": "elephant", "legs": 4},
])
df.ai.responses("what is the animal's name?")

This method returns a Series of strings, each containing the assistant's response to the corresponding input. Each row is serialised to JSON before being sent to the assistant. The model used is set by the responses_model function. The default model is gpt-4o-mini.

Parameters:

Name	Type	Description	Default
instructions	str	System prompt for the assistant.	required
response_format	Type[T]	Desired Python type of the responses. Defaults to str.	str
batch_size	int	Number of requests sent in one batch. Defaults to 128.	128
temperature	float	Sampling temperature. Defaults to 0.0.	0.0
top_p	float	Nucleus sampling parameter. Defaults to 1.0.	1.0

Returns:

Type	Description
Series	pandas.Series: Responses aligned with the DataFrame’s original index.

Source code in src/openaivec/pandas_ext.py

def responses(
    self,
    instructions: str,
    response_format: Type[T] = str,
    batch_size: int = 128,
    temperature: float = 0.0,
    top_p: float = 1.0,
) -> pd.Series:
    """Generate a response for each row after serialising it to JSON.

    Example:
        ```python
        df = pd.DataFrame([
            {"name": "cat", "legs": 4},
            {"name": "dog", "legs": 4},
            {"name": "elephant", "legs": 4},
        ])
        df.ai.responses("what is the animal's name?")
        ```
        This method returns a Series of strings, each containing the
        assistant's response to the corresponding input.
        Each row is serialised to JSON before being sent to the assistant.
        The model used is set by the `responses_model` function.
        The default model is `gpt-4o-mini`.

    Args:
        instructions (str): System prompt for the assistant.
        response_format (Type[T], optional): Desired Python type of the
            responses. Defaults to ``str``.
        batch_size (int, optional): Number of requests sent in one batch.
            Defaults to ``128``.
        temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
        top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.

    Returns:
        pandas.Series: Responses aligned with the DataFrame’s original index.
    """
    return self._obj.pipe(
        lambda df: (
            df.pipe(lambda df: pd.Series(df.to_dict(orient="records"), index=df.index, name="record"))
            .map(lambda x: json.dumps(x, ensure_ascii=False))
            .ai.responses(
                instructions=instructions,
                response_format=response_format,
                batch_size=batch_size,
                temperature=temperature,
                top_p=top_p,
            )
        )
    )

OpenAIVecSeriesAccessor

pandas Series accessor (.ai) that adds OpenAI helpers.

Source code in src/openaivec/pandas_ext.py

@pd.api.extensions.register_series_accessor("ai")
class OpenAIVecSeriesAccessor:
    """pandas Series accessor (``.ai``) that adds OpenAI helpers."""

    def __init__(self, series_obj: pd.Series):
        self._obj = series_obj

    def responses(
        self,
        instructions: str,
        response_format: Type[T] = str,
        batch_size: int = 128,
        temperature: float = 0.0,
        top_p: float = 1.0,
    ) -> pd.Series:
        """Call an LLM once for every Series element.

        Example:
            ```python
            animals = pd.Series(["cat", "dog", "elephant"])
            animals.ai.responses("translate to French")
            ```
            This method returns a Series of strings, each containing the
            assistant's response to the corresponding input.
            The model used is set by the `responses_model` function.
            The default model is `gpt-4o-mini`.

        Args:
            instructions (str): System prompt prepended to every user message.
            response_format (Type[T], optional): Pydantic model or built‑in
                type the assistant should return. Defaults to ``str``.
            batch_size (int, optional): Number of prompts grouped into a single
                request. Defaults to ``128``.
            temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
            top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.

        Returns:
            pandas.Series: Series whose values are instances of ``response_format``.
        """
        client: BatchResponses = BatchResponses(
            client=_get_openai_client(),
            model_name=_RESPONSES_MODEL_NAME,
            system_message=instructions,
            response_format=response_format,
            temperature=temperature,
            top_p=top_p,
        )

        return pd.Series(
            client.parse(self._obj.tolist(), batch_size=batch_size),
            index=self._obj.index,
            name=self._obj.name,
        )

    def embeddings(self, batch_size: int = 128) -> pd.Series:
        """Compute OpenAI embeddings for every Series element.

        Example:
            ```python
            animals = pd.Series(["cat", "dog", "elephant"])
            animals.ai.embeddings()
            ```
            This method returns a Series of numpy arrays, each containing the
            embedding vector for the corresponding input.
            The embedding model is set by the `embeddings_model` function.
            The default embedding model is `text-embedding-3-small`.

        Args:
            batch_size (int, optional): Number of inputs grouped into a
                single request. Defaults to ``128``.

        Returns:
            pandas.Series: Series whose values are ``np.ndarray`` objects
                (dtype ``float32``).
        """
        client: BatchEmbeddings = BatchEmbeddings(
            client=_get_openai_client(),
            model_name=_EMBEDDINGS_MODEL_NAME,
        )

        return pd.Series(
            client.create(self._obj.tolist(), batch_size=batch_size),
            index=self._obj.index,
            name=self._obj.name,
        )

    def count_tokens(self) -> pd.Series:
        """Count `tiktoken` tokens per row.

        Example:
            ```python
            animals = pd.Series(["cat", "dog", "elephant"])
            animals.ai.count_tokens()
            ```
            This method uses the `tiktoken` library to count tokens based on the
            model name set by `responses_model`.

        Returns:
            pandas.Series: Token counts for each element.
        """
        return self._obj.map(_TIKTOKEN_ENCODING.encode).map(len).rename("num_tokens")

    def extract(self) -> pd.DataFrame:
        """Expand a Series of Pydantic models/dicts into columns.

        Example:
            ```python
            animals = pd.Series([
                {"name": "cat", "legs": 4},
                {"name": "dog", "legs": 4},
                {"name": "elephant", "legs": 4},
            ])
            animals.ai.extract()
            ```
            This method returns a DataFrame with the same index as the Series,
            where each column corresponds to a key in the dictionaries.
            If the Series has a name, extracted columns are prefixed with it.

        Returns:
            pandas.DataFrame: Expanded representation.
        """
        extracted = pd.DataFrame(
            self._obj.map(lambda x: _extract_value(x, self._obj.name)).tolist(),
            index=self._obj.index,
        )

        if self._obj.name:
            # If the Series has a name and all elements are dict or BaseModel, use it as the prefix for the columns
            extracted.columns = [f"{self._obj.name}_{col}" for col in extracted.columns]
        return extracted

count_tokens()

Count tiktoken tokens per row.

Example

animals = pd.Series(["cat", "dog", "elephant"])
animals.ai.count_tokens()

This method uses the tiktoken library to count tokens based on the model name set by responses_model.

Returns:

Type	Description
Series	pandas.Series: Token counts for each element.

Source code in src/openaivec/pandas_ext.py

def count_tokens(self) -> pd.Series:
    """Count `tiktoken` tokens per row.

    Example:
        ```python
        animals = pd.Series(["cat", "dog", "elephant"])
        animals.ai.count_tokens()
        ```
        This method uses the `tiktoken` library to count tokens based on the
        model name set by `responses_model`.

    Returns:
        pandas.Series: Token counts for each element.
    """
    return self._obj.map(_TIKTOKEN_ENCODING.encode).map(len).rename("num_tokens")

embeddings(batch_size=128)

Compute OpenAI embeddings for every Series element.

Example

animals = pd.Series(["cat", "dog", "elephant"])
animals.ai.embeddings()

This method returns a Series of numpy arrays, each containing the embedding vector for the corresponding input. The embedding model is set by the embeddings_model function. The default embedding model is text-embedding-3-small.

Parameters:

Name	Type	Description	Default
batch_size	int	Number of inputs grouped into a single request. Defaults to 128.	128

Returns:

Type	Description
Series	pandas.Series: Series whose values are np.ndarray objects (dtype float32).

Source code in src/openaivec/pandas_ext.py

def embeddings(self, batch_size: int = 128) -> pd.Series:
    """Compute OpenAI embeddings for every Series element.

    Example:
        ```python
        animals = pd.Series(["cat", "dog", "elephant"])
        animals.ai.embeddings()
        ```
        This method returns a Series of numpy arrays, each containing the
        embedding vector for the corresponding input.
        The embedding model is set by the `embeddings_model` function.
        The default embedding model is `text-embedding-3-small`.

    Args:
        batch_size (int, optional): Number of inputs grouped into a
            single request. Defaults to ``128``.

    Returns:
        pandas.Series: Series whose values are ``np.ndarray`` objects
            (dtype ``float32``).
    """
    client: BatchEmbeddings = BatchEmbeddings(
        client=_get_openai_client(),
        model_name=_EMBEDDINGS_MODEL_NAME,
    )

    return pd.Series(
        client.create(self._obj.tolist(), batch_size=batch_size),
        index=self._obj.index,
        name=self._obj.name,
    )

extract()

Expand a Series of Pydantic models/dicts into columns.

Example

animals = pd.Series([
    {"name": "cat", "legs": 4},
    {"name": "dog", "legs": 4},
    {"name": "elephant", "legs": 4},
])
animals.ai.extract()

This method returns a DataFrame with the same index as the Series, where each column corresponds to a key in the dictionaries. If the Series has a name, extracted columns are prefixed with it.

Returns:

Type	Description
DataFrame	pandas.DataFrame: Expanded representation.

Source code in src/openaivec/pandas_ext.py

def extract(self) -> pd.DataFrame:
    """Expand a Series of Pydantic models/dicts into columns.

    Example:
        ```python
        animals = pd.Series([
            {"name": "cat", "legs": 4},
            {"name": "dog", "legs": 4},
            {"name": "elephant", "legs": 4},
        ])
        animals.ai.extract()
        ```
        This method returns a DataFrame with the same index as the Series,
        where each column corresponds to a key in the dictionaries.
        If the Series has a name, extracted columns are prefixed with it.

    Returns:
        pandas.DataFrame: Expanded representation.
    """
    extracted = pd.DataFrame(
        self._obj.map(lambda x: _extract_value(x, self._obj.name)).tolist(),
        index=self._obj.index,
    )

    if self._obj.name:
        # If the Series has a name and all elements are dict or BaseModel, use it as the prefix for the columns
        extracted.columns = [f"{self._obj.name}_{col}" for col in extracted.columns]
    return extracted

responses(instructions, response_format=str, batch_size=128, temperature=0.0, top_p=1.0)

Call an LLM once for every Series element.

Example

animals = pd.Series(["cat", "dog", "elephant"])
animals.ai.responses("translate to French")

This method returns a Series of strings, each containing the assistant's response to the corresponding input. The model used is set by the responses_model function. The default model is gpt-4o-mini.

Parameters:

Name	Type	Description	Default
instructions	str	System prompt prepended to every user message.	required
response_format	Type[T]	Pydantic model or built‑in type the assistant should return. Defaults to str.	str
batch_size	int	Number of prompts grouped into a single request. Defaults to 128.	128
temperature	float	Sampling temperature. Defaults to 0.0.	0.0
top_p	float	Nucleus sampling parameter. Defaults to 1.0.	1.0

Returns:

Type	Description
Series	pandas.Series: Series whose values are instances of response_format.

Source code in src/openaivec/pandas_ext.py

def responses(
    self,
    instructions: str,
    response_format: Type[T] = str,
    batch_size: int = 128,
    temperature: float = 0.0,
    top_p: float = 1.0,
) -> pd.Series:
    """Call an LLM once for every Series element.

    Example:
        ```python
        animals = pd.Series(["cat", "dog", "elephant"])
        animals.ai.responses("translate to French")
        ```
        This method returns a Series of strings, each containing the
        assistant's response to the corresponding input.
        The model used is set by the `responses_model` function.
        The default model is `gpt-4o-mini`.

    Args:
        instructions (str): System prompt prepended to every user message.
        response_format (Type[T], optional): Pydantic model or built‑in
            type the assistant should return. Defaults to ``str``.
        batch_size (int, optional): Number of prompts grouped into a single
            request. Defaults to ``128``.
        temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
        top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.

    Returns:
        pandas.Series: Series whose values are instances of ``response_format``.
    """
    client: BatchResponses = BatchResponses(
        client=_get_openai_client(),
        model_name=_RESPONSES_MODEL_NAME,
        system_message=instructions,
        response_format=response_format,
        temperature=temperature,
        top_p=top_p,
    )

    return pd.Series(
        client.parse(self._obj.tolist(), batch_size=batch_size),
        index=self._obj.index,
        name=self._obj.name,
    )

embeddings_model(name)

Override the model used for text embeddings.

Parameters:

Name	Type	Description	Default
name	str	Embedding model name, e.g. text-embedding-3-small.	required

Source code in src/openaivec/pandas_ext.py

def embeddings_model(name: str) -> None:
    """Override the model used for text embeddings.

    Args:
        name (str): Embedding model name, e.g. ``text-embedding-3-small``.
    """
    global _EMBEDDINGS_MODEL_NAME
    _EMBEDDINGS_MODEL_NAME = name

responses_model(name)

Override the model used for text responses.

Parameters:

Name	Type	Description	Default
name	str	Model name as listed in the OpenAI API (for example, gpt-4o-mini).	required

Source code in src/openaivec/pandas_ext.py

def responses_model(name: str) -> None:
    """Override the model used for text responses.

    Args:
        name (str): Model name as listed in the OpenAI API
            (for example, ``gpt-4o-mini``).
    """
    global _RESPONSES_MODEL_NAME, _TIKTOKEN_ENCODING
    _RESPONSES_MODEL_NAME = name

    try:
        _TIKTOKEN_ENCODING = tiktoken.encoding_for_model(name)

    except KeyError:
        _LOGGER.info(
            "The model name '%s' is not supported by tiktoken. Instead, using the 'o200k_base' encoding.",
            name,
        )
        _TIKTOKEN_ENCODING = tiktoken.get_encoding("o200k_base")

use(client)

Register a custom OpenAI‑compatible client.

Parameters:

Name	Type	Description	Default
client	OpenAI	A pre‑configured openai.OpenAI or openai.AzureOpenAI instance. The same instance is reused by every helper in this module.	required

Source code in src/openaivec/pandas_ext.py

def use(client: OpenAI) -> None:
    """Register a custom OpenAI‑compatible client.

    Args:
        client (OpenAI): A pre‑configured `openai.OpenAI` or
            `openai.AzureOpenAI` instance.
            The same instance is reused by every helper in this module.
    """
    global _CLIENT
    _CLIENT = client

use_async(client)

Register a custom asynchronous OpenAI‑compatible client.

Parameters:

Name	Type	Description	Default
client	AsyncOpenAI	A pre‑configured openai.AsyncOpenAI or openai.AsyncAzureOpenAI instance. The same instance is reused by every helper in this module.	required

Source code in src/openaivec/pandas_ext.py

def use_async(client: AsyncOpenAI) -> None:
    """Register a custom asynchronous OpenAI‑compatible client.

    Args:
        client (AsyncOpenAI): A pre‑configured `openai.AsyncOpenAI` or
            `openai.AsyncAzureOpenAI` instance.
            The same instance is reused by every helper in this module.
    """
    global _ASYNC_CLIENT
    _ASYNC_CLIENT = client

use_azure_openai(api_key, endpoint, api_version)

Create and register an openai.AzureOpenAI client.

Parameters:

Name	Type	Description	Default
api_key	str	Azure OpenAI subscription key.	required
endpoint	str	Resource endpoint, e.g. https://<resource>.openai.azure.com.	required
api_version	str	REST API version such as 2024‑02‑15-preview.	required

Source code in src/openaivec/pandas_ext.py

def use_azure_openai(api_key: str, endpoint: str, api_version: str) -> None:
    """Create and register an `openai.AzureOpenAI` client.

    Args:
        api_key (str): Azure OpenAI subscription key.
        endpoint (str): Resource endpoint, e.g.
            ``https://<resource>.openai.azure.com``.
        api_version (str): REST API version such as ``2024‑02‑15-preview``.
    """
    global _CLIENT, _ASYNC_CLIENT
    _CLIENT = AzureOpenAI(
        api_key=api_key,
        azure_endpoint=endpoint,
        api_version=api_version,
    )
    _ASYNC_CLIENT = AsyncAzureOpenAI(
        api_key=api_key,
        azure_endpoint=endpoint,
        api_version=api_version,
    )

use_openai(api_key)

Create and register a default openai.OpenAI client.

Parameters:

Name	Type	Description	Default
api_key	str	Value forwarded to the api_key parameter of openai.OpenAI.	required

Source code in src/openaivec/pandas_ext.py

def use_openai(api_key: str) -> None:
    """Create and register a default `openai.OpenAI` client.

    Args:
        api_key (str): Value forwarded to the ``api_key`` parameter of
            `openai.OpenAI`.
    """
    global _CLIENT, _ASYNC_CLIENT
    _CLIENT = OpenAI(api_key=api_key)
    _ASYNC_CLIENT = AsyncOpenAI(api_key=api_key)