openaivec 0.13.4__py3-none-any.whl → 0.13.6__py3-none-any.whl

openaivec/proxy.py

@@ -2,7 +2,7 @@ import asyncio
 import threading
 from collections.abc import Hashable
 from dataclasses import dataclass, field
-from typing import Awaitable, Callable, Dict, Generic, List, Optional, TypeVar
+from typing import Any, Awaitable, Callable, Dict, Generic, List, TypeVar
 
 from openaivec.optimize import BatchSizeSuggester
 
@@ -22,9 +22,9 @@ class ProxyBase(Generic[S, T]):
             should process the entire input in a single call.
     """
 
-    batch_size: Optional[int] = None  # subclasses may override via dataclass
-    show_progress: bool = False  # Enable progress bar display
-    suggester: BatchSizeSuggester = None  # Batch size optimization, initialized by subclasses
+    batch_size: int | None  # subclasses may override via dataclass
+    show_progress: bool  # Enable progress bar display
+    suggester: BatchSizeSuggester  # Batch size optimization, initialized by subclasses
 
     def _is_notebook_environment(self) -> bool:
         """Check if running in a Jupyter notebook environment.
@@ -33,7 +33,7 @@ class ProxyBase(Generic[S, T]):
             bool: True if running in a notebook, False otherwise.
         """
         try:
-            from IPython import get_ipython
+            from IPython.core.getipython import get_ipython
 
             ipython = get_ipython()
             if ipython is not None:
@@ -89,7 +89,7 @@ class ProxyBase(Generic[S, T]):
 
         return False
 
-    def _create_progress_bar(self, total: int, desc: str = "Processing batches") -> Optional[object]:
+    def _create_progress_bar(self, total: int, desc: str = "Processing batches") -> Any:
         """Create a progress bar if conditions are met.
 
         Args:
@@ -97,7 +97,7 @@ class ProxyBase(Generic[S, T]):
             desc (str): Description for the progress bar.
 
         Returns:
-            Optional[object]: Progress bar instance or None if not available.
+            Any: Progress bar instance or None if not available.
         """
         try:
             from tqdm.auto import tqdm as tqdm_progress
@@ -108,21 +108,21 @@ class ProxyBase(Generic[S, T]):
             pass
         return None
 
-    def _update_progress_bar(self, progress_bar: Optional[object], increment: int) -> None:
+    def _update_progress_bar(self, progress_bar: Any, increment: int) -> None:
         """Update progress bar with the given increment.
 
         Args:
-            progress_bar (Optional[object]): Progress bar instance.
+            progress_bar (Any): Progress bar instance.
             increment (int): Number of items to increment.
         """
         if progress_bar:
             progress_bar.update(increment)
 
-    def _close_progress_bar(self, progress_bar: Optional[object]) -> None:
+    def _close_progress_bar(self, progress_bar: Any) -> None:
         """Close the progress bar.
 
         Args:
-            progress_bar (Optional[object]): Progress bar instance.
+            progress_bar (Any): Progress bar instance.
         """
         if progress_bar:
             progress_bar.close()
@@ -179,6 +179,10 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
     not duplicate work via an in-flight registry. All public behavior is preserved
     while minimizing redundant requests and maintaining input order in the output.
 
+    When ``batch_size=None``, automatic batch size optimization is enabled,
+    dynamically adjusting batch sizes based on execution time to maintain optimal
+    performance (targeting 30-60 seconds per batch).
+
     Example:
         >>> from typing import List
         >>> p = BatchingMapProxy[int, str](batch_size=3)
@@ -188,15 +192,19 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         ['v:1', 'v:2', 'v:2', 'v:3', 'v:4']
     """
 
-    # Number of items to process per call to map_func. If None or <= 0, process all at once.
-    batch_size: Optional[int] = None
+    # Number of items to process per call to map_func.
+    # - If None (default): Enables automatic batch size optimization, dynamically adjusting
+    #   based on execution time (targeting 30-60 seconds per batch)
+    # - If positive integer: Fixed batch size
+    # - If <= 0: Process all items at once
+    batch_size: int | None = None
     show_progress: bool = False
     suggester: BatchSizeSuggester = field(default_factory=BatchSizeSuggester, repr=False)
 
     # internals
