Compare commits

...
Author SHA1 Message Date
Nicolas Mowen ef0de74fa5 Refactor to use async openai 2026-11-04 10:42:02 -07:00
GeoandGitHub 169c5e1303 fix(genai): run blocking OpenAI chat calls in worker threads (#24868)
CI / AMD64 Build (push) Waiting to run
CI / AMD64 Smoke Test (push) Blocked by required conditions
CI / ARM Build (push) Waiting to run
CI / Jetson Jetpack 6 (push) Waiting to run
CI / AMD64 Extra Build (push) Blocked by required conditions
CI / ARM Extra Build (push) Blocked by required conditions
CI / Synaptics Build (push) Blocked by required conditions
CI / Assemble and push default build (push) Blocked by required conditions
* Don't block the event loop in OpenAI chat streaming

* Close OpenAI stream opened after the caller disconnected

* Show GenAI worker thread usage in system metrics

* Remove GenAI thread stats per review

* Use Starlette iterate_in_threadpool for OpenAI stream
2026-11-04 10:32:16 -07:00
4 changed files with 235 additions and 65 deletions
+3 -1
View File
@@ -1,5 +1,6 @@
"""Chat and LLM tool calling APIs.""" """Chat and LLM tool calling APIs."""
import asyncio
import base64 import base64
import json import json
import logging import logging
@@ -1653,7 +1654,8 @@ async def chat_completion(
f"Calling LLM (iteration {tool_iterations + 1}/{max_iterations}) " f"Calling LLM (iteration {tool_iterations + 1}/{max_iterations}) "
f"with {len(conversation)} message(s) in conversation" f"with {len(conversation)} message(s) in conversation"
) )
response = genai_client.chat_with_tools( response = await asyncio.to_thread(
genai_client.chat_with_tools,
messages=conversation, messages=conversation,
tools=llm_tools if llm_tools else None, tools=llm_tools if llm_tools else None,
tool_choice="auto", tool_choice="auto",
+18 -10
View File
@@ -7,7 +7,8 @@ overrides what is genuinely Azure-specific:
- Client construction: parses ``api-version`` out of the configured - Client construction: parses ``api-version`` out of the configured
``base_url`` query string and instantiates :class:`openai.AzureOpenAI` ``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 malformed; :class:`GenAIClientManager` catches the exception and
disables the provider. disables the provider.
- Context size: Azure does not expose a per-model ``max_model_len`` field - Context size: Azure does not expose a per-model ``max_model_len`` field
@@ -29,9 +30,10 @@ the ``transcribe`` role:
""" """
import logging import logging
from typing import Any
from urllib.parse import parse_qs, urlparse from urllib.parse import parse_qs, urlparse
from openai import AzureOpenAI from openai import AsyncAzureOpenAI, AzureOpenAI
from frigate.config import GenAIProviderEnum from frigate.config import GenAIProviderEnum
from frigate.genai import register_genai_provider from frigate.genai import register_genai_provider
@@ -44,8 +46,8 @@ logger = logging.getLogger(__name__)
class AzureOpenAIClient(OpenAIClient): class AzureOpenAIClient(OpenAIClient):
"""Generative AI client for Frigate using Azure OpenAI.""" """Generative AI client for Frigate using Azure OpenAI."""
def _init_provider(self) -> AzureOpenAI: def _client_options(self) -> dict[str, Any]:
"""Initialize the AzureOpenAI client from the configured base_url.""" """Build the Azure client options from the configured base_url."""
parsed_url = urlparse(self.genai_config.base_url or "") parsed_url = urlparse(self.genai_config.base_url or "")
query_params = parse_qs(parsed_url.query) query_params = parse_qs(parsed_url.query)
api_version = query_params.get("api-version", [None])[0] api_version = query_params.get("api-version", [None])[0]
@@ -53,13 +55,19 @@ class AzureOpenAIClient(OpenAIClient):
if not api_version: if not api_version:
raise ValueError("Azure OpenAI base_url is missing 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( def _init_provider(self) -> AzureOpenAI:
api_key=self.genai_config.api_key, """Initialize the AzureOpenAI client from the configured base_url."""
api_version=api_version, return AzureOpenAI(**self._client_options())
azure_endpoint=azure_endpoint,
) 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: def get_context_size(self) -> int:
"""Azure does not reliably surface per-model context size; use 128K.""" """Azure does not reliably surface per-model context size; use 128K."""
+27 -10
View File
@@ -7,7 +7,7 @@ from collections.abc import AsyncGenerator
from typing import Any from typing import Any
from httpx import TimeoutException from httpx import TimeoutException
from openai import OpenAI from openai import AsyncOpenAI, OpenAI
from frigate.config import GenAIProviderEnum from frigate.config import GenAIProviderEnum
from frigate.genai import GenAIClient, register_genai_provider from frigate.genai import GenAIClient, register_genai_provider
@@ -15,6 +15,7 @@ from frigate.genai.utils import interleave_images
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
# gpt-transcribe replaced the singular `language` field with a `languages` array # gpt-transcribe replaced the singular `language` field with a `languages` array
# and rejects a request that sends both. Older transcription models # and rejects a request that sends both. Older transcription models
# (gpt-4o-transcribe, gpt-4o-mini-transcribe, whisper-1) still take the singular # (gpt-4o-transcribe, gpt-4o-mini-transcribe, whisper-1) still take the singular
@@ -45,8 +46,8 @@ class OpenAIClient(GenAIClient):
provider: OpenAI provider: OpenAI
context_size: int | None = None context_size: int | None = None
def _init_provider(self) -> OpenAI: def _client_options(self) -> dict[str, Any]:
"""Initialize the client. """Build the constructor options shared by the sync and async clients.
Subclasses (e.g. Azure) should raise on configuration errors; the Subclasses (e.g. Azure) should raise on configuration errors; the
manager catches construction failures and disables the provider. manager catches construction failures and disables the provider.
@@ -62,7 +63,19 @@ class OpenAIClient(GenAIClient):
if self.genai_config.base_url: if self.genai_config.base_url:
provider_opts["base_url"] = 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( def _send(
self, self,
@@ -412,9 +425,13 @@ class OpenAIClient(GenAIClient):
finish_reason = "stop" finish_reason = "stop"
usage_stats: dict[str, Any] | None = None usage_stats: dict[str, Any] | None = None
stream = self.provider.chat.completions.create(**request_params) # Exiting these blocks, including when the caller disconnects and
# the generator is cancelled, closes the response and its connection
for chunk in stream: 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) chunk_usage = getattr(chunk, "usage", None)
if chunk_usage is not None: if chunk_usage is not None:
usage_stats = _stats_from_openai_usage(chunk_usage) usage_stats = _stats_from_openai_usage(chunk_usage)
@@ -431,9 +448,9 @@ class OpenAIClient(GenAIClient):
# Extract reasoning deltas (reasoning_content or reasoning, # Extract reasoning deltas (reasoning_content or reasoning,
# depending on the server) # depending on the server)
reasoning_delta = getattr(delta, "reasoning_content", None) or getattr( reasoning_delta = getattr(
delta, "reasoning", None delta, "reasoning_content", None
) ) or getattr(delta, "reasoning", None)
if reasoning_delta: if reasoning_delta:
reasoning_parts.append(reasoning_delta) reasoning_parts.append(reasoning_delta)
yield ("reasoning_delta", reasoning_delta) yield ("reasoning_delta", reasoning_delta)
+150 -7
View File
@@ -145,12 +145,72 @@ def _openai_chunk(content=None, tool_calls=None, finish_reason=None, usage=None)
return SimpleNamespace(choices=[choice], usage=usage) return SimpleNamespace(choices=[choice], usage=usage)
class _FakeAsyncOpenAIStream:
"""Mimics openai.AsyncStream: async iterable and async context manager.
With ``stall_after`` set, the stream stops producing after that many
chunks, like a model that stops sending tokens, and sets ``stalled``.
"""
def __init__(self, chunks, stall_after=None):
self._chunks = chunks
self._stall_after = stall_after
self.stalled = asyncio.Event()
self.close = AsyncMock()
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
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): class TestOpenAIProvider(unittest.TestCase):
def _client(self): def _client(self):
return _make_client( return _make_client(
"openai", model="gpt-4o", api_key="k", base_url="http://localhost:9999/v1" "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): def test_stream_tool_call_arguments_are_dict(self):
# Arguments arrive split across chunks, as the real API streams them. # Arguments arrive split across chunks, as the real API streams them.
chunks = [ chunks = [
@@ -162,8 +222,9 @@ class TestOpenAIProvider(unittest.TestCase):
_openai_chunk(tool_calls=[_openai_tc(0, arguments=' "person"}')]), _openai_chunk(tool_calls=[_openai_tc(0, arguments=' "person"}')]),
_openai_chunk(finish_reason="tool_calls"), _openai_chunk(finish_reason="tool_calls"),
] ]
client = self._client() client, _ = self._streaming_client(
client.provider.chat.completions.create = MagicMock(return_value=iter(chunks)) AsyncMock(return_value=_FakeAsyncOpenAIStream(chunks))
)
final = _final_message(_collect(client, SIMPLE_MESSAGES)) final = _final_message(_collect(client, SIMPLE_MESSAGES))
self.assertEqual(final["finish_reason"], "tool_calls") self.assertEqual(final["finish_reason"], "tool_calls")
@@ -177,8 +238,9 @@ class TestOpenAIProvider(unittest.TestCase):
_openai_chunk(content="lo"), _openai_chunk(content="lo"),
_openai_chunk(finish_reason="stop"), _openai_chunk(finish_reason="stop"),
] ]
client = self._client() client, _ = self._streaming_client(
client.provider.chat.completions.create = MagicMock(return_value=iter(chunks)) AsyncMock(return_value=_FakeAsyncOpenAIStream(chunks))
)
events = _collect(client, SIMPLE_MESSAGES) events = _collect(client, SIMPLE_MESSAGES)
deltas = [v for (k, v) in events if k == "content_delta"] deltas = [v for (k, v) in events if k == "content_delta"]
@@ -186,14 +248,84 @@ class TestOpenAIProvider(unittest.TestCase):
self.assertEqual(_final_message(events)["content"], "hello") self.assertEqual(_final_message(events)["content"], "hello")
def test_multimodal_message_does_not_crash(self): def test_multimodal_message_does_not_crash(self):
client = self._client() client, _ = self._streaming_client(
client.provider.chat.completions.create = MagicMock( AsyncMock(
return_value=iter([_openai_chunk(content="ok", finish_reason="stop")]) return_value=_FakeAsyncOpenAIStream(
[_openai_chunk(content="ok", finish_reason="stop")]
)
)
) )
# Passing the OpenAI-native multimodal list through must not raise. # Passing the OpenAI-native multimodal list through must not raise.
final = _final_message(_collect(client, MULTIMODAL_MESSAGES)) final = _final_message(_collect(client, MULTIMODAL_MESSAGES))
self.assertEqual(final["content"], "ok") self.assertEqual(final["content"], "ok")
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():
opening = asyncio.Event()
async def hang(**kwargs):
opening.set()
await asyncio.Event().wait()
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 opening.wait()
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
return async_client
async_client = asyncio.run(_run())
async_client.close.assert_awaited_once()
def test_stream_closed_when_caller_leaves_mid_stream(self):
# 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 stream.stalled.wait()
task.cancel()
with self.assertRaises(asyncio.CancelledError):
await asyncio.wait_for(task, timeout=1)
return first, stream, async_client
first, stream, async_client = asyncio.run(_run())
self.assertEqual(first, ("content_delta", "a"))
stream.close.assert_awaited_once()
async_client.close.assert_awaited_once()
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Gemini # Gemini
@@ -940,6 +1072,17 @@ class TestAzureOpenAITranscribe(unittest.TestCase):
self.assertIsInstance(client.provider, AzureOpenAI) self.assertIsInstance(client.provider, AzureOpenAI)
self.assertTrue(client.supports_transcription) 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): def test_transcribe_inherited(self):
client = self._client() client = self._client()
create = MagicMock(return_value="azure text") create = MagicMock(return_value="azure text")