feat(providers): route gemini through the native AI Studio API
- add a native Gemini adapter over generateContent/streamGenerateContent - switch the built-in gemini provider off the OpenAI-compatible endpoint - preserve thought signatures and native functionResponse replay - route auxiliary Gemini clients through the same adapter - add focused unit coverage plus native-provider integration checks
This commit is contained in:
kshitijk4poor
212
tests/agent/test_gemini_native_adapter.py
Normal file
212
tests/agent/test_gemini_native_adapter.py
Normal file
@@ -0,0 +1,212 @@
|
||||
"""Tests for the native Google AI Studio Gemini adapter."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
|
||||
class DummyResponse:
|
||||
def __init__(self, status_code=200, payload=None, headers=None, text=None):
|
||||
self.status_code = status_code
|
||||
self._payload = payload or {}
|
||||
self.headers = headers or {}
|
||||
self.text = text if text is not None else json.dumps(self._payload)
|
||||
|
||||
def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
def test_build_native_request_preserves_thought_signature_on_tool_replay():
|
||||
from agent.gemini_native_adapter import build_gemini_request
|
||||
|
||||
request = build_gemini_request(
|
||||
messages=[
|
||||
{"role": "system", "content": "Be helpful."},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
},
|
||||
"extra_content": {
|
||||
"google": {"thought_signature": "sig-123"}
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
],
|
||||
tools=[],
|
||||
tool_choice=None,
|
||||
)
|
||||
|
||||
parts = request["contents"][0]["parts"]
|
||||
assert parts[0]["functionCall"]["name"] == "get_weather"
|
||||
assert parts[0]["thoughtSignature"] == "sig-123"
|
||||
|
||||
|
||||
def test_build_native_request_uses_original_function_name_for_tool_result():
|
||||
from agent.gemini_native_adapter import build_gemini_request
|
||||
|
||||
request = build_gemini_request(
|
||||
messages=[
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Paris"}',
|
||||
},
|
||||
}
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": '{"forecast": "sunny"}',
|
||||
},
|
||||
],
|
||||
tools=[],
|
||||
tool_choice=None,
|
||||
)
|
||||
|
||||
tool_response = request["contents"][1]["parts"][0]["functionResponse"]
|
||||
assert tool_response["name"] == "get_weather"
|
||||
|
||||
|
||||
def test_translate_native_response_surfaces_reasoning_and_tool_calls():
|
||||
from agent.gemini_native_adapter import translate_gemini_response
|
||||
|
||||
payload = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{"thought": True, "text": "thinking..."},
|
||||
{"functionCall": {"name": "search", "args": {"q": "hermes"}}},
|
||||
]
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 10,
|
||||
"candidatesTokenCount": 5,
|
||||
"totalTokenCount": 15,
|
||||
},
|
||||
}
|
||||
|
||||
response = translate_gemini_response(payload, model="gemini-2.5-flash")
|
||||
choice = response.choices[0]
|
||||
assert choice.finish_reason == "tool_calls"
|
||||
assert choice.message.reasoning == "thinking..."
|
||||
assert choice.message.tool_calls[0].function.name == "search"
|
||||
assert json.loads(choice.message.tool_calls[0].function.arguments) == {"q": "hermes"}
|
||||
|
||||
|
||||
def test_native_client_uses_x_goog_api_key_and_native_models_endpoint(monkeypatch):
|
||||
from agent.gemini_native_adapter import GeminiNativeClient
|
||||
|
||||
recorded = {}
|
||||
|
||||
class DummyHTTP:
|
||||
def post(self, url, json=None, headers=None, timeout=None):
|
||||
recorded["url"] = url
|
||||
recorded["json"] = json
|
||||
recorded["headers"] = headers
|
||||
return DummyResponse(
|
||||
payload={
|
||||
"candidates": [
|
||||
{
|
||||
"content": {"parts": [{"text": "hello"}]},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 1,
|
||||
"candidatesTokenCount": 1,
|
||||
"totalTokenCount": 2,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
def close(self):
|
||||
return None
|
||||
|
||||
monkeypatch.setattr("agent.gemini_native_adapter.httpx.Client", lambda *a, **k: DummyHTTP())
|
||||
|
||||
client = GeminiNativeClient(api_key="AIza-test", base_url="https://generativelanguage.googleapis.com/v1beta")
|
||||
response = client.chat.completions.create(
|
||||
model="gemini-2.5-flash",
|
||||
messages=[{"role": "user", "content": "Hello"}],
|
||||
)
|
||||
|
||||
assert recorded["url"] == "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash:generateContent"
|
||||
assert recorded["headers"]["x-goog-api-key"] == "AIza-test"
|
||||
assert "Authorization" not in recorded["headers"]
|
||||
assert response.choices[0].message.content == "hello"
|
||||
|
||||
|
||||
def test_native_http_error_keeps_status_and_retry_after():
|
||||
from agent.gemini_native_adapter import gemini_http_error
|
||||
|
||||
response = DummyResponse(
|
||||
status_code=429,
|
||||
headers={"Retry-After": "17"},
|
||||
payload={
|
||||
"error": {
|
||||
"code": 429,
|
||||
"message": "quota exhausted",
|
||||
"status": "RESOURCE_EXHAUSTED",
|
||||
"details": [
|
||||
{
|
||||
"@type": "type.googleapis.com/google.rpc.ErrorInfo",
|
||||
"reason": "RESOURCE_EXHAUSTED",
|
||||
"metadata": {"service": "generativelanguage.googleapis.com"},
|
||||
}
|
||||
],
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
err = gemini_http_error(response)
|
||||
assert getattr(err, "status_code", None) == 429
|
||||
assert getattr(err, "retry_after", None) == 17.0
|
||||
assert "quota exhausted" in str(err)
|
||||
|
||||
|
||||
def test_stream_event_translation_emits_tool_call_delta_with_stable_index():
|
||||
from agent.gemini_native_adapter import translate_stream_event
|
||||
|
||||
tool_call_indices = {}
|
||||
event = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{"functionCall": {"name": "search", "args": {"q": "abc"}}}
|
||||
]
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
first = translate_stream_event(event, model="gemini-2.5-flash", tool_call_indices=tool_call_indices)
|
||||
second = translate_stream_event(event, model="gemini-2.5-flash", tool_call_indices=tool_call_indices)
|
||||
|
||||
assert first[0].choices[0].delta.tool_calls[0].index == 0
|
||||
assert second[0].choices[0].delta.tool_calls[0].index == 0
|
||||
assert first[0].choices[0].delta.tool_calls[0].id == second[0].choices[0].delta.tool_calls[0].id
|
||||
assert first[-1].choices[0].finish_reason == "tool_calls"
|
||||
@@ -22,7 +22,7 @@ class TestGeminiProviderRegistry:
|
||||
assert pconfig.id == "gemini"
|
||||
assert pconfig.name == "Google AI Studio"
|
||||
assert pconfig.auth_type == "api_key"
|
||||
assert pconfig.inference_base_url == "https://generativelanguage.googleapis.com/v1beta/openai"
|
||||
assert pconfig.inference_base_url == "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
def test_gemini_env_vars(self):
|
||||
pconfig = PROVIDER_REGISTRY["gemini"]
|
||||
@@ -99,7 +99,7 @@ class TestGeminiCredentials:
|
||||
creds = resolve_api_key_provider_credentials("gemini")
|
||||
assert creds["provider"] == "gemini"
|
||||
assert creds["api_key"] == "google-secret"
|
||||
assert creds["base_url"] == "https://generativelanguage.googleapis.com/v1beta/openai"
|
||||
assert creds["base_url"] == "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
def test_resolve_with_gemini_api_key(self, monkeypatch):
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "gemini-secret")
|
||||
@@ -119,7 +119,7 @@ class TestGeminiCredentials:
|
||||
assert result["provider"] == "gemini"
|
||||
assert result["api_mode"] == "chat_completions"
|
||||
assert result["api_key"] == "google-key"
|
||||
assert result["base_url"] == "https://generativelanguage.googleapis.com/v1beta/openai"
|
||||
assert result["base_url"] == "https://generativelanguage.googleapis.com/v1beta"
|
||||
|
||||
|
||||
# ── Model Catalog ──
|
||||
@@ -193,50 +193,45 @@ class TestGeminiAgentInit:
|
||||
importlib.reload(run_agent)
|
||||
|
||||
def test_gemini_agent_uses_chat_completions(self, monkeypatch):
|
||||
"""Gemini falls through to chat_completions — no special elif needed."""
|
||||
"""Gemini still reports chat_completions even though the transport is native."""
|
||||
monkeypatch.setenv("GOOGLE_API_KEY", "test-key")
|
||||
with patch("run_agent.OpenAI") as mock_openai:
|
||||
mock_openai.return_value = MagicMock()
|
||||
with patch("agent.gemini_native_adapter.GeminiNativeClient") as mock_client:
|
||||
mock_client.return_value = MagicMock()
|
||||
from run_agent import AIAgent
|
||||
agent = AIAgent(
|
||||
model="gemini-2.5-flash",
|
||||
provider="gemini",
|
||||
api_key="test-key",
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta",
|
||||
)
|
||||
assert agent.api_mode == "chat_completions"
|
||||
assert agent.provider == "gemini"
|
||||
|
||||
def test_gemini_uses_bearer_auth(self, monkeypatch):
|
||||
"""Gemini OpenAI-compatible endpoint should receive the real API key."""
|
||||
def test_gemini_agent_uses_native_client(self, monkeypatch):
|
||||
monkeypatch.setenv("GOOGLE_API_KEY", "AIzaSy_REAL_KEY")
|
||||
real_key = "AIzaSy_REAL_KEY"
|
||||
with patch("run_agent.OpenAI") as mock_openai:
|
||||
mock_openai.return_value = MagicMock()
|
||||
with patch("agent.gemini_native_adapter.GeminiNativeClient") as mock_client, \
|
||||
patch("run_agent.OpenAI") as mock_openai:
|
||||
mock_client.return_value = MagicMock()
|
||||
from run_agent import AIAgent
|
||||
AIAgent(
|
||||
model="gemini-2.5-flash",
|
||||
provider="gemini",
|
||||
api_key=real_key,
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta/openai",
|
||||
api_key="AIzaSy_REAL_KEY",
|
||||
base_url="https://generativelanguage.googleapis.com/v1beta",
|
||||
)
|
||||
call_kwargs = mock_openai.call_args[1]
|
||||
assert call_kwargs.get("api_key") == real_key
|
||||
headers = call_kwargs.get("default_headers", {})
|
||||
assert "x-goog-api-key" not in headers
|
||||
assert mock_client.called
|
||||
mock_openai.assert_not_called()
|
||||
|
||||
def test_gemini_resolve_provider_client_auth(self, monkeypatch):
|
||||
"""resolve_provider_client('gemini') should pass the real API key through."""
|
||||
def test_gemini_resolve_provider_client_uses_native_client(self, monkeypatch):
|
||||
"""resolve_provider_client('gemini') should build GeminiNativeClient."""
|
||||
monkeypatch.setenv("GEMINI_API_KEY", "AIzaSy_TEST_KEY")
|
||||
real_key = "AIzaSy_TEST_KEY"
|
||||
with patch("agent.auxiliary_client.OpenAI") as mock_openai:
|
||||
mock_openai.return_value = MagicMock()
|
||||
with patch("agent.gemini_native_adapter.GeminiNativeClient") as mock_client, \
|
||||
patch("agent.auxiliary_client.OpenAI") as mock_openai:
|
||||
mock_client.return_value = MagicMock()
|
||||
from agent.auxiliary_client import resolve_provider_client
|
||||
resolve_provider_client("gemini")
|
||||
call_kwargs = mock_openai.call_args[1]
|
||||
assert call_kwargs.get("api_key") == real_key
|
||||
headers = call_kwargs.get("default_headers", {})
|
||||
assert "x-goog-api-key" not in headers
|
||||
assert mock_client.called
|
||||
mock_openai.assert_not_called()
|
||||
|
||||
|
||||
# ── models.dev Integration ──
|
||||
|
||||
Reference in New Issue
Block a user