-    __cache: Dict[S, T] = field(default_factory=dict)
-    __lock: threading.RLock = field(default_factory=threading.RLock, repr=False)
-    __inflight: Dict[S, threading.Event] = field(default_factory=dict, repr=False)
+    _cache: Dict[S, T] = field(default_factory=dict)
+    _lock: threading.RLock = field(default_factory=threading.RLock, repr=False)
+    _inflight: Dict[S, threading.Event] = field(default_factory=dict, repr=False)
 
     def __all_cached(self, items: List[S]) -> bool:
         """Check whether all items are present in the cache.
@@ -209,8 +217,8 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         Returns:
             bool: True if every item is already cached, False otherwise.
         """
-        with self.__lock:
-            return all(x in self.__cache for x in items)
+        with self._lock:
+            return all(x in self._cache for x in items)
 
     def __values(self, items: List[S]) -> List[T]:
         """Fetch cached values for ``items`` preserving the given order.
@@ -224,8 +232,8 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             list[T]: The cached values corresponding to ``items`` in the same
             order.
         """
-        with self.__lock:
-            return [self.__cache[x] for x in items]
+        with self._lock:
+            return [self._cache[x] for x in items]
 
     def __acquire_ownership(self, items: List[S]) -> tuple[List[S], List[S]]:
         """Acquire ownership for missing items and identify keys to wait for.
@@ -245,14 +253,14 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         """
         owned: List[S] = []
         wait_for: List[S] = []
-        with self.__lock:
+        with self._lock:
             for x in items:
-                if x in self.__cache:
+                if x in self._cache:
                     continue
-                if x in self.__inflight:
+                if x in self._inflight:
                     wait_for.append(x)
                 else:
-                    self.__inflight[x] = threading.Event()
+                    self._inflight[x] = threading.Event()
                     owned.append(x)
         return owned, wait_for
 
@@ -268,10 +276,10 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             # Release waiters and surface a clear error.
             self.__finalize_failure(to_call)
             raise ValueError("map_func must return a list of results with the same length and order as inputs")
-        with self.__lock:
+        with self._lock:
             for x, y in zip(to_call, results):
-                self.__cache[x] = y
-                ev = self.__inflight.pop(x, None)
+                self._cache[x] = y
+                ev = self._inflight.pop(x, None)
                 if ev:
                     ev.set()
 
@@ -282,9 +290,9 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             to_call (list[S]): Items that were intended to be computed when an
             error occurred.
         """
-        with self.__lock:
+        with self._lock:
             for x in to_call:
-                ev = self.__inflight.pop(x, None)
+                ev = self._inflight.pop(x, None)
                 if ev:
                     ev.set()
 
@@ -296,11 +304,11 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             - Do not call concurrently with active map() calls to avoid
               unnecessary recomputation or racy wake-ups.
         """
-        with self.__lock:
-            for ev in self.__inflight.values():
+        with self._lock:
+            for ev in self._inflight.values():
                 ev.set()
-            self.__inflight.clear()
-            self.__cache.clear()
+            self._inflight.clear()
+            self._cache.clear()
 
     def close(self) -> None:
         """Alias for clear()."""
@@ -337,8 +345,8 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             current_batch_size = self._normalized_batch_size(len(owned))
             batch = owned[i : i + current_batch_size]
             # Double-check cache right before processing
-            with self.__lock:
-                uncached_in_batch = [x for x in batch if x not in self.__cache]
+            with self._lock:
+                uncached_in_batch = [x for x in batch if x not in self._cache]
 
             pending_to_call.extend(uncached_in_batch)
 
@@ -400,13 +408,13 @@ class BatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         rescued: List[S] = []  # keys we claim to batch-process
         for x in keys:
             while True:
-                with self.__lock:
-                    if x in self.__cache:
+                with self._lock:
+                    if x in self._cache:
                         break
-                    ev = self.__inflight.get(x)
+                    ev = self._inflight.get(x)
                     if ev is None:
                         # Not cached and no one computing; claim ownership to batch later.
-                        self.__inflight[x] = threading.Event()
+                        self._inflight[x] = threading.Event()
                         rescued.append(x)
                         break
                 # Someone else is computing; wait for completion.
