Refactor to use async openai

frigate/genai/plugins/azure-openai.py

@@ -7,7 +7,8 @@ overrides what is genuinely Azure-specific:
 
 - Client construction: parses ``api-version`` out of the configured
   ``base_url`` query string and instantiates :class:`openai.AzureOpenAI`
-  with ``azure_endpoint`` instead of ``base_url``. Raises if the URL is
+  (or :class:`openai.AsyncAzureOpenAI` for streaming chat) with
+  ``azure_endpoint`` instead of ``base_url``. Raises if the URL is
   malformed; :class:`GenAIClientManager` catches the exception and
   disables the provider.
 - Context size: Azure does not expose a per-model ``max_model_len`` field
@@ -29,9 +30,10 @@ the ``transcribe`` role:
 """
 
 import logging
+from typing import Any
 from urllib.parse import parse_qs, urlparse
 
-from openai import AzureOpenAI
+from openai import AsyncAzureOpenAI, AzureOpenAI
 
 from frigate.config import GenAIProviderEnum
 from frigate.genai import register_genai_provider
@@ -44,8 +46,8 @@ logger = logging.getLogger(__name__)
 class AzureOpenAIClient(OpenAIClient):
     """Generative AI client for Frigate using Azure OpenAI."""
 
-    def _init_provider(self) -> AzureOpenAI:
-        """Initialize the AzureOpenAI client from the configured base_url."""
+    def _client_options(self) -> dict[str, Any]:
+        """Build the Azure client options from the configured base_url."""
         parsed_url = urlparse(self.genai_config.base_url or "")
         query_params = parse_qs(parsed_url.query)
         api_version = query_params.get("api-version", [None])[0]
@@ -53,13 +55,19 @@ class AzureOpenAIClient(OpenAIClient):
         if not api_version:
             raise ValueError("Azure OpenAI base_url is missing api-version.")
 
-        azure_endpoint = f"{parsed_url.scheme}://{parsed_url.netloc}/"
+        return {
+            "api_key": self.genai_config.api_key,
+            "api_version": api_version,
+            "azure_endpoint": f"{parsed_url.scheme}://{parsed_url.netloc}/",
+        }
 
-        return AzureOpenAI(
-            api_key=self.genai_config.api_key,
-            api_version=api_version,
-            azure_endpoint=azure_endpoint,
-        )
+    def _init_provider(self) -> AzureOpenAI:
+        """Initialize the AzureOpenAI client from the configured base_url."""
+        return AzureOpenAI(**self._client_options())
+
+    def _init_async_provider(self) -> AsyncAzureOpenAI:
+        """Create an async Azure client for streaming chat."""
+        return AsyncAzureOpenAI(**self._client_options())
 
     def get_context_size(self) -> int:
         """Azure does not reliably surface per-model context size; use 128K."""

frigate/genai/plugins/openai.py

@@ -1,15 +1,13 @@
 """OpenAI Provider for Frigate AI."""
 
-import asyncio
 import base64
 import json
 import logging
-from collections.abc import AsyncGenerator, Callable
+from collections.abc import AsyncGenerator
 from typing import Any
 
 from httpx import TimeoutException
-from openai import OpenAI
-from starlette.concurrency import iterate_in_threadpool
+from openai import AsyncOpenAI, OpenAI
 
 from frigate.config import GenAIProviderEnum
 from frigate.genai import GenAIClient, register_genai_provider
@@ -18,18 +16,6 @@ from frigate.genai.utils import interleave_images
 logger = logging.getLogger(__name__)
 
 
-async def _open_stream(create: Callable[..., Any], **kwargs: Any) -> Any:
-    """Open a stream in a worker thread, closing it if the caller goes away."""
-    task = asyncio.ensure_future(asyncio.to_thread(create, **kwargs))
-    try:
-        return await asyncio.shield(task)
-    except asyncio.CancelledError:
-        task.add_done_callback(
-            lambda t: t.cancelled() or t.exception() or t.result().close()
-        )
-        raise
-
-
 # gpt-transcribe replaced the singular `language` field with a `languages` array
 # and rejects a request that sends both. Older transcription models
 # (gpt-4o-transcribe, gpt-4o-mini-transcribe, whisper-1) still take the singular
@@ -60,8 +46,8 @@ class OpenAIClient(GenAIClient):
     provider: OpenAI
     context_size: int | None = None
 
-    def _init_provider(self) -> OpenAI:
-        """Initialize the client.
+    def _client_options(self) -> dict[str, Any]:
+        """Build the constructor options shared by the sync and async clients.
 
         Subclasses (e.g. Azure) should raise on configuration errors; the
         manager catches construction failures and disables the provider.
@@ -77,7 +63,19 @@ class OpenAIClient(GenAIClient):
         if self.genai_config.base_url:
             provider_opts["base_url"] = self.genai_config.base_url
 
-        return OpenAI(api_key=self.genai_config.api_key, **provider_opts)
+        return {"api_key": self.genai_config.api_key, **provider_opts}
+
+    def _init_provider(self) -> OpenAI:
+        """Initialize the client."""
+        return OpenAI(**self._client_options())
+
+    def _init_async_provider(self) -> AsyncOpenAI:
+        """Create an async client for streaming chat.
+
+        A new client is created per stream because its connection pool binds
+        to the event loop that first uses it.
+        """
+        return AsyncOpenAI(**self._client_options())
 
     def _send(
         self,
@@ -427,12 +425,13 @@ class OpenAIClient(GenAIClient):
             finish_reason = "stop"
             usage_stats: dict[str, Any] | None = None
 
-            stream = await _open_stream(
-                self.provider.chat.completions.create, **request_params
-            )
-
-            with stream:
-                async for chunk in iterate_in_threadpool(stream):
+            # Exiting these blocks, including when the caller disconnects and
+            # the generator is cancelled, closes the response and its connection
+            async with (
+                self._init_async_provider() as client,
+                await client.chat.completions.create(**request_params) as stream,
+            ):
+                async for chunk in stream:
                     chunk_usage = getattr(chunk, "usage", None)
                     if chunk_usage is not None:
                         usage_stats = _stats_from_openai_usage(chunk_usage)

frigate/test/test_genai_providers.py

@@ -15,7 +15,6 @@ crashing message conversion.
 import asyncio
 import base64
 import json
-import time
 import unittest
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -146,21 +145,45 @@ def _openai_chunk(content=None, tool_calls=None, finish_reason=None, usage=None)
     return SimpleNamespace(choices=[choice], usage=usage)
 
 
-class _FakeOpenAIStream:
-    """Mimics openai.Stream: iterable, closable, and a context manager."""
+class _FakeAsyncOpenAIStream:
+    """Mimics openai.AsyncStream: async iterable and async context manager.
 
