fix(llm): set HTTP timeout on Gemini clients (#903)
* fix(llm): set HTTP timeout on Gemini clients (#785) * fix(embedding): set HTTP timeout on Gemini embedding client (#785) Same wedge-class failure as the LLM client: a stalled Gemini embedding socket hangs the in-process reconciler, which shares the deriver worker's uvloop event loop. Apply the same 10-minute timeout here, in lockstep with src/llm/registry.py's _build_gemini_http_options. * style(test): drop extra blank line in test_registry imports
This commit is contained in:
parent
d815c8b8dc
commit
5c32bd10ec
|
|
@ -182,10 +182,13 @@ class _EmbeddingClient:
|
|||
if self.transport == "gemini":
|
||||
if not config.api_key:
|
||||
raise ValueError("Gemini API key is required")
|
||||
http_options = (
|
||||
genai_types.HttpOptions(base_url=config.base_url)
|
||||
if config.base_url
|
||||
else None
|
||||
# 10-minute HTTP timeout, in lockstep with the LLM registry's Gemini
|
||||
# client (`src/llm/registry.py:_build_gemini_http_options`). Without
|
||||
# this, a stalled Gemini embedding socket wedges the deriver worker
|
||||
# exactly the way #785 describes for the LLM client.
|
||||
http_options = genai_types.HttpOptions(
|
||||
base_url=config.base_url,
|
||||
timeout=600_000,
|
||||
)
|
||||
self.client: genai.Client | AsyncOpenAI = genai.Client(
|
||||
api_key=config.api_key,
|
||||
|
|
|
|||
|
|
@ -32,6 +32,14 @@ from .history_adapters import (
|
|||
)
|
||||
from .types import ProviderClient
|
||||
|
||||
# Default client-level HTTP timeouts. Anthropic accepts seconds (float);
|
||||
# google-genai's HttpOptions.timeout is an int in milliseconds, so the Gemini
|
||||
# value is kept separately. Both default to 10 minutes to match the existing
|
||||
# Anthropic behavior — long enough for slow streamed responses, short enough
|
||||
# that a stalled socket can no longer wedge the deriver worker (see #785).
|
||||
_ANTHROPIC_TIMEOUT_S = 600.0
|
||||
_GEMINI_TIMEOUT_MS = 600_000
|
||||
|
||||
# Client-level ``default_headers`` applied to OpenAI-compatible clients, keyed by
|
||||
# base-URL prefix. Currently only OpenRouter, which uses them for app attribution
|
||||
# (https://openrouter.ai/docs/app-attribution); add a prefix here to tag another
|
||||
|
|
@ -54,13 +62,28 @@ def _default_headers_for(base_url: str | None) -> dict[str, str]:
|
|||
return {}
|
||||
|
||||
|
||||
def _build_gemini_http_options(base_url: str | None) -> genai_types.HttpOptions:
|
||||
"""Build Gemini ``HttpOptions`` carrying a default HTTP timeout.
|
||||
|
||||
google-genai's ``HttpOptions.timeout`` is an int in milliseconds. A stalled
|
||||
Gemini socket without this value wedges the entire deriver process because
|
||||
all deriver workers share one uvloop event loop (see #785). Keep the
|
||||
timeout even when no ``base_url`` is configured — that's the path the
|
||||
default ``get_gemini_client`` takes and it's the one that was hanging.
|
||||
"""
|
||||
return genai_types.HttpOptions(
|
||||
base_url=base_url,
|
||||
timeout=_GEMINI_TIMEOUT_MS,
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_anthropic_client() -> AsyncAnthropic:
|
||||
"""Default Anthropic client built from settings.LLM.ANTHROPIC_API_KEY."""
|
||||
return AsyncAnthropic(
|
||||
api_key=settings.LLM.ANTHROPIC_API_KEY,
|
||||
base_url=settings.LLM.ANTHROPIC_BASE_URL,
|
||||
timeout=600.0,
|
||||
timeout=_ANTHROPIC_TIMEOUT_S,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -77,12 +100,10 @@ def get_openai_client() -> AsyncOpenAI:
|
|||
@lru_cache(maxsize=1)
|
||||
def get_gemini_client() -> genai.Client:
|
||||
"""Default Gemini client built from settings.LLM.GEMINI_API_KEY."""
|
||||
http_options = (
|
||||
genai_types.HttpOptions(base_url=settings.LLM.GEMINI_BASE_URL)
|
||||
if settings.LLM.GEMINI_BASE_URL
|
||||
else None
|
||||
return genai.Client(
|
||||
api_key=settings.LLM.GEMINI_API_KEY,
|
||||
http_options=_build_gemini_http_options(settings.LLM.GEMINI_BASE_URL),
|
||||
)
|
||||
return genai.Client(api_key=settings.LLM.GEMINI_API_KEY, http_options=http_options)
|
||||
|
||||
|
||||
# Bounded cache — in practice the (base_url, api_key) key space is small
|
||||
|
|
@ -105,7 +126,9 @@ def get_anthropic_override_client(
|
|||
api_key: str | None,
|
||||
) -> AsyncAnthropic:
|
||||
"""Anthropic client for a specific (base_url, api_key) pair. Cached by key."""
|
||||
return AsyncAnthropic(api_key=api_key, base_url=base_url, timeout=600.0)
|
||||
return AsyncAnthropic(
|
||||
api_key=api_key, base_url=base_url, timeout=_ANTHROPIC_TIMEOUT_S
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=128)
|
||||
|
|
@ -113,8 +136,10 @@ def get_gemini_override_client(
|
|||
base_url: str | None, api_key: str | None
|
||||
) -> genai.Client:
|
||||
"""Gemini client for a specific (base_url, api_key) pair. Cached by key."""
|
||||
http_options = genai_types.HttpOptions(base_url=base_url) if base_url else None
|
||||
return genai.Client(api_key=api_key, http_options=http_options)
|
||||
return genai.Client(
|
||||
api_key=api_key,
|
||||
http_options=_build_gemini_http_options(base_url),
|
||||
)
|
||||
|
||||
|
||||
# Module-level default-client registry, populated at import time. Tests patch
|
||||
|
|
@ -125,7 +150,7 @@ if settings.LLM.ANTHROPIC_API_KEY:
|
|||
CLIENTS["anthropic"] = AsyncAnthropic(
|
||||
api_key=settings.LLM.ANTHROPIC_API_KEY,
|
||||
base_url=settings.LLM.ANTHROPIC_BASE_URL,
|
||||
timeout=600.0,
|
||||
timeout=_ANTHROPIC_TIMEOUT_S,
|
||||
)
|
||||
|
||||
if settings.LLM.OPENAI_API_KEY:
|
||||
|
|
@ -136,14 +161,9 @@ if settings.LLM.OPENAI_API_KEY:
|
|||
)
|
||||
|
||||
if settings.LLM.GEMINI_API_KEY:
|
||||
http_options = (
|
||||
genai_types.HttpOptions(base_url=settings.LLM.GEMINI_BASE_URL)
|
||||
if settings.LLM.GEMINI_BASE_URL
|
||||
else None
|
||||
)
|
||||
CLIENTS["gemini"] = genai.Client(
|
||||
api_key=settings.LLM.GEMINI_API_KEY,
|
||||
http_options=http_options,
|
||||
http_options=_build_gemini_http_options(settings.LLM.GEMINI_BASE_URL),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from types import SimpleNamespace
|
||||
from typing import Any
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
|
|
@ -141,6 +141,12 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
|
|||
embedding = await client.embed("hello world")
|
||||
|
||||
assert embedding == [0.2] * 12
|
||||
# 10-minute HTTP timeout, in lockstep with the LLM registry's Gemini client
|
||||
# (see #785). Without this, a stalled Gemini embedding socket wedges the
|
||||
# deriver worker — the same failure mode the LLM fix addresses.
|
||||
gemini_client = cast(Any, client.client)
|
||||
assert gemini_client.http_options.base_url == "https://gemini-proxy.example/v1beta"
|
||||
assert gemini_client.http_options.timeout == 600_000
|
||||
assert calls == [
|
||||
{
|
||||
"model": "gemini-embedding-001",
|
||||
|
|
@ -150,6 +156,49 @@ async def test_gemini_embedding_client_uses_output_dimensionality(
|
|||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gemini_embedding_client_keeps_timeout_without_base_url(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""No-base-url Gemini embedding client must still carry an HTTP timeout."""
|
||||
|
||||
class FakeGeminiModels:
|
||||
async def embed_content(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
contents: str,
|
||||
config: dict[str, Any],
|
||||
) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
embeddings=[SimpleNamespace(values=[0.1] * 8)],
|
||||
)
|
||||
|
||||
class FakeGeminiClient:
|
||||
def __init__(self, *, api_key: str | None, http_options: Any) -> None:
|
||||
self.api_key = api_key
|
||||
self.http_options = http_options
|
||||
self.aio = SimpleNamespace(models=FakeGeminiModels())
|
||||
|
||||
monkeypatch.setattr("src.embedding_client.genai.Client", FakeGeminiClient)
|
||||
|
||||
client = _EmbeddingClient(
|
||||
EmbeddingModelConfig(
|
||||
transport="gemini",
|
||||
model="gemini-embedding-001",
|
||||
api_key="gemini-key",
|
||||
),
|
||||
vector_dimensions=8,
|
||||
max_input_tokens=4096,
|
||||
max_tokens_per_request=300_000,
|
||||
send_dimensions=False,
|
||||
)
|
||||
|
||||
gemini_client = cast(Any, client.client)
|
||||
assert gemini_client.http_options.base_url is None
|
||||
assert gemini_client.http_options.timeout == 600_000
|
||||
|
||||
|
||||
def _build_openai_client(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
*,
|
||||
|
|
|
|||
|
|
@ -1,22 +1,139 @@
|
|||
"""Tests for src.llm.registry helpers."""
|
||||
"""Tests for the provider-client registry in src/llm/registry.py.
|
||||
|
||||
Locks the HTTP-timeout behavior added for the Gemini transport (#785) and
|
||||
pins the existing 600s timeout on the Anthropic clients so regressions on
|
||||
either side are caught at unit-test time.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from src.llm.registry import _default_headers_for # pyright: ignore[reportPrivateUsage]
|
||||
from collections.abc import Iterator
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from google.genai import types as genai_types
|
||||
|
||||
from src import config as app_config
|
||||
from src.llm import registry as registry_module
|
||||
|
||||
# Gemini's HttpOptions.timeout is an int in milliseconds; keep it in lockstep
|
||||
# with the Anthropic client's 600s timeout to match the rest of the registry.
|
||||
_GEMINI_TIMEOUT_MS = 600_000
|
||||
_ANTHROPIC_TIMEOUT_S = 600.0
|
||||
|
||||
|
||||
def test_default_headers_for_openrouter_base_url() -> None:
|
||||
"""OpenRouter base URLs get the app-attribution headers."""
|
||||
headers = _default_headers_for("https://openrouter.ai/api/v1")
|
||||
assert headers["HTTP-Referer"] == "https://honcho.dev"
|
||||
assert headers["X-Openrouter-Title"] == "Honcho"
|
||||
@pytest.fixture(autouse=True)
|
||||
def _patch_settings(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
|
||||
"""Default the LLM settings so the registry reads valid values."""
|
||||
monkeypatch.setenv("PYTHON_DOTENV_DISABLED", "1")
|
||||
monkeypatch.setenv("ANTHROPIC_API_KEY", "test-anthropic-key")
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "test-gemini-key")
|
||||
yield
|
||||
|
||||
|
||||
def test_default_headers_for_non_openrouter_base_url() -> None:
|
||||
"""Other OpenAI-compatible providers get no extra headers."""
|
||||
assert _default_headers_for("https://api.openai.com/v1") == {}
|
||||
@pytest.fixture
|
||||
def _fresh_lru_caches() -> Iterator[None]:
|
||||
"""Drop lru_cache state so each test exercises a fresh client build."""
|
||||
registry_module.get_anthropic_client.cache_clear()
|
||||
registry_module.get_gemini_client.cache_clear()
|
||||
registry_module.get_anthropic_override_client.cache_clear()
|
||||
registry_module.get_gemini_override_client.cache_clear()
|
||||
yield
|
||||
registry_module.get_anthropic_client.cache_clear()
|
||||
registry_module.get_gemini_client.cache_clear()
|
||||
registry_module.get_anthropic_override_client.cache_clear()
|
||||
registry_module.get_gemini_override_client.cache_clear()
|
||||
|
||||
|
||||
def test_default_headers_for_none_base_url() -> None:
|
||||
"""A missing base URL (default OpenAI) gets no extra headers."""
|
||||
assert _default_headers_for(None) == {}
|
||||
def test_get_gemini_client_sets_http_timeout(
|
||||
_fresh_lru_caches: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Default Gemini client must carry an HttpOptions timeout, not None."""
|
||||
monkeypatch.setattr(app_config.settings.LLM, "GEMINI_BASE_URL", None)
|
||||
|
||||
with patch("src.llm.registry.genai.Client") as mock_client:
|
||||
registry_module.get_gemini_client()
|
||||
|
||||
assert mock_client.call_count == 1
|
||||
http_options = mock_client.call_args.kwargs["http_options"]
|
||||
assert isinstance(http_options, genai_types.HttpOptions)
|
||||
assert http_options.timeout == _GEMINI_TIMEOUT_MS
|
||||
|
||||
|
||||
def test_get_gemini_client_preserves_custom_base_url(
|
||||
_fresh_lru_caches: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Base URL and timeout must coexist on the default Gemini client."""
|
||||
monkeypatch.setattr(
|
||||
app_config.settings.LLM, "GEMINI_BASE_URL", "https://gemini-proxy.example.com"
|
||||
)
|
||||
|
||||
with patch("src.llm.registry.genai.Client") as mock_client:
|
||||
registry_module.get_gemini_client()
|
||||
|
||||
http_options = mock_client.call_args.kwargs["http_options"]
|
||||
assert isinstance(http_options, genai_types.HttpOptions)
|
||||
assert http_options.base_url == "https://gemini-proxy.example.com"
|
||||
assert http_options.timeout == _GEMINI_TIMEOUT_MS
|
||||
|
||||
|
||||
def test_get_gemini_override_client_sets_http_timeout(
|
||||
_fresh_lru_caches: None,
|
||||
) -> None:
|
||||
"""Override Gemini client must also carry a timeout."""
|
||||
with patch("src.llm.registry.genai.Client") as mock_client:
|
||||
registry_module.get_gemini_override_client(
|
||||
"https://gemini-proxy.example.com", "sk-override"
|
||||
)
|
||||
|
||||
http_options = mock_client.call_args.kwargs["http_options"]
|
||||
assert isinstance(http_options, genai_types.HttpOptions)
|
||||
assert http_options.base_url == "https://gemini-proxy.example.com"
|
||||
assert http_options.timeout == _GEMINI_TIMEOUT_MS
|
||||
|
||||
|
||||
def test_get_gemini_override_client_handles_missing_base_url(
|
||||
_fresh_lru_caches: None,
|
||||
) -> None:
|
||||
"""Override Gemini client with no base URL still carries a timeout."""
|
||||
with patch("src.llm.registry.genai.Client") as mock_client:
|
||||
registry_module.get_gemini_override_client(None, "sk-override")
|
||||
|
||||
http_options = mock_client.call_args.kwargs["http_options"]
|
||||
assert isinstance(http_options, genai_types.HttpOptions)
|
||||
assert http_options.timeout == _GEMINI_TIMEOUT_MS
|
||||
|
||||
|
||||
def test_get_anthropic_client_keeps_600s_timeout(
|
||||
_fresh_lru_caches: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""Anthropic timeout is the established behavior — lock it."""
|
||||
monkeypatch.setattr(app_config.settings.LLM, "ANTHROPIC_BASE_URL", None)
|
||||
|
||||
with patch("src.llm.registry.AsyncAnthropic") as mock_anthropic:
|
||||
registry_module.get_anthropic_client()
|
||||
|
||||
assert mock_anthropic.call_args.kwargs["timeout"] == _ANTHROPIC_TIMEOUT_S
|
||||
|
||||
|
||||
def test_get_anthropic_override_client_keeps_600s_timeout(
|
||||
_fresh_lru_caches: None,
|
||||
) -> None:
|
||||
"""Override Anthropic client also keeps the 600s timeout."""
|
||||
with patch("src.llm.registry.AsyncAnthropic") as mock_anthropic:
|
||||
registry_module.get_anthropic_override_client(None, "sk-override")
|
||||
|
||||
assert mock_anthropic.call_args.kwargs["timeout"] == _ANTHROPIC_TIMEOUT_S
|
||||
|
||||
|
||||
def test_gemini_http_options_builder_applies_timeout() -> None:
|
||||
"""The shared helper must always set a timeout, even with no base_url."""
|
||||
options = registry_module._build_gemini_http_options(None)
|
||||
assert isinstance(options, genai_types.HttpOptions)
|
||||
assert options.timeout == _GEMINI_TIMEOUT_MS
|
||||
assert options.base_url is None
|
||||
|
||||
options = registry_module._build_gemini_http_options("https://example.com")
|
||||
assert isinstance(options, genai_types.HttpOptions)
|
||||
assert options.timeout == _GEMINI_TIMEOUT_MS
|
||||
assert options.base_url == "https://example.com"
|
||||
|
|
|
|||
Loading…
Reference in New Issue