"""Gemini provider implementation."""
from __future__ import annotations
import copy
import json
import os
import urllib.error
import urllib.request
from typing import Any, Dict, Iterator, List, Optional, Set
from ..core.exceptions import ProviderError
from ..model_runtime import (
call_with_retries,
execute_legacy_tool_call,
max_provider_retries,
)
from ..models import (
ContentPart,
ModelEvent,
ModelRequest,
ModelResponse,
ProviderCapabilities,
ToolCall,
ToolResult,
Usage,
)
from .errors import map_gemini_http_error, map_provider_exception
from .registry import reasoning_parameters
from .usage import gemini_usage
def _redact_secret(message: str, secret: Optional[str]) -> str:
if message and secret:
return message.replace(secret, "***")
return message
[docs]
class GeminiProvider:
"""Google Gemini provider using the public REST API."""
provider_name = "gemini"
capabilities = ProviderCapabilities(
tools=True,
streaming=True,
native_streaming=True,
structured_outputs=True,
json_schema_mode="json_schema",
multimodal=True,
image_input=True,
file_input=True,
audio_input=True,
video_input=True,
reasoning=True,
embeddings=True,
)
[docs]
def __init__(self, config: Any):
self.config = config
api_key_env = getattr(config, "api_key_env", None) or "GEMINI_API_KEY"
# Strip whitespace from keys read from env files: http.client rejects
# a header value ending in a newline and quotes it in the error.
self.api_key = (
os.getenv(api_key_env) or os.getenv("GOOGLE_API_KEY") or ""
).strip() or None
if not self.api_key and not getattr(config, "base_url", None):
raise ProviderError(
f"{api_key_env} or GOOGLE_API_KEY environment variable not set"
)
self.base_url = (
getattr(config, "base_url", None)
or "https://generativelanguage.googleapis.com/v1beta"
)
[docs]
def generate(
self,
messages: List[Dict[str, Any]],
tools: Optional[List[Dict[str, Any]]] = None,
hitl_context: Optional[Dict[str, Any]] = None,
) -> str:
"""Generate a response from Gemini."""
payload = self._build_payload(messages, tools)
try:
data = self._post_json("generateContent", payload)
except Exception as e:
raise self._provider_error(e, "Gemini API error") from e
function_calls = self._extract_function_calls(data)
if function_calls:
return self._handle_function_calls(
function_calls=function_calls,
available_tools=tools or [],
messages=messages,
request=None,
hitl_context=hitl_context,
).content
return self._extract_text(data)
[docs]
def invoke(
self,
request: ModelRequest,
tools: Optional[List[Dict[str, Any]]] = None,
) -> ModelResponse:
"""Invoke Gemini through the provider-neutral adapter surface."""
payload = self._build_payload(
[message.model_dump(exclude_none=True) for message in request.messages],
tools,
request=request,
)
try:
data = self._post_json("generateContent", payload, timeout=request.timeout)
except Exception as e:
raise self._provider_error(e, "Gemini API error") from e
function_calls = self._extract_function_calls(data)
if function_calls:
return self._runtime_tool_call_response(data, payload, function_calls)
return ModelResponse(
content=self._extract_text(data),
provider=self.provider_name,
model=self._model_name(),
raw=data,
usage=gemini_usage(data),
)
[docs]
def stream(self, request: ModelRequest) -> Iterator[ModelEvent]:
"""Stream Gemini responses as provider-neutral events."""
payload = self._build_payload(
[message.model_dump(exclude_none=True) for message in request.messages],
None,
request=request,
)
content_parts: List[str] = []
usage: Optional[Usage] = None
try:
for data in self._post_stream(
"streamGenerateContent",
payload,
timeout=request.timeout,
):
usage = gemini_usage(data) or usage
text = self._extract_text(data)
if text:
content_parts.append(text)
yield ModelEvent(type="delta", delta=text)
except Exception as e:
error = self._provider_error(e, "Gemini streaming error")
yield ModelEvent(type="error", metadata={"message": error.message})
raise error from e
response = ModelResponse(
content="".join(content_parts),
provider=self.provider_name,
model=self._model_name(),
usage=usage,
)
if usage is not None:
yield ModelEvent(type="usage", usage=usage)
yield ModelEvent(type="final", response=response, usage=usage)
[docs]
def close(self) -> None:
"""Gemini REST provider does not hold persistent resources."""
def _provider_error(self, exc: BaseException, prefix: str) -> ProviderError:
"""Map a REST failure, reading and redacting an HTTP error body."""
if isinstance(exc, urllib.error.HTTPError):
return map_gemini_http_error(
exc,
prefix=prefix,
model=self._model_name(),
redact=lambda text: _redact_secret(text, self.api_key),
)
return map_provider_exception(
exc,
provider=self.provider_name,
model=self._model_name(),
message=f"{prefix}: {_redact_secret(str(exc), self.api_key)}",
)
def _model_name(self) -> str:
return str(getattr(self.config, "model", None) or "gemini-3.5-flash")
def _max_output_tokens(self) -> int:
return int(
getattr(self.config, "max_output_tokens", None)
or getattr(self.config, "max_tokens", 1000)
)
def _build_payload(
self,
messages: List[Dict[str, Any]],
tools: Optional[List[Dict[str, Any]]],
request: Optional[ModelRequest] = None,
) -> Dict[str, Any]:
contents = []
system_parts: List[str] = []
for message in messages:
role = message.get("role", "user")
content = message.get("content", "")
if role == "system":
system_parts.append(self._content_to_text(content))
continue
gemini_role = "model" if role == "assistant" else "user"
contents.append(
{"role": gemini_role, "parts": self._content_to_parts(content)}
)
payload: Dict[str, Any] = {
"contents": contents,
"generationConfig": {
"temperature": (
request.temperature
if request is not None and request.temperature is not None
else getattr(self.config, "temperature", 0.7)
),
"maxOutputTokens": self._max_output_tokens(),
},
}
if request is not None and request.max_output_tokens is not None:
payload["generationConfig"]["maxOutputTokens"] = request.max_output_tokens
if payload["generationConfig"].get("temperature") is None:
payload["generationConfig"].pop("temperature", None)
if request is not None and request.response_schema is not None:
payload["generationConfig"]["responseMimeType"] = "application/json"
payload["generationConfig"]["responseSchema"] = (
request.response_schema.json_schema or {}
)
if request is not None and request.reasoning is not None:
thinking = dict(reasoning_parameters(request))
if request.reasoning.budget_tokens is not None:
thinking.pop("thinkingLevel", None)
thinking["thinkingBudget"] = request.reasoning.budget_tokens
if thinking:
payload["generationConfig"]["thinkingConfig"] = thinking
if request is not None:
generation_config = request.provider_options.get("generation_config")
if isinstance(generation_config, dict):
payload["generationConfig"].update(generation_config)
if system_parts:
payload["systemInstruction"] = {
"parts": [{"text": "\n\n".join(system_parts)}]
}
formatted_tools = self._format_tools(tools or [])
if not formatted_tools and request is not None and request.tools:
formatted_tools = [
{
"name": spec.name,
"description": spec.description,
"parametersJsonSchema": copy.deepcopy(spec.parameters),
}
for spec in request.tools
]
if formatted_tools:
payload["tools"] = [{"functionDeclarations": formatted_tools}]
return payload
def _post_json(
self,
method: str,
payload: Dict[str, Any],
*,
timeout: Optional[float] = None,
) -> Dict[str, Any]:
request = urllib.request.Request(
self._method_url(method),
data=json.dumps(payload).encode("utf-8"),
headers=self._request_headers(),
method="POST",
)
with urllib.request.urlopen(
request,
timeout=timeout or getattr(self.config, "timeout", None) or 60,
) as response:
return json.loads(response.read().decode("utf-8"))
def _post_stream(
self,
method: str,
payload: Dict[str, Any],
*,
timeout: Optional[float] = None,
) -> Iterator[Dict[str, Any]]:
request = urllib.request.Request(
self._method_url(method, stream=True),
data=json.dumps(payload).encode("utf-8"),
headers=self._request_headers(),
method="POST",
)
with urllib.request.urlopen(
request,
timeout=timeout or getattr(self.config, "timeout", None) or 60,
) as response:
for raw_line in response:
line = raw_line.decode("utf-8").strip()
if not line:
continue
if line.startswith("data:"):
line = line[len("data:") :].strip()
if line in {"[DONE]", "DONE"}:
break
yield json.loads(line)
def _method_url(self, method: str, *, stream: bool = False) -> str:
"""Build the endpoint URL; the API key travels in a header, never here."""
url = f"{self.base_url}/models/{self._model_name()}:{method}"
if stream:
separator = "&" if "?" in url else "?"
url = f"{url}{separator}alt=sse"
return url
def _request_headers(self) -> Dict[str, str]:
"""JSON headers plus ``x-goog-api-key`` when a key is configured.
Sending the key as a header keeps it out of the URL, so exceptions,
logs and proxies that quote the URL never carry it.
"""
headers = {"Content-Type": "application/json"}
if self.api_key:
headers["x-goog-api-key"] = self.api_key
return headers
def _content_to_parts(self, content: Any) -> List[Dict[str, Any]]:
if not isinstance(content, list):
return [{"text": str(content)}]
parts: List[Dict[str, Any]] = []
for item in content:
part = item if isinstance(item, ContentPart) else ContentPart(**item)
if part.type == "text":
parts.append({"text": part.text or ""})
elif part.type in {"image_url", "file_url", "audio_url", "video_url"}:
parts.append(
{
"fileData": {
"mimeType": self._default_mime_type(part),
"fileUri": part.url or "",
}
}
)
elif part.type in {
"image_base64",
"file",
"audio_base64",
"video_base64",
}:
parts.append(
{
"inlineData": {
"mimeType": self._default_mime_type(part),
"data": part.data or "",
}
}
)
else:
raise ProviderError(
f"Gemini provider cannot serialize content part type: {part.type}"
)
return parts
def _default_mime_type(self, part: ContentPart) -> str:
if part.mime_type:
return part.mime_type
if part.type.startswith("image"):
return "image/png"
if part.type.startswith("audio"):
return "audio/wav"
if part.type.startswith("video"):
return "video/mp4"
return "application/octet-stream"
def _content_to_text(self, content: Any) -> str:
if isinstance(content, str):
return content
if isinstance(content, list):
text_parts = []
for item in content:
part = item if isinstance(item, ContentPart) else ContentPart(**item)
if part.type == "text":
text_parts.append(part.text or "")
return "".join(text_parts)
return str(content)
def _format_tools(self, tools: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
declarations = []
for tool in tools:
func = tool.get("function")
if not callable(func):
continue
parameters = tool.get("parameters") or {}
if parameters.get("type") == "object":
# The REST API accepts JSON Schema directly. Using its dedicated
# field preserves nested types/constraints without a lossy
# conversion to the older Google Schema representation.
declarations.append(
{
"name": getattr(func, "__name__", ""),
"description": tool.get("description", ""),
"parametersJsonSchema": copy.deepcopy(parameters),
}
)
continue
declarations.append(
{
"name": getattr(func, "__name__", ""),
"description": tool.get("description", ""),
"parameters": {
"type": "OBJECT",
"properties": {
name: self._gemini_parameter_schema(param_info)
for name, param_info in (
tool.get("parameters") or {}
).items()
},
"required": [
name
for name, param_info in (
tool.get("parameters") or {}
).items()
if isinstance(param_info, dict)
and param_info.get("required", False)
],
},
}
)
return declarations
def _gemini_parameter_schema(self, param_info: Any) -> Dict[str, Any]:
if not isinstance(param_info, dict):
return {"type": "STRING"}
json_type = self._python_type_to_gemini_schema_type(
str(param_info.get("type", "str"))
)
schema: Dict[str, Any] = {"type": json_type}
if param_info.get("description"):
schema["description"] = str(param_info["description"])
if param_info.get("enum"):
schema["enum"] = list(param_info["enum"])
return schema
def _python_type_to_gemini_schema_type(self, python_type: str) -> str:
mapping = {
"string": "STRING",
"integer": "INTEGER",
"number": "NUMBER",
"boolean": "BOOLEAN",
"array": "ARRAY",
"object": "OBJECT",
"str": "STRING",
"int": "INTEGER",
"float": "NUMBER",
"bool": "BOOLEAN",
"list": "ARRAY",
"dict": "OBJECT",
"List": "ARRAY",
"Dict": "OBJECT",
}
return mapping.get(python_type, "STRING")
def _extract_text(self, data: Dict[str, Any]) -> str:
candidates = data.get("candidates") or []
if not candidates:
return ""
content = candidates[0].get("content") or {}
parts = content.get("parts") or []
return "".join(str(part.get("text", "")) for part in parts)
def _extract_function_calls(
self,
data: Dict[str, Any],
*,
id_offset: int = 0,
) -> List[Dict[str, Any]]:
"""Return the first candidate's function calls with run-unique ids.
The API's ``functionCall.id`` is used when present; otherwise ids are
numbered from ``id_offset``, the count of earlier calls in the run.
"""
calls: List[Dict[str, Any]] = []
for part in self._candidate_content(data).get("parts") or []:
function_call = self._part_function_call(part)
if function_call is None:
continue
name = str(function_call.get("name") or "")
if not name:
continue
native_id = function_call.get("id")
calls.append(
{
"id": (
str(native_id)
if native_id
else f"gemini-call-{id_offset + len(calls)}"
),
"name": name,
"args": function_call.get("args") or {},
"raw": function_call,
}
)
return calls
def _candidate_content(self, data: Dict[str, Any]) -> Dict[str, Any]:
candidates = data.get("candidates") or []
if not candidates or not isinstance(candidates[0], dict):
return {}
content = candidates[0].get("content")
return content if isinstance(content, dict) else {}
def _part_function_call(self, part: Any) -> Optional[Dict[str, Any]]:
if not isinstance(part, dict):
return None
function_call = part.get("functionCall") or part.get("function_call")
return function_call if isinstance(function_call, dict) else None
def _content_function_call_count(self, contents: List[Dict[str, Any]]) -> int:
return sum(
1
for content in contents
if isinstance(content, dict) and content.get("role") == "model"
for part in content.get("parts") or []
if self._part_function_call(part) is not None
)
def _native_function_call_ids(self, contents: List[Dict[str, Any]]) -> Set[str]:
"""Return API-issued call ids from the latest model turn."""
for content in reversed(contents):
if isinstance(content, dict) and content.get("role") == "model":
return {
str(function_call["id"])
for function_call in (
self._part_function_call(part)
for part in content.get("parts") or []
)
if function_call is not None and function_call.get("id")
}
return set()
def _runtime_tool_call_response(
self,
data: Dict[str, Any],
payload: Dict[str, Any],
function_calls: List[Dict[str, Any]],
) -> ModelResponse:
tool_calls = [
ToolCall(
id=str(call["id"]),
name=str(call["name"]),
arguments=(
call.get("args") if isinstance(call.get("args"), dict) else {}
),
raw=call.get("raw"),
)
for call in function_calls
]
template = {key: value for key, value in payload.items() if key != "contents"}
contents = list(payload.get("contents") or [])
# Return the model turn exactly as received: Gemini 3 rejects a
# continuation whose functionCall parts lost their thoughtSignature,
# and text or thought parts belong to the same turn.
model_content = copy.deepcopy(self._candidate_content(data))
model_content["role"] = model_content.get("role") or "model"
contents.append(model_content)
return ModelResponse(
provider=self.provider_name,
model=self._model_name(),
tool_calls=tool_calls,
raw=data,
usage=gemini_usage(data),
metadata={
"gemini_payload": template,
"gemini_contents": contents,
},
)
def _function_response_part(
self,
result: ToolResult,
*,
call_id: Optional[str] = None,
) -> Dict[str, Any]:
response: Dict[str, Any] = {"result": result.content}
if result.is_error:
response["is_error"] = True
function_response: Dict[str, Any] = {
"name": result.name,
"response": response,
}
if call_id:
# Echo only ids the API issued; generated ids are Praval-local.
function_response["id"] = call_id
return {"functionResponse": function_response}
def _handle_function_calls(
self,
*,
function_calls: List[Dict[str, Any]],
available_tools: List[Dict[str, Any]],
messages: List[Dict[str, Any]],
request: Optional[ModelRequest],
hitl_context: Optional[Dict[str, Any]],
) -> ModelResponse:
function_response_parts: List[Dict[str, Any]] = []
tool_calls: List[ToolCall] = []
tool_results: List[Dict[str, Any]] = []
for idx, function_call in enumerate(function_calls):
name = function_call["name"]
raw_args = function_call.get("args") or {}
tool_call_id = function_call["id"]
continuation_state = {
"schema": "gemini_tool_v1",
"messages": messages,
"function_calls": function_calls,
"current_index": idx,
"tool_results": list(tool_results),
}
result_content = execute_legacy_tool_call(
hitl_context=hitl_context,
tool_call_id=tool_call_id,
function_name=name,
raw_args=raw_args,
available_tools=available_tools or [],
continuation_state=continuation_state,
)
tool_calls.append(
ToolCall(
id=tool_call_id,
name=name,
arguments=raw_args if isinstance(raw_args, dict) else {},
raw=function_call.get("raw"),
)
)
tool_results.append(
{
"tool_call_id": tool_call_id,
"name": name,
"content": result_content,
}
)
function_response_parts.append(
{
"functionResponse": {
"name": name,
"response": {"result": result_content},
}
}
)
followup_payload = self._build_followup_payload(
messages=messages,
function_calls=function_calls,
function_response_parts=function_response_parts,
available_tools=available_tools,
request=request,
)
try:
# The tools have already run: retry only this request.
data = call_with_retries(
"follow_up",
lambda: self._post_json(
"generateContent",
followup_payload,
timeout=request.timeout if request is not None else None,
),
retries=max_provider_retries(self.config),
map_error=lambda exc: self._provider_error(exc, "Gemini API error"),
)
content = self._extract_text(data)
except Exception:
data = None
content = "\n".join(result["content"] for result in tool_results)
return ModelResponse(
content=content,
provider=self.provider_name,
model=self._model_name(),
tool_calls=tool_calls,
raw=data,
usage=gemini_usage(data),
metadata={"tool_results": tool_results},
)
def _build_followup_payload(
self,
*,
messages: List[Dict[str, Any]],
function_calls: List[Dict[str, Any]],
function_response_parts: List[Dict[str, Any]],
available_tools: List[Dict[str, Any]],
request: Optional[ModelRequest],
) -> Dict[str, Any]:
payload = self._build_payload(messages, available_tools, request=request)
payload["contents"].append(
{
"role": "model",
"parts": [
{
"functionCall": {
"name": call["name"],
"args": call.get("args") or {},
}
}
for call in function_calls
],
}
)
payload["contents"].append({"role": "user", "parts": function_response_parts})
return payload