@@ -463,6 +471,10 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
     coordinates concurrent coroutines to avoid duplicate work via an in-flight
     registry of asyncio events.
 
+    When ``batch_size=None``, automatic batch size optimization is enabled,
+    dynamically adjusting batch sizes based on execution time to maintain optimal
+    performance (targeting 30-60 seconds per batch).
+
     Example:
         >>> import asyncio
         >>> from typing import List
@@ -476,16 +488,21 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         ['v:1', 'v:2', 'v:3']
     """
 
-    batch_size: Optional[int] = None
+    # Number of items to process per call to map_func.
+    # - If None (default): Enables automatic batch size optimization, dynamically adjusting
+    #   based on execution time (targeting 30-60 seconds per batch)
+    # - If positive integer: Fixed batch size
+    # - If <= 0: Process all items at once
+    batch_size: int | None = None
     max_concurrency: int = 8
     show_progress: bool = False
     suggester: BatchSizeSuggester = field(default_factory=BatchSizeSuggester, repr=False)
 
     # internals
-    __cache: Dict[S, T] = field(default_factory=dict, repr=False)
-    __lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
-    __inflight: Dict[S, asyncio.Event] = field(default_factory=dict, repr=False)
-    __sema: Optional[asyncio.Semaphore] = field(default=None, init=False, repr=False)
+    _cache: Dict[S, T] = field(default_factory=dict, repr=False)
+    _lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
+    _inflight: Dict[S, asyncio.Event] = field(default_factory=dict, repr=False)
+    __sema: asyncio.Semaphore | None = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
         """Initialize internal semaphore based on ``max_concurrency``.
@@ -517,8 +534,8 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         Returns:
             bool: True if every item in ``items`` is already cached, False otherwise.
         """
-        async with self.__lock:
-            return all(x in self.__cache for x in items)
+        async with self._lock:
+            return all(x in self._cache for x in items)
 
     async def __values(self, items: List[S]) -> List[T]:
         """Get cached values for ``items`` preserving their given order.
@@ -532,8 +549,8 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         Returns:
             list[T]: Cached values corresponding to ``items`` in the same order.
         """
-        async with self.__lock:
-            return [self.__cache[x] for x in items]
+        async with self._lock:
+            return [self._cache[x] for x in items]
 
     async def __acquire_ownership(self, items: List[S]) -> tuple[List[S], List[S]]:
         """Acquire ownership for missing keys and identify keys to wait for.
@@ -548,14 +565,14 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         """
         owned: List[S] = []
         wait_for: List[S] = []
-        async with self.__lock:
+        async with self._lock:
             for x in items:
-                if x in self.__cache:
+                if x in self._cache:
                     continue
-                if x in self.__inflight:
+                if x in self._inflight:
                     wait_for.append(x)
                 else:
-                    self.__inflight[x] = asyncio.Event()
+                    self._inflight[x] = asyncio.Event()
                     owned.append(x)
         return owned, wait_for
 
@@ -570,10 +587,10 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             # Prevent deadlocks if map_func violates the contract.
             await self.__finalize_failure(to_call)
             raise ValueError("map_func must return a list of results with the same length and order as inputs")
-        async with self.__lock:
+        async with self._lock:
             for x, y in zip(to_call, results):
-                self.__cache[x] = y
-                ev = self.__inflight.pop(x, None)
+                self._cache[x] = y
+                ev = self._inflight.pop(x, None)
                 if ev:
                     ev.set()
 
@@ -584,9 +601,9 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             to_call (list[S]): Items whose computation failed; their waiters will
             be released.
         """
-        async with self.__lock:
+        async with self._lock:
             for x in to_call:
-                ev = self.__inflight.pop(x, None)
+                ev = self._inflight.pop(x, None)
                 if ev:
                     ev.set()
 
