refactor: move blocking provider I/O off event loop (#5835)

# What does this PR do?
This PR moves blocking operations to worker threads, parallelizes
independent async calls, and adds reusable OpenAI client caching keyed
by `(api_key, base_url)` with shutdown cleanup.
It updates zstd request decompression, localfs and S3 file operations,
sqlite-vec hybrid search execution, and image localization paths in both
the OpenAI mixin and VertexAI adapter to reduce event-loop blocking and
improve request throughput.

## Test Plan
- `uv run pytest -q
tests/unit/providers/utils/inference/test_openai_mixin_models.py
tests/unit/providers/utils/inference/test_openai_mixin_provider_data.py
tests/unit/providers/inference/vertexai/test_adapter_chat.py
tests/unit/files/test_files.py -x --tb=short`
  - Result: `118 passed in 1.29s`

Signed-off-by: Sébastien Han <seb@redhat.com>

Author: leseb

src/ogx/core/server/server.py

@@ -357,18 +357,27 @@ async def __call__(self, scope: Scope, receive: Receive, send: Send) -> Any:
 
             try:
                 max_decompressed_size = 100 * 1024 * 1024  # 100 MB
-                decompressor = zstandard.ZstdDecompressor()
-                # Use streaming decompression to handle frames without content size
-                reader = decompressor.stream_reader(compressed_body)
-                decompressed_body = reader.read(max_decompressed_size)
-                if reader.read(1):
-                    reader.close()
+
+                def _decompress_zstd(compressed: bytes, max_size: int) -> tuple[bytes | None, bool]:
+                    decompressor = zstandard.ZstdDecompressor()
+                    reader = decompressor.stream_reader(compressed)
+                    try:
+                        data = reader.read(max_size)
+                        is_oversized = bool(reader.read(1))
+                        return (None, True) if is_oversized else (data, False)
+                    finally:
+                        reader.close()
+
+                decompressed_body, oversized = await asyncio.to_thread(
+                    _decompress_zstd, compressed_body, max_decompressed_size
+                )
+
+                if oversized:
                     return await _send_error_response(
                         send,
                         status=413,
                         message=f"Decompressed request body exceeds maximum allowed size of {max_decompressed_size} bytes",
                     )
-                reader.close()
 
                 # Strip content-encoding header and update content-length
                 new_headers = [

src/ogx/providers/inline/files/localfs/files.py

@@ -21,6 +21,7 @@
   download as a belt-and-suspenders measure.
 """
 
+import asyncio
 import re
 import time
 import uuid
@@ -165,8 +166,11 @@ async def openai_upload_file(
         content = await file.read()
         file_size = len(content)
 
-        with open(file_path, "wb") as f:
-            f.write(content)
+        def _write_file() -> None:
+            with open(file_path, "wb") as f:
+                f.write(content)
+
+        await asyncio.to_thread(_write_file)
 
         created_at = int(time.time())
         expires_at = created_at + self.config.ttl_secs
@@ -255,8 +259,12 @@ async def openai_delete_file(self, request: DeleteFileRequest) -> OpenAIFileDele
         file_id = request.file_id
         # Delete physical file
         _, file_path = await self._lookup_file_id(file_id, action=Action.DELETE)
-        if file_path.exists():
-            file_path.unlink()
+
+        def _delete_if_exists() -> None:
+            if file_path.exists():
+                file_path.unlink()
+
+        await asyncio.to_thread(_delete_if_exists)
 
         # Delete metadata from database
         assert self.sql_store is not None, "Files provider not initialized"
@@ -278,9 +286,11 @@ async def openai_retrieve_file_content(self, request: RetrieveFileContentRequest
             await self.openai_delete_file(DeleteFileRequest(file_id=file_id))
             raise OpenAIFileObjectNotFoundError(file_id)
 
+        file_content = await asyncio.to_thread(file_path.read_bytes)
+
         # Return as binary response with appropriate content type
         return Response(
-            content=file_path.read_bytes(),
+            content=file_content,
             media_type="application/octet-stream",
             headers={
                 "Content-Disposition": f'attachment; filename="{sanitize_content_disposition_filename(file_obj.filename)}"'

src/ogx/providers/inline/vector_io/sqlite_vec/sqlite_vec.py

@@ -460,8 +460,10 @@ async def query_hybrid(
             reranker_params = {}
 
         # Get results from both search methods, passing filters to each
-        vector_response = await self.query_vector(embedding, k, score_threshold, filters)
-        keyword_response = await self.query_keyword(query_string, k, score_threshold, filters)
+        vector_response, keyword_response = await asyncio.gather(
+            self.query_vector(embedding, k, score_threshold, filters),
+            self.query_keyword(query_string, k, score_threshold, filters),
+        )
 
         # Convert responses to score dictionaries using chunk_id (EmbeddedChunk inherits from Chunk)
         vector_scores = {

src/ogx/providers/remote/files/s3/files.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
 import uuid
 from datetime import UTC, datetime
 from typing import TYPE_CHECKING, Any, cast
@@ -73,33 +74,36 @@ def _create_s3_client(config: S3FilesImplConfig) -> "S3Client":
 
 
 async def _create_bucket_if_not_exists(client: "S3Client", config: S3FilesImplConfig) -> None:
-    try:
-        client.head_bucket(Bucket=config.bucket_name)
-    except ClientError as e:
-        error_code = e.response["Error"]["Code"]
-        if error_code == "404":
-            if not config.auto_create_bucket:
-                raise RuntimeError(
-                    f"S3 bucket '{config.bucket_name}' does not exist. "
-                    f"Either create the bucket manually or set 'auto_create_bucket: true' in your configuration."
-                ) from e
-            try:
-                # For us-east-1, we can't specify LocationConstraint
-                if config.region == "us-east-1":
-                    client.create_bucket(Bucket=config.bucket_name)
-                else:
-                    client.create_bucket(
-                        Bucket=config.bucket_name,
-                        CreateBucketConfiguration=cast(Any, {"LocationConstraint": config.region}),
-                    )
-            except ClientError as create_error:
-                raise RuntimeError(
-                    f"Failed to create S3 bucket '{config.bucket_name}': {create_error}"
-                ) from create_error
-        elif error_code == "403":
-            raise RuntimeError(f"Access denied to S3 bucket '{config.bucket_name}'") from e
-        else:
-            raise RuntimeError(f"Failed to access S3 bucket '{config.bucket_name}': {e}") from e
+    def _check_and_create() -> None:
+        try:
+            client.head_bucket(Bucket=config.bucket_name)
+        except ClientError as e:
+            error_code = e.response["Error"]["Code"]
+            if error_code == "404":
+                if not config.auto_create_bucket:
+                    raise RuntimeError(
+                        f"S3 bucket '{config.bucket_name}' does not exist. "
+                        f"Either create the bucket manually or set 'auto_create_bucket: true' in your configuration."
+                    ) from e
+                try:
+                    # For us-east-1, we can't specify LocationConstraint
+                    if config.region == "us-east-1":
+                        client.create_bucket(Bucket=config.bucket_name)
+                    else:
+                        client.create_bucket(
+                            Bucket=config.bucket_name,
+                            CreateBucketConfiguration=cast(Any, {"LocationConstraint": config.region}),
+                        )
+                except ClientError as create_error:
+                    raise RuntimeError(
+                        f"Failed to create S3 bucket '{config.bucket_name}': {create_error}"
+                    ) from create_error
+            elif error_code == "403":
+                raise RuntimeError(f"Access denied to S3 bucket '{config.bucket_name}'") from e
+            else:
+                raise RuntimeError(f"Failed to access S3 bucket '{config.bucket_name}': {e}") from e
+
+    await asyncio.to_thread(_check_and_create)
 
 
 def _make_file_object(
@@ -164,7 +168,8 @@ async def _get_file(
     async def _delete_file(self, file_id: str) -> None:
         """Delete a file from S3 and the database."""
         try:
-            self.client.delete_object(
+            await asyncio.to_thread(
+                self.client.delete_object,
                 Bucket=self._config.bucket_name,
                 Key=file_id,
             )
@@ -251,7 +256,8 @@ async def openai_upload_file(
         await self.sql_store.insert("openai_files", entry)
 
         try:
-            self.client.put_object(
+            await asyncio.to_thread(
+                self.client.put_object,
                 Bucket=self._config.bucket_name,
                 Key=file_id,
                 Body=content,
@@ -319,12 +325,17 @@ async def openai_retrieve_file_content(self, request: RetrieveFileContentRequest
         row = await self._get_file(file_id)
 
         try:
-            response = self.client.get_object(
-                Bucket=self._config.bucket_name,
-                Key=row["id"],
-            )
-            # TODO: can we stream this instead of loading it into memory
-            content = response["Body"].read()
+
+            def _download_from_s3() -> bytes:
+                response = self.client.get_object(
+                    Bucket=self._config.bucket_name,
+                    Key=row["id"],
+                )
+                # TODO: can we stream this instead of loading it into memory
+                body: bytes = response["Body"].read()
+                return body
+
+            content = await asyncio.to_thread(_download_from_s3)
         except ClientError as e:
             if e.response["Error"]["Code"] == "NoSuchKey":
                 await self._delete_file(file_id)

src/ogx/providers/remote/inference/vertexai/vertexai.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import base64
 import struct
 import time
@@ -742,7 +743,7 @@ async def openai_chat_completion(
         self._warn_unsupported_chat_params(params)
         tools, tool_choice = self._resolve_deprecated_tools(params)
 
-        messages = [await self._localize_image_url(message) for message in params.messages]
+        messages = list(await asyncio.gather(*[self._localize_image_url(message) for message in params.messages]))
         system_instruction, contents = converters.convert_openai_messages_to_gemini(messages)
         tools_input = converters.convert_openai_tools_to_gemini(tools)
         config = self._build_generation_config(

src/ogx/providers/utils/inference/openai_mixin.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
+import asyncio
 import base64
 import ssl
 import uuid
@@ -14,7 +15,7 @@
 import httpx
 from openai import AsyncOpenAI, DefaultAsyncHttpxClient
 from openai.types.chat import ChatCompletionChunk
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 
 from ogx.core.request_headers import NeedsRequestProviderData
 from ogx.log import get_logger
@@ -106,9 +107,7 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
     # Format: {"model_id": {"embedding_dimension": 1536, "context_length": 8192}}
     embedding_model_metadata: dict[str, dict[str, int]] = {}
 
-    # Cache of available models keyed by model ID
-    # This is set in list_models() and used in check_model_availability()
-    _model_cache: dict[str, Model] = {}
+    _model_cache: dict[str, Model] = PrivateAttr(default_factory=dict)
 
     # Optional field name in provider data to look for API key, which takes precedence
     provider_data_api_key_field: str | None = None
@@ -118,6 +117,9 @@ class OpenAIMixin(NeedsRequestProviderData, ABC, BaseModel):
     # Trade-off: SSL context changes require server restart
     shared_ssl_context: ssl.SSLContext | bool = Field(default_factory=ssl.create_default_context, exclude=True)
 
+    _cached_client: AsyncOpenAI | None = PrivateAttr(default=None)
+    _cached_client_key: tuple[str, str] | None = PrivateAttr(default=None)
+
     def get_api_key(self) -> str | None:
         """
         Get the API key.
@@ -199,10 +201,7 @@ async def list_provider_model_ids(self) -> Iterable[str]:
 
         :return: An iterable of model IDs or None if not implemented
         """
-        client = self.client
-        async with client:
-            model_ids = [m.id async for m in client.models.list()]
-        return model_ids
+        return [m.id async for m in self.client.models.list()]
 
     async def initialize(self) -> None:
         """
@@ -215,26 +214,20 @@ async def initialize(self) -> None:
         pass
 
     async def shutdown(self) -> None:
-        """
-        Shutdown the OpenAI mixin.
-
-        This method provides a default implementation that does nothing.
-        Subclasses can override this method to perform cleanup tasks
-        such as closing connections, releasing resources, etc.
-        """
-        pass
+        """Shutdown the OpenAI mixin, closing the cached HTTP client."""
+        if self._cached_client is not None:
+            await self._cached_client.close()
+            self._cached_client = None
+            self._cached_client_key = None
 
     @property
     def client(self) -> AsyncOpenAI:
         """
         Get an AsyncOpenAI client instance.
 
-        Uses the abstract methods get_api_key() and get_base_url() which must be
-        implemented by child classes.
-
-        Network configuration from config.network is automatically applied.
-        Users can also provide the API key via the provider data header, which
-        is used instead of any config API key.
+        Caches the client keyed by (api_key, base_url) for connection reuse.
+        When the key changes (e.g. per-request provider_data), a new client
+        is created and cached in its place.
         """
 
         api_key = self._get_api_key_from_config_or_provider_data()
@@ -244,19 +237,15 @@ def client(self) -> AsyncOpenAI:
                 message += f' Please provide a valid API key in the provider data header, e.g. x-ogx-provider-data: {{"{self.provider_data_api_key_field}": "<API_KEY>"}}.'
             raise ValueError(message)
 
+        base_url = self.get_base_url()
+        cache_key = (api_key, base_url)
+
+        if self._cached_client is not None and self._cached_client_key == cache_key:
+            return self._cached_client
+
         extra_params = self.get_extra_client_params()
         network_kwargs = build_network_client_kwargs(self.config.network)
 
-        # Handle http_client creation/merging:
-        # - If get_extra_client_params() provides an http_client (e.g., OCI with custom auth),
-        #   merge network config into it. The merge behavior:
-        #   * Preserves auth from get_extra_client_params() (provider-specific auth like OCI signer)
-        #   * Preserves headers from get_extra_client_params() as base
-        #   * Applies network config (TLS, proxy, timeout, headers) on top
-        #   * Network config headers take precedence over provider headers (allows override)
-        # - Otherwise, if network config exists, create http_client from it
-        # - Otherwise, use a cached SSL context for performance
-        # This allows providers with custom auth to still use standard network settings
         if "http_client" in extra_params:
             if network_kwargs:
                 extra_params["http_client"] = _merge_network_config_into_client(
@@ -267,12 +256,16 @@ def client(self) -> AsyncOpenAI:
         else:
             extra_params["http_client"] = DefaultAsyncHttpxClient(verify=self.shared_ssl_context)
 
-        return AsyncOpenAI(
+        client = AsyncOpenAI(
             api_key=api_key,
-            base_url=self.get_base_url(),
+            base_url=base_url,
             **extra_params,
         )
 
+        self._cached_client = client
+        self._cached_client_key = cache_key
+        return client
+
     def _get_api_key_from_config_or_provider_data(self) -> str | None:
         api_key = self.get_api_key()
 
@@ -430,7 +423,7 @@ async def _localize_image_url(m: OpenAIMessageParam) -> OpenAIMessageParam:
                 # else it's a string and we don't need to modify it
                 return m
 
-            messages = [await _localize_image_url(m) for m in messages]
+            messages = list(await asyncio.gather(*[_localize_image_url(m) for m in messages]))
 
         request_params = await prepare_openai_completion_params(
             model=provider_model_id,