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 2026-04-20 00:00:50 +05:30 committed by Teknium
parent aa5bd09232
commit 3dea497b20
7 changed files with 1070 additions and 29 deletions

View 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"