@@ -598,11 +615,11 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
             - Do not call concurrently with active map() calls to avoid
               unnecessary recomputation or racy wake-ups.
         """
-        async with self.__lock:
-            for ev in self.__inflight.values():
+        async with self._lock:
+            for ev in self._inflight.values():
                 ev.set()
-            self.__inflight.clear()
-            self.__cache.clear()
+            self._inflight.clear()
+            self._cache.clear()
 
     async def aclose(self) -> None:
         """Alias for clear()."""
@@ -686,13 +703,13 @@ class AsyncBatchingMapProxy(ProxyBase[S, T], Generic[S, T]):
         rescued: List[S] = []  # keys we claim to batch-process
         for x in keys:
             while True:
-                async with self.__lock:
-                    if x in self.__cache:
+                async with self._lock:
+                    if x in self._cache:
                         break
-                    ev = self.__inflight.get(x)
+                    ev = self._inflight.get(x)
                     if ev is None:
                         # Not cached and no one computing; claim ownership to batch later.
-                        self.__inflight[x] = asyncio.Event()
+                        self._inflight[x] = asyncio.Event()
                         rescued.append(x)
                         break
                 # Someone else is computing; wait for completion.

openaivec/responses.py

@@ -165,8 +165,8 @@ class BatchResponses(Generic[ResponseFormat]):
     system_message: str
     temperature: float | None = 0.0
     top_p: float = 1.0
-    response_format: Type[ResponseFormat] = str
-    cache: BatchingMapProxy[str, ResponseFormat] = field(default_factory=lambda: BatchingMapProxy(batch_size=128))
+    response_format: Type[ResponseFormat] = str  # type: ignore[assignment]
+    cache: BatchingMapProxy[str, ResponseFormat] = field(default_factory=lambda: BatchingMapProxy(batch_size=None))
     _vectorized_system_message: str = field(init=False)
     _model_json_schema: dict = field(init=False)
 
@@ -179,7 +179,7 @@ class BatchResponses(Generic[ResponseFormat]):
         temperature: float | None = 0.0,
         top_p: float = 1.0,
         response_format: Type[ResponseFormat] = str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
     ) -> "BatchResponses":
         """Factory constructor.
 
@@ -190,7 +190,8 @@ class BatchResponses(Generic[ResponseFormat]):
             temperature (float, optional): Sampling temperature. Defaults to 0.0.
             top_p (float, optional): Nucleus sampling parameter. Defaults to 1.0.
             response_format (Type[ResponseFormat], optional): Expected output type. Defaults to ``str``.
-            batch_size (int, optional): Max unique prompts per API call. Defaults to 128.
+            batch_size (int | None, optional): Max unique prompts per API call. Defaults to None
+                (automatic batch size optimization). Set to a positive integer for fixed batch size.
 
         Returns:
             BatchResponses: Configured instance backed by a batching proxy.
@@ -206,14 +207,17 @@ class BatchResponses(Generic[ResponseFormat]):
         )
 
     @classmethod
-    def of_task(cls, client: OpenAI, model_name: str, task: PreparedTask, batch_size: int = 128) -> "BatchResponses":
+    def of_task(
+        cls, client: OpenAI, model_name: str, task: PreparedTask[ResponseFormat], batch_size: int | None = None
+    ) -> "BatchResponses":
         """Factory from a PreparedTask.
 
         Args:
             client (OpenAI): OpenAI client.
             model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
             task (PreparedTask): Prepared task with instructions and response format.
-            batch_size (int, optional): Max unique prompts per API call. Defaults to 128.
+            batch_size (int | None, optional): Max unique prompts per API call. Defaults to None
+                (automatic batch size optimization). Set to a positive integer for fixed batch size.
 
         Returns:
             BatchResponses: Configured instance backed by a batching proxy.
@@ -294,8 +298,10 @@ class BatchResponses(Generic[ResponseFormat]):
         """
         messages = [Message(id=i, body=message) for i, message in enumerate(user_messages)]
         responses: ParsedResponse[Response[ResponseFormat]] = self._request_llm(messages)
+        if not responses.output_parsed:
+            return [None] * len(messages)
         response_dict = {message.id: message.body for message in responses.output_parsed.assistant_messages}
-        sorted_responses = [response_dict.get(m.id, None) for m in messages]
+        sorted_responses: List[ResponseFormat | None] = [response_dict.get(m.id, None) for m in messages]
         return sorted_responses
 
     @observe(_LOGGER)
