openaivec 0.10.0__py3-none-any.whl → 1.0.10__py3-none-any.whl

openaivec/spark.py

@@ -1,53 +1,52 @@
 """Asynchronous Spark UDFs for the OpenAI and Azure OpenAI APIs.
 
-This module provides builder classes (`ResponsesUDFBuilder`, `EmbeddingsUDFBuilder`)
+This module provides functions (`responses_udf`, `task_udf`, `embeddings_udf`,
+`count_tokens_udf`, `split_to_chunks_udf`, `similarity_udf`, `parse_udf`)
 for creating asynchronous Spark UDFs that communicate with either the public
 OpenAI API or Azure OpenAI using the `openaivec.spark` subpackage.
-It supports UDFs for generating responses and creating embeddings asynchronously.
-The UDFs operate on Spark DataFrames and leverage asyncio for potentially
-improved performance in I/O-bound operations.
+It supports UDFs for generating responses, creating embeddings, parsing text,
+and computing similarities asynchronously. The UDFs operate on Spark DataFrames
+and leverage asyncio for improved performance in I/O-bound operations.
+
+**Performance Optimization**: All AI-powered UDFs (`responses_udf`, `task_udf`, `embeddings_udf`, `parse_udf`)
+automatically cache duplicate inputs within each partition, significantly reducing
+API calls and costs when processing datasets with overlapping content.
+
 
 ## Setup
 
-First, obtain a Spark session:
+First, obtain a Spark session and configure authentication:
 
 ```python
 from pyspark.sql import SparkSession
+from openaivec.spark import setup, setup_azure
 
 spark = SparkSession.builder.getOrCreate()
-```
-
-Next, instantiate UDF builders with your OpenAI API key (or Azure credentials)
-and model/deployment names, then register the desired UDFs:
-
-```python
-import os
-from openaivec.spark import ResponsesUDFBuilder, EmbeddingsUDFBuilder
-from pydantic import BaseModel
 
 # Option 1: Using OpenAI
-resp_builder = ResponsesUDFBuilder.of_openai(
-    api_key=os.getenv("OPENAI_API_KEY"),
-    model_name="gpt-4o-mini", # Model for responses
-)
-emb_builder = EmbeddingsUDFBuilder.of_openai(
-    api_key=os.getenv("OPENAI_API_KEY"),
-    model_name="text-embedding-3-small", # Model for embeddings
+setup(
+    spark,
+    api_key="your-openai-api-key",
+    responses_model_name="gpt-4.1-mini",  # Optional: set default model
+    embeddings_model_name="text-embedding-3-small"  # Optional: set default model
 )
 
 # Option 2: Using Azure OpenAI
-# resp_builder = ResponsesUDFBuilder.of_azure_openai(
-#     api_key=os.getenv("AZURE_OPENAI_KEY"),
-#     endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
-#     api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
-#     model_name="your-resp-deployment-name", # Deployment for responses
-# )
-# emb_builder = EmbeddingsUDFBuilder.of_azure_openai(
-#     api_key=os.getenv("AZURE_OPENAI_KEY"),
-#     endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
-#     api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
-#     model_name="your-emb-deployment-name", # Deployment for embeddings
+# setup_azure(
+#     spark,
+#     api_key="your-azure-openai-api-key",
+#     base_url="https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/",
+#     api_version="preview",
+#     responses_model_name="my-gpt4-deployment",  # Optional: set default deployment
+#     embeddings_model_name="my-embedding-deployment"  # Optional: set default deployment
 # )
+```
+
+Next, create UDFs and register them:
+
+```python
+from openaivec.spark import responses_udf, task_udf, embeddings_udf, count_tokens_udf, split_to_chunks_udf
+from pydantic import BaseModel
 
 # Define a Pydantic model for structured responses (optional)
 class Translation(BaseModel):
@@ -55,27 +54,39 @@ class Translation(BaseModel):
     fr: str
     # ... other languages
 
-# Register the asynchronous responses UDF
+# Register the asynchronous responses UDF with performance tuning
 spark.udf.register(
     "translate_async",
-    resp_builder.build(
+    responses_udf(
         instructions="Translate the text to multiple languages.",
         response_format=Translation,
+        model_name="gpt-4.1-mini",  # For Azure: deployment name, for OpenAI: model name
+        batch_size=64,              # Rows per API request within partition
+        max_concurrency=8           # Concurrent requests PER EXECUTOR
     ),
 )
 
-# Or use a predefined task with build_from_task method
+# Or use a predefined task with task_udf
 from openaivec.task import nlp
 spark.udf.register(
     "sentiment_async",
-    resp_builder.build_from_task(nlp.SENTIMENT_ANALYSIS),
+    task_udf(nlp.SENTIMENT_ANALYSIS),
 )
 
-# Register the asynchronous embeddings UDF
+# Register the asynchronous embeddings UDF with performance tuning
 spark.udf.register(
     "embed_async",
-    emb_builder.build(),
+    embeddings_udf(
+        model_name="text-embedding-3-small",  # For Azure: deployment name, for OpenAI: model name
+        batch_size=128,                       # Larger batches for embeddings
+        max_concurrency=8                     # Concurrent requests PER EXECUTOR
+    ),
 )
+
+# Register token counting, text chunking, and similarity UDFs
+spark.udf.register("count_tokens", count_tokens_udf())
+spark.udf.register("split_chunks", split_to_chunks_udf(max_tokens=512, sep=[".", "!", "?"]))
+spark.udf.register("compute_similarity", similarity_udf())
 ```
 
 You can now invoke the UDFs from Spark SQL:
@@ -85,79 +96,194 @@ SELECT
     text,
     translate_async(text) AS translation,
     sentiment_async(text) AS sentiment,
-    embed_async(text) AS embedding
+    embed_async(text) AS embedding,
+    count_tokens(text) AS token_count,
+    split_chunks(text) AS chunks,
+    compute_similarity(embed_async(text1), embed_async(text2)) AS similarity
 FROM your_table;
 ```
 
+## Performance Considerations
+
+When using these UDFs in distributed Spark environments:
+
+- **`batch_size`**: Controls rows processed per API request within each partition.
+  Recommended: 32-128 for responses, 64-256 for embeddings.
+
+- **`max_concurrency`**: Sets concurrent API requests **PER EXECUTOR**, not per cluster.
+  Total cluster concurrency = max_concurrency × number_of_executors.
+  Recommended: 4-12 per executor to avoid overwhelming OpenAI rate limits.
+
+- **Rate Limit Management**: Monitor OpenAI API usage when scaling executors.
+  Consider your OpenAI tier limits and adjust max_concurrency accordingly.
+
+Example for a 5-executor cluster with max_concurrency=8:
+Total concurrent requests = 8 × 5 = 40 simultaneous API calls.
+
 Note: This module provides asynchronous support through the pandas extensions.
 """
 
 import asyncio