-    def __init__(self, chunks):
-        self._chunks = iter(chunks)
-        self.close = MagicMock()
+    With ``stall_after`` set, the stream stops producing after that many
+    chunks, like a model that stops sending tokens, and sets ``stalled``.
+    """
 
-    def __iter__(self):
-        return self._chunks
+    def __init__(self, chunks, stall_after=None):
+        self._chunks = chunks
+        self._stall_after = stall_after
+        self.stalled = asyncio.Event()
+        self.close = AsyncMock()
 
-    def __enter__(self):
+    async def __aiter__(self):
+        for i, chunk in enumerate(self._chunks):
+            if i == self._stall_after:
+                self.stalled.set()
+                await asyncio.Event().wait()
+            yield chunk
+
+    async def __aenter__(self):
         return self
 
-    def __exit__(self, *exc):
-        self.close()
+    async def __aexit__(self, *exc):
+        await self.close()
+
+
+class _FakeAsyncOpenAI:
+    """Mimics openai.AsyncOpenAI: async context manager with a completions API."""
+
+    def __init__(self, create):
+        self.chat = SimpleNamespace(completions=SimpleNamespace(create=create))
+        self.close = AsyncMock()
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, *exc):
+        await self.close()
 
 
 class TestOpenAIProvider(unittest.TestCase):
@@ -169,6 +192,25 @@ class TestOpenAIProvider(unittest.TestCase):
             "openai", model="gpt-4o", api_key="k", base_url="http://localhost:9999/v1"
         )
 
+    def _streaming_client(self, create):
+        """Build a client whose async completions call is ``create``."""
+        client = self._client()
+        async_client = _FakeAsyncOpenAI(create)
+        client._init_async_provider = MagicMock(return_value=async_client)
+        return client, async_client
+
+    def test_async_client_uses_configured_options(self):
+        from openai import AsyncOpenAI
+
+        client = self._client()
+        async_client = client._init_async_provider()
+        try:
+            self.assertIsInstance(async_client, AsyncOpenAI)
+            self.assertEqual(async_client.base_url, client.provider.base_url)
+            self.assertEqual(async_client.api_key, "k")
+        finally:
+            asyncio.run(async_client.close())
+
     def test_stream_tool_call_arguments_are_dict(self):
         # Arguments arrive split across chunks, as the real API streams them.
         chunks = [
@@ -180,9 +222,8 @@ class TestOpenAIProvider(unittest.TestCase):
             _openai_chunk(tool_calls=[_openai_tc(0, arguments=' "person"}')]),
             _openai_chunk(finish_reason="tool_calls"),
         ]
-        client = self._client()
-        client.provider.chat.completions.create = MagicMock(
-            return_value=_FakeOpenAIStream(chunks)
+        client, _ = self._streaming_client(
+            AsyncMock(return_value=_FakeAsyncOpenAIStream(chunks))
         )
 
         final = _final_message(_collect(client, SIMPLE_MESSAGES))
@@ -197,9 +238,8 @@ class TestOpenAIProvider(unittest.TestCase):
             _openai_chunk(content="lo"),
             _openai_chunk(finish_reason="stop"),
         ]
-        client = self._client()
-        client.provider.chat.completions.create = MagicMock(
-            return_value=_FakeOpenAIStream(chunks)
+        client, _ = self._streaming_client(
+            AsyncMock(return_value=_FakeAsyncOpenAIStream(chunks))
         )
 
         events = _collect(client, SIMPLE_MESSAGES)
@@ -208,105 +248,83 @@ class TestOpenAIProvider(unittest.TestCase):
         self.assertEqual(_final_message(events)["content"], "hello")
 
     def test_multimodal_message_does_not_crash(self):
-        client = self._client()
-        client.provider.chat.completions.create = MagicMock(
-            return_value=_FakeOpenAIStream(
+        client, _ = self._streaming_client(
+            AsyncMock(
+                return_value=_FakeAsyncOpenAIStream(
                     [_openai_chunk(content="ok", finish_reason="stop")]
                 )
             )
+        )
         # Passing the OpenAI-native multimodal list through must not raise.
         final = _final_message(_collect(client, MULTIMODAL_MESSAGES))
         self.assertEqual(final["content"], "ok")
 
-    def test_stream_does_not_block_event_loop(self):
-        # Other coroutines must keep running while the stream blocks.
-        def slow_chunks():
-            for content in ("a", "b", "c"):
-                time.sleep(0.1)
-                yield _openai_chunk(content=content)
-            yield _openai_chunk(finish_reason="stop")
-
-        client = self._client()
-        client.provider.chat.completions.create = MagicMock(
-            return_value=_FakeOpenAIStream(slow_chunks())
+    def test_stream_uses_async_client_and_closes_it(self):
+        # The blocking sync client must never be used on the event loop.
+        stream = _FakeAsyncOpenAIStream(
+            [_openai_chunk(content="ok", finish_reason="stop")]
         )
+        create = AsyncMock(return_value=stream)
+        client, async_client = self._streaming_client(create)
+        client.provider.chat.completions.create = MagicMock()
 
+        final = _final_message(_collect(client, SIMPLE_MESSAGES))
+
+        self.assertEqual(final["content"], "ok")
+        self.assertTrue(create.await_args.kwargs["stream"])
+        client.provider.chat.completions.create.assert_not_called()
+        stream.close.assert_awaited_once()
+        async_client.close.assert_awaited_once()
+
+    def test_client_closed_when_caller_leaves_while_opening(self):
+        # Cancelling before the response arrives must still close the client.
         async def _run():
-            ticks = 0
-            done = asyncio.Event()
+            opening = asyncio.Event()
 
-            async def ticker():
-                nonlocal ticks
-                while not done.is_set():
-                    ticks += 1
-                    await asyncio.sleep(0.01)
+            async def hang(**kwargs):
+                opening.set()
+                await asyncio.Event().wait()
 
-            ticker_task = asyncio.create_task(ticker())
-            events = [
-                event
-                async for event in client.chat_with_tools_stream(
-                    messages=SIMPLE_MESSAGES, tools=None, tool_choice="auto"
-                )
-            ]
-            done.set()
-            await ticker_task
-            return events, ticks
-
-        events, ticks = asyncio.run(_run())
-        self.assertEqual(_final_message(events)["content"], "abc")
-        self.assertGreater(ticks, 10)
-
-    def test_stream_closed_when_caller_leaves_while_opening(self):
-        # A stream opened after the caller left must still be closed.
-        stream = MagicMock()
-
-        def slow_create(**kwargs):
-            time.sleep(0.2)
-            return stream
-
-        client = self._client()
-        client.provider.chat.completions.create = slow_create
-
-        async def _run():
+            client, async_client = self._streaming_client(hang)
             gen = client.chat_with_tools_stream(
                 messages=SIMPLE_MESSAGES, tools=None, tool_choice="auto"
             )
             task = asyncio.create_task(gen.__anext__())
-            await asyncio.sleep(0.05)
+            await opening.wait()
             task.cancel()
             with self.assertRaises(asyncio.CancelledError):
-                await task
-            await asyncio.sleep(0.3)
+                await asyncio.wait_for(task, timeout=1)
+            return async_client
 
-        asyncio.run(_run())
-        stream.close.assert_called_once()
+        async_client = asyncio.run(_run())
+        async_client.close.assert_awaited_once()
 
     def test_stream_closed_when_caller_leaves_mid_stream(self):
-        # A stream abandoned after the first chunk must be closed once.
-        def chunks_then_block():
-            yield _openai_chunk(content="a")
-            time.sleep(0.5)
-            yield _openai_chunk(finish_reason="stop")
-
-        stream = _FakeOpenAIStream(chunks_then_block())
-        client = self._client()
-        client.provider.chat.completions.create = MagicMock(return_value=stream)
-
+        # Cancelling while waiting on a stalled model must close the stream
+        # right away instead of waiting for the next chunk.
         async def _run():
+            stream = _FakeAsyncOpenAIStream(
+                [_openai_chunk(content="a"), _openai_chunk(finish_reason="stop")],
+                stall_after=1,
+            )
+            client, async_client = self._streaming_client(
+                AsyncMock(return_value=stream)
+            )
             gen = client.chat_with_tools_stream(
                 messages=SIMPLE_MESSAGES, tools=None, tool_choice="auto"
             )
             first = await gen.__anext__()
             task = asyncio.create_task(gen.__anext__())
-            await asyncio.sleep(0.05)
+            await stream.stalled.wait()
             task.cancel()
             with self.assertRaises(asyncio.CancelledError):
-                await task
-            await gen.aclose()
-            return first
+                await asyncio.wait_for(task, timeout=1)
+            return first, stream, async_client
 
-        self.assertEqual(asyncio.run(_run()), ("content_delta", "a"))
-        stream.close.assert_called_once()
+        first, stream, async_client = asyncio.run(_run())
+        self.assertEqual(first, ("content_delta", "a"))
+        stream.close.assert_awaited_once()
+        async_client.close.assert_awaited_once()
 
 
 # ---------------------------------------------------------------------------
@@ -1054,6 +1072,17 @@ class TestAzureOpenAITranscribe(unittest.TestCase):
         self.assertIsInstance(client.provider, AzureOpenAI)
         self.assertTrue(client.supports_transcription)
 
+    def test_streaming_routes_through_async_azure_client(self):
+        from openai import AsyncAzureOpenAI
+
+        client = self._client()
+        async_client = client._init_async_provider()
+        try:
+            self.assertIsInstance(async_client, AsyncAzureOpenAI)
+            self.assertEqual(async_client.base_url, client.provider.base_url)
+        finally:
+            asyncio.run(async_client.close())
+
     def test_transcribe_inherited(self):
         client = self._client()
         create = MagicMock(return_value="azure text")