From 169c5e13035cd88167c5d0ad2e6c06f4a1a15dbb Mon Sep 17 00:00:00 2001 From: Geo Date: Wed, 4 Nov 2026 14:32:16 -0300 Subject: [PATCH] fix(genai): run blocking OpenAI chat calls in worker threads (#24868) * 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 --- frigate/api/chat.py | 4 +- frigate/genai/plugins/openai.py | 104 +++++++++++++---------- frigate/test/test_genai_providers.py | 120 ++++++++++++++++++++++++++- 3 files changed, 181 insertions(+), 47 deletions(-) diff --git a/frigate/api/chat.py b/frigate/api/chat.py index 7e1c9aa3da..a16ef3e117 100644 --- a/frigate/api/chat.py +++ b/frigate/api/chat.py @@ -1,5 +1,6 @@ """Chat and LLM tool calling APIs.""" +import asyncio import base64 import json import logging @@ -1653,7 +1654,8 @@ async def chat_completion( f"Calling LLM (iteration {tool_iterations + 1}/{max_iterations}) " 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, tools=llm_tools if llm_tools else None, tool_choice="auto", diff --git a/frigate/genai/plugins/openai.py b/frigate/genai/plugins/openai.py index f69ccbe774..1558a648ef 100644 --- a/frigate/genai/plugins/openai.py +++ b/frigate/genai/plugins/openai.py @@ -1,13 +1,15 @@ """OpenAI Provider for Frigate AI.""" +import asyncio import base64 import json import logging -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Callable from typing import Any from httpx import TimeoutException from openai import OpenAI +from starlette.concurrency import iterate_in_threadpool from frigate.config import GenAIProviderEnum from frigate.genai import GenAIClient, register_genai_provider @@ -15,6 +17,19 @@ 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 @@ -412,57 +427,60 @@ class OpenAIClient(GenAIClient): finish_reason = "stop" usage_stats: dict[str, Any] | None = None - stream = self.provider.chat.completions.create(**request_params) + stream = await _open_stream( + self.provider.chat.completions.create, **request_params + ) - for chunk in stream: - chunk_usage = getattr(chunk, "usage", None) - if chunk_usage is not None: - usage_stats = _stats_from_openai_usage(chunk_usage) + with stream: + async for chunk in iterate_in_threadpool(stream): + chunk_usage = getattr(chunk, "usage", None) + if chunk_usage is not None: + usage_stats = _stats_from_openai_usage(chunk_usage) - if not chunk or not chunk.choices: - continue + if not chunk or not chunk.choices: + continue - choice = chunk.choices[0] - delta = choice.delta + choice = chunk.choices[0] + delta = choice.delta - # Check for finish reason - if choice.finish_reason: - finish_reason = choice.finish_reason + # Check for finish reason + if choice.finish_reason: + finish_reason = choice.finish_reason - # Extract reasoning deltas (reasoning_content or reasoning, - # depending on the server) - reasoning_delta = getattr(delta, "reasoning_content", None) or getattr( - delta, "reasoning", None - ) - if reasoning_delta: - reasoning_parts.append(reasoning_delta) - yield ("reasoning_delta", reasoning_delta) + # Extract reasoning deltas (reasoning_content or reasoning, + # depending on the server) + reasoning_delta = getattr( + delta, "reasoning_content", None + ) or getattr(delta, "reasoning", None) + if reasoning_delta: + reasoning_parts.append(reasoning_delta) + yield ("reasoning_delta", reasoning_delta) - # Extract content deltas - if delta.content: - content_parts.append(delta.content) - yield ("content_delta", delta.content) + # Extract content deltas + if delta.content: + content_parts.append(delta.content) + yield ("content_delta", delta.content) - # Extract tool calls - if delta.tool_calls: - for tc in delta.tool_calls: - idx = tc.index - fn = tc.function + # Extract tool calls + if delta.tool_calls: + for tc in delta.tool_calls: + idx = tc.index + fn = tc.function - if idx not in tool_calls_by_index: - tool_calls_by_index[idx] = { - "id": tc.id or "", - "name": fn.name if fn and fn.name else "", - "arguments": "", - } + if idx not in tool_calls_by_index: + tool_calls_by_index[idx] = { + "id": tc.id or "", + "name": fn.name if fn and fn.name else "", + "arguments": "", + } - t = tool_calls_by_index[idx] - if tc.id: - t["id"] = tc.id - if fn and fn.name: - t["name"] = fn.name - if fn and fn.arguments: - t["arguments"] += fn.arguments + t = tool_calls_by_index[idx] + if tc.id: + t["id"] = tc.id + if fn and fn.name: + t["name"] = fn.name + if fn and fn.arguments: + t["arguments"] += fn.arguments # Build final message full_content = "".join(content_parts).strip() or None diff --git a/frigate/test/test_genai_providers.py b/frigate/test/test_genai_providers.py index fab2e98e1b..2d0aae6112 100644 --- a/frigate/test/test_genai_providers.py +++ b/frigate/test/test_genai_providers.py @@ -15,6 +15,7 @@ crashing message conversion. import asyncio import base64 import json +import time import unittest from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -145,6 +146,23 @@ 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.""" + + def __init__(self, chunks): + self._chunks = iter(chunks) + self.close = MagicMock() + + def __iter__(self): + return self._chunks + + def __enter__(self): + return self + + def __exit__(self, *exc): + self.close() + + class TestOpenAIProvider(unittest.TestCase): def _client(self): return _make_client( @@ -163,7 +181,9 @@ class TestOpenAIProvider(unittest.TestCase): _openai_chunk(finish_reason="tool_calls"), ] client = self._client() - client.provider.chat.completions.create = MagicMock(return_value=iter(chunks)) + client.provider.chat.completions.create = MagicMock( + return_value=_FakeOpenAIStream(chunks) + ) final = _final_message(_collect(client, SIMPLE_MESSAGES)) self.assertEqual(final["finish_reason"], "tool_calls") @@ -178,7 +198,9 @@ class TestOpenAIProvider(unittest.TestCase): _openai_chunk(finish_reason="stop"), ] client = self._client() - client.provider.chat.completions.create = MagicMock(return_value=iter(chunks)) + client.provider.chat.completions.create = MagicMock( + return_value=_FakeOpenAIStream(chunks) + ) events = _collect(client, SIMPLE_MESSAGES) deltas = [v for (k, v) in events if k == "content_delta"] @@ -188,12 +210,104 @@ class TestOpenAIProvider(unittest.TestCase): def test_multimodal_message_does_not_crash(self): client = self._client() client.provider.chat.completions.create = MagicMock( - return_value=iter([_openai_chunk(content="ok", finish_reason="stop")]) + return_value=_FakeOpenAIStream( + [_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()) + ) + + async def _run(): + ticks = 0 + done = asyncio.Event() + + async def ticker(): + nonlocal ticks + while not done.is_set(): + ticks += 1 + await asyncio.sleep(0.01) + + 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(): + 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) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + await asyncio.sleep(0.3) + + asyncio.run(_run()) + stream.close.assert_called_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) + + async def _run(): + 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) + task.cancel() + with self.assertRaises(asyncio.CancelledError): + await task + await gen.aclose() + return first + + self.assertEqual(asyncio.run(_run()), ("content_delta", "a")) + stream.close.assert_called_once() + # --------------------------------------------------------------------------- # Gemini