-from dataclasses import dataclass
-from typing import Dict, Iterator, List, Type, TypeVar, Union, get_args, get_origin, Optional
-from typing_extensions import Literal
-from enum import Enum
 import logging
-from pyspark.sql.pandas.functions import pandas_udf
-from pyspark.sql.udf import UserDefinedFunction
-from pyspark.sql.types import BooleanType, IntegerType, StringType, ArrayType, FloatType, StructField, StructType
-from openai import AsyncOpenAI, AsyncAzureOpenAI
-import tiktoken
-from . import pandas_ext
+import os
+from collections.abc import Iterator
+from enum import Enum
+from typing import Union, get_args, get_origin
+
+import numpy as np
 import pandas as pd
+import tiktoken
 from pydantic import BaseModel
+from pyspark import SparkContext
+from pyspark.sql import SparkSession
+from pyspark.sql.pandas.functions import pandas_udf
+from pyspark.sql.types import ArrayType, BooleanType, FloatType, IntegerType, StringType, StructField, StructType
+from pyspark.sql.udf import UserDefinedFunction
+from typing_extensions import Literal
 
-from .serialize import deserialize_base_model, serialize_base_model
-from .util import TextChunker
-from .task.model import PreparedTask
+from openaivec import pandas_ext
+from openaivec._cache import AsyncBatchingMapProxy
+from openaivec._model import EmbeddingsModelName, PreparedTask, ResponseFormat, ResponsesModelName
+from openaivec._provider import CONTAINER
+from openaivec._schema import SchemaInferenceInput, SchemaInferenceOutput, SchemaInferer
+from openaivec._serialize import deserialize_base_model, serialize_base_model
+from openaivec._util import TextChunker
 
 __all__ = [
-    "ResponsesUDFBuilder",
-    "EmbeddingsUDFBuilder",
+    "setup",
+    "setup_azure",
+    "responses_udf",
+    "task_udf",
+    "embeddings_udf",
+    "infer_schema",
+    "parse_udf",
     "split_to_chunks_udf",
     "count_tokens_udf",
     "similarity_udf",
 ]
 
-ResponseFormat = BaseModel | Type[str]
-T = TypeVar("T", bound=BaseModel)
 
-_INITIALIZED: bool = False
 _LOGGER: logging.Logger = logging.getLogger(__name__)
-_TIKTOKEN_ENC: tiktoken.Encoding | None = None
 
 
-def _initialize(api_key: str, endpoint: str | None, api_version: str | None) -> None:
-    """Initializes the OpenAI client for asynchronous operations.
+def setup(
+    spark: SparkSession, api_key: str, responses_model_name: str | None = None, embeddings_model_name: str | None = None
+):
+    """Setup OpenAI authentication and default model names in Spark environment.
+    1. Configures OpenAI API key in SparkContext environment.
+    2. Configures OpenAI API key in local process environment.
+    3. Optionally registers default model names for responses and embeddings in the DI container.
+
+    Args:
+        spark (SparkSession): The Spark session to configure.
+        api_key (str): OpenAI API key for authentication.
+        responses_model_name (str | None): Default model name for response generation.
+            If provided, registers `ResponsesModelName` in the DI container.
+        embeddings_model_name (str | None): Default model name for embeddings.
+            If provided, registers `EmbeddingsModelName` in the DI container.
+
+    Example:
+        ```python
+        from pyspark.sql import SparkSession
+        from openaivec.spark import setup
+
+        spark = SparkSession.builder.getOrCreate()
+        setup(
+            spark,
+            api_key="sk-***",
+            responses_model_name="gpt-4.1-mini",
+            embeddings_model_name="text-embedding-3-small",
+        )
+        ```
+    """
+
+    CONTAINER.register(SparkSession, lambda: spark)
+    CONTAINER.register(SparkContext, lambda: CONTAINER.resolve(SparkSession).sparkContext)
+
+    sc = CONTAINER.resolve(SparkContext)
+    sc.environment["OPENAI_API_KEY"] = api_key
+
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    if responses_model_name:
+        CONTAINER.register(ResponsesModelName, lambda: ResponsesModelName(responses_model_name))
+
+    if embeddings_model_name:
+        CONTAINER.register(EmbeddingsModelName, lambda: EmbeddingsModelName(embeddings_model_name))
+
+    CONTAINER.clear_singletons()
 
-    This function sets up the global asynchronous OpenAI client instance
-    (either `AsyncOpenAI` or `AsyncAzureOpenAI`) used by the UDFs in this
-    module. It ensures the client is initialized only once.
 
+def setup_azure(
+    spark: SparkSession,
+    api_key: str,
+    base_url: str,
+    api_version: str = "preview",
+    responses_model_name: str | None = None,
+    embeddings_model_name: str | None = None,
+):
+    """Setup Azure OpenAI authentication and default model names in Spark environment.
+    1. Configures Azure OpenAI API key, base URL, and API version in SparkContext environment.
+    2. Configures Azure OpenAI API key, base URL, and API version in local process environment.
+    3. Optionally registers default model names for responses and embeddings in the DI container.
     Args:
