mirror of
https://github.com/headroomlabs-ai/headroom.git
synced 2026-08-27 14:17:10 -04:00
fix(gemini): resolve Google model capabilities through ModelRegistry (#1276)
## Description Google model capability lookup was still tied to static provider tables for support checks and context limits. That made plausible future Gemini model ids fail token counting or context lookup even when they clearly belonged to the Google provider family. This change adds a tolerant `ModelRegistry.resolve()` runtime lookup path and routes the Google provider through it. Exact built-in registry matches still win first, LiteLLM pricing metadata can supply live limits when available, and provider-scoped family fallbacks cover future Gemini ids without letting Google claim unrelated models. ## Type of Change - [x] Bug fix (non-breaking change that fixes an issue) - [ ] New feature (non-breaking change that adds functionality) - [ ] Breaking change (fix or feature that would cause existing functionality to change) - [ ] Documentation update - [ ] Performance improvement - [ ] Code refactoring (no functional changes) ## Changes Made - Added `ModelRegistry.resolve()` as a tolerant runtime capability resolver. - Added provider-scoped Google/Gemini family fallbacks for plausible future model ids. - Added support for LiteLLM-style `gemini/gemini-...` model ids in provider inference and family fallback matching. - Updated `GoogleProvider.supports_model()` and `GoogleProvider.get_context_limit()` to use the shared model registry path. - Added regression tests for future Gemini ids, legacy Gemini context limits, and unrelated model rejection. ## Testing - [x] Unit tests pass (`pytest`) - [x] Linting passes (`ruff check .`) - [ ] Type checking passes (`mypy headroom`) - [x] New tests added for new functionality - [x] Manual testing performed ### Test Output ```text uv run --no-project --with pytest --with opentelemetry-api --with pydantic --with tiktoken --with litellm --with click --with rich python -B -m pytest tests/test_provider_model_fallback.py tests/test_models.py 65 passed uv run --no-project --with ruff ruff check headroom/models/registry.py headroom/providers/google.py tests/test_provider_model_fallback.py tests/test_models.py All checks passed! uv run --no-project --with ruff ruff format --check headroom/models/registry.py headroom/providers/google.py tests/test_provider_model_fallback.py tests/test_models.py 4 files already formatted ``` ## Real Behavior Proof - Environment: macOS arm64 local checkout, Python 3.13 virtualenv for editable install; deployed smoke test in a Cloud Run staging service using an earlier commit from this fork branch before the review follow-up. - Exact command / steps: installed `headroom-ai[langchain]` from the fork branch in the staging service, triggered long-context requests that activate Headroom's LangChain compression path, then checked Cloud Run logs after 2026-06-22 12:20 Europe/Paris. - Observed result: Headroom initialized successfully, compressed conversation memory (`23255 -> 5618 chars`), and no logs matched the previous model-resolution failure signatures (`not recognized as a Google model`, `Unknown context limit`). - Not tested: staging was not rerun after the `gemini/gemini-...` review follow-up; that prefix path is covered by local regression tests. Full repository `uv run pytest` on local macOS is currently blocked by a native `maturin`/`esaxx-rs` compile failure (`fatal error: 'cstdint' file not found`). Type checking was not run. ## Review Readiness - [x] I have performed a self-review - [x] This PR is ready for human review ## Checklist - [x] My code follows the project's style guidelines - [x] I have performed a self-review of my code - [x] I have commented my code, particularly in hard-to-understand areas - [ ] I have made corresponding changes to the documentation - [x] My changes generate no new warnings - [x] I have added tests that prove my fix is effective or that my feature works - [ ] New and existing unit tests pass locally with my changes - [x] I have updated the CHANGELOG.md if applicable ## Screenshots (if applicable) N/A ## Additional Notes - Documentation changes are not included because this is a runtime compatibility fix with no public API or user-facing configuration change. - Full local test execution should be retried in CI or a Linux environment where the native Rust extension build is healthy. Co-authored-by: Julien Guarino <julien.guarino@fashiondata.io>
This commit is contained in:
julienguarino
GitHub
parent
c632023cc1
commit
17ecad9d89
5 changed files with 259 additions and 48 deletions
|
|
@ -68,6 +68,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
|
|||
* **proxy:** route Codex OAuth image generation and edit requests through the ChatGPT Codex image backend, while preserving OpenAI API-key image passthrough ([#1215](https://github.com/chopratejas/headroom/pull/1215)).
|
||||
* **wrap (codex):** keep RTK guidance in the global Codex `AGENTS.md` instead of modifying the shared project `AGENTS.md` ([#1235](https://github.com/chopratejas/headroom/issues/1235)).
|
||||
* **subscription:** run the transcript token-window scan off the event loop (`asyncio.to_thread`). The subscription tracker's poll loop scanned every `~/.claude/projects/**/*.jsonl` transcript and `json.loads`'d each line inline on the proxy's single asyncio event loop; on large or long-running sessions this took seconds and froze `/health` and every in-flight proxied request — a periodic "wedge" recurring on the poll interval. The scan now runs in a worker thread so the loop stays responsive.
|
||||
* **gemini:** resolve future Gemini model capabilities through the shared model registry so token counting and context lookup no longer reject new Gemini families.
|
||||
* **proxy:** enable SSO credential resolution in the native Bedrock route via the `aws-config` `sso` feature flag, making the credential chain match what `docs/bedrock.md` already documented ([#999](https://github.com/chopratejas/headroom/pull/999)).
|
||||
* **proxy:** route native Bedrock `/model/{id}/converse` requests to the upstream Converse endpoint instead of the hard-coded `/invoke` action — the non-streaming handler now resolves the action from the inbound path, matching the streaming handler ([#999](https://github.com/chopratejas/headroom/pull/999)).
|
||||
* **proxy:** preserve byte-faithful `/v1/messages` forwarding when Anthropic tool arrays are already canonical, and only canonicalize-and-mutate tool lists when sorting changes ordering ([#1042](https://github.com/chopratejas/headroom/issues/1042)).
|
||||
|
|
|
|||
|
|
@ -487,6 +487,108 @@ for model_name, info in _MODELS.items():
|
|||
_ALIASES[alias.lower()] = model_name
|
||||
|
||||
|
||||
_PROVIDER_TOKENIZER_BACKENDS = {
|
||||
"anthropic": "anthropic",
|
||||
"google": "google",
|
||||
"openai": "tiktoken",
|
||||
}
|
||||
|
||||
_PROVIDER_FAMILY_DEFAULTS: dict[str, tuple[tuple[str, dict[str, Any]], ...]] = {
|
||||
"google": (
|
||||
(
|
||||
"gemini-1.0",
|
||||
{
|
||||
"context_window": 32768,
|
||||
"max_output_tokens": 4096,
|
||||
"supports_vision": False,
|
||||
"tokenizer_backend": "google",
|
||||
"notes": "Google Gemini 1.0 family fallback",
|
||||
},
|
||||
),
|
||||
(
|
||||
"gemini-pro",
|
||||
{
|
||||
"context_window": 32768,
|
||||
"max_output_tokens": 4096,
|
||||
"supports_vision": False,
|
||||
"tokenizer_backend": "google",
|
||||
"notes": "Legacy Google Gemini Pro fallback",
|
||||
},
|
||||
),
|
||||
(
|
||||
"gemini-1.5-pro",
|
||||
{
|
||||
"context_window": 2000000,
|
||||
"max_output_tokens": 8192,
|
||||
"supports_vision": True,
|
||||
"tokenizer_backend": "google",
|
||||
"notes": "Google Gemini 1.5 Pro family fallback",
|
||||
},
|
||||
),
|
||||
(
|
||||
"gemini-1.5-flash",
|
||||
{
|
||||
"context_window": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"supports_vision": True,
|
||||
"tokenizer_backend": "google",
|
||||
"notes": "Google Gemini 1.5 Flash family fallback",
|
||||
},
|
||||
),
|
||||
(
|
||||
"gemini-2",
|
||||
{
|
||||
"context_window": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"supports_vision": True,
|
||||
"tokenizer_backend": "google",
|
||||
"notes": "Google Gemini 2 family fallback",
|
||||
},
|
||||
),
|
||||
(
|
||||
"gemini-",
|
||||
{
|
||||
"context_window": 1000000,
|
||||
"max_output_tokens": 8192,
|
||||
"supports_vision": True,
|
||||
"tokenizer_backend": "google",
|
||||
"notes": "Future Google Gemini family fallback",
|
||||
},
|
||||
),
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _infer_provider(model: str) -> str | None:
|
||||
"""Infer a provider from common model id prefixes."""
|
||||
model_lower = model.lower()
|
||||
if model_lower.startswith(("gemini-", "gemini/gemini-", "google/gemini-")):
|
||||
return "google"
|
||||
if model_lower.startswith(("claude", "anthropic/claude")):
|
||||
return "anthropic"
|
||||
if model_lower.startswith(("gpt-", "o1", "o3", "o4", "openai/gpt-")):
|
||||
return "openai"
|
||||
return None
|
||||
|
||||
|
||||
def _unprefixed_model_id(model: str) -> str:
|
||||
"""Drop a common LiteLLM provider prefix before family matching."""
|
||||
model_lower = model.lower()
|
||||
for prefix in ("anthropic/", "gemini/", "google/", "openai/"):
|
||||
if model_lower.startswith(prefix):
|
||||
return model_lower[len(prefix) :]
|
||||
return model_lower
|
||||
|
||||
|
||||
def _family_fallback(model: str, provider: str) -> ModelInfo | None:
|
||||
"""Return a provider-scoped family fallback for plausible future models."""
|
||||
model_lower = _unprefixed_model_id(model)
|
||||
for prefix, defaults in _PROVIDER_FAMILY_DEFAULTS.get(provider, ()):
|
||||
if model_lower.startswith(prefix):
|
||||
return ModelInfo(name=model, provider=provider, **defaults)
|
||||
return None
|
||||
|
||||
|
||||
class ModelRegistry:
|
||||
"""Registry of LLM models and their capabilities.
|
||||
|
||||
|
|
@ -536,6 +638,73 @@ class ModelRegistry:
|
|||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def resolve(
|
||||
cls,
|
||||
model: str,
|
||||
provider: str | None = None,
|
||||
default_context_window: int = 128000,
|
||||
) -> ModelInfo | None:
|
||||
"""Resolve model capabilities for a runtime provider path.
|
||||
|
||||
This is the tolerant runtime counterpart to :meth:`get`: it first
|
||||
checks the built-in registry, then LiteLLM metadata, then
|
||||
provider-scoped family fallbacks. Unknown models for unrelated
|
||||
providers still return ``None`` instead of being claimed globally.
|
||||
|
||||
Args:
|
||||
model: Model identifier from a request/provider.
|
||||
provider: Optional provider hint (for example ``"google"``).
|
||||
default_context_window: Conservative context window when a
|
||||
plausible family fallback has no exact catalog hit.
|
||||
|
||||
Returns:
|
||||
Resolved ModelInfo when the model belongs to the hinted or
|
||||
inferred provider, otherwise None.
|
||||
"""
|
||||
provider_hint = provider.lower() if provider else None
|
||||
|
||||
info = cls.get(model)
|
||||
if info is not None and (provider_hint is None or info.provider == provider_hint):
|
||||
return info
|
||||
|
||||
inferred_provider = _infer_provider(model)
|
||||
if (
|
||||
provider_hint is not None
|
||||
and inferred_provider is not None
|
||||
and inferred_provider != provider_hint
|
||||
):
|
||||
return None
|
||||
|
||||
resolved_provider = provider_hint or inferred_provider
|
||||
if resolved_provider is None:
|
||||
return None
|
||||
|
||||
fallback = _family_fallback(model, resolved_provider)
|
||||
if provider_hint is not None and inferred_provider is None and fallback is None:
|
||||
return None
|
||||
|
||||
pricing = get_model_pricing(model)
|
||||
if pricing is not None:
|
||||
context_window = (
|
||||
pricing.max_input_tokens or pricing.max_tokens or default_context_window
|
||||
)
|
||||
max_output_tokens = pricing.max_output_tokens or 4096
|
||||
return ModelInfo(
|
||||
name=model,
|
||||
provider=resolved_provider,
|
||||
context_window=int(context_window),
|
||||
max_output_tokens=int(max_output_tokens),
|
||||
supports_vision=pricing.supports_vision,
|
||||
tokenizer_backend=_PROVIDER_TOKENIZER_BACKENDS.get(resolved_provider),
|
||||
notes="Resolved from LiteLLM pricing metadata",
|
||||
)
|
||||
|
||||
if fallback is not None:
|
||||
return fallback
|
||||
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def register(
|
||||
cls,
|
||||
|
|
|
|||
|
|
@ -25,6 +25,7 @@ import warnings
|
|||
from datetime import date
|
||||
from typing import Any
|
||||
|
||||
from headroom.models.registry import ModelRegistry
|
||||
from headroom.tokenizers import EstimatingTokenCounter
|
||||
|
||||
from .base import Provider, TokenCounter
|
||||
|
|
@ -32,13 +33,11 @@ from .base import Provider, TokenCounter
|
|||
# Check if litellm is available for pricing/context limit lookups
|
||||
try:
|
||||
import litellm
|
||||
from litellm import get_model_info as litellm_get_model_info
|
||||
|
||||
LITELLM_AVAILABLE = True
|
||||
except ImportError:
|
||||
LITELLM_AVAILABLE = False
|
||||
litellm = None # type: ignore[assignment]
|
||||
litellm_get_model_info = None # type: ignore[assignment]
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -271,15 +270,15 @@ class GoogleProvider(Provider):
|
|||
return "google"
|
||||
|
||||
def supports_model(self, model: str) -> bool:
|
||||
"""Check if model is a known Gemini model."""
|
||||
model_lower = model.lower()
|
||||
if model_lower in _CONTEXT_LIMITS:
|
||||
return True
|
||||
# Check prefix match
|
||||
for prefix in ["gemini-2", "gemini-1.5", "gemini-1.0", "gemini-pro"]:
|
||||
if model_lower.startswith(prefix):
|
||||
return True
|
||||
return False
|
||||
"""Check if this Google provider can handle a Gemini model."""
|
||||
return (
|
||||
ModelRegistry.resolve(
|
||||
model,
|
||||
provider="google",
|
||||
default_context_window=1000000,
|
||||
)
|
||||
is not None
|
||||
)
|
||||
|
||||
def get_token_counter(self, model: str) -> TokenCounter:
|
||||
"""Get token counter for a Gemini model.
|
||||
|
|
@ -296,44 +295,17 @@ class GoogleProvider(Provider):
|
|||
def get_context_limit(self, model: str) -> int:
|
||||
"""Get context limit for a Gemini model.
|
||||
|
||||
Tries LiteLLM first for up-to-date limits, falls back to hardcoded values.
|
||||
Note: Gemini 1.5 Pro has 2M token context!
|
||||
Runtime capability lookup goes through the shared ModelRegistry so
|
||||
future Gemini families can use catalog or family fallback metadata
|
||||
instead of hard-failing on the provider's static table.
|
||||
"""
|
||||
model_lower = model.lower()
|
||||
|
||||
# Try LiteLLM first for up-to-date context limits
|
||||
if LITELLM_AVAILABLE and litellm_get_model_info is not None:
|
||||
# Try different model name formats that LiteLLM might recognize
|
||||
model_variants = [
|
||||
f"gemini/{model_lower}", # gemini/gemini-1.5-pro
|
||||
model_lower, # gemini-1.5-pro
|
||||
]
|
||||
for variant in model_variants:
|
||||
try:
|
||||
info = litellm_get_model_info(variant)
|
||||
if info:
|
||||
if "max_input_tokens" in info and info["max_input_tokens"]:
|
||||
return info["max_input_tokens"]
|
||||
if "max_tokens" in info and info["max_tokens"]:
|
||||
return info["max_tokens"]
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
# Fallback to hardcoded limits
|
||||
# Direct match
|
||||
if model_lower in _CONTEXT_LIMITS:
|
||||
return _CONTEXT_LIMITS[model_lower]
|
||||
|
||||
# Prefix match
|
||||
for prefix, limit in [
|
||||
("gemini-2.0", 1000000),
|
||||
("gemini-1.5-pro", 2000000),
|
||||
("gemini-1.5-flash", 1000000),
|
||||
("gemini-1.0", 32768),
|
||||
("gemini-pro", 32768),
|
||||
]:
|
||||
if model_lower.startswith(prefix):
|
||||
return limit
|
||||
info = ModelRegistry.resolve(
|
||||
model,
|
||||
provider="google",
|
||||
default_context_window=1000000,
|
||||
)
|
||||
if info is not None:
|
||||
return info.context_window
|
||||
|
||||
raise ValueError(
|
||||
f"Unknown context limit for model '{model}'. "
|
||||
|
|
|
|||
|
|
@ -2,6 +2,8 @@
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from headroom.models import (
|
||||
|
|
@ -88,6 +90,31 @@ class TestModelRegistry:
|
|||
assert info is not None
|
||||
assert info.name == "gpt-4o"
|
||||
|
||||
def test_resolve_future_google_family_fallback(self):
|
||||
"""Resolve should return provider-scoped fallbacks for plausible future models."""
|
||||
with patch("headroom.models.registry.get_model_pricing", return_value=None):
|
||||
info = ModelRegistry.resolve("gemini-3-pro-preview", provider="google")
|
||||
|
||||
assert info is not None
|
||||
assert info.provider == "google"
|
||||
assert info.context_window == 1000000
|
||||
assert info.tokenizer_backend == "google"
|
||||
|
||||
def test_resolve_google_litellm_prefixed_family_fallback(self):
|
||||
"""Resolve should support LiteLLM-style Gemini provider prefixes."""
|
||||
with patch("headroom.models.registry.get_model_pricing", return_value=None):
|
||||
info = ModelRegistry.resolve("gemini/gemini-3-pro-preview", provider="google")
|
||||
|
||||
assert info is not None
|
||||
assert info.provider == "google"
|
||||
assert info.context_window == 1000000
|
||||
assert info.tokenizer_backend == "google"
|
||||
|
||||
def test_resolve_does_not_claim_unrelated_models_for_google(self):
|
||||
"""Provider-scoped resolution should not mask unrelated model catalogs."""
|
||||
assert ModelRegistry.resolve("not-a-google-model", provider="google") is None
|
||||
assert ModelRegistry.resolve("gpt-4o", provider="google") is None
|
||||
|
||||
def test_register_custom_model(self):
|
||||
"""Test registering custom model."""
|
||||
info = ModelRegistry.register(
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ from headroom.providers.anthropic import (
|
|||
from headroom.providers.anthropic import (
|
||||
_load_custom_model_config as anthropic_load_config,
|
||||
)
|
||||
from headroom.providers.google import GeminiTokenCounter, GoogleProvider
|
||||
from headroom.providers.openai import (
|
||||
OpenAIProvider,
|
||||
_infer_model_family,
|
||||
|
|
@ -24,6 +25,47 @@ from headroom.providers.openai import (
|
|||
)
|
||||
|
||||
|
||||
class TestGoogleModelFallback:
|
||||
"""Tests for Google provider model fallback."""
|
||||
|
||||
def test_future_gemini_model_uses_registry_family_fallback(self):
|
||||
"""Future Gemini models should not hard-fail token counting."""
|
||||
provider = GoogleProvider()
|
||||
|
||||
with patch("headroom.models.registry.get_model_pricing", return_value=None):
|
||||
assert provider.supports_model("gemini-3-pro-preview")
|
||||
assert provider.get_context_limit("gemini-3-pro-preview") == 1000000
|
||||
assert isinstance(
|
||||
provider.get_token_counter("gemini-3-pro-preview"),
|
||||
GeminiTokenCounter,
|
||||
)
|
||||
|
||||
def test_litellm_prefixed_gemini_model_uses_registry_family_fallback(self):
|
||||
"""LiteLLM-style Gemini ids should resolve through the Google provider."""
|
||||
provider = GoogleProvider()
|
||||
|
||||
with patch("headroom.models.registry.get_model_pricing", return_value=None):
|
||||
assert provider.supports_model("gemini/gemini-3-pro-preview")
|
||||
assert provider.get_context_limit("gemini/gemini-3-pro-preview") == 1000000
|
||||
|
||||
def test_google_legacy_context_limits_are_preserved(self):
|
||||
"""Moving lookup through ModelRegistry must keep legacy Gemini limits."""
|
||||
provider = GoogleProvider()
|
||||
|
||||
with patch("headroom.models.registry.get_model_pricing", return_value=None):
|
||||
assert provider.get_context_limit("gemini-1.5-pro-latest") == 2000000
|
||||
assert provider.get_context_limit("gemini-1.0-pro") == 32768
|
||||
|
||||
def test_unknown_non_gemini_model_still_rejected(self):
|
||||
"""The Google provider should not claim unrelated unknown models."""
|
||||
provider = GoogleProvider()
|
||||
|
||||
assert not provider.supports_model("not-a-google-model")
|
||||
assert not provider.supports_model("gpt-4o")
|
||||
with pytest.raises(ValueError):
|
||||
provider.get_token_counter("not-a-google-model")
|
||||
|
||||
|
||||
class TestAnthropicModelFallback:
|
||||
"""Tests for Anthropic provider model fallback."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue