openaivec 0.13.5__tar.gz → 0.13.6__tar.gz

{openaivec-0.13.5 → openaivec-0.13.6}/.github/workflows/python-test.yml

@@ -27,5 +27,8 @@ jobs:
       - name: Lint with ruff
         run: uv run ruff check .
 
+      - name: Type check with pyright
+        run: uv run pyright src/openaivec || echo "Type check completed with issues - see above"
+
       - name: Run tests
         run: uv run pytest

{openaivec-0.13.5 → openaivec-0.13.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openaivec
-Version: 0.13.5
+Version: 0.13.6
 Summary: Generative mutation for tabular calculation
 Project-URL: Homepage, https://microsoft.github.io/openaivec/
 Project-URL: Repository, https://github.com/microsoft/openaivec
@@ -159,7 +159,7 @@ client = BatchResponses.of(
     client=OpenAI(),
     model_name="gpt-4.1-mini",
     system_message="Please answer only with 'xx family' and do not output anything else.",
-    batch_size=32,
+    # batch_size defaults to None (automatic optimization)
 )
 
 result = client.parse(["panda", "rabbit", "koala"])
@@ -304,7 +304,7 @@ async def process_data():
     # Asynchronous processing with fine-tuned concurrency control
     results = await df["text"].aio.responses(
         "Analyze sentiment and classify as positive/negative/neutral",
-        batch_size=64,        # Process 64 items per API request
+        # batch_size defaults to None (automatic optimization)
         max_concurrency=12    # Allow up to 12 concurrent requests
     )
     return results
@@ -315,7 +315,7 @@ sentiments = asyncio.run(process_data())
 
 **Key Parameters for Performance Tuning:**
 
-- **`batch_size`** (default: 128): Controls how many inputs are grouped into a single API request. Higher values reduce API call overhead but increase memory usage and request processing time.
+- **`batch_size`** (default: None): Controls how many inputs are grouped into a single API request. When None (default), automatic batch size optimization adjusts based on execution time. Set to a positive integer for fixed batch size. Higher values reduce API call overhead but increase memory usage and request processing time.
 - **`max_concurrency`** (default: 8): Limits the number of concurrent API requests. Higher values increase throughput but may hit rate limits or overwhelm the API.
 
 **Performance Benefits:**
@@ -460,12 +460,12 @@ When using openaivec with Spark, proper configuration of `batch_size` and `max_c
 - **Transparent**: Works automatically without code changes - your existing UDFs become more efficient
 - **Partition-Level**: Each partition maintains its own cache, optimal for distributed processing patterns
 
-**`batch_size`** (default: 128):
+**`batch_size`** (default: None):
 
 - Controls how many rows are processed together in each API request within a partition
-- **Larger values**: Fewer API calls per partition, reduced overhead
-- **Smaller values**: More granular processing, better memory management
-- **Recommendation**: 32-128 depending on data complexity and partition size
+- **Default (None)**: Automatic batch size optimization adjusts based on execution time
+- **Positive integer**: Fixed batch size - larger values reduce API calls but increase memory usage
+- **Recommendation**: Use default automatic optimization, or set 32-128 for fixed batch size
 
 **`max_concurrency`** (default: 8):
 
@@ -483,7 +483,7 @@ spark.udf.register(
     "analyze_sentiment",
     responses_udf(
         instructions="Analyze sentiment as positive/negative/neutral",
-        batch_size=64,        # Good balance for most use cases
+        # batch_size defaults to None (automatic optimization)
         max_concurrency=8     # 80 total concurrent requests across cluster
     )
 )

{openaivec-0.13.5 → openaivec-0.13.6}/README.md

@@ -133,7 +133,7 @@ client = BatchResponses.of(
     client=OpenAI(),
     model_name="gpt-4.1-mini",
     system_message="Please answer only with 'xx family' and do not output anything else.",
-    batch_size=32,
+    # batch_size defaults to None (automatic optimization)
 )
 
 result = client.parse(["panda", "rabbit", "koala"])
@@ -278,7 +278,7 @@ async def process_data():
     # Asynchronous processing with fine-tuned concurrency control
     results = await df["text"].aio.responses(
         "Analyze sentiment and classify as positive/negative/neutral",
-        batch_size=64,        # Process 64 items per API request
+        # batch_size defaults to None (automatic optimization)
         max_concurrency=12    # Allow up to 12 concurrent requests
     )
     return results
@@ -289,7 +289,7 @@ sentiments = asyncio.run(process_data())
 
 **Key Parameters for Performance Tuning:**
 
-- **`batch_size`** (default: 128): Controls how many inputs are grouped into a single API request. Higher values reduce API call overhead but increase memory usage and request processing time.
+- **`batch_size`** (default: None): Controls how many inputs are grouped into a single API request. When None (default), automatic batch size optimization adjusts based on execution time. Set to a positive integer for fixed batch size. Higher values reduce API call overhead but increase memory usage and request processing time.
 - **`max_concurrency`** (default: 8): Limits the number of concurrent API requests. Higher values increase throughput but may hit rate limits or overwhelm the API.
 
 **Performance Benefits:**
@@ -434,12 +434,12 @@ When using openaivec with Spark, proper configuration of `batch_size` and `max_c
 - **Transparent**: Works automatically without code changes - your existing UDFs become more efficient
 - **Partition-Level**: Each partition maintains its own cache, optimal for distributed processing patterns
 
-**`batch_size`** (default: 128):
+**`batch_size`** (default: None):
 
 - Controls how many rows are processed together in each API request within a partition
-- **Larger values**: Fewer API calls per partition, reduced overhead
-- **Smaller values**: More granular processing, better memory management
-- **Recommendation**: 32-128 depending on data complexity and partition size
+- **Default (None)**: Automatic batch size optimization adjusts based on execution time
+- **Positive integer**: Fixed batch size - larger values reduce API calls but increase memory usage
+- **Recommendation**: Use default automatic optimization, or set 32-128 for fixed batch size
 
 **`max_concurrency`** (default: 8):
 
@@ -457,7 +457,7 @@ spark.udf.register(
     "analyze_sentiment",
     responses_udf(
         instructions="Analyze sentiment as positive/negative/neutral",
-        batch_size=64,        # Good balance for most use cases
+        # batch_size defaults to None (automatic optimization)
         max_concurrency=8     # 80 total concurrent requests across cluster
     )
 )

{openaivec-0.13.5 → openaivec-0.13.6}/pyproject.toml

@@ -39,6 +39,7 @@ dev = [
     "ipykernel>=6.29.5",
     "langdetect>=1.0.9",
     "pyarrow>=19.0.1",
+    "pyright>=1.1.403",
     "pyspark>=3.5.5",
     "pytest>=8.3.5",
     "pytest-asyncio",

{openaivec-0.13.5 → openaivec-0.13.6}/src/openaivec/embeddings.py

@@ -31,16 +31,17 @@ class BatchEmbeddings:
 
     client: OpenAI
     model_name: str
-    cache: BatchingMapProxy[str, NDArray[np.float32]] = field(default_factory=lambda: BatchingMapProxy(batch_size=128))
+    cache: BatchingMapProxy[str, NDArray[np.float32]] = field(default_factory=lambda: BatchingMapProxy(batch_size=None))
 
     @classmethod
-    def of(cls, client: OpenAI, model_name: str, batch_size: int = 128) -> "BatchEmbeddings":
+    def of(cls, client: OpenAI, model_name: str, batch_size: int | None = None) -> "BatchEmbeddings":
         """Factory constructor.
 
         Args:
             client (OpenAI): OpenAI client.
             model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
-            batch_size (int, optional): Max unique inputs per API call. Defaults to 128.
+            batch_size (int | None, optional): Max unique inputs per API call. Defaults to None
+                (automatic batch size optimization). Set to a positive integer for fixed batch size.
 
         Returns:
             BatchEmbeddings: Configured instance backed by a batching proxy.
@@ -127,7 +128,7 @@ class AsyncBatchEmbeddings:
     client: AsyncOpenAI
     model_name: str
     cache: AsyncBatchingMapProxy[str, NDArray[np.float32]] = field(
-        default_factory=lambda: AsyncBatchingMapProxy(batch_size=128, max_concurrency=8)
+        default_factory=lambda: AsyncBatchingMapProxy(batch_size=None, max_concurrency=8)
     )
 
     @classmethod
@@ -135,7 +136,7 @@ class AsyncBatchEmbeddings:
         cls,
         client: AsyncOpenAI,
         model_name: str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
         max_concurrency: int = 8,
     ) -> "AsyncBatchEmbeddings":
         """Factory constructor.
@@ -143,7 +144,8 @@ class AsyncBatchEmbeddings:
         Args:
             client (AsyncOpenAI): OpenAI async client.
             model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
-            batch_size (int, optional): Max unique inputs per API call. Defaults to 128.
+            batch_size (int | None, optional): Max unique inputs per API call. Defaults to None
+                (automatic batch size optimization). Set to a positive integer for fixed batch size.
             max_concurrency (int, optional): Max concurrent API calls. Defaults to 8.
 
         Returns:
@@ -155,8 +157,8 @@ class AsyncBatchEmbeddings:
             cache=AsyncBatchingMapProxy(batch_size=batch_size, max_concurrency=max_concurrency),
         )
 
-    @observe(_LOGGER)
     @backoff_async(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
+    @observe(_LOGGER)
     async def _embed_chunk(self, inputs: List[str]) -> List[NDArray[np.float32]]:
         """Embed one minibatch of strings asynchronously.
 
@@ -186,4 +188,4 @@ class AsyncBatchEmbeddings:
         Returns:
             List[NDArray[np.float32]]: Embedding vectors aligned to ``inputs``.
         """
-        return await self.cache.map(inputs, self._embed_chunk)
+        return await self.cache.map(inputs, self._embed_chunk)  # type: ignore[arg-type]

{openaivec-0.13.5 → openaivec-0.13.6}/src/openaivec/model.py

@@ -1,13 +1,11 @@
 from dataclasses import dataclass
-from typing import Type, TypeVar
+from typing import Generic, Type, TypeVar
 
-from pydantic import BaseModel
-
-ResponseFormat = TypeVar("ResponseFormat", bound=BaseModel | str)
+ResponseFormat = TypeVar("ResponseFormat")
 
 
 @dataclass(frozen=True)
-class PreparedTask:
+class PreparedTask(Generic[ResponseFormat]):
     """A data class representing a complete task configuration for OpenAI API calls.
 
     This class encapsulates all the necessary parameters for executing a task,
@@ -84,10 +82,10 @@ class OpenAIAPIKey:
     """Container for OpenAI API key configuration.
 
     Attributes:
-        value (str): The API key for OpenAI services.
+        value (str | None): The API key for OpenAI services.
     """
 
-    value: str
+    value: str | None
 
 
 @dataclass(frozen=True)
@@ -95,10 +93,10 @@ class AzureOpenAIAPIKey:
     """Container for Azure OpenAI API key configuration.
 
     Attributes:
-        value (str): The API key for Azure OpenAI services.
+        value (str | None): The API key for Azure OpenAI services.
     """
 
-    value: str
+    value: str | None
 
 
 @dataclass(frozen=True)
@@ -106,10 +104,10 @@ class AzureOpenAIBaseURL:
     """Container for Azure OpenAI base URL configuration.
 
     Attributes:
-        value (str): The base URL for Azure OpenAI services.
+        value (str | None): The base URL for Azure OpenAI services.
     """
 
-    value: str
+    value: str | None
 
 
 @dataclass(frozen=True)

{openaivec-0.13.5 → openaivec-0.13.6}/src/openaivec/optimize.py

@@ -21,7 +21,7 @@ class BatchSizeSuggester:
     min_duration: float = 30.0
     max_duration: float = 60.0
     step_ratio: float = 0.1
-    sample_size: int = 10
+    sample_size: int = 4
     _history: List[PerformanceMetric] = field(default_factory=list)
     _lock: threading.RLock = field(default_factory=threading.RLock, repr=False)
     _batch_size_changed_at: datetime | None = field(default=None, init=False)

{openaivec-0.13.5 → openaivec-0.13.6}/src/openaivec/pandas_ext.py

@@ -42,7 +42,7 @@ to easily interact with OpenAI APIs for tasks like generating responses or embed
 import inspect
 import json
 import logging
-from typing import Any, Awaitable, Callable, List, Type, TypeVar
+from typing import Awaitable, Callable, List, Type, TypeVar
 
 import numpy as np
 import pandas as pd
@@ -184,6 +184,7 @@ class OpenAIVecSeriesAccessor:
         Args:
             cache (BatchingMapProxy[str, np.ndarray]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
 
         Returns:
             pandas.Series: Series whose values are ``np.ndarray`` objects
@@ -217,7 +218,7 @@ class OpenAIVecSeriesAccessor:
         self,
         instructions: str,
         response_format: Type[ResponseFormat] = str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
         temperature: float | None = 0.0,
         top_p: float = 1.0,
         show_progress: bool = False,
@@ -247,8 +248,9 @@ class OpenAIVecSeriesAccessor:
             instructions (str): System prompt prepended to every user message.
             response_format (Type[ResponseFormat], 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``.
+            batch_size (int | None, optional): Number of prompts grouped into a single
+                request. Defaults to ``None`` (automatic batch size optimization
+                based on execution time). Set to a positive integer for fixed batch size.
             temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
             top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
@@ -266,7 +268,7 @@ class OpenAIVecSeriesAccessor:
 
     def task_with_cache(
         self,
-        task: PreparedTask,
+        task: PreparedTask[ResponseFormat],
         cache: BatchingMapProxy[str, ResponseFormat],
     ) -> pd.Series:
         """Execute a prepared task on every Series element using a provided cache.
@@ -280,6 +282,7 @@ class OpenAIVecSeriesAccessor:
                 response format, and other parameters for processing the inputs.
             cache (BatchingMapProxy[str, ResponseFormat]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
 
         Returns:
             pandas.Series: Series whose values are instances of the task's
@@ -311,7 +314,7 @@ class OpenAIVecSeriesAccessor:
         )
         return pd.Series(client.parse(self._obj.tolist()), index=self._obj.index, name=self._obj.name)
 
-    def task(self, task: PreparedTask, batch_size: int = 128, show_progress: bool = False) -> pd.Series:
+    def task(self, task: PreparedTask, batch_size: int | None = None, show_progress: bool = False) -> pd.Series:
         """Execute a prepared task on every Series element.
 
         This method applies a pre-configured task to each element in the Series,
@@ -343,8 +346,9 @@ class OpenAIVecSeriesAccessor:
         Args:
             task (PreparedTask): A pre-configured task containing instructions,
                 response format, and other parameters for processing the inputs.
-            batch_size (int, optional): Number of prompts grouped into a single
-                request to optimize API usage. Defaults to 128.
+            batch_size (int | None, optional): Number of prompts grouped into a single
+                request to optimize API usage. Defaults to ``None`` (automatic batch size
+                optimization based on execution time). Set to a positive integer for fixed batch size.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
 
         Returns:
@@ -356,7 +360,7 @@ class OpenAIVecSeriesAccessor:
             cache=BatchingMapProxy(batch_size=batch_size, show_progress=show_progress),
         )
 
-    def embeddings(self, batch_size: int = 128, show_progress: bool = False) -> pd.Series:
+    def embeddings(self, batch_size: int | None = None, show_progress: bool = False) -> pd.Series:
         """Compute OpenAI embeddings for every Series element.
 
         Example:
@@ -378,8 +382,9 @@ class OpenAIVecSeriesAccessor:
             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``.
+            batch_size (int | None, optional): Number of inputs grouped into a
+                single request. Defaults to ``None`` (automatic batch size optimization
+                based on execution time). Set to a positive integer for fixed batch size.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
 
         Returns:
@@ -494,6 +499,7 @@ class OpenAIVecDataFrameAccessor:
             instructions (str): System prompt for the assistant.
             cache (BatchingMapProxy[str, ResponseFormat]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
             response_format (Type[ResponseFormat], optional): Desired Python type of the
                 responses. Defaults to ``str``.
             temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
@@ -538,7 +544,7 @@ class OpenAIVecDataFrameAccessor:
         self,
         instructions: str,
         response_format: Type[ResponseFormat] = str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
         temperature: float | None = 0.0,
         top_p: float = 1.0,
         show_progress: bool = False,
@@ -573,8 +579,9 @@ class OpenAIVecDataFrameAccessor:
             instructions (str): System prompt for the assistant.
             response_format (Type[ResponseFormat], optional): Desired Python type of the
                 responses. Defaults to ``str``.
-            batch_size (int, optional): Number of requests sent in one batch.
-                Defaults to ``128``.
+            batch_size (int | None, optional): Number of requests sent in one batch.
+                Defaults to ``None`` (automatic batch size optimization
+                based on execution time). Set to a positive integer for fixed batch size.
             temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
             top_p (float, optional): Nucleus sampling parameter. Defaults to ``1.0``.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
@@ -590,7 +597,7 @@ class OpenAIVecDataFrameAccessor:
             top_p=top_p,
         )
 
-    def task(self, task: PreparedTask, batch_size: int = 128, show_progress: bool = False) -> pd.Series:
+    def task(self, task: PreparedTask, batch_size: int | None = None, show_progress: bool = False) -> pd.Series:
         """Execute a prepared task on each DataFrame row after serialising it to JSON.
 
         This method applies a pre-configured task to each row in the DataFrame,
@@ -618,8 +625,9 @@ class OpenAIVecDataFrameAccessor:
         Args:
             task (PreparedTask): A pre-configured task containing instructions,
                 response format, and other parameters for processing the inputs.
-            batch_size (int, optional): Number of requests sent in one batch
-                to optimize API usage. Defaults to 128.
+            batch_size (int | None, optional): Number of requests sent in one batch
+                to optimize API usage. Defaults to ``None`` (automatic batch size
+                optimization based on execution time). Set to a positive integer for fixed batch size.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
 
         Returns:
@@ -634,7 +642,7 @@ class OpenAIVecDataFrameAccessor:
             )
         )
 
-    def fillna(self, target_column_name: str, max_examples: int = 500, batch_size: int = 128) -> pd.DataFrame:
+    def fillna(self, target_column_name: str, max_examples: int = 500, batch_size: int | None = None) -> pd.DataFrame:
         """Fill missing values in a DataFrame column using AI-powered inference.
 
         This method uses machine learning to intelligently fill missing (NaN) values
@@ -648,8 +656,9 @@ class OpenAIVecDataFrameAccessor:
             max_examples (int, optional): The maximum number of example rows to use
                 for context when predicting missing values. Higher values may improve
                 accuracy but increase API costs and processing time. Defaults to 500.
-            batch_size (int, optional): Number of requests sent in one batch
-                to optimize API usage. Defaults to 128.
+            batch_size (int | None, optional): Number of requests sent in one batch
+                to optimize API usage. Defaults to ``None`` (automatic batch size
+                optimization based on execution time). Set to a positive integer for fixed batch size.
 
         Returns:
             pandas.DataFrame: A new DataFrame with missing values filled in the target
@@ -721,7 +730,7 @@ class OpenAIVecDataFrameAccessor:
         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")
+        ).rename("similarity")  # type: ignore[arg-type]
 
 
 @pd.api.extensions.register_series_accessor("aio")
@@ -750,6 +759,7 @@ class AsyncOpenAIVecSeriesAccessor:
             instructions (str): System prompt prepended to every user message.
             cache (AsyncBatchingMapProxy[str, ResponseFormat]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
             response_format (Type[ResponseFormat], optional): Pydantic model or built‑in
                 type the assistant should return. Defaults to ``str``.
             temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
@@ -804,6 +814,7 @@ class AsyncOpenAIVecSeriesAccessor:
         Args:
             cache (AsyncBatchingMapProxy[str, np.ndarray]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
 
         Returns:
             pandas.Series: Series whose values are ``np.ndarray`` objects
@@ -844,7 +855,7 @@ class AsyncOpenAIVecSeriesAccessor:
 
     async def task_with_cache(
         self,
-        task: PreparedTask,
+        task: PreparedTask[ResponseFormat],
         cache: AsyncBatchingMapProxy[str, ResponseFormat],
     ) -> pd.Series:
         """Execute a prepared task on every Series element using a provided cache (asynchronously).
@@ -859,6 +870,7 @@ class AsyncOpenAIVecSeriesAccessor:
                 response format, and other parameters for processing the inputs.
             cache (AsyncBatchingMapProxy[str, ResponseFormat]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
 
         Returns:
             pandas.Series: Series whose values are instances of the task's
@@ -902,7 +914,7 @@ class AsyncOpenAIVecSeriesAccessor:
         self,
         instructions: str,
         response_format: Type[ResponseFormat] = str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
         temperature: float | None = 0.0,
         top_p: float = 1.0,
         max_concurrency: int = 8,
@@ -934,8 +946,9 @@ class AsyncOpenAIVecSeriesAccessor:
             instructions (str): System prompt prepended to every user message.
             response_format (Type[ResponseFormat], 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``.
+            batch_size (int | None, optional): Number of prompts grouped into a single
+                request. Defaults to ``None`` (automatic batch size optimization
+                based on execution time). Set to a positive integer for fixed batch size.
             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
@@ -959,7 +972,7 @@ class AsyncOpenAIVecSeriesAccessor:
         )
 
     async def embeddings(
-        self, batch_size: int = 128, max_concurrency: int = 8, show_progress: bool = False
+        self, batch_size: int | None = None, max_concurrency: int = 8, show_progress: bool = False
     ) -> pd.Series:
         """Compute OpenAI embeddings for every Series element (asynchronously).
 
@@ -983,8 +996,9 @@ class AsyncOpenAIVecSeriesAccessor:
             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``.
+            batch_size (int | None, optional): Number of inputs grouped into a
+                single request. Defaults to ``None`` (automatic batch size optimization
+                based on execution time). Set to a positive integer for fixed batch size.
             max_concurrency (int, optional): Maximum number of concurrent
                 requests. Defaults to ``8``.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
@@ -1003,7 +1017,7 @@ class AsyncOpenAIVecSeriesAccessor:
         )
 
     async def task(
-        self, task: PreparedTask, batch_size: int = 128, max_concurrency: int = 8, show_progress: bool = False
+        self, task: PreparedTask, batch_size: int | None = None, max_concurrency: int = 8, show_progress: bool = False
     ) -> pd.Series:
         """Execute a prepared task on every Series element (asynchronously).
 
@@ -1037,8 +1051,9 @@ class AsyncOpenAIVecSeriesAccessor:
         Args:
             task (PreparedTask): A pre-configured task containing instructions,
                 response format, and other parameters for processing the inputs.
-            batch_size (int, optional): Number of prompts grouped into a single
-                request to optimize API usage. Defaults to 128.
+            batch_size (int | None, optional): Number of prompts grouped into a single
+                request to optimize API usage. Defaults to ``None`` (automatic batch size
+                optimization based on execution time). Set to a positive integer for fixed batch size.
             max_concurrency (int, optional): Maximum number of concurrent
                 requests. Defaults to 8.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
@@ -1084,6 +1099,7 @@ class AsyncOpenAIVecDataFrameAccessor:
             instructions (str): System prompt for the assistant.
             cache (AsyncBatchingMapProxy[str, ResponseFormat]): Pre-configured cache
                 instance for managing API call batching and deduplication.
+                Set cache.batch_size=None to enable automatic batch size optimization.
             response_format (Type[ResponseFormat], optional): Desired Python type of the
                 responses. Defaults to ``str``.
             temperature (float, optional): Sampling temperature. Defaults to ``0.0``.
@@ -1134,7 +1150,7 @@ class AsyncOpenAIVecDataFrameAccessor:
         self,
         instructions: str,
         response_format: Type[ResponseFormat] = str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
         temperature: float | None = 0.0,
         top_p: float = 1.0,
         max_concurrency: int = 8,
@@ -1171,8 +1187,9 @@ class AsyncOpenAIVecDataFrameAccessor:
             instructions (str): System prompt for the assistant.
             response_format (Type[ResponseFormat], optional): Desired Python type of the
                 responses. Defaults to ``str``.
-            batch_size (int, optional): Number of requests sent in one batch.
-                Defaults to ``128``.
+            batch_size (int | None, optional): Number of requests sent in one batch.
+                Defaults to ``None`` (automatic batch size optimization
+                based on execution time). Set to a positive integer for fixed batch size.
             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
@@ -1196,7 +1213,7 @@ class AsyncOpenAIVecDataFrameAccessor:
         )
 
     async def task(
-        self, task: PreparedTask, batch_size: int = 128, max_concurrency: int = 8, show_progress: bool = False
+        self, task: PreparedTask, batch_size: int | None = None, max_concurrency: int = 8, show_progress: bool = False
     ) -> pd.Series:
         """Execute a prepared task on each DataFrame row after serialising it to JSON (asynchronously).
 
@@ -1235,8 +1252,9 @@ class AsyncOpenAIVecDataFrameAccessor:
         Args:
             task (PreparedTask): A pre-configured task containing instructions,
                 response format, and other parameters for processing the inputs.
-            batch_size (int, optional): Number of requests sent in one batch
-                to optimize API usage. Defaults to 128.
+            batch_size (int | None, optional): Number of requests sent in one batch
+                to optimize API usage. Defaults to ``None`` (automatic batch size
+                optimization based on execution time). Set to a positive integer for fixed batch size.
             max_concurrency (int, optional): Maximum number of concurrent
                 requests. Defaults to 8.
             show_progress (bool, optional): Show progress bar in Jupyter notebooks. Defaults to ``False``.
@@ -1286,7 +1304,7 @@ class AsyncOpenAIVecDataFrameAccessor:
         else:
             return result
 
-    async def assign(self, **kwargs: Any) -> pd.DataFrame:
+    async def assign(self, **kwargs) -> pd.DataFrame:
         """Asynchronously assign new columns to the DataFrame, evaluating sequentially.
 
         This method extends pandas' `assign` method by supporting asynchronous
@@ -1321,7 +1339,7 @@ class AsyncOpenAIVecDataFrameAccessor:
             ```
 
         Args:
-            **kwargs: Any. Column names as keys and either static values or callables
+            **kwargs: Column names as keys and either static values or callables
                 (synchronous or asynchronous) as values.
 
         Returns:
@@ -1346,7 +1364,7 @@ class AsyncOpenAIVecDataFrameAccessor:
         return df_current
 
     async def fillna(
-        self, target_column_name: str, max_examples: int = 500, batch_size: int = 128, max_concurrency: int = 8
+        self, target_column_name: str, max_examples: int = 500, batch_size: int | None = None, max_concurrency: int = 8
     ) -> pd.DataFrame:
         """Fill missing values in a DataFrame column using AI-powered inference (asynchronously).
 
@@ -1361,8 +1379,9 @@ class AsyncOpenAIVecDataFrameAccessor:
             max_examples (int, optional): The maximum number of example rows to use
                 for context when predicting missing values. Higher values may improve
                 accuracy but increase API costs and processing time. Defaults to 500.
-            batch_size (int, optional): Number of requests sent in one batch
-                to optimize API usage. Defaults to 128.
+            batch_size (int | None, optional): Number of requests sent in one batch
+                to optimize API usage. Defaults to ``None`` (automatic batch size
+                optimization based on execution time). Set to a positive integer for fixed batch size.
             max_concurrency (int, optional): Maximum number of concurrent
                 requests. Defaults to 8.