-        api_key: The OpenAI or Azure OpenAI API key.
-        endpoint: The Azure OpenAI endpoint URL. Required for Azure.
-        api_version: The Azure OpenAI API version. Required for Azure.
+        spark (SparkSession): The Spark session to configure.
+        api_key (str): Azure OpenAI API key for authentication.
+        base_url (str): Base URL for the Azure OpenAI resource.
+        api_version (str): API version to use. Defaults to "preview".
+        responses_model_name (str | None): Default model name for response generation.
+            If provided, registers `ResponsesModelName` in the DI container.
+        embeddings_model_name (str | None): Default model name for embeddings.
+            If provided, registers `EmbeddingsModelName` in the DI container.
+
+    Example:
+        ```python
+        from pyspark.sql import SparkSession
+        from openaivec.spark import setup_azure
+
+        spark = SparkSession.builder.getOrCreate()
+        setup_azure(
+            spark,
+            api_key="azure-key",
+            base_url="https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/",
+            api_version="preview",
+            responses_model_name="gpt4-deployment",
+            embeddings_model_name="embedding-deployment",
+        )
+        ```
     """
-    global _INITIALIZED
-    if not _INITIALIZED:
-        if endpoint and api_version:
-            pandas_ext.use_async(AsyncAzureOpenAI(api_key=api_key, azure_endpoint=endpoint, api_version=api_version))
-        else:
-            pandas_ext.use_async(AsyncOpenAI(api_key=api_key))
-        _INITIALIZED = True
+
+    CONTAINER.register(SparkSession, lambda: spark)
+    CONTAINER.register(SparkContext, lambda: CONTAINER.resolve(SparkSession).sparkContext)
+
+    sc = CONTAINER.resolve(SparkContext)
+    sc.environment["AZURE_OPENAI_API_KEY"] = api_key
+    sc.environment["AZURE_OPENAI_BASE_URL"] = base_url
+    sc.environment["AZURE_OPENAI_API_VERSION"] = api_version
+
+    os.environ["AZURE_OPENAI_API_KEY"] = api_key
+    os.environ["AZURE_OPENAI_BASE_URL"] = base_url
+    os.environ["AZURE_OPENAI_API_VERSION"] = api_version
+
+    if responses_model_name:
+        CONTAINER.register(ResponsesModelName, lambda: ResponsesModelName(responses_model_name))
+
+    if embeddings_model_name:
+        CONTAINER.register(EmbeddingsModelName, lambda: EmbeddingsModelName(embeddings_model_name))
+
+    CONTAINER.clear_singletons()
 
 
 def _python_type_to_spark(python_type):
     origin = get_origin(python_type)
 
-    # For list types (e.g., List[int])
-    if origin is list or origin is List:
+    # For list types (e.g., list[int])
+    if origin is list:
         # Retrieve the inner type and recursively convert it
         inner_type = get_args(python_type)[0]
         return ArrayType(_python_type_to_spark(inner_type))
 
-    # For Optional types (Union[..., None])
+    # For Optional types (T | None via Union internally)
     elif origin is Union:
         non_none_args = [arg for arg in get_args(python_type) if arg is not type(None)]
         if len(non_none_args) == 1:
@@ -170,7 +296,7 @@ def _python_type_to_spark(python_type):
         return StringType()
 
     # For Enum types - also treat as StringType since Spark doesn't have enum types
-    elif hasattr(python_type, '__bases__') and Enum in python_type.__bases__:
+    elif hasattr(python_type, "__bases__") and Enum in python_type.__bases__:
         return StringType()
 
     # For nested Pydantic models (to be treated as Structs)
@@ -190,7 +316,7 @@ def _python_type_to_spark(python_type):
         raise ValueError(f"Unsupported type: {python_type}")
 
 
-def _pydantic_to_spark_schema(model: Type[BaseModel]) -> StructType:
+def _pydantic_to_spark_schema(model: type[BaseModel]) -> StructType:
     fields = []
     for field_name, field in model.model_fields.items():
         field_type = field.annotation
@@ -201,7 +327,7 @@ def _pydantic_to_spark_schema(model: Type[BaseModel]) -> StructType:
     return StructType(fields)
 
 
-def _safe_cast_str(x: Optional[str]) -> Optional[str]:
+def _safe_cast_str(x: str | None) -> str | None:
     try:
         if x is None:
             return None
@@ -212,7 +338,7 @@ def _safe_cast_str(x: Optional[str]) -> Optional[str]:
         return None
 
 
-def _safe_dump(x: Optional[BaseModel]) -> Dict:
+def _safe_dump(x: BaseModel | None) -> dict:
     try:
         if x is None:
             return {}
@@ -223,347 +349,490 @@ def _safe_dump(x: Optional[BaseModel]) -> Dict:
         return {}
 
 
-@dataclass(frozen=True)
-class ResponsesUDFBuilder:
-    """Builder for asynchronous Spark pandas UDFs for generating responses.
+def responses_udf(
+    instructions: str,
+    response_format: type[ResponseFormat] = str,
+    model_name: str | None = None,
+    batch_size: int | None = None,
+    max_concurrency: int = 8,
+    **api_kwargs,
+) -> UserDefinedFunction:
+    """Create an asynchronous Spark pandas UDF for generating responses.
 
-    Configures and builds UDFs that leverage `pandas_ext.aio.responses`
+    Configures and builds UDFs that leverage `pandas_ext.aio.responses_with_cache`
     to generate text or structured responses from OpenAI models asynchronously.
-    An instance stores authentication parameters and the model name.
+    Each partition maintains its own cache to eliminate duplicate API calls within
+    the partition, significantly reducing API usage and costs when processing
+    datasets with overlapping content.
+
+    Note:
+        Authentication must be configured via SparkContext environment variables.
+        Set the appropriate environment variables on the SparkContext:
+
+        For OpenAI:
+            sc.environment["OPENAI_API_KEY"] = "your-openai-api-key"
 
-    This builder supports two main methods:
-    - `build()`: Creates UDFs with custom instructions and response formats
-    - `build_from_task()`: Creates UDFs from predefined tasks (e.g., sentiment analysis)
+        For Azure OpenAI:
+            sc.environment["AZURE_OPENAI_API_KEY"] = "your-azure-openai-api-key"
+            sc.environment["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
+            sc.environment["AZURE_OPENAI_API_VERSION"] = "preview"
+
+    Args:
+        instructions (str): The system prompt or instructions for the model.
+        response_format (type[ResponseFormat]): The desired output format. Either `str` for plain text
+            or a Pydantic `BaseModel` for structured JSON output. Defaults to `str`.
+        model_name (str | None): For Azure OpenAI, use your deployment name (e.g., "my-gpt4-deployment").
+            For OpenAI, use the model name (e.g., "gpt-4.1-mini"). Defaults to configured model in DI container
+            via ResponsesModelName if not provided.
+        batch_size (int | None): Number of rows per async batch request within each partition.
+            Larger values reduce API call overhead but increase memory usage.
+            Defaults to None (automatic batch size optimization that dynamically
+            adjusts based on execution time, targeting 30-60 seconds per batch).
+            Set to a positive integer (e.g., 32-128) for fixed batch size.
+        max_concurrency (int): Maximum number of concurrent API requests **PER EXECUTOR**.
+            Total cluster concurrency = max_concurrency × number_of_executors.
+            Higher values increase throughput but may hit OpenAI rate limits.
+            Recommended: 4-12 per executor. Defaults to 8.
+        **api_kwargs: Additional OpenAI API parameters (e.g. ``temperature``, ``top_p``,
+            ``frequency_penalty``, ``presence_penalty``, ``seed``, ``max_output_tokens``, etc.)
+            forwarded verbatim to the underlying API calls. These parameters are applied to
+            all API requests made by the UDF.
 
-    Attributes:
-        api_key (str): OpenAI or Azure API key.
-        endpoint (Optional[str]): Azure endpoint base URL. None for public OpenAI.
-        api_version (Optional[str]): Azure API version. Ignored for public OpenAI.
-        model_name (str): Deployment name (Azure) or model name (OpenAI) for responses.
+    Returns:
+        UserDefinedFunction: A Spark pandas UDF configured to generate responses asynchronously.
+            Output schema is `StringType` or a struct derived from `response_format`.
+
+    Raises:
+        ValueError: If `response_format` is not `str` or a Pydantic `BaseModel`.
+
+    Example:
+        ```python
+        from pyspark.sql import SparkSession
+        from openaivec.spark import responses_udf, setup
+
+        spark = SparkSession.builder.getOrCreate()
+        setup(spark, api_key="sk-***", responses_model_name="gpt-4.1-mini")
+        udf = responses_udf("Reply with one word.")
+        spark.udf.register("short_answer", udf)
+        df = spark.createDataFrame([("hello",), ("bye",)], ["text"])
+        df.selectExpr("short_answer(text) as reply").show()
+        ```
+
+    Note:
+        For optimal performance in distributed environments:
+        - **Automatic Caching**: Duplicate inputs within each partition are cached,
+          reducing API calls and costs significantly on datasets with repeated content
+        - Monitor OpenAI API rate limits when scaling executor count
+        - Consider your OpenAI tier limits: total_requests = max_concurrency × executors
+        - Use Spark UI to optimize partition sizes relative to batch_size
     """
