Skip to content

Commit 4f517a4

Browse files
authored
UNS-480 [FEAT] Add Gemini LLM adapter for Google AI Studio (#1890)
* UNS-480 [FEAT] Add Gemini LLM adapter for Google AI Studio Add a new LLM adapter for Google's Gemini models using LiteLLM's gemini/ provider prefix. The adapter follows the established SDK adapter pattern and is auto-discovered by register_adapters(). * UNS-480 [FEAT] Add defaults for max_tokens and max_retries in Gemini JSON schema * UNS-480 [FIX] Address PR review: avoid dict mutation, validate blank model, update default model * UNS-482 [FEAT] Add Gemini thinking mode support with tests Extends the Gemini LLM adapter with optional thinking mode, mirroring the Anthropic/Bedrock pattern. enable_thinking and budget_tokens are consumed from adapter metadata (not Pydantic fields); when enabled, temperature is forced to 1. Schema gains an allOf/if-then conditional so budget_tokens (min 1024) is only required when thinking is on. * UNS-482 [FIX] Use pytest.approx for temperature float comparison * UNS-482 [FIX] Validate budget_tokens when Gemini thinking mode is enabled Raise ValueError if budget_tokens is missing, not an integer, or below 1024 when enable_thinking=True. Previously these cases silently produced incomplete or invalid thinking configs that would fail at the API level. * UNS-480 [FIX] Add gemini-2.5-flash to model description in JSON schema * UNS-480 [FIX] Clean up Gemini JSON schema descriptions per PR review Remove internal implementation detail (LiteLLM) from model description and remove experimental model from thinking mode supported models list.
1 parent 740834d commit 4f517a4

5 files changed

Lines changed: 359 additions & 0 deletions

File tree

8.04 KB
Loading

unstract/sdk1/src/unstract/sdk1/adapters/base1.py

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -639,6 +639,72 @@ def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
639639
return f"anthropic/{model}"
640640

641641

642+
class GeminiLLMParameters(BaseChatCompletionParameters):
643+
"""See https://docs.litellm.ai/docs/providers/gemini."""
644+
645+
api_key: str
646+
647+
@staticmethod
648+
def validate(adapter_metadata: dict[str, "Any"]) -> dict[str, "Any"]:
649+
result_metadata = adapter_metadata.copy()
650+
result_metadata["model"] = GeminiLLMParameters.validate_model(adapter_metadata)
651+
652+
# Handle Gemini thinking configuration
653+
enable_thinking = adapter_metadata.get("enable_thinking", False)
654+
655+
# If enable_thinking is not explicitly provided but thinking config is present,
656+
# assume thinking was enabled in a previous validation
657+
has_thinking_config = (
658+
"thinking" in adapter_metadata
659+
and adapter_metadata.get("thinking") is not None
660+
)
661+
if not enable_thinking and has_thinking_config:
662+
enable_thinking = True
663+
664+
if enable_thinking:
665+
if has_thinking_config:
666+
result_metadata["thinking"] = adapter_metadata["thinking"]
667+
else:
668+
budget_tokens = adapter_metadata.get("budget_tokens")
669+
if budget_tokens is None:
670+
raise ValueError(
671+
"budget_tokens is required when thinking mode is enabled"
672+
)
673+
if not isinstance(budget_tokens, int) or budget_tokens < 1024:
674+
raise ValueError(
675+
f"budget_tokens must be an integer >= 1024, got {budget_tokens}"
676+
)
677+
result_metadata["thinking"] = {
678+
"type": "enabled",
679+
"budget_tokens": budget_tokens,
680+
}
681+
# Gemini thinking mode requires temperature=1
682+
result_metadata["temperature"] = 1
683+
684+
# Exclude control fields from pydantic validation
685+
exclude_fields = ("enable_thinking", "budget_tokens", "thinking")
686+
validation_metadata = {
687+
k: v for k, v in result_metadata.items() if k not in exclude_fields
688+
}
689+
690+
validated = GeminiLLMParameters(**validation_metadata).model_dump()
691+
692+
if enable_thinking and "thinking" in result_metadata:
693+
validated["thinking"] = result_metadata["thinking"]
694+
695+
return validated
696+
697+
@staticmethod
698+
def validate_model(adapter_metadata: dict[str, "Any"]) -> str:
699+
model = str(adapter_metadata.get("model", "")).strip()
700+
if not model:
701+
raise ValueError("model is required")
702+
if model.startswith("gemini/"):
703+
return model
704+
else:
705+
return f"gemini/{model}"
706+
707+
642708
class AnyscaleLLMParameters(BaseChatCompletionParameters):
643709
"""See https://docs.litellm.ai/docs/providers/anyscale."""
644710

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
from typing import Any
2+
3+
from unstract.sdk1.adapters.base1 import BaseAdapter, GeminiLLMParameters
4+
from unstract.sdk1.adapters.enums import AdapterTypes
5+
6+
7+
class GeminiLLMAdapter(GeminiLLMParameters, BaseAdapter):
8+
@staticmethod
9+
def get_id() -> str:
10+
return "gemini|085f6c03-b57e-4594-85bb-40e2616c2736"
11+
12+
@staticmethod
13+
def get_metadata() -> dict[str, Any]:
14+
return {
15+
"name": "Gemini",
16+
"version": "1.0.0",
17+
"adapter": GeminiLLMAdapter,
18+
"description": "Google Gemini LLM adapter via Google AI Studio",
19+
"is_active": True,
20+
}
21+
22+
@staticmethod
23+
def get_name() -> str:
24+
return "Gemini"
25+
26+
@staticmethod
27+
def get_description() -> str:
28+
return "Google Gemini LLM adapter via Google AI Studio"
29+
30+
@staticmethod
31+
def get_provider() -> str:
32+
return "gemini"
33+
34+
@staticmethod
35+
def get_icon() -> str:
36+
return "/icons/adapter-icons/Gemini.png"
37+
38+
@staticmethod
39+
def get_adapter_type() -> AdapterTypes:
40+
return AdapterTypes.LLM
Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
{
2+
"title": "Gemini LLM",
3+
"type": "object",
4+
"required": ["adapter_name", "api_key", "model"],
5+
"properties": {
6+
"adapter_name": {
7+
"type": "string",
8+
"title": "Name",
9+
"default": "",
10+
"description": "Provide a unique name for this adapter instance. Example: gemini-group-1"
11+
},
12+
"api_key": {
13+
"type": "string",
14+
"title": "API Key",
15+
"default": "",
16+
"description": "Google AI Studio API key",
17+
"format": "password"
18+
},
19+
"model": {
20+
"type": "string",
21+
"title": "Model",
22+
"default": "gemini-2.0-flash",
23+
"description": "Supported: gemini-2.0-flash, gemini-2.5-pro, gemini-2.5-flash, gemini-1.5-pro, gemini-1.5-flash. The gemini/ prefix will be added automatically if omitted."
24+
},
25+
"temperature": {
26+
"type": "number",
27+
"minimum": 0,
28+
"maximum": 2,
29+
"title": "Temperature",
30+
"default": 0.1,
31+
"description": "Sampling temperature between 0 and 2"
32+
},
33+
"max_tokens": {
34+
"type": "number",
35+
"minimum": 0,
36+
"multipleOf": 1,
37+
"default": 8192,
38+
"title": "Maximum Output Tokens",
39+
"description": "Maximum number of output tokens to limit LLM replies, the maximum possible differs from model to model."
40+
},
41+
"timeout": {
42+
"type": "number",
43+
"minimum": 0,
44+
"multipleOf": 1,
45+
"title": "Timeout",
46+
"default": 600,
47+
"description": "Timeout in seconds"
48+
},
49+
"max_retries": {
50+
"type": "number",
51+
"minimum": 0,
52+
"multipleOf": 1,
53+
"default": 3,
54+
"title": "Max Retries",
55+
"description": "Maximum number of retries"
56+
},
57+
"enable_thinking": {
58+
"type": "boolean",
59+
"title": "Enable Thinking Mode",
60+
"default": false,
61+
"description": "Enable extended thinking for supported models. Thinking mode is only supported on: gemini-2.5-pro, gemini-2.5-flash. When enabled, temperature is forced to 1."
62+
}
63+
},
64+
"allOf": [
65+
{
66+
"if": {
67+
"properties": {
68+
"enable_thinking": { "const": true }
69+
},
70+
"required": ["enable_thinking"]
71+
},
72+
"then": {
73+
"required": ["budget_tokens"],
74+
"properties": {
75+
"budget_tokens": {
76+
"type": "integer",
77+
"minimum": 1024,
78+
"default": 1024,
79+
"title": "Budget Tokens",
80+
"description": "Number of tokens allocated for the thinking process. Minimum: 1024."
81+
}
82+
}
83+
}
84+
}
85+
]
86+
}
Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,167 @@
1+
"""Unit tests for the Gemini LLM adapter (UNS-480 / UNS-482)."""
2+
3+
import json
4+
from pathlib import Path
5+
6+
import pytest
7+
from unstract.sdk1.adapters.base1 import GeminiLLMParameters
8+
from unstract.sdk1.adapters.llm1.gemini import GeminiLLMAdapter
9+
10+
BASE_METADATA = {"api_key": "test-key", "model": "gemini-2.5-flash"}
11+
12+
13+
# ── validate_model ───────────────────────────────────────────────────────────
14+
15+
16+
def test_validate_model_prefixes_when_missing() -> None:
17+
assert (
18+
GeminiLLMParameters.validate_model({"model": "gemini-2.5-flash"})
19+
== "gemini/gemini-2.5-flash"
20+
)
21+
22+
23+
def test_validate_model_does_not_double_prefix() -> None:
24+
assert (
25+
GeminiLLMParameters.validate_model({"model": "gemini/gemini-2.5-pro"})
26+
== "gemini/gemini-2.5-pro"
27+
)
28+
29+
30+
def test_validate_model_blank_raises() -> None:
31+
with pytest.raises(ValueError, match="model is required"):
32+
GeminiLLMParameters.validate_model({"model": " "})
33+
34+
35+
# ── validate: thinking disabled ──────────────────────────────────────────────
36+
37+
38+
def test_validate_thinking_disabled_by_default() -> None:
39+
result = GeminiLLMParameters.validate({**BASE_METADATA, "temperature": 0.3})
40+
assert result["model"] == "gemini/gemini-2.5-flash"
41+
assert "thinking" not in result
42+
assert result["temperature"] == pytest.approx(0.3)
43+
44+
45+
def test_validate_excludes_control_fields_from_model() -> None:
46+
result = GeminiLLMParameters.validate(BASE_METADATA.copy())
47+
assert "enable_thinking" not in result
48+
assert "budget_tokens" not in result
49+
50+
51+
# ── validate: thinking enabled ───────────────────────────────────────────────
52+
53+
54+
def test_validate_thinking_enabled_with_budget() -> None:
55+
result = GeminiLLMParameters.validate(
56+
{**BASE_METADATA, "enable_thinking": True, "budget_tokens": 2048}
57+
)
58+
assert result["thinking"] == {"type": "enabled", "budget_tokens": 2048}
59+
assert result["temperature"] == 1
60+
61+
62+
def test_validate_thinking_overrides_user_temperature() -> None:
63+
result = GeminiLLMParameters.validate(
64+
{
65+
**BASE_METADATA,
66+
"temperature": 0.7,
67+
"enable_thinking": True,
68+
"budget_tokens": 1024,
69+
}
70+
)
71+
assert result["temperature"] == 1
72+
73+
74+
def test_validate_thinking_enabled_without_budget_raises() -> None:
75+
with pytest.raises(ValueError, match="budget_tokens is required"):
76+
GeminiLLMParameters.validate({**BASE_METADATA, "enable_thinking": True})
77+
78+
79+
def test_validate_thinking_budget_tokens_invalid_type_raises() -> None:
80+
with pytest.raises(ValueError, match="budget_tokens must be an integer >= 1024"):
81+
GeminiLLMParameters.validate(
82+
{**BASE_METADATA, "enable_thinking": True, "budget_tokens": "hello"}
83+
)
84+
85+
86+
def test_validate_thinking_budget_tokens_too_small_raises() -> None:
87+
with pytest.raises(ValueError, match="budget_tokens must be an integer >= 1024"):
88+
GeminiLLMParameters.validate(
89+
{**BASE_METADATA, "enable_thinking": True, "budget_tokens": 512}
90+
)
91+
92+
93+
def test_validate_preserves_existing_thinking_config() -> None:
94+
existing = {"type": "enabled", "budget_tokens": 4096}
95+
result = GeminiLLMParameters.validate({**BASE_METADATA, "thinking": existing})
96+
assert result["thinking"] == existing
97+
assert result["temperature"] == 1
98+
99+
100+
def test_validate_does_not_mutate_input() -> None:
101+
metadata = {**BASE_METADATA, "enable_thinking": True, "budget_tokens": 2048}
102+
snapshot = metadata.copy()
103+
GeminiLLMParameters.validate(metadata)
104+
assert metadata == snapshot
105+
106+
107+
# ── Pydantic field surface ───────────────────────────────────────────────────
108+
109+
110+
def test_thinking_controls_not_pydantic_fields() -> None:
111+
fields = GeminiLLMParameters.model_fields
112+
assert "enable_thinking" not in fields
113+
assert "budget_tokens" not in fields
114+
assert "thinking" not in fields
115+
assert "api_key" in fields
116+
117+
118+
def test_api_key_is_required() -> None:
119+
from pydantic import ValidationError
120+
121+
with pytest.raises(ValidationError):
122+
GeminiLLMParameters(model="gemini/gemini-2.5-flash")
123+
124+
125+
# ── Adapter identity ─────────────────────────────────────────────────────────
126+
127+
128+
def test_adapter_identity() -> None:
129+
assert GeminiLLMAdapter.get_name() == "Gemini"
130+
assert GeminiLLMAdapter.get_provider() == "gemini"
131+
assert GeminiLLMAdapter.get_id().startswith("gemini|")
132+
metadata = GeminiLLMAdapter.get_metadata()
133+
assert metadata["is_active"] is True
134+
assert metadata["name"] == "Gemini"
135+
136+
137+
# ── JSON schema ──────────────────────────────────────────────────────────────
138+
139+
140+
@pytest.fixture
141+
def gemini_schema() -> dict:
142+
schema_path = (
143+
Path(__file__).parent.parent
144+
/ "src/unstract/sdk1/adapters/llm1/static/gemini.json"
145+
)
146+
return json.loads(schema_path.read_text())
147+
148+
149+
def test_schema_required_fields(gemini_schema: dict) -> None:
150+
assert set(gemini_schema["required"]) >= {"adapter_name", "api_key", "model"}
151+
152+
153+
def test_schema_enable_thinking_default_false(gemini_schema: dict) -> None:
154+
assert gemini_schema["properties"]["enable_thinking"]["default"] is False
155+
156+
157+
def test_schema_budget_tokens_conditional(gemini_schema: dict) -> None:
158+
all_of = gemini_schema["allOf"]
159+
assert len(all_of) == 1
160+
conditional = all_of[0]
161+
assert conditional["if"]["properties"]["enable_thinking"]["const"] is True
162+
then_block = conditional["then"]
163+
assert "budget_tokens" in then_block["required"]
164+
budget = then_block["properties"]["budget_tokens"]
165+
assert budget["minimum"] == 1024
166+
assert budget["default"] == 1024
167+
assert "maximum" not in budget

0 commit comments

Comments
 (0)