@@ -308,7 +314,8 @@ class BatchResponses(Generic[ResponseFormat]):
         Returns:
             List[ResponseFormat | None]: Assistant responses aligned to ``inputs``.
         """
-        return self.cache.map(inputs, self._predict_chunk)
+        result = self.cache.map(inputs, self._predict_chunk)
+        return result  # type: ignore[return-value]
 
 
 @dataclass(frozen=True)
@@ -362,9 +369,9 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
     system_message: str
     temperature: float | None = 0.0
     top_p: float = 1.0
-    response_format: Type[ResponseFormat] = str
+    response_format: Type[ResponseFormat] = str  # type: ignore[assignment]
     cache: AsyncBatchingMapProxy[str, ResponseFormat] = field(
-        default_factory=lambda: AsyncBatchingMapProxy(batch_size=128, max_concurrency=8)
+        default_factory=lambda: AsyncBatchingMapProxy(batch_size=None, max_concurrency=8)
     )
     _vectorized_system_message: str = field(init=False)
     _model_json_schema: dict = field(init=False)
@@ -378,7 +385,7 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
         temperature: float | None = 0.0,
         top_p: float = 1.0,
         response_format: Type[ResponseFormat] = str,
-        batch_size: int = 128,
+        batch_size: int | None = None,
         max_concurrency: int = 8,
     ) -> "AsyncBatchResponses":
         """Factory constructor.
@@ -390,7 +397,8 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
             temperature (float, optional): Sampling temperature. Defaults to 0.0.
             top_p (float, optional): Nucleus sampling parameter. Defaults to 1.0.
             response_format (Type[ResponseFormat], optional): Expected output type. Defaults to ``str``.
-            batch_size (int, optional): Max unique prompts per API call. Defaults to 128.
+            batch_size (int | None, optional): Max unique prompts 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:
@@ -408,7 +416,12 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
 
     @classmethod
     def of_task(
-        cls, client: AsyncOpenAI, model_name: str, task: PreparedTask, batch_size: int = 128, max_concurrency: int = 8
+        cls,
+        client: AsyncOpenAI,
+        model_name: str,
+        task: PreparedTask[ResponseFormat],
+        batch_size: int | None = None,
+        max_concurrency: int = 8,
     ) -> "AsyncBatchResponses":
         """Factory from a PreparedTask.
 
@@ -416,7 +429,8 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
             client (AsyncOpenAI): OpenAI async client.
             model_name (str): For Azure OpenAI, use your deployment name. For OpenAI, use the model name.
             task (PreparedTask): Prepared task with instructions and response format.
-            batch_size (int, optional): Max unique prompts per API call. Defaults to 128.
+            batch_size (int | None, optional): Max unique prompts 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:
@@ -439,8 +453,8 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
             _vectorize_system_message(self.system_message),
         )
 
-    @observe(_LOGGER)
     @backoff_async(exceptions=[RateLimitError, InternalServerError], scale=1, max_retries=12)
+    @observe(_LOGGER)
     async def _request_llm(self, user_messages: List[Message[str]]) -> ParsedResponse[Response[ResponseFormat]]:
         """Make a single async call to the OpenAI JSON‑mode endpoint.
 
@@ -493,10 +507,12 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
         The function is pure – it has no side‑effects and the result depends only on its arguments.
         """
         messages = [Message(id=i, body=message) for i, message in enumerate(user_messages)]
-        responses: ParsedResponse[Response[ResponseFormat]] = await self._request_llm(messages)
+        responses: ParsedResponse[Response[ResponseFormat]] = await self._request_llm(messages)  # type: ignore[call-issue]
+        if not responses.output_parsed:
+            return [None] * len(messages)
         response_dict = {message.id: message.body for message in responses.output_parsed.assistant_messages}
         # Ensure proper handling for missing IDs - this shouldn't happen in normal operation
-        sorted_responses = [response_dict.get(m.id, None) for m in messages]
+        sorted_responses: List[ResponseFormat | None] = [response_dict.get(m.id, None) for m in messages]
         return sorted_responses
 
     @observe(_LOGGER)
@@ -509,4 +525,5 @@ class AsyncBatchResponses(Generic[ResponseFormat]):
         Returns:
             List[ResponseFormat | None]: Assistant responses aligned to ``inputs``.
         """
-        return await self.cache.map(inputs, self._predict_chunk)
+        result = await self.cache.map(inputs, self._predict_chunk)
+        return result  # type: ignore[return-value]