+    _model_name = model_name or CONTAINER.resolve(ResponsesModelName).value
 
-    # Params for OpenAI SDK
-    api_key: str
-    endpoint: str | None
-    api_version: str | None
-
-    # Params for Responses API
-    model_name: str
-
-    @classmethod
-    def of_openai(cls, api_key: str, model_name: str) -> "ResponsesUDFBuilder":
-        """Creates a builder configured for the public OpenAI API.
-
-        Args:
-            api_key (str): The OpenAI API key.
-            model_name (str): The OpenAI model name for responses (e.g., "gpt-4o-mini").
-
-        Returns:
-            ResponsesUDFBuilder: A builder instance configured for OpenAI responses.
-        """
-        return cls(api_key=api_key, endpoint=None, api_version=None, model_name=model_name)
-
-    @classmethod
-    def of_azure_openai(cls, api_key: str, endpoint: str, api_version: str, model_name: str) -> "ResponsesUDFBuilder":
-        """Creates a builder configured for Azure OpenAI.
-
-        Args:
-            api_key (str): The Azure OpenAI API key.
-            endpoint (str): The Azure OpenAI endpoint URL.
-            api_version (str): The Azure OpenAI API version (e.g., "2024-02-01").
-            model_name (str): The Azure OpenAI deployment name for responses.
-
-        Returns:
-            ResponsesUDFBuilder: A builder instance configured for Azure OpenAI responses.
-        """
-        return cls(api_key=api_key, endpoint=endpoint, api_version=api_version, model_name=model_name)
-
-    def build(
-        self,
-        instructions: str,
-        response_format: Type[T] = str,
-        batch_size: int = 128,  # Default batch size for async might differ
-        temperature: float = 0.0,
-        top_p: float = 1.0,
-        max_concurrency: int = 8,
-    ) -> UserDefinedFunction:
-        """Builds the asynchronous pandas UDF for generating responses.
-
-        Args:
-            instructions (str): The system prompt or instructions for the model.
-            response_format (Type[T]): The desired output format. Either `str` for plain text
-                or a Pydantic `BaseModel` for structured JSON output. Defaults to `str`.
-            batch_size (int): Number of rows per async batch request passed to the underlying
-                `pandas_ext` function. Defaults to 128.
-            temperature (float): Sampling temperature (0.0 to 2.0). Defaults to 0.0.
-            top_p (float): Nucleus sampling parameter. Defaults to 1.0.
-
-        Returns:
-            UserDefinedFunction: A Spark pandas UDF configured to generate responses asynchronously.
-                Output schema is `StringType` or a struct derived from `response_format`.
-
-        Raises:
-            ValueError: If `response_format` is not `str` or a Pydantic `BaseModel`.
-        """
-        if issubclass(response_format, BaseModel):
-            spark_schema = _pydantic_to_spark_schema(response_format)
-            json_schema_string = serialize_base_model(response_format)
-
-            @pandas_udf(returnType=spark_schema)
-            def structure_udf(col: Iterator[pd.Series]) -> Iterator[pd.DataFrame]:
-                _initialize(self.api_key, self.endpoint, self.api_version)
-                pandas_ext.responses_model(self.model_name)
+    if issubclass(response_format, BaseModel):
+        spark_schema = _pydantic_to_spark_schema(response_format)
+        json_schema_string = serialize_base_model(response_format)
+
+        @pandas_udf(returnType=spark_schema)  # type: ignore[call-overload]
+        def structure_udf(col: Iterator[pd.Series]) -> Iterator[pd.DataFrame]:
+            pandas_ext.set_responses_model(_model_name)
+            response_format = deserialize_base_model(json_schema_string)
+            cache = AsyncBatchingMapProxy[str, response_format](
+                batch_size=batch_size,
+                max_concurrency=max_concurrency,
+            )
 
