UNS-480 [FEAT] Add Gemini LLM adapter for Google AI Studio#1890
UNS-480 [FEAT] Add Gemini LLM adapter for Google AI Studio#1890jaseemjaskp wants to merge 9 commits intomainfrom
Conversation
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().
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a Gemini LLM integration: new Pydantic parameter model Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
| Filename | Overview |
|---|---|
| unstract/sdk1/src/unstract/sdk1/adapters/base1.py | Adds GeminiLLMParameters class with thinking-mode handling; correctly makes a copy of adapter_metadata (unlike the Anthropic sibling). Minor: isinstance(int) check accepts bool values. |
| unstract/sdk1/src/unstract/sdk1/adapters/llm1/gemini.py | New adapter class following established pattern; implements all required identity static methods cleanly. |
| unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json | Well-formed JSON schema with allOf conditional for budget_tokens; model default correctly omits the gemini/ prefix. UI hint for thinking mode is missing gemini-2.0-flash-thinking-exp. |
| unstract/sdk1/tests/test_gemini_adapter.py | Good test coverage across 16 cases; covers thinking enabled/disabled, mutation, re-entrant behavior, Pydantic surface, and schema shape. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[adapter_metadata dict] --> B[GeminiLLMParameters.validate]
B --> C[Copy metadata, prefix model with gemini/]
C --> D{enable_thinking?}
D -- "No, and no existing thinking key" --> E[Pass to Pydantic validation]
D -- "Yes or pre-existing thinking key" --> F{has_thinking_config?}
F -- "Yes: preserve existing dict" --> G[Copy thinking dict as-is]
F -- "No: build new config" --> H{budget_tokens present and >= 1024?}
H -- No --> I[Raise ValueError]
H -- Yes --> J[Build thinking dict, force temperature=1]
G --> K[Force temperature=1]
J --> K
K --> E
E --> L[model_dump - validated fields]
L --> M{thinking enabled?}
M -- Yes --> N[Re-attach thinking dict to output]
M -- No --> O[Return validated dict]
N --> O
Prompt To Fix All With AI
This is a comment left during a code review.
Path: unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json
Line: 57-62
Comment:
**UI hint for thinking mode omits `gemini-2.0-flash-thinking-exp`**
The PR description explicitly lists `gemini-2.0-flash-thinking-exp` as a supported thinking model, but the `enable_thinking` description in the schema only mentions `gemini-2.5-pro` and `gemini-2.5-flash`. Users who want to use the experimental model won't see it in the hint.
```suggestion
"description": "Enable extended thinking for supported models. Thinking mode is only supported on: gemini-2.5-pro, gemini-2.5-flash, gemini-2.0-flash-thinking-exp. When enabled, temperature is forced to 1."
```
How can I resolve this? If you propose a fix, please make it concise.
---
This is a comment left during a code review.
Path: unstract/sdk1/src/unstract/sdk1/adapters/base1.py
Line: 673-676
Comment:
**`bool` values pass the `isinstance(int)` guard**
In Python `bool` is a subclass of `int`, so `isinstance(True, int)` returns `True`. A caller that passes a boolean for `budget_tokens` will get the `>= 1024` branch rather than the type-error branch — producing a misleading error message. Add an explicit `bool` exclusion alongside the existing guard:
```python
if (
not isinstance(budget_tokens, int)
or isinstance(budget_tokens, bool)
or budget_tokens < 1024
):
raise ValueError(
f"budget_tokens must be an integer >= 1024, got {budget_tokens}"
)
```
How can I resolve this? If you propose a fix, please make it concise.Reviews (8): Last reviewed commit: "UNS-480 [FIX] Clean up Gemini JSON schem..." | Re-trigger Greptile
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json (1)
33-53: Preferintegerfor discrete numeric inputs.
max_tokens,timeout, andmax_retriesare discrete/count-like settings. Usingintegeris clearer and prevents fractional inputs from slipping through lenient validators/clients.Proposed schema tightening
"max_tokens": { - "type": "number", + "type": "integer", "minimum": 0, - "multipleOf": 1, "title": "Maximum Output Tokens", "description": "Maximum number of output tokens to limit LLM replies, the maximum possible differs from model to model." }, "timeout": { - "type": "number", + "type": "integer", "minimum": 0, - "multipleOf": 1, "title": "Timeout", "default": 600, "description": "Timeout in seconds" }, "max_retries": { - "type": "number", + "type": "integer", "minimum": 0, - "multipleOf": 1, "title": "Max Retries", "description": "Maximum number of retries" }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json` around lines 33 - 53, The schema uses "type": "number" for discrete/count fields; change the types for max_tokens, timeout, and max_retries to "integer" (keeping their existing constraints like minimum, multipleOf, default, titles and descriptions) so validators/clients reject fractional values—update the entries for the keys "max_tokens", "timeout", and "max_retries" in the Gemini schema accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py`:
- Around line 653-658: The validate_model function should fail fast on missing
or blank model values instead of returning "gemini/"; update
validate_model(adapter_metadata) to read the raw model, strip whitespace, and if
it's empty or missing raise a ValueError (or appropriate ValidationError) with a
clear message; otherwise preserve the existing prefixing logic (if model already
startswith "gemini/" return it, else return "gemini/{model}"). Ensure the check
uses adapter_metadata.get("model") and applies .strip() before prefixing so
blank strings are rejected.
---
Nitpick comments:
In `@unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json`:
- Around line 33-53: The schema uses "type": "number" for discrete/count fields;
change the types for max_tokens, timeout, and max_retries to "integer" (keeping
their existing constraints like minimum, multipleOf, default, titles and
descriptions) so validators/clients reject fractional values—update the entries
for the keys "max_tokens", "timeout", and "max_retries" in the Gemini schema
accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 36d4675d-43c9-4c73-92cf-aa2667969cf6
⛔ Files ignored due to path filters (1)
frontend/public/icons/adapter-icons/Gemini.pngis excluded by!**/*.png
📒 Files selected for processing (3)
unstract/sdk1/src/unstract/sdk1/adapters/base1.pyunstract/sdk1/src/unstract/sdk1/adapters/llm1/gemini.pyunstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json
…model, update default model
|
@jaseemjaskp Please include thinking options in the adapter settings. This would enable the user to modify the thinking settings according to their needs. You can refer our VertexAI adapter for this. |
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.
|
Tip: Greploop — Automatically fix all review issues by running Use the Greptile plugin for Claude Code to query reviews, search comments, and manage custom context directly from your terminal. |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@unstract/sdk1/src/unstract/sdk1/adapters/base1.py`:
- Around line 664-672: The code sets thinking mode without enforcing the
thinking-budget contract: in the enable_thinking branch (use symbols
enable_thinking, has_thinking_config, result_metadata, adapter_metadata,
thinking_config, budget_tokens) validate that budget_tokens is present and >=
1024 instead of silently accepting missing/invalid values; if adapter_metadata
provides thinking config, validate its budget_tokens meets the >=1024
requirement and otherwise raise/return a clear validation error (or set a safe
default only if spec allows); apply the same validation logic to the other
analogous block (the block around lines 676-687) so both code paths enforce the
required budget_tokens existence and minimum constraint.
In `@unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json`:
- Line 23: Update the JSON "description" field to include gemini-2.5-flash so
supported-model text is consistent with the thinking-mode description;
specifically, edit the "description" value (the LiteLLM model string
description) to list gemini-2.0-flash, gemini-2.5-pro, gemini-2.5-flash,
gemini-1.5-pro, gemini-1.5-flash (and note that the gemini/ prefix is added
automatically) to match the other reference to thinking mode.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: bda4bd74-4b96-4b4b-8a10-e55c94f3e390
📒 Files selected for processing (3)
unstract/sdk1/src/unstract/sdk1/adapters/base1.pyunstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.jsonunstract/sdk1/tests/test_gemini_adapter.py
unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json
Outdated
Show resolved
Hide resolved
…bled 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.
unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json
Outdated
Show resolved
Hide resolved
chandrasekharan-zipstack
left a comment
There was a problem hiding this comment.
LGTM, left a comment about avoiding mentioning litellm in the description
unstract/sdk1/src/unstract/sdk1/adapters/llm1/static/gemini.json
Outdated
Show resolved
Hide resolved
Remove internal implementation detail (LiteLLM) from model description and remove experimental model from thinking mode supported models list.
Frontend Lint Report (Biome)✅ All checks passed! No linting or formatting issues found. |
Test ResultsSummary
Runner Tests - Full Report
SDK1 Tests - Full Report
|
|



What
Why
How
GeminiLLMParametersclass inbase1.pywithapi_keyfield andvalidate()/validate_model()(idempotentgemini/prefix)GeminiLLMParameters.validate()with thinking-mode handling:enable_thinkingandbudget_tokensare consumed from adapter metadata (not Pydantic fields), excluded from model validation, and re-attached to the final dict. When thinking is enabled, temperature is forced to1(Gemini API requirement). Re-entrant: pre-existingthinkingconfig in metadata is preserved.GeminiLLMAdapterinllm1/gemini.pyfollowing the established Anthropic adapter patterngemini.jsonschema with fields: adapter_name, api_key, model, temperature, max_tokens, timeout, max_retries, enable_thinking. Added anallOf/if/thenconditional sobudget_tokens(min 1024, default 1024, no max) is only required whenenable_thinking=true. UI hint onenable_thinkinglists supported models (gemini-2.5-pro,gemini-2.5-flash,gemini-2.0-flash-thinking-exp).Gemini.png)tests/test_gemini_adapter.py— 16 tests covering prefix logic, thinking enabled/disabled, temperature override, input non-mutation, preservation of existing thinking config, Pydantic field surface, adapter identity, and schema shaperegister_adapters()— no manual registration or__init__.pychanges neededCan this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)
GeminiLLMParametersclass. No other adapters are touched. Existing SDK tests continue to pass alongside the 16 new Gemini adapter tests.Database Migrations
Env Config
Relevant Docs
Related Issues or PRs
Dependencies Versions
litellmpackage forgemini/provider routing and thethinkingkwarg.Notes on Testing
uv run pytest tests/test_gemini_adapter.py→ 16 passedvalidate_model(prefix add, no double-prefix, blank rejected); thinking disabled (nothinkingkey, user temperature preserved, control fields excluded); thinking enabled (config built, temperature forced to 1, missingbudget_tokenshandled, existing config preserved, input not mutated); Pydantic field surface (enable_thinking/budget_tokens/thinkingnot fields,api_keyrequired); adapter identity; JSON schemaallOfshapemodelfield is free-text (not enum dropdown) — LiteLLM handles routing errors for unsupported model strings. Thinking on unsupported models is not blocked adapter-side; the schema UI hint is the only guardrail (by design).Checklist
I have read and understood the Contribution Guidelines.