From 99de9d8e8b5fe0601496307e40c2d0897502c818 Mon Sep 17 00:00:00 2001 From: charchit7 Date: Mon, 22 Jun 2026 13:41:21 +0000 Subject: [PATCH] adds support for gemini API --- agent/core/llm_params.py | 32 +++++++++++++++++++ agent/core/model_switcher.py | 13 +++++++- tests/unit/test_llm_params.py | 59 +++++++++++++++++++++++++++++++++++ 3 files changed, 103 insertions(+), 1 deletion(-) diff --git a/agent/core/llm_params.py b/agent/core/llm_params.py index d2f821c2..7b441c02 100644 --- a/agent/core/llm_params.py +++ b/agent/core/llm_params.py @@ -32,6 +32,15 @@ def _resolve_hf_router_token(session_hf_token: str | None = None) -> str | None: # ``extra_body`` field. The probe cascade walks down when a provider rejects # an accepted-looking value, so this stays intentionally small and generic. _HF_EFFORTS = {"low", "medium", "high"} +# Gemini 2.5+ thinking models. LiteLLM maps reasoning_effort → thinking +# budget for both the Google AI Studio (``gemini/``) and Vertex AI +# (``vertex_ai/``) routes; "disable" turns thinking off but we model that +# as "no effort" (None) rather than an effort level. +_GEMINI_EFFORTS = {"low", "medium", "high"} + +# Prefixes routed directly through LiteLLM's Google Gemini adapters rather +# than the HuggingFace router catch-all. +_GEMINI_PREFIXES = ("gemini/", "vertex_ai/") def _hf_router_effort_level(reasoning_effort: str) -> str: @@ -96,6 +105,16 @@ def _resolve_llm_params( """ Build LiteLLM kwargs for a given model id. + • ``gemini/`` / ``vertex_ai/`` — Google Gemini via the AI + Studio API (``GEMINI_API_KEY``) or Vertex AI (GCP creds from + ``VERTEX_PROJECT`` / ``VERTEX_LOCATION`` or application-default + credentials). Routed directly through LiteLLM's Google adapters rather + than the HF Router (which doesn't serve Gemini). ``reasoning_effort`` + is forwarded as a top-level kwarg; LiteLLM's Gemini adapter translates + it into the thinking budget for 2.5+ thinking models. "minimal" + normalizes to "low". Models that don't support thinking reject it and + the probe cascade drops it. + • ``ollama/``, ``vllm/``, ``lm_studio/``, and ``llamacpp/`` — local OpenAI-compatible endpoints. The id prefix selects a configurable localhost base URL, and the model suffix is sent @@ -129,6 +148,19 @@ def _resolve_llm_params( if local_model_provider(normalized_model) is not None: return _resolve_local_model_params(normalized_model, reasoning_effort, strict) + if normalized_model.startswith(_GEMINI_PREFIXES): + params = {"model": normalized_model} + if reasoning_effort: + level = "low" if reasoning_effort == "minimal" else reasoning_effort + if level not in _GEMINI_EFFORTS: + if strict: + raise UnsupportedEffortError( + f"Gemini doesn't accept effort={level!r}" + ) + else: + params["reasoning_effort"] = level + return params + hf_model = normalized_model api_key = _resolve_hf_router_token(session_hf_token) params = { diff --git a/agent/core/model_switcher.py b/agent/core/model_switcher.py index 5ece764d..88521367 100644 --- a/agent/core/model_switcher.py +++ b/agent/core/model_switcher.py @@ -20,7 +20,7 @@ from litellm import acompletion from agent.core.effort_probe import ProbeInconclusive, probe_effort -from agent.core.llm_params import _resolve_llm_params +from agent.core.llm_params import _GEMINI_PREFIXES, _resolve_llm_params from agent.core.local_models import ( LOCAL_MODEL_PREFIXES, is_local_model_id, @@ -48,6 +48,8 @@ {"id": KIMI_K27_CODE_MODEL_ID, "label": "Kimi K2.7 Code"}, {"id": GLM_52_MODEL_ID, "label": "GLM 5.2"}, {"id": DEEPSEEK_V4_PRO_MODEL_ID, "label": "DeepSeek V4 Pro"}, + {"id": "gemini/gemini-2.5-pro", "label": "Gemini 2.5 Pro"}, + {"id": "gemini/gemini-2.5-flash", "label": "Gemini 2.5 Flash"}, ] @@ -59,6 +61,7 @@ def is_valid_model_id(model_id: str) -> bool: """Loose format check — lets users pick any model id. Accepts: + • gemini/, vertex_ai/ (direct Google API) • ollama/, vllm/, lm_studio/, llamacpp/ • /[:] (HF router; tag = provider or policy) • huggingface//[:] (same, optional LiteLLM prefix) @@ -95,6 +98,12 @@ def _print_hf_routing_info(model_id: str, console) -> bool: if is_local_model_id(model_id): return True + # Gemini goes direct to Google, not through the HF router, so the + # catalog has nothing to say about it. The probe below covers "does + # this model exist". + if model_id.startswith(_GEMINI_PREFIXES): + return True + from agent.core import hf_router_catalog as cat bare, _, tag = model_id.partition(":") @@ -163,6 +172,7 @@ def print_model_listing(config, console) -> None: console.print( "\n[dim]Paste any HF model id (e.g. 'MiniMaxAI/MiniMax-M3:novita').\n" "Add ':fastest', ':cheapest', ':preferred', or ':' to override routing.\n" + "Use 'gemini/' or 'vertex_ai/' for Google Gemini.\n" "Use 'ollama/', 'vllm/', 'lm_studio/', or " "'llamacpp/' for local OpenAI-compatible endpoints.[/dim]" ) @@ -173,6 +183,7 @@ def print_invalid_id(arg: str, console) -> None: console.print( "[dim]Expected:\n" " • /[:tag] (HF router — paste from huggingface.co)\n" + " • gemini/ | vertex_ai/\n" " • ollama/ | vllm/ | lm_studio/ | llamacpp/[/dim]" ) diff --git a/tests/unit/test_llm_params.py b/tests/unit/test_llm_params.py index a985025a..a6352763 100644 --- a/tests/unit/test_llm_params.py +++ b/tests/unit/test_llm_params.py @@ -77,6 +77,65 @@ def test_huggingface_prefix_is_stripped_for_router_calls(): assert params["api_base"] == HF_ROUTER_BASE_URL +def test_gemini_effort_is_forwarded_as_reasoning_effort(): + params = _resolve_llm_params( + "gemini/gemini-2.5-pro", + reasoning_effort="high", + strict=True, + ) + + assert params == {"model": "gemini/gemini-2.5-pro", "reasoning_effort": "high"} + + +def test_vertex_ai_effort_is_forwarded_as_reasoning_effort(): + params = _resolve_llm_params( + "vertex_ai/gemini-2.5-flash", + reasoning_effort="medium", + strict=True, + ) + + assert params == { + "model": "vertex_ai/gemini-2.5-flash", + "reasoning_effort": "medium", + } + + +def test_gemini_minimal_effort_normalizes_to_low(): + params = _resolve_llm_params( + "gemini/gemini-2.5-pro", + reasoning_effort="minimal", + strict=True, + ) + + assert params["reasoning_effort"] == "low" + + +def test_gemini_max_effort_is_rejected_in_strict_mode(): + with pytest.raises(UnsupportedEffortError, match="Gemini doesn't accept"): + _resolve_llm_params( + "gemini/gemini-2.5-pro", + reasoning_effort="max", + strict=True, + ) + + +def test_gemini_unsupported_effort_is_dropped_in_non_strict_mode(): + params = _resolve_llm_params( + "gemini/gemini-2.5-pro", + reasoning_effort="xhigh", + strict=False, + ) + + assert params == {"model": "gemini/gemini-2.5-pro"} + + +def test_gemini_is_not_routed_through_hf_router(): + params = _resolve_llm_params("gemini/gemini-2.5-pro") + + assert params["model"] == "gemini/gemini-2.5-pro" + assert "api_base" not in params + + def test_resolve_ollama_params_adds_v1_and_uses_default_key(monkeypatch): monkeypatch.delenv("OLLAMA_API_KEY", raising=False) monkeypatch.setenv("OLLAMA_BASE_URL", "http://localhost:11434")