+            try:
                 for part in col:
                     predictions: pd.Series = asyncio.run(
-                        part.aio.responses(
+                        part.aio.responses_with_cache(
                             instructions=instructions,
-                            response_format=deserialize_base_model(json_schema_string),
-                            batch_size=batch_size,
-                            temperature=temperature,
-                            top_p=top_p,
-                            max_concurrency=max_concurrency,
+                            response_format=response_format,
+                            cache=cache,
+                            **api_kwargs,
                         )
                     )
                     yield pd.DataFrame(predictions.map(_safe_dump).tolist())
+            finally:
+                asyncio.run(cache.clear())
 
-            return structure_udf
+        return structure_udf  # type: ignore[return-value]
 
-        elif issubclass(response_format, str):
+    elif issubclass(response_format, str):
 
-            @pandas_udf(returnType=StringType())
-            def string_udf(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
-                _initialize(self.api_key, self.endpoint, self.api_version)
-                pandas_ext.responses_model(self.model_name)
+        @pandas_udf(returnType=StringType())  # type: ignore[call-overload]
+        def string_udf(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
+            pandas_ext.set_responses_model(_model_name)
+            cache = AsyncBatchingMapProxy[str, str](
+                batch_size=batch_size,
+                max_concurrency=max_concurrency,
+            )
 
+            try:
                 for part in col:
                     predictions: pd.Series = asyncio.run(
-                        part.aio.responses(
+                        part.aio.responses_with_cache(
                             instructions=instructions,
                             response_format=str,
-                            batch_size=batch_size,
-                            temperature=temperature,
-                            top_p=top_p,
-                            max_concurrency=max_concurrency,
+                            cache=cache,
+                            **api_kwargs,
                         )
                     )
                     yield predictions.map(_safe_cast_str)
+            finally:
+                asyncio.run(cache.clear())
 
-            return string_udf
+        return string_udf  # type: ignore[return-value]
 
-        else:
-            raise ValueError(f"Unsupported response_format: {response_format}")
-
-    def build_from_task(
-        self,
-        task: PreparedTask,
-        batch_size: int = 128,
-        max_concurrency: int = 8,
-    ) -> UserDefinedFunction:
-        """Builds the asynchronous pandas UDF from a predefined task.
-
-        This method allows users to create UDFs from predefined tasks such as sentiment analysis,
-        translation, or other common NLP operations defined in the openaivec.task module.
-
-        Args:
-            task (PreparedTask): A predefined task configuration containing instructions,
-                response format, temperature, and top_p settings.
-            batch_size (int): Number of rows per async batch request passed to the underlying
-                `pandas_ext` function. Defaults to 128.
-            max_concurrency (int): Maximum number of concurrent requests. Defaults to 8.
-
-        Returns:
-            UserDefinedFunction: A Spark pandas UDF configured to execute the specified task
-                asynchronously, returning a struct derived from the task's response format.
-
-        Example:
-            ```python
-            from openaivec.task import nlp
-            
-            builder = ResponsesUDFBuilder.of_openai(
-                api_key="your-api-key",
-                model_name="gpt-4o-mini"
-            )
-            
-            sentiment_udf = builder.build_from_task(nlp.SENTIMENT_ANALYSIS)
-            
-            spark.udf.register("analyze_sentiment", sentiment_udf)
-            ```
-        """
-        # Serialize task parameters for Spark serialization compatibility
-        task_instructions = task.instructions
-        task_response_format_json = serialize_base_model(task.response_format)
-        task_temperature = task.temperature
-        task_top_p = task.top_p
-
-        # Deserialize the response format from JSON
-        response_format = deserialize_base_model(task_response_format_json)
-        spark_schema = _pydantic_to_spark_schema(response_format)
+    else:
+        raise ValueError(f"Unsupported response_format: {response_format}")
 
-        @pandas_udf(returnType=spark_schema)
-        def task_udf(col: Iterator[pd.Series]) -> Iterator[pd.DataFrame]:
-            _initialize(self.api_key, self.endpoint, self.api_version)
-            pandas_ext.responses_model(self.model_name)
 
-            for part in col:
-                predictions: pd.Series = asyncio.run(
-                    part.aio.responses(
-                        instructions=task_instructions,
-                        response_format=response_format,
-                        batch_size=batch_size,
-                        temperature=task_temperature,
-                        top_p=task_top_p,
-                        max_concurrency=max_concurrency,
-                    )
-                )
-                yield pd.DataFrame(predictions.map(_safe_dump).tolist())
+def task_udf(
+    task: PreparedTask[ResponseFormat],
+    model_name: str | None = None,
+    batch_size: int | None = None,
+    max_concurrency: int = 8,
+    **api_kwargs,
+) -> UserDefinedFunction:
+    """Create an asynchronous Spark pandas UDF from a predefined task.
+
+    This function allows users to create UDFs from predefined tasks such as sentiment analysis,
+    translation, or other common NLP operations defined in the openaivec.task module.
+    Each partition maintains its own cache to eliminate duplicate API calls within
+    the partition, significantly reducing API usage and costs when processing
+    datasets with overlapping content.
 
-        return task_udf
+    Args:
+        task (PreparedTask): A predefined task configuration containing instructions
+            and response format.
+        model_name (str | None): For Azure OpenAI, use your deployment name (e.g., "my-gpt4-deployment").
+            For OpenAI, use the model name (e.g., "gpt-4.1-mini"). Defaults to configured model in DI container
+            via ResponsesModelName if not provided.
+        batch_size (int | None): Number of rows per async batch request within each partition.
+            Larger values reduce API call overhead but increase memory usage.
+            Defaults to None (automatic batch size optimization that dynamically
+            adjusts based on execution time, targeting 30-60 seconds per batch).
+            Set to a positive integer (e.g., 32-128) for fixed batch size.
+        max_concurrency (int): Maximum number of concurrent API requests **PER EXECUTOR**.
+            Total cluster concurrency = max_concurrency × number_of_executors.
+            Higher values increase throughput but may hit OpenAI rate limits.
+            Recommended: 4-12 per executor. Defaults to 8.
+
+    Additional Keyword Args:
+        Arbitrary OpenAI Responses API parameters (e.g. ``temperature``, ``top_p``,
+        ``frequency_penalty``, ``presence_penalty``, ``seed``, ``max_output_tokens``, etc.)
+        are forwarded verbatim to the underlying API calls. These parameters are applied to
+        all API requests made by the UDF.
 
+    Returns:
+        UserDefinedFunction: A Spark pandas UDF configured to execute the specified task
+            asynchronously with automatic caching for duplicate inputs within each partition.
+            Output schema is StringType for str response format or a struct derived from
+            the task's response format for BaseModel.
 
-@dataclass(frozen=True)
-class EmbeddingsUDFBuilder:
-    """Builder for asynchronous Spark pandas UDFs for creating embeddings.
+    Example:
+        ```python
+        from openaivec.task import nlp
 
-    Configures and builds UDFs that leverage `pandas_ext.aio.embeddings`
-    to generate vector embeddings from OpenAI models asynchronously.
-    An instance stores authentication parameters and the model name.
+        sentiment_udf = task_udf(nlp.SENTIMENT_ANALYSIS)
 
-    Attributes:
-        api_key (str): OpenAI or Azure API key.
-        endpoint (Optional[str]): Azure endpoint base URL. None for public OpenAI.
-        api_version (Optional[str]): Azure API version. Ignored for public OpenAI.
-        model_name (str): Deployment name (Azure) or model name (OpenAI) for embeddings.
+        spark.udf.register("analyze_sentiment", sentiment_udf)
+        ```
+
+    Note:
+        **Automatic Caching**: Duplicate inputs within each partition are cached,
+        reducing API calls and costs significantly on datasets with repeated content.
     """
+    return responses_udf(
+        instructions=task.instructions,
+        response_format=task.response_format,
+        model_name=model_name,
+        batch_size=batch_size,
+        max_concurrency=max_concurrency,
+        **api_kwargs,
+    )
+
+
+def infer_schema(
+    instructions: str,
+    example_table_name: str,
+    example_field_name: str,
+    max_examples: int = 100,
+) -> SchemaInferenceOutput:
+    """Infer the schema for a response format based on example data.
+
+    This function retrieves examples from a Spark table and infers the schema
+    for the response format using the provided instructions. It is useful when
+    you want to dynamically generate a schema based on existing data.
 
-    # Params for OpenAI SDK
-    api_key: str
-    endpoint: str | None
-    api_version: str | None
+    Args:
+        instructions (str): Instructions for the model to infer the schema.
+        example_table_name (str | None): Name of the Spark table containing example data.
+        example_field_name (str | None): Name of the field in the table to use as examples.
+        max_examples (int): Maximum number of examples to retrieve for schema inference.
 
-    # Params for Embeddings API
-    model_name: str
+    Returns:
+        InferredSchema: An object containing the inferred schema and response format.
+
+    Example:
+        ```python
+        from pyspark.sql import SparkSession
+
+        spark = SparkSession.builder.getOrCreate()
+        spark.createDataFrame([("great product",), ("bad service",)], ["text"]).createOrReplaceTempView("examples")
+        infer_schema(
+            instructions="Classify sentiment as positive or negative.",
+            example_table_name="examples",
+            example_field_name="text",
+            max_examples=2,
+        )
+        ```
+    """
 
-    @classmethod
-    def of_openai(cls, api_key: str, model_name: str) -> "EmbeddingsUDFBuilder":
-        """Creates a builder configured for the public OpenAI API.
+    spark = CONTAINER.resolve(SparkSession)
+    examples: list[str] = (
+        spark.table(example_table_name).rdd.map(lambda row: row[example_field_name]).takeSample(False, max_examples)
+    )
+
+    input = SchemaInferenceInput(
+        instructions=instructions,
+        examples=examples,
+    )
+    inferer = CONTAINER.resolve(SchemaInferer)
+    return inferer.infer_schema(input)
+
+
+def parse_udf(
+    instructions: str,
+    response_format: type[ResponseFormat] | None = None,
+    example_table_name: str | None = None,
+    example_field_name: str | None = None,
+    max_examples: int = 100,
+    model_name: str | None = None,
+    batch_size: int | None = None,
+    max_concurrency: int = 8,
+    **api_kwargs,
+) -> UserDefinedFunction:
+    """Create an asynchronous Spark pandas UDF for parsing responses.
+    This function allows users to create UDFs that parse responses based on
+    provided instructions and either a predefined response format or example data.
+    It supports both structured responses using Pydantic models and plain text responses.
+    Each partition maintains its own cache to eliminate duplicate API calls within
+    the partition, significantly reducing API usage and costs when processing
+    datasets with overlapping content.
 
-        Args:
-            api_key (str): The OpenAI API key.
-            model_name (str): The OpenAI model name for embeddings (e.g., "text-embedding-3-small").
+    Args:
+        instructions (str): The system prompt or instructions for the model.
+        response_format (type[ResponseFormat] | None): The desired output format.
+            Either `str` for plain text or a Pydantic `BaseModel` for structured JSON output.
+            If not provided, the schema will be inferred from example data.
+        example_table_name (str | None): Name of the Spark table containing example data.
+            If provided, `example_field_name` must also be specified.
+        example_field_name (str | None): Name of the field in the table to use as examples.
+            If provided, `example_table_name` must also be specified.
+        max_examples (int): Maximum number of examples to retrieve for schema inference.
+            Defaults to 100.
+        model_name (str | None): For Azure OpenAI, use your deployment name (e.g., "my-gpt4-deployment").
+            For OpenAI, use the model name (e.g., "gpt-4.1-mini"). Defaults to configured model in DI container
+            via ResponsesModelName if not provided.
+        batch_size (int | None): Number of rows per async batch request within each partition.
+            Larger values reduce API call overhead but increase memory usage.
+            Defaults to None (automatic batch size optimization that dynamically
+            adjusts based on execution time, targeting 30-60 seconds per batch).
+            Set to a positive integer (e.g., 32-128) for fixed batch size
+        max_concurrency (int): Maximum number of concurrent API requests **PER EXECUTOR**.
+            Total cluster concurrency = max_concurrency × number_of_executors.
+            Higher values increase throughput but may hit OpenAI rate limits.
+            Recommended: 4-12 per executor. Defaults to 8.
+        **api_kwargs: Additional OpenAI API parameters (e.g. ``temperature``, ``top_p``,
+            ``frequency_penalty``, ``presence_penalty``, ``seed``, ``max_output_tokens``, etc.)
+            forwarded verbatim to the underlying API calls. These parameters are applied to
+            all API requests made by the UDF and override any parameters set in the
+            response_format or example data.
+    Example:
+        ```python
+        from pyspark.sql import SparkSession
+
+        spark = SparkSession.builder.getOrCreate()
+        spark.createDataFrame(
+            [("Order #123 delivered",), ("Order #456 delayed",)],
+            ["body"],
+        ).createOrReplaceTempView("messages")
+        udf = parse_udf(
+            instructions="Extract order id as `order_id` and status as `status`.",
+            example_table_name="messages",
+            example_field_name="body",
+        )
+        spark.udf.register("parse_ticket", udf)
+        spark.sql("SELECT parse_ticket(body) AS parsed FROM messages").show()
+        ```
+    Returns:
+        UserDefinedFunction: A Spark pandas UDF configured to parse responses asynchronously.
+            Output schema is `StringType` for str response format or a struct derived from
+            the response_format for BaseModel.
+    Raises:
+        ValueError: If neither `response_format` nor `example_table_name` and `example_field_name` are provided.
+    """
 
-        Returns:
-            EmbeddingsUDFBuilder: A builder instance configured for OpenAI embeddings.
-        """
-        return cls(api_key=api_key, endpoint=None, api_version=None, model_name=model_name)
+    if not response_format and not (example_field_name and example_table_name):
+        raise ValueError("Either response_format or example_table_name and example_field_name must be provided.")
+
+    schema: SchemaInferenceOutput | None = None
+
+    if not response_format:
+        schema = infer_schema(
+            instructions=instructions,
+            example_table_name=example_table_name,
+            example_field_name=example_field_name,
+            max_examples=max_examples,
+        )
+
+    return responses_udf(
+        instructions=schema.inference_prompt if schema else instructions,
+        response_format=schema.model if schema else response_format,
+        model_name=model_name,
+        batch_size=batch_size,
+        max_concurrency=max_concurrency,
+        **api_kwargs,
+    )
+
+
+def embeddings_udf(
+    model_name: str | None = None,
+    batch_size: int | None = None,
+    max_concurrency: int = 8,
+    **api_kwargs,
+) -> UserDefinedFunction:
+    """Create an asynchronous Spark pandas UDF for generating embeddings.
+
+    Configures and builds UDFs that leverage `pandas_ext.aio.embeddings_with_cache`
+    to generate vector embeddings from OpenAI models asynchronously.
+    Each partition maintains its own cache to eliminate duplicate API calls within
+    the partition, significantly reducing API usage and costs when processing
+    datasets with overlapping content.
 
-    @classmethod
-    def of_azure_openai(cls, api_key: str, endpoint: str, api_version: str, model_name: str) -> "EmbeddingsUDFBuilder":
-        """Creates a builder configured for Azure OpenAI.
+    Note:
+        Authentication must be configured via SparkContext environment variables.
+        Set the appropriate environment variables on the SparkContext:
 
-        Args:
-            api_key (str): The Azure OpenAI API key.
-            endpoint (str): The Azure OpenAI endpoint URL.
-            api_version (str): The Azure OpenAI API version (e.g., "2024-02-01").
-            model_name (str): The Azure OpenAI deployment name for embeddings.
+        For OpenAI:
+            sc.environment["OPENAI_API_KEY"] = "your-openai-api-key"
 
-        Returns:
-            EmbeddingsUDFBuilder: A builder instance configured for Azure OpenAI embeddings.
-        """
-        return cls(api_key=api_key, endpoint=endpoint, api_version=api_version, model_name=model_name)
+        For Azure OpenAI:
+            sc.environment["AZURE_OPENAI_API_KEY"] = "your-azure-openai-api-key"
+            sc.environment["AZURE_OPENAI_BASE_URL"] = "https://YOUR-RESOURCE-NAME.services.ai.azure.com/openai/v1/"
+            sc.environment["AZURE_OPENAI_API_VERSION"] = "preview"
 
-    def build(self, batch_size: int = 128, max_concurrency: int = 8) -> UserDefinedFunction:
-        """Builds the asynchronous pandas UDF for generating embeddings.
+    Args:
+        model_name (str | None): For Azure OpenAI, use your deployment name (e.g., "my-embedding-deployment").
+            For OpenAI, use the model name (e.g., "text-embedding-3-small").
+            Defaults to configured model in DI container via EmbeddingsModelName if not provided.
+        batch_size (int | None): Number of rows per async batch request within each partition.
+            Larger values reduce API call overhead but increase memory usage.
+            Defaults to None (automatic batch size optimization that dynamically
+            adjusts based on execution time, targeting 30-60 seconds per batch).
+            Set to a positive integer (e.g., 64-256) for fixed batch size.
+            Embeddings typically handle larger batches efficiently.
+        max_concurrency (int): Maximum number of concurrent API requests **PER EXECUTOR**.
+            Total cluster concurrency = max_concurrency × number_of_executors.
+            Higher values increase throughput but may hit OpenAI rate limits.
+            Recommended: 4-12 per executor. Defaults to 8.
+        **api_kwargs: Additional OpenAI API parameters (e.g., dimensions for text-embedding-3 models).
 
-        Args:
-            batch_size (int): Number of rows per async batch request passed to the underlying
-                `pandas_ext` function. Defaults to 128.
+    Returns:
+        UserDefinedFunction: A Spark pandas UDF configured to generate embeddings asynchronously
+            with automatic caching for duplicate inputs within each partition,
+            returning an `ArrayType(FloatType())` column.
+
+    Note:
+        For optimal performance in distributed environments:
+        - **Automatic Caching**: Duplicate inputs within each partition are cached,
+          reducing API calls and costs significantly on datasets with repeated content
+        - Monitor OpenAI API rate limits when scaling executor count
+        - Consider your OpenAI tier limits: total_requests = max_concurrency × executors
+        - Embeddings API typically has higher throughput than chat completions
+        - Use larger batch_size for embeddings compared to response generation
+    """
 
-        Returns:
-            UserDefinedFunction: A Spark pandas UDF configured to generate embeddings asynchronously,
-                returning an `ArrayType(FloatType())` column.
-        """
+    _model_name = model_name or CONTAINER.resolve(EmbeddingsModelName).value
 
-        @pandas_udf(returnType=ArrayType(FloatType()))
-        def embeddings_udf(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
-            _initialize(self.api_key, self.endpoint, self.api_version)
-            pandas_ext.embeddings_model(self.model_name)
+    @pandas_udf(returnType=ArrayType(FloatType()))  # type: ignore[call-overload,misc]
+    def _embeddings_udf(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
+        pandas_ext.set_embeddings_model(_model_name)
+        cache = AsyncBatchingMapProxy[str, np.ndarray](
+            batch_size=batch_size,
+            max_concurrency=max_concurrency,
+        )
 
+        try:
             for part in col:
-                embeddings: pd.Series = asyncio.run(
-                    part.aio.embeddings(batch_size=batch_size, max_concurrency=max_concurrency)
-                )
+                embeddings: pd.Series = asyncio.run(part.aio.embeddings_with_cache(cache=cache, **api_kwargs))
                 yield embeddings.map(lambda x: x.tolist())
+        finally:
+            asyncio.run(cache.clear())
 
-        return embeddings_udf
+    return _embeddings_udf  # type: ignore[return-value]
 
 
-
-def split_to_chunks_udf(model_name: str, max_tokens: int, sep: List[str]) -> UserDefinedFunction:
+def split_to_chunks_udf(max_tokens: int, sep: list[str]) -> UserDefinedFunction:
     """Create a pandas‑UDF that splits text into token‑bounded chunks.
 
     Args:
-        model_name: Model identifier passed to *tiktoken*.
-        max_tokens: Maximum tokens allowed per chunk.
-        sep: Ordered list of separator strings used by ``TextChunker``.
+        max_tokens (int): Maximum tokens allowed per chunk.
+        sep (list[str]): Ordered list of separator strings used by ``TextChunker``.
 
     Returns:
         A pandas UDF producing an ``ArrayType(StringType())`` column whose
             values are lists of chunks respecting the ``max_tokens`` limit.
     """
 
-    @pandas_udf(ArrayType(StringType()))
+    @pandas_udf(ArrayType(StringType()))  # type: ignore[call-overload,misc]
     def fn(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
-        global _TIKTOKEN_ENC
-        if _TIKTOKEN_ENC is None:
-            _TIKTOKEN_ENC = tiktoken.encoding_for_model(model_name)
-
-        chunker = TextChunker(_TIKTOKEN_ENC)
+        encoding = tiktoken.get_encoding("o200k_base")
+        chunker = TextChunker(encoding)
 
         for part in col:
             yield part.map(lambda x: chunker.split(x, max_tokens=max_tokens, sep=sep) if isinstance(x, str) else [])
 
-    return fn
+    return fn  # type: ignore[return-value]
 
 
-def count_tokens_udf(model_name: str = "gpt-4o") -> UserDefinedFunction:
+def count_tokens_udf() -> UserDefinedFunction:
     """Create a pandas‑UDF that counts tokens for every string cell.
 
     The UDF uses *tiktoken* to approximate tokenisation and caches the
     resulting ``Encoding`` object per executor.
 
-    Args:
-        model_name: Model identifier understood by ``tiktoken``.
-
     Returns:
         A pandas UDF producing an ``IntegerType`` column with token counts.
     """
 
-    @pandas_udf(IntegerType())
+    @pandas_udf(IntegerType())  # type: ignore[call-overload]
     def fn(col: Iterator[pd.Series]) -> Iterator[pd.Series]:
-        global _TIKTOKEN_ENC
-        if _TIKTOKEN_ENC is None:
-            _TIKTOKEN_ENC = tiktoken.encoding_for_model(model_name)
+        encoding = tiktoken.get_encoding("o200k_base")
 
         for part in col:
-            yield part.map(lambda x: len(_TIKTOKEN_ENC.encode(x)) if isinstance(x, str) else 0)
+            yield part.map(lambda x: len(encoding.encode(x)) if isinstance(x, str) else 0)
 
-    return fn
+    return fn  # type: ignore[return-value]
 
 
 def similarity_udf() -> UserDefinedFunction:
-    @pandas_udf(FloatType())
+    """Create a pandas-UDF that computes cosine similarity between embedding vectors.
+
+    Returns:
+        UserDefinedFunction: A Spark pandas UDF that takes two embedding vector columns
+            and returns their cosine similarity as a FloatType column.
+    """
+
+    @pandas_udf(FloatType())  # type: ignore[call-overload]
     def fn(a: pd.Series, b: pd.Series) -> pd.Series:
-        """Compute cosine similarity between two vectors.
+        # Import pandas_ext to ensure .ai accessor is available in Spark workers
+        from openaivec import pandas_ext
 
-        Args:
-            a: First vector.
-            b: Second vector.
+        # Explicitly reference pandas_ext to satisfy linters
+        assert pandas_ext is not None
 
-        Returns:
-            Cosine similarity between the two vectors.
-        """
-        pandas_ext._wakeup()
         return pd.DataFrame({"a": a, "b": b}).ai.similarity("a", "b")
 
-    return fn
+    return fn  # type: ignore[return-value]