From ef0de74fa55c5f5273b6f899f5ea3374482387bd Mon Sep 17 00:00:00 2001 From: Nicolas Mowen Date: Wed, 4 Nov 2026 10:42:02 -0700 Subject: [PATCH] Refactor to use async openai --- frigate/genai/plugins/azure-openai.py | 28 ++-- frigate/genai/plugins/openai.py | 49 +++---- frigate/test/test_genai_providers.py | 203 +++++++++++++++----------- 3 files changed, 158 insertions(+), 122 deletions(-) diff --git a/frigate/genai/plugins/azure-openai.py b/frigate/genai/plugins/azure-openai.py index ce6a906231..16afa764b8 100644 --- a/frigate/genai/plugins/azure-openai.py +++ b/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.""" diff --git a/frigate/genai/plugins/openai.py b/frigate/genai/plugins/openai.py index 1558a648ef..fe4b63a636 100644 --- a/frigate/genai/plugins/openai.py +++ b/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) diff --git a/frigate/test/test_genai_providers.py b/frigate/test/test_genai_providers.py index 2d0aae6112..8ae8fb34f1 100644 --- a/frigate/test/test_genai_providers.py +++ b/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( - [_openai_chunk(content="ok", finish_reason="